App

The imperative shell: the diffbloch CLI, the default experiment runner it delegates to, and the logger backends (console/CSV always available; Weights & Biases and Comet behind optional extras — vendor SDKs never reach the core).

Thin command-line entry point.

This is the orchestration / SLURM boundary: a workflow engine (e.g. Dagster or Prefect) shells out to it, or a SLURM job runs it. Kept deliberately thin — it delegates to the library and holds no science.

diffBloch.app.cli.main(argv: list[str] | None = None) → int

Parse arguments and dispatch to a subcommand. Returns a process exit code.

The default experiment runner the run infer CLI exposes.

run_experiment() encodes the default recipe as one ordered Plan -> Plan pipeline: plan-shaping (build_orientation_plans, selecting coupled SOLVE beams from g_max/sg_max and building rocking geometry plus its reduction) followed by the config-enabled parameter fitting stages (orientation, then thickness), then run_inference evaluates it – so a caller with an experiment directory gets a full result in one call. It is a convenience, not the only path: every step is ordinary public API, so a Python user who wants a different composition composes their own pipeline([...]) with from_experiment + run_inference directly. The CLI stays thin by delegating here and holds no science.

Checkpoint / resume. The preprocess (the coupled fit especially) is the run’s expensive phase, so each dataset’s settled Plan is checkpointed to <experiment_dir>/reproducibility/ (plan.<stem>.npz + plan.<stem>.lock per inputs.exp_data entry, see _reproducibility_dir() and dataset_checkpoint_stem()) and reused when still valid. The lock binds the checkpoint to per-dataset axes – the structure and that dataset’s input bytes, the experiment’s authoritative PETS cell, the dataset-scoped config projection, its file-local ignored rotations, the software version, and the recipe (the Plan’s own provenance) – so it is reused only when all match, and never restaled by changes to other datasets in a pooled experiment unless the first dataset’s authoritative cell changes. The match is a longest-prefix over the recipe: an identical recipe is a full reuse; a recipe that appends steps resumes from the snapshot and runs only the suffix (the append-only increment). Any run that computes – fresh, stale-recompute, refresh, or resume – regenerates the .npz + .lock pair, so the lock always describes the on-disk checkpoint. With inputs.multi_dataset the settled per-dataset plans are pooled in memory (pool()) just before refinement; the pooled plan itself is never checkpointed. Diagnostics about load/resume go through stdlib logging (not the domain-observation logger).

diffBloch.app.program.converge_experiment(experiment_dir: str | Path, *, logger: Logger = NULL_LOGGER, device: Device = 'cuda', n_orientations: int = 1) → ConvergenceState

Run the standard numerical-convergence test for an experiment.

Starting g_max prefers each PETS file’s own dstarmax (its processing-resolution cutoff, dstar_max) when present; a PETS version that doesn’t record it (or a file where the tag is otherwise absent) falls back to the experiment’s configured blochwave.g_max, the previous manual-only behaviour. sg_max and rocking-curve tilt steps still start from the configured simulation settings. The sweep then follows the defaults owned by ConvergenceTest and ConvergenceTolerance, returning the smallest settled values found.

With inputs.multi_dataset the sweep runs per dataset – each file with its own integration geometry, its own dstar_max-or-fallback starting point, and its first n_orientations rotations – and the returned state is the elementwise maximum over the per-dataset settled states: the tightest single setting adequate for every pooled file. The per-dataset settled values remain visible through the convergence logger events.

diffBloch.app.program.preprocess_experiment(experiment_dir: str | Path, *, logger: Logger = NULL_LOGGER, checkpoint: bool = True, refresh: bool = False, device: Device | None = 'cuda', workers: int = 1, max_batch: int | None = None, plot_thickness: bool = False, plot_thickness_dir: str | Path | None = None) → Plan

Load and preprocess the experiment at experiment_dir, returning the settled Plan.

