Graph and Geometry Layers#

This layer is a compilable, ball-query operation.

By default, it will project a grid of points to a 1D set of points.

It does not support batch size > 1.

class physicsnemo.nn.module.ball_query.BQWarp(radius: float = 0.25, neighbors_in_radius: int | None = 10)[source]#

Bases: Module

Warp-based ball-query layer for finding neighboring points within a specified radius.

This layer uses an accelerated ball query implementation to efficiently find points within a specified radius of query points.

Only supports batch size 1.

forward(
x: Tensor,
p_grid: Tensor,
reverse_mapping: bool = True,
) → tuple[Tensor, Tensor][source]#

Performs ball query operation to find neighboring points and their features.

This method uses the Warp-accelerated ball query implementation to find points within a specified radius. It can operate in two modes: - Forward mapping: Find points from x that are near p_grid points (reverse_mapping=False) - Reverse mapping: Find points from p_grid that are near x points (reverse_mapping=True)

Parameters:

  • x – Tensor of shape (batch_size, num_points, 3+features) containing point coordinates and their features

  • p_grid – Tensor of shape (batch_size, grid_x, grid_y, grid_z, 3) containing grid point coordinates

  • reverse_mapping – Boolean flag to control the direction of the mapping: - True: Find p_grid points near x points - False: Find x points near p_grid points

Returns:

  • mapping: Tensor containing indices of neighboring points

  • outputs: Tensor containing coordinates of the neighboring points

Return type:

tuple containing

class physicsnemo.nn.module.dgm_layers.DGMLayer(
in_features_1: int,
in_features_2: int,
out_features: int,
activation_fn: Module | Callable[[Tensor], Tensor] | None = None,
weight_norm: bool = False,
activation_par: Parameter | None = None,
)[source]#

Bases: Module

Deep Galerkin Model layer.

Parameters:

  • in_features_1 (int) – Number of input features for first input.

  • in_features_2 (int) – Number of input features for second input.

  • out_features (int) – Number of output features.

  • activation_fn (Union[nn.Module, Callable[[Tensor], Tensor]], optional) – Activation function, by default Activation.IDENTITY

  • weight_norm (bool, optional) – Apply weight normalization, by default False

  • activation_par (Optional[nn.Parameter], optional) – Activation parameter, by default None

Notes

Reference: DGM: A deep learning algorithm for solving partial differential equations, https://arxiv.org/pdf/1708.07469.pdf

forward(
input_1: Tensor,
input_2: Tensor,
) → Tensor[source]#

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.