credit.models.crossformer
=========================

.. py:module:: credit.models.crossformer


Attributes
----------

.. autoapisummary::

   credit.models.crossformer.logger
   credit.models.crossformer.image_height


Classes
-------

.. autoapisummary::

   credit.models.crossformer.CubeEmbedding
   credit.models.crossformer.UpBlock
   credit.models.crossformer.CrossEmbedLayer
   credit.models.crossformer.DynamicPositionBias
   credit.models.crossformer.LayerNorm
   credit.models.crossformer.FeedForward
   credit.models.crossformer.Attention
   credit.models.crossformer.Transformer
   credit.models.crossformer.CrossFormer


Functions
---------

.. autoapisummary::

   credit.models.crossformer.cast_tuple
   credit.models.crossformer.apply_spectral_norm


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

.. py:data:: logger

.. py:function:: cast_tuple(val, length=1)

.. py:function:: apply_spectral_norm(model)

.. py:class:: CubeEmbedding(img_size, patch_size, in_chans, embed_dim, norm_layer=nn.LayerNorm)

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


   :param img_size: T, Lat, Lon
   :param patch_size: T, Lat, Lon


   .. py:attribute:: img_size


   .. py:attribute:: patches_resolution


   .. py:attribute:: embed_dim


   .. py:attribute:: proj


   .. py:method:: forward(x: torch.Tensor)


.. py:class:: UpBlock(in_chans, out_chans, num_groups, num_residuals=2, upsample_v_conv=False, attention_type=None, reduction=32, spatial_kernel=7)

   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:: upsample_v_conv
      :value: False



   .. py:attribute:: output_channels


   .. py:attribute:: b


   .. py:attribute:: attention
      :value: None



   .. py:method:: forward(x)


.. py:class:: CrossEmbedLayer(dim_in, dim_out, kernel_sizes, stride=2)

   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:: convs


   .. py:method:: forward(x)


.. py:class:: DynamicPositionBias(dim)

   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:: layers


   .. py:method:: forward(x)


.. py:class:: LayerNorm(dim, eps=1e-05)

   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:: eps
      :value: 1e-05



   .. py:attribute:: g


   .. py:attribute:: b


   .. py:method:: forward(x)


.. py:class:: FeedForward(dim, mult=4, dropout=0.0)

   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:: layers


   .. py:method:: forward(x)


.. py:class:: Attention(dim, attn_type, window_size, dim_head=32, dropout=0.0)

   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:: heads


   .. py:attribute:: scale
      :value: 0.1767766952966369



   .. py:attribute:: attn_type


   .. py:attribute:: window_size


   .. py:attribute:: norm


   .. py:attribute:: dropout


   .. py:attribute:: to_qkv


   .. py:attribute:: to_out


   .. py:attribute:: dpb


   .. py:method:: forward(x)


.. py:class:: Transformer(dim, *, local_window_size, global_window_size, depth=4, dim_head=32, attn_dropout=0.0, ff_dropout=0.0)

   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:: layers


   .. py:method:: forward(x)


.. py:class:: CrossFormer(image_height: int = 640, patch_height: int = 1, image_width: int = 1280, patch_width: int = 1, frames: int = 2, output_frames: int = 1, channels: int = 4, surface_channels: int = 7, input_only_channels: int = 3, output_only_channels: int = 0, levels: int = 15, dim: tuple = (64, 128, 256, 512), depth: tuple = (2, 2, 8, 2), dim_head: int = 32, global_window_size: tuple = (5, 5, 2, 1), local_window_size: int = 10, cross_embed_kernel_sizes: tuple = ((4, 8, 16, 32), (2, 4), (2, 4), (2, 4)), cross_embed_strides: tuple = (4, 2, 2, 2), attn_dropout: float = 0.0, ff_dropout: float = 0.0, use_spectral_norm: bool = True, attention_type: str = None, interp: bool = True, upsample_v_conv: bool = False, padding_conf: dict = None, post_conf: dict = None, **kwargs)

   Bases: :py:obj:`credit.models.base_model.BaseModel`


   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:: image_height
      :value: 640



   .. py:attribute:: image_width
      :value: 1280



   .. py:attribute:: patch_height
      :value: 1



   .. py:attribute:: patch_width
      :value: 1



   .. py:attribute:: upsample_v_conv
      :value: False



   .. py:attribute:: frames
      :value: 2



   .. py:attribute:: output_frames
      :value: 1



   .. py:attribute:: channels
      :value: 4



   .. py:attribute:: surface_channels
      :value: 7



   .. py:attribute:: levels
      :value: 15



   .. py:attribute:: use_spectral_norm
      :value: True



   .. py:attribute:: use_interp
      :value: True



   .. py:attribute:: use_padding


   .. py:attribute:: use_post_block


   .. py:attribute:: input_only_channels
      :value: 3



   .. py:attribute:: base_input_channels
      :value: 70



   .. py:attribute:: input_channels
      :value: 140



   .. py:attribute:: base_output_channels
      :value: 67



   .. py:attribute:: output_channels
      :value: 67



   .. py:attribute:: layers


   .. py:attribute:: cube_embedding


   .. py:attribute:: up_block1


   .. py:attribute:: up_block2


   .. py:attribute:: up_block3


   .. py:method:: forward(x)


   .. py:method:: rk4(x)


.. py:data:: image_height
   :value: 180