Added memory estimates in struphy and feectools

.github/actions/install/install-struphy-editable/action.yml

@@ -14,4 +14,4 @@ runs:
         pip install -U --upgrade-strategy eager -e ".[phys,mpi]"
         pip install -e ".[dev,doc]" 
         pip list
-        struphy -h
+        # struphy -h

.github/actions/install/install-struphy/action.yml

@@ -26,6 +26,6 @@ runs:
         pip list
         struphy -h
         PYTHON=$(which python)
-        STRUPHY_PATH=$($PYTHON -c "import struphy; print(struphy.__path__[0])")
+        STRUPHY_PATH=$($PYTHON -c 'import importlib.util; import os; print(os.path.dirname(importlib.util.find_spec("struphy").origin))')
         echo "Struphy is installed at: $STRUPHY_PATH"
         echo "STRUPHY_PATH=${STRUPHY_PATH}" >> $GITHUB_ENV

.github/actions/install/struphy_in_container/action.yml

@@ -29,6 +29,6 @@ runs:
           pip install -U --upgrade-strategy eager -e ".[phys,mpi]"
           pip install -e ".[doc]"
           PYTHON=$(which python)
-          STRUPHY_PATH=$($PYTHON -c "import struphy; print(struphy.__path__[0])")
+          STRUPHY_PATH=$($PYTHON -c 'import importlib.util; import os; print(os.path.dirname(importlib.util.find_spec("struphy").origin))')
           echo "Struphy is installed at: $STRUPHY_PATH"
           echo "STRUPHY_PATH=${STRUPHY_PATH}" >> $GITHUB_ENV

src/struphy/console/test.py

@@ -57,6 +57,7 @@ def struphy_test(
             f"{LIBPATH}/polar/tests/",
             f"{LIBPATH}/post_processing/tests/",
             f"{LIBPATH}/propagators/tests/",
+            f"{LIBPATH}/simulation/tests/",
         ]
 
         if mpi > 1:

src/struphy/feec/mass.py

@@ -62,6 +62,7 @@ def __init__(
         self._domain = domain
         self._matrix_free = matrix_free
         self._eq_mhd = eq_mhd
+        self._dry_run = False
 
         if self._eq_mhd is None:
             self._eq_mhd = equils.HomogenSlab()
@@ -91,6 +92,77 @@ def matrix_free(self) -> bool:
         """If set to true will not compute the matrix associated with the operators but directly compute the dot product when called."""
         return self._matrix_free
 
+    @property
+    def dry_run(self) -> bool:
+        """If True, mass operators created from now on do not allocate (nor assemble) their
+        stencil matrices; only their sizes are computed. Set temporarily by :meth:`estimate_mem`."""
+        return self._dry_run
+
+    def estimate_mem(
+        self,
+        names: tuple[str] = ("M0", "M1", "M2", "M3", "Mv"),
+        print_report: bool = False,
+    ) -> dict[str, int]:
+        """Estimate the local (per-MPI-rank) memory footprint of mass matrices, in bytes,
+        without allocating them.
+
+        Each requested operator is created exactly as by the corresponding property (same weights,
+        hence the same zero-block detection), but with ``dry_run=True``, so that only the sizes of
+        its stencil matrices are computed, see
+        :attr:`~struphy.feec.mass.WeightedMassOperator.nbytes`. Operators that have already been
+        created (and hence allocated) report their actual size instead; dry-run operators are not
+        kept in the cache.
+
+        Parameters
+        ----------
+        names : tuple[str]
+            Names of the mass operator properties to estimate, e.g. ``("M0", "M1")``.
+
+        print_report : bool
+            Whether to print the breakdown on MPI rank 0.
+
+        Returns
+        -------
+        dict
+            Mapping ``{name: local_bytes}``.
+        """
+        mem = {}
+
+        self._dry_run = True
+        try:
+            for name in names:
+                assert isinstance(getattr(type(self), name, None), property), (
+                    f"'{name}' is not a mass operator property of {type(self).__name__}."
+                )
+                cached = "_" + name
+                was_cached = hasattr(self, cached)
+
+                mem[name] = getattr(self, name).nbytes
+
+                # do not keep a dry-run (unusable) operator in the cache
+                if not was_cached:
+                    delattr(self, cached)
+        finally:
+            self._dry_run = False
+
+        if print_report and (self.derham.comm is None or self.derham.comm.Get_rank() == 0):
+            print("\nESTIMATED MASS MATRIX MEMORY (local, rank 0):")
+            for name, nbytes in mem.items():
+                print(f"  {name}: {nbytes / 1e6:.2f} MB")
+
+        return mem
+
+    def allocated_mem(self) -> dict[str, int]:
+        """Local (per-MPI-rank) memory footprint, in bytes, of the mass matrices that have
+        actually been created so far (i.e. those whose property has been accessed)."""
+        mem = {}
+        for name, method in inspect.getmembers(type(self), predicate=inspect.isdatadescriptor):
+            if isinstance(method, property) and hasattr(self, "_" + name):
+                op = getattr(self, "_" + name)
+                if isinstance(op, WeightedMassOperator):
+                    mem[name] = op.nbytes
+        return mem
+
     def info(self):
         print("The mass matrices of the Derham complex are:")
         self.M0.info()
@@ -897,6 +969,7 @@ def create_weighted_mass(
         weights: tuple | list | str | None = None,
         assemble: bool = False,
         transposed: bool = False,
+        dry_run: bool = None,
     ):
         r"""Weighted mass matrix :math:`V^\alpha_h \to V^\beta_h` with given (matrix-valued) weight function :math:`W(\boldsymbol \eta)`:
 
@@ -950,11 +1023,19 @@ def create_weighted_mass(
         transposed: bool
             Whether to assemble the transposed operator.
 
+        dry_run: bool
+            Whether to create the operator without allocating (and assembling) its stencil matrices,
+            for memory estimation only. If None (default), the value of the ``dry_run`` attribute of
+            this :class:`WeightedMassOperators` object is used, see :meth:`estimate_mem`.
+
         Returns
         -------
         out : A WeightedMassOperator object.
         """
-        logger.debug(f"\nCreating weighted mass matrix {name} from {V_id} to {W_id}.")
+        if dry_run is None:
+            dry_run = self.dry_run
+
+        logger.debug(f"\nCreating weighted mass matrix {name} from {V_id} to {W_id} ({dry_run = }).")
 
         spline_functions = {}
         if isinstance(weights, tuple):  # Case 3 (1D tuple)
@@ -1175,9 +1256,10 @@ def f_call_matrix(e1, e2, e3):
             spline_functions=spline_functions,
             transposed=transposed,
             matrix_free=self.matrix_free,
+            dry_run=dry_run,
         )
 
