upt
===

.. py:module:: upt


Classes
-------

.. autoapisummary::

   upt.UPT


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

.. py:class:: UPT(model_config)

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


   Implementation of the UPT (Universal Physics Transformer) model.

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

   :param model_config: Configuration for the UPT model.


   .. py:attribute:: encoder


   .. py:attribute:: use_rope


   .. py:attribute:: pos_embed


   .. py:attribute:: approximator_blocks


   .. py:attribute:: decoder


   .. py:attribute:: norm


   .. py:attribute:: prediction_layer


   .. py:attribute:: bias_layers
      :value: False



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

      Apply separate bias layers for surface and volume points.



   .. py:method:: compute_rope_args(surface_position_batch_idx, surface_position, surface_position_supernode_idx, query_position)

      Compute the RoPE frequency arguments for the surface_position and query_position. If we don't uise RoPE, return empty dicts.
      :param surface_position_batch_idx: Batch indices for the surface positions.
      :param surface_position: Surface position coordinates.
      :param surface_position_supernode_idx: Supernode indices for the surface positions.
      :param query_position: Query position coordinates.

      :returns: Dictionary containing the RoPE frequency arguments.
      :rtype: dict[str, torch.Tensor]



   .. py:method:: forward(surface_position_batch_idx, surface_position_supernode_idx, surface_position, query_position, surface_mask_query, input_features = None)

      Forward pass of the UPT model.

      :param input_position: input coordinates of the surface points.
      :param surface_mask_query: surface mask for the query points, indicating which points are surface points.
      :param surface_position_batch_idx: batch indices for the surface positions, since the surface positions are a sparse tensor for the supernode pooling.
      :param surface_position_supernode_idx: supernode indices for the surface positions.
      :param geometry_position: geometry position information.
      :param query_position: input coordinates of the query points.
      :param surface_features: surface 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]