agedi.api._registry =================== .. py:module:: agedi.api._registry .. autoapi-nested-parse:: Model-backend registry and builder helpers. Attributes ---------- .. autoapisummary:: agedi.api._registry._MODEL_REGISTRY Functions --------- .. autoapisummary:: agedi.api._registry.register_model agedi.api._registry._resolve_sde agedi.api._registry._build_type_map_from_data agedi.api._registry._build_noisers agedi.api._registry._build_conditioning agedi.api._registry._build_score_components agedi.api._registry._painn_factory agedi.api._registry._build_regressor Module Contents --------------- .. py:data:: _MODEL_REGISTRY :type: Dict[str, Callable] .. py:function:: 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 :func:`create_diffusion`. :param name: Alias used to select this backend (e.g. ``"PaiNN"``). :type name: str :param factory: Factory function with signature:: factory(cutoff, heads, feature_size, n_blocks, head_dim, n_rbf) -> Tuple[Translator, nn.Module, List[Head]] :type factory: Callable .. rubric:: 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) .. py:function:: _resolve_sde(alias: Union[str, agedi.diffusion.sdes.SDE]) -> agedi.diffusion.sdes.SDE Resolve an SDE alias string to an :class:`~agedi.diffusion.sdes.SDE` instance. :param alias: A short alias string or an already-instantiated SDE object. Recognised aliases are ``"ve"`` and ``"vp"``. :type alias: str or SDE :returns: The resolved SDE instance. :rtype: SDE .. py:function:: _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. :param data: List of ASE :class:`~ase.Atoms` objects to inspect. :type data: Sequence[Atoms] :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). :rtype: List[int] .. py:function:: _build_noisers(noisers: Sequence[Union[str, agedi.diffusion.noisers.Noiser]], sde: Union[str, SDE] = 've', type_map: Optional[List[int]] = None) -> List[agedi.diffusion.noisers.Noiser] Build a list of Noiser objects from a sequence of noiser names or objects. :param noisers: A sequence of noiser identifiers or already-instantiated :class:`~agedi.diffusion.noisers.Noiser` objects. String identifiers are resolved via the noiser registry (see :meth:`~agedi.diffusion.noisers.Noiser.register`). Built-in identifiers (CamelCase preferred; snake_case aliases also accepted): * ``"Positions"`` – :class:`~agedi.diffusion.noisers.Positions` (StandardNormal prior + Normal distribution, for clusters). * ``"CellPositions"`` – :class:`~agedi.diffusion.noisers.CellPositions` (UniformCell prior + Normal distribution, for periodic systems). * ``"ConfinedCellPositions"`` – :class:`~agedi.diffusion.noisers.ConfinedCellPositions` (UniformCellConfined prior + TruncatedNormal distribution, for Z-confined surfaces/porous materials). * ``"Types"`` – :class:`~agedi.diffusion.noisers.Types`. :type noisers: Sequence[Union[str, Noiser]] :param sde: Stochastic differential equation to use for position noisers. Either a short alias (``"ve"``, ``"vp"``) or an already-instantiated :class:`~agedi.diffusion.sdes.SDE` object. Defaults to ``"ve"``. :type sde: str or SDE, optional :param type_map: Compact type map (see :func:`_build_type_map_from_data`) to pass to the :class:`~agedi.diffusion.noisers.Types` noiser when building it from a string identifier. Ignored for all other noiser types and for already-instantiated noiser objects. :type type_map: List[int], optional :returns: Instantiated noisers in the same order as *noisers*. :rtype: List[Noiser] .. py:function:: _build_conditioning(condition: str, type: Optional[str] = None) -> List[Conditioning] Build a list of conditioning modules. Always includes a :class:`~agedi.models.conditionings.TimeConditioning`. When *condition* is not ``"none"``, an additional property-conditioning module is appended. :param condition: Name of the property to condition on, or ``"none"`` for time-only conditioning. :type condition: str :param type: Type of the conditioning module: ``"scalar"`` or ``"integer"``. Required when *condition* is not ``"none"``. :type type: str, optional :returns: The list of conditioning modules. :rtype: List[Conditioning] .. py:function:: _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. :param model: Name of the GNN backbone. The name is looked up in the model registry populated via :func:`register_model`. Use ``"PaiNN"`` for the built-in SchNetPack PaiNN backend. :type model: str :param cutoff: Cutoff radius (Å) for the neighbour list. :type cutoff: float :param heads: Names of score heads to build (``"positions"``, ``"types"``). :type heads: Sequence[str] :param feature_size: Embedding/feature dimension for the backbone. :type feature_size: int :param n_blocks: Number of interaction blocks in the backbone. :type n_blocks: int :param head_dim: Input dimension for each score head (typically ``feature_size + conditioning output dims``). :type head_dim: int :param n_rbf: Number of radial basis functions. Default is 30. :type n_rbf: int, optional :returns: A 3-tuple of the translator, the representation backbone, and the list of score-head modules. :rtype: Tuple[Translator, nn.Module, List[Head]] .. py:function:: _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. .. py:function:: _build_regressor(translator: Translator, representation: Representation, feature_size: int) -> agedi.models.regressor.RegressorModel Build a :class:`~agedi.models.regressor.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 :class:`~agedi.models.schnetpack.regressor_heads.Forces` and :class:`~agedi.models.schnetpack.regressor_heads.Energy` heads are added on top of the shared representation. The resulting model is attached to the :class:`~agedi.Agedi` object as ``regressor_model`` so that force-field guidance can be used during sampling (see :class:`~agedi.diffusion.ForcefieldGuidanceConfig`). :param translator: The translator from the score model (shared, not copied). :type translator: Translator :param representation: The representation from the score model (shared, not copied). :type representation: Representation :param feature_size: Embedding/feature dimension of the shared representation. :type feature_size: int :returns: An initialised force-regression model (not yet trained). :rtype: RegressorModel