Adjust weights to account for the more general case

Author: filippozacchei

pysindy/_core.py

@@ -944,63 +944,56 @@ def comprehend_and_validate(arr, t):
     return x, x_dot, u
 
 def _expand_sample_weights(sample_weight, trajectories):
-    """Expand trajectory-level weights to per-sample or per-component weights.
+    """Expand trajectory-level weights to per-sample or per-component weights."""
 
-    Parameters
-    ----------
-    sample_weight : array-like
-        Can be
-        - None
-        - scalar
-        - shape (n_traj,)                : one weight per trajectory
-        - shape (n_samples_total,)       : one weight per sample
-        - shape (n_samples_total, n_tgt) : per-sample, per-target weights
-        - shape (n_traj, n_tgt)          : one weight per trajectory, per target
-
-    trajectories : list of arrays
-        Each trajectory, shape (n_samples_i, n_features).
-
-    Returns
-    -------
-    ndarray
-        Expanded weights:
-        - (n_samples_total,)
-        - (n_samples_total, n_tgt)
-    """
     if sample_weight is None:
         return None
 
     sample_weight = np.asarray(sample_weight)
     n_traj = len(trajectories)
     n_samples_total = sum(len(traj) for traj in trajectories)
 
-    # case: one weight per trajectory
-    if sample_weight.ndim == 1 and len(sample_weight) == n_traj:
+    # (1) trajectory-level
+    if sample_weight.ndim == 1 and sample_weight.shape[0] == n_traj:
         expanded = []
         for w, traj in zip(sample_weight, trajectories):
             expanded.extend([w] * len(traj))
         return np.asarray(expanded)
 
-    # case: one weight per sample
-    if sample_weight.ndim == 1 and len(sample_weight) == n_samples_total:
+    # (2) sample-level
+    if sample_weight.ndim == 1 and sample_weight.shape[0] == n_samples_total:
         return sample_weight
 
-    # case: per-sample, per-target
+    # (3) per-sample, per-target
     if sample_weight.ndim == 2 and sample_weight.shape[0] == n_samples_total:
         return sample_weight
 
-    # case: per-trajectory, per-target
+    # (4) per-trajectory, per-target
     if sample_weight.ndim == 2 and sample_weight.shape[0] == n_traj:
         expanded = []
         for w_vec, traj in zip(sample_weight, trajectories):
+            expanded.append(np.tile(w_vec, (len(traj), 1)))
+        return np.vstack(expanded)
+
+    # (5) per-trajectory, per-time-step, per-target
+    if sample_weight.ndim == 3 and sample_weight.shape[0] == n_traj:
+        expanded = []
+        for w_block, traj in zip(sample_weight, trajectories):
             n = len(traj)
-            expanded.append(np.tile(w_vec, (n, 1)))  # repeat per sample in that traj
+            if w_block.shape[0] != n:
+                raise ValueError(
+                    f"sample_weight time dimension {w_block.shape[0]} "
+                    f"does not match trajectory length {n}"
+                )
+            expanded.append(w_block)  # shape (n, n_targets)
         return np.vstack(expanded)
 
     raise ValueError(
-        f"sample_weight must be length {n_traj} (per trajectory), "
-        f"{n_samples_total} (per sample), "
-        f"({n_samples_total}, n_targets) (per sample/target), or "
-        f"({n_traj}, n_targets) (per trajectory/target). "
-        f"Got {sample_weight.shape}"
+        f"sample_weight must be one of:\n"
+        f"  (n_traj,), (n_samples_total,), "
+        f"  (n_samples_total, n_targets), "
+        f"  (n_traj, n_targets), "
+        f"  (n_traj, n_time, n_targets).\n"
+        f"Got {sample_weight.shape}."
     )
+