Skip to content

adding ball query based local features #6

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 1 commit into
base: transolver-volume
Choose a base branch
from
Open

adding ball query based local features #6

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 4 additions & 0 deletions examples/cfd/external_aerodynamics/transolver/requirements.txt
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,3 +7,7 @@ einops
transformer_engine[pytorch]
tensorstore
zarr>=3.0
s3fs
treelib
warp-lang
timm
2 changes: 1 addition & 1 deletion examples/cfd/external_aerodynamics/transolver/src/compute_normalizations.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -103,7 +103,7 @@ def compute_mean_std_min_max(
return mean, std, min_val, max_val


@hydra.main(version_base="1.3", config_path="conf", config_name="train_surface")
@hydra.main(version_base="1.3", config_path="conf", config_name="train_surfaceX")
def main(cfg: DictConfig) -> None:
"""
Script entry point for computing normalization statistics for a specified field
Expand Down
15 changes: 9 additions & 6 deletions examples/cfd/external_aerodynamics/transolver/src/conf/train_surfaceX.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,12 +18,12 @@

output_dir: "runs"
checkpoint_dir: null # Optional: set custom checkpoint path, defaults to output_dir
run_id: "surfaceX/bfloat16"
run_id: "surfaceX/runs1"
profile: false

# Training configuration
training:
precision: bfloat16 # float32, float16, bfloat16, or float8
precision: float32 # float32, float16, bfloat16, or float8
num_epochs: 501 # Add one to save at 250
save_interval: 25 # Save checkpoint every N epochs
compile: false
Expand All @@ -44,7 +44,10 @@ model:
slice_num: 512 # Number of slices in physics attention
use_te: false # Use transformer engine
plus: false

include_local_features: true # use local features
radii: [0.05, 0.25] # radius for local features
neighbors_in_radius: [8, 32] # neighbors in radius for local features
n_hidden_local: 512 # hidden dimension for local features

# StepLR scheduler: Decays the learning rate by gamma every step_size epochs
scheduler:
Expand All @@ -65,14 +68,14 @@ optimizer:
# Data configuration
data:
train:
data_path: /lustre/fsw/portfolios/coreai/projects/coreai_modulus_cae/drivaer_aws/domino/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/drivaer_aws/domino/val/
data_path: /lustre/fsw/portfolios/coreai/projects/coreai_modulus_cae/datasets/drivaer_aws/domino/val/
max_workers: 8
normalization_dir: "src/" # Directory for normalization files
preload_depth: 1
pin_memory: true
resolution: 300_000
resolution: 80_000
mode: surface
# Preprocessing switches:
# (Changing thes will change the embedding dim)
Expand Down
16 changes: 10 additions & 6 deletions examples/cfd/external_aerodynamics/transolver/src/conf/train_volumeX.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,14 +18,14 @@

output_dir: "runs"
checkpoint_dir: null # Optional: set custom checkpoint path, defaults to output_dir
run_id: "volumeX/fake-name"
run_id: "volumeX/runs1"
profile: false

# Training configuration
training:
precision: bfloat16 # float32, float16, bfloat16, or float8
num_epochs: 500 # Add one to save at 250
save_interval: 10 # Save checkpoint every N epochs
precision: float32 # float32, float16, bfloat16, or float8
num_epochs: 501 # Add one to save at 250
save_interval: 25 # Save checkpoint every N epochs
compile: false

# Model configuration
Expand All @@ -44,6 +44,10 @@ model:
slice_num: 256 # Number of slices in physics attention
use_te: false # Use transformer engine
plus: false
include_local_features: true # use local features
radii: [0.05, 0.25] # radius for local features
neighbors_in_radius: [8, 32] # neighbors in radius for local features
n_hidden_local: 512 # hidden dimension for local features

# scheduler:
# name: "OneCycleLR"
Expand All @@ -70,9 +74,9 @@ optimizer:
# Data configuration
data:
train:
data_path: /lustre/fsw/portfolios/coreai/projects/coreai_modulus_cae/drivaer_aws/domino/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/drivaer_aws/domino/val/
data_path: /lustre/fsw/portfolios/coreai/projects/coreai_modulus_cae/datasets/drivaer_aws/domino/val/
max_workers: 8
normalization_dir: "src/" # Directory for normalization files
preload_depth: 1
Expand Down
17 changes: 10 additions & 7 deletions examples/cfd/external_aerodynamics/transolver/src/inference_on_zarr.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -40,7 +40,7 @@
create_transolver_dataset,
TransolverDataPipe,
)
from train import forward_pass
from train import forward_passX
from tabulate import tabulate

# import transformer_engine.pytorch as te
Expand Down Expand Up @@ -162,7 +162,7 @@ def batched_inference_loop(
local_batch["geometry"] = batch["geometry"]

# Run the forward inference pass:
local_loss, local_metrics, local_preds_targets = forward_pass(
local_loss, local_metrics, local_preds_targets = forward_passX(
local_batch,
model,
precision,
Expand Down Expand Up @@ -309,7 +309,7 @@ def inference(cfg: DictConfig) -> None:
pred_pressure = pred_pressure.reshape(-1)
pred_drag_coeff, _, _ = compute_force_coefficients(
batch["surface_normals"][0],
batch["surface_areas"],
batch["surface_areas"][0],
coeff,
pred_pressure,
pred_shear,
Expand All @@ -318,7 +318,7 @@ def inference(cfg: DictConfig) -> None:

pred_lift_coeff, _, _ = compute_force_coefficients(
batch["surface_normals"][0],
batch["surface_areas"],
batch["surface_areas"][0],
coeff,
pred_pressure,
pred_shear,
Expand All @@ -333,7 +333,7 @@ def inference(cfg: DictConfig) -> None:
true_pressure = true_pressure.reshape(-1)
true_drag_coeff, _, _ = compute_force_coefficients(
batch["surface_normals"][0],
batch["surface_areas"],
batch["surface_areas"][0],
coeff,
true_pressure,
true_shear,
Expand All @@ -342,7 +342,7 @@ def inference(cfg: DictConfig) -> None:

true_lift_coeff, _, _ = compute_force_coefficients(
batch["surface_normals"][0],
batch["surface_areas"],
batch["surface_areas"][0],
coeff,
true_pressure,
true_shear,
Expand All @@ -352,6 +352,9 @@ def inference(cfg: DictConfig) -> None:
pred_lift_coeff = pred_lift_coeff.item()
pred_drag_coeff = pred_drag_coeff.item()

true_lift_coeff = true_lift_coeff.item()
true_drag_coeff = true_drag_coeff.item()

# Extract metric values and convert tensors to floats
l2_pressure = (
metrics["l2_pressure_surf"].item()
Expand Down Expand Up @@ -431,7 +434,7 @@ def inference(cfg: DictConfig) -> None:
)


@hydra.main(version_base=None, config_path="conf", config_name="train_surface")
@hydra.main(version_base=None, config_path="conf", config_name="train_surfaceX")
def launch(cfg: DictConfig) -> None:
"""
Launch inference with Hydra configuration.
Expand Down
Binary file modified examples/cfd/external_aerodynamics/transolver/src/surface_fields_normalization.npz
View file
Open in desktop
Binary file not shown.
Loading