feectools 0.1.6__tar.gz → 0.1.7__tar.gz

{feectools-0.1.6/feectools.egg-info → feectools-0.1.7}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: feectools
-Version: 0.1.6
+Version: 0.1.7
 Summary: Slimmed-down fork of Psydac (https://github.com/pyccel/psydac) with less functionality and fewer dependencies.
 Author-email: Psydac development team <psydac@googlegroups.com>
 Maintainer-email: Stefan Possanner <stefan.possanner@ipp.mpg.de>, Max Lindqvist <max.lindqvist@ipp.mpg.de>, Yaman Güçlü <yaman.guclu@gmail.com>, Martin Campos Pinto <martin.campos-pinto@ipp.mpg.de>, Ahmed Ratnani <ratnaniahmed@gmail.com>

{feectools-0.1.6 → feectools-0.1.7}/feectools/linalg/block.py

@@ -886,6 +886,14 @@ class BlockLinearOperator(LinearOperator):
         """
         return tuple(self._blocks)
 
+    # ...
+    @property
+    def nbytes(self):
+        """Local (per-MPI-rank) memory footprint of all non-zero blocks, in bytes.
+        Blocks which do not expose an 'nbytes' attribute (e.g. matrix-free operators)
+        are counted as zero."""
+        return int(sum(getattr(Lij, 'nbytes', 0) for Lij in self._blocks.values()))
+
     # ...
     def update_ghost_regions(self):
         for Lij in self._blocks.values():

{feectools-0.1.6 → feectools-0.1.7}/feectools/linalg/direct_solvers.py

@@ -23,7 +23,7 @@ def to_bnd(A):
     ua   = dmat.offsets.max()
     cmat = dmat.tocsr()
 
-    A_bnd = xp.zeros((1+ua+2*la, cmat.shape[1]), A.dtype)
+    A_bnd = xp.zeros((1+ua+2*la, cmat.shape[1]), dtype=A.dtype)
 
     for i,j in zip(*cmat.nonzero()):
         A_bnd[la+ua+i-j, j] = cmat[i,j]

{feectools-0.1.6 → feectools-0.1.7}/feectools/linalg/kron.py

@@ -73,6 +73,12 @@ class KroneckerStencilMatrix(LinearOperator):
     def mats( self ):
         return self._mats
 
+    # ...
+    @property
+    def nbytes( self ):
+        """Local (per-MPI-rank) memory footprint of the 1d factor matrices, in bytes."""
+        return int(sum(getattr(mat, 'nbytes', 0) for mat in self._mats))
+
     # ...
     def dot(self, x, out=None):
 

feectools-0.1.7/feectools/linalg/memory.py

@@ -0,0 +1,41 @@
+# coding: utf-8
+"""
+Bookkeeping of the memory occupied by the stencil matrices that are currently alive.
+
+Every :class:`~feectools.linalg.stencil.StencilMatrix` that allocates its data array registers
+itself (weakly) in the module-level :data:`stencil_matrix_memory` tracker, so that an application
+can report how much memory its matrices actually take, without having to walk its own data
+structures. Matrices created with ``dry_run=True`` do not allocate anything and are not registered.
+"""
+
+import weakref
+
+__all__ = ('MatrixMemoryTracker', 'stencil_matrix_memory')
+
+
+class MatrixMemoryTracker:
+    """Weak registry of allocated matrices; matrices that are garbage collected drop out of it."""
+
+    def __init__(self):
+        self._matrices = weakref.WeakSet()
+
+    def register(self, matrix):
+        """Add a matrix to the registry (does not keep it alive)."""
+        self._matrices.add(matrix)
+
+    def clear(self):
+        """Forget all registered matrices."""
+        self._matrices.clear()
+
+    @property
+    def n_matrices(self):
+        """Number of currently alive registered matrices."""
+        return len(self._matrices)
+
+    @property
+    def nbytes(self):
+        """Local (per-MPI-rank) memory footprint, in bytes, of all currently alive registered matrices."""
+        return int(sum(matrix.nbytes for matrix in tuple(self._matrices)))
+
+
+stencil_matrix_memory = MatrixMemoryTracker()

{feectools-0.1.6 → feectools-0.1.7}/feectools/linalg/stencil.py

@@ -14,6 +14,7 @@ from scipy.sparse import coo_matrix, diags as sp_diags
 
 from feectools.ddm.mpi import mpi as MPI
 from feectools.linalg.basic  import VectorSpace, Vector, LinearOperator
+from feectools.linalg.memory import stencil_matrix_memory
 from feectools.ddm.cart      import find_mpi_type, CartDecomposition, InterfaceCartDecomposition
 from feectools.ddm.utilities import get_data_exchanger
 from feectools.api.settings  import PSYDAC_BACKENDS
@@ -637,7 +638,7 @@ class StencilVector(Vector):
 
     # ...
     def _toarray_parallel_no_pads(self, order='C'):
-        a         = xp.zeros( self.space.npts, self.dtype )
+        a         = xp.zeros( self.space.npts, dtype=self.dtype )
         idx_from  = tuple( slice(m*p,-m*p) if p != 0 else slice(0, None) for p,m in zip(self.pads, self.space.shifts) )
         idx_to    = tuple( slice(s,e+1) for s,e in zip(self.starts,self.ends) )
         a[idx_to] = self._data[idx_from]
@@ -649,7 +650,7 @@ class StencilVector(Vector):
         pads = [m*p for m,p in zip(self.space.shifts, self.pads)]
         # Step 0: create extended n-dimensional array with zero values
         shape = tuple( n+2*p for n,p in zip( self.space.npts, pads ) )
-        a = xp.zeros( shape, self.dtype )
+        a = xp.zeros( shape, dtype=self.dtype )
 
         # Step 1: write extended data chunk (local to process) onto array
         idx = tuple( slice(s,e+2*p+1) for s,e,p in
@@ -905,8 +906,14 @@ class StencilMatrix(LinearOperator):
 
     precompiled : bool
         Whether to use precompiled kernels for .dot() and .transpose()
+
+    dry_run : bool
+        If True, only compute the shape of the data array (:attr:`data_shape`) and return early,
+        without allocating the (potentially large) data array and without setting up the
+        dot/transpose kernels. The resulting object is *not* usable as a linear operator; its only
+        purpose is to report the memory footprint the matrix would have via :attr:`nbytes`.
     """
-    def __init__( self, V, W, pads=None , backend=None, precompiled=True):
+    def __init__( self, V, W, pads=None , backend=None, precompiled=True, dry_run=False):
 
         assert isinstance(V, StencilVectorSpace)
         assert isinstance(W, StencilVectorSpace)
@@ -918,13 +925,21 @@ class StencilMatrix(LinearOperator):
             for p,vp in zip(pads, V.pads):
                 assert p<=vp
 
-        self._pads     = pads or tuple(V.pads)
-        dims           = list(W.shape)
-        diags          = [compute_diag_len(p, md, mc) for p,md,mc in zip(self._pads, V.shifts, W.shifts)]
-        self._data     = xp.zeros(tuple(int(d) for d in (dims + diags)), dtype=W.dtype)
-        self._domain   = V
-        self._codomain = W
-        self._ndim     = len(dims)
+        self._pads       = pads or tuple(V.pads)
+        dims             = list(W.shape)
+        diags            = [compute_diag_len(p, md, mc) for p,md,mc in zip(self._pads, V.shifts, W.shifts)]
+        self._data_shape = tuple(int(d) for d in (dims + diags))
+        self._domain     = V
+        self._codomain   = W
+        self._ndim       = len(dims)
+        self._dry_run    = dry_run
+
+        # memory estimation only: do not allocate the data array, see the nbytes property
+        if dry_run:
+            return
+
+        self._data       = xp.zeros(self._data_shape, dtype=W.dtype)
+        stencil_matrix_memory.register(self)
         self._backend  = backend
         self._precompiled = precompiled
         self._is_T     = False
@@ -975,6 +990,17 @@ class StencilMatrix(LinearOperator):
             backend = PSYDAC_BACKENDS.get(os.environ.get('PSYDAC_BACKEND')) or PSYDAC_BACKENDS['python']
         self.set_backend(backend, precompiled)
 
+    # ...
+    def __getattr__(self, name):
+        # only called when the attribute was not found the usual way; give a helpful
+        # message for the attributes that are missing on a dry-run matrix
+        if self.__dict__.get('_dry_run', False):
+            raise AttributeError(
+                f"'{type(self).__name__}.{name}' is not available because the matrix was created with "
+                "dry_run=True (memory estimation only, no data allocated)."
+            )
+        raise AttributeError(f"'{type(self).__name__}' object has no attribute '{name}'")
+
     #--------------------------------------
     # Abstract interface
     #--------------------------------------
@@ -992,6 +1018,28 @@ class StencilMatrix(LinearOperator):
     def dtype(self):
         return self._domain.dtype
 
+    # ...
+    @property
+    def dry_run(self):
+        """Whether the matrix was created for memory estimation only, i.e. without allocating data."""
+        return self._dry_run
+
+    # ...
+    @property
+    def data_shape(self):
+        """Shape of the local data array (n_rows in each direction + n_diagonals in each direction)."""
+        return self._data_shape
+
+    # ...
+    @property
+    def nbytes(self):
+        """Local (per-MPI-rank) memory footprint of the data array, in bytes. Also available
+        for matrices created with ``dry_run=True``, i.e. before/without allocating the data."""
+        nbytes = xp.dtype(self._codomain.dtype).itemsize
+        for n in self._data_shape:
+            nbytes *= n
+        return int(nbytes)
+
     # ...
     def dot(self, v, out=None):
         """
@@ -2076,6 +2124,11 @@ class StencilDiagonalMatrix(LinearOperator):
     def dtype(self):
         return self._data.dtype
 
+    @property
+    def nbytes(self):
+        """Local (per-MPI-rank) memory footprint of the data array, in bytes."""
+        return int(self._data.nbytes)
+
     def tosparse(self):
         return sp_diags(self._data.ravel())
 
@@ -2369,6 +2422,12 @@ class StencilInterfaceMatrix(LinearOperator):
     def dtype(self):
         return self.domain.dtype
 
+    # ...
+    @property
+    def nbytes(self):
+        """Local (per-MPI-rank) memory footprint of the data array, in bytes."""
+        return int(self._data.nbytes)
+
     # ...
     def dot(self, v, out=None):
 

{feectools-0.1.6 → feectools-0.1.7}/feectools/linalg/tests/test_linalg.py

@@ -4,6 +4,7 @@ import cunumpy as xp
 from feectools.linalg.block import BlockLinearOperator, BlockVector, BlockVectorSpace
 from feectools.linalg.basic import LinearOperator, ZeroOperator, IdentityOperator, ComposedLinearOperator, SumLinearOperator, PowerLinearOperator, ScaledLinearOperator
 from feectools.linalg.stencil import StencilVectorSpace, StencilVector, StencilMatrix
+from feectools.linalg.memory import stencil_matrix_memory
 from feectools.linalg.solvers import ConjugateGradient, inverse
 from feectools.ddm.cart       import DomainDecomposition, CartDecomposition
 
@@ -1057,6 +1058,60 @@ def test_dot_inner():
     assert r0 == r2
     assert r0 == r3
 
+#===============================================================================
+# DRY-RUN / MEMORY REGISTRATION TESTS
+#===============================================================================
+
+def test_stencil_matrix_dry_run_no_data_no_register():
+    """A dry-run StencilMatrix must not allocate _data and must not register."""
+    V = get_StencilVectorSpace(npts=[4, 5], pads=[1, 2], periods=[False, False])
+    stencil_matrix_memory.clear()
+
+    m = StencilMatrix(V, V, dry_run=True)
+
+    assert m.dry_run
+    assert not hasattr(m, '_data'), "dry_run matrix must not allocate _data"
+    assert stencil_matrix_memory.n_matrices == 0, "dry_run matrix must not register in tracker"
+
+
+def test_stencil_matrix_dry_run_data_shape_and_nbytes():
+    """data_shape and nbytes must be usable on a dry-run matrix without error."""
+    V = get_StencilVectorSpace(npts=[4, 5], pads=[1, 2], periods=[False, False])
+
+    m = StencilMatrix(V, V, dry_run=True)
+
+    assert isinstance(m.data_shape, tuple)
+    assert len(m.data_shape) > 0
+    assert isinstance(m.nbytes, int)
+    assert m.nbytes > 0
+
+
+def test_stencil_matrix_allocated_registers_and_contributes_nbytes():
+    """An allocated StencilMatrix must register and contribute to tracker nbytes."""
+    import gc
+    stencil_matrix_memory.clear()
+    gc.collect()
+
+    V = get_StencilVectorSpace(npts=[4, 5], pads=[1, 2], periods=[False, False])
+    m = StencilMatrix(V, V)
+
+    assert not m.dry_run
+    assert hasattr(m, '_data'), "allocated matrix must have _data"
+    assert stencil_matrix_memory.n_matrices >= 1
+    assert stencil_matrix_memory.nbytes >= m.nbytes
+
+
+def test_stencil_matrix_dry_run_nbytes_matches_allocated():
+    """nbytes reported by a dry-run matrix must equal the nbytes of the equivalent allocated one."""
+    V = get_StencilVectorSpace(npts=[4, 5], pads=[1, 2], periods=[False, False])
+
+    m_dry  = StencilMatrix(V, V, dry_run=True)
+    m_alloc = StencilMatrix(V, V)
+
+    assert m_dry.nbytes == m_alloc.nbytes
+    assert m_dry.data_shape == m_alloc.data_shape
+
+
 #===============================================================================
 # SCRIPT FUNCTIONALITY
 #===============================================================================

{feectools-0.1.6 → feectools-0.1.7}/feectools/linalg/utilities.py

@@ -184,17 +184,29 @@ def _sym_ortho(a, b):
            http://www.stanford.edu/group/SOL/dissertations/sou-cheng-choi-thesis.pdf
     """
     if b == 0:
-        return xp.sign(a), 0, abs(a)
+        return _scalar_sign(a), 0, abs(a)
     elif a == 0:
-        return 0, xp.sign(b), abs(b)
+        return 0, _scalar_sign(b), abs(b)
     elif abs(b) > abs(a):
         tau = a / b
-        s = xp.sign(b) / sqrt(1 + tau * tau)
+        s = _scalar_sign(b) / sqrt(1 + tau * tau)
         c = s * tau
         r = b / s
     else:
         tau = b / a
-        c = xp.sign(a) / sqrt(1+tau*tau)
+        c = _scalar_sign(a) / sqrt(1+tau*tau)
         s = c * tau
         r = a / c
     return c, s, r
+
+#==============================================================================
+def _scalar_sign(x):
+    """
+    Sign of a real Python scalar. `xp.sign` (array_api_compat) requires its
+    argument to expose a `.dtype` attribute, which plain Python floats don't have.
+    """
+    if x > 0:
+        return 1.0
+    elif x < 0:
+        return -1.0
+    return 0.0

{feectools-0.1.6 → feectools-0.1.7/feectools.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: feectools
-Version: 0.1.6
+Version: 0.1.7
 Summary: Slimmed-down fork of Psydac (https://github.com/pyccel/psydac) with less functionality and fewer dependencies.
 Author-email: Psydac development team <psydac@googlegroups.com>
 Maintainer-email: Stefan Possanner <stefan.possanner@ipp.mpg.de>, Max Lindqvist <max.lindqvist@ipp.mpg.de>, Yaman Güçlü <yaman.guclu@gmail.com>, Martin Campos Pinto <martin.campos-pinto@ipp.mpg.de>, Ahmed Ratnani <ratnaniahmed@gmail.com>

{feectools-0.1.6 → feectools-0.1.7}/feectools.egg-info/SOURCES.txt

@@ -75,6 +75,7 @@ feectools/linalg/block.py
 feectools/linalg/direct_solvers.py
 feectools/linalg/fft.py
 feectools/linalg/kron.py
+feectools/linalg/memory.py
 feectools/linalg/solvers.py
 feectools/linalg/sparse.py
 feectools/linalg/stencil.py

{feectools-0.1.6 → feectools-0.1.7}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name            = "feectools"
-version         = "0.1.6"
+version         = "0.1.7"
 description     = "Slimmed-down fork of Psydac (https://github.com/pyccel/psydac) with less functionality and fewer dependencies."
 readme          = "README.md"
 requires-python = ">= 3.10"