emmi_inference.models.base_model
================================

.. py:module:: emmi_inference.models.base_model

.. 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.base_model.BaseModel


Functions
---------

.. autoapisummary::

   emmi_inference.models.base_model.validate_required_batch_input


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

.. py:function:: validate_required_batch_input(func)

   Decorator to validate batch input for methods that require it.


.. py:class:: BaseModel(input_dim, output_dim, required_batch_modes, dim = None, output_projection = False, use_physics_features = False, physics_dim = None, bias_layers = True, position_projection = 'sincos', **kwargs)

   Bases: :py:obj:`ksuit.models.Model`


   Base class for all model we use in this tutorial.

   :param Model: Base class for single models in ksuit.

   :param input_dim: dimension of the input features.
   :param output_dim: dimension of the output prediction.
   :param required_batch_keys: defines the required batch keys for the model. This is used to validate the input batch.
   :param dim: dimension of the latents of the model.
   :param output_projection: Boolean to indicate to use the output projection. Defaults to False.
   :param use_physics_features: Boolean to indicate to use additional physics features (next to the input coordiates). Defaults to False.
   :param physics_dim: dimension of the physics features. Defaults to None.
   :param bias_layers: Boolean to indicate to use bias layers. Defaults to True.
   :param position_projection: String to indicate the type of position projection to use. Can be "sincos" or "linear". Defaults to "sincos".


   .. py:attribute:: input_dim


   .. py:attribute:: output_dim


   .. py:attribute:: physics_dim
      :value: None



   .. py:attribute:: use_physics_features
      :value: False



   .. py:attribute:: required_batch_modes


   .. py:attribute:: position_projection
      :value: 'sincos'



   .. py:method:: output_projection(x)

      Most model implementations will have an output projection layer that maps the last latent vector into the output physics space.
      We have a unified projection layer that can be used in all models.

      :param x: tensor of shape (batch_size, num_points, dim) containing the features for each point.

      :returns: tensor of shape (batch_size, num_points, output_dim) containing the projected features into (normalized) physics space.



   .. py:method:: surface_and_volume_bias(x, surface_mask)

      For some of the models, the surface and volume are concatenated into a single input tesnor (e.g., Pointnet, Transolver, Transformer).
      For AB-UPT, we shared weight for the physics blocks. Hence, we need to indicate which points are surface and which are volume points.
      We do this by applying a bias (i.e., an MLP) to the surface and volume points separately. The surface mask indicates which points are surface points.
      This function only works for tensors where surface and volume points are concatenated along the second dimension (not for AB-UPT).
      Howerver, for other models, self.surface_bias and self.volume_bias can be called directly in the child class.

      :param x: tensor of shape (batch_size, num_points, input_dim) containing the features for each point.
      :param surface_mask: Boolean tensor of shape (batch_size, num_points) indicating which points are surface points.

      :returns: biased tensor x of shape (batch_size, num_points, input_dim) where the surface points have been processed by the surface bias and the volume points by the volume bias.
      :rtype: torch.Tensor



   .. py:method:: gather_outputs(x, surface_mask)

      The output gathering function is used to extract the relevant outputs from the model's output tensor.
      It assumes that the output tensor has a specific structure, where the first dimension corresponds to the batch size,
      the second dimension corresponds to the surface/volume points, and the third dimension corresponds to the output features

      The surface pressure is expected to be at index 0, the volume velocity at indices 1:4,
      and if the output dimension is 11, the surface wall shear stress is at indices 4:7,
      the volume total pressure coefficient at index 7, and the volume vorticity at indices 8:11. These last three are only available
      for AhmedML and DriverML datasets.

      :param x: output tensor from the model, shape (batch_size, num_points, output_dim)
      :param surface_mask: Indicator boolean tensor for surface points, shape (batch_size, num_points). All surface points should be True, and all volume points should be False.

      :returns:

                A dictionary containing the gathered outputs:
                    - "surface_pressure": Tensor of shape (batch_size, num_surface_points, 1)
                    - "volume_velocity": Tensor of shape (batch_size, num_volume_points, 3)
                    - "surface_wallshearstress": Tensor of shape (batch_size, num_surface_points, 3) if output_dim is 11
                    - "volume_totalpcoeff": Tensor of shape (batch_size, num_volume_points, 1) if output_dim is 11
                    - "volume_vorticity": Tensor of shape (batch_size, num_volume_points, 3) if output_dim is 11
      :rtype: dict[str, torch.Tensor]