diff --git a/.github/actions/install/install-struphy-editable/action.yml b/.github/actions/install/install-struphy-editable/action.yml index 5b1c975ca..50e088c17 100644 --- a/.github/actions/install/install-struphy-editable/action.yml +++ b/.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 diff --git a/.github/actions/install/install-struphy/action.yml b/.github/actions/install/install-struphy/action.yml index 13ea329a6..0050adebe 100644 --- a/.github/actions/install/install-struphy/action.yml +++ b/.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 diff --git a/.github/actions/install/struphy_in_container/action.yml b/.github/actions/install/struphy_in_container/action.yml index fb591491c..5f4405ccb 100644 --- a/.github/actions/install/struphy_in_container/action.yml +++ b/.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 \ No newline at end of file diff --git a/feectools b/feectools index 74c88399a..efcfe31f3 160000 --- a/feectools +++ b/feectools @@ -1 +1 @@ -Subproject commit 74c88399a7adb72a07c5bdc30d6e352f10e851e4 +Subproject commit efcfe31f39f0e5cf2c7f331058dd4a8ce1c6eb82 diff --git a/src/struphy/console/test.py b/src/struphy/console/test.py index 25718fb83..fc4d190a7 100644 --- a/src/struphy/console/test.py +++ b/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: diff --git a/src/struphy/feec/mass.py b/src/struphy/feec/mass.py index 908ab0762..c5ea7cc1d 100644 --- a/src/struphy/feec/mass.py +++ b/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: diff --git a/src/struphy/feec/memory.py b/src/struphy/feec/memory.py new file mode 100644 index 000000000..bbfc17c3e --- /dev/null +++ b/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)) diff --git a/src/struphy/feec/tests/test_estimate_mem.py b/src/struphy/feec/tests/test_estimate_mem.py new file mode 100644 index 000000000..75753b4a0 --- /dev/null +++ b/src/struphy/feec/tests/test_estimate_mem.py @@ -0,0 +1,77 @@ +import pytest +from feectools.linalg.memory import stencil_matrix_memory +from feectools.linalg.stencil import StencilMatrix + +from struphy import DerhamOptions, domains, grids +from struphy.feec.mass import WeightedMassOperators +from struphy.feec.memory import linop_nbytes +from struphy.feec.psydac_derham import Derham + + +@pytest.fixture +def derham(): + return Derham( + grids.TensorProductGrid(num_elements=(4, 5, 6)), + DerhamOptions(degree=(2, 2, 3)), + ) + + +def test_stencil_matrix_dry_run(derham): + """A dry-run StencilMatrix reports the same size as an allocated one, but allocates nothing.""" + V = derham.coeff_spaces["3"] + + n_before = stencil_matrix_memory.n_matrices + dry = StencilMatrix(V, V, dry_run=True) + assert dry.dry_run + assert stencil_matrix_memory.n_matrices == n_before, "a dry-run matrix must not be registered" + + # the data array is not there, and saying so is part of the deal + with pytest.raises(AttributeError, match="dry_run"): + dry._data + + mat = StencilMatrix(V, V) + assert not mat.dry_run + assert mat.data_shape == dry.data_shape + assert mat.nbytes == dry.nbytes == mat._data.nbytes + + +def test_mass_ops_estimate_mem_matches_allocation(derham): + """The dry-run estimate of the standard mass matrices must match the real allocation.""" + mass_ops = WeightedMassOperators(derham, domains.Cuboid()) + + names = ("M0", "M1", "M2", "M3", "Mv") + estimated = mass_ops.estimate_mem(names=names) + + # nothing was created (nor allocated) by the estimate + assert mass_ops.allocated_mem() == {} + assert not mass_ops.dry_run + + # now allocate and assemble for real + for name in names: + getattr(mass_ops, name) + + allocated = mass_ops.allocated_mem() + assert set(allocated) == set(names) + for name in names: + assert estimated[name] == allocated[name], f"{name}: {estimated[name]} != {allocated[name]}" + assert estimated[name] > 0 + + # and the estimate of an already allocated operator is its actual size + assert mass_ops.estimate_mem(names=("M1",))["M1"] == allocated["M1"] + + +def test_mass_ops_zero_blocks_are_not_counted(derham): + """On a Cartesian domain the metric is diagonal, so M1 must be a 3x3 block matrix with + only 3 non-zero blocks -- the estimate has to see that, too.""" + mass_ops = WeightedMassOperators(derham, domains.Cuboid()) + + estimated = mass_ops.estimate_mem(names=("M1",))["M1"] + blocks = [block for row in mass_ops.M1._mat.blocks for block in row] + + assert sum(block is not None for block in blocks) == 3 + assert estimated == sum(block.nbytes for block in blocks if block is not None) + + +def test_derivative_matrices_are_matrix_free(derham): + """The derivative operators of the Derham sequence do not store any matrix data.""" + assert sum(linop_nbytes(op) for op in (derham.grad, derham.curl, derham.div)) == 0 diff --git a/src/struphy/models/tests/test_estimate_mem.py b/src/struphy/models/tests/test_estimate_mem.py new file mode 100644 index 000000000..52dd021d6 --- /dev/null +++ b/src/struphy/models/tests/test_estimate_mem.py @@ -0,0 +1,117 @@ +from types import SimpleNamespace +from unittest.mock import patch + +from struphy.models.variables import PICVariable +from struphy.simulation import sim as sim_module +from struphy.simulation.sim import Simulation + + +def test_simulation_estimate_mem_returns_total(): + class DummyDerham: + def __init__(self, grid, derham_opts, comm=None, domain=None): + self.grid = grid + self.derham_opts = derham_opts + self.comm = comm + self.domain = domain + # derivative operators are matrix-free (no data), see linop_nbytes + self.grad = None + self.curl = None + self.div = None + + class DummyMassOperators: + def __init__(self, derham, domain, eq_mhd=None): + assert isinstance(derham, DummyDerham) + + def estimate_mem(self): + return {"M1": 7} + + class DummyFEECVariable: + def estimate_mem(self, derham): + assert isinstance(derham, DummyDerham) + return 10 + + class DummyPICVariable: + def estimate_mem(self, clone_config=None, derham=None, domain=None, equil=None): + assert isinstance(derham, DummyDerham) + return 20 + + class DummySPHVariable: + def estimate_mem(self, derham=None, domain=None, equil=None): + assert isinstance(derham, DummyDerham) + return 30 + + sim = Simulation.__new__(Simulation) + sim._grid = object() + sim._derham_opts = object() + sim._clone_config = None + sim._domain = object() + sim._equil = object() + sim.comm = None + sim.rank = 0 + sim.comm_size = 1 + sim._model = SimpleNamespace( + field_species={"f": SimpleNamespace(variables={"u": DummyFEECVariable()})}, + fluid_species={}, + particle_species={"p": SimpleNamespace(variables={"v": DummyPICVariable(), "w": DummySPHVariable()})}, + diagnostic_species={"d": SimpleNamespace(variables={"z": DummyFEECVariable()})}, + ) + + with ( + patch.object(sim_module, "Derham", DummyDerham), + patch.object(sim_module, "WeightedMassOperators", DummyMassOperators), + patch.object(sim_module, "FEECVariable", DummyFEECVariable), + patch.object(sim_module, "PICVariable", DummyPICVariable), + patch.object(sim_module, "SPHVariable", DummySPHVariable), + ): + mem = sim.estimate_mem(print_report=False) + + assert mem["f.u"] == 10 + assert mem["p.v"] == 20 + assert mem["p.w"] == 30 + assert mem["d.z"] == 10 + assert mem["matrices.derivatives"] == 0 + assert mem["matrices.M1"] == 7 + assert mem["total"] == 77 + assert mem["total"] >= 0 + + +def test_picvariable_estimate_mem_uses_dry_run_particles(): + class DummyKineticBackground: + pass + + class DummyParticles: + last_instance = None + last_kwargs = None + + def __init__(self, **kwargs): + DummyParticles.last_kwargs = kwargs + DummyParticles.last_instance = self + self.Np = 16 + self.n_cols = 12 + self.nbytes_local = 128 + if not kwargs.get("dry_run", False): + self.markers = object() + + var = PICVariable(space="Particles6D") + var._species = SimpleNamespace( + loading_params=SimpleNamespace(), + weights_params=SimpleNamespace(), + boundary_params=SimpleNamespace(), + sorting_params=SimpleNamespace(), + bufsize=0.25, + equation_params={}, + saving_params=SimpleNamespace(n_markers=0), + ) + var._backgrounds = DummyKineticBackground() + var._initial_condition = var._backgrounds + var._n_as_volume_form = False + + with ( + patch("struphy.models.variables.KineticBackground", DummyKineticBackground), + patch("struphy.models.variables.particles.Particles6D", DummyParticles), + ): + nbytes = var.estimate_mem() + + assert DummyParticles.last_kwargs["dry_run"] is True + assert not hasattr(DummyParticles.last_instance, "markers") + assert nbytes == 128 diff --git a/src/struphy/models/variables.py b/src/struphy/models/variables.py index 2fa50e6a3..2ceb8e71c 100644 --- a/src/struphy/models/variables.py +++ b/src/struphy/models/variables.py @@ -10,6 +10,7 @@ from feectools.ddm.mpi import mpi as MPI from struphy.feec.linear_operators import BoundaryOperator +from struphy.feec.memory import coeff_space_nbytes from struphy.feec.psydac_derham import Derham, SplineFunction from struphy.fields_background.base import FluidEquilibrium from struphy.fields_background.projected_equils import ProjectedFluidEquilibrium @@ -84,6 +85,11 @@ def space(self): def allocate(self): """Alocate object and memory for variable.""" + @abstractmethod + def estimate_mem(self) -> int: + """Estimate the local (per-MPI-rank) memory footprint of this variable, in bytes, + without actually allocating it. Can be called before :meth:`allocate`.""" + def __repr__(self): return f"{self.__class__.__name__} ({self.space})" @@ -396,6 +402,18 @@ def allocate( self.compute_boundary_spline() + def estimate_mem(self, derham: Derham) -> int: + """Estimate the local (per-MPI-rank) memory footprint of the spline coefficient vector(s) + of this variable, in bytes, without creating a :class:`~struphy.feec.psydac_derham.SplineFunction`. + + Uses the (cheap, metadata-only) local array shape of ``derham.coeff_spaces[self.space]``. + If a lifting function is set, ``allocate()`` additionally creates ``spline_lift``, ``spline_0`` + and ``boundary_spline`` of the same space, so the estimate is scaled by a factor of 4.""" + nbytes = coeff_space_nbytes(derham.coeff_spaces[self.space]) + if self.lifting_function is not None: + nbytes *= 4 # spline + spline_lift + spline_0 + boundary_spline + return nbytes + def compute_boundary_spline(self, spline_lift: SplineFunction | None = None): """Compute boundary_spline = spline_lift - spline_0. If spline_lift is None, uses self.spline_lift from the initial condition. This method can be used to update the boundary spline during the simulation if the lifting function changes in time.""" @@ -608,6 +626,74 @@ def allocate( # other data (wave-particle power exchange, etc.) # TODO + def estimate_mem( + self, + clone_config: CloneConfig = None, + derham: Derham = None, + domain: Domain = None, + equil: FluidEquilibrium = None, + projected_equil: ProjectedFluidEquilibrium = None, + ) -> int: + """Estimate the local (per-MPI-rank) memory footprint of this variable's marker arrays, in bytes, + without allocating them. + + Constructs the same :class:`~struphy.pic.base.Particles` object as :meth:`allocate` would, but with + ``dry_run=True`` so that only the marker array sizing (:attr:`~struphy.pic.base.Particles.n_rows`, + :attr:`~struphy.pic.base.Particles.n_cols`) is computed.""" + assert isinstance(self.backgrounds, KineticBackground), ( + "List input not allowed, you can sum Kineticbackgrounds before passing them to add_background." + ) + + if derham is None: + domain_decomp = None + else: + domain_array = derham.domain_array + nprocs = derham.domain_decomposition.nprocs + domain_decomp = (domain_array, nprocs) + + kinetic_class = getattr(particles, self.space) + + comm_world = MPI.COMM_WORLD + if comm_world.Get_size() == 1: + comm_world = None + + dummy_particles: Particles = kinetic_class( + comm_world=comm_world, + clone_config=clone_config, + domain_decomp=domain_decomp, + name=self.species.__class__.__name__, + loading_params=self.species.loading_params, + weights_params=self.species.weights_params, + boundary_params=self.species.boundary_params, + sorting_params=self.species.sorting_params, + bufsize=self.species.bufsize, + domain=domain, + equil=equil, + projected_equil=projected_equil, + background=self.backgrounds, + initial_condition=self.initial_condition, + n_as_volume_form=self.n_as_volume_form, + equation_params=self.species.equation_params, + dry_run=True, + ) + + nbytes = dummy_particles.nbytes_local + + # marker array for saving trajectories (approximated with Np since n_mks_global + # is only known after markers have actually been drawn) + n_markers = self.species.saving_params.n_markers + if isinstance(n_markers, float): + if n_markers > 1.0: + n_to_save = int(n_markers) + else: + n_to_save = int(dummy_particles.Np * n_markers) + else: + n_to_save = n_markers + if n_to_save > 0: + nbytes += n_to_save * dummy_particles.n_cols * 8 + + return nbytes + @property def n_to_save(self) -> int: return self._n_to_save @@ -833,6 +919,70 @@ def allocate( # other data (wave-particle power exchange, etc.) # TODO + def estimate_mem( + self, + derham: Derham = None, + domain: Domain = None, + equil: FluidEquilibrium = None, + projected_equil: ProjectedFluidEquilibrium = None, + ) -> int: + """Estimate the local (per-MPI-rank) memory footprint of this variable's marker arrays, in bytes, + without allocating them. + + Constructs the same :class:`~struphy.pic.particles.ParticlesSPH` object as :meth:`allocate` would, + but with ``dry_run=True`` so that only the marker array sizing + (:attr:`~struphy.pic.base.Particles.n_rows`, :attr:`~struphy.pic.base.Particles.n_cols`) is computed.""" + assert isinstance(self.backgrounds, FluidEquilibrium), ( + "List input not allowed; you can sum FluidEquilibrium objects before passing them to add_background." + ) + + if derham is None: + domain_decomp = None + else: + domain_array = derham.domain_array + nprocs = derham.domain_decomposition.nprocs + domain_decomp = (domain_array, nprocs) + + comm_world = MPI.COMM_WORLD + if comm_world.Get_size() == 1: + comm_world = None + + dummy_particles: ParticlesSPH = ParticlesSPH( + comm_world=comm_world, + domain_decomp=domain_decomp, + name=self.species.__class__.__name__, + loading_params=self.species.loading_params, + weights_params=self.species.weights_params, + boundary_params=self.species.boundary_params, + sorting_params=self.species.sorting_params, + bufsize=self.species.bufsize, + domain=domain, + equil=equil, + projected_equil=projected_equil, + background=self.backgrounds, + n_as_volume_form=self.n_as_volume_form, + perturbations=self.perturbations, + equation_params=self.species.equation_params, + dry_run=True, + ) + + nbytes = dummy_particles.nbytes_local + + # marker array for saving trajectories (approximated with Np since n_mks_global + # is only known after markers have actually been drawn) + n_markers = self.species.saving_params.n_markers + if isinstance(n_markers, float): + if n_markers > 1.0: + n_to_save = int(n_markers) + else: + n_to_save = int(dummy_particles.Np * n_markers) + else: + n_to_save = n_markers + if n_to_save > 0: + nbytes += n_to_save * dummy_particles.n_cols * 8 + + return nbytes + @property def n_to_save(self) -> int: return self._n_to_save diff --git a/src/struphy/pic/base.py b/src/struphy/pic/base.py index e86b9f743..be325e985 100644 --- a/src/struphy/pic/base.py +++ b/src/struphy/pic/base.py @@ -97,6 +97,7 @@ def __init__( perturbations: dict[str, Perturbation] = None, n_as_volume_form: bool = False, equation_params: dict = None, + dry_run: bool = False, ): r""" The marker information is stored in a 2D numpy array. @@ -172,6 +173,12 @@ def __init__( equation_params : dict Normalization parameters (epsilon, alpha, ...) + + dry_run : bool + If True, only compute the sizing of the marker array (:attr:`n_rows`, :attr:`n_cols`, ...) + and return early, without allocating any of the (potentially large) marker/sorting/buffer + arrays. Used by :attr:`nbytes_local` to estimate the memory footprint before actually + allocating the particles, see :meth:`~struphy.models.variables.PICVariable.estimate_mem`. """ self._clone_config = clone_config @@ -274,7 +281,10 @@ def __init__( # create marker array self._bufsize = bufsize - self._allocate_marker_array() + self._allocate_marker_array(dry_run=dry_run) + + if dry_run: + return # boundary conditions bc = boundary_params.bc @@ -363,7 +373,7 @@ def __init__( self._generate_sampling_moments() # create buffers for mpi_sort_markers - self._sorting_etas = xp.zeros(self.markers.shape, dtype=float) + self._sorting_etas = xp.zeros((self.markers.shape[0], 3), dtype=float) self._is_on_proc_domain = xp.zeros((self.markers.shape[0], 3), dtype=bool) self._can_stay = xp.zeros(self.markers.shape[0], dtype=bool) self._reqs = [None] * self.mpi_size @@ -471,6 +481,27 @@ def n_rows(self): self._allocate_marker_array() return self._n_rows + @property + def nbytes_local(self) -> int: + """Estimated local (per-MPI-rank) memory footprint, in bytes, of all marker-related arrays + (markers, sorting buffers, lost-marker container). Only depends on :attr:`n_rows` and + :attr:`n_cols`, so it is valid whether or not the arrays were actually allocated + (see the ``dry_run`` argument of :meth:`__init__`).""" + float_size = 8 # dtype=float + bool_size = 1 # dtype=bool + n_rows = self.n_rows + n_cols = self.n_cols + + nbytes = 0 + nbytes += n_rows * n_cols * float_size # markers + nbytes += n_rows * 3 * float_size # sorting_etas (mpi_sort_markers buffer) + nbytes += n_rows * 3 * bool_size # is_on_proc_domain + nbytes += n_rows * bool_size # can_stay + # holes, ghost_particles, valid_mks, is_outside_right, is_outside_left, is_outside + nbytes += n_rows * bool_size * 6 + nbytes += int(n_rows * 0.5) * 10 * float_size # lost_markers + return int(nbytes) + @property def kinds(self): """Name of the class.""" @@ -1106,8 +1137,11 @@ def _n_mks_load_and_Np_per_clone(self): return n_mks_load, Np_per_clone - def _allocate_marker_array(self): - """Create marker array :attr:`~struphy.pic.base.Particles.markers`.""" + def _allocate_marker_array(self, dry_run: bool = False): + """Create marker array :attr:`~struphy.pic.base.Particles.markers`. + + If dry_run is True, only :attr:`n_rows` (and :attr:`n_cols`) are computed and no array + is actually allocated; see :attr:`nbytes_local`.""" if not hasattr(self, "_n_mks_load"): self._n_mks_load, self._Np_per_clone = self._n_mks_load_and_Np_per_clone() @@ -1115,8 +1149,17 @@ def _allocate_marker_array(self): n_mks_load_loc = self.n_mks_load[self._mpi_rank] bufsize = self.bufsize + 1.0 / xp.sqrt(n_mks_load_loc) - # allocate markers array (3 x positions, vdim x velocities, weight, s0, w0, ..., ID) with buffer + # size of markers array (3 x positions, vdim x velocities, weight, s0, w0, ..., ID) with buffer self._n_rows = round(n_mks_load_loc * (1 + bufsize)) + + # Have at least 3 spare places in markers array + assert self.first_free_idx + 2 < self.n_cols - 2, ( + f"{self.first_free_idx + 2} is not smaller than {self.n_cols - 2 =}; not enough columns in marker array !!" + ) + + if dry_run: + return + self._markers = xp.zeros((self.n_rows, self.n_cols), dtype=float) # allocate auxiliary arrays @@ -1145,11 +1188,6 @@ def _allocate_marker_array(self): self.first_free_idx, ) - # Have at least 3 spare places in markers array - assert self.args_markers.first_free_idx + 2 < self.n_cols - 1, ( - f"{self.args_markers.first_free_idx + 2} is not smaller than {self.n_cols - 1 =}; not enough columns in marker array !!" - ) - def _initialize_sorting_boxes(self): """Initializes the sorting boxes. @@ -4333,10 +4371,11 @@ def sendrecv_determine_mtbs( assert alpha.size == 3 assert xp.all(alpha >= 0.0) and xp.all(alpha <= 1.0) bi = self.first_pusher_idx - self._sorting_etas = xp.mod( + xp.mod( alpha * (self.markers[:, :3] + self.markers[:, bi + 3 + self.vdim : bi + 3 + self.vdim + 3]) + (1.0 - alpha) * self.markers[:, bi : bi + 3], 1.0, + out=self._sorting_etas, ) # check which particles are on the current process domain diff --git a/src/struphy/pic/tests/test_estimate_mem.py b/src/struphy/pic/tests/test_estimate_mem.py new file mode 100644 index 000000000..389343977 --- /dev/null +++ b/src/struphy/pic/tests/test_estimate_mem.py @@ -0,0 +1,110 @@ +import pytest +from feectools.ddm.mpi import mpi as MPI + +from struphy import BoundaryParameters, LoadingParameters, SortingParameters, WeightsParameters, domains +from struphy.feec.psydac_derham import Derham +from struphy.io.options import DerhamOptions +from struphy.pic.particles import Particles6D +from struphy.topology.grids import TensorProductGrid + + +def _make_domain_decomp(mpi_comm, num_elements=(8, 6, 4), degree=(2, 2, 2)): + domain = domains.Cuboid() + derham = Derham(TensorProductGrid(num_elements=num_elements), DerhamOptions(degree=degree), comm=mpi_comm) + domain_decomp = (derham.domain_array, derham.domain_decomposition.nprocs) + return domain, domain_decomp + + +@pytest.mark.parametrize("Np", [1000, 12345]) +def test_dry_run_does_not_allocate_markers(Np): + """dry_run=True must compute the marker array sizing (n_rows, n_cols) without + allocating the (potentially large) marker/sorting/buffer arrays.""" + mpi_comm = MPI.COMM_WORLD + domain, domain_decomp = _make_domain_decomp(mpi_comm) + + loading_params = LoadingParameters(Np=Np, seed=1234) + boundary_params = BoundaryParameters() + sorting_params = SortingParameters(do_sort=False) + + particles = Particles6D( + comm_world=mpi_comm, + loading_params=loading_params, + boundary_params=boundary_params, + sorting_params=sorting_params, + domain_decomp=domain_decomp, + domain=domain, + dry_run=True, + ) + + # sizing info is available ... + assert particles.n_rows > 0 + assert particles.n_cols > 0 + + # ... but none of the (large) arrays that allocate() would create exist + for attr in ( + "_markers", + "_holes", + "_ghost_particles", + "_valid_mks", + "_is_outside_right", + "_is_outside_left", + "_is_outside", + "_lost_markers", + "_sorting_etas", + "_is_on_proc_domain", + "_can_stay", + ): + assert not hasattr(particles, attr), f"dry_run=True should not create '{attr}'" + + +@pytest.mark.parametrize("Np", [1000, 12345]) +def test_nbytes_local_matches_real_allocation(Np): + """The dry_run estimate (nbytes_local) must equal the real memory footprint once the + particles are actually allocated (dry_run=False), since both use the exact same + n_rows/n_cols sizing and the same (fixed) list of marker-related arrays.""" + mpi_comm = MPI.COMM_WORLD + domain, domain_decomp = _make_domain_decomp(mpi_comm) + + loading_params = LoadingParameters(Np=Np, seed=1234) + boundary_params = BoundaryParameters() + sorting_params = SortingParameters(do_sort=False) + weights_params = WeightsParameters() + + common_kwargs = dict( + comm_world=mpi_comm, + loading_params=loading_params, + weights_params=weights_params, + boundary_params=boundary_params, + sorting_params=sorting_params, + domain_decomp=domain_decomp, + domain=domain, + ) + + dry = Particles6D(**common_kwargs, dry_run=True) + real = Particles6D(**common_kwargs, dry_run=False) + real.draw_markers(sort=False) + + assert dry.n_rows == real.n_rows + assert dry.n_cols == real.n_cols + + real_nbytes = ( + real.markers.nbytes + + real._sorting_etas.nbytes + + real._is_on_proc_domain.nbytes + + real._can_stay.nbytes + + real._holes.nbytes + + real._ghost_particles.nbytes + + real._valid_mks.nbytes + + real._is_outside_right.nbytes + + real._is_outside_left.nbytes + + real._is_outside.nbytes + + real._lost_markers.nbytes + ) + + assert dry.nbytes_local == real_nbytes + assert real.nbytes_local == real_nbytes + + +if __name__ == "__main__": + test_dry_run_does_not_allocate_markers(Np=1000) + test_nbytes_local_matches_real_allocation(Np=1000) diff --git a/src/struphy/simulation/sim.py b/src/struphy/simulation/sim.py index 2c683fab3..6127a2b5e 100644 --- a/src/struphy/simulation/sim.py +++ b/src/struphy/simulation/sim.py @@ -14,6 +14,7 @@ import yaml from feectools.ddm.mpi import MockMPI from feectools.ddm.mpi import mpi as MPI +from feectools.linalg.memory import stencil_matrix_memory from feectools.linalg.stencil import StencilVector from line_profiler import profile from pyevtk.hl import gridToVTK @@ -36,6 +37,7 @@ # core imports from struphy.feec.basis_projection_ops import BasisProjectionOperators from struphy.feec.mass import WeightedMassOperators +from struphy.feec.memory import linop_nbytes, vector_nbytes from struphy.feec.psydac_derham import Derham from struphy.fields_background.base import ( FluidEquilibrium, @@ -290,6 +292,197 @@ def allocate(self): logger.debug("... Done.") + def estimate_mem(self, print_report: bool = True) -> dict: + """Estimate the memory footprint of all model variables and FEEC matrices, in bytes, + BEFORE calling :meth:`allocate`. + + Builds a throwaway Derham sequence (cheap: metadata only, no mass/basis operators) to obtain + the FEEC coefficient space sizes and MPI domain decomposition, then calls ``estimate_mem()`` on + every model variable, mirroring the loop in :meth:`_allocate_variables`. The throwaway Derham is + discarded afterward -- calling :meth:`allocate` still performs the full allocation from scratch. + + On top of the variables, the following FEEC matrices are estimated (they are usually much + larger than the spline coefficient vectors): + + * the derivative matrices (grad, curl, div) of the Derham sequence (matrix-free, hence + essentially free of charge), + * the standard mass matrices M0, M1, M2, M3 and Mv, created with ``dry_run=True`` so that + only their sizes (including the zero-block structure) are computed, see + :meth:`~struphy.feec.mass.WeightedMassOperators.estimate_mem`. + + Note + ---- + Which matrices a model really allocates is only known once the propagators have been + allocated (they fetch and build their operators in ``allocate()``, partly under names + assembled at run time). The standard mass matrices are therefore used as a baseline: + a model may not need all of them, but it may also build additional weighted mass matrices + (e.g. ``M2n``), basis projection operators and preconditioners which are *not* included here. + Use :meth:`report_mem` after :meth:`allocate` for the exact numbers. + + Parameters + ---------- + print_report : bool + If True (default), print a breakdown of the estimated memory usage of each variable on + MPI rank 0: both local (this rank) and global (summed over all ranks) values. + + Returns + ------- + dict + Mapping ``{"species.variable": local_bytes}`` for the model variables and + ``{"matrices.": local_bytes}`` for the FEEC matrices (bytes on the *current* + MPI rank), plus a ``"total"`` entry with the local sum over all entries. + """ + logger.debug("\nEstimating memory usage ...") + + if self.grid is None or self.derham_opts is None: + raise RuntimeError( + "Simulation.estimate_mem() requires 'grid' and 'derham_opts' to be set (needed to build a Derham sequence for FEEC variable sizing)." + ) + + if self.clone_config is None: + derham_comm = MPI.COMM_WORLD + else: + derham_comm = self.clone_config.sub_comm + + derham = Derham( + self.grid, + self.derham_opts, + comm=derham_comm, + domain=self.domain, + ) + + mem = {} + + if self.model.field_species: + for species, spec in self.model.field_species.items(): + for k, v in spec.variables.items(): + assert isinstance(v, FEECVariable) + mem[f"{species}.{k}"] = v.estimate_mem(derham=derham) + + if self.model.fluid_species: + for species, spec in self.model.fluid_species.items(): + for k, v in spec.variables.items(): + assert isinstance(v, FEECVariable) + mem[f"{species}.{k}"] = v.estimate_mem(derham=derham) + + if self.model.particle_species: + for species, spec in self.model.particle_species.items(): + for k, v in spec.variables.items(): + if isinstance(v, PICVariable): + mem[f"{species}.{k}"] = v.estimate_mem( + clone_config=self.clone_config, + derham=derham, + domain=self.domain, + equil=self.equil, + ) + elif isinstance(v, SPHVariable): + mem[f"{species}.{k}"] = v.estimate_mem( + derham=derham, + domain=self.domain, + equil=self.equil, + ) + + if self.model.diagnostic_species: + for species, spec in self.model.diagnostic_species.items(): + for k, v in spec.variables.items(): + assert isinstance(v, FEECVariable) + mem[f"{species}.{k}"] = v.estimate_mem(derham=derham) + + # FEEC matrices: derivative matrices (already allocated by the throwaway Derham) ... + mem["matrices.derivatives"] = sum(linop_nbytes(op) for op in (derham.grad, derham.curl, derham.div)) + + # ... and the standard mass matrices (dry run, i.e. sized but not allocated) + mass_ops = WeightedMassOperators(derham, self.domain, eq_mhd=self.equil) + for name, nbytes in mass_ops.estimate_mem().items(): + mem[f"matrices.{name}"] = nbytes + + total_local = sum(mem.values()) + mem["total"] = total_local + + if print_report: + if self.comm is not None: + total_global = self.comm.allreduce(total_local, op=MPI.SUM) + else: + total_global = total_local + + if self.rank == 0: + print("\nESTIMATED MEMORY USAGE (before allocate()):") + for name, nbytes in mem.items(): + if name == "total": + continue + print(f" {name}: {nbytes / 1e6:.2f} MB (local, rank 0)") + print( + f" TOTAL: {total_local / 1e6:.2f} MB (local, rank 0), " + f"{total_global / 1e6:.2f} MB (global, summed over {self.comm_size} rank(s))" + ) + + logger.debug("... Done.") + + return mem + + def report_mem(self, print_report: bool = True) -> dict: + """Actual local (per-MPI-rank) memory footprint of the big arrays, in bytes, + AFTER calling :meth:`allocate`. + + In contrast to :meth:`estimate_mem`, nothing is estimated here: the spline coefficient + vectors and marker arrays are measured on the allocated objects, and the FEEC matrices are + obtained from the (weak) registry of all allocated stencil matrices, + :data:`~feectools.linalg.memory.stencil_matrix_memory`. The matrix number therefore + includes mass matrices, basis projection operators, preconditioners and any other + stencil matrix built by the propagators. + + Parameters + ---------- + print_report : bool + If True (default), print the breakdown on MPI rank 0. + + Returns + ------- + dict + Mapping ``{"feec_matrices": ..., "spline_coeffs": ..., "markers": ..., "total": ...}`` + with the local number of bytes. + """ + mem = {"feec_matrices": stencil_matrix_memory.nbytes, "spline_coeffs": 0, "markers": 0} + + for species in ( + self.model.field_species, + self.model.fluid_species, + self.model.particle_species, + self.model.diagnostic_species, + ): + if not species: + continue + for spec in species.values(): + for v in spec.variables.values(): + if isinstance(v, FEECVariable): + mem["spline_coeffs"] += vector_nbytes(v.spline.vector) + elif isinstance(v, (PICVariable, SPHVariable)): + mem["markers"] += v.particles.nbytes_local + if v.n_to_save > 0: + mem["markers"] += v.saved_markers.nbytes + + total_local = sum(mem.values()) + mem["total"] = total_local + + if print_report: + if self.comm is not None: + total_global = self.comm.allreduce(total_local, op=MPI.SUM) + else: + total_global = total_local + + if self.rank == 0: + print("\nMEMORY USAGE (after allocate()):") + for name, nbytes in mem.items(): + if name == "total": + continue + print(f" {name}: {nbytes / 1e6:.2f} MB (local, rank 0)") + print( + f" TOTAL: {total_local / 1e6:.2f} MB (local, rank 0), " + f"{total_global / 1e6:.2f} MB (global, summed over {self.comm_size} rank(s))" + ) + + return mem + def save_geometry_and_equil_vtk(self): """Write a VTK file with geometry and (projected) equilibrium fields. diff --git a/src/struphy/simulation/tests/__init__.py b/src/struphy/simulation/tests/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/src/struphy/simulation/tests/test_estimate_mem.py b/src/struphy/simulation/tests/test_estimate_mem.py new file mode 100644 index 000000000..34a3aa74f --- /dev/null +++ b/src/struphy/simulation/tests/test_estimate_mem.py @@ -0,0 +1,165 @@ +import os +import shutil + +import pytest + +from struphy import ( + BoundaryParameters, + DerhamOptions, + EnvironmentOptions, + LoadingParameters, + SavingParameters, + Simulation, + SortingParameters, + Time, + WeightsParameters, + domains, + grids, + maxwellians, + perturbations, +) +from struphy.models import Maxwell, VlasovAmpereOneSpecies + + +def _real_vector_nbytes(vector): + """Real (backing-array, incl. ghost/padding regions) memory footprint of a psydac + StencilVector or BlockVector, i.e. what actually got allocated in RAM.""" + if hasattr(vector, "_data"): + return vector._data.nbytes + return sum(_real_vector_nbytes(block) for block in vector.blocks) + + +@pytest.fixture +def out_folders(tmp_path): + out = os.path.join(str(tmp_path), "struphy_estimate_mem_tests") + yield out + shutil.rmtree(out, ignore_errors=True) + + +def test_estimate_mem_feec_only_matches_allocation(out_folders): + """estimate_mem() must be callable before allocate(), and its FEEC estimates must + exactly match the real (backing-array) memory footprint after allocate().""" + model = Maxwell() + + env = EnvironmentOptions(out_folders=out_folders, sim_folder="light_wave_1d") + time_opts = Time(dt=0.05, Tend=50.0) + domain = domains.Cuboid(r3=20.0) + grid = grids.TensorProductGrid(num_elements=(1, 1, 32)) + derham_opts = DerhamOptions(degree=(1, 1, 3)) + + model.propagators.maxwell.options = model.propagators.maxwell.Options(algo="explicit") + model.em_fields.e_field.add_perturbation(perturbations.Noise(amp=0.1, comp=0, seed=123)) + + sim = Simulation( + model=model, + env=env, + time_opts=time_opts, + domain=domain, + grid=grid, + derham_opts=derham_opts, + ) + + # estimate_mem() must not require (or trigger) allocate() first + assert not hasattr(sim, "_derham") + mem_before = sim.estimate_mem(print_report=False) + # and must not leave the simulation in an allocated state + assert not hasattr(sim, "_derham") + + sim.allocate() + + for species, spec in sim.model.field_species.items(): + for k, v in spec.variables.items(): + estimated = mem_before[f"{species}.{k}"] + actual = _real_vector_nbytes(v.spline.vector) + assert estimated == actual, f"{species}.{k}: estimated {estimated} != actual {actual}" + + assert mem_before["total"] == sum(v for k, v in mem_before.items() if k != "total") + + # the FEEC matrices (much bigger than the coefficient vectors) are part of the estimate ... + matrices = {k: v for k, v in mem_before.items() if k.startswith("matrices.")} + assert set(matrices) == {f"matrices.{name}" for name in ("derivatives", "M0", "M1", "M2", "M3", "Mv")} + assert matrices["matrices.M1"] > mem_before["em_fields.e_field"] + + # ... and the mass matrices this model really uses are estimated exactly + allocated = sim.mass_ops.allocated_mem() + assert set(allocated) == {"M1", "M2"} + for name, nbytes in allocated.items(): + assert mem_before[f"matrices.{name}"] == nbytes + + # report_mem() sees all allocated stencil matrices, i.e. at least the mass matrices + report = sim.report_mem(print_report=False) + assert report["feec_matrices"] >= sum(allocated.values()) + assert report["spline_coeffs"] == sum( + _real_vector_nbytes(v.spline.vector) + for spec in sim.model.field_species.values() + for v in spec.variables.values() + ) + assert report["markers"] == 0 + assert report["total"] == sum(v for k, v in report.items() if k != "total") + + +def test_estimate_mem_hybrid_feec_and_pic(out_folders): + """estimate_mem() on a model with both FEEC and PIC variables: FEEC estimates match + exactly, and the PIC estimate matches the real marker-array footprint plus the + (separately estimated) marker-saving buffer.""" + model = VlasovAmpereOneSpecies(alpha=1.0, epsilon=-1.0, with_B0=False) + + env = EnvironmentOptions(out_folders=out_folders, sim_folder="weak_Landau") + time_opts = Time(dt=0.05, Tend=15) + domain = domains.Cuboid(r1=12.56) + grid = grids.TensorProductGrid(num_elements=(16, 1, 1)) + derham_opts = DerhamOptions(degree=(3, 1, 1)) + + loading_params = LoadingParameters(ppc=200, seed=1234) + weights_params = WeightsParameters(control_variate=True) + boundary_params = BoundaryParameters() + sorting_params = SortingParameters(boxes_per_dim=(8, 1, 1), do_sort=True) + saving_params = SavingParameters(n_markers=50) + + model.kinetic_ions.set_markers( + loading_params=loading_params, + weights_params=weights_params, + boundary_params=boundary_params, + sorting_params=sorting_params, + saving_params=saving_params, + bufsize=0.4, + ) + + model.propagators.push_eta.options = model.propagators.push_eta.Options() + model.propagators.coupling_va.options = model.propagators.coupling_va.Options() + model.initial_poisson.options = model.initial_poisson.Options(stab_mat="M0") + + background = maxwellians.Maxwellian3D(n=(1.0, None)) + model.kinetic_ions.var.add_background(background) + + sim = Simulation( + model=model, + env=env, + time_opts=time_opts, + domain=domain, + grid=grid, + derham_opts=derham_opts, + ) + + mem_before = sim.estimate_mem(print_report=False) + sim.allocate() + + for species, spec in sim.model.field_species.items(): + for k, v in spec.variables.items(): + estimated = mem_before[f"{species}.{k}"] + actual = _real_vector_nbytes(v.spline.vector) + assert estimated == actual, f"{species}.{k}: estimated {estimated} != actual {actual}" + + markers = 0 + for species, spec in sim.model.particle_species.items(): + for k, v in spec.variables.items(): + estimated = mem_before[f"{species}.{k}"] + actual = v.particles.nbytes_local + if v.n_to_save > 0: + actual += v.saved_markers.nbytes + assert estimated == actual, f"{species}.{k}: estimated {estimated} != actual {actual}" + markers += actual + + report = sim.report_mem(print_report=False) + assert report["markers"] == markers + assert report["feec_matrices"] >= sum(sim.mass_ops.allocated_mem().values()) > 0