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.

Supports arbitrary batch sizes.

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.

Supports arbitrary batch sizes.

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) containing point coordinates

  • p_grid – Tensor of shape (batch_size, …, 3) containing grid point coordinates. Can be 3D, 4D, or 5D.

  • 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. Shape (B, Q, max_points) when neighbors_in_radius is set.

  • outputs: Tensor containing coordinates of the neighboring points. Shape (B, Q, max_points, 3) when neighbors_in_radius is set.

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.