The preprocess half of run_experiment() with the terminal scoring stripped off: it loads experiment.yaml (verifying the input lock), reads the structure + experimental data, builds the geometry via from_experiment, and runs the default integrated recipe – returning the settled coupled Plan (fitted orientations, tilt-segment couplings, pinned scored sets). This is the entry point for callers who want only the calibrated Plan (to checkpoint it, or to drive their own downstream refinement) without paying for the terminal inference pass.

The recipe includes per-trial beam coupling (the fit re-derives the SOLVE union + SCORED set at every trial orientation). Its coupling policy, orientation-search bounds, thickness grid, and whether hydrogens are loaded all come from config (blochwave / preprocess.orientation / preprocess.thickness / inputs.load_hydrogens). A caller wanting a different composition (e.g. the cheaper tilt-independent fit) composes their own pipeline([...]) with the public steps.

checkpoint (default True) reuses/resumes a valid plan.npz in the experiment dir and writes a fresh one after computing; refresh forces a full recompute (ignoring any snapshot) while still regenerating the checkpoint. checkpoint=False neither reads nor writes.

device (default "cuda") runs the forward solve on the selected accelerator for the coupled preprocess fits. If CUDA is requested on a host without CUDA, the app falls back to CPU and emits a device-selection event. The preprocess geometry (beam selection, coupling unions) stays CPU-side numpy; only the eigensolve – the O(N^3) cost – moves to the device (the fits move the seed params, so fgb and every per-trial score co-locate there). Device is execution-only: it does not enter the checkpoint lock, so a committed CPU checkpoint is still reused when a run moves to GPU.

workers (default 1, sequential) fans the per-rotation orientation search over a thread pool (rotations are independent). On a GPU run the per-trial cost is host-bound around a small eigensolve, so overlapping rotations across cores is the main wall-clock lever (a small worker count is usually the sweet spot; gains flatten as the solves serialise on one GPU stream). Like device it is execution-only – the results are identical to a sequential run and it does not enter the checkpoint lock. Cap host threads to 1 when using it – OMP_NUM_THREADS/MKL_NUM_THREADS/TORCH_NUM_THREADS (or torch.set_num_threads(1)); in a pod torch/BLAS size their pools from the node core count, not the cgroup limit, so an uncapped run oversubscribes the cores the workers need (capping alone can dominate the speedup, before any parallelism).

max_batch (default None) caps the matrix_exp propagator block for the coupled fits. None lets each solve derive a memory-safe block from its beam count; raise it to fill a larger accelerator’s memory budget. Execution-only (memory, bit-for-bit to machine precision), out of the checkpoint lock like device/workers. See build_engine().

