talmolab · jf514 · Nov 1, 2024 · Oct 8, 2024 · Oct 25, 2024 · Oct 25, 2024
diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
@@ -101,7 +101,7 @@ jobs:
           verbose: false
           token: ${{ secrets.CODECOV_TOKEN }}
 
-      # Test probably delete this
+      # Smoke test. Shows end to end run with out crashing.
       - name: Smoke Test
         shell: bash -l {0}
-        run: python run_stac.py
+        run: python run_stac.py --config-name config_synth_data.yaml
diff --git a/Mat-to-Nwb-Synth-Data.ipynb b/Mat-to-Nwb-Synth-Data.ipynb
diff --git a/configs/config_synth_data.yaml b/configs/config_synth_data.yaml
@@ -0,0 +1,5 @@
+defaults:
+  - stac: stac_synth_data
+  - model: synth_data
+  - _self_
+
diff --git a/configs/model/fly_tethered.yaml b/configs/model/fly_tethered.yaml
@@ -176,5 +176,3 @@ N_SAMPLE_FRAMES: 100
 # If you have reason to believe that the initial offsets are correct for particular keypoints, 
 # you can regularize those sites using _SITES_TO_REGULARIZE. 
 M_REG_COEF: 1
-
-TIME_BINS: 0.02
diff --git a/configs/model/fly_treadmill.yaml b/configs/model/fly_treadmill.yaml
@@ -50,6 +50,8 @@ KEYPOINT_INITIAL_OFFSETS:
   l3: 0. 0. 0.
   r3: 0. 0. 0.
 
+ROOT_OPTIMIZATION_KEYPOINT: head
+
 TRUNK_OPTIMIZATION_KEYPOINTS: 
   - 'head'
   - 'thorax'
@@ -101,5 +103,3 @@ N_SAMPLE_FRAMES: 100
 # If you have reason to believe that the initial offsets are correct for particular keypoints, 
 # you can regularize those sites using _SITES_TO_REGULARIZE. 
 M_REG_COEF: 1
-
-TIME_BINS: 0.02
diff --git a/configs/model/mouse.yaml b/configs/model/mouse.yaml
@@ -122,6 +122,8 @@ KEYPOINT_INITIAL_OFFSETS:
   Lisfranc_L: 0.0 0.0 0.0
   MTP_R: 0.0 0.0 0.0
 
+ROOT_OPTIMIZATION_KEYPOINT: Trunk
+
 TRUNK_OPTIMIZATION_KEYPOINTS:
  - "Trunk"
  - "HipL"

diff --git a/configs/model/rodent.yaml b/configs/model/rodent.yaml
@@ -89,6 +89,8 @@ KEYPOINT_INITIAL_OFFSETS:
   WristL: 0. 0. 0.0
   WristR: 0. 0. 0.0
 
+ROOT_OPTIMIZATION_KEYPOINT: SpineL
+
 TRUNK_OPTIMIZATION_KEYPOINTS:
   - "SpineF"
   - "SpineL"

diff --git a/configs/model/synth_data.yaml b/configs/model/synth_data.yaml
@@ -0,0 +1,58 @@
+
+MJCF_PATH: 'models/synth_model.xml'
+
+# Frames per clip for transform.
+N_FRAMES_PER_CLIP: 1
+
+# Tolerance for the optimizations of the full model, limb, and root. 
+# TODO: Re-implement optimizer loops to use these tolerances
+FTOL: 5.0e-03
+ROOT_FTOL: 1.0e-05
+LIMB_FTOL: 1.0e-06
+
+# Number of alternating pose and offset optimization rounds. 
+N_ITERS: 1
+
+KP_NAMES: 
+  - part_0
+
+ROOT_OPTIMIZATION_KEYPOINT: part_0
+
+# The model sites used to register the keypoints.
+KEYPOINT_MODEL_PAIRS:
+  part_0: base
+
+# The initial offsets for each keypoint in meters.
+KEYPOINT_INITIAL_OFFSETS:
+  part_0: 0 0 0.01
+
+TRUNK_OPTIMIZATION_KEYPOINTS:
+  - part_0
+
+INDIVIDUAL_PART_OPTIMIZATION:
+  model_base: [base]
+
+# Color to use for each keypoint when visualizing the results
+KEYPOINT_COLOR_PAIRS:
+  part_0: 0 .5 1 1
+
+# What is the size of the animal you'd like to register, relative to the model?
+SCALE_FACTOR: 1
+
+# Multiplier to put the mocap data into the same scale as the data. Eg, if the
+# mocap data is known to be in millimeters and the model is in meters, this is
+# .001
+MOCAP_SCALE_FACTOR: 1
+
+# If you have reason to believe that the initial offsets are correct for particular keypoints, 
+# you can regularize those sites using this with M_REG_COEF. 
+SITES_TO_REGULARIZE: 
+  - part_0
+
+RENDER_FPS: 200
+
+N_SAMPLE_FRAMES: 1
+
+# If you have reason to believe that the initial offsets are correct for particular keypoints, 
+# you can regularize those sites using _SITES_TO_REGULARIZE. 
+M_REG_COEF: 1
diff --git a/configs/stac/stac_fly_treadmill.yaml b/configs/stac/stac_fly_treadmill.yaml
@@ -3,7 +3,7 @@ ik_only_path: "transform_treadmill.p"
 # File is too large to commit
 # DL from: https://datadryad.org/stash/dataset/doi:10.5061/dryad.mpg4f4r73
 # Actual file: https://datadryad.org/stash/downloads/file_stream/3361804
