credit.postblock
================

.. py:module:: credit.postblock


Submodules
----------

.. toctree::
   :maxdepth: 1

   /autoapi/credit/postblock/_postblock/index
   /autoapi/credit/postblock/wet_mask_samudra/index


Classes
-------

.. autoapisummary::

   credit.postblock.PostBlock
   credit.postblock.TracerFixer
   credit.postblock.GlobalMassFixer
   credit.postblock.GlobalWaterFixer
   credit.postblock.GlobalEnergyFixer
   credit.postblock.WetMaskBlock


Package Contents
----------------

.. py:class:: PostBlock(post_conf)

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


   Base class for all neural network modules.

   Your models should also subclass this class.

   Modules can also contain other Modules, allowing them to be nested in
   a tree structure. You can assign the submodules as regular attributes::

       import torch.nn as nn
       import torch.nn.functional as F


       class Model(nn.Module):
           def __init__(self) -> None:
               super().__init__()
               self.conv1 = nn.Conv2d(1, 20, 5)
               self.conv2 = nn.Conv2d(20, 20, 5)

           def forward(self, x):
               x = F.relu(self.conv1(x))
               return F.relu(self.conv2(x))

   Submodules assigned in this way will be registered, and will also have their
   parameters converted when you call :meth:`to`, etc.

   .. note::
       As per the example above, an ``__init__()`` call to the parent class
       must be made before assignment on the child.

   :ivar training: Boolean represents whether this module is in training or
                   evaluation mode.
   :vartype training: bool


   .. py:attribute:: operations


   .. py:method:: forward(x)


.. py:class:: TracerFixer(post_conf)

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


   This module fixes tracer values by replacing their values to a given threshold
   (e.g., `tracer[tracer<thres] = thres`).

   :param post_conf: config dictionary that includes all specs for the tracer fixer.
   :type post_conf: dict


   .. py:attribute:: tracer_indices


   .. py:attribute:: tracer_thres


   .. py:attribute:: tracer_thres_max


   .. py:method:: forward(x)


.. py:class:: GlobalMassFixer(post_conf)

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


   This module applies global mass conservation fixes for both dry air and water budget.
   The output ensures that the global dry air mass and global water budgets are conserved
   through correction ratios applied during model runs. Variables `specific total water`
   and `precipitation` will be corrected to close the budget. All corrections are done
   using float32 PyTorch tensors.

   :param post_conf: config dictionary that includes all specs for the global mass fixer.
   :type post_conf: dict


   .. py:attribute:: q_ind_start


   .. py:attribute:: q_ind_end


   .. py:method:: forward(x)


.. py:class:: GlobalWaterFixer(post_conf)

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


   Base class for all neural network modules.

   Your models should also subclass this class.

   Modules can also contain other Modules, allowing them to be nested in
   a tree structure. You can assign the submodules as regular attributes::

       import torch.nn as nn
       import torch.nn.functional as F


       class Model(nn.Module):
           def __init__(self) -> None:
               super().__init__()
               self.conv1 = nn.Conv2d(1, 20, 5)
               self.conv2 = nn.Conv2d(20, 20, 5)

           def forward(self, x):
               x = F.relu(self.conv1(x))
               return F.relu(self.conv2(x))

   Submodules assigned in this way will be registered, and will also have their
   parameters converted when you call :meth:`to`, etc.

   .. note::
       As per the example above, an ``__init__()`` call to the parent class
       must be made before assignment on the child.

   :ivar training: Boolean represents whether this module is in training or
                   evaluation mode.
   :vartype training: bool


   .. py:attribute:: q_ind_start


   .. py:attribute:: q_ind_end


   .. py:attribute:: precip_ind


   .. py:attribute:: evapor_ind


   .. py:method:: forward(x)


.. py:class:: GlobalEnergyFixer(post_conf)

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


   This module applys global energy conservation fixes. The output ensures that the global sum
   of total energy in the atmosphere is balanced by radiantion and energy fluxes at the top of
   the atmosphere and the surface. Variables `air temperature` will be modified to close the
   budget. All corrections are done using float32 Pytorch tensors.

   :param post_conf: config dictionary that includes all specs for the global energy fixer.
   :type post_conf: dict


   .. py:attribute:: T_ind_start


   .. py:attribute:: T_ind_end


   .. py:attribute:: q_ind_start


   .. py:attribute:: q_ind_end


   .. py:attribute:: U_ind_start


   .. py:attribute:: U_ind_end


   .. py:attribute:: V_ind_start


   .. py:attribute:: V_ind_end


   .. py:attribute:: TOA_solar_ind


   .. py:attribute:: TOA_OLR_ind


   .. py:attribute:: surf_solar_ind


   .. py:attribute:: surf_LR_ind


   .. py:attribute:: surf_SH_ind


   .. py:attribute:: surf_LH_ind


   .. py:method:: forward(x)


.. py:class:: WetMaskBlock(conf)

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


   Post-processing layer that applies wet mask to ocean predictions
   Zero trainable parameters, but mask influences gradients

   Masks predictions so land points = 0, ocean points preserve values.
   This encourages the model to focus learning on ocean regions.


   .. py:method:: forward(predictions)

      Apply wet mask to predictions with gradient influence

      :param predictions: tensor of shape (batch, n_vars, time, lat, lon)

      :returns: same shape, with land values set to zero
      :rtype: masked_predictions