agedi.api._registry

Model-backend registry and builder helpers.

Attributes

_MODEL_REGISTRY

Functions

register_model(→ None)

Register a custom score model backbone factory under name.

_resolve_sde(→ agedi.diffusion.sdes.SDE)

Resolve an SDE alias string to an SDE instance.

_build_type_map_from_data(→ List[int])

Build a compact type map from the element types present in training data.

_build_noisers(→ List[agedi.diffusion.noisers.Noiser])

Build a list of Noiser objects from a sequence of noiser names or objects.

_build_conditioning(→ List[Conditioning])

Build a list of conditioning modules.

_build_score_components(→ Tuple[Translator, ...)

Instantiate the translator, representation, and score heads for a model.

_painn_factory(→ Tuple[Translator, torch.nn.Module, ...)

Factory for the SchNetPack PaiNN score model backend.

_build_regressor(→ agedi.models.regressor.RegressorModel)

Build a RegressorModel with an Energy and a Forces head.

Module Contents

agedi.api._registry._MODEL_REGISTRY: Dict[str, Callable]
agedi.api._registry.register_model(name: str, factory: Callable) → None

Register a custom score model backbone factory under name.

The factory is called with the keyword arguments cutoff, heads, feature_size, n_blocks, head_dim, and n_rbf and must return a 3-tuple (translator, representation, List[Head]).

Registered models can be selected by passing model=name to create_diffusion().

Parameters:

  • name (str) – Alias used to select this backend (e.g. "PaiNN").

  • factory (Callable) –

    Factory function with signature:

    factory(cutoff, heads, feature_size, n_blocks, head_dim, n_rbf)
    -> Tuple[Translator, nn.Module, List[Head]]

Examples

from agedi.functional import register_model

def my_factory(cutoff, heads, feature_size, n_blocks, head_dim, n_rbf):
    ...
    return translator, representation, head_list

register_model("MyModel", my_factory)
agedi.api._registry._resolve_sde(alias: str | agedi.diffusion.sdes.SDE) → agedi.diffusion.sdes.SDE

Resolve an SDE alias string to an SDE instance.

Parameters:

alias (str or SDE) – A short alias string or an already-instantiated SDE object. Recognised aliases are "ve" and "vp".

Returns:

The resolved SDE instance.

Return type:

SDE

agedi.api._registry._build_type_map_from_data(data: Sequence[Atoms]) → List[int]

Build a compact type map from the element types present in training data.

The map is [0, z1, z2, ...] where z1 < z2 < ... are the sorted unique atomic numbers found in data. Index 0 is reserved for the absorbing state.

Parameters:

data (Sequence[Atoms]) – List of ASE Atoms objects to inspect.

Returns:

A list where type_map[i] is the atomic number corresponding to compact index i (and type_map[0] == 0 for the absorbing state).

Return type:

List[int]

agedi.api._registry._build_noisers(noisers: Sequence[str | agedi.diffusion.noisers.Noiser], sde: str | SDE = 've', type_map: List[int] | None = None) → List[agedi.diffusion.noisers.Noiser]

Build a list of Noiser objects from a sequence of noiser names or objects.

Parameters:

  • noisers (Sequence[Union[str, Noiser]]) –

    A sequence of noiser identifiers or already-instantiated Noiser objects. String identifiers are resolved via the noiser registry (see register()). Built-in identifiers (CamelCase preferred; snake_case aliases also accepted):

    • "Positions" – Positions (StandardNormal prior + Normal distribution, for clusters).

    • "CellPositions" – CellPositions (UniformCell prior + Normal distribution, for periodic systems).

    • "ConfinedCellPositions" – ConfinedCellPositions (UniformCellConfined prior + TruncatedNormal distribution, for Z-confined surfaces/porous materials).

    • "Types" – Types.

  • sde (str or SDE, optional) – Stochastic differential equation to use for position noisers. Either a short alias ("ve", "vp") or an already-instantiated SDE object. Defaults to "ve".

  • type_map (List[int], optional) – Compact type map (see _build_type_map_from_data()) to pass to the Types noiser when building it from a string identifier. Ignored for all other noiser types and for already-instantiated noiser objects.

Returns:

Instantiated noisers in the same order as noisers.

Return type:

List[Noiser]

agedi.api._registry._build_conditioning(condition: str, type: str | None = None) → List[Conditioning]

Build a list of conditioning modules.

Always includes a TimeConditioning. When condition is not "none", an additional property-conditioning module is appended.

Parameters:

  • condition (str) – Name of the property to condition on, or "none" for time-only conditioning.

  • type (str, optional) – Type of the conditioning module: "scalar" or "integer". Required when condition is not "none".

Returns:

The list of conditioning modules.

Return type:

List[Conditioning]

agedi.api._registry._build_score_components(model: str, cutoff: float, heads: Sequence[str], feature_size: int, n_blocks: int, head_dim: int, n_rbf: int = 30) → Tuple[Translator, torch.nn.Module, List[Head]]

Instantiate the translator, representation, and score heads for a model.

Parameters:

  • model (str) – Name of the GNN backbone. The name is looked up in the model registry populated via register_model(). Use "PaiNN" for the built-in SchNetPack PaiNN backend.

  • cutoff (float) – Cutoff radius (Å) for the neighbour list.

  • heads (Sequence[str]) – Names of score heads to build ("positions", "types").

  • feature_size (int) – Embedding/feature dimension for the backbone.

  • n_blocks (int) – Number of interaction blocks in the backbone.

  • head_dim (int) – Input dimension for each score head (typically feature_size + conditioning output dims).

  • n_rbf (int, optional) – Number of radial basis functions. Default is 30.

Returns:

A 3-tuple of the translator, the representation backbone, and the list of score-head modules.

Return type:

Tuple[Translator, nn.Module, List[Head]]

agedi.api._registry._painn_factory(cutoff: float, heads: Sequence[str], feature_size: int, n_blocks: int, head_dim: int, n_rbf: int) → Tuple[Translator, torch.nn.Module, List[Head]]

Factory for the SchNetPack PaiNN score model backend.

agedi.api._registry._build_regressor(translator: Translator, representation: Representation, feature_size: int) → agedi.models.regressor.RegressorModel

Build a RegressorModel with an Energy and a Forces head.

The force field regressor shares the translator and representation from the score model so that the atomic embeddings are learned jointly with the diffusion score. Only the Forces and Energy heads are added on top of the shared representation.

The resulting model is attached to the Agedi object as regressor_model so that force-field guidance can be used during sampling (see ForcefieldGuidanceConfig).

Parameters:

  • translator (Translator) – The translator from the score model (shared, not copied).

  • representation (Representation) – The representation from the score model (shared, not copied).

  • feature_size (int) – Embedding/feature dimension of the shared representation.

Returns:

An initialised force-regression model (not yet trained).

Return type:

RegressorModel