Change parameterization of LogNormal and fix parsing bug

Author: 4ment

torchtree/cli/evolution.py

@@ -1458,19 +1458,19 @@ def create_ucln_prior(branch_model_id):
     mean = Parameter.json_factory(
         f"{branch_model_id}.rates.prior.mean", **{"tensor": [0.001]}
     )
-    scale = Parameter.json_factory(
-        f"{branch_model_id}.rates.prior.scale", **{"tensor": [1.0]}
+    stdev = Parameter.json_factory(
+        f"{branch_model_id}.rates.prior.stdev", **{"tensor": [0.1]}
     )
     mean[CONSTRAINT.LOWER.value] = 0.0
-    scale[CONSTRAINT.LOWER.value] = 0.0
+    stdev[CONSTRAINT.LOWER.value] = 0.0
     joint_list.append(
         Distribution.json_factory(
             f"{branch_model_id}.rates.prior",
             "LogNormal",
             f"{branch_model_id}.rates",
             {
                 "mean": mean,
-                "scale": scale,
+                "stdev": stdev,
             },
         )
     )
@@ -1485,7 +1485,7 @@ def create_ucln_prior(branch_model_id):
         Distribution.json_factory(
             f"{branch_model_id}.rates.scale.prior",
             "torch.distributions.Gamma",
-            f"{branch_model_id}.rates.prior.scale",
+            f"{branch_model_id}.rates.prior.stdev",
             {
                 "concentration": 0.5396,
                 "rate": 2.6184,

torchtree/core/utils.py

@@ -181,9 +181,9 @@ def process_object(data, dic):
                 raise JSONParseError("Missing `id' and `type' keys") from None
 
         if id_ in dic:
-            raise JSONParseError("Object with ID `{id_}' already exists")
+            raise JSONParseError(f"Object with ID `{id_}' already exists")
         if "type" not in data:
-            raise JSONParseError("Object with ID `{id_}' does not have a type")
+            raise JSONParseError(f"Object with ID `{id_}' does not have a type")
 
         try:
             klass = get_class(data["type"])

torchtree/distributions/log_normal.py

@@ -18,26 +18,26 @@ class LogNormal(torch.distributions.LogNormal):
     :attr:`scale`.
 
     :param mean: mean of the distribution
-    :param scale: standard deviation of log of the distribution
-    :param scale_real: standard deviation of the distribution
+    :param scale: scale (sigma) parameter of log of the distribution
+    :param stdev: standard deviation of the distribution
     """
 
     def __init__(
         self,
         mean: Union[Tensor, float],
         scale: Union[Tensor, float, None] = None,
-        scale_real: Union[Tensor, float, None] = None,
+        stdev: Union[Tensor, float, None] = None,
         validate_args=None,
     ) -> None:
-        if (scale is not None) + (scale_real is not None) != 1:
-            raise ValueError("Exactly one of scale or scale_real may be specified.")
+        if (scale is not None) + (stdev is not None) != 1:
+            raise ValueError("Exactly one of scale or stdev may be specified.")
 
         if scale is not None:
             log_mean = log(mean) if isinstance(mean, Number) else mean.log()
             var = scale * scale if isinstance(scale, Number) else scale.pow(2)
             loc = log_mean - var * 0.5
         else:
-            temp = 1.0 + scale_real * scale_real / (mean * mean)
+            temp = 1.0 + stdev * stdev / (mean * mean)
             if isinstance(temp, Number):
                 loc = math.log(mean / math.sqrt(temp))
                 scale = math.sqrt(math.log(temp))