Geotransolver Model

examples/cfd/external_aerodynamics/transolver/src/compute_normalizations.py → examples/cfd/external_aerodynamics/transformer_models/src/compute_normalizations.py

@@ -89,7 +89,6 @@ def compute_mean_std_min_max(
 
         # Update running mean and M2 (Welford's algorithm)
         delta = batch_mean - mean
-        N += batch_n
         mean = mean + delta * (batch_n / N)
         M2 = M2 + batch_M2 + delta**2 * (batch_n * N) / N
         time_end = time.time()

examples/cfd/external_aerodynamics/transformer_models/src/conf/data/core.yaml

@@ -0,0 +1,60 @@
+# SPDX-FileCopyrightText: Copyright (c) 2023 - 2025 NVIDIA CORPORATION & AFFILIATES.
+# SPDX-FileCopyrightText: All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# You may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Paths to your data:
+train:
+  data_path:  /lustre/fsw/portfolios/coreai/projects/coreai_modulus_cae/datasets/drivaer_aws/domino/train/
+val:
+  data_path: /lustre/fsw/portfolios/coreai/projects/coreai_modulus_cae/datasets/drivaer_aws/domino/val/
+
+# You can set a normalization factor directory:
+normalization_dir: "src/"
+
+# How many events in advance should we be preloading?
+preload_depth: 1
+
+# Pin memory for GPU transfers?
+pin_memory: true
+
+# Sampling resolution of the point clouds:
+resolution: 200_000
+
+# Surface / Volume / (combined, if supported)
+mode: ???
+
+# For building embeddings: include normal directions for each point?
+include_normals: true
+# Include SDF?  (It's 0 for surface data...)
+include_sdf: true
+# Apply translation invariance via center-of-mass subtraction?
+translational_invariance: true
+# Rescale x/y/z inputs to the model for scale invariance?
+scale_invariance: true
+reference_scale: [12.0, 4.5, 3.25]
+
+# Which parts of the data files to read?  No need to read everything, all the time.
+data_keys: ???
+
+# Load and return the STL geometry info in the dataloader?
+include_geometry: false
+
+# Broadcast global features to the same resolution as points?
+broadcast_global_features: true
+
+# Return the mesh areas and normals?  You don't usually want this for training.
+# We switch it on automatically for inference.
+return_mesh_features: false
+

examples/cfd/external_aerodynamics/transformer_models/src/conf/data/surface.yaml

@@ -0,0 +1,33 @@
+# SPDX-FileCopyrightText: Copyright (c) 2023 - 2025 NVIDIA CORPORATION & AFFILIATES.
+# SPDX-FileCopyrightText: All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# You may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+defaults:
+  -  core
+
+# Overrides for surface data:
+mode: surface
+
+# Surface-speficic needs:
+data_keys:
+  - "surface_fields"
+  - "surface_mesh_centers"
+  - "surface_normals"
+  - "surface_areas"
+  - "stl_faces"
+  - "stl_centers"
+  - "stl_coordinates"
+  - "air_density"
+  - "stream_velocity"

examples/cfd/external_aerodynamics/transformer_models/src/conf/data/volume.yaml

@@ -0,0 +1,30 @@
+# SPDX-FileCopyrightText: Copyright (c) 2023 - 2025 NVIDIA CORPORATION & AFFILIATES.
+# SPDX-FileCopyrightText: All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# You may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+defaults:
+  -  core
+
+# Overrides for volume data:
+mode: volume
+
+# volume-specific needs:
+data_keys:
+  - "volume_fields"
+  - "volume_mesh_centers"
+  - "stl_faces"
+  - "stl_centers"
+  - "stl_coordinates"
+

examples/cfd/external_aerodynamics/transformer_models/src/conf/geotransolver_surface.yaml

@@ -0,0 +1,48 @@
+# SPDX-FileCopyrightText: Copyright (c) 2023 - 2025 NVIDIA CORPORATION & AFFILIATES.
+# SPDX-FileCopyrightText: All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+defaults:
+  - training: base
+  - model: geotransolver
+  - data: surface
+
+output_dir: "runs"
+checkpoint_dir: null  # Optional: set custom checkpoint path, defaults to output_dir
+run_id: "geotransolver/surface/bq"
+
+# Performance considerations:
+precision: float32 # float32, float16, bfloat16, or float8
+compile: true
+profile: false
+
+model:
+  functional_dim: 6
+  include_local_features: true # use local features
+  radii: [0.01, 0.05, 0.25, 1.0, 2.5, 5.0] # radius for local features
+  neighbors_in_radius: [4, 8, 16, 64, 128, 256]  # neighbors in radius for local features
+  n_hidden_local: 32 # hidden dimension for local features
+
+data:
+  include_sdf: false
+  include_geometry: true
+  geometry_sampling: 300_000
+  broadcast_global_features: true
+
+
+# Logging configuration
+logging:
+  level: INFO
+  format: '%(asctime)s - %(name)s - %(levelname)s - %(message)s'

examples/cfd/external_aerodynamics/transformer_models/src/conf/geotransolver_volume.yaml

@@ -0,0 +1,49 @@
+# SPDX-FileCopyrightText: Copyright (c) 2023 - 2025 NVIDIA CORPORATION & AFFILIATES.
+# SPDX-FileCopyrightText: All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+defaults:
+  - training: base
+  - model: geotransolver
+  - data: volume
+
+output_dir: "runs"
+checkpoint_dir: null  # Optional: set custom checkpoint path, defaults to output_dir
+run_id: "geotransolver/volume/bq"
+
+# Performance considerations:
+precision: float32 # float32, float16, bfloat16, or float8
+compile: true 
+profile: false
+
+data:
+  include_geometry: true
+  geometry_sampling: 300_000
+  broadcast_global_features: false
+  volume_sample_from_disk: true
+
+
+model:
+  functional_dim: 7
+  out_dim: 5
+  include_local_features: true # use local features
+  radii: [0.01, 0.05, 0.25, 1.0, 2.5, 5.0] # radius for local features
+  neighbors_in_radius: [4, 8, 16, 64, 128, 256]  # neighbors in radius for local features
+  n_hidden_local: 32 # hidden dimension for local features
+
+# Logging configuration
+logging:
+  level: INFO
+  format: '%(asctime)s - %(name)s - %(levelname)s - %(message)s'

examples/cfd/external_aerodynamics/transformer_models/src/conf/model/geotransolver.yaml

@@ -0,0 +1,35 @@
+# SPDX-FileCopyrightText: Copyright (c) 2023 - 2025 NVIDIA CORPORATION & AFFILIATES.
+# SPDX-FileCopyrightText: All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+_target_: physicsnemo.experimental.models.geotransolver.GeoTransolver
+functional_dim: 6
+global_dim: 2
+geometry_dim: 3
+out_dim: 4
+n_layers: 20
+n_hidden: 256
+dropout: 0.0
+n_head: 8
+act: "gelu"
+mlp_ratio: 2
+slice_num: 128
+use_te: false
+plus: false
+include_local_features: true # use local features
+radii: [0.05, 0.25, 1.0, 2.5] # radius for local features
+neighbors_in_radius: [8, 32, 64, 128] # neighbors in radius for local features
+n_hidden_local: 32 # hidden dimension for local features
+

examples/cfd/external_aerodynamics/transformer_models/src/conf/model/transolver.yaml

@@ -0,0 +1,34 @@
+# SPDX-FileCopyrightText: Copyright (c) 2023 - 2025 NVIDIA CORPORATION & AFFILIATES.
+# SPDX-FileCopyrightText: All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+_target_: physicsnemo.models.transolver.Transolver
+functional_dim: 2
+out_dim: 4
+embedding_dim: 6
+n_layers: 8
+n_hidden: 256
+dropout: 0.0
+n_head: 8
+act: "gelu"
+mlp_ratio: 2
+slice_num: 512
+unified_pos: false
+ref: 8
+structured_shape: null
+use_te: false
+time_input: false
+plus: false
+

examples/cfd/external_aerodynamics/transformer_models/src/conf/training/base.yaml

@@ -0,0 +1,32 @@
+# SPDX-FileCopyrightText: Copyright (c) 2023 - 2025 NVIDIA CORPORATION & AFFILIATES.
+# SPDX-FileCopyrightText: All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+num_epochs: 501
+save_interval: 25
+
+scheduler:
+  name: "StepLR"
+  params:
+    step_size: 100
+    gamma: 0.5
+
+optimizer:
+  _target_: torch.optim.AdamW
+  lr: 1.0e-3
+  weight_decay: 1.0e-4
+  betas: [0.9, 0.999]
+  eps: 1.0e-8
+

examples/cfd/external_aerodynamics/transformer_models/src/conf/transolver_surface.yaml

@@ -0,0 +1,38 @@
+# SPDX-FileCopyrightText: Copyright (c) 2023 - 2025 NVIDIA CORPORATION & AFFILIATES.
+# SPDX-FileCopyrightText: All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+defaults:
+  - training: base
+  - model: transolver
+  - data: surface
+
+output_dir: "runs"
+checkpoint_dir: null  # Optional: set custom checkpoint path, defaults to output_dir
+run_id: "surface/float32"
+
+# Performance considerations:
+precision: float32 # float32, float16, bfloat16, or float8
+compile: true
+profile: false
+
+data:
+  include_sdf: false
+
+# Logging configuration
+logging:
+  level: INFO
+  format: '%(asctime)s - %(name)s - %(levelname)s - %(message)s'
+