Merge pull request #52 from minaskar/dev

periodic/reflective parameters

docs/source/flow.ipynb

@@ -33,7 +33,7 @@
     "\n",
     "2) **Neural Spline Flows (NSF):** Neural Spline Flow (NSF) extends the concept of normalizing flows by using neural networks to parameterize piecewise monotonic rational quadratic splines as the invertible transformations. These splines provide a flexible way to model complex distributions while ensuring smooth and differentiable transformations. NSF combines the expressive power of neural networks with the efficiency of spline-based transformations, allowing for efficient sampling and exact likelihood computation. This makes NSF particularly effective for modeling high-dimensional data with complex, multimodal distributions, enhancing the flexibility and accuracy of normalizing flow-based generative models.\n",
     "\n",
-    "The predefined MAF and NSF flows are ``'maf3'``, ``'maf6'``, ``'maf12'``, ``'nsf3'``, ``'nsf6'``, and ``'nsf12'``. By default, ``pocoMC`` uses the ``'nsf3'`` flow, meaning a Neural Spline Flow with 3 transformations. This balances flexibility and computational cost. The user can change the flow by setting the ``flow`` parameter in the ``pocoMC`` ``Sampler`` class as follows:"
+    "The predefined MAF and NSF flows are ``'maf3'``, ``'maf6'``, ``'maf12'``, ``'nsf3'``, ``'nsf6'``, and ``'nsf12'``. By default, ``pocoMC`` uses the ``'nsf6'`` flow, meaning a Neural Spline Flow with 6 transformations. This balances flexibility and computational cost. The user can change the flow by setting the ``flow`` parameter in the ``pocoMC`` ``Sampler`` class as follows:"
    ]
   },
   {

docs/source/index.rst

@@ -114,6 +114,11 @@ Copyright 2022-2024 Minas Karamanis and contributors.
 Changelog
 =========
 
+**1.2.3 (27/08/24)**
+
+- Added support for periodic and reflective parameters.
+- Changed default normalizing flow to ``nsf6``.
+
 **1.2.2 (20/06/24)**
 
 - Fixed bug in ``posterior`` method related to blobs.

docs/source/priors.rst

@@ -74,4 +74,45 @@ This would look like::
                 c = np.random.normal(loc=m, scale=s, size=size)
                 return np.array([m, s, c]).T
 
-        prior = CustomPrior()
+        prior = CustomPrior()
+
+
+Boundary conditions
+-------------------
+
+By default, ``pocoMC`` assumes that all parameters specified in the prior have hard bounds. In other words, each
+parameter is free to vary in a prespecified range. Anytime a value is proposed by ``pocoMC`` that lies outside of
+this range, it is automatically rejected. This is the desired behavior for most problems, since individual parameters 
+are often either defined everywhere (i.e. from negative infinity to infinity) or over a finite range (e.g., from -1 to 
++1).
+
+However, there are problems in which specific parameters may behave differently. ``pocoMC`` supports two such cases:
+
+- **Periodic boundary conditions**. In this case, ``pocoMC`` assumes that the parameter is periodic. For example, 
+  if the parameter is on the interval ``[0, 2*np.pi]``, then the parameter can be wrapped around to the other side
+  of the interval. This can be useful for phase parameters that might be periodic e.g. on a range ``[0,2*np.pi]``.
+- **Reflective boundary conditions**. In this case, ``pocoMC`` assumes that the parameter is reflective. For example,
+  if the parameter is on the interval ``[0, 1]``, then the parameter can be flipped around to the other side of the
+  interval. This can be useful for parameters that are ratios where ``a/b`` and  ``b/a`` are equivalent.
+
+Given the above, it is possible to set the ``periodic`` and ``reflective`` attributes of the prior. For example, in 
+a five-parameter model, if we want the first two parameters to be periodic, and the third and fourth to be reflective, 
+we would do::
+
+    from scipy.stats import uniform, norm
+
+    prior = pc.Prior([
+        uniform(loc=0.0, scale=2*np.pi), # this parameter is periodic
+        uniform(loc=0.0, scale=2*np.pi), # this parameter is periodic
+        uniform(loc=0.0, scale=1.0), # this parameter is reflective
+        uniform(loc=0.0, scale=1.0), # this parameter is reflective
+        norm(loc=0.0, scale=3.0), # this parameter is neither periodic nor reflective
+    ])
+
+    sampler = pc.Sampler(prior, 
+                        loglike, 
+                        periodic=[0,1], 
+                        reflective=[2,3])
+
+As you can see, nothing changes in the definition of the prior. Instead, we just need to provide the indices of the
+parameters that should be periodic and reflective to the sampler.

pocomc/_version.py

@@ -1 +1 @@
-version = "1.2.2"
+version = "1.2.3"

pocomc/mcmc.py

@@ -90,6 +90,12 @@ def preconditioned_pcn(state_dict: dict,
         # Transform to x space
         x_prime, logdetj_prime = scaler.inverse(u_prime)
 
+        # Apply boundary conditions
+        if (scaler.periodic is not None) or (scaler.reflective is not None):
+            x_prime = scaler.apply_boundary_conditions_x(x_prime)
+            u_prime = scaler.forward(x_prime)
+            x_prime, logdetj_prime = scaler.inverse(u_prime)
+
         # Compute finite mask
         finite_mask_logdetj_prime = np.isfinite(logdetj_prime)
         finite_mask_x_prime = np.isfinite(x_prime).all(axis=1)
@@ -223,7 +229,6 @@ def preconditioned_rwm(state_dict: dict,
     # Get number of particles and parameters/dimensions
     n_walkers, n_dim = x.shape
 
-
     cov = geometry.normal_cov
     chol = np.linalg.cholesky(cov)
 
@@ -248,6 +253,12 @@ def preconditioned_rwm(state_dict: dict,
         # Transform to x space
         x_prime, logdetj_prime = scaler.inverse(u_prime)
 
+        # Apply boundary conditions
+        if (scaler.periodic is not None) or (scaler.reflective is not None):
+            x_prime = scaler.apply_boundary_conditions_x(x_prime)
+            u_prime = scaler.forward(x_prime)
+            x_prime, logdetj_prime = scaler.inverse(u_prime)
+
         # Compute finite mask
         finite_mask_logdetj_prime = np.isfinite(logdetj_prime)
         finite_mask_x_prime = np.isfinite(x_prime).all(axis=1)
@@ -393,11 +404,17 @@ def pcn(state_dict: dict,
         # Propose new points in u space
         u_prime = np.empty((n_walkers, n_dim))
         for k in range(n_walkers):
-            u_prime[k] = mu + (1.0 - sigma ** 2.0) ** 0.5 * diff[k] + sigma * np.sqrt(s[k]) * np.dot(chol_cov, np.random.randn(n_dim))        
+            u_prime[k] = mu + (1.0 - sigma ** 2.0) ** 0.5 * diff[k] + sigma * np.sqrt(s[k]) * np.dot(chol_cov, np.random.randn(n_dim)) 
 
         # Transform to x space
         x_prime, logdetj_prime = scaler.inverse(u_prime)
 
+        # Apply boundary conditions
+        if (scaler.periodic is not None) or (scaler.reflective is not None):
+            x_prime = scaler.apply_boundary_conditions_x(x_prime)
+            u_prime = scaler.forward(x_prime)
+            x_prime, logdetj_prime = scaler.inverse(u_prime)
+
         # Compute finite mask
         finite_mask_logdetj_prime = np.isfinite(logdetj_prime)
         finite_mask_x_prime = np.isfinite(x_prime).all(axis=1)
@@ -541,6 +558,12 @@ def rwm(state_dict: dict,
         # Transform to x space
         x_prime, logdetj_prime = scaler.inverse(u_prime)
 
+        # Apply boundary conditions
+        if (scaler.periodic is not None) or (scaler.reflective is not None):
+            x_prime = scaler.apply_boundary_conditions_x(x_prime)
+            u_prime = scaler.forward(x_prime)
+            x_prime, logdetj_prime = scaler.inverse(u_prime)
+
         # Compute finite mask
         finite_mask_logdetj_prime = np.isfinite(logdetj_prime)
         finite_mask_x_prime = np.isfinite(x_prime).all(axis=1)

pocomc/sampler.py

@@ -59,6 +59,17 @@ class Sampler:
         returned by the likelihood function (e.g., chi-squared values, residuals, etc.). Blobs are stored as a
         structured array with named fields when the data type is provided. Currently, the blobs feature is not
         compatible with vectorized likelihood calculations.
+    periodic : list or ``None``
+        List of parameter indeces that should be wrapped around the domain (default is ``periodic=None``).
+        This can be useful for phase parameters that might be periodic e.g. on a range ``[0,2*np.pi]``. For example,
+        ``periodic=[0,1]`` will wrap around the first and second parameters.
+    reflective : list or ``None``
+        List of parameter indeces that should be reflected around the domain (default is ``reflective=None``).
+        This can arise in cases where parameters are ratios where ``a/b`` and  ``b/a`` are equivalent. For example,
+        ``reflective=[0,1]`` will reflect the first and second parameters.
+    transform : str
+        Type of transformation to apply to bounded parameters (default is ``transform='probit'``). Available options
+        are ``'probit'`` and ``'logit'``. See ``Reparameterize`` for more details.
     pool : pool or int
         Number of processes to use for parallelisation (default is ``pool=None``). If ``pool`` is an integer
         greater than 1, a ``multiprocessing`` pool is created with the specified number of processes (e.g., ``pool=8``). 
@@ -70,7 +81,7 @@ class Sampler:
         Maximum number of threads to use for torch. If ``None`` torch uses all
         available threads while training the normalizing flow (default is ``pytorch_threads=1``). 
     flow : ``zuko.flow.Flow`` or str
-        Normalizing flow to use for preconditioning (default is ``flow='nsf3'``). Available options are
+        Normalizing flow to use for preconditioning (default is ``flow='nsf6'``). Available options are
         ``'nsf3'``, ``'nsf6'``, ``'nsf12'``, ``'maf3'``, ``'maf6'``, and ``'maf12'``. 'nsf' stands for
         Neural Spline Flows and 'maf' stands for Masked Autoregressive Flows. The number indicates the
         number of transformations in the flow. More transformations lead to more flexibility but also
@@ -99,7 +110,7 @@ class Sampler:
         are not multimodal or have strong non-linear correlations between parameters.
     dynamic : bool
         If True, dynamically adjust the effective sample size (ESS) threshold based on the
-        number of unique particles (default is ``dynamic=False``). This can be useful for
+        number of unique particles (default is ``dynamic=True``). This can be useful for
         targets with a large number of modes or strong non-linear correlations between parameters.
     metric : str
         Metric used for determining the next temperature (``beta``) level (default is ``metric="ess"``).
@@ -150,9 +161,12 @@ def __init__(self,
                  likelihood_kwargs: dict = None,
                  vectorize: bool = False,
                  blobs_dtype: str = None,
+                 periodic: list = None,
+                 reflective: list = None,
+                 transform: str = "probit",
                  pool=None,
                  pytorch_threads=1,
-                 flow='nsf3',
+                 flow='nsf6',
                  train_config: dict = None,
                  train_frequency: int = None,
                  precondition: bool = True,
@@ -295,7 +309,13 @@ def __init__(self,
         self.flow_untrained = True
 
         # Scaler
-        self.scaler = Reparameterize(self.n_dim, bounds=self.bounds)
+        if transform not in ['probit', 'logit']:
+            raise ValueError(f"Invalid transform {transform}. Options are 'probit' or 'logit'.")
+        self.scaler = Reparameterize(self.n_dim, 
+                                     bounds=self.bounds, 
+                                     periodic=periodic, 
+                                     reflective=reflective,
+                                     transform=transform,)
 
         # Output
         if output_dir is None:

pocomc/scaler.py

@@ -18,7 +18,7 @@ class Reparameterize:
     periodic : ``list``
         List of indices corresponding to parameters with periodic boundary conditions
     reflective : ``list``
-        List of indices corresponding to parameters with periodic boundary conditions
+        List of indices corresponding to parameters with reflective boundary conditions
     transform : ``str``
         Type of transform to use for bounded parameters. Options are ``"probit"``
         (default) and ``"logit"``.
@@ -81,9 +81,9 @@ def __init__(self,
 
         self._create_masks()
 
-    def apply_boundary_conditions(self, x: np.ndarray):
+    def apply_boundary_conditions_x(self, x: np.ndarray):
         """
-        Apply boundary conditions (i.e. periodic or reflective) to input.
+        Apply boundary conditions (i.e. periodic or reflective) to input ``x``.
         The first kind include phase parameters that might be periodic
         e.g. on a range ``[0,2*np.pi]``. The latter can arise in cases
         where parameters are ratios where ``a/b`` and  ``b/a`` are equivalent.
@@ -100,15 +100,15 @@ def apply_boundary_conditions(self, x: np.ndarray):
         if (self.periodic is None) and (self.reflective is None):
             return x
         elif self.periodic is None:
-            return self._apply_reflective_boundary_conditions(x)
+            return self._apply_reflective_boundary_conditions_x(x)
         elif self.reflective is None:
-            return self._apply_periodic_boundary_conditions(x)
+            return self._apply_periodic_boundary_conditions_x(x)
         else:
-            return self._apply_reflective_boundary_conditions(self._apply_periodic_boundary_conditions(x))
+            return self._apply_reflective_boundary_conditions_x(self._apply_periodic_boundary_conditions_x(x))
 
-    def _apply_periodic_boundary_conditions(self, x: np.ndarray):
+    def _apply_periodic_boundary_conditions_x(self, x: np.ndarray):
         """
-        Apply periodic boundary conditions to input.
+        Apply periodic boundary conditions to input ``x``.
         This can be useful for phase parameters that might be periodic
         e.g. on a range ``[0,2*np.pi]``
         
@@ -131,9 +131,9 @@ def _apply_periodic_boundary_conditions(self, x: np.ndarray):
                         x[j, i] = self.high[i] + x[j, i] - self.low[i]
         return x
 
-    def _apply_reflective_boundary_conditions(self, x: np.ndarray):
+    def _apply_reflective_boundary_conditions_x(self, x: np.ndarray):
         """
-        Apply reflective boundary conditions to input. This can arise in cases
+        Apply reflective boundary conditions to input ``x``. This can arise in cases
         where parameters are ratios where ``a/b`` and  ``b/a`` are equivalent.
         
         Parameters