Migrate to SciPy sparse arrays (#3613)

Author: paulromano

openmc/_sparse_compat.py

@@ -0,0 +1,43 @@
+"""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
+    csr_array = sp.csr_array
+    csc_array = sp.csc_array
+    dok_array = sp.dok_array
+    lil_array = sp.lil_array
+    eye_array = sp.eye_array
+    block_array = sp.block_array
+else:
+    # Fall back to sparse matrices
+    csr_array = sp.csr_matrix
+    csc_array = sp.csc_matrix
+    dok_array = sp.dok_matrix
+    lil_array = sp.lil_matrix
+    eye_array = sp.eye
+    block_array = sp.bmat
+
+__all__ = [
+    'csr_array',
+    'csc_array',
+    'dok_array',
+    'lil_array',
+    'eye_array',
+    'block_array',
+]

openmc/cmfd.py

@@ -25,6 +25,7 @@
 from .checkvalue import (check_type, check_length, check_value,
                          check_greater_than, check_less_than)
 from .exceptions import OpenMCError
+from ._sparse_compat import csr_array
 
 # See if mpi4py module can be imported, define have_mpi global variable
 try:
@@ -980,8 +981,7 @@ def _initialize_linsolver(self):
         loss_row = self._loss_row
         loss_col = self._loss_col
         temp_data = np.ones(len(loss_row))
-        temp_loss = sparse.csr_matrix((temp_data, (loss_row, loss_col)),
-                                      shape=(n, n))
+        temp_loss = csr_array((temp_data, (loss_row, loss_col)), shape=(n, n))
         temp_loss.sort_indices()
 
         # Pass coremap as 1-d array of 32-bit integers
@@ -1585,7 +1585,7 @@ def _build_loss_matrix(self, adjoint):
         # Create csr matrix
         loss_row = self._loss_row
         loss_col = self._loss_col
-        loss = sparse.csr_matrix((data, (loss_row, loss_col)), shape=(n, n))
+        loss = csr_array((data, (loss_row, loss_col)), shape=(n, n))
         loss.sort_indices()
         return loss
 
@@ -1612,7 +1612,7 @@ def _build_prod_matrix(self, adjoint):
         # Create csr matrix
         prod_row = self._prod_row
         prod_col = self._prod_col
-        prod = sparse.csr_matrix((data, (prod_row, prod_col)), shape=(n, n))
+        prod = csr_array((data, (prod_row, prod_col)), shape=(n, n))
         prod.sort_indices()
         return prod
 

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
@@ -1134,7 +1134,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
@@ -1297,7 +1297,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

@@ -17,13 +17,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
 
 
@@ -619,7 +619,7 @@ def form_matrix(self, rates, fission_yields=None):
 
         Returns
         -------
-        scipy.sparse.csc_matrix
+        scipy.sparse.csc_array
             Sparse matrix representing depletion.
 
         See Also
@@ -713,7 +713,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 add_redox_term(self, matrix, buffer, oxidation_states):
         r"""Adds a redox term to the depletion matrix from data contained in
@@ -731,7 +731,7 @@ def add_redox_term(self, matrix, buffer, oxidation_states):
 
         Parameters
         ----------
-        matrix : scipy.sparse.csc_matrix
+        matrix : scipy.sparse.csc_array
             Sparse matrix representing depletion
         buffer : dict
             Dictionary of buffer nuclides used to maintain anoins net balance.
@@ -743,7 +743,7 @@ def add_redox_term(self, matrix, buffer, oxidation_states):
             states as integers (e.g., +1, 0).
         Returns
         -------
-        matrix : scipy.sparse.csc_matrix
+        matrix : scipy.sparse.csc_array
             Sparse matrix with redox term added
         """
         # Elements list with the same size as self.nuclides
@@ -769,7 +769,7 @@ def add_redox_term(self, matrix, buffer, oxidation_states):
         for nuc, idx in buffer_idx.items():
             array[idx] -= redox_change * buffer[nuc] / os[idx]
 
-        return sp.csc_matrix(array)
+        return csc_array(array)
 
     def form_rr_term(self, tr_rates, current_timestep, mats):
         """Function to form the transfer rate term matrices.
@@ -800,13 +800,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]
@@ -857,15 +857,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
@@ -159,7 +160,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)
@@ -194,7 +195,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))

openmc/tallies.py

