nvalchemi.distributed.DomainParallel

class nvalchemi.distributed.DomainParallel(dynamics, config, **kwargs)

Wraps any BaseDynamics subclass with spatial domain decomposition.

Flow per step:

1. Outer BEFORE_STEP hooks on owned batch.

2. Inner dynamics pre_update (velocity-Verlet half-kick) on owned batch.

3. Wrap positions into the periodic box.

4. Sync the updated positions back into the persistent ShardedBatch (update_from_batch).

5. DistributedModel(sharded) — the adapter rebuilds the halo block, rebuilds NL, runs the wrapper, consolidates owned-shape outputs.

6. Write the consolidated outputs back to the owned batch in-place.

7. Inner dynamics post_update (velocity-Verlet finalize) on owned batch.

8. Atom migration (reshard_by_destination) for atoms that crossed domain boundaries.

9. Outer AFTER_STEP hooks on owned batch.

Parameters:

• dynamics (BaseDynamics) – The underlying single-GPU dynamics integrator or optimizer.

• config (DomainConfig) – Domain decomposition configuration.

• **kwargs (Any) – Forwarded to BaseDynamics.__init__ (hooks, n_steps, device_type, …).

close()

Release resources held by the adapter (restores any state its distributed_setup mutated on the inner wrapper). Safe to call multiple times.

Return type:

None

gather(local_batch, dst=0)

Gather the distributed system back into a full Batch on rank dst. Returns None on other ranks.

Single-process fallback: returns local_batch unchanged.

Parameters:

• local_batch (Batch)

• dst (int)

Return type:

Batch | None

partition(batch)

Scatter the full-system batch across ranks and build the per-step machinery (ShardedBatch + DistributedModel).

Must be called once before run() / step().

Parameters:

batch (Batch | None) – Full-system batch on rank 0; None elsewhere. In the single-process fallback (no distributed init), passes through.

Returns:

This rank’s owned local batch (per-atom fields are .to_local() views of the ShardedBatch’s ShardTensors).

Return type:

Batch

run(batch, n_steps=None)

Run the domain-decomposed simulation for n_steps steps.

Parameters:

• batch (Batch)

• n_steps (int | None)

Return type:

Batch

step(batch)

Execute one domain-decomposed dynamics step.

Parameters:

batch (Batch)

Return type:

tuple[Batch, torch.Tensor | None]