Source code for pynucastro.reduction.reduction_utils
[docs]
class MPIImportError(Exception):
pass
[docs]
class FailedMPIImport:
"""Class that can replace an mpi4py.MPI import and will throw an error if used."""
def __init__(self, error=None, msg=None):
self.error_obj = error
if msg is None:
msg = ("Failed to import MPI from mpi4py. Check your mpi4py installation if you"
" want to run with MPI.")
self.msg = msg
def __getattr__(self, attr):
if self.error_obj is not None:
raise MPIImportError(self.msg) from self.error_obj
raise MPIImportError(self.msg)
[docs]
def mpi_importer():
"""
Lazy MPI import, where we only throw an error if the import failed and then we attempt to use
the object.
"""
try:
from mpi4py import MPI # pylint: disable=import-outside-toplevel
except (ModuleNotFoundError, ImportError) as e:
MPI = FailedMPIImport(e)
return MPI
[docs]
def mpi_numpy_decomp(MPI_N, MPI_rank, n):
"""
Decompose a set of conditions for *MPI_N* MPI processes, where the conditions are a sequence of
3 sequences with ordering (composition_sequence, density_sequence, temperature_sequence). This
structure for the dataset is necessary for the vectorized reduction algorithms.
"""
if MPI_N <= n[0]:
comp_idx = MPI_rank
comp_step = MPI_N
rho_idx = T_idx = 0
rho_step = T_step = 1
elif MPI_N <= n[0]*n[1]:
m = MPI_N // n[0]
if MPI_rank <= n[0]*m:
comp_idx = MPI_rank % n[0]
comp_step = n[0]
rho_idx = MPI_rank // n[0]
rho_step = m
else:
comp_idx = n[0]
comp_step = 1
rho_idx = n[1]
rho_step = 1
T_idx = 0
T_step = 1
elif MPI_N <= n[0]*n[1]*n[2]:
m = MPI_N // (n[0] * n[1])
if MPI_rank <= n[0]*n[1]*m:
comp_idx = MPI_rank % n[0]
comp_step = n[0]
rho_idx = (MPI_rank // n[0]) % n[1]
rho_step = n[1]
T_idx = MPI_rank // (n[0] * n[1])
T_step = m
else:
comp_idx = n[0]
comp_step = 1
rho_idx = n[1]
rho_step = 1
T_idx = n[2]
T_step = 1
else:
m = MPI_N // (n[0] * n[1])
if MPI_rank <= n[0]*n[1]*n[2]:
comp_idx = MPI_rank % n[0]
comp_step = n[0]
rho_idx = (MPI_rank // n[0]) % n[1]
rho_step = n[1]
T_idx = MPI_rank // (n[0] * n[1])
T_step = n[2]
else:
comp_idx = n[0]
comp_step = 1
rho_idx = n[1]
rho_step = 1
T_idx = n[2]
T_step = 1
return comp_idx, comp_step, rho_idx, rho_step, T_idx, T_step
[docs]
def to_list(x, n=1):
"""
Convert a sequence or non-iterable to a list. In the non-iterable case, the supplied object will
be repeated *n* times (default 1).
"""
try:
return list(x)
except TypeError:
return [x] * n