-        if assemble:
+        if assemble and not dry_run:
             out.assemble()
 
         return out
@@ -1336,6 +1418,14 @@ class WeightedMassOperator(LinOpWithTransp):
 
     matrix_free : bool
         If set to true will not compute the matrix associated with the operator but directly compute the product when called
+
+    dry_run : bool
+        If True, the (potentially large) stencil matrices of the operator are not allocated;
+        only their sizes are computed, see :attr:`nbytes`. The block structure (which blocks are
+        non-zero) is determined in exactly the same way as for a regular operator, but the operator
+        can neither be assembled nor applied. Used to estimate the memory footprint of the FEEC
+        matrices before allocating them, see
+        :meth:`~struphy.feec.mass.WeightedMassOperators.estimate_mem`.
     """
 
     def __init__(
@@ -1353,6 +1443,7 @@ def __init__(
         transposed: bool = False,
         matrix_free: bool = False,
         nquads: tuple | list = None,
+        dry_run: bool = False,
     ):
         logger.debug(f"{derham = }")
         logger.debug(f"{V = }")
@@ -1377,6 +1468,9 @@ def __init__(
         self._V = V
         self._W = W
         self._name = name
+        self._dry_run = dry_run
+
+        assert not (dry_run and transposed), "dry_run=True is not supported for transposed operators."
 
         # spline functions that are used as weights in the operator, to be evaluated at quadrature points
         self._spline_functions = spline_functions if spline_functions is not None else {}
@@ -1507,6 +1601,7 @@ def __init__(
                                 Ws.coeff_space,
                                 backend=PSYDAC_BACKEND_GPYCCEL,
                                 precompiled=True,
+                                dry_run=dry_run,
                             )
                             for Vs in V.spaces
                         ]
@@ -1520,6 +1615,7 @@ def __init__(
                                 Ws.coeff_space,
                                 backend=PSYDAC_BACKEND_GPYCCEL,
                                 precompiled=True,
+                                dry_run=dry_run,
                             )
                             if i != j
                             else None
@@ -1535,6 +1631,7 @@ def __init__(
                                 Ws.coeff_space,
                                 backend=PSYDAC_BACKEND_GPYCCEL,
                                 precompiled=True,
+                                dry_run=dry_run,
                             )
                             if i == j
                             else None
@@ -1613,6 +1710,7 @@ def __init__(
                                     wspace.coeff_space,
                                     backend=PSYDAC_BACKEND_GPYCCEL,
                                     precompiled=True,
+                                    dry_run=dry_run,
                                 ),
                             ]
                         self._weights[-1] += [lambda *etas: 0 * etas[0]]
@@ -1669,6 +1767,7 @@ def __init__(
                                         wspace.coeff_space,
                                         backend=PSYDAC_BACKEND_GPYCCEL,
                                         precompiled=True,
+                                        dry_run=dry_run,
                                     )
                                 ]
 
@@ -1692,6 +1791,7 @@ def __init__(
                             wspace.coeff_space,
                             backend=PSYDAC_BACKEND_GPYCCEL,
                             precompiled=True,
+                            dry_run=dry_run,
                         )
                 else:
                     self._mat = blocks[0][0]
@@ -1726,6 +1826,14 @@ def __init__(
         self._V_extraction_op_T = self._V_extraction_op.T
         self._V_boundary_op_T = self._V_boundary_op.T
 
+        if self._dry_run:
+            # memory estimation only (see the nbytes property): skip the composite operators,
+            # the .dot() temporaries and the assembly kernel; none of them is needed for sizing
+            # and all of them would allocate memory.
+            self._domain = self._mat.domain
+            self._codomain = self._mat.codomain
+            return
+
         # TODO: maybe remove since this is done in the .dot() explicitly
         # build composite linear operators BW * EW * M * EV^T * BV^T, resp. IDV * EV * M^T * EW^T * IDW^T
         if self._transposed:
@@ -1791,6 +1899,19 @@ def codomain_femspace(self):
     def spline_functions(self):
         return self._spline_functions
 
+    @property
+    def dry_run(self) -> bool:
+        """Whether the operator was created for memory estimation only, i.e. without allocating
+        its stencil matrices (in which case it can neither be assembled nor applied)."""
+        return self._dry_run
+
+    @property
+    def nbytes(self) -> int:
+        """Local (per-MPI-rank) memory footprint of the stencil matrices of this operator, in bytes.
+        Also available for operators created with ``dry_run=True``, i.e. before/without allocation.
+        Matrix-free operators do not store a matrix and return 0."""
+        return int(getattr(self._mat, "nbytes", 0))
+
     @property
     def dtype(self):
         return self._dtype
@@ -1987,6 +2108,9 @@ def assemble(self, weights=None, clear=True):
             Whether to first set all data to zero before assembly. If False,
             the new contributions are added to existing ones.
         """
