foxes.engines.mpi¶

Classes¶

`MPIEngine`	The MPI engine for foxes calculations.
`MPIEngineRunner`	Engine runner for MPIEngine.

Module Contents¶

class foxes.engines.mpi.MPIEngine(*args: Any, share_cstore: bool = False, pool_args: dict[str, Any] | None = None, supports_shared_data: bool = True, min_shared_array_bytes: int = 65536, **kwargs: Any)[source]¶

Bases: foxes.engines.process.ProcessEngine

The MPI engine for foxes calculations.

Notes¶

Builds one MPI-backed shared-memory copy of shared mdata per node and lets chunk tasks attach to that cache by token. This reduces repeated transfers to MPI workers, but still replicates the shared input once per shared-memory domain because MPI shared windows do not span multiple nodes.

Examples¶

Run command, e.g. for 12 processors and a script run.py:

>>> mpiexec -n 12 -m mpi4py.futures run.py

Parameters:

args: Arguments for the base class
pool_args: Arguments for the pool constructor
share_cstore: Whether to share the chunk store between chunks.
supports_shared_data: Flag for whether this engine supports shared data for chunk calculations.
min_shared_array_bytes: Minimum array size in bytes for placing model data into process shared storage. Arrays with nbytes less than or equal to this threshold are transferred inline to workers. Supported top-level extra_data values use their complete payload size.
kwargs: Additional arguments for the base class

class ChunkResultsManager(algo: foxes.core.algorithm.Algorithm, engine: Engine, chunk_store: Any, goal_data: xarray.Dataset, n_chunks_states: int, n_chunks_targets: int, out_vars: list[str], out_dims: tuple[str, ...], coords: dict[str, Any], iterative: bool, write_nc: dict[str, Any] | None)¶

Helper class for results management during chunk calculations

Parameters:¶

algo: The algorithm object.
engine: The engine object.
chunk_store: The chunk store.
goal_data: The goal dataset.
n_chunks_states: The number of state chunks.
n_chunks_targets: The number of target chunks.
out_vars: The output variables.
out_dims: The output dimensions.
coords: The coordinates.
iterative: Whether the calculation is iterative.
write_nc: NetCDF output parameters, or None.

update(results: dict[tuple[int, int], Any], futures: list[Any] | None = None) → None¶

Update chunk calculation progress and accumulate the results.

Parameters:¶

results: A dictionary of chunk results.
futures: The current futures for asynchronous writing, or None.

algo¶

base_name : str | None = None¶

chunk_store¶

ci_states = 0¶

ci_targets = 0¶

coords¶

counter = 0¶

data_vars : dict[str, Any]¶

engine¶

fcounter = 0¶

gen_size : Iterator[Any] | None = None¶

goal_data¶

iterative¶

n_chunks_all¶

n_chunks_states¶

n_chunks_targets¶

name¶

out_dims¶

out_dir : Any = None¶

out_vars¶

pack : bool | None = None¶

pbar : Any = None¶

pdone = -1¶

res_vars : list[str] | None = None¶

results : xarray.Dataset | None = None¶

ret_data = True¶

scount = 0¶

split_size = None¶

tres : dict[str, list[numpy.ndarray]] | None = None¶

verbosity¶

wcount = 0¶

wfutures : list[Any] = []¶

write_from_ds = False¶

write_on_fly = False¶

await_result(future: Any) → Any¶

Waits for result from a future

Parameters:¶

future: The future

Returns:¶

result: The calculation result

calc_chunk_sizes(n_states: int, n_targets: int = 1) → tuple[numpy.ndarray, numpy.ndarray]¶

Compute the sizes of the state and target chunks.

Parameters:¶

n_states: The number of states.
n_targets: The number of point targets.

Returns:¶

chunk_sizes_states: The sizes of all state chunks, with shape (n_chunks_states,).
chunk_sizes_targets: The sizes of all target chunks, with shape (n_chunks_targets,).

future_is_done(future: Any) → bool¶

Checks if a future is done

Parameters:¶

future: The future

Returns:¶

