diff --git a/openmc/_sparse_compat.py b/openmc/_sparse_compat.py new file mode 100644 index 000000000..c00777e19 --- /dev/null +++ b/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', +] diff --git a/openmc/cmfd.py b/openmc/cmfd.py index eff6a151f..8595d1e02 100644 --- a/openmc/cmfd.py +++ b/openmc/cmfd.py @@ -25,6 +25,7 @@ import openmc.lib 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 @@ class CMFDRun: 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 @@ class CMFDRun: # 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 @@ class CMFDRun: # 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 diff --git a/openmc/deplete/abc.py b/openmc/deplete/abc.py index 32f468306..ee742f22b 100644 --- a/openmc/deplete/abc.py +++ b/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 @@ class DepSystemSolver(ABC): 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 diff --git a/openmc/deplete/chain.py b/openmc/deplete/chain.py index 873d7ca89..a835face7 100644 --- a/openmc/deplete/chain.py +++ b/openmc/deplete/chain.py @@ -17,13 +17,13 @@ from warnings import warn 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 @@ class Chain: Returns ------- - scipy.sparse.csc_matrix + scipy.sparse.csc_array Sparse matrix representing depletion. See Also @@ -713,7 +713,7 @@ class Chain: 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 @@ class Chain: 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 @@ class Chain: 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 @@ class Chain: 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 @@ class Chain: 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,7 +857,7 @@ class Chain: Returns ------- - scipy.sparse.csc_matrix + scipy.sparse.csc_array Sparse vector representing external source term. """ @@ -865,7 +865,7 @@ class Chain: 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 diff --git a/openmc/deplete/cram.py b/openmc/deplete/cram.py index 53de83bb6..cecc388f4 100644 --- a/openmc/deplete/cram.py +++ b/openmc/deplete/cram.py @@ -6,11 +6,11 @@ Implements two different forms of CRAM for use in openmc.deplete. 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 @@ class IPFCramSolver(DepSystemSolver): 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 @@ class IPFCramSolver(DepSystemSolver): 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 diff --git a/openmc/deplete/pool.py b/openmc/deplete/pool.py index aa348c02a..58f90894b 100644 --- a/openmc/deplete/pool.py +++ b/openmc/deplete/pool.py @@ -5,10 +5,11 @@ Provided to avoid some circular imports 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)) diff --git a/openmc/tallies.py b/openmc/tallies.py index 09365e525..7b1ef0219 100644 --- a/openmc/tallies.py +++ b/openmc/tallies.py @@ -12,11 +12,11 @@ import lxml.etree as ET 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 @@ class Tally(IDManagerMixin): # 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 @@ class Tally(IDManagerMixin): # 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 @@ class Tally(IDManagerMixin): # 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 @@ class Tally(IDManagerMixin): 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 @@ class Tally(IDManagerMixin): # 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