+        assert not self._dry_run, (
+            "A dry-run operator has no matrix data and cannot be assembled (memory estimation only)."
+        )
 
         if self._matrix_free:
             if weights is not None:

src/struphy/feec/memory.py

@@ -0,0 +1,70 @@
+"""Memory footprint of FEEC objects (coefficient spaces, vectors and matrices).
+
+The functions in this module never allocate any of the (potentially large) arrays they measure:
+they either read the metadata of an existing object or the ``nbytes`` of an already allocated one.
+They are used to estimate the memory usage of a simulation before allocating it, see
+:meth:`struphy.simulation.sim.Simulation.estimate_mem`.
+"""
+
+import logging
+
+from feectools.linalg.block import BlockLinearOperator
+
+logger = logging.getLogger("struphy")
+
+
+def vector_nbytes(vector) -> int:
+    """Actual local (per-MPI-rank) memory footprint, in bytes, of an *allocated*
+    :class:`~feectools.linalg.stencil.StencilVector` or
+    :class:`~feectools.linalg.block.BlockVector`, including the ghost/padding regions."""
+    if hasattr(vector, "_data"):
+        return int(vector._data.nbytes)
+    if hasattr(vector, "blocks"):
+        return sum(vector_nbytes(block) for block in vector.blocks)
+    logger.debug(f"Cannot determine the memory footprint of a {type(vector).__name__}, counting 0 bytes.")
+    return 0
+
+
+def coeff_space_nbytes(space, float_size: int = 8) -> int:
+    """Local (per-MPI-rank) memory footprint, in bytes, of a coefficient space
+    (:class:`~feectools.linalg.stencil.StencilVectorSpace` or
+    :class:`~feectools.linalg.block.BlockVectorSpace`), computed from its (metadata-only)
+    local array ``shape`` -- no array is allocated."""
+    if hasattr(space, "spaces"):
+        return sum(coeff_space_nbytes(s, float_size=float_size) for s in space.spaces)
+    nbytes = float_size
+    for n in space.shape:
+        nbytes *= n
+    return int(nbytes)
+
+
+def linop_nbytes(op, _seen: set = None) -> int:
+    """Local (per-MPI-rank) memory footprint, in bytes, of the matrices stored in an
+    (already allocated) linear operator.
+
+    Composite operators (compositions, sums, scalings, block operators, polar operators)
+    are traversed recursively; operators appearing more than once in the tree (for example
+    ``curl.T @ M2 @ curl``) are counted once. Operators that do not store a matrix
+    (identity, matrix-free, boundary/extraction masks) contribute zero.
+    """
+    if _seen is None:
+        _seen = set()
+
+    if op is None or id(op) in _seen:
+        return 0
+    _seen.add(id(op))
+
+    # composite operators: recurse into the children
+    if isinstance(op, BlockLinearOperator):
+        return sum(linop_nbytes(block, _seen) for row in op.blocks for block in row)
+
+    for attr in ("multiplicants", "addends", "mats"):
+        if hasattr(op, attr):
+            return sum(linop_nbytes(child, _seen) for child in getattr(op, attr))
+
+    for attr in ("operator", "tp_operator"):
+        if hasattr(op, attr):
+            return linop_nbytes(getattr(op, attr), _seen)
+
+    # leaf operator: has its own data array (StencilMatrix, StencilDiagonalMatrix, ...)
+    return int(getattr(op, "nbytes", 0))