Update geotransolver for 2d and 3d use cases

examples/cfd/darcy_transolver/README.md

@@ -89,6 +89,15 @@ To train Transolver on your own physics problem, modify the `dataloader` in
 `train_transolver_darcy.py` to use your own pre-computed data or on-the-fly
 solver.
 
+#### Other models
+
+You can also use GeoTransolver, Flare, or GeoTransolver with flare attention
+for the darcy example with the "fixed" usage.  Change the `model` parameter
+to update it!  Note that the inputs to geotransolver here uses the positions
+and the `x` diffusivity to both the model trunk and encoder paths.  For these
+problems, GeoTransolver is not yet extensively tested, and your results
+may vary.
+
 ## Additional Information
 
 More components are added for convenience. `Validators` calculate the loss

examples/cfd/darcy_transolver/config_fix.yaml

@@ -20,33 +20,31 @@
 # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 # SOFTWARE.
 
-output_dir: ./output/darcy_transolver_fix
-run_id: bf16_dev_r85_b8_s64
+defaults:
+  - model: transolver   # transolver | flare | geotransolver | geotransolver_fa
+  - _self_
+
+hydra:
+  output_subdir: null
+  run:
+    dir: .
+
+output_dir: ./output/
+run_id: ${hydra:runtime.choices.model}-muon_${precision}_r${resolution}_b${data.batch_size}_s${model.slice_num}
+
+resolution: 85 #421, 211, 141, 106, 85 all viable
 
 data:
-  train_path: /user_data/datasets/darcy_fix/example_data/piececonst_r421_N1024_smooth1.npz
-  test_path: /user_data/datasets/darcy_fix/example_data/piececonst_r421_N1024_smooth2.npz
-  resolution: 85 #421, 211, 141, 106, 85 all viable
-  batch_size: 8 # This is the GLOBAL batch size
+  train_path: //lustre/fsw/portfolios/coreai/users/coreya/datasets/darcy_fix/example_data/piececonst_r421_N1024_smooth1.npz
+  test_path: //lustre/fsw/portfolios/coreai/users/coreya/datasets/darcy_fix/example_data/piececonst_r421_N1024_smooth2.npz
+  resolution: ${resolution}
+  batch_size: 64 # This is the GLOBAL batch size
 
-model:
-  functional_dim: 1
-  out_dim: 1
-  embedding_dim: 2
-  n_layers: 4
-  n_hidden: 128
-  dropout: 0.0
-  n_head: 4
-  act: gelu
-  mlp_ratio: 4
-  unified_pos: False
-  ref: 8
-  slice_num: 64
-  use_te: False
-  time_input: False
-
 precision: bf16
 
+optimizer:
+  type: muon   # adamw | muon (muon requires PyTorch >= 2.9)
+
 normaliser:
   permeability:
     mean: 1.25
@@ -69,3 +67,4 @@ training:
 validation:
   sample_size: 200
   validation_pseudo_epochs: 1
+  save_plots: false

examples/cfd/darcy_transolver/model/flare.yaml

@@ -0,0 +1,39 @@
+# SPDX-FileCopyrightText: Copyright (c) 2023 - 2026 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.
+
+# FLARE inherits from Transolver and replaces physics attention blocks with
+# FLARE (Fast Low-rank Attention Routing Engine) blocks.
+# use_te and plus are forced off internally by FLARE.
+
+_target_: physicsnemo.experimental.models.flare.FLARE
+_convert_: all
+
+functional_dim: 1
+out_dim: 1
+embedding_dim: 2
+n_layers: 4
+n_hidden: 128
+dropout: 0.0
+n_head: 4
+act: gelu
+mlp_ratio: 4
+slice_num: 64
+unified_pos: False
+ref: 8
+structured_shape:
+  - ${resolution}
+  - ${resolution}
+time_input: False

examples/cfd/darcy_transolver/model/geotransolver.yaml

@@ -0,0 +1,40 @@
+# SPDX-FileCopyrightText: Copyright (c) 2023 - 2026 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.
+
+# GeoTransolver with GALE attention (default).
+# Uses geometry-aware cross-attention instead of Transolver's structured grid.
+
+_target_: physicsnemo.experimental.models.geotransolver.GeoTransolver
+_convert_: all
+
+functional_dim: 3
+out_dim: 1
+geometry_dim: 3
+n_layers: 4
+n_hidden: 128
+dropout: 0.0
+n_head: 4
+act: gelu
+mlp_ratio: 4
+slice_num: 64
+use_te: False
+time_input: False
+attention_type: GALE
+include_local_features: False
+state_mixing_mode: weighted
+structured_shape:
+  - ${resolution}
+  - ${resolution}

examples/cfd/darcy_transolver/model/geotransolver_fa.yaml

@@ -0,0 +1,41 @@
+# SPDX-FileCopyrightText: Copyright (c) 2023 - 2026 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.
+
+# GeoTransolver with GALE_FA (FLARE attention backend).
+# Combines FLARE-style global-query self-attention with GALE geometry/context
+# cross-attention.  use_te must be False for GALE_FA.
+
+_target_: physicsnemo.experimental.models.geotransolver.GeoTransolver
+_convert_: all
+
+functional_dim: 3
+out_dim: 1
+geometry_dim: 3
+n_layers: 4
+n_hidden: 128
+dropout: 0.0
+n_head: 4
+act: gelu
+mlp_ratio: 4
+slice_num: 64
+use_te: False
+time_input: False
+attention_type: GALE_FA
+include_local_features: False
+state_mixing_mode: concat_project
+structured_shape:
+  - ${resolution}
+  - ${resolution}

examples/cfd/darcy_transolver/model/transolver.yaml

@@ -0,0 +1,36 @@
+# SPDX-FileCopyrightText: Copyright (c) 2023 - 2026 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
+_convert_: all
+
+functional_dim: 1
+out_dim: 1
+embedding_dim: 2
+n_layers: 4
+n_hidden: 128
+dropout: 0.0
+n_head: 4
+act: gelu
+mlp_ratio: 4
+slice_num: 64
+unified_pos: False
+ref: 8
+structured_shape:
+  - ${resolution}
+  - ${resolution}
+use_te: False
+time_input: False