plot_thickness (default False) ORs with cfg.preprocess.thickness.plot – either turns on one wR2-vs-thickness PNG per rotation from optimize_thickness’s grid search, saved under plot_thickness_dir (default <inputs.structure's directory>/thickness_optim). Execution-only like device/workers, out of the checkpoint lock.

diffBloch.app.program.refine_experiment(experiment_dir: str | Path, *, logger: Logger = NULL_LOGGER, checkpoint: bool = True, refresh: bool = False, device: Device | None = 'cuda', workers: int = 1, max_batch: int | None = None, verbose: bool = False, profile: bool = False, checkpoint_activations: bool = True, plot_thickness: bool = False, plot_thickness_dir: str | Path | None = None) → ModelRefinementResult

Settle the coupled Plan and gradient-refine the structure against the observed data.

preprocess_experiment() for the geometry (checkpoint reuse for free – see it for the recipe and checkpoint/refresh/device/workers/ plot_thickness/plot_thickness_dir semantics), then run the default single-stage refinement on that settled Plan. This is the boring config-knobs path: the residual (to_loss()), the trainable selection (to_spec()), and the optimizer/step budget all come from experiment.yaml. It composes no hard constraints or penalties – scientific composition (hydrogen riding, freeze-H, penalties, multi-stage) is a Python/API concern, built with build_refinement_model(), build_refinement_problem(), and with_hydrogen_riding(), then run via run_refinement_model. The RefinementProblem here is pure optimization-definition data; the imperative loop lives in run_refinement_model. Returns the ModelRefinementResult (per-step losses + best snapshot); _write_refinement_outputs() persists the best structure/parameters/summary to experiment_dir alongside a refinement.lock binding them to the settled Plan (plan.lock), the refinement-determining config, and the code version that produced them – the refinement-stage counterpart to the preprocess checkpoint’s own lock.

device places the refinement solve on the accelerator: the seed params move there and the forward co-locates onto them (as in the preprocess fits).

verbose (“verbose refinement”) reports one per-rotation wR2/R_obs/diffraction-loss line per step in addition to the epoch mean; see run_refinement_model(). Execution-only, like logger.

profile logs per-phase wall time (structure factors, each rotation’s solve, backward, optimizer step) via stdlib diagnostics logging; see run_refinement_model() and build_engine(). Execution-only and off by default.

checkpoint_activations (default True) trades peak memory for one extra forward recompute per solve on backward; disabling it removes that recompute at the cost of retaining solve intermediates until backward. Execution-only – gradients are unaffected. See RefinementEngine.

diffBloch.app.program.run_experiment(experiment_dir: str | Path, *, logger: Logger = NULL_LOGGER, checkpoint: bool = True, refresh: bool = False, device: Device | None = 'cuda', workers: int = 1, max_batch: int | None = None, plot_thickness: bool = False, plot_thickness_dir: str | Path | None = None) → InferenceResult

Load, preprocess, and score every rotation of the experiment at experiment_dir.

preprocess_experiment() followed by the terminal forward model: it settles the coupled Plan (see that function for the recipe, checkpoint/refresh, device/workers, and plot_thickness/plot_thickness_dir semantics – all shared), then evaluates every rotation with run_inference – emitting per-rotation observations to logger (the null default discards them). Returns the InferenceResult (per-rotation R_obs + aggregate). device also runs the terminal eigensolve on the accelerator.

Logger backends – the app/ boundary where sinks and vendor SDKs live, never the core.

Each backend consumes the uniform Event surface (channel + measurements), so none knows the concrete event types and new events need no backend change. ConsoleLogger (here, no vendor dependency) routes events to stdlib logging and CSVLogger (here too) appends them to a file; each third-party backend lives in its own confined submodule that imports its SDK lazily, so importing this package never requires an optional dependency:

• WandbLogger (diffBloch.app.loggers.wandb)

• CometLogger (diffBloch.app.loggers.comet)

• ThicknessPlotLogger (diffBloch.app.loggers.plotting)

Writing your own backend is a single method: implement report(event) for the events you care about.

class diffBloch.app.loggers.CSVLogger(path: Path)

Bases: object

Append each event’s measurements to a CSV file in long format (Lightning-style sink).

One row per measurement – channel, step, metric, value – so a heterogeneous event stream (rotation, inference, refinement) shares a single flat table with no sparse columns, ready to filter by channel or pivot by step. The header is written once at construction (a fresh file per run); each report() appends and flushes, so the file is crash-safe and tailable.

This is an observation log, not persistence: run state is checkpointed by serialising the whole Plan, never reconstructed from these rows.

path: Path

report(event: Event) → None

class diffBloch.app.loggers.ConsoleLogger(level: int = 20, per_rotation: bool = True)

Bases: object

Log each event to stdlib logging at level (console/file handlers attached by app).

The bridge from the domain-observation channel to the diagnostics channel: handy for a live scroll of per-rotation R_obs while chasing a residual. Attach a handler (or call logging.basicConfig()) at the app boundary to see it.

On a real terminal (sys.stdout.isatty()), refinement epochs, orientation-search rotations, and thickness-search rotations render as an in-place progress bar (with a linearly-extrapolated ETA) instead of one scrolling line per event – their respective *Started event supplies the total up front. Off a terminal (piped to a file, CI logs) this falls back to the plain per-event line, since \r-based in-place updates are meaningless there and would just corrupt a log file with control characters.

level: int = 20

per_rotation: bool = True

report(event: Event) → None

class diffBloch.app.loggers.EarlyAbortLogger(wr2_ceiling: float = 0.6, patience: int = 5, inner: Logger = <diffBloch.observability.NullLogger object>)

Bases: object

Watch the per-rotation fit stream and abort a run that is not tracking the data.

A fit-quality guard for a long, oracle-less from-scratch fit – one with no committed checkpoint and no reference R_obs to pin against, so a mis-set-up run (wrong energy / g_max, bad data lineage) would otherwise burn its whole budget producing a bad answer. optimize_orientation emits one OrientationOptimized per rotation as it finishes, carrying the scaling-optimised wr2 at the fitted orientation. A healthy run reaches a low wr2 on essentially every rotation; a fundamentally broken one stays high on all of them. This guard gives the run patience rotations to show at least one orientation reaching wr2 <= wr2_ceiling; if none does, it raises FitAbortedError, unwinding the fit before the remaining rotations run. Pick wr2_ceiling generously (well above a healthy fit, well below a garbage one) so a real run clears it within the first rotation or two and never false-aborts.

Only the fit stream drives the decision; every event is forwarded verbatim to inner (default NULL_LOGGER), so this composes with a ConsoleLogger / CSVLogger for the live scroll – EarlyAbortLogger(inner=ConsoleLogger()). Raising from report() is the abort mechanism: the fit loop’s only per-rotation hook is the logger, and both the sequential and workers > 1 paths call report from the driving thread, so the raise unwinds the run cleanly. Compute saved: all remaining rotations under workers = 1 (sequential – nothing further starts); under workers > 1 the queued rotations are cancelled but the <= workers already running cannot be interrupted and finish first (optimize_orientation cancels the rest on abort).

wr2_ceiling: float = 0.6

patience: int = 5

inner: Logger = <diffBloch.observability.NullLogger object>

report(event: Event) → None

exception diffBloch.app.loggers.FitAbortedError

Bases: RuntimeError

Raised by EarlyAbortLogger to unwind a fit judged unpromising and stop it early.

Carries the diagnostic that triggered the abort (rotations seen, best wr2, the ceiling), so the caller running the fit sees why it stopped, not just that it did.

diffBloch.app.loggers.format_measurements(event: Event) → str

Render an event’s measurements as space-joined name=value pairs (shared by backends).

diffBloch.app.loggers.namespaced_measurements(event: Event) → dict[str, float]

Map an event to {channel}/{metric}: value (the series convention shared by the wandb/comet backends).

diffBloch.app.loggers.residual_label(residual: str) → str

Weights & Biases logger backend – the only module that touches the wandb SDK.

wandb is imported lazily inside WandbLogger.report(), so it is an optional dependency (uv sync --extra wandb): importing diffBloch.app.loggers never requires it, and the pure core never touches it. The app owns the run lifecycle (wandb.init / wandb.finish); this logger only forwards each event’s measurements to the active run as {channel}/{metric} series.

class diffBloch.app.loggers.wandb.WandbLogger

Bases: object

Log each event’s measurements to Weights & Biases as {channel}/{metric} series.

report(event: Event) → None

Comet ML logger backend – confined home for the comet_ml integration.

Unlike Weights & Biases (a module-level wandb.log after wandb.init), Comet ML logs to an Experiment object. So the app owns the run lifecycle – it creates the comet_ml.Experiment (uv sync --extra comet) and passes it in – and this logger only forwards each event’s measurements to that experiment as {channel}/{metric} metric series. The logger itself imports no vendor SDK (it duck-types experiment.log_metrics), so importing diffBloch.app.loggers never requires comet_ml and the pure core never touches it.

class diffBloch.app.loggers.comet.CometLogger(experiment: object)

Bases: object

Log each event’s measurements to a Comet ML experiment as {channel}/{metric} series.

experiment is a comet_ml.Experiment created and owned by the app (Experiment.end() at the end of the run); this logger only calls experiment.log_metrics per event.

experiment: object

report(event: Event) → None