106 lines
3.2 KiB
Python
106 lines
3.2 KiB
Python
from __future__ import annotations
|
|
|
|
import unittest
|
|
|
|
from airfrans_frontier.runtime import remove_pythonpath_entries
|
|
|
|
remove_pythonpath_entries()
|
|
|
|
import torch
|
|
|
|
from airfrans_frontier.models import (
|
|
DeepONetBranchTrunk,
|
|
LocalPointTransformer,
|
|
NeRFCFDMultiRes,
|
|
PointContextPerceiver,
|
|
PointwiseMLP,
|
|
RasterFNOUNet,
|
|
SirenConditionedINR,
|
|
)
|
|
|
|
|
|
class PointwiseMLPTests(unittest.TestCase):
|
|
def test_forward_pass_returns_batch_by_target_dim(self) -> None:
|
|
model = PointwiseMLP(input_dim=5, output_dim=4, hidden_width=16, depth=2, activation="gelu")
|
|
batch = torch.randn(7, 5)
|
|
|
|
output = model(batch)
|
|
|
|
self.assertEqual(tuple(output.shape), (7, 4))
|
|
|
|
def test_frontier_models_return_batch_by_target_dim(self) -> None:
|
|
feature_names = ("x", "y", "sdf", "u_inf", "log_re", "aoa_deg")
|
|
batch = torch.randn(8, len(feature_names))
|
|
models = [
|
|
NeRFCFDMultiRes(
|
|
feature_names=feature_names,
|
|
output_dim=4,
|
|
coordinate_features=("x", "y", "sdf"),
|
|
encoding_levels=2,
|
|
hidden_width=16,
|
|
depth=2,
|
|
condition_width=12,
|
|
condition_depth=2,
|
|
activation="gelu",
|
|
),
|
|
DeepONetBranchTrunk(
|
|
feature_names=feature_names,
|
|
output_dim=4,
|
|
coordinate_features=("x", "y", "sdf"),
|
|
fourier_scales=(1.0, 2.0),
|
|
hidden_width=16,
|
|
depth=2,
|
|
condition_width=12,
|
|
condition_depth=2,
|
|
activation="gelu",
|
|
),
|
|
PointContextPerceiver(
|
|
input_dim=len(feature_names),
|
|
output_dim=4,
|
|
hidden_width=16,
|
|
latent_width=12,
|
|
context_points=6,
|
|
attention_depth=2,
|
|
activation="gelu",
|
|
),
|
|
LocalPointTransformer(
|
|
feature_names=feature_names,
|
|
output_dim=4,
|
|
coordinate_features=("x", "y", "sdf"),
|
|
hidden_width=16,
|
|
depth=2,
|
|
neighbors=3,
|
|
activation="gelu",
|
|
),
|
|
RasterFNOUNet(
|
|
feature_names=feature_names,
|
|
output_dim=4,
|
|
coordinate_features=("x", "y", "sdf"),
|
|
grid_resolution=4,
|
|
hidden_width=16,
|
|
depth=2,
|
|
condition_width=12,
|
|
condition_depth=2,
|
|
activation="gelu",
|
|
),
|
|
SirenConditionedINR(
|
|
feature_names=feature_names,
|
|
output_dim=4,
|
|
coordinate_features=("x", "y", "sdf"),
|
|
hidden_width=16,
|
|
depth=2,
|
|
condition_width=12,
|
|
condition_depth=2,
|
|
omega0=10.0,
|
|
activation="gelu",
|
|
),
|
|
]
|
|
|
|
for model in models:
|
|
with self.subTest(model=type(model).__name__):
|
|
output = model(batch)
|
|
self.assertEqual(tuple(output.shape), (8, 4))
|
|
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|