document iterative solver and Monte Carlo caching options

- add README and Sphinx docs for the jacobi-gmres solver workflow
  - document aggregate_by and postprocess_multiprocessing behavior
  - update examples and API docstrings to reflect aggregated MC outputs

Author: romainsacchi

README.md

@@ -154,6 +154,44 @@ If not specified, all the methods, years, regions and scenarios
 defined in the datapackage.json file are used, which can be very
 time-consuming.
 
+For larger Monte Carlo studies, ``Pathways.calculate(...)`` also supports an
+experimental iterative solver and cache-shaping options:
+
+```python
+
+p.calculate(
+    methods=methods,
+    models=models,
+    scenarios=scenarios,
+    regions=regions,
+    years=years,
+    variables=variables,
+    use_distributions=300,
+    solver="jacobi-gmres",
+    iterative_rtol=1e-8,
+    aggregate_by=["act_category", "location"],
+    multiprocessing=True,
+    postprocess_multiprocessing=True,
+)
+
+```
+
+- `solver="direct"` is the default and uses `bw2calc.MultiLCA`.
+- `solver="jacobi-gmres"` uses an experimental iterative `MultiLCA` backend.
+  It can reuse previous supply arrays as warm starts and falls back to the
+  direct solve if GMRES does not converge.
+- `aggregate_by` currently supports `act_category` and `location`. These
+  dimensions are collapsed to a single `"aggregated"` label before Monte Carlo
+  iteration arrays are cached, which reduces cache size and shortens
+  post-processing.
+- `postprocess_multiprocessing=True` parallelizes final cached-result assembly.
+  During this phase, Pathways logs per-year assembly progress so long runs do
+  not appear stalled.
+
+Use `aggregate_by` only when you do not need detailed attribution along those
+dimensions in the final results. If you relax `iterative_rtol` for speed, check
+the result against a direct-solver reference first.
+
 Once calculated, the results of the LCA calculations are stored in the `.lcia_results`
 attribute of the `Pathways` object as an ``xarray.DataArray``. 
 

docs/api_reference.rst

@@ -14,12 +14,14 @@ Pathways
    - :py:meth:`pathways.Pathways.calculate` —
      Compute LCA results for selected methods, models, scenarios, regions, years, and variables.
      Parameters include ``demand_cutoff``, ``use_distributions``, ``subshares``,
-     ``remove_uncertainty``, ``seed``, ``multiprocessing``,
+     ``remove_uncertainty``, ``seed``, ``solver``, the ``iterative_*`` tuning
+     arguments, ``aggregate_by``, ``multiprocessing``,
      ``postprocess_multiprocessing``, and ``double_accounting``.
+     ``solver="direct"`` uses ``bw2calc.MultiLCA`` and
+     ``solver="jacobi-gmres"`` selects the experimental iterative backend.
 
    - :py:meth:`pathways.Pathways.aggregate_results` —
      Aggregate low-contribution activity categories under ``"other"``; optional interpolation.
 
    - :py:meth:`pathways.Pathways.export_results` —
      Export **non-zero** results to compressed Parquet (``.gzip``).
-

docs/examples.rst

@@ -45,9 +45,11 @@ Monte Carlo sampling and export
 -------------------------------
 
 Set ``use_distributions`` to a non-zero integer to trigger Monte Carlo draws
-for the technosphere uncertainty parameters stored in the datapackage.  The
-example below runs five iterations and writes both the aggregated results and
-per-iteration outputs to disk.
+for the technosphere uncertainty parameters stored in the datapackage. For
+larger runs, you can switch to the experimental iterative solver and collapse
+selected dimensions before cached iteration arrays are written. The example
+below runs five iterations, aggregates ``act_category`` and ``location`` during
+caching, and exports the final tensor.
 
 .. code-block:: python
 
@@ -59,6 +61,11 @@ per-iteration outputs to disk.
        years=[2030],
        variables=["Electricity|Generation"],
        use_distributions=5,
