Added lotvol_model on regular scale

Author: mohanwu

examples/lotvol.ipynb

@@ -0,0 +1,227 @@
+"""
+Lotka-Volterra predator-prey model on the regular scale.
+
+The model is:
+
+```
+x_m0 \propto 1
+H_mt ~ N(H_{m,t-1} + (alpha H_{m,t-1} - beta H_{m,t-1}L_{m,t-1}) dt/m,
+            sigma_H^2 dt/m)
+L_mt ~ N(L_{m,t-1} + (-gamma L_{m,t-1} + delta H_{m,t-1}L_{m,t-1}) dt/m,
+            sigma_L^2 dt/m)
+y_t ~ N(x_{m,mt}, diag(tau_H^2, tau_L^2) )
+```
+
+- Model parameters: `theta = (alpha, beta, gamma, delta, sigma_H, sigma_L, tau_H, tau_L)`.
+- Global constants: `dt` and `n_res`, i.e., `m`.
+- State dimensions: `n_state = (n_res, 2)`.
+- Measurement dimensions: `n_meas = 2`.
+
+**Notes:**
+
+- The measurement `y_t` corresponds to `x_t = (x_{m,(t-1)m+1}, ..., x_{m,tm})`, i.e., aligns with the last element of `x_t`.
+- The prior is such that `p(x_0 | y_0, theta)` is given by:
+
+    ```
+    x_{m,n} = 0 for n = -m+1, ..., -1,
+    x_{m0} ~ TruncatedNormal( y_0, diag(tau_H^2, tau_L^2) ),
+    ```
+
+    where
+
+    ```
+    z ~ TruncatedNormal(mu, diag(sigma^2)) <=>
+    z = mu + diag(sigma) Z_0,   Z_0 ~iid N(0,1) truncated at -mu.
+    ```
+
+"""
+
+import jax
+import jax.numpy as jnp
+import jax.scipy as jsp
+from jax import random
+from jax import lax
+from pfjax import sde as sde
+
+# --- helper functions ---------------------------------------------------------
+
+
+def lotvol_drift(x, dt, theta):
+    r"""
+    Calculates the SDE drift function.
+    """
+    alpha = theta[0]
+    beta = theta[1]
+    gamma = theta[2]
+    delta = theta[3]
+    return x + jnp.array([alpha * x[0] - beta * x[0] * x[1],
+                          -gamma * x[1] + delta * x[0] * x[1]]) * dt
+
+
+# --- main functions -----------------------------------------------------------
+
+class RegLotVolModel(sde.SDEModel):
+    def __init__(self, dt, n_res):
+        r"""
+        Class constructor for the Lotka-Volterra model.
+
+        Args:
+            dt: SDE interobservation time.
+            n_res: SDE resolution number.  There are `n_res` latent variables per observation, equally spaced with interobservation time `dt/n_res`.
+        """
+        # creates "private" variables self._dt and self._n_res
+        super().__init__(dt, n_res, diff_diag=True)
+        # self.dt = dt
+        # self.n_res = n_res
+        # the following variable is mainly used for testing, i.e.,
+        # in the `_for` versions of certain methods.
+        # it does contain the number of SDE dimensions, which is used
+        # outside of testing, but can be circumvented by pulling shape
+        # from input arguments.  however, this may fail less informatively
+        # than if using prespecified SDE dimensions...
+        self._n_state = (self._n_res, 2)
+
+    def drift(self, x, theta):
+        r"""
+        Calculates the SDE drift function.
+        """
+        alpha = theta[0]
+        beta = theta[1]
+        gamma = theta[2]
+        delta = theta[3]
+        return jnp.array([alpha * x[0] - beta * x[0] * x[1],
+                          -gamma * x[1] + delta * x[0] * x[1]])
+
+    def diff(self, x, theta):
+        r"""
+        Calculates the SDE diffusion function.
+        """
+        return theta[4:6]
+
+    def state_lpdf_for(self, x_curr, x_prev, theta):
+        r"""
+        Calculates the log-density of `p(x_curr | x_prev, theta)`.
+
+        For-loop version for testing.
+
+        Args:
+            x_curr: State variable at current time `t`.
+            x_prev: State variable at previous time `t-1`.
+            theta: Parameter value.
+        Returns:
+            The log-density of `p(x_curr | x_prev, theta)`.
+        """
+        dt_res = self._dt/self._n_res
+        x0 = jnp.append(jnp.expand_dims(
+            x_prev[self._n_res-1], axis=0), x_curr[:self._n_res-1], axis=0)
+        x1 = x_curr
+        sigma = theta[4:6] * jnp.sqrt(dt_res)
+        lp = jnp.array(0.0)
+        for t in range(self._n_res):
+            lp = lp + jnp.sum(jsp.stats.norm.logpdf(
+                x1[t],
+                loc=lotvol_drift(x0[t], dt_res, theta),
+                scale=sigma
+            ))
+        return lp
+
+    def state_sample_for(self, key, x_prev, theta):
+        r"""
+        Samples from `x_curr ~ p(x_curr | x_prev, theta)`.
+
+        For-loop version for testing.
+
+        Args:
+            key: PRNG key.
+            x_prev: State variable at previous time `t-1`.
+            theta: Parameter value.
+
+        Returns:
+            Sample of the state variable at current time `t`: `x_curr ~ p(x_curr | x_prev, theta)`.
+        """
+        dt_res = self._dt/self._n_res
+        sigma = theta[4:6] * jnp.sqrt(dt_res)
+        x_curr = jnp.zeros(self._n_state)
+        x_state = x_prev[self._n_res-1]
+        for t in range(self._n_res):
+            key, subkey = random.split(key)
+            x_state = lotvol_drift(x_state, dt_res, theta) + \
+                random.normal(subkey, (self._n_state[1],)) * sigma
+            x_curr = x_curr.at[t].set(x_state)
+        return x_curr
+
+    def meas_lpdf(self, y_curr, x_curr, theta):
+        r"""
+        Log-density of `p(y_curr | x_curr, theta)`.
+
+        Args:
+            y_curr: Measurement variable at current time `t`.
+            x_curr: State variable at current time `t`.
+            theta: Parameter value.
+
+        Returns
+            The log-density of `p(y_curr | x_curr, theta)`.
+        """
+        tau = theta[6:8]
+        return jnp.sum(
+            jsp.stats.norm.logpdf(y_curr,
+                                  loc=x_curr[-1], scale=tau)
+        )
+
+    def meas_sample(self, key, x_curr, theta):
+        r"""
+        Sample from `p(y_curr | x_curr, theta)`.
+
+        Args:
+            key: PRNG key.
+            x_curr: State variable at current time `t`.
+            theta: Parameter value.
+
+        Returns:
+            Sample of the measurement variable at current time `t`: `y_curr ~ p(y_curr | x_curr, theta)`.
+        """
+        tau = theta[6:8]
+        return x_curr[-1] + \
+            tau * random.normal(key, (self._n_state[1],))
+
+    def pf_init(self, key, y_init, theta):
+        r"""
+        Importance sampler for `x_init`.  
+
+        See file comments for exact sampling distribution of `p(x_init | y_init, theta)`, i.e., we have a "perfect" importance sampler with `logw = CONST(theta)`.
+
+        Args:
+            key: PRNG key.
+            y_init: Measurement variable at initial time `t = 0`.
+            theta: Parameter value.
+
+        Returns:
+            - x_init: A sample from the proposal distribution for `x_init`.
+            - logw: The log-weight of `x_init`.
+        """
+        tau = theta[6:8]
+        key, subkey = random.split(key)
+        x_init = y_init + tau * random.truncated_normal(
+            subkey,
+            lower=-y_init/tau,
+            upper=jnp.inf,
+            shape=(self._n_state[1],)
+        )
+        logw = jnp.sum(jsp.stats.norm.logcdf(y_init/tau))
+        return \
+            jnp.append(jnp.zeros((self._n_res-1,) + x_init.shape),
+                       jnp.expand_dims(x_init, axis=0), axis=0), \
+            logw
+    
+    def is_valid(self, x, theta):
+        """
+        Checks whether SDE observations are valid.
+
+        Args:
+            x: SDE variables.  A vector of size `n_dims`.
+            theta: Parameter value.
+
+        Returns:
+            Whether or not `x>=0`.
+        """
+        return x >= 0

