Experimented with pgnet model with no DNA

Author: mohanwu

examples/pgnet.ipynb

@@ -263,3 +263,17 @@ def pf_step(self, key, x_prev, y_curr, theta):
                 jnp.log(y_curr), jnp.eye(4)[:-1, :], jnp.diag(omega)
             )
         return x_curr, logw
+
+    def prop_lpdf(self, x_curr, x_prev, y_curr, theta):
+        r"""
+        Calculate the log-density of the proposal distribution `q(x_curr) = q(x_curr | x_prev, y_curr, theta)`.
+
+        In this case we have a bootstrap filter, so `q(x_curr) = p(x_curr | x_prev, theta)`.
+
+        Args:
+            x_curr: State variable at current time `t`.
+            x_prev: State variable at previous time `t-1`.
+            y_curr: Measurement variable at current time `t`.
+            theta: Parameter value.
+        """
+        return super().state_lpdf(x_curr=x_curr, x_prev=x_prev, theta=theta)

src/pfjax/experimental/models/pgnet_model_no_DNA.py

@@ -222,4 +222,18 @@ def is_valid(self, x, theta):
         Returns:
             Whether or not `x>=0`.
         """
-        return (x >= 0) & (x[3] <= self._K)
+        return jnp.alltrue(x >= 0) & (x[3] <= self._K)
+
+    def prop_lpdf(self, x_curr, x_prev, y_curr, theta):
+        r"""
+        Calculate the log-density of the proposal distribution `q(x_curr) = q(x_curr | x_prev, y_curr, theta)`.
+
+        In this case we have a bootstrap filter, so `q(x_curr) = p(x_curr | x_prev, theta)`.
+
+        Args:
+            x_curr: State variable at current time `t`.
+            x_prev: State variable at previous time `t-1`.
+            y_curr: Measurement variable at current time `t`.
+            theta: Parameter value.
+        """
+        return super().state_lpdf(x_curr=x_curr, x_prev=x_prev, theta=theta)

src/pfjax/experimental/models/pgnet_model_reg_no_DNA.py

@@ -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):