emmi_inference.models.transolver
================================

.. py:module:: emmi_inference.models.transolver

.. autoapi-nested-parse::

   WARNING
   This file is a 1:1 duplicate from the ``tutorial`` folder.

   It is here to avoid installation of a tutorial as a package and keep it isolated.



Classes
-------

.. autoapisummary::

   emmi_inference.models.transolver.Transolver


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

.. py:class:: Transolver(dim = 256, num_heads = 8, depth = 8, num_slices = 64, mlp_expansion_factor = 2, use_placeholder = True, attn_ctor = TransolverAttention, **kwargs)

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


   Implementation of the Transolver model.
   Reference code: https://github.com/thuml/Transolver/
   Paper: https://arxiv.org/abs/2402.02366

   :param BaseModel: Base model class that contains the utilities for all models we use in this tutorial.

   :param dim: hidden dimensionality of the network. Defaults to 256.
   :param num_heads: number of heads of the multi-head self-attention. Defaults to 8.
   :param depth: number of Transolver blocks. Defaults to 8.
   :param num_slices: number of latent tokens/physics tokens used for down projecting the input points. Defaults to 64.
   :param mlp_expansion_factor: up projection ratio/factor of the MLP in the Transolver blocks. Defaults to 2.
   :param use_placeholder: left over from the original code, no idea what the function is. Defaults to True.
   :param attn_ctor: constructor for the attention module. Defaults to TransolverAttention.


   .. py:attribute:: dim
      :value: 256



   .. py:attribute:: num_heads
      :value: 8



   .. py:attribute:: depth
      :value: 8



   .. py:attribute:: mlp_expansion_factor
      :value: 2



   .. py:attribute:: embed


   .. py:attribute:: blocks


   .. py:method:: forward(input_position, surface_mask_input, physics_features = None)

      forward pass of the Transolver model.

      :param input_position: conncatenated surface and volume points coodindates.
      :param surface_mask_input: mask for the input points, indicating which points are surface points.
      :param physics_features: physics features for the input points. Defaults to None.

      :returns: dictionary with the output tensors, containing the surface pressure and volume velocity.
      :rtype: dict[str, torch.Tensor]



   .. py:method:: kwarg_forward(input_position, surface_mask_input, physics_features = None)

      forward pass of the Transolver model.

      :param input_position: conncatenated surface and volume points coodindates.
      :param surface_mask_input: mask for the input points, indicating which points are surface points.
      :param physics_features: physics features for the input points. Defaults to None.

      :returns: dictionary with the output tensors, containing the surface pressure and volume velocity.
      :rtype: dict[str, torch.Tensor]