Fix shape handling for small libraries and recent PySINDy

src/deepymod/data/base.py

@@ -158,10 +158,10 @@ def apply_normalize(X):
             X (torch.tensor): data to be minmax normalized
         Returns:
             (torch.tensor): minmaxed data"""
-        X_norm = (X - X.view(-1, X.shape[-1]).min(dim=0).values) / (
-            X.view(-1, X.shape[-1]).max(dim=0).values
-            - X.view(-1, X.shape[-1]).min(dim=0).values
-        ) * 2 - 1
+        X_flat = X.reshape(-1, X.shape[-1])
+        minimum = X_flat.min(dim=0).values
+        maximum = X_flat.max(dim=0).values
+        X_norm = (X - minimum) / (maximum - minimum) * 2 - 1
         return X_norm
 
     @staticmethod

src/deepymod/model/deepmod.py

@@ -137,12 +137,12 @@ def forward(self, thetas: TensorList, time_derivs: TensorList) -> TensorList:
             ]
 
         self.coeff_vectors = [
-            self.fit(theta, time_deriv.squeeze())[:, None]
+            np.asarray(self.fit(theta, time_deriv.reshape(-1))).reshape(-1, 1)
             for theta, time_deriv in zip(normed_thetas, normed_time_derivs)
         ]
         sparsity_masks = [
             torch.tensor(coeff_vector != 0.0, dtype=torch.bool)
-            .squeeze()
+            .reshape(-1)
             .to(thetas[0].device)  # move to gpu if required
             for coeff_vector in self.coeff_vectors
         ]
@@ -293,7 +293,7 @@ def constraint_coeffs(self, scaled=False, sparse=False) -> TensorList:
         coeff_vectors = self.constraint.coeff_vectors
         if scaled:  # perform normalization
             coeff_vectors = [
-                coeff / norm[:, None]
+                coeff / norm.reshape(-1)[:, None]
                 for coeff, norm, mask in zip(
                     coeff_vectors, self.library.norms, self.sparsity_masks
                 )

src/deepymod/model/sparse_estimators.py

@@ -177,9 +177,7 @@ def TrainSTLSQ(
         delta_t = delta_threshold  # for interal use, can be updated
 
         # Initial estimate
-        optimizer = STLSQ(
-            threshold=0, alpha=0.0, fit_intercept=False
-        )  # Now similar to LSTSQ
+        optimizer = STLSQ(threshold=0, alpha=0.0)  # Now similar to LSTSQ
         y_predict = optimizer.fit(X_train, y_train).predict(X_test)
         min_loss = np.linalg.norm(y_predict - y_test, 2) + l0 * np.count_nonzero(
             optimizer.coef_

src/deepymod/training/training.py

@@ -117,11 +117,12 @@ def train(
                 )
 
             # ================= Checking convergence
-            l1_norm = torch.sum(
-                torch.abs(
-                    torch.cat(model.constraint_coeffs(sparse=True, scaled=True), dim=1)
-                )
-            )
+            l1_norm = torch.stack(
+                [
+                    torch.sum(torch.abs(coeffs))
+                    for coeffs in model.constraint_coeffs(sparse=True, scaled=True)
+                ]
+            ).sum()
             converged = convergence(
                 iteration, l1_norm
             )  # Check if change is smaller than delta and if we've exceeded patience

src/deepymod/utils/logger.py

@@ -18,6 +18,12 @@ def __init__(self, exp_ID, log_dir):
         )
         self.log_dir = self.writer.get_logdir()
 
+    @staticmethod
+    def coefficient_items(values) -> np.ndarray:
+        if torch.is_tensor(values):
+            values = values.detach().cpu().numpy()
+        return np.asarray(values).reshape(-1)
+
     def __call__(
         self,
         iteration,
@@ -29,7 +35,9 @@ def __call__(
         estimator_coeffs,
         **kwargs,
     ):
-        l1_norm = torch.sum(torch.abs(torch.cat(constraint_coeffs, dim=1)), dim=0)
+        l1_norm = torch.stack(
+            [torch.sum(torch.abs(coeffs)) for coeffs in constraint_coeffs]
+        )
 
         self.update_tensorboard(
             iteration,
@@ -94,22 +102,25 @@ def update_tensorboard(
         ):
             self.writer.add_scalars(
                 f"coeffs/output_{output_idx}",
-                {f"coeff_{idx}": val for idx, val in enumerate(coeffs.squeeze())},
+                {
+                    f"coeff_{idx}": val
+                    for idx, val in enumerate(self.coefficient_items(coeffs))
+                },
                 iteration,
             )
             self.writer.add_scalars(
                 f"unscaled_coeffs/output_{output_idx}",
                 {
                     f"coeff_{idx}": val
-                    for idx, val in enumerate(unscaled_coeffs.squeeze())
+                    for idx, val in enumerate(self.coefficient_items(unscaled_coeffs))
                 },
                 iteration,
             )
             self.writer.add_scalars(
                 f"estimator_coeffs/output_{output_idx}",
                 {
                     f"coeff_{idx}": val
-                    for idx, val in enumerate(estimator_coeffs.squeeze())
+                    for idx, val in enumerate(self.coefficient_items(estimator_coeffs))
                 },
                 iteration,
             )

tests/test_compatibility.py

@@ -0,0 +1,106 @@
+import numpy as np
+import torch
+
+from deepymod.data.base import Dataset
+from deepymod.model.deepmod import DeepMoD, Estimator
+from deepymod.model.sparse_estimators import PDEFIND
+from deepymod.utils.logger import Logger
+
+
+class ScalarEstimator(Estimator):
+    def fit(self, X: np.ndarray, y: np.ndarray) -> np.ndarray:
+        return np.float64(2.0)
+
+
+class MatrixEstimator(Estimator):
+    def fit(self, X: np.ndarray, y: np.ndarray) -> np.ndarray:
+        return np.array([[1.0], [0.0]])
+
+
+def test_apply_normalize_accepts_non_contiguous_tensors():
+    values = torch.arange(24, dtype=torch.float32).reshape(2, 3, 4).permute(1, 0, 2)
+
+    normalized = Dataset.apply_normalize(values)
+    flat = normalized.reshape(-1, normalized.shape[-1])
+
+    assert normalized.shape == values.shape
+    assert torch.allclose(flat.min(dim=0).values, -torch.ones(values.shape[-1]))
+    assert torch.allclose(flat.max(dim=0).values, torch.ones(values.shape[-1]))
+
+
+def test_estimator_keeps_single_coefficient_and_mask_one_dimensional():
+    theta = torch.ones((5, 1))
+    time_deriv = 2.0 * torch.ones((5, 1))
+
+    estimator = ScalarEstimator()
+    masks = estimator([theta], [time_deriv])
+
+    assert estimator.coeff_vectors[0].shape == (1, 1)
+    assert masks[0].shape == (1,)
+    assert masks[0].dtype == torch.bool
+
+
+def test_estimator_flattens_matrix_coefficients_without_extra_dimension():
+    theta = torch.ones((5, 2))
+    time_deriv = torch.ones((5, 1))
+
+    estimator = MatrixEstimator()
+    masks = estimator([theta], [time_deriv])
+
+    assert estimator.coeff_vectors[0].shape == (2, 1)
+    assert masks[0].shape == (2,)
+
+
+def test_constraint_coeffs_scaled_accepts_single_term_norms():
+    model = DeepMoD.__new__(DeepMoD)
+    model.constraint = type(
+        "ConstraintStub",
+        (),
+        {
+            "coeff_vectors": [torch.tensor([[4.0]])],
+            "sparsity_masks": [torch.tensor([True])],
+        },
+    )()
+    model.library = type("LibraryStub", (), {"norms": [torch.tensor(2.0)]})()
+
+    coeffs = model.constraint_coeffs(scaled=True)
+
+    assert coeffs[0].shape == (1, 1)
+    assert torch.allclose(coeffs[0], torch.tensor([[2.0]]))
+
+
+def test_logger_handles_coefficients_with_single_or_heterogeneous_lengths():
+    assert Logger.coefficient_items(torch.tensor([[1.0]])).shape == (1,)
+    assert Logger.coefficient_items(np.array([[1.0], [2.0]])).shape == (2,)
+
+    logger = Logger.__new__(Logger)
+    captured = {}
+
+    def capture_tensorboard(*args, **kwargs):
+        captured["l1"] = args[4]
+
+    logger.update_tensorboard = capture_tensorboard
+    logger.update_terminal = lambda *args, **kwargs: None
+
+    coeffs = [torch.ones((2, 1)), torch.ones((3, 1))]
+    logger(
+        iteration=0,
+        loss=torch.tensor(0.0),
+        MSE=torch.tensor([0.0, 0.0]),
+        Reg=torch.tensor([0.0, 0.0]),
+        constraint_coeffs=coeffs,
+        unscaled_constraint_coeffs=coeffs,
+        estimator_coeffs=coeffs,
+    )
+
+    assert torch.allclose(captured["l1"], torch.tensor([2.0, 3.0]))
+
+
+def test_pdefind_stlsq_runs_with_installed_pysindy_signature():
+    x = np.linspace(-1.0, 1.0, 30)
+    X = np.column_stack([x, np.ones_like(x)])
+    y = (2.0 * x)[:, None]
+
+    coeffs = PDEFIND.TrainSTLSQ(X, y, alpha=0.0, delta_threshold=0.01, max_iterations=2)
+
+    assert np.asarray(coeffs).reshape(-1).shape == (2,)