add other normalization option

thoglu · Jan 23, 2024 · 4ca7aec · 4ca7aec
1 parent 601209c
commit 4ca7aec
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 1 deletion.
diff --git a/jammy_flows/flow_options.py b/jammy_flows/flow_options.py
@@ -187,7 +187,7 @@
 opts_dict["f"]["kwargs"]["vertical_restrict_max_min_width_height_ratio"] = (-1.0, lambda x: (x==-1.0) or (x>0.0))
 opts_dict["f"]["kwargs"]["vertical_fix_boundary_derivative"] = (1, lambda x: [0,1])
 opts_dict["f"]["kwargs"]["min_kappa"] = (1e-10, lambda x: x>0)
-opts_dict["f"]["kwargs"]["kappa_prediction"] = ("direct_log_real_bounded", ["direct_log_real_bounded", "log_bounded"])
+opts_dict["f"]["kwargs"]["kappa_prediction"] = ("direct_log_real_bounded", ["direct_log_real_bounded", "softplus_real_bounded", "log_bounded"])
 opts_dict["f"]["kwargs"]["add_extra_rotation_inbetween"] = (0, [0,1])
 
 """

diff --git a/jammy_flows/layers/spheres/fvm_2d.py b/jammy_flows/layers/spheres/fvm_2d.py
@@ -97,6 +97,8 @@ def __init__(self,
 
         if(kappa_prediction=="direct_log_real_bounded"):
             self.kappa_fn=lambda x: x.exp()+self.min_kappa
+        elif(kappa_prediction=="softplus_real_bounded"):
+            self.kappa_fn=lambda x: F.softplus(x)+self.min_kappa
         else:
             log_min_kappa=numpy.log(self.min_kappa)