🐛[BUG]: GeoTransolver raises ValueError: The last dimension of p_grid and x must be 3

[chychen]

Version

2.1.0a0

On which installation method(s) does this occur?

Docker

Describe the issue

I am encountering a ValueError when initializing GeoTransolver and attempting a forward pass. It seems that the model is passing a tensor with a dimension size other than 3 (specifically, likely related to the geometry input) into the ball_query module, which strictly enforces a last dimension of 3 for spatial coordinates.

Expected Behavior

The model should process the geometry features independently of the spatial coordinates, or the ball_query logic should be resilient to the input dimensions provided.

Actual Behavior

The code crashes at physicsnemo/nn/module/ball_query.py, line 89:
ValueError: The last dimension of p_grid and x must be 3

Minimum reproducible example

import torch
from physicsnemo.experimental.models.geotransolver import GeoTransolver
model = GeoTransolver(
    functional_dim=64,
    out_dim=3,
    geometry_dim=5,
    global_dim=16,
    n_hidden=256,
    n_layers=4,
    use_te=False,
    include_local_features=True
)
local_emb = torch.randn(1, 1000, 64)
coords = torch.randn(1, 1000, 3)  # (batch, nodes, spatial_dim)
geometry = torch.randn(1, 1000, 5)  # (batch, nodes, spatial_dim)
global_emb = torch.randn(1, 1, 16)  # (batch, 1, global_features)
output = model(local_emb, local_positions=coords, global_embedding=global_emb, geometry=geometry)

Relevant log output

Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1776, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1787, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/media/jaych/Data/physicsnemo/physicsnemo/experimental/models/geotransolver/geotransolver.py", line 532, in forward
    embedding_states, local_embedding_bq = self.context_builder.build_context(
                                           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/media/jaych/Data/physicsnemo/physicsnemo/experimental/models/geotransolver/context_projector.py", line 801, in build_context
    context_feats = self.local_extractors[i].extract_context_features(
                    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/media/jaych/Data/physicsnemo/physicsnemo/experimental/models/geotransolver/context_projector.py", line 552, in extract_context_features
    tokenizer(processor(spatial_coords, geometry))
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1776, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1787, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/media/jaych/Data/physicsnemo/physicsnemo/experimental/models/geotransolver/context_projector.py", line 423, in forward
    _, neighbors = self.bq_warp(query_points, key_features)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1776, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/usr/local/lib/python3.12/dist-packages/torch/nn/modules/module.py", line 1787, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/media/jaych/Data/physicsnemo/physicsnemo/nn/module/ball_query.py", line 89, in forward
    raise ValueError("The last dimension of p_grid and x must be 3")
ValueError: The last dimension of p_grid and x must be 3

Environment details

docker run -it --rm nvcr.io/nvidia/physicsnemo/physicsnemo:26.03 bash

[coreyjadams]

Hi @chychen , Thanks for opening an issue.

In GeoTransolver, the ContextProjector is, admittedly, a little complex. There are three stages to it that are relevant to you here, which are represented as the three columns in the purple box in Figure 1 in the arxiv paper here.

Working from left to right:

• The "Multi-scale Feature Extractor" takes your geometry (shape [1, 1000, 5]) and your input positions (shape [1, 1000, 3]) and produces two outputs. (This is where your crash is coming from, I'll elaborate below)
  • One output takes, for every input point, a query of nearby geometry points. These points are projected to a latent feature via an MLP, and returned to the input of the model and concatenated with the rest of the input embeddings.
  • The other output takes, for every geometry point, a query of the nearby input points. These points are again projected to a latent space, but this time the latent space is tokenize to the same low-rank attention space as the main transolver backbone for a cross-attention layer, later.
• The Global Tokenizer takes the geometry point cloud ( shape [1, 1000, 5]) and tokenizes it onto the low-rank latent attention space.
• The global tokenize does the same operation for the global embedding (shape [1, 1, 16]), projecting it on to the low-rank latent attention space.

At the end of the context projector, the various contexts are concatenated and returned.

The fundamental issue you're hitting is that it's not well defined to compute the distance between a point from a 3D position and a 5D geometry. It will probably work if you set include_local_features=False above - the MultiScale Feature Extractor will be skipped, the 5D geometry information should be projected just fine, though if there is a way to represent the geometry in 3D you may find the model is more expressive with the multi scale features enabled.

[coreyjadams]

I just tested it locally in the development branch:

>>> import torch
... from physicsnemo.experimental.models.geotransolver import GeoTransolver
... model = GeoTransolver(
...     functional_dim=64,
...     out_dim=3,
...     geometry_dim=5,
...     global_dim=16,
...     n_hidden=256,
...     n_layers=4,
...     use_te=False,
...     include_local_features=False # <<<<<<<<< CHANGED TO FALSE
... )
...
local_emb = torch.randn(1, 1000, 64)
coords = torch.randn(1, 1000, 3)  # (batch, nodes, spatial_dim)
geometry = torch.randn(1, 1000, 5)  # (batch, nodes, spatial_dim)
global_emb = torch.randn(1, 1, 16)  # (batch, 1, global_features)
output = model(local_emb, local_positions=coords, global_embedding=global_emb, geometry=geometry)
<python-input-0>:2: ExperimentalFeatureWarning: You are importing from 'physicsnemo.experimental'. The APIs in this namespace are experimental, under active development, and may change without notice. Expect possible back-compatibility breaking changes and only partial test coverage.


Warp DeprecationWarning: The symbol `warp.context.Device` will soon be removed from the public API. Use `warp.Device` instead.
>>> local_emb = torch.randn(1, 1000, 64)
>>> coords = torch.randn(1, 1000, 3)  # (batch, nodes, spatial_dim)
>>> geometry = torch.randn(1, 1000, 5)  # (batch, nodes, spatial_dim)
>>> global_emb = torch.randn(1, 1, 16)  # (batch, 1, global_features)
>>> output = model(local_emb, local_positions=coords, global_embedding=global_emb, geometry=geometry)
>>>
>>>
>>> print(output.shape)
torch.Size([1, 1000, 3])
>>>

We're still looking at ways to improve the geometry encodings, feel free to elaborate on what's going in to your geometry and we might be able to come up with some ideas to help?

[coreyjadams]

Hi @chychen - does this help your issue?

[chychen]

thanks clear explanation here. I created this issue based on some researchers' feedback. I have shared the link to them, they might add more details here if they found embedded geometry with >3 dims really matters to them.