-data_path: "/tests/data/wt_berlin_linear_treadmill_dataset.csv"
+data_path: "../tests/data/wt_berlin_linear_treadmill_dataset.csv"
 
 n_fit_frames: 1800
 skip_fit: False

diff --git a/configs/stac/stac_synth_data.yaml b/configs/stac/stac_synth_data.yaml
@@ -0,0 +1,12 @@
+fit_offsets_path: "synth_fit.p"
+ik_only_path: "synth_ik_only.p"
+data_path: "tests/data/test_synth_1_frames.nwb"
+
+n_fit_frames: 1
+skip_fit_offsets: False
+skip_ik_only: False
+
+mujoco:
+  solver: newton
+  iterations: 1
+  ls_iterations: 4
diff --git a/demos/create_synth_data.ipynb b/demos/create_synth_data.ipynb
diff --git a/models/synth_model.xml b/models/synth_model.xml
@@ -0,0 +1,18 @@
+<mujoco>
+  <option timestep=".001">
+  </option>
+
+  <default>
+    <joint type="hinge" axis="0 -1 0"/>
+    <geom type="capsule" size=".02"/>
+  </default>
+
+  <worldbody>
+    <light pos="0 -.4 1"/>
+    <camera name="fixed" pos="0 -1 0" xyaxes="1 0 0 0 0 1"/>
+    <body name="base" pos="0 0 .2">
+      <joint type="free" name="root"/>
+      <geom fromto="0 0 0 0 0 -.25" rgba="1 1 0 1"/>
+    </body>
+  </worldbody>
+</mujoco>
diff --git a/stac_mjx/compute_stac.py b/stac_mjx/compute_stac.py
@@ -14,6 +14,7 @@ def root_optimization(
     mjx_model,
     mjx_data,
     kp_data: jp.ndarray,
+    root_kp_idx: int,
     lb: jp.ndarray,
     ub: jp.ndarray,
     site_idxs: jp.ndarray,
@@ -50,7 +51,7 @@ def root_optimization(
     # necessarily exactly so. The value of 3*18 is chosen for the
     # rodent.xml, corresponding to the index of 'SpineL' keypoint.
     # For the mouse model this should be 3*5, corresponding 'Trunk'
-    root_kp_idx = 3 * 18
+    # root_kp_idx = 3 * 18
     # FLY_MODEL:
     # root_kp_idx = 0
     q0.at[:3].set(kp_data[frame, :][root_kp_idx : root_kp_idx + 3])

diff --git a/stac_mjx/stac.py b/stac_mjx/stac.py
@@ -80,15 +80,22 @@
             self._mj_model.body(i).name for i in range(self._mj_model.nbody)
         ]
 
-        joint_names = [self._mj_model.joint(i).name for i in range(self._mj_model.njnt)]
+        if "ROOT_OPTIMIZATION_KEYPOINT" in self.cfg.model:
+            self._root_kp_idx = self._kp_names.index(
+                self.cfg.model.ROOT_OPTIMIZATION_KEYPOINT
+            )
+        else:
+            self._root_kp_idx = -1
 
         # Set up bounds and part_names based on joint ranges, taking into account the dimensionality of parameters
+        joint_names = [self._mj_model.joint(i).name for i in range(self._mj_model.njnt)]
         self._lb, self._ub, self._part_names = _align_joint_dims(
             self._mj_model.jnt_type, self._mj_model.jnt_range, joint_names
         )
 
         self._indiv_parts = self.part_opt_setup()
 
+        # Generate boolean flags for keypoints included in trunk optimization.
         self._trunk_kps = jp.array(
             [n in self.cfg.model.TRUNK_OPTIMIZATION_KEYPOINTS for n in kp_names],
         )
@@ -113,7 +120,7 @@
                 [any(part in name for part in parts) for name in self._part_names]
             )
 
