1 lightning.one_step.experiment

lightning.one_step.experiment

Compose and execute reproducible Lightning VIN experiments.

This module provides the top-level config-as-factory boundary that owns run mode dispatch, deterministic seeding, artifact paths, trainer/module/datamodule construction, checkpoint resume, ordinal-binner fitting, summaries, and Optuna studies. AriaNBVExperimentConfig preserves nested config ownership so one TOML snapshot identifies the dataset source, scorer/loss, optimizer, callbacks, and logger without flattening their contracts.

The orchestrator does not move oracle labels into actor-visible model inputs. The datamodule yields :class:aria_nbv.data_handling.VinOracleBatch, while VinLightningModule performs the explicit post-forward CORAL loss and metric use of those labels.

1.1 Classes

Name	Description
AriaNBVExperimentConfig	Own one reproducible VIN Lightning run and its nested factories.

1.1.1 AriaNBVExperimentConfig

lightning.one_step.experiment.AriaNBVExperimentConfig()

Own one reproducible VIN Lightning run and its nested factories.

Calling :meth:setup_target constructs but does not execute the three runtime owners. :meth:setup_target_and_run then delegates loop ownership to Lightning. Run artifacts, checkpoint paths, and fitted binner state are resolved through :class:PathConfig before construction.

1.1.1.1 Attributes

Name	Description
run_mode	Requested experiment action, normalized from kebab or snake case.
run_name	Human-readable identifier used by loggers, checkpoints, and snapshots.
stage	Lightning stage constructed and executed by :meth:setup_target_and_run.
summary_stage	Dataset stage sampled by the non-training VIN summary action.
summary_num_batches	Positive number of real oracle-labelled batches included in a summary.
summary_include_torchsummary	Append external torchsummary module traversal to scorer diagnostics.
summary_torchsummary_depth	Positive maximum module-tree depth for optional torchsummary output.
plot_out_dir	Path for non-training VIN encoding figures, resolved by the plot action.
ckpt_path	Optional checkpoint for strict validation/test loading or training resume.
seed	Torch CPU/CUDA seed set before dataset, scorer, and trainer construction.
out_dir	Run-artifact root for config snapshots, logs, checkpoints, and binner state.
fit_binner_only	Compatibility flag that rewrites run_mode to "fit_binner".
fit_binner_overwrite	Overwrite module_config.binner_path when run_mode="fit_binner".
inspect_config	Print rich datamodule and LightningModule config summaries before running.
paths	Repository-aware resolver for configs, run roots, checkpoints, and artifacts.
trainer_config	Trainer factory including callback ownership, logger, devices, and loop limits.
datamodule_config	VIN source, stage-reuse, worker, and candidate-collation configuration.
module_config	One-step scorer, CORAL loss, binner, optimizer, scheduler, and metric config.
optuna_config	Optional study and search-space config used only by the Optuna action.
verbosity	Experiment-orchestration logging verbosity.
verbose	Legacy Optuna-facing verbosity value retained for config compatibility.
resolved_out_dir	Return out_dir resolved under the configured repository run root.
default_config_path	Default path for saving the experiment TOML.

1.1.1.2 Methods

Name	Description
save_config	Save this config (and nested configs) as TOML.
setup_target	Construct the trainer, scorer module, and stage-aware datamodule.
setup_target_and_run	Construct runtime owners and delegate one stage loop to Lightning.
run_optuna_study	Run an Optuna hyperparameter sweep for this experiment config.
summarize_vin	Run VIN summary on a real oracle batch from the datamodule.
fit_binner_and_save	Fit RriOrdinalBinner from oracle batches and persist it.
run	Execute the configured action.

1.1.1.2.1 save_config

lightning.one_step.experiment.AriaNBVExperimentConfig.save_config(
    path=None,
    *,
    include_comments=True,
    include_type_hints=True,
)

Save this config (and nested configs) as TOML.

1.1.1.2.1.1 Parameters

Name	Type	Description	Default
path	Path | str | None	Destination TOML path. If None, uses default_config_path.	None
include_comments	bool	Include docstring comments in the TOML.	True
include_type_hints	bool	Include type hints in the TOML comments.	True

1.1.1.2.1.2 Returns

Name	Type	Description
	Path	Resolved path that was written.

1.1.1.2.2 setup_target

lightning.one_step.experiment.AriaNBVExperimentConfig.setup_target(
    setup_stage=None,
    *,
    trial=None,
)

Construct the trainer, scorer module, and stage-aware datamodule.

1.1.1.2.2.1 Parameters

Name	Type	Description	Default
setup_stage	Stage | str | None	Dataset stage to initialize; defaults to stage.	None
trial	'optuna.Trial | None'	Optional Optuna trial forwarded to trainer callback setup.	None

1.1.1.2.2.2 Returns

Name	Type	Description
	pl.Trainer	Tuple of external Lightning trainer, VinLightningModule,
and	VinLightningModule	class:VinDataModule. No fit/validation/test loop has run.

1.1.1.2.2.3 Notes

Torch seeds and artifact directories are established before any runtime owner is constructed. Checkpoint weights remain Lightning’s responsibility when the returned trainer starts its loop.

1.1.1.2.3 setup_target_and_run

lightning.one_step.experiment.AriaNBVExperimentConfig.setup_target_and_run(
    stage=None,
)

Construct runtime owners and delegate one stage loop to Lightning.

Training verifies the fitted binner class count before fit; validation and test pass the resolved checkpoint directly to Lightning. A globally ranked interrupted trainer writes one explicit recovery checkpoint.

1.1.1.2.3.1 Returns

Name	Type	Description
	pl.Trainer	Trainer after the requested loop finishes.

1.1.1.2.4 run_optuna_study

lightning.one_step.experiment.AriaNBVExperimentConfig.run_optuna_study()

Run an Optuna hyperparameter sweep for this experiment config.

1.1.1.2.5 summarize_vin

lightning.one_step.experiment.AriaNBVExperimentConfig.summarize_vin()

Run VIN summary on a real oracle batch from the datamodule.

1.1.1.2.6 fit_binner_and_save

lightning.one_step.experiment.AriaNBVExperimentConfig.fit_binner_and_save(
    datamodule=None,
    *,
    overwrite=False,
)

Fit RriOrdinalBinner from oracle batches and persist it.

1.1.1.2.6.1 Parameters

Name	Type	Description	Default
datamodule	VinDataModule | None	Optional pre-built datamodule (avoids re-instantiation).	None
overwrite	bool	Overwrite existing binner JSON when True.	False

1.1.1.2.6.2 Returns

Name	Type	Description
	Path	Path to the saved rri_binner.json.

1.1.1.2.7 run

lightning.one_step.experiment.AriaNBVExperimentConfig.run()

Execute the configured action.

This method is intended to be called from CLI entry points.