+       solver="jacobi-gmres",
+       iterative_rtol=1e-8,
+       aggregate_by=["act_category", "location"],
+       multiprocessing=True,
+       postprocess_multiprocessing=True,
        remove_uncertainty=False,
    )
 
@@ -74,3 +81,7 @@ per-iteration outputs to disk.
 
    mc_book = STATS_DIR / "ModelX_baseline_2030.xlsx"
    print(mc_book.exists())
+
+In this example, ``pw.lca_results`` keeps the same dimension names, but the
+``act_category`` and ``location`` coordinates each contain only the single value
+``"aggregated"``.

docs/index.rst

@@ -11,7 +11,10 @@ Pathways provides tools for **prospective life cycle assessment (LCA)** driven b
 You can then compute multi-year, multi-region impact results with a single call to
 :py:meth:`pathways.Pathways.calculate`, aggregate the results with
 :py:meth:`pathways.Pathways.aggregate_results`, and export them to a compact
-Parquet file with :py:meth:`pathways.Pathways.export_results`.
+Parquet file with :py:meth:`pathways.Pathways.export_results`. For larger Monte
+Carlo studies, ``calculate`` can also use an experimental iterative solver,
+collapse selected dimensions before caching, and parallelize the final
+cache-assembly step.
 
 Contents
 --------

docs/quickstart.rst

@@ -50,3 +50,32 @@ and export the non-zero entries to Parquet:
    # 5) Export non-zero cells to compressed Parquet
    out = pw.export_results("results_baseline")
    print("Wrote:", out)  # e.g., results_baseline.gzip
+
+Large Monte Carlo runs
+----------------------
+
+When ``use_distributions`` is greater than zero, you can switch from the
+default direct solver to the experimental iterative backend and reduce cache
+size by collapsing dimensions you do not need later:
+
+.. code-block:: python
+
+   pw.calculate(
+       methods=["AWARE"],
+       models=["REMIND"],
+       scenarios=["SSP2-NPi"],
+       regions=["World"],
+       years=[2050],
+       variables=["Electricity|Generation"],
+       use_distributions=300,
+       solver="jacobi-gmres",
+       iterative_rtol=1e-8,
+       aggregate_by=["act_category", "location"],
+       multiprocessing=True,
+       postprocess_multiprocessing=True,
+   )
+
+``solver="direct"`` remains the default. ``aggregate_by`` currently supports
+``"act_category"`` and ``"location"`` and replaces those coordinates with a
+single ``"aggregated"`` label in the resulting array, so use it only when you
+do not need detailed attribution along those dimensions.

docs/user_guide.rst

@@ -54,6 +54,43 @@ Pathways workflow
 
    ``(act_category, variable, year, region, location, model, scenario, impact_category)``
 
+   When ``aggregate_by`` is used during Monte Carlo runs, the dimensions stay in
+   the output, but the collapsed coordinates become ``"aggregated"``.
+
+   **Advanced Monte Carlo options**
+
+   .. code-block:: python
+
+      pw.calculate(
+          methods=["AWARE"],
+          models=["REMIND"],
+          scenarios=["SSP2-NPi"],
+          regions=["World"],
+          years=[2050],
+          variables=["Electricity|Generation"],
+          use_distributions=300,
+          solver="jacobi-gmres",
+          iterative_rtol=1e-8,
+          aggregate_by=["act_category", "location"],
+          multiprocessing=True,
+          postprocess_multiprocessing=True,
+      )
+
+   - ``solver="direct"`` is the default and uses ``bw2calc.MultiLCA``.
+   - ``solver="jacobi-gmres"`` uses an experimental iterative backend. It
+     reuses prior supply arrays as warm starts when possible and falls back to
+     the direct solve if GMRES does not converge.
+   - The iterative backend exposes ``iterative_rtol``, ``iterative_atol``,
+     ``iterative_restart``, ``iterative_maxiter``, and
+     ``iterative_use_guess`` on :meth:`pathways.Pathways.calculate`.
+   - ``aggregate_by`` currently supports ``"act_category"`` and
+     ``"location"``. It collapses those dimensions before Monte Carlo
+     iteration arrays are cached, which reduces cache size and shortens
+     post-processing.
+   - ``postprocess_multiprocessing=True`` parallelizes final cached-result
+     assembly after the solver stage. Pathways also logs yearly assembly
+     progress during this step.
+
 3. **Aggregate for display**
 
    .. code-block:: python
