Pre-commit

Author: filippozacchei

pysindy/_core.py

@@ -1,5 +1,4 @@
 import warnings
-
 from abc import ABC
 from abc import abstractmethod
 from itertools import product
@@ -12,11 +11,12 @@
 from scipy.integrate import odeint
 from scipy.integrate import solve_ivp
 from scipy.interpolate import interp1d
+from sklearn import set_config
 from sklearn.base import BaseEstimator
 from sklearn.metrics import r2_score
 from sklearn.pipeline import Pipeline
 from sklearn.utils.validation import check_is_fitted
-from sklearn import set_config
+
 set_config(enable_metadata_routing=True)
 from typing_extensions import Self
 
@@ -348,7 +348,7 @@ def fit(
         feature_names : list of string, length n_input_features, optional
             Names for the input features (e.g. :code:`['x', 'y', 'z']`).
             If None, will use :code:`['x0', 'x1', ...]`.
-            
+
         sample_weight : float or array-like of shape (n_samples,), optional
             Per-sample weights for the regression. Passed internally to
             the optimizer (e.g. STLSQ). Supports compatibility with
@@ -383,11 +383,15 @@ def fit(
         self.feature_names = feature_names
 
         if sample_weight is not None:
-            mode = "weak" if "Weak" in self.feature_library.__class__.__name__ else "standard"
+            mode = (
+                "weak"
+                if "Weak" in self.feature_library.__class__.__name__
+                else "standard"
+            )
             sample_weight = _expand_sample_weights(
                 sample_weight, x, feature_library=self.feature_library, mode=mode
             )
-                                           
+
         steps = [
             ("features", self.feature_library),
             ("shaping", SampleConcatter()),
@@ -429,7 +433,7 @@ def predict(self, x, u=None):
         x, _, u = _comprehend_and_validate_inputs(x, 1, None, u, self.feature_library)
 
         check_is_fitted(self, "model")
-        
+
         if self.n_control_features_ > 0 and u is None:
             raise TypeError("Model was fit using control variables, so u is required")
         if self.n_control_features_ == 0 and u is not None:
@@ -485,7 +489,16 @@ def print(self, lhs=None, precision=3, **kwargs):
                 names = f"{lhs[i]}"
             print(f"{names} = {eqn}", **kwargs)
 
-    def score(self, x, t, x_dot=None, u=None, metric=r2_score, sample_weight=None, **metric_kws):
+    def score(
+        self,
+        x,
+        t,
+        x_dot=None,
+        u=None,
+        metric=r2_score,
+        sample_weight=None,
+        **metric_kws,
+    ):
         """
         Returns a score for the time derivative prediction produced by the model.
 
@@ -518,14 +531,14 @@ def score(self, x, t, x_dot=None, u=None, metric=r2_score, sample_weight=None, *
             See `Scikit-learn \
             <https://scikit-learn.org/stable/modules/model_evaluation.html>`_
             for more options.
-        
+
         sample_weight : array-like of shape (n_samples,), optional
             Per-sample weights passed directly to the metric. This is the
             preferred way to supply weights.
 
         metric_kws: dict, optional
             Optional keyword arguments to pass to the metric function.
-            
+
 
         Returns
         -------
@@ -548,7 +561,7 @@ def score(self, x, t, x_dot=None, u=None, metric=r2_score, sample_weight=None, *
 
         if sample_weight is not None:
             sample_weight = _expand_sample_weights(sample_weight, x)
-    
+
         x_dot = concat_sample_axis(x_dot)
         x_dot_predict = concat_sample_axis(x_dot_predict)
 
@@ -557,7 +570,7 @@ def score(self, x, t, x_dot=None, u=None, metric=r2_score, sample_weight=None, *
                 x_dot, x_dot_predict, return_indices=True
             )
             sample_weight = sample_weight[good_idx]
-            metric_kws = {**metric_kws, "sample_weight": sample_weight}  
+            metric_kws = {**metric_kws, "sample_weight": sample_weight}
         else:
             x_dot, x_dot_predict = drop_nan_samples(x_dot, x_dot_predict)
 
@@ -945,78 +958,115 @@ def comprehend_and_validate(arr, t):
         u = [comprehend_and_validate(ui, ti) for ui, ti in _zip_like_sequence(u, t)]
     return x, x_dot, u
 
-def _expand_sample_weights(sample_weight, trajectories, feature_library=None, mode="standard"):
-    """Expand per-trajectory sample weights for estimators or weak-form libraries.
+
+def _expand_sample_weights(
+    sample_weight, trajectories, feature_library=None, mode="standard"
+):
+    """
+    Expand per-trajectory or per-sample weights for use in SINDy estimators.
 
     Parameters
     ----------
-    sample_weight : sequence of array-like or None
-        Per-trajectory sample weights. Each element corresponds to one trajectory.
+    sample_weight : sequence of scalars or array-like
+        Weights for each trajectory. In "standard" mode, each entry can be:
+          - a scalar weight (applied to all samples in that trajectory), or
+          - an array of length equal to the number of samples (n_time) for that
+          trajectory.
+        In "weak" mode, each entry must be a single scalar weight per trajectory.
+
     trajectories : sequence
-        Sequence of trajectory objects, each with attributes `n_time` and `n_coord`.
+        Sequence of trajectory-like objects, each having attributes `n_time` and
+        `n_coord`.
+
     feature_library : object, optional
-        Library instance, required when mode='weak'.
+        Library instance used in weak-form mode. Must define attribute `K`
+        (the number of weak test functions). If missing, assumes K=1 with a warning.
+
     mode : {'standard', 'weak'}, default='standard'
-        Expansion mode:
-        - 'standard' : Concatenate weights per sample or per coordinate.
-        - 'weak'     : Expand weights for weak-form (integral) test functions.
+        - "standard": Expand per-sample weights to match concatenated samples.
+        - "weak": Repeat each trajectory’s single scalar weight `K` times.
 
     Returns
     -------
     np.ndarray or None
-        Concatenated and expanded sample weights, or None if no weights are given.
+        A 1D numpy array of concatenated and expanded sample weights,
+        or None if `sample_weight` is None.
     """
+    # -------------------------------------------------------------
+    # Early exit for None
+    # -------------------------------------------------------------
     if sample_weight is None:
         return None
 
-    if not (isinstance(sample_weight, Sequence) and not isinstance(sample_weight, np.ndarray)):
-        raise ValueError("sample_weight must be a list or tuple, not a scalar or numpy array.")
+    if not (
+        isinstance(sample_weight, Sequence)
+        and not isinstance(sample_weight, np.ndarray)
+    ):
+        raise ValueError(
+            "sample_weight must be a list or tuple, not a scalar or numpy array."
+        )
 
     if len(sample_weight) != len(trajectories):
         raise ValueError("sample_weight length must match number of trajectories.")
 
-    # --- Validate shape consistency ---
-    validated = []
-    for sw, traj in zip(sample_weight, trajectories):
-        arr = np.asarray(sw)
-        if arr.ndim == 0:
-            validated.append(arr)
-            continue
-        if arr.shape[0] != traj.n_time:
-            raise ValueError("sample_weight entry length does not match trajectory length.")
-        if arr.ndim == 2 and arr.shape[1] not in (1, traj.n_coord):
-            raise ValueError("sample_weight 2D second dim must be 1 or equal to n_coord.")
-        validated.append(arr)
-
-    # --- Weak-form expansion ---
+    # -------------------------------------------------------------
+    # Weak mode: one weight per trajectory, repeated K times
+    # -------------------------------------------------------------
     if mode == "weak":
-        n_funcs = getattr(feature_library, "K", 1)
-        if n_funcs is None:
-            warnings.warn("feature_library missing 'K'; assuming 1 test function.")
-            n_funcs = 1
-        return np.concatenate([np.repeat(np.asarray(sw), n_funcs, axis=0) for sw in validated])
-
-    # --- Standard expansion ---
-    n_coords = {int(t.n_coord) for t in trajectories}
-    if len(n_coords) != 1:
-        raise ValueError("All trajectories must have the same n_coord.")
-    n_coord = n_coords.pop()
-
-    processed = []
-    for arr in validated:
-        arr = np.asarray(arr)
-        if arr.ndim == 1:
-            arr = arr.reshape(-1, 1)
-        elif arr.ndim == 2 and arr.shape[1] == 1:
-            pass  # already correct shape
-        processed.append(arr)
-
-    # Promote to n_coord if any arrays have multiple coordinates
-    is_scalar_weight = all(a.shape[1] == 1 for a in processed)
-    if is_scalar_weight:
-        return np.concatenate([a.ravel() for a in processed])
-    expanded = [
-        np.broadcast_to(a, (a.shape[0], n_coord)) if a.shape[1] == 1 else a
-        for a in processed
-    ]
-    return np.concatenate(expanded, axis=0)
+        if feature_library is None:
+            raise ValueError("feature_library is required in weak mode.")
+
+        K = getattr(feature_library, "K", None)
+        if K is None:
+            warnings.warn("feature_library missing 'K'; assuming K=1.", UserWarning)
+            K = 1
+
+        validated = []
+        for w, traj in zip(sample_weight, trajectories):
+            arr = np.asarray(w)
+            if arr.ndim > 0 and arr.size > 1:
+                raise ValueError(
+                    "Weak mode expects exactly one weight per trajectory (scalar), "
+                    f"but got shape {arr.shape} for trajectory with {traj.n_time}"
+                    f"samples."
+                )
+            validated.append(float(arr))
+        return np.repeat(validated, K)
+
+    # -------------------------------------------------------------
+    # Standard mode: expand scalars or per-sample arrays
+    # -------------------------------------------------------------
+    expanded = []
+    for w, traj in zip(sample_weight, trajectories):
+        arr = np.asarray(w)
+
+        # Scalar → expand to all samples in trajectory
+        if arr.ndim == 0:
+            arr = np.full(traj.n_time, arr, dtype=float)
+
+        # 1D array → must match number of samples
+        elif arr.ndim == 1:
+            if arr.shape[0] != traj.n_time:
+                raise ValueError(
+                    f"sample_weight length {arr.shape[0]} does"
+                    f" not match trajectory length {traj.n_time}."
+                )
+
+        # 2D array → only (n,1) allowed
+        elif arr.ndim == 2:
+            if arr.shape[1] != 1:
+                raise ValueError(
+                    "sample_weight 2D arrays must have second dimension = 1."
+                )
+            if arr.shape[0] != traj.n_time:
+                raise ValueError(
+                    "sample_weight 2D array length does not match trajectory length."
+                )
+            arr = arr.ravel()
+
+        else:
+            raise ValueError("Invalid sample_weight shape.")
+
+        expanded.append(arr.ravel())
+
+    return np.concatenate(expanded)

pysindy/feature_library/weighted_weak_pde_library.py

@@ -1,6 +1,7 @@
 import numpy as np
-from .weak_pde_library import WeakPDELibrary
+
 from ..utils import AxesArray
+from .weak_pde_library import WeakPDELibrary
 
 
 class WeightedWeakPDELibrary(WeakPDELibrary):
@@ -39,14 +40,15 @@ def _build_whitener_from_variance(self):
 
         # --- robust weight-field shape handling ---
         base_grid = np.asarray(self.spatiotemporal_grid)
-        expected = tuple(base_grid.shape[:-1])           # e.g. (Nx, Nt) for a 2D grid
+        expected = tuple(base_grid.shape[:-1])  # e.g. (Nx, Nt) for a 2D grid
         var_grid = np.asarray(self.spatiotemporal_weights)
 
         if var_grid.shape == expected + (1,):
             var_grid = var_grid[..., 0]
         elif var_grid.shape != expected:
             raise ValueError(
-                f"spatiotemporal_weights must have shape {expected} or {expected + (1,)}, "
+                f"spatiotemporal_weights must have \
+                shape {expected} or {expected + (1,)}, "
                 f"got {var_grid.shape}"
             )
 
@@ -86,7 +88,7 @@ def _build_whitener_from_variance(self):
             vk = val_lists[k]
             Cov[k, k] = np.dot(vk, vk)
             idx_k = idx_lists[k]
-            
+
             map_k = dict(zip(idx_k.tolist(), vk.tolist()))
             for ell in range(k + 1, K):
                 s = 0.0
@@ -133,18 +135,17 @@ def _weak_form_setup(self):
             self._build_whitener_from_variance()
 
     def convert_u_dot_integral(self, u):
-        Vy = super().convert_u_dot_integral(u)     # (K, 1)
+        Vy = super().convert_u_dot_integral(u)  # (K, 1)
         Vy_w = self._apply_whitener(np.asarray(Vy))
         return AxesArray(Vy_w, {"ax_sample": 0, "ax_coord": 1})
 
     def transform(self, x_full):
-        VTheta_list = super().transform(x_full)    # list of (K, n_features)
+        VTheta_list = super().transform(x_full)  # list of (K, n_features)
         if self._L_chol is None:
             return VTheta_list
         out = []
         for VTheta in VTheta_list:
             A = np.asarray(VTheta)
-            A_w = self._apply_whitener(A)          # (K, m)
+            A_w = self._apply_whitener(A)  # (K, m)
             out.append(AxesArray(A_w, {"ax_sample": 0, "ax_coord": 1}))
         return out
-

pysindy/optimizers/base.py

@@ -16,7 +16,6 @@
 from sklearn.utils.extmath import safe_sparse_dot
 from sklearn.utils.validation import check_is_fitted
 from sklearn.utils.validation import check_X_y
-from sklearn.base import clone
 
 from .._typing import Float2D
 from .._typing import FloatDType
@@ -173,7 +172,7 @@ def fit(self, x_, y, sample_weight=None, **reduce_kws):
         y = AxesArray(np.asarray(y), y_axes)
         x_, y = drop_nan_samples(x_, y)
         x_, y = check_X_y(x_, y, accept_sparse=[], y_numeric=True, multi_output=True)
-          
+
         x, y, X_offset, y_offset, X_scale = _preprocess_data(
             x_,
             y,

pysindy/optimizers/stlsq.py

@@ -78,8 +78,8 @@ class STLSQ(BaseOptimizer):
     history_ : list
         History of ``coef_``. ``history_[k]`` contains the values of
         ``coef_`` at iteration k of sequentially thresholded least-squares.
-        
-            
+
+
     Notes
     -----
     - Supports ``sample_weight`` during :meth:`fit`. Sample weights are applied

pysindy/optimizers/trapping_sr3.py

@@ -517,7 +517,8 @@ def _objective(self, x, y, coef_sparse, A, PW, k):
         if self.verbose and k % max(1, self.max_iter // 10) == 0:
             print(
                 f"{k:5d} ... {sindy_loss:8.3e} ... {relax_loss:8.3e} ... {L1:8.2e}"
-                f" ... {nonlin_ens_loss:8.2e} ... {cubic_ens_loss:8.2e} ... {obj:8.2e}"
+                f" ... {nonlin_ens_loss:8.2e} ... {cubic_ens_loss:8.2e}"
+                f" ... {obj:8.2e}"
             )
         return obj
 

pysindy/utils/_axes.py

@@ -69,7 +69,8 @@
 
 import numpy as np
 from numpy.typing import NDArray
-from sklearn.base import BaseEstimator, TransformerMixin
+from sklearn.base import BaseEstimator
+from sklearn.base import TransformerMixin
 
 HANDLED_FUNCTIONS = {}
 
@@ -839,6 +840,7 @@ def __sklearn_is_fitted__(self):
     def transform(self, x_list):
         return concat_sample_axis(x_list)
 
+
 def concat_sample_axis(x_list: List[AxesArray]):
     """Concatenate all trajectories and axes used to create samples."""
     new_arrs = []