agedi.models.schnetpack.regressor_heads
=======================================

.. py:module:: agedi.models.schnetpack.regressor_heads


Classes
-------

.. autoapisummary::

   agedi.models.schnetpack.regressor_heads.Forces
   agedi.models.schnetpack.regressor_heads.Energy


Functions
---------

.. autoapisummary::

   agedi.models.schnetpack.regressor_heads.build_gated_equivariant_mlp


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

.. py:function:: build_gated_equivariant_mlp(s_in: int, v_in: int, n_out: int, n_layers: int = 2, activation: Callable = F.silu, sactivation: Callable = F.silu) -> torch.nn.Sequential

   Build neural network analog to MLP with `GatedEquivariantBlock`s instead of dense layers.

   :param n_in: Number of input nodes.
   :type n_in: int
   :param n_out: Number of output nodes.
   :type n_out: int
   :param n_layers: Number of layers.
   :type n_layers: int
   :param activation: Activation function.
   :type activation: Callable
   :param sactivation: Activation function for the skip connection.
   :type sactivation: Callable
   :param n_hidden: Number of hidden nodes.
   :type n_hidden: int

   :rtype: nn.Module


.. py:class:: Forces(input_dim_scalar: int = 64, input_dim_vector: int = 64, gated_blocks: int = 3, **kwargs)

   Bases: :py:obj:`agedi.models.head.Head`


   Predict the atomic force on the atoms in the structure.

   :param input_dim_scalar: The dimension of the scalar input.
   :type input_dim_scalar: int
   :param input_dim_vector: The dimension of the vector input.
   :type input_dim_vector: int
   :param gated_blocks: The number of gated blocks in the network.
   :type gated_blocks: int

   :rtype: Head


   .. py:attribute:: _key
      :value: 'forces'



   .. py:attribute:: input_dim_scalar
      :value: 64



   .. py:attribute:: input_dim_vector
      :value: 64



   .. py:attribute:: gated_blocks
      :value: 3



   .. py:attribute:: net


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

      Return hyperparameters sufficient to reconstruct this head.

      :returns: Hyperparameter dictionary.
      :rtype: dict



   .. py:method:: _score(translated_batch: dict) -> torch.Tensor

      Predict the force on the atoms in the structure.

      :param translated_batch: The translated input batch.
      :type translated_batch: dict

      :returns: The predicted forces tensor.
      :rtype: torch.Tensor



   .. py:method:: predict(translated_batch: dict) -> torch.Tensor

      Predict forces – alias for :meth:`Forces._score` kept for backwards compatibility.

      :param translated_batch: The translated input batch.
      :type translated_batch: dict

      :returns: The predicted forces tensor.
      :rtype: torch.Tensor



.. py:class:: Energy(input_dim_scalar: int = 64, **kwargs)

   Bases: :py:obj:`agedi.models.head.Head`


   Predict the potential energy of the structure.

   :param input_dim_scalar: The dimension of the scalar input.
   :type input_dim_scalar: int

   :rtype: Head


   .. py:attribute:: _key
      :value: 'energy'



   .. py:attribute:: input_dim_scalar
      :value: 64



   .. py:attribute:: net


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

      Return hyperparameters sufficient to reconstruct this head.

      :returns: Hyperparameter dictionary.
      :rtype: dict



   .. py:method:: _score(translated_batch: dict) -> torch.Tensor

      Predict the force on the atoms in the structure.

      :param translated_batch: The translated input batch.
      :type translated_batch: dict

      :returns: The predicted forces tensor.
      :rtype: torch.Tensor



   .. py:method:: predict(translated_batch: dict) -> torch.Tensor

      Predict forces – alias for :meth:`Forces._score` kept for backwards compatibility.

      :param translated_batch: The translated input batch.
      :type translated_batch: dict

      :returns: The predicted forces tensor.
      :rtype: torch.Tensor