mirror of
https://github.com/openmc-dev/openmc.git
synced 2026-07-21 06:25:30 -04:00
Migrate to SciPy sparse arrays (#3613)
This commit is contained in:
Paul Romano
GitHub
parent
932f36f411
commit
9c91bddf04
7 changed files with 82 additions and 45 deletions
43
openmc/_sparse_compat.py
Normal file
43
openmc/_sparse_compat.py
Normal file
|
|
@ -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',
|
||||
]
|
||||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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))
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue