chemprop.nn.ffn
===============

.. py:module:: chemprop.nn.ffn


Classes
-------

.. autoapisummary::

   chemprop.nn.ffn.FFN
   chemprop.nn.ffn.MLP
   chemprop.nn.ffn.ConstrainerFFN


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

.. py:class:: FFN(*args, **kwargs)

   Bases: :py:obj:`torch.nn.Module`


   A :class:`FFN` is a differentiable function
   :math:`f_\theta : \mathbb R^i \mapsto \mathbb R^o`


   .. py:attribute:: input_dim
      :type:  int


   .. py:attribute:: output_dim
      :type:  int


   .. py:method:: forward(X)
      :abstractmethod:



.. py:class:: MLP(*args: torch.nn.modules.module.Module)
              MLP(arg)

   Bases: :py:obj:`torch.nn.Sequential`, :py:obj:`FFN`


   An :class:`MLP` is an FFN that implements the following function:

   .. math::
       \mathbf h_0 &= \mathbf W_0 \mathbf x \,+ \mathbf b_{0} \\
       \mathbf h_l &= \mathbf W_l \left( \mathtt{dropout} \left( \sigma ( \,\mathbf h_{l-1}\, ) \right) \right) + \mathbf b_l\\

   where :math:`\mathbf x` is the input tensor, :math:`\mathbf W_l` and :math:`\mathbf b_l`
   are the learned weight matrix and bias, respectively, of the :math:`l`-th layer,
   :math:`\mathbf h_l` is the hidden representation after layer :math:`l`, and :math:`\sigma`
   is the activation function.


   .. py:method:: build(input_dim, output_dim, hidden_dim = 300, n_layers = 1, dropout = 0.0, activation = 'relu')
      :classmethod:



   .. py:property:: input_dim
      :type: int



   .. py:property:: output_dim
      :type: int



.. py:class:: ConstrainerFFN(n_constraints = 1, fp_dim = DEFAULT_HIDDEN_DIM, hidden_dim = 300, n_layers = 1, dropout = 0.0, activation = 'relu')

   Bases: :py:obj:`torch.nn.Module`, :py:obj:`chemprop.nn.hparams.HasHParams`, :py:obj:`lightning.pytorch.core.mixins.HyperparametersMixin`


   A :class:`ConstrainerFFN` adjusts atom or bond property predictions to satisfy molecular
   constraints by using an :class:`MLP` to map learned atom or bond embeddings to weights that
   determine how much of the total adjustment needed is added to each atom or bond prediction.


   .. py:attribute:: ffn


   .. py:method:: forward(fp, preds, batch, constraints)

      Performs a weighted adjustment to the predictions to satisfy the constraints, with the
      weights being determined from the learned atom or bond fingerprints via an :class:`MLP`.

      :param fp: a tensor of shape ``b x h`` containing the atom or bond-level fingerprints, where ``b``
                 is the number of atoms or bonds and ``h`` is the length of each fingerprint.
      :type fp: Tensor
      :param preds: a tensor of shape ``b x t`` containing the atom or bond-level predictions, where ``t``
                    is the number of predictions per atom or bond.
      :type preds: Tensor
      :param batch: a tensor of shape ``b`` containing indices of which molecule each atom or bond belongs to
      :type batch: Tensor
      :param constraints: a tensor of shape ``m x t`` containing the values to which the atom or bond-level
                          predictions should sum to for each molecule, where ``m`` is the number of molecules in
                          the batch.
      :type constraints: Tensor

      :returns: a tensor of shape ``b x t`` containing the atom or bond-level predictions adjusted to
                satisfy the molecule-level constraints
      :rtype: Tensor