Add possibility to change distribution of the magnitudes of the slow …

…modulation

icomo/slow_modulation.py

@@ -67,6 +67,7 @@ def priors_for_cps(
     name_durations,
     beta_magnitude=1,
     sigma_magnitude_fix=None,
+    dist_magnitudes=pm.Normal,
     model=None,
 ):
     """Create priors for changepoints.
@@ -80,21 +81,25 @@ def priors_for_cps(
     Parameters
     ----------
     cp_dim : str
-        dimension of the :class:`pymc.Model` for the changepoints. Define it by passing
+        Dimension of the :class:`pymc.Model` for the changepoints. Define it by passing
         `coords={cp_dim: np.arange(num_cps)}` to :class:`pymc.Model` at creation. The length of this
         dimension determines the number of changepoints.
     time_dim : str
-        dimension of the :class:`pymc.Model` for the time.
+        Dimension of the :class:`pymc.Model` for the time.
     name_positions : str
-        name under which the positions of the changepoints are stored in :class:`pymc.Model`
+        Name under which the positions of the changepoints are stored in :class:`pymc.Model`
     name_magnitudes : str
-        name under which the magnitudes of the changepoints are stored in :class:`pymc.Model`
+        Name under which the magnitudes of the changepoints are stored in :class:`pymc.Model`
     name_durations : str
-        name under which the durations of the changepoints are stored in :class:`pymc.Model`
+        Name under which the durations of the changepoints are stored in :class:`pymc.Model`
     beta_magnitude : float, default=1
-        beta parameter of the hierarchical prior for the magnitudes
+        Beta parameter of the hierarchical prior for the magnitudes
     sigma_magnitude_fix : float, default=None
-        if not `None`, the standard deviation from which the magnitudes are sampled is fixed
+        If not `None`, the standard deviation from which the magnitudes are sampled is fixed
+    dist_magnitudes : :class:`pymc.Distribution`, default=pm.Normal
+        Distribution from which the magnitudes are sampled. Can for example be
+        functools.partial(pm.StudentT, nu=4) to sample from a StudentT distribution for
+        a more robust model.
     model : :class:`pymc.Model`, default=None
         pm.Model in which the priors are created. If None, the pm.Model is taken from the
         the context.
@@ -123,6 +128,7 @@ def priors_for_cps(
         dims=(cp_dim,),
         beta=beta_magnitude,
         fix_hyper_sigma=sigma_magnitude_fix,
+        dist_values=dist_magnitudes,
     )
 
     ### Durations:

icomo/tools.py

@@ -3,7 +3,9 @@
 import pymc as pm
 
 
-def hierarchical_priors(name, dims, beta=1, fix_hyper_sigma=None):
+def hierarchical_priors(
+    name, dims, beta=1, fix_hyper_sigma=None, dist_values=pm.Normal
+):
     """Create hierarchical priors for a variable.
 
     Create an n-dimensional variable with hierarchical prior with name `name` and
@@ -13,14 +15,18 @@ def hierarchical_priors(name, dims, beta=1, fix_hyper_sigma=None):
     Parameters
     ----------
     name : str
-        name under which the variable is stored in pm.Model
+        Name under which the variable is stored in pm.Model
     dims : tuple of str
-        dimensions over which the variable is defined. Define a dimension by passing
+        Dimensions over which the variable is defined. Define a dimension by passing
         coords={dim_name: np.arange(size)} to pm.Model at creation.
     beta : float, default=1
-        beta parameter of the half-cauchy distribution
+        Beta parameter of the half-cauchy distribution
     fix_hyper_sigma : float, default=None
-        if not None, the standard deviation from which the variable is sampled is fixed
+        If not None, the standard deviation from which the variable is sampled is fixed
+    dist_values : :class:`pymc.Distribution`, default=pm.Normal
+        Distribution from which the values are sampled. Can for example be
+        functools.partial(pm.StudentT, nu=4) to sample from a StudentT distribution for
+        a more robust model.
 
     Returns
     -------
@@ -33,6 +39,6 @@ def hierarchical_priors(name, dims, beta=1, fix_hyper_sigma=None):
         if fix_hyper_sigma is None
         else fix_hyper_sigma
     )
-    values = (pm.Normal(f"{name}_raw", 0, 1, dims=dims)) * sigma
+    values = (dist_values(f"{name}_raw", mu=0, sigma=1, dims=dims)) * sigma
     values = pm.Deterministic(f"{name}", values, dims=dims)
     return values