diff --git a/pina/model/layers/pod.py b/pina/model/layers/pod.py index 3e5627487..be8e137d8 100644 --- a/pina/model/layers/pod.py +++ b/pina/model/layers/pod.py @@ -4,7 +4,7 @@ import torch from .stride import Stride from .utils_convolution import optimizing - +import warnings class PODBlock(torch.nn.Module): """ @@ -85,7 +85,7 @@ def scale_coefficients(self): """ return self.__scale_coefficients - def fit(self, X): + def fit(self, X, randomized=True): """ Set the POD basis by performing the singular value decomposition of the given tensor. If `self.scale_coefficients` is True, the coefficients @@ -93,7 +93,7 @@ def fit(self, X): :param torch.Tensor X: The tensor to be reduced. """ - self._fit_pod(X) + self._fit_pod(X, randomized) if self.__scale_coefficients: self._fit_scaler(torch.matmul(self._basis, X.T)) @@ -112,16 +112,24 @@ def _fit_scaler(self, coeffs): "mean": torch.mean(coeffs, dim=1), } - def _fit_pod(self, X): + def _fit_pod(self, X, randomized): """ Private method that computes the POD basis of the given tensor and stores it in the private member `_basis`. :param torch.Tensor X: The tensor to be reduced. """ if X.device.type == "mps": # svd_lowrank not arailable for mps + warnings.warn( + "svd_lowrank not available for mps, using svd instead." + "This may slow down computations.", ResourceWarning + ) self._basis = torch.svd(X.T)[0].T else: - self._basis = torch.svd_lowrank(X.T, q=X.shape[0])[0].T + if randomized: + warnings.warn("Considering a randomized algorithm to compute the POD basis") + self._basis = torch.svd_lowrank(X.T, q=X.shape[0])[0].T + else: + self._basis = torch.svd(X.T)[0].T def forward(self, X): """ diff --git a/tests/test_layers/test_pod.py b/tests/test_layers/test_pod.py index 23b30ceef..433fcafed 100644 --- a/tests/test_layers/test_pod.py +++ b/tests/test_layers/test_pod.py @@ -25,9 +25,10 @@ def test_fit(rank, scale): @pytest.mark.parametrize("scale", [True, False]) @pytest.mark.parametrize("rank", [1, 2, 10]) -def test_fit(rank, scale): +@pytest.mark.parametrize("randomized", [True, False]) +def test_fit(rank, scale, randomized): pod = PODBlock(rank, scale) - pod.fit(toy_snapshots) + pod.fit(toy_snapshots, randomized) n_snap = toy_snapshots.shape[0] dof = toy_snapshots.shape[1] assert pod.basis.shape == (rank, dof) @@ -65,18 +66,20 @@ def test_forward(): @pytest.mark.parametrize("scale", [True, False]) @pytest.mark.parametrize("rank", [1, 2, 10]) -def test_expand(rank, scale): +@pytest.mark.parametrize("randomized", [True, False]) +def test_expand(rank, scale, randomized): pod = PODBlock(rank, scale) - pod.fit(toy_snapshots) + pod.fit(toy_snapshots, randomized) c = pod(toy_snapshots) torch.testing.assert_close(pod.expand(c), toy_snapshots) torch.testing.assert_close(pod.expand(c[0]), toy_snapshots[0].unsqueeze(0)) @pytest.mark.parametrize("scale", [True, False]) @pytest.mark.parametrize("rank", [1, 2, 10]) -def test_reduce_expand(rank, scale): +@pytest.mark.parametrize("randomized", [True, False]) +def test_reduce_expand(rank, scale, randomized): pod = PODBlock(rank, scale) - pod.fit(toy_snapshots) + pod.fit(toy_snapshots, randomized) torch.testing.assert_close( pod.expand(pod.reduce(toy_snapshots)), toy_snapshots)