folivo/PINA
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Fix basis device transfer in PODBlock (#650)

Browse Source 
* fix gpu data moving
This commit is contained in:
Filippo Olivo
2025-09-19 13:42:32 +02:00
committed by GitHub
parent 87c5c6a674
commit 4a6e73fa54
2 changed files with 20 additions and 15 deletions
Download Patch File Download Diff File Expand all files Collapse all files

28
pina/model/block/pod_block.py
View File

@@ -1,7 +1,7 @@
"""Module for Base Continuous Convolution class.""" """Module for Base Continuous Convolution class."""
import torch
import warnings import warnings
import torch
class PODBlock(torch.nn.Module): class PODBlock(torch.nn.Module):
@@ -29,9 +29,10 @@ class PODBlock(torch.nn.Module):
""" """
super().__init__() super().__init__()
self.__scale_coefficients = scale_coefficients self.__scale_coefficients = scale_coefficients
self._basis = None self.register_buffer("_basis", None)
self._singular_values = None self._singular_values = None
self._scaler = None self.register_buffer("_std", None)
self.register_buffer("_mean", None)
self._rank = rank self._rank = rank
@property @property
@@ -94,12 +95,12 @@ class PODBlock(torch.nn.Module):
:return: The scaler dictionary. :return: The scaler dictionary.
:rtype: dict :rtype: dict
""" """
if self._scaler is None: if self._std is None:
return None return None
return { return {
"mean": self._scaler["mean"][: self.rank], "mean": self._mean[: self.rank],
"std": self._scaler["std"][: self.rank], "std": self._std[: self.rank],
} }
@property @property
@@ -119,6 +120,10 @@ class PODBlock(torch.nn.Module):
are scaled after the projection to have zero mean and unit variance. are scaled after the projection to have zero mean and unit variance.
:param torch.Tensor X: The input tensor to be reduced. :param torch.Tensor X: The input tensor to be reduced.
:param bool randomized: If ``True``, a randomized algorithm is used to
compute the POD basis. In general, this leads to faster
computations, but the results may be less accurate. Default is
``True``.
""" """
self._fit_pod(X, randomized) self._fit_pod(X, randomized)
@@ -132,10 +137,8 @@ class PODBlock(torch.nn.Module):
:param torch.Tensor coeffs: The coefficients to be scaled. :param torch.Tensor coeffs: The coefficients to be scaled.
""" """
self._scaler = { self._std = torch.std(coeffs, dim=1) # pylint: disable=W0201
"std": torch.std(coeffs, dim=1), self._mean = torch.mean(coeffs, dim=1) # pylint: disable=W0201
"mean": torch.mean(coeffs, dim=1),
}
def _fit_pod(self, X, randomized): def _fit_pod(self, X, randomized):
""" """
@@ -154,13 +157,14 @@ class PODBlock(torch.nn.Module):
else: else:
if randomized: if randomized:
warnings.warn( warnings.warn(
"Considering a randomized algorithm to compute the POD basis" "Considering a randomized algorithm to compute the POD "
"basis"
) )
u, s, _ = torch.svd_lowrank(X.T, q=X.shape[0]) u, s, _ = torch.svd_lowrank(X.T, q=X.shape[0])
else: else:
u, s, _ = torch.svd(X.T) u, s, _ = torch.svd(X.T)
self._basis = u.T self._basis = u.T # pylint: disable=W0201
self._singular_values = s self._singular_values = s
def forward(self, X): def forward(self, X):

7
tests/test_block/test_pod.py
View File

@@ -42,13 +42,14 @@ def test_fit(rank, scale, randomized):
assert pod.singular_values.shape == (rank,) assert pod.singular_values.shape == (rank,)
assert pod._singular_values.shape == (n_snap,) assert pod._singular_values.shape == (n_snap,)
if scale is True: if scale is True:
assert pod._scaler["mean"].shape == (n_snap,) assert pod._mean.shape == (n_snap,)
assert pod._scaler["std"].shape == (n_snap,) assert pod._std.shape == (n_snap,)
assert pod.scaler["mean"].shape == (rank,) assert pod.scaler["mean"].shape == (rank,)
assert pod.scaler["std"].shape == (rank,) assert pod.scaler["std"].shape == (rank,)
assert pod.scaler["mean"].shape[0] == pod.basis.shape[0] assert pod.scaler["mean"].shape[0] == pod.basis.shape[0]
else: else:
assert pod._scaler == None assert pod._std == None
assert pod._mean == None
assert pod.scaler == None assert pod.scaler == None