RF3 has full array, deal with that during saving

Author: k-chrispens

src/sampleworks/utils/guidance_script_utils.py

@@ -69,6 +69,33 @@ def save_trajectory(
         raise ValueError(f"Invalid scaler type: {scaler_type}")
 
 
+def _assign_coords_to_array(
+    array_copy: AtomArrayStack,
+    coords: np.ndarray,
+    reward_param_mask: np.ndarray,
+) -> None:
+    """Assign trajectory coords into an AtomArrayStack, handling shape mismatches.
+
+    When the trajectory spans all atoms in the array (e.g. model
+    trajectories saved during a has_mismatch run (see pure_guidance.py or fk_steering.py)), coords
+    are assigned directly to ``.coord``.  Otherwise the ``reward_param_mask``
+    is used to index into the correct atom subset.
+    """
+    n_atoms_array = array_copy.coord.shape[-2]  # pyright: ignore[reportOptionalMemberAccess]
+    n_atoms_coords = coords.shape[-2]
+
+    if n_atoms_coords == n_atoms_array:
+        array_copy.coord = coords
+    elif n_atoms_coords == int(reward_param_mask.sum()):
+        array_copy.coord[:, reward_param_mask] = coords  # pyright: ignore[reportOptionalSubscript]
+    else:
+        raise ValueError(
+            f"Trajectory coords ({n_atoms_coords} atoms) match neither "
+            f"the full atom array ({n_atoms_array}) nor the masked subset "
+            f"({int(reward_param_mask.sum())})"
+        )
+
+
 def _save_trajectory(
     trajectory, atom_array, output_dir, reward_param_mask, subdir_name, save_every
 ):
@@ -89,7 +116,7 @@ def _save_trajectory(
             continue
         array_copy = atom_array.copy()
         array_copy = stack([array_copy] * ensemble_size)
-        array_copy.coord[:, reward_param_mask] = coords.detach().numpy()  # type: ignore[reportOptionalSubscript] coords will be subscriptable
+        _assign_coords_to_array(array_copy, coords.detach().numpy(), reward_param_mask)
         save_structure(str(output_dir / f"trajectory_{i}.cif"), array_copy)
 
 
@@ -115,7 +142,7 @@ def _save_fk_steering_trajectory(
         array_copy = stack([array_copy] * ensemble_size)
         # we save only the first ensemble out of n_particles, since saving
         # each particle at every step would clog trajectory saving
-        array_copy.coord[:, reward_param_mask] = coords[0].detach().numpy()  # type: ignore[reportOptionalSubscript] coords will be subscriptable
+        _assign_coords_to_array(array_copy, coords[0].detach().numpy(), reward_param_mask)
         save_structure(str(output_dir / f"trajectory_{i}.cif"), array_copy)