@@ -74,3 +111,5 @@ Notes
 - Supported ``ecoinvent_version`` values are ``"3.10"``, ``"3.11"``, and ``"3.12"``.
 - Matrix files ``A_matrix.csv`` and ``B_matrix.csv`` are semicolon-delimited and may include a header row.
 - Index files ``A_matrix_index.csv`` and ``B_matrix_index.csv`` may include a header row.
+- If you loosen ``iterative_rtol`` for speed, compare the result against a
+  direct-solver run before using it for production Monte Carlo analysis.

pathways/jacobi_gmres_multi_lca.py

@@ -0,0 +1,196 @@
+"""Experimental iterative ``MultiLCA`` backend for Pathways.
+
+This backend keeps Pathways on the standard ``bw2calc.MultiLCA`` lifecycle,
+but swaps the direct sparse solve for GMRES with a Jacobi preconditioner.
+Whenever available, it reuses the matrix-preparation and preconditioner
+helpers from ``bw2calc.JacobiGMRESLCA``.
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import Optional
+
+import bw2calc as bc
+import matrix_utils as mu
+import numpy as np
+from scipy import sparse
+from scipy.sparse.linalg import LinearOperator, gmres
+
+logger = logging.getLogger(__name__)
+
+try:
+    _BW_JACOBI_GMRES_LCA = bc.JacobiGMRESLCA
+except AttributeError:  # pragma: no cover - older bw2calc versions
+    _BW_JACOBI_GMRES_LCA = None
+
+
+class JacobiGMRESMultiLCA(bc.MultiLCA):
+    """Solve multi-demand LCI systems with GMRES and Jacobi preconditioning."""
+
+    def __init__(
+        self,
+        *args,
+        rtol: float = 1e-8,
+        atol: float = 0.0,
+        restart: Optional[int] = 50,
+        maxiter: Optional[int] = 300,
+        use_guess: bool = True,
+        direct_fallback: bool = True,
+        **kwargs,
+    ):
+        super().__init__(*args, **kwargs)
+        self.rtol = rtol
+        self.atol = atol
+        self.restart = restart
+        self.maxiter = maxiter
+        self.use_guess = use_guess
+        self.direct_fallback = direct_fallback
+
+        self._matrix_prepared = False
+        self._cached_preconditioner: Optional[LinearOperator] = None
+        self.guesses: dict[str, np.ndarray] = {}
+
+    def __next__(self) -> None:
+        # Matrix values can change on each Monte Carlo draw.
+        self._matrix_prepared = False
+        self._cached_preconditioner = None
+        self.guesses = {}
+        super().__next__()
+
+    def load_lci_data(self, nonsquare_ok=False) -> None:
+        super().load_lci_data(nonsquare_ok=nonsquare_ok)
+        self._matrix_prepared = False
+        self._cached_preconditioner = None
+        self.guesses = {}
+
+    def _prepare_matrix(self) -> None:
+        # ``MappedMatrix`` updates ``technosphere_mm.matrix`` across MC draws.
+        # Rebind here so GMRES always sees the current technosphere values instead
+        # of a stale CSC conversion from an earlier draw.
+        if hasattr(self, "technosphere_mm"):
+            self.technosphere_matrix = self.technosphere_mm.matrix
+
+        if _BW_JACOBI_GMRES_LCA is not None:
+            _BW_JACOBI_GMRES_LCA._prepare_matrix(self)
+            return
+
+        if self._matrix_prepared:
+            return
+
+        self.technosphere_matrix = self.technosphere_matrix.tocsc(copy=False)
+        self.technosphere_matrix.sum_duplicates()
+        self.technosphere_matrix.eliminate_zeros()
+        self.technosphere_matrix.sort_indices()
+        self._matrix_prepared = True
+
+    def _build_jacobi_preconditioner(self) -> Optional[LinearOperator]:
+        if _BW_JACOBI_GMRES_LCA is not None:
+            return _BW_JACOBI_GMRES_LCA._build_jacobi_preconditioner(self)
+
+        if self._cached_preconditioner is not None:
+            return self._cached_preconditioner
+
+        diagonal = self.technosphere_matrix.diagonal()
+        if np.any(diagonal == 0):
+            return None
+
+        inverse_diagonal = 1.0 / diagonal
+        self._cached_preconditioner = LinearOperator(
+            shape=self.technosphere_matrix.shape,
+            matvec=lambda x: inverse_diagonal * x,
+            dtype=self.technosphere_matrix.dtype,
+        )
+        return self._cached_preconditioner
+
+    def _solve_with_gmres(
+        self,
+        demand: np.ndarray,
+        *,
+        x0: np.ndarray | None = None,
+        demand_name: str | None = None,
+    ) -> np.ndarray:
+        self._prepare_matrix()
+        preconditioner = self._build_jacobi_preconditioner()
+
+        try:
+            solution, info = gmres(
+                self.technosphere_matrix,
+                demand,
+                x0=x0,
+                rtol=self.rtol,
+                atol=self.atol,
+                restart=self.restart,
+                maxiter=self.maxiter,
+                M=preconditioner,
+            )
+        except TypeError:  # pragma: no cover - SciPy compatibility fallback
+            solution, info = gmres(
+                self.technosphere_matrix,
+                demand,
+                x0=x0,
+                tol=self.rtol,
+                atol=self.atol,
+                restart=self.restart,
+                maxiter=self.maxiter,
+                M=preconditioner,
+            )
+
+        solution = np.asarray(solution, dtype=np.float64)
+        if not solution.shape:
+            solution = solution.reshape((1,))
+
+        if info != 0:
+            if not self.direct_fallback:
+                raise RuntimeError(
+                    "GMRES failed to converge "
+                    f"(demand={demand_name!r}, info={info}, rtol={self.rtol}, maxiter={self.maxiter})"
+                )
+
+            logger.warning(
+                "GMRES failed to converge for demand %s; falling back to direct solve.",
+                demand_name,
+            )
+            solution = np.asarray(bc.spsolve(self.technosphere_matrix, demand))
+            if not solution.shape:
+                solution = solution.reshape((1,))
+
+        return solution
+
+    def lci_calculation(self) -> None:
+        """Calculate inventories for many demands using iterative solves."""
+        count = len(self.dicts.activity)
+        demand_items = list(self.demand_arrays.items())
+        if not demand_items:
+            self.supply_arrays = {}
+            self.inventories = mu.SparseMatrixDict([])
+            return
+
+        supply_arrays: dict[str, np.ndarray] = {}
+        previous_solution: np.ndarray | None = None
+
+        for name, demand in demand_items:
+            x0 = None
+            if self.use_guess:
+                x0 = self.guesses.get(name)
+                if x0 is None:
+                    x0 = previous_solution
+
+            solution = self._solve_with_gmres(demand, x0=x0, demand_name=name)
+            supply_arrays[name] = solution
+            previous_solution = solution
+
+            if self.use_guess:
+                self.guesses[name] = solution
+
+        self.supply_arrays = supply_arrays
+        self.inventories = mu.SparseMatrixDict(
+            [
+                (
+                    name,
+                    self.biosphere_matrix
+                    @ sparse.spdiags([arr], [0], count, count),
+                )
+                for name, arr in self.supply_arrays.items()
+            ]
+        )