-        if self.cfg.model.INDIVIDUAL_PART_OPTIMIZATION is None:
+        if "INDIVIDUAL_PART_OPTIMIZATION" not in self.cfg.model:
             indiv_parts = []
         else:
             indiv_parts = jp.array(
@@ -224,11 +231,16 @@
 
         # Begin optimization steps
         # Skip root optimization if model is fixed (no free joint at root)
-        if self._mj_model.jnt_type[0] == mujoco.mjtJoint.mjJNT_FREE:
+        if self._root_kp_idx == -1:
+            print(
+                "ROOT_OPTIMIZATION_KEYPOINT not specified, skipping Root Optimization."
+            )
+        elif self._mj_model.jnt_type[0] == mujoco.mjtJoint.mjJNT_FREE:
             mjx_data = compute_stac.root_optimization(
                 mjx_model,
                 mjx_data,
                 kp_data,
+                self._root_kp_idx,
                 self._lb,
                 self._ub,
                 self._body_site_idxs,
@@ -339,15 +351,20 @@
         )
 
         # q_phase - root
-        if self._mj_model.jnt_type[0] == mujoco.mjtJoint.mjJNT_FREE:
+        if self._root_kp_idx == -1:
+            print(
+                "Missing or invalid ROOT_OPTIMIZATION_KEYPOINT, skipping root_optimization()"
+            )
+        elif self._mj_model.jnt_type[0] == mujoco.mjtJoint.mjJNT_FREE:
             vmap_root_opt = jax.vmap(
                 compute_stac.root_optimization,
-                in_axes=(0, 0, 0, None, None, None, None),
+                in_axes=(0, 0, 0, None, None, None, None, None),
             )
             mjx_data = vmap_root_opt(
                 mjx_model,
                 mjx_data,
                 batched_kp_data,
+                self._root_kp_idx,
                 self._lb,
                 self._ub,
                 self._body_site_idxs,

diff --git a/tests/configs/model/test_mouse.yaml b/tests/configs/model/test_mouse.yaml
@@ -122,6 +122,8 @@ KEYPOINT_INITIAL_OFFSETS:
   Lisfranc_L: 0.0 0.0 0.0
   MTP_R: 0.0 0.0 0.0
 
+ROOT_OPTIMIZATION_KEYPOINT: Trunk
+
 TRUNK_OPTIMIZATION_KEYPOINTS:
  - "Trunk"
  - "HipL"

diff --git a/tests/configs/model/test_rodent.yaml b/tests/configs/model/test_rodent.yaml
@@ -80,6 +80,8 @@ KEYPOINT_INITIAL_OFFSETS:
   WristL: 0. 0. 0.0
   WristR: 0. 0. 0.0
 
+ROOT_OPTIMIZATION_KEYPOINT: SpineL
+
 TRUNK_OPTIMIZATION_KEYPOINTS:
   - "Spine"
   - "Hip"

diff --git a/tests/configs/model/test_rodent_label3d.yaml b/tests/configs/model/test_rodent_label3d.yaml
@@ -8,7 +8,6 @@ N_FRAMES_PER_CLIP: 360
 # presumed to be derived from label3d:
 KP_NAMES_LABEL3D_PATH: "tests/data/rat23.mat"
 
-
 # The model sites used to register the keypoints.
 KEYPOINT_MODEL_PAIRS:
   AnkleL: lower_leg_L
@@ -61,6 +60,8 @@ KEYPOINT_INITIAL_OFFSETS:
   WristL: 0. 0. 0.0
   WristR: 0. 0. 0.0
 
+ROOT_OPTIMIZATION_KEYPOINT: SpineL
+
 TRUNK_OPTIMIZATION_KEYPOINTS:
   - "Spine"
   - "Hip"

diff --git a/tests/configs/model/test_rodent_less_kp_names.yaml b/tests/configs/model/test_rodent_less_kp_names.yaml
@@ -59,6 +59,8 @@ KEYPOINT_INITIAL_OFFSETS:
   WristL: 0. 0. 0.0
   WristR: 0. 0. 0.0
 
+ROOT_OPTIMIZATION_KEYPOINT: SpineL
+
 TRUNK_OPTIMIZATION_KEYPOINTS:
   - "Spine"
   - "Hip"

diff --git a/tests/configs/model/test_rodent_no_kp_names.yaml b/tests/configs/model/test_rodent_no_kp_names.yaml
@@ -58,6 +58,8 @@ KEYPOINT_INITIAL_OFFSETS:
   WristL: 0. 0. 0.0
   WristR: 0. 0. 0.0
 
+ROOT_OPTIMIZATION_KEYPOINT: SpineL
+
 TRUNK_OPTIMIZATION_KEYPOINTS:
   - "Spine"
   - "Hip"

diff --git a/tests/data/test_synth_1_frames.nwb b/tests/data/test_synth_1_frames.nwb