Added memory estimates

feectools/linalg/block.py

@@ -882,6 +882,14 @@ def nonzero_block_indices(self):
         """
         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/linalg/kron.py

@@ -73,6 +73,12 @@ def ndim( self ):
     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/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 self._matrices))
+
+
+stencil_matrix_memory = MatrixMemoryTracker()

feectools/linalg/stencil.py

@@ -12,6 +12,7 @@
 
 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
@@ -899,8 +900,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)
@@ -912,13 +919,21 @@ def __init__( self, V, W, pads=None , backend=None, precompiled=True):
             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._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(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     = np.zeros(dims+diags, dtype=W.dtype)
-        self._domain   = V
-        self._codomain = W
-        self._ndim     = len(dims)
+        stencil_matrix_memory.register(self)
         self._backend  = backend
         self._precompiled = precompiled
         self._is_T     = False
@@ -969,6 +984,17 @@ def __init__( self, V, W, pads=None , backend=None, precompiled=True):
             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
     #--------------------------------------
@@ -986,6 +1012,28 @@ def codomain(self):
     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 = np.dtype(self._codomain.dtype).itemsize
+        for n in self._data_shape:
+            nbytes *= n
+        return int(nbytes)
+
     # ...
     def dot(self, v, out=None):
         """
@@ -2048,6 +2096,11 @@ def codomain(self):
     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())
 
@@ -2336,6 +2389,12 @@ def codomain(self):
     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):
 

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name            = "feectools"
-version         = "0.1.4"
+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"