@@ -12,11 +12,11 @@
 import h5py
 import numpy as np
 import pandas as pd
-import scipy.sparse as sps
 from scipy.stats import chi2, norm
 
 import openmc
 import openmc.checkvalue as cv
+from ._sparse_compat import lil_array
 from ._xml import clean_indentation, get_elem_list, get_text
 from .mixin import IDManagerMixin
 from .mesh import MeshBase
@@ -435,10 +435,10 @@ def _read_results(self):
 
             # Convert NumPy arrays to SciPy sparse LIL matrices
             if self.sparse:
-                self._sum = sps.lil_matrix(self._sum.flatten(), self._sum.shape)
-                self._sum_sq = sps.lil_matrix(self._sum_sq.flatten(), self._sum_sq.shape)
-                self._sum_third = sps.lil_matrix(self._sum_third.flatten(), self._sum_third.shape)
-                self._sum_fourth = sps.lil_matrix(self.sum_fourth.flatten(), self._sum_fourth.shape)
+                self._sum = lil_array(self._sum.flatten(), self._sum.shape)
+                self._sum_sq = lil_array(self._sum_sq.flatten(), self._sum_sq.shape)
+                self._sum_third = lil_array(self._sum_third.flatten(), self._sum_third.shape)
+                self._sum_fourth = lil_array(self._sum_fourth.flatten(), self._sum_fourth.shape)
 
             # Read simulation time (needed for figure of merit)
             self._simulation_time = f["runtime"]["simulation"][()]
@@ -534,8 +534,7 @@ def mean(self):
 
             # Convert NumPy array to SciPy sparse LIL matrix
             if self.sparse:
-                self._mean = sps.lil_matrix(self._mean.flatten(),
-                                            self._mean.shape)
+                self._mean = lil_array(self._mean.flatten(), self._mean.shape)
 
         if self.sparse:
             return np.reshape(self._mean.toarray(), self.shape)
@@ -556,8 +555,7 @@ def std_dev(self):
 
             # Convert NumPy array to SciPy sparse LIL matrix
             if self.sparse:
-                self._std_dev = sps.lil_matrix(self._std_dev.flatten(),
-                                               self._std_dev.shape)
+                self._std_dev = lil_array(self._std_dev.flatten(), self._std_dev.shape)
 
             self.with_batch_statistics = True
 
@@ -588,7 +586,7 @@ def vov(self):
             self._vov[mask] = numerator[mask]/denominator[mask] - 1.0/n
 
             if self.sparse:
-                self._vov = sps.lil_matrix(self._vov.flatten(), self._vov.shape)
+                self._vov = lil_array(self._vov.flatten(), self._vov.shape)
 
         if self.sparse:
             return np.reshape(self._vov.toarray(), self.shape)
@@ -963,22 +961,17 @@ def sparse(self, sparse):
         # Convert NumPy arrays to SciPy sparse LIL matrices
         if sparse and not self.sparse:
             if self._sum is not None:
-                self._sum = sps.lil_matrix(self._sum.flatten(), self._sum.shape)
+                self._sum = lil_array(self._sum.flatten(), self._sum.shape)
             if self._sum_sq is not None:
-                self._sum_sq = sps.lil_matrix(self._sum_sq.flatten(),
-                                              self._sum_sq.shape)
+                self._sum_sq = lil_array(self._sum_sq.flatten(), self._sum_sq.shape)
             if self._sum_third is not None:
-                self._sum_third = sps.lil_matrix(self._sum_third.flatten(),
-                                                 self._sum_third.shape)
+                self._sum_third = lil_array(self._sum_third.flatten(), self._sum_third.shape)
             if self._sum_fourth is not None:
-                self._sum_fourth = sps.lil_matrix(self._sum_fourth.flatten(),
-                                                  self._sum_fourth.shape)
+                self._sum_fourth = lil_array(self._sum_fourth.flatten(), self._sum_fourth.shape)
             if self._mean is not None:
-                self._mean = sps.lil_matrix(self._mean.flatten(),
-                                            self._mean.shape)
+                self._mean = lil_array(self._mean.flatten(), self._mean.shape)
             if self._std_dev is not None:
-                self._std_dev = sps.lil_matrix(self._std_dev.flatten(),
-                                               self._std_dev.shape)
+                self._std_dev = lil_array(self._std_dev.flatten(), self._std_dev.shape)
 
             self._sparse = True