MolecularDiffusion.runmodes.train.eval¶

Attributes¶

`ANGLE_RELAX`
`DISTRIBUTED_DEFAULT_TIMEOUT_SEC`
`DIST_RELAX_BOND`
`DIST_THRESHOLD`
`SCALE_FACTOR`
`TASK_REGISTRY`

Classes¶

TaskEvalConfig

Declarative evaluation config for a task type.

Functions¶

`analyze_and_save`(→ Dict[str, Any])	Samples molecules from a generative model, saves them as XYZ files,
`evaluate`(task, solver[, epoch, current_best_metric, ...])	Unified evaluation entry point for all task types.
`get_versioned_output_path`(→ str)	Get next available version folder (engine_logs/version_X).

Module Contents¶

class MolecularDiffusion.runmodes.train.eval.TaskEvalConfig¶

Declarative evaluation config for a task type.

higher_is_better: bool¶

mae_eval: bool = False¶

metric_key: str¶

needs_generative: bool¶

split_reset_hook: str | None¶

MolecularDiffusion.runmodes.train.eval.analyze_and_save(model, epoch: int, n_samples: int = 1000, batch_size: int = 100, logger: Literal['wandb', 'logging'] = 'logging', path_save: str = 'samples', use_posebuster: bool = False, postbuster_timeout: int = 60) → Dict[str, Any]¶

Samples molecules from a generative model, saves them as XYZ files, and computes structural validity statistics.

Parameters:

model – The generative model used for sampling.
epoch (int) – The current training epoch (for logging purposes).
n_samples (int) – Total number of molecules to sample.
batch_size (int) – Number of molecules sampled per batch.
logger (str) – Logging backend, either “wandb” or “logging”.
path_save (str) – Directory to save the sampled XYZ files and CSV.

Returns:

Dictionary summarizing validity and connectivity statistics.

Return type:

Dict[str, Any]

MolecularDiffusion.runmodes.train.eval.evaluate(task: str, solver: MolecularDiffusion.core.Engine, epoch: int = 0, current_best_metric: float = torch.inf, best_checkpoints: list = None, logger: Literal['wandb', 'logging'] = 'logging', output_path: str = None, use_amp: bool = False, precision: str = 'bf16', **kwargs)¶

Unified evaluation entry point for all task types.

Resolves task → TaskEvalConfig via TASK_REGISTRY, runs _run_eval(), then executes one shared compare-and-checkpoint block.

Returns:: (current_best_metric, best_checkpoints)
Return type:: Tuple[float, list]

MolecularDiffusion.runmodes.train.eval.get_versioned_output_path(base_output_path: str) → str¶

Get next available version folder (engine_logs/version_X).

This mimics Lightning’s behavior of creating version_0, version_1, etc.

Parameters:: base_output_path (str) – The base output directory (e.g., training_outputs/my_model)
Returns:: Path to the versioned checkpoint folder (e.g., training_outputs/my_model/engine_logs/version_0)
Return type:: str

MolecularDiffusion.runmodes.train.eval.ANGLE_RELAX = 20¶

MolecularDiffusion.runmodes.train.eval.DISTRIBUTED_DEFAULT_TIMEOUT_SEC = 1800¶

MolecularDiffusion.runmodes.train.eval.DIST_RELAX_BOND = 0.25¶

MolecularDiffusion.runmodes.train.eval.DIST_THRESHOLD = 3¶

MolecularDiffusion.runmodes.train.eval.SCALE_FACTOR = 1.2¶

MolecularDiffusion.runmodes.train.eval.TASK_REGISTRY: dict[str, TaskEvalConfig]¶