agedi.diffusion.distributions.uniform
=====================================

.. py:module:: agedi.diffusion.distributions.uniform


Classes
-------

.. autoapisummary::

   agedi.diffusion.distributions.uniform.Uniform
   agedi.diffusion.distributions.uniform.UniformCell
   agedi.diffusion.distributions.uniform.UniformCellConfined


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

.. py:class:: Uniform(low: float = 0.0, high: float = 1.0, key: str = 'x', **kwargs)

   Bases: :py:obj:`agedi.diffusion.distributions.Distribution`


   Uniform Distribution

   :param low: The lower bound of the distribution
   :type low: float
   :param high: The upper bound of the distribution
   :type high: float


   .. py:attribute:: low
      :value: 0.0



   .. py:attribute:: high
      :value: 1.0



   .. py:method:: get_hparams() -> Dict

      Return hyperparameters for this distribution.



   .. py:method:: _setup(batch: agedi.data.AtomsGraph) -> None

      Prepare the distribution for sampling from *batch*.

      Sets ``self.shape`` to the shape of the target attribute in the batch.

      :param batch: Batch of atomistic data.
      :type batch: AtomsGraph



   .. py:method:: _sample(shape: Optional[torch.Size] = None, **kwargs) -> torch.Tensor

      Sample from the uniform distribution

      :param shape: The shape of the sample
      :type shape: torch.Size

      :returns: Sampled tensor
      :rtype: torch.Tensor



.. py:class:: UniformCell(low: float = 0.0, high: float = 1.0, key: str = 'x', **kwargs)

   Bases: :py:obj:`Uniform`


   Uniform Prior Distribution for cell parameters


   .. py:method:: _setup(batch: agedi.data.AtomsGraph) -> None

      Prepare the distribution for sampling of the batch

      :param batch: Batch of data
      :type batch: AtomsGraph

      :rtype: None



   .. py:method:: _sample(**kwargs) -> torch.Tensor

      Sample from the uniform distribution

      :param mu: Mean of the distribution
      :type mu: torch.Tensor
      :param sigma: Standard deviation of the distribution
      :type sigma: torch.Tensor

      :returns: Sampled tensor
      :rtype: torch.Tensor



.. py:class:: UniformCellConfined(low: float = 0.0, high: float = 1.0, key: str = 'x', **kwargs)

   Bases: :py:obj:`UniformCell`


   Uniform Prior Distribution for cell parameters with Z-directional confinement


   .. py:method:: _setup(batch: agedi.data.AtomsGraph) -> None

      Prepare the distribution for sampling of the batch

      :param batch: Batch of data
      :type batch: AtomsGraph

      :rtype: None