agedi.api.sampling
==================

.. py:module:: agedi.api.sampling

.. autoapi-nested-parse::

   Sampling from a trained diffusion model.



Functions
---------

.. autoapisummary::

   agedi.api.sampling.sample


Module Contents
---------------

.. py:function:: sample(diffusion: Agedi, *, n_samples: int, n_atoms: Optional[int] = None, atomic_numbers: Optional[List[int]] = None, formula: Optional[str] = None, positions: Optional[numpy.ndarray] = None, cell: Optional[numpy.ndarray] = None, pbc: Optional[numpy.ndarray] = None, template: Optional[Union[agedi.data.AtomsGraph, ase.Atoms]] = None, confinement: Optional[Tuple[float, float]] = None, compile: bool = False, steps: int = 500, eps: float = 0.001, batch_size: int = 64, ff_guidance: Optional[agedi.diffusion.ForcefieldGuidanceConfig] = None, property: Optional[Dict[str, float]] = None, progress_bar: bool = False, save_trajectory: bool = False, print_timings: bool = False, as_atoms: bool = True) -> Union[List[agedi.data.AtomsGraph], List[ase.Atoms], List[List[agedi.data.AtomsGraph]], List[List[ase.Atoms]]]

   Sample structures from a trained diffusion model.

   :param diffusion: A trained :class:`~agedi.Agedi` model.
   :param n_samples: Number of structures to generate.
   :param n_atoms: Number of atoms per structure. Automatically determined from
                   ``formula`` if provided, or from the length of ``atomic_numbers``
                   when ``n_atoms`` is not explicitly given.
   :param atomic_numbers: Atomic numbers of the generated atoms.  Not required when the model
                          has a types-noiser or when ``formula`` is provided.
   :param formula: Chemical formula (e.g. ``"H2O"``).  Used to derive ``n_atoms`` and
                   ``atomic_numbers`` when they are not provided explicitly.
   :param positions: Fixed positions of the atoms (shape ``(n_atoms, 3)``).  Required
                     when no positions-noiser is configured (type-only diffusion).
                     Positions will not be modified during sampling.
   :param cell: Unit-cell matrix (3×3 array or flat length-9 array).  Not required
                when ``template`` is provided (the template's cell is used instead).
   :param pbc: Periodic boundary conditions as a length-3 boolean array (e.g.
               ``[True, True, False]``).  When ``template`` is provided its ``pbc``
               is used unless this argument is given explicitly.  Defaults to
               ``[True, True, True]`` (fully periodic) when neither ``template``
               nor ``pbc`` is supplied.
   :param template: Template structure.  May be an :class:`~agedi.AtomsGraph` or an
                    ASE :class:`~ase.Atoms` object; the latter is automatically converted
                    to an :class:`~agedi.AtomsGraph` (with ``confinement`` applied when
                    provided).  When given, ``cell`` and ``pbc`` are taken from the
                    template unless explicitly provided.
   :param ff_guidance: Force-field guidance configuration.  When ``None`` (default) a
                       :class:`~agedi.diffusion.ForcefieldGuidanceConfig` with default
                       values is used (i.e. guidance is disabled).
   :param compile: When ``True``, use ``torch.compile`` on the reverse diffusion step
                   for faster sampling.  Before the sampling loop starts, the maximum
                   number of neighbors and cell-list dimensions are estimated
                   automatically via NVIDIA nvalchemiops
                   (``estimate_max_neighbors`` and ``estimate_cell_list_sizes``), and
                   all neighbor-list buffers are pre-allocated with fixed shapes.
                   Requires NVIDIA nvalchemiops.  Defaults to ``False``.
   :param print_timings: When ``True``, print a per-stage timing breakdown at the end of
                         each sampling batch (graph init, score model, denoise, neighbor
                         list, etc.).  Defaults to ``False``.