is_done: True if the future is done

get_chunk_input_data(algo: foxes.core.algorithm.Algorithm, model_data: xarray.Dataset, farm_data: xarray.Dataset | None, point_data: xarray.Dataset | None, states_i0_i1: tuple[int, int], targets_i0_i1: tuple[int, int], out_vars: list[str], chunki_states: int, chunki_points: int, n_chunks_states: int, n_chunks_points: int) → tuple[foxes.core.data.MData, foxes.core.data.FData] | tuple[foxes.core.data.MData, foxes.core.data.FData, foxes.core.data.TData]¶

Extract the data for a single chunk calculation.

Parameters:¶

algo: The algorithm object.
model_data: The initial model data.
farm_data: The initial farm data.
point_data: The initial point data.
states_i0_i1: The start and end indices of the state slice.
targets_i0_i1: The start and end indices of the target slice.
out_vars: Names of the output variables.
chunki_states: The index of the states chunk.
chunki_points: The index of the points chunk.
n_chunks_states: The number of state chunks.
n_chunks_points: The number of point chunks.

Returns:¶

data: The input data for the chunk calculation, either (mdata, fdata) or (mdata, fdata, tdata).

get_start_calc_message(n_chunks_states: int, n_chunks_targets: int) → str¶: Helper function for start calculation message

init_shared_memory(shared_memory: list[Any], mdata: foxes.core.MData, shared_mdata: foxes.core.MData | None, verbosity: int = 0) → dict[str, Any] | None[source]¶

Sets the shared memory for the chunk calculation

Parameters:¶

shared_memory: The shared memory object for the chunk calculation
mdata: The mdata to be used in the chunk calculation
shared_mdata: The shared mdata to be used in the chunk calculation
verbosity: The verbosity level, 0=silent

Returns:¶

handle: The handle for accessing the shared data

map(func: Any, inputs: Any, *args: Any, **kwargs: Any) → list[Any]¶

Runs a function on a list of files

Parameters:¶

func: Function to be called on each file, func(input, *args, **kwargs) -> data
inputs: The input data list
args: Arguments for func
kwargs: Keyword arguments for func

Returns:¶

results: Results for the submitted inputs

classmethod new(engine_type: str | None, *args: Any, **kwargs: Any) → Engine¶

Create an engine instance at runtime.

Parameters:¶

engine_type: The selected derived class name.
args: Additional positional arguments for the constructor.
kwargs: Additional keyword arguments for the constructor.

new_chunk_results_manager(algo: foxes.core.algorithm.Algorithm, **kwargs: Any) → ChunkResultsManager¶

Create a new chunk results manager.

Parameters:¶

algo: The algorithm object.
kwargs: Additional keyword arguments.

Returns:¶

crm: The chunk results manager.

new_runner() → MPIEngineRunner[source]¶

Creates a new EngineRunner for running calculations in this engine.

Returns:¶

runner: The engine runner

prepare_chunk_mdata_for_shared(mdata: foxes.core.MData, shared_handle: dict[str, Any] | None) → None[source]¶: Remove entries that worker recombination restores from MPI shared cache.

print(*args: Any, level: int = 1, **kwargs: Any) → None¶: Print output based on the configured verbosity.

release_shared_memory(shared_memory: list[Any], shared_handle: dict[str, Any] | None) → None[source]¶

Releases the shared memory after the chunk calculation

Parameters:¶

shared_memory: The shared memory object for the chunk calculation
shared_handle: The handle for accessing the shared data

run_calculation(algo: foxes.core.Algorithm, model: foxes.core.DataCalcModel, model_data: xarray.Dataset | None = None, farm_data: xarray.Dataset | None = None, point_data: xarray.Dataset | None = None, extra_data: dict[str, Any] | None = None, out_vars: list[str] | None = None, chunk_store: dict[Any, Any] | None = None, sel: dict[str, Any] | None = None, isel: dict[str, Any] | None = None, iterative: bool = False, write_nc: dict[str, Any] | None = None, write_chunk_ani: dict[str, Any] | None = None, **calc_pars: Any) → xarray.Dataset¶