examples/pgnet.ipynb

@@ -198,7 +198,8 @@ def is_valid_state(self, x, theta):
         """
         valid_x = jax.vmap(self.is_valid, in_axes=(0, None))(x, theta)
         nan_x = jnp.any(jnp.isnan(x), axis=1)
-        return jnp.alltrue(valid_x, where=~nan_x) and jnp.alltrue(~nan_x)
+        return jnp.alltrue(valid_x, where=~nan_x) & jnp.alltrue(~nan_x)
+        # return jnp.alltrue(valid_x) and jnp.alltrue(~nan_x)
 
     def state_lpdf(self, x_curr, x_prev, theta):
         """
@@ -353,9 +354,8 @@ def pf_step(self, key, x_prev, y_curr, theta):
         x_curr = self.state_sample(key, x_prev, theta)
         logw = lax.cond(
             self.is_valid_state(x_curr, theta),
-            lambda _x: self.meas_lpdf(y_curr, x_curr, theta),
-            lambda _x: -jnp.inf,
-            0.0
+            lambda: self.meas_lpdf(y_curr, x_curr, theta),
+            lambda: -jnp.inf
         )
         # logw = self.meas_lpdf(y_curr, x_curr, theta)
         return x_curr, logw
@@ -436,12 +436,12 @@ def scan_fun(carry, n):
             theta=theta
         )
         logw = logw + self.meas_lpdf(y_curr, x_prop, theta) - last["lp"]
-        logw = lax.cond(
-            self.is_valid_state(x_prop, theta),
-            lambda _x: logw,
-            lambda _x: -jnp.inf,
-            0.0
-        )
+        # logw = lax.cond(
+        #     self.is_valid_state(x_prop, theta),
+        #     lambda _x: logw,
+        #     lambda _x: -jnp.inf,
+        #     0.0
+        # )
         return x_prop, logw
 
     def bridge_prop_for(self, key, x_prev, y_curr, theta, Y, A, Omega):