diff --git a/src/deepymod/data/base.py b/src/deepymod/data/base.py index b8b98e118..e839ae0c8 100644 --- a/src/deepymod/data/base.py +++ b/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 diff --git a/src/deepymod/model/deepmod.py b/src/deepymod/model/deepmod.py index f3705bb13..1dd62f051 100644 --- a/src/deepymod/model/deepmod.py +++ b/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 ) diff --git a/src/deepymod/model/sparse_estimators.py b/src/deepymod/model/sparse_estimators.py index 774a16417..18cee5b84 100644 --- a/src/deepymod/model/sparse_estimators.py +++ b/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_ diff --git a/src/deepymod/training/training.py b/src/deepymod/training/training.py index edeeaa724..9fb47f2bf 100644 --- a/src/deepymod/training/training.py +++ b/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 diff --git a/src/deepymod/utils/logger.py b/src/deepymod/utils/logger.py index 56c890685..b820c1830 100644 --- a/src/deepymod/utils/logger.py +++ b/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,14 +102,17 @@ 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, ) @@ -109,7 +120,7 @@ def update_tensorboard( 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, ) diff --git a/tests/test_compatibility.py b/tests/test_compatibility.py new file mode 100644 index 000000000..c3327878c --- /dev/null +++ b/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,)