Migrate to SciPy sparse arrays

openmc/deplete/_sparse_compat.py

@@ -0,0 +1,37 @@
+"""Compatibility module for scipy.sparse arrays
+
+This module provides a compatibility layer for working with scipy.sparse arrays
+across different scipy versions. Sparse arrays were introduced gradually in
+scipy, with full support arriving in scipy 1.15. This module provides a unified
+API that uses sparse arrays when available and falls back to sparse matrices for
+older scipy versions.
+
+For more information on the migration from sparse matrices to sparse arrays,
+see: https://docs.scipy.org/doc/scipy/reference/sparse.migration_to_sparray.html
+"""
+
+import scipy
+from scipy import sparse as sp
+
+# Check scipy version for feature availability
+_SCIPY_VERSION = tuple(map(int, scipy.__version__.split('.')[:2]))
+
+if _SCIPY_VERSION >= (1, 15):
+    # Use sparse arrays
+    csc_array = sp.csc_array
+    dok_array = sp.dok_array
+    eye_array = sp.eye_array
+    block_array = sp.block_array
+else:
+    # Fall back to sparse matrices
+    csc_array = sp.csc_matrix
+    dok_array = sp.dok_matrix
+    eye_array = sp.eye
+    block_array = sp.bmat
+
+__all__ = [
+    'csc_array',
+    'dok_array',
+    'eye_array',
+    'block_array',
+]

openmc/deplete/abc.py

@@ -600,7 +600,7 @@ class Integrator(ABC):
         User-supplied functions are expected to have the following signature:
         ``solver(A, n0, t) -> n1`` where
 
-            * ``A`` is a :class:`scipy.sparse.csc_matrix` making up the
+            * ``A`` is a :class:`scipy.sparse.csc_array` making up the
               depletion matrix
             * ``n0`` is a 1-D :class:`numpy.ndarray` of initial compositions
               for a given material in atoms/cm3
@@ -1099,7 +1099,7 @@ class SIIntegrator(Integrator):
         User-supplied functions are expected to have the following signature:
         ``solver(A, n0, t) -> n1`` where
 
-            * ``A`` is a :class:`scipy.sparse.csc_matrix` making up the
+            * ``A`` is a :class:`scipy.sparse.csc_array` making up the
               depletion matrix
             * ``n0`` is a 1-D :class:`numpy.ndarray` of initial compositions
               for a given material in atoms/cm3
@@ -1216,7 +1216,7 @@ def __call__(self, A, n0, dt):
 
         Parameters
         ----------
-        A : scipy.sparse.csc_matrix
+        A : scipy.sparse.csc_array
             Sparse transmutation matrix ``A[j, i]`` describing rates at
             which isotope ``i`` transmutes to isotope ``j``
         n0 : numpy.ndarray

openmc/deplete/chain.py

@@ -16,13 +16,13 @@
 from typing import List
 
 import lxml.etree as ET
-import scipy.sparse as sp
 
 from openmc.checkvalue import check_type, check_greater_than, PathLike
 from openmc.data import gnds_name, zam
 from openmc.exceptions import DataError
 from .nuclide import FissionYieldDistribution, Nuclide
 from .._xml import get_text
+from ._sparse_compat import csc_array, dok_array
 import openmc.data
 
 
@@ -618,7 +618,7 @@ def form_matrix(self, rates, fission_yields=None):
 
         Returns
         -------
-        scipy.sparse.csc_matrix
+        scipy.sparse.csc_array
             Sparse matrix representing depletion.
 
         See Also
@@ -629,7 +629,7 @@ def form_matrix(self, rates, fission_yields=None):
 
         n = len(self)
 
-        # we accumulate indices and value entries for everything and create the matrix 
+        # we accumulate indices and value entries for everything and create the matrix
         # in one step at the end to avoid expensive index checks scipy otherwise does.
         rows, cols, vals = [], [], []
         def setval(i, j, val):
@@ -712,7 +712,7 @@ def setval(i, j, val):
                 reactions.clear()
 
         # Return CSC representation instead of DOK
-        return sp.csc_matrix((vals, (rows, cols)), shape=(n, n))
+        return csc_array((vals, (rows, cols)), shape=(n, n))
 
     def form_rr_term(self, tr_rates, current_timestep, mats):
         """Function to form the transfer rate term matrices.
@@ -743,13 +743,13 @@ def form_rr_term(self, tr_rates, current_timestep, mats):
 
         Returns
         -------
-        scipy.sparse.csc_matrix
+        scipy.sparse.csc_array
             Sparse matrix representing transfer term.
 
         """
         # Use DOK as intermediate representation
         n = len(self)
-        matrix = sp.dok_matrix((n, n))
+        matrix = dok_array((n, n))
 
         for i, nuc in enumerate(self.nuclides):
             elm = re.split(r'\d+', nuc.name)[0]
@@ -800,15 +800,15 @@ def form_ext_source_term(self, ext_source_rates, current_timestep, mat):
 
         Returns
         -------
-        scipy.sparse.csc_matrix
+        scipy.sparse.csc_array
             Sparse vector representing external source term.
 
         """
         if not ext_source_rates.get_components(mat, current_timestep):
             return
         # Use DOK as intermediate representation
         n = len(self)
-        vector = sp.dok_matrix((n, 1))
+        vector = dok_array((n, 1))
 
         for i, nuc in enumerate(self.nuclides):
             # Build source term vector

openmc/deplete/cram.py

@@ -6,11 +6,11 @@
 import numbers
 
 import numpy as np
-import scipy.sparse as sp
 import scipy.sparse.linalg as sla
 
 from openmc.checkvalue import check_type, check_length
 from .abc import DepSystemSolver
+from ._sparse_compat import csc_array, eye_array
 
 __all__ = ["CRAM16", "CRAM48", "Cram16Solver", "Cram48Solver", "IPFCramSolver"]
 
@@ -60,7 +60,7 @@ def __call__(self, A, n0, dt):
 
         Parameters
         ----------
-        A : scipy.sparse.csr_matrix
+        A : scipy.sparse.csc_array
             Sparse transmutation matrix ``A[j, i]`` desribing rates at
             which isotope ``i`` transmutes to isotope ``j``
         n0 : numpy.ndarray
@@ -75,9 +75,9 @@ def __call__(self, A, n0, dt):
             Final compositions after ``dt``
 
         """
-        A = dt * sp.csc_matrix(A, dtype=np.float64)
+        A = dt * csc_array(A, dtype=np.float64)
         y = n0.copy()
-        ident = sp.eye(A.shape[0], format='csc')
+        ident = eye_array(A.shape[0], format='csc')
         for alpha, theta in zip(self.alpha, self.theta):
             y += 2*np.real(alpha*sla.spsolve(A - theta*ident, y))
         return y * self.alpha0

openmc/deplete/pool.py

@@ -5,10 +5,11 @@
 from itertools import repeat, starmap
 from multiprocessing import Pool
 
-from scipy.sparse import bmat, hstack, vstack, csc_matrix
 import numpy as np
+from scipy.sparse import hstack
 
 from openmc.mpi import comm
+from ._sparse_compat import block_array
 
 # Configurable switch that enables / disables the use of
 # multiprocessing routines during depletion
@@ -146,7 +147,7 @@ def deplete(func, chain, n, rates, dt, current_timestep=None, matrix_func=None,
                             cols.append(None)
 
                     rows.append(cols)
-                matrix = bmat(rows)
+                matrix = block_array(rows)
 
                 # Concatenate vectors of nuclides in one
                 n_multi = np.concatenate(n)
@@ -181,7 +182,7 @@ def deplete(func, chain, n, rates, dt, current_timestep=None, matrix_func=None,
         # of the nuclide vectors
         for i, matrix in enumerate(matrices):
             if not np.equal(*matrix.shape):
-                matrices[i] = vstack([matrix, csc_matrix([0]*matrix.shape[1])])
+                matrix.resize(matrix.shape[1], matrix.shape[1])
                 n[i] = np.append(n[i], 1.0)
 
     inputs = zip(matrices, n, repeat(dt))