variable length unfold model

Author: georgematheos

b3d/chisight/shared/particle_system.py

@@ -8,7 +8,7 @@
 from b3d.chisight.dense.dense_likelihood import make_dense_observation_model, DenseImageLikelihoodArgs
 from b3d import Pose, Mesh
 from b3d.chisight.sparse.gps_utils import add_dummy_var
-from b3d.chisight.sparse.pose_utils import uniform_pose_in_ball
+from b3d.pose.pose_utils import uniform_pose_in_ball
 dummy_mapped_uniform_pose = add_dummy_var(uniform_pose_in_ball).vmap(in_axes=(0,None,None,None))
 
 
@@ -113,7 +113,8 @@ def particle_system_state_step(carried_state, _):
 
 @gen
 def latent_particle_model(
-        num_timesteps, # const object
+        max_num_timesteps, # const object
+        num_timesteps,
         num_particles, # const object
         num_clusters, # const object
         relative_particle_poses_prior_params,
@@ -132,32 +133,49 @@ def latent_particle_model(
         camera_pose_prior_params
     ) @ "state0"
 
-    final_state, scan_retvals = particle_system_state_step.scan(n=(num_timesteps.const - 1))(state0, None) @ "states1+"
+    masked_final_state, masked_scan_retvals = b3d.modeling_utils.masked_scan_combinator(
+        particle_system_state_step,
+        n=(max_num_timesteps.const-1)
+    )(
+        state0,
+        genjax.Mask(
+            # This next line tells the scan combinator how many timesteps to run
+            jnp.arange(max_num_timesteps.const - 1) < num_timesteps - 1,
+            jnp.zeros(max_num_timesteps.const - 1)
+        )
+    ) @ "states1+"
+
 
     # concatenate each element of init_retval, scan_retvals
-    return jax.tree.map(
+    concatenated_states_possibly_invalid = jax.tree.map(
         lambda t1, t2: jnp.concatenate([t1[None, :], t2], axis=0),
-        init_retval, scan_retvals
+        init_retval, masked_scan_retvals.value
+    )
+    masked_concatenated_states = genjax.Mask(
+        jnp.concatenate([jnp.array([True]), masked_scan_retvals.flag]),
+        concatenated_states_possibly_invalid
     )
+    return masked_concatenated_states
 
 @genjax.gen
 def sparse_observation_model(particle_absolute_poses, camera_pose, visibility, instrinsics, sigma):
     # TODO: add visibility
     uv = b3d.camera.screen_from_world(particle_absolute_poses.pos, camera_pose, instrinsics.const)
-    uv_ = genjax.normal(uv, jnp.tile(sigma, uv.shape)) @ "sensor_coordinates"
+    uv_ = b3d.modeling_utils.normal(uv, jnp.tile(sigma, uv.shape)) @ "sensor_coordinates"
     return uv_
 
 @genjax.gen
 def sparse_gps_model(latent_particle_model_args, obs_model_args):
-    # (b3d.camera.Intrinsics.from_array(jnp.array([1.0, 1.0, 1.0, 1.0])), 0.1)
-    particle_dynamics_summary = latent_particle_model(*latent_particle_model_args) @ "particle_dynamics"
-    obs = sparse_observation_model.vmap(in_axes=(0, 0, 0, None, None))(
-        particle_dynamics_summary["absolute_particle_poses"],
-        particle_dynamics_summary["camera_pose"],
-        particle_dynamics_summary["vis_mask"],
+    masked_particle_dynamics_summary = latent_particle_model(*latent_particle_model_args) @ "particle_dynamics"
+    _UNSAFE_particle_dynamics_summary = masked_particle_dynamics_summary.value
+    masked_obs = sparse_observation_model.mask().vmap(in_axes=(0, 0, 0, 0, None, None))(
+        masked_particle_dynamics_summary.flag,
+        _UNSAFE_particle_dynamics_summary["absolute_particle_poses"],
+        _UNSAFE_particle_dynamics_summary["camera_pose"],
+        _UNSAFE_particle_dynamics_summary["vis_mask"],
         *obs_model_args
     ) @ "observation"
-    return (particle_dynamics_summary, obs)
+    return (masked_particle_dynamics_summary, masked_obs)
 
 
 
@@ -166,15 +184,17 @@ def make_dense_gps_model(renderer):
 
     @genjax.gen
     def dense_gps_model(latent_particle_model_args, dense_likelihood_args):
-        # (b3d.camera.Intrinsics.from_array(jnp.array([1.0, 1.0, 1.0, 1.0])), 0.1)
-        particle_dynamics_summary = latent_particle_model(*latent_particle_model_args) @ "particle_dynamics"
-        absolute_particle_poses_last_frame = particle_dynamics_summary["absolute_particle_poses"][-1]
-        camera_pose_last_frame = particle_dynamics_summary["camera_pose"][-1]
+        masked_particle_dynamics_summary = latent_particle_model(*latent_particle_model_args) @ "particle_dynamics"
+        _UNSAFE_particle_dynamics_summary = masked_particle_dynamics_summary.value
+
+        last_timestep_index = jnp.sum(masked_particle_dynamics_summary.flag) - 1
+        absolute_particle_poses_last_frame = _UNSAFE_particle_dynamics_summary["absolute_particle_poses"][last_timestep_index]
+        camera_pose_last_frame = _UNSAFE_particle_dynamics_summary["camera_pose"][last_timestep_index]
         absolute_particle_poses_in_camera_frame = camera_pose_last_frame.inv() @ absolute_particle_poses_last_frame
 
         (meshes, likelihood_args) = dense_likelihood_args
         merged_mesh = Mesh.transform_and_merge_meshes(meshes, absolute_particle_poses_in_camera_frame)
         image = dense_observation_model(merged_mesh, likelihood_args) @ "observation"
-        return (particle_dynamics_summary, image)
+        return (masked_particle_dynamics_summary, image)
 
     return dense_gps_model

b3d/modeling_utils.py

@@ -3,6 +3,7 @@
 import jax
 import jax.numpy as jnp
 from tensorflow_probability.substrates import jax as tfp
+from genjax import Mask
 
 uniform_discrete = genjax.exact_density(
     lambda key, vals: jax.random.choice(key, vals),
@@ -35,4 +36,56 @@ def logpdf(v, *args, **kwargs):
         d = dist(*args, **kwargs)
         return jnp.sum(d.log_prob(v))
 
-    return genjax.exact_density(sampler, logpdf)
+    return genjax.exact_density(sampler, logpdf)
+
+normal = tfp_distribution(tfp.distributions.Normal)
+
+def masked_scan_combinator(step, **scan_kwargs):
+    """
+    Given a generative function `step` so that `step.scan(n=N)` is valid,
+    return a generative function accepting an input
+    `(initial_state, masked_input_values_array)` and returning a pair
+    `(masked_final_state, masked_returnvalue_sequence)`.
+    This operates similarly to `step.scan`, but the input values can be masked.
+    """
+    mstep = step.mask().dimap(
+        pre=lambda masked_state, masked_inval: (
+            jnp.logical_and(masked_state.flag, masked_inval.flag),
+            masked_state.value,
+            masked_inval.value
+        ),
+        post=lambda args, masked_retval: (
+            Mask(masked_retval.flag, masked_retval.value[0]),
+            Mask(masked_retval.flag, masked_retval.value[1])
+        )
+    )
+
+    # This should be given a pair (
+    #     Mask(True, initial_state), 
+    #     Mask(bools_indicating_active, input_vals)
+    # ).
+    # It wll output a pair (masked_final_state, masked_returnvalue_sequence).
+    scanned = mstep.scan(**scan_kwargs)
+
+    scanned_nice = scanned.dimap(
+        pre=lambda initial_state, masked_input_values: (
+            Mask(True, initial_state),
+            Mask(masked_input_values.flag, masked_input_values.value)
+        ),
+        post=lambda args, retval: retval
+    )
+
+    return scanned_nice
+
+def variable_length_unfold_combinator(step, **scan_kwargs):
+    """
+    Step should accept one arg, `state`, as input,
+    and should return a pair `(new_state, retval_for_this_timestep)`.
+    """
+    scanned = masked_scan_combinator(step, **scan_kwargs)
+    return scanned.dimap(
+        pre=lambda initial_state, n_steps: (
+            initial_state,
+            Mask(jnp.array())
+        )
+    )