perf: store Solution attributes in slots and tighten the memory gate

CHANGELOG.md

@@ -3,6 +3,7 @@
 ## Breaking changes
 
 - `Symbol.children`, `Symbol.domain` (and `secondary_domain`, `tertiary_domain`, `quaternary_domain`) and `FunctionParameter.input_names` return immutable tuples that compare equal to lists, and `Symbol.domains` returns a read-only copy, as do `to_json()`, `get_children_domains()`, `read_domain_or_domains()` and `pybamm.EMPTY_DOMAINS`. Editing them raises `TypeError` instead of emitting `SymbolMutationDeprecationWarning`, and not even `list` or `dict` methods can change a symbol through them. This ends the deprecation of these edits early; before 26.9 they left the symbol's identity stale. The sequences are no longer `list` instances and slicing them gives a plain tuple; tuples are accepted wherever domains are given. Build a new symbol with `with_domains`, `create_copy(new_children=...)` or `pybamm.replace`. The deprecated setters (`name`, `domains`, `scale`, ...) are unchanged. ([#5825](https://github.com/pybamm-team/PyBaMM/pull/5825))
+- `Solution` and `EISSolution` no longer have an instance `__dict__`, so attributes they do not define cannot be attached to them. Subclass `pybamm.Solution` to add attributes. A solution pickled with an attached attribute still loads, without that attribute and with a logged warning. ([#5827](https://github.com/pybamm-team/PyBaMM/pull/5827))
 
 ## Deprecated
 
@@ -19,9 +20,11 @@
 - `IDAKLUSolver`'s `print_stats` with an iterative linear solver and `preconditioner="none"` no longer reports a garbage "Number of calls to residual function in preconditioner": without a BBD preconditioner, SUNDIALS read that count out of the model's own data. ([#5782](https://github.com/pybamm-team/PyBaMM/pull/5782))
 - Discretising a model whose differential variables are all scalars or on zero-dimensional domains (such as a lumped current collector) no longer raises SciPy 1.18's `block_diag` `DeprecationWarning`, and its mass matrix stays a sparse matrix rather than becoming a sparse array once SciPy changes `block_diag`'s return type. ([#5782](https://github.com/pybamm-team/PyBaMM/pull/5782))
 - The CI test matrix now installs every dependency from `uv.lock`, as local nox sessions do, instead of resolving the latest releases, so only the prebuilt `pybammsolvers` wheel is installed outside the lock. ([#5782](https://github.com/pybamm-team/PyBaMM/pull/5782))
+- An editable `pybammsolvers` install, such as the one `uv sync` creates, can be imported on its own: `import pybammsolvers` no longer fails to load `libcasadi` unless `casadi` was imported first. ([#5827](https://github.com/pybamm-team/PyBaMM/pull/5827))
 
 ## Optimizations
 
+- `Solution` stores its attributes in `__slots__`, so an experiment, which keeps four solutions per step, retains about 40% less memory per step on Python 3.13. Solutions pickled by earlier versions still load. ([#5827](https://github.com/pybamm-team/PyBaMM/pull/5827))
 - Symbol mutation is no longer guarded by intercepting every attribute write, and `children` and `domain` return the stored immutable sequences instead of building a view on each access. Writes to a symbol's slots after construction are instead rejected by a static check of PyBaMM's code, tests, examples and documentation. Under the test suite's mutation guard, constructing an expression is twice as fast and parameterising it is 25% faster, which recovers most of the `test_parameterise` benchmark regression from [#5779](https://github.com/pybamm-team/PyBaMM/pull/5779). ([#5825](https://github.com/pybamm-team/PyBaMM/pull/5825))
 
 # [v26.9.0.0](https://github.com/pybamm-team/PyBaMM/tree/pybamm-v26.9.0.0) - 2026-09-28

packages/pybamm/src/pybamm/solvers/solution.py

@@ -7,7 +7,6 @@
 import numbers
 import pickle
 from dataclasses import astuple, dataclass
-from functools import cached_property
 from itertools import chain
 
 import casadi
@@ -96,11 +95,39 @@ class SolutionBase:
     across solution types.
     """
 
+    # An experiment keeps several solutions per step, so solutions have no
+    # __dict__; set_up_time and solve_time are set by solvers and simulations
+    __slots__ = ("__weakref__", "_data", "set_up_time", "solve_time")
+
     def __init__(self):
         self.set_up_time = None
         self.solve_time = None
         self._data = pybamm.FuzzyDict()
 
+    def __setstate__(self, state: dict | tuple) -> None:
+        """Restore a pickled solution.
+
+        Parameters
+        ----------
+        state : dict or tuple
+            The ``(instance dict, slots)`` pair a solution pickles, or the
+            single dict of a solution pickled before solutions had slots.
+        """
+        if isinstance(state, tuple):
+            instance_state, slot_state = state
+            state = {**(instance_state or {}), **(slot_state or {})}
+        dropped = []
+        for name, value in state.items():
+            try:
+                setattr(self, name, value)
+            except AttributeError:
+                dropped.append(name)
+        if dropped:
+            pybamm.logger.warning(
+                f"Dropped attributes {dropped} that {type(self).__name__} "
+                "does not define when loading a pickled solution"
+            )
+
     def __getitem__(self, key):
         """Access a variable by name."""
         return self._data[key]
@@ -183,6 +210,8 @@ class EISSolution(SolutionBase):
         Complex impedance values at each frequency.
     """
 
+    __slots__ = ()
+
     def __init__(self, frequencies, impedance):
         super().__init__()
         impedance = np.asarray(impedance, dtype=complex)
@@ -297,8 +326,47 @@ class Solution(SolutionBase):
 
     """
 
-    # A class default, so solutions pickled before this attribute existed load
-    solver_statistics: SolverStatistics | None = None
+    # Solvers, Simulation and make_cycle_solution set all_first_states, closest_event_idx,
+    # cycle_summary_variables, extrap_events, integration_time, solver_statistics and steps
+    __slots__ = (
+        "_all_inputs_casadi",
+        "_all_inputs_stacked",
+        "_all_models",
+        "_all_sensitivities",
+        "_all_t_evals",
+        "_all_ts",
+        "_all_yps",
+        "_all_ys",
+        "_all_ys_and_sens",
+        "_cycles",
+        "_first_state",
+        "_initial_start_time",
+        "_last_state",
+        "_observable",
+        "_options",
+        "_sensitivities",
+        "_sub_solutions",
+        "_summary_variables",
+        "_t",
+        "_t_eval",
+        "_t_event",
+        "_termination",
+        "_user_options",
+        "_variables",
+        "_y",
+        "_y_event",
+        "_yp",
+        "all_first_states",
+        "all_inputs",
+        "all_summary_variables",
+        "closest_event_idx",
+        "cycle_summary_variables",
+        "extrap_events",
+        "integration_time",
+        "solver_statistics",
+        "steps",
+        "variables_returned",
+    )
 
     def __init__(
         self,
@@ -379,6 +447,7 @@ def __init__(
 
         super().__init__()
         self.integration_time = None
+        self.solver_statistics = None
 
         self._all_inputs_stacked = None
         self._all_inputs_casadi = None
@@ -406,13 +475,29 @@ def __init__(
         # Initialise initial start time
         self.initial_start_time = None
 
+        self._first_state = None
+        self._last_state = None
+
         # Check no ys are too large
         if check_solution:
             self.check_ys_are_not_too_large()
 
         # Solution now uses CasADi
         pybamm.citations.register("Andersson2019")
 
+    def __setstate__(self, state: dict | tuple) -> None:
+        # An older pickle may lack attributes added since, and caches the
+        # first and last states under their property names
+        self.solver_statistics = None
+        self._first_state = None
+        self._last_state = None
+        if isinstance(state, dict):
+            state = {
+                f"_{name}" if name in ("first_state", "last_state") else name: value
+                for name, value in state.items()
+            }
+        super().__setstate__(state)
+
     def has_sensitivities(self) -> bool:
         return len(self._all_sensitivities) > 0
 
@@ -639,13 +724,18 @@ def user_options(self) -> dict:
     def options(self) -> dict:
         return self._options
 
-    @cached_property
+    @property
     def first_state(self):
         """
         A Solution object that only contains the first state. This is faster to evaluate
         than the full solution when only the first state is needed (e.g. to initialize
         a model with the solution)
         """
+        if self._first_state is None:
+            self._first_state = self._build_first_state()
+        return self._first_state
+
+    def _build_first_state(self):
         sensitivities = {}
         n_states = self.all_models[0].len_rhs_and_alg
         for key in self._all_sensitivities:
@@ -684,13 +774,18 @@ def first_state(self):
 
         return new_sol
 
-    @cached_property
+    @property
     def last_state(self):
         """
         A Solution object that only contains the final state. This is faster to evaluate
         than the full solution when only the final state is needed (e.g. to initialize
         a model with the solution)
         """
+        if self._last_state is None:
+            self._last_state = self._build_last_state()
+        return self._last_state
+
+    def _build_last_state(self):
         sensitivities = {}
         n_states = self.all_models[-1].len_rhs_and_alg
         for key in self._all_sensitivities:

packages/pybamm/tests/memory/test_memory_leaks.py

@@ -220,12 +220,12 @@ def test_cycle_count_memory_scaling(self):
         mem_20_cycles, _ = tracemalloc.get_traced_memory()
         tracemalloc.stop()
 
-        # Linear growth (4x) means each step rebuilds the model. On a ~1.9 MB fixed
-        # footprint, ~48 KB/cycle growth gives ~1.37x, so 1.4 bounds allocator noise.
+        # Linear growth (4x) means each step rebuilds the model. On a ~1.86 MB fixed
+        # footprint, ~33 KB/cycle growth gives ~1.25x, so 1.3 bounds allocator noise.
         ratio = mem_20_cycles / mem_5_cycles
-        assert ratio < 1.4, (
-            f"Memory grew {ratio:.1f}x for 4x more cycles. "
-            f"Expected sub-linear growth (<1.4x). "
+        assert ratio < 1.3, (
+            f"Memory grew {ratio:.2f}x for 4x more cycles. "
+            f"Expected sub-linear growth (<1.3x). "
             f"This may indicate termination hashing is broken (see #5453)."
         )
 

packages/pybamm/tests/unit/test_eis_simulation.py

@@ -1,3 +1,5 @@
+import pickle  # nosec B403 - used in tests with trusted input
+
 import numpy as np
 import pytest
 
@@ -146,6 +148,23 @@ def test_eis_solution_save_pickle(self, tmp_path):
         loaded = pybamm.load(str(filepath))
         np.testing.assert_array_equal(loaded.impedance, z)
 
+    def test_eis_solution_pickled_before_slots_loads(self):
+        freqs = np.array([1.0, 10.0])
+        z = np.array([1 + 0.5j, 2 + 1j])
+        data = pybamm.EISSolution(freqs, z).data
+        # A solution pickled before solutions had slots holds one dict
+        state = {"_data": data, "set_up_time": 0.5, "solve_time": 1.5}
+
+        class PreSlotsPickle:
+            def __reduce__(self):
+                return object.__new__, (pybamm.EISSolution,), state
+
+        loaded = pickle.loads(pickle.dumps(PreSlotsPickle()))  # nosec B301
+        np.testing.assert_array_equal(loaded.impedance, z)
+        np.testing.assert_array_equal(loaded.frequencies, freqs)
+        assert loaded.data.keys() == data.keys()
+        assert loaded.total_time == pytest.approx(2.0)
+
     def test_eis_solution_save_data_invalid_format(self, tmp_path):
         sol = pybamm.EISSolution(np.array([1.0]), np.array([1 + 1j]))
         with pytest.raises(ValueError, match="Unrecognised format"):

packages/pybamm/tests/unit/test_solvers/test_solution.py

@@ -727,6 +727,41 @@ def test_cycles(self):
         np.testing.assert_array_equal(sol.cycles[1].t, sol.t[len_cycle_1:])
         np.testing.assert_allclose(sol.cycles[1].y, sol.y[:, len_cycle_1:])
 
+    def test_solutions_reject_undefined_attributes(self):
+        # An experiment keeps several solutions per step, so none carries a __dict__
+        t = np.linspace(0, 1)
+        solution = pybamm.Solution(t, np.tile(t, (2, 1)), pybamm.BaseModel(), {})
+        eis_solution = pybamm.EISSolution(np.array([1.0]), np.array([1 + 1j]))
+        for sol in (solution, solution.first_state, eis_solution):
+            assert not hasattr(sol, "__dict__")
+            with pytest.raises(AttributeError, match=r"label"):
+                sol.label = "extra"
+
+    @pytest.mark.skipif(
+        sys.version_info < (3, 13), reason="__static_attributes__ is new in 3.13"
+    )
+    def test_attributes_set_by_methods_are_slots(self):
+        # An attribute a method sets without a slot raises AttributeError, which
+        # a rarely run path could otherwise hide
+        classes = pybamm.Solution.__mro__[:-1]
+        slots = {name for cls in classes for name in cls.__slots__}
+        properties = {
+            name
+            for cls in classes
+            for name, value in vars(cls).items()
+            if isinstance(value, property)
+        }
+        for cls in classes:
+            assert set(cls.__static_attributes__) <= slots | properties
+
+    def test_pickle_keeps_slots(self):
+        t = np.linspace(0, 1)
+        solution = pybamm.Solution(t, np.tile(t, (2, 1)), pybamm.BaseModel(), {})
+        first_state = solution.first_state
+        loaded = pickle.loads(pickle.dumps(solution))  # nosec B301
+        np.testing.assert_array_equal(loaded.t, t)
+        np.testing.assert_array_equal(loaded.first_state.y, first_state.y)
+
     def test_total_time(self):
         sol = pybamm.Solution(np.array([0]), np.array([[1, 2]]), pybamm.BaseModel(), {})
         sol.set_up_time = 0.5
@@ -1942,8 +1977,31 @@ def test_statistics_survive_pickling(self):
         solution = self._solution(0, statistics)
         assert pickle.loads(pickle.dumps(solution)).solver_statistics == statistics  # nosec B301
 
-    def test_solution_pickled_without_statistics_loads(self):
+    def test_solution_pickled_before_slots_loads(self, caplog):
         solution = self._solution(0, pybamm.SolverStatistics(1, 2, 3, 4, 5))
-        # A solution pickled before the attribute existed has no such state
-        del solution.__dict__["solver_statistics"]
-        assert pickle.loads(pickle.dumps(solution)).solver_statistics is None  # nosec B301
+        first_state = solution.first_state
+        # A solution pickled before Solution had slots holds one dict, without
+        # solver_statistics, with the cached first state under its name
+        state = {
+            name: getattr(solution, name)
+            for cls in type(solution).__mro__
+            for name in getattr(cls, "__slots__", ())
+            if not name.startswith("__") and hasattr(solution, name)
+        }
+        del state["solver_statistics"], state["_first_state"], state["_last_state"]
+        state["first_state"] = first_state
+        # and perhaps with an attribute a caller attached
+        state["label"] = "extra"
+
+        class PreSlotsPickle:
+            def __reduce__(self):
+                return object.__new__, (pybamm.Solution,), state
+
+        with caplog.at_level(logging.WARNING, logger="pybamm.logger"):
+            loaded = pickle.loads(pickle.dumps(PreSlotsPickle()))  # nosec B301
+        assert "Dropped attributes ['label']" in caplog.text
+        assert not hasattr(loaded, "label")
+        assert loaded.solver_statistics is None
+        np.testing.assert_array_equal(loaded.t, solution.t)
+        np.testing.assert_array_equal(loaded.first_state.y, first_state.y)
+        assert loaded.last_state.t[0] == pytest.approx(solution.t[-1])

packages/pybammsolvers/CMakeLists.txt

@@ -205,6 +205,11 @@ if (${USE_PYTHON_CASADI})
                 "${SuiteSparse_ROOT}/lib64"
             )
         endif()
+        # An editable install loads idaklu from the build tree, which has no
+        # ../casadi, so it also searches the environment's CasADi directory
+        if (SKBUILD_STATE STREQUAL "editable")
+            list(APPEND IDAKLU_INSTALL_RPATH "${CASADI_LIB_DIR}")
+        endif()
         set_target_properties(idaklu PROPERTIES
             INSTALL_RPATH "${IDAKLU_INSTALL_RPATH}"
             BUILD_WITH_INSTALL_RPATH TRUE