Cohorts across an allocation¶

These functions support a fit that runs on several processes or across several allocations. See fitting a cohort across an allocation.

diff_soup.inference.balance_owners ¶

balance_owners(experiments: Any, lanes: int, *, weight: Callable | None = None) -> Any

Reorder a study's experiments so that equal consecutive blocks of owners carry balanced loads.

The owners are the members of the study's outermost level or cluster, or else the experiments themselves. Each owner's experiments stay together, and the owners are assigned heaviest first, each to the lightest block with room.

Parameters:

Name	Type	Description	Default
`experiments`	`Any`	A `Study` or a sequence of `Experiment`.	required
`lanes`	`int`	The number of consecutive blocks, filled in order with `ceil(n_owners / lanes)` owners each.	required
`weight`	`Callable \| None`	A function giving the load of one experiment, or `None` to count its measured values.	`None`

Returns:

Type	Description
`Any`	The reordered study, or a tuple of the reordered experiments.

diff_soup.inference.prepare_windows ¶

prepare_windows(experiment: Experiment, trajectory: Any = None, **problem_kwargs: Any) -> Any

Compute and store the window starts of a multiple-shooting experiment before a fit.

A later fit of the same experiment loads them instead of solving again. Set DIFFSOUP_CACHE_DIR to a shared directory for a run of several processes.

Parameters:

Name	Type	Description	Default
`experiment`	`Experiment`	The experiment whose windows to initialize.	required
`trajectory`	`Any`	A `Multishoot` applied to the experiment, or `None` to keep its own.	`None`
`**problem_kwargs`	`Any`	Further keywords of `InferenceProblem`, such as `parameters=`.	`{}`

Returns:

Type	Description
`Any`	The key the windows are stored under, or `None` when there is nothing to store.

diff_soup.inference.checkpointed_map ¶

checkpointed_map(fn: Callable, items: Sequence[Any], *, ledger_path: Any, resume: bool = True, execution: str = 'serial', device_groups: Any = None, on_conflict: str = 'error') -> CampaignResult

Map a function over independent tasks, recording each outcome in a JSON ledger so a rerun resumes.

Each item carries a task_id that stays the same across restarts. A result is stored as JSON and comes back in that form. A task that raises is recorded as failed, and the others still run.

Parameters:

Name	Type	Description	Default
`fn`	`Callable`	The task, a function of one item returning its result.	required
`items`	`Sequence[Any]`	The tasks, mappings or objects with a unique non-empty string `task_id`.	required
`ledger_path`	`Any`	The JSON file the ledger is written to, whole, after every change.	required
`resume`	`bool`	Whether to read the ledger first, keep its accepted tasks, and run the failed and interrupted ones again.	`True`
`execution`	`str`	`"serial"` to run the tasks one after another, or `"processes"` to run each on a device group, on the processes that own it.	`'serial'`
`device_groups`	`Any`	Disjoint JAX submeshes for `"processes"`, or `None` to give each device its own group.	`None`
`on_conflict`	`str`	`"error"` to refuse a ledger that recorded a `task_id` for a different item, or `"restart"` to discard that ledger.	`'error'`

Returns:

Type	Description
`CampaignResult`	A `CampaignResult` with the accepted `results` in item order, the `failures`, and one entry of `records` per task.

diff_soup.inference.world ¶

world() -> World

Return this process's rank and the number of processes in the run.

A process not launched under MPI, or one using jax.distributed, is process 0 of 1. Launch several processes with python -m mpi4py, which needs mpi4py installed.

Returns:

Type	Description
`World`	A named tuple `(rank, size)`.