Runs the model calculation

Parameters:¶

algo

The algorithm object

model

The model that whose calculate function should be run

model_data

The initial model data

farm_data

The initial farm data

point_data

The initial point data

extra_data

Additional non-array input data from the models

out_vars

Names of the output variables

chunk_store

The chunk store

sel

Selection of coordinate subsets

isel

Selection of coordinate subsets index values

iterative

Flag for use within the iterative algorithm

write_nc

Parameters for writing results to netCDF files, e.g. {‘out_dir’: ‘results’, ‘base_name’: ‘calc_results’, ‘ret_data’: False, ‘split’: 1000}.

The split parameter controls how the output is split: - ‘chunks’: one file per chunk (fastest method), - ‘input’: split according to sizes of multiple states input files, - int: split with this many states per file, - None: create a single output file.

Use ret_data = False together with non-single file writing to avoid constructing the full Dataset in memory.

write_chunk_ani

Parameters for writing chunk animations, e.g. {‘fpath_base’: ‘results/chunk_animation’, ‘vars’: [‘WS’], ‘resolution’: 100, ‘chunk’: 5}.’} The chunk is either an integer that refers to a states chunk, a tuple (states_chunk_index, points_chunk_index), or a list of chunk indices.

calc_pars

Additional parameters for the model.calculate()

Returns:¶

results: The model results

select_subsets(*datasets: Any, sel: dict[str, Any] | None = None, isel: dict[str, Any] | None = None, default_n_states: int | None = None) → tuple[list[Any], int | None]¶

Take subsets of datasets.

Parameters:¶

datasets: The xarray dataset or data array objects.
sel: The selection dictionary.
isel: The index selection dictionary.
default_n_states: The fallback number of states if no dataset has a state dimension.

Returns:¶

subsets: The subsets of the input data.
n_states: The number of states after subset selection, or the fallback value.

submit(f: Any, *args: Any, **kwargs: Any) → Any¶

Submits a job to worker, obtaining a future

Parameters:¶

f: The function f(*args, **kwargs) to be submitted
args: Arguments for the function
kwargs: Arguments for the function

Returns:¶

future: The future object

chunk_size_points = None¶

chunk_size_states = None¶

property entered : bool¶

Return whether this engine has been entered.

Returns:¶

flag: True if the engine has been entered.

property has_progress_bar : bool¶

Return whether a progress bar is active.

Returns:¶

has_pbar: True if a progress bar is active.

property loop_dims : list[str]¶

Return the loop dimensions, including chunking when applicable.

Returns:¶

dims: The loop dimensions, possibly chunked.

min_shared_array_bytes = 65536¶

property n_procs : int¶

Return the number of processes.

Returns:¶

n_procs: The number of processes.

property n_workers : int¶

Return the number of worker processes.

Returns:¶

n_workers: The number of worker processes.

property name : str¶

Return the engine name.

Returns:¶

nme: The engine name.

pool_args¶

property prints_progress : bool¶

Return whether progress printing is active.

Returns:¶

has_pbar: True if progress printing is active.

progress_bar = True¶

property running_chunk_calc : bool¶

Return whether a chunk calculation is running.

Returns:¶

flag: True if a chunk calculation is running.

share_cstore = False¶

supports_shared_data = True¶

verbosity = 1¶

class foxes.engines.mpi.MPIEngineRunner[source]¶

Bases: foxes.engines.process.ProcessEngineRunner

Engine runner for MPIEngine.

run(algo: foxes.core.Algorithm, model: foxes.core.DataCalcModel, mdata: foxes.core.MData, fdata: foxes.core.FData, tdata: foxes.core.TData | None = None, *, shared: Any, chunk_store: dict[Any, Any], chunk_key: Any, out_dims: tuple[str, ...], write_nc: dict[str, Any] | None, write_chunk_ani: dict[str, Any] | None = None, utm_zone: tuple[int, str] | None = None, **cpars: Any) → tuple[dict[str, Any] | None, dict[Any, Any]]¶: Helper function for running in a single process