371 lines
13 KiB
Python
371 lines
13 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Generate reference inputs and outputs for Python vs Rust comparison testing.
|
|
|
|
This script creates:
|
|
1. A deterministic input tensor
|
|
2. FNO model weights in SafeTensors format
|
|
3. Reference output from neuraloperator
|
|
|
|
The Rust test can then load these files and verify numerical equivalence.
|
|
|
|
Usage:
|
|
python generate_reference_outputs.py --output weights/fno/fno2d_darcy/reference
|
|
"""
|
|
|
|
import argparse
|
|
import json
|
|
import os
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
import torch
|
|
from safetensors.torch import save_file, load_file
|
|
|
|
# Check if neuraloperator is available
|
|
NEURALOP_AVAILABLE = False
|
|
try:
|
|
from neuralop.models import FNO
|
|
NEURALOP_AVAILABLE = True
|
|
except ImportError:
|
|
pass
|
|
|
|
|
|
def create_grid_positional_encoding(height: int, width: int) -> torch.Tensor:
|
|
"""Create 2D grid positional encoding matching neuraloperator."""
|
|
# Create coordinate grids in [-1, 1]
|
|
y = torch.linspace(-1, 1, height)
|
|
x = torch.linspace(-1, 1, width)
|
|
|
|
# Create meshgrid (neuraloperator uses 'ij' indexing)
|
|
grid_y, grid_x = torch.meshgrid(y, x, indexing='ij')
|
|
|
|
# Stack to [2, H, W]
|
|
grid = torch.stack([grid_x, grid_y], dim=0)
|
|
|
|
return grid
|
|
|
|
|
|
def create_deterministic_input(batch_size: int, in_channels: int,
|
|
height: int, width: int, seed: int = 42) -> torch.Tensor:
|
|
"""Create a deterministic input tensor for reproducible testing."""
|
|
torch.manual_seed(seed)
|
|
np.random.seed(seed)
|
|
|
|
# Create input with known pattern for easier debugging
|
|
# Use a combination of smooth functions
|
|
y = torch.linspace(0, 2 * np.pi, height)
|
|
x = torch.linspace(0, 2 * np.pi, width)
|
|
yy, xx = torch.meshgrid(y, x, indexing='ij')
|
|
|
|
# Create smooth input pattern
|
|
input_tensor = torch.zeros(batch_size, in_channels, height, width)
|
|
for b in range(batch_size):
|
|
for c in range(in_channels):
|
|
# Different frequency for each channel
|
|
freq = (c + 1) * (b + 1)
|
|
input_tensor[b, c] = torch.sin(freq * xx) * torch.cos(freq * yy)
|
|
|
|
return input_tensor
|
|
|
|
|
|
def create_fno_weights_manual(config: dict, seed: int = 42) -> dict:
|
|
"""
|
|
Create FNO weights manually that match both neuraloperator and rtx-neural-operator.
|
|
|
|
Returns weights in rtx-neural-operator naming convention.
|
|
"""
|
|
torch.manual_seed(seed)
|
|
|
|
data_channels = config["data_channels"]
|
|
out_channels = config["out_channels"]
|
|
width = config["width"]
|
|
n_modes = config["n_modes"]
|
|
n_layers = config["n_layers"]
|
|
|
|
weights = {}
|
|
|
|
# Lifting MLP: (data_channels + 2) -> 2*width -> width
|
|
lifting_in = data_channels + 2 # +2 for positional encoding
|
|
hidden_dim = 2 * width
|
|
|
|
# Use Xavier/Glorot initialization for better numerical stability
|
|
scale_fc1 = np.sqrt(2.0 / (lifting_in + hidden_dim))
|
|
scale_fc2 = np.sqrt(2.0 / (hidden_dim + width))
|
|
|
|
weights["lifting.fcs.0.weight"] = torch.randn(hidden_dim, lifting_in) * scale_fc1
|
|
weights["lifting.fcs.0.bias"] = torch.zeros(hidden_dim)
|
|
weights["lifting.fcs.1.weight"] = torch.randn(width, hidden_dim) * scale_fc2
|
|
weights["lifting.fcs.1.bias"] = torch.zeros(width)
|
|
|
|
# Spectral convolution layers
|
|
for i in range(n_layers):
|
|
# SpectralConv weights: [in_ch, out_ch, modes_h, modes_w]
|
|
# Note: neuraloperator uses [in_ch, out_ch, ...] convention, not PyTorch's [out_ch, in_ch, ...]
|
|
# Use small initialization for spectral weights
|
|
scale_spectral = 1.0 / (width * n_modes[0] * n_modes[1])
|
|
|
|
# Shape is [in_channels, out_channels, modes_h, modes_w] to match Rust implementation
|
|
weights[f"spectral_conv.{i}.weights1_real"] = torch.randn(width, width, n_modes[0], n_modes[1]) * scale_spectral
|
|
weights[f"spectral_conv.{i}.weights1_imag"] = torch.randn(width, width, n_modes[0], n_modes[1]) * scale_spectral
|
|
weights[f"spectral_conv.{i}.weights2_real"] = torch.randn(width, width, n_modes[0], n_modes[1]) * scale_spectral
|
|
weights[f"spectral_conv.{i}.weights2_imag"] = torch.randn(width, width, n_modes[0], n_modes[1]) * scale_spectral
|
|
|
|
# Skip connection (1x1 conv as linear)
|
|
scale_conv = np.sqrt(2.0 / (width + width))
|
|
weights[f"conv.{i}.weight"] = torch.randn(width, width) * scale_conv
|
|
weights[f"conv.{i}.bias"] = torch.zeros(width)
|
|
|
|
# Projection layers: width -> 128 -> out_channels
|
|
proj_hidden = 128
|
|
scale_proj1 = np.sqrt(2.0 / (width + proj_hidden))
|
|
scale_proj2 = np.sqrt(2.0 / (proj_hidden + out_channels))
|
|
|
|
weights["projection.0.weight"] = torch.randn(proj_hidden, width) * scale_proj1
|
|
weights["projection.0.bias"] = torch.zeros(proj_hidden)
|
|
weights["projection.1.weight"] = torch.randn(out_channels, proj_hidden) * scale_proj2
|
|
weights["projection.1.bias"] = torch.zeros(out_channels)
|
|
|
|
return weights
|
|
|
|
|
|
def run_fno_forward_manual(input_tensor: torch.Tensor, weights: dict, config: dict) -> torch.Tensor:
|
|
"""
|
|
Run FNO forward pass manually to match rtx-neural-operator implementation.
|
|
|
|
This implements the same algorithm as the Rust code for direct comparison.
|
|
"""
|
|
batch_size, in_channels, height, width_dim = input_tensor.shape
|
|
model_width = config["width"]
|
|
n_modes = config["n_modes"]
|
|
n_layers = config["n_layers"]
|
|
|
|
# Step 1: Add positional encoding
|
|
grid = create_grid_positional_encoding(height, width_dim) # [2, H, W]
|
|
grid = grid.unsqueeze(0).expand(batch_size, -1, -1, -1) # [B, 2, H, W]
|
|
x = torch.cat([input_tensor, grid], dim=1) # [B, in_ch+2, H, W]
|
|
|
|
# Step 2: Lifting MLP
|
|
# Reshape for linear: [B, C, H, W] -> [B, H, W, C]
|
|
x = x.permute(0, 2, 3, 1)
|
|
|
|
# FC1 with GELU
|
|
fc1_w = weights["lifting.fcs.0.weight"] # [hidden, in]
|
|
fc1_b = weights["lifting.fcs.0.bias"] # [hidden]
|
|
x = torch.nn.functional.linear(x, fc1_w, fc1_b)
|
|
x = torch.nn.functional.gelu(x)
|
|
|
|
# FC2
|
|
fc2_w = weights["lifting.fcs.1.weight"] # [width, hidden]
|
|
fc2_b = weights["lifting.fcs.1.bias"] # [width]
|
|
x = torch.nn.functional.linear(x, fc2_w, fc2_b)
|
|
|
|
# Reshape back: [B, H, W, C] -> [B, C, H, W]
|
|
x = x.permute(0, 3, 1, 2)
|
|
|
|
# Step 3: Fourier layers
|
|
for i in range(n_layers):
|
|
# Spectral convolution path
|
|
x_ft = torch.fft.fft2(x) # [B, C, H, W] complex
|
|
|
|
# Get spectral weights
|
|
w1_real = weights[f"spectral_conv.{i}.weights1_real"]
|
|
w1_imag = weights[f"spectral_conv.{i}.weights1_imag"]
|
|
w2_real = weights[f"spectral_conv.{i}.weights2_real"]
|
|
w2_imag = weights[f"spectral_conv.{i}.weights2_imag"]
|
|
|
|
# Create complex weight tensors
|
|
# Weight shape: [in_ch, out_ch, modes_h, modes_w] (neuraloperator convention)
|
|
w1 = torch.complex(w1_real, w1_imag)
|
|
w2 = torch.complex(w2_real, w2_imag)
|
|
|
|
modes_h, modes_w = n_modes
|
|
|
|
# Initialize output in frequency domain
|
|
out_ft = torch.zeros_like(x_ft)
|
|
|
|
# Apply weights to upper left corner (positive frequencies)
|
|
# x_ft[:, :, :modes_h, :modes_w] with w1
|
|
x_upper = x_ft[:, :, :modes_h, :modes_w] # [B, in_ch, modes_h, modes_w]
|
|
# Einstein sum: batch, in_ch, h, w with in_ch, out_ch, h, w -> batch, out_ch, h, w
|
|
# Contract over in_ch, element-wise multiply h,w
|
|
out_upper = torch.einsum('bihw,iohw->bohw', x_upper, w1)
|
|
out_ft[:, :, :modes_h, :modes_w] = out_upper
|
|
|
|
# Apply weights to lower left corner (negative frequencies in height)
|
|
# x_ft[:, :, -modes_h:, :modes_w] with w2
|
|
x_lower = x_ft[:, :, -modes_h:, :modes_w]
|
|
out_lower = torch.einsum('bihw,iohw->bohw', x_lower, w2)
|
|
out_ft[:, :, -modes_h:, :modes_w] = out_lower
|
|
|
|
# Inverse FFT
|
|
x_spectral = torch.fft.ifft2(out_ft).real
|
|
|
|
# Skip connection path (1x1 conv as pointwise linear)
|
|
conv_w = weights[f"conv.{i}.weight"] # [out, in]
|
|
conv_b = weights[f"conv.{i}.bias"] # [out]
|
|
|
|
# Reshape for linear
|
|
x_skip = x.permute(0, 2, 3, 1) # [B, H, W, C]
|
|
x_skip = torch.nn.functional.linear(x_skip, conv_w, conv_b)
|
|
x_skip = x_skip.permute(0, 3, 1, 2) # [B, C, H, W]
|
|
|
|
# Combine and activate
|
|
x = x_spectral + x_skip
|
|
if i < n_layers - 1: # No activation on last layer
|
|
x = torch.nn.functional.gelu(x)
|
|
|
|
# Step 4: Projection
|
|
x = x.permute(0, 2, 3, 1) # [B, H, W, C]
|
|
|
|
# Projection layer 1 with GELU
|
|
proj1_w = weights["projection.0.weight"]
|
|
proj1_b = weights["projection.0.bias"]
|
|
x = torch.nn.functional.linear(x, proj1_w, proj1_b)
|
|
x = torch.nn.functional.gelu(x)
|
|
|
|
# Projection layer 2
|
|
proj2_w = weights["projection.1.weight"]
|
|
proj2_b = weights["projection.1.bias"]
|
|
x = torch.nn.functional.linear(x, proj2_w, proj2_b)
|
|
|
|
x = x.permute(0, 3, 1, 2) # [B, C, H, W]
|
|
|
|
return x
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser(
|
|
description="Generate reference inputs/outputs for Python vs Rust comparison"
|
|
)
|
|
parser.add_argument(
|
|
"--output",
|
|
type=str,
|
|
default="./weights/fno/fno2d_darcy",
|
|
help="Output directory for reference files"
|
|
)
|
|
parser.add_argument(
|
|
"--height",
|
|
type=int,
|
|
default=64,
|
|
help="Input height (default: 64)"
|
|
)
|
|
parser.add_argument(
|
|
"--width",
|
|
type=int,
|
|
default=64,
|
|
help="Input width (default: 64)"
|
|
)
|
|
parser.add_argument(
|
|
"--batch-size",
|
|
type=int,
|
|
default=1,
|
|
help="Batch size (default: 1)"
|
|
)
|
|
parser.add_argument(
|
|
"--seed",
|
|
type=int,
|
|
default=42,
|
|
help="Random seed (default: 42)"
|
|
)
|
|
|
|
args = parser.parse_args()
|
|
|
|
# Model configuration (matching fno2d_darcy)
|
|
config = {
|
|
"description": "FNO2d trained on Darcy Flow equation",
|
|
"url": "https://github.com/neuraloperator/neuraloperator",
|
|
"in_channels": 3, # data_channels + 2 (positional encoding)
|
|
"data_channels": 1, # Raw data channels
|
|
"out_channels": 1,
|
|
"width": 32,
|
|
"n_modes": [12, 12],
|
|
"n_layers": 4,
|
|
"pde_type": "darcy",
|
|
}
|
|
|
|
output_dir = Path(args.output)
|
|
output_dir.mkdir(parents=True, exist_ok=True)
|
|
|
|
print(f"Generating reference data with seed={args.seed}")
|
|
print(f" Input shape: [{args.batch_size}, {config['data_channels']}, {args.height}, {args.width}]")
|
|
print(f" Model: width={config['width']}, modes={config['n_modes']}, layers={config['n_layers']}")
|
|
|
|
# Create deterministic input
|
|
input_tensor = create_deterministic_input(
|
|
args.batch_size,
|
|
config["data_channels"],
|
|
args.height,
|
|
args.width,
|
|
seed=args.seed
|
|
)
|
|
print(f"\nInput tensor shape: {input_tensor.shape}")
|
|
print(f"Input range: [{input_tensor.min():.4f}, {input_tensor.max():.4f}]")
|
|
|
|
# Create weights
|
|
weights = create_fno_weights_manual(config, seed=args.seed)
|
|
print(f"\nCreated {len(weights)} weight tensors")
|
|
|
|
# Run forward pass
|
|
print("\nRunning forward pass...")
|
|
with torch.no_grad():
|
|
output_tensor = run_fno_forward_manual(input_tensor, weights, config)
|
|
|
|
print(f"Output tensor shape: {output_tensor.shape}")
|
|
print(f"Output range: [{output_tensor.min():.4f}, {output_tensor.max():.4f}]")
|
|
print(f"Output mean: {output_tensor.mean():.6f}")
|
|
print(f"Output std: {output_tensor.std():.6f}")
|
|
|
|
# Save weights
|
|
weights_path = output_dir / "model.safetensors"
|
|
save_file(weights, weights_path)
|
|
print(f"\nSaved weights to {weights_path}")
|
|
|
|
# Save config
|
|
config_path = output_dir / "config.json"
|
|
with open(config_path, "w") as f:
|
|
json.dump(config, f, indent=2)
|
|
print(f"Saved config to {config_path}")
|
|
|
|
# Save input tensor
|
|
input_path = output_dir / "reference_input.safetensors"
|
|
save_file({"input": input_tensor}, input_path)
|
|
print(f"Saved input to {input_path}")
|
|
|
|
# Save output tensor
|
|
output_path = output_dir / "reference_output.safetensors"
|
|
save_file({"output": output_tensor}, output_path)
|
|
print(f"Saved output to {output_path}")
|
|
|
|
# Save metadata
|
|
metadata = {
|
|
"seed": args.seed,
|
|
"batch_size": args.batch_size,
|
|
"height": args.height,
|
|
"width": args.width,
|
|
"input_shape": list(input_tensor.shape),
|
|
"output_shape": list(output_tensor.shape),
|
|
"output_mean": float(output_tensor.mean()),
|
|
"output_std": float(output_tensor.std()),
|
|
"output_min": float(output_tensor.min()),
|
|
"output_max": float(output_tensor.max()),
|
|
}
|
|
metadata_path = output_dir / "reference_metadata.json"
|
|
with open(metadata_path, "w") as f:
|
|
json.dump(metadata, f, indent=2)
|
|
print(f"Saved metadata to {metadata_path}")
|
|
|
|
print("\n" + "=" * 60)
|
|
print("Reference data generation complete!")
|
|
print("=" * 60)
|
|
print(f"\nFiles created in {output_dir}:")
|
|
print(" - model.safetensors (weights)")
|
|
print(" - config.json (model config)")
|
|
print(" - reference_input.safetensors (input tensor)")
|
|
print(" - reference_output.safetensors (expected output)")
|
|
print(" - reference_metadata.json (test metadata)")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|