Native TorchLean FNO1D Burgers

Read this after the basic CNN/MLP examples if you want the operator-learning path. The Python scripts do the two jobs Lean should not own here: download and reshape the public burgers_data_R10.mat file, then plot the prediction CSV. The model, loss, optimizer, and training loop stay in TorchLean.

This executable uses real-split Fourier arithmetic because the Burgers data are real-valued. The portable path evaluates the dense multidimensional DFT with separate real and imaginary tensors. On CUDA, the fused spectralConv1dRfft autograd primitive uses the same representation for its weights and executes the transforms through cuFFT.

The training task follows the standard FNO Burgers setup: learn the operator $u_0(x)\mapsto u(x,T)$ on a fixed periodic grid. The default grid and row counts are modest enough for a local run while still exercising the real operator-learning path. Larger runs can raise --steps, export more rows, and bump the constants below.

References for the dataset/training convention:

• Li et al., “Fourier Neural Operator for Parametric Partial Differential Equations”, 2020/2021.
• MathWorks’ Burgers FNO example and the burgers_data_R10.mat public dataset.
• SciML FNO tutorials using fields a for initial conditions and u for final solutions.

def NN.Examples.Models.Operators.Fno1dBurgers.exeName : String

CLI subcommand name used in terminal banners and errors.

def NN.Examples.Models.Operators.Fno1dBurgers.grid : ℕ

Spatial grid resolution used by the prepared Burgers .npy slices.

def NN.Examples.Models.Operators.Fno1dBurgers.width : ℕ

Channel width inside the compact FNO block.

def NN.Examples.Models.Operators.Fno1dBurgers.modes : ℕ

Number of Fourier modes retained on each side of the real FFT spectrum.

def NN.Examples.Models.Operators.Fno1dBurgers.blocks : ℕ

Number of spectral blocks used by the compact training run.

def NN.Examples.Models.Operators.Fno1dBurgers.defaultTrainRows : ℕ

Default number of training rows expected from the preparation script.

def NN.Examples.Models.Operators.Fno1dBurgers.defaultTestRows : ℕ

Default number of held-out rows expected from the preparation script.

def NN.Examples.Models.Operators.Fno1dBurgers.modelCfg : TorchLean.nn.models.FNOConfig 1

Spec.Shape-level FNO configuration shared by the constructor and sample loaders.

@[reducible, inline]

abbrev NN.Examples.Models.Operators.Fno1dBurgers.σ : Spec.Shape

Model input shape: one sampled initial condition on the fixed grid.

@[reducible, inline]

abbrev NN.Examples.Models.Operators.Fno1dBurgers.τ : Spec.Shape

Model output shape: one predicted terminal solution on the same grid.

def NN.Examples.Models.Operators.Fno1dBurgers.defaultDir : System.FilePath

Directory where the preparation script writes Burgers tensors by default.

def NN.Examples.Models.Operators.Fno1dBurgers.trainXPath : System.FilePath

Default training input tensor path.

def NN.Examples.Models.Operators.Fno1dBurgers.trainYPath : System.FilePath

Default training target tensor path.

def NN.Examples.Models.Operators.Fno1dBurgers.testXPath : System.FilePath

Default held-out input tensor path.

def NN.Examples.Models.Operators.Fno1dBurgers.testYPath : System.FilePath

Default held-out target tensor path.

def NN.Examples.Models.Operators.Fno1dBurgers.defaultPlotCsv : System.FilePath

Default CSV path for the prediction-vs-target plot script.

def NN.Examples.Models.Operators.Fno1dBurgers.defaultLogJson : System.FilePath

Default JSON training-log path.

def NN.Examples.Models.Operators.Fno1dBurgers.missingDataHint : String

User-facing hint printed when the prepared Burgers tensors are missing.

structure NN.Examples.Models.Operators.Fno1dBurgers.BurgersOptionsextends NN.Examples.ModelZoo.SeededTrainFlags, NN.Examples.ModelZoo.PairedNpyEvalFlags, NN.Examples.ModelZoo.CsvArtifactFlags : Type

FNO Burgers command-line options: training flags, data paths, and artifact paths.

Seeded optimizer/log flags come from ModelZoo, the Burgers tensor paths use ModelZoo.PairedNpyEvalFlags, and the plot path uses ModelZoo.CsvArtifactFlags.

• steps : ℕ
• batchSize : ℕ
• log : Training.LogDestination
• logPath : System.FilePath
• cudaMemWatch : ℕ
• lr : Float
• seed : ℕ
• trainRows : ℕ
• testRows : ℕ
• evalRows : ℕ
• trainX : System.FilePath
• trainY : System.FilePath
• testX : System.FilePath
• testY : System.FilePath
• plotCsv : System.FilePath

def NN.Examples.Models.Operators.Fno1dBurgers.instReprBurgersOptions.repr : BurgersOptions → ℕ → Std.Format

@[implicit_reducible]

instance NN.Examples.Models.Operators.Fno1dBurgers.instReprBurgersOptions : Repr BurgersOptions

def NN.Examples.Models.Operators.Fno1dBurgers.dataPaths (cfg : BurgersOptions) : List System.FilePath

All required dataset files for this run.

def NN.Examples.Models.Operators.Fno1dBurgers.BurgersOptions.parse (args : List String) : Except String (BurgersOptions × List String)

Parse the FNO Burgers command-line options.

def NN.Examples.Models.Operators.Fno1dBurgers.BurgersOptions.effectiveCudaMemWatch (cfg : BurgersOptions) (opts : Runtime.Autograd.Torch.Options) : ℕ

Effective CUDA-memory-watch cadence for this run.

def NN.Examples.Models.Operators.Fno1dBurgers.BurgersOptions.logNotes (cfg : BurgersOptions) (spectralPath device : String) : Array String

TrainLog note fields for the fused CUDA execution path.

def NN.Examples.Models.Operators.Fno1dBurgers.mkModel : TorchLean.nn.M (TorchLean.nn.Sequential σ τ)

def NN.Examples.Models.Operators.Fno1dBurgers.loadDataset (xPath yPath : System.FilePath) (n : ℕ) : IO (Runtime.Autograd.Train.Dataset (TorchLean.SupervisedSample Float σ τ))

Load one fixed-grid Burgers split as supervised TorchLean samples.

def NN.Examples.Models.Operators.Fno1dBurgers.writePredictionProbe (plotCsv : System.FilePath) (x target prediction : Spec.Tensor Float σ) : IO Unit

Write one FNO prediction row to CSV for the companion plotting script.

def NN.Examples.Models.Operators.Fno1dBurgers.metricHistory : Runtime.Training.MetricHistory

def NN.Examples.Models.Operators.Fno1dBurgers.writeMetricLog (dest : Runtime.Training.LogDestination) (hist : Runtime.Training.MetricHistory) (cfg : BurgersOptions) (spectralPath device : String) : IO Unit

Persist the train/test MSE history with model/data metadata attached.

structure NN.Examples.Models.Operators.Fno1dBurgers.LoadedData : Type

Loaded train/test splits before evaluation prefixes and cycling streams are derived.

• train : Runtime.Autograd.Train.Dataset (TorchLean.SupervisedSample Float σ τ)

  Training split as supervised samples.

• test : Runtime.Autograd.Train.Dataset (TorchLean.SupervisedSample Float σ τ)

  Held-out split as supervised samples.

def NN.Examples.Models.Operators.Fno1dBurgers.loadData (cfg : BurgersOptions) : IO LoadedData

Validate paths and load both Burgers splits.

structure NN.Examples.Models.Operators.Fno1dBurgers.EvalData : Type

Deterministic evaluation prefixes and cycling stream derived from the loaded train/test sets.

• trainDatasetSamples : List (TorchLean.SupervisedSample Float σ τ)
• testDatasetSamples : List (TorchLean.SupervisedSample Float σ τ)
• trainSamples : List (Spec.Tensor Float σ × Spec.Tensor Float τ)
• testSamples : List (Spec.Tensor Float σ × Spec.Tensor Float τ)
• reportTrainDatasetSamples : List (TorchLean.SupervisedSample Float σ τ)
• reportTestDatasetSamples : List (TorchLean.SupervisedSample Float σ τ)
• reportTrainSamples : List (Spec.Tensor Float σ × Spec.Tensor Float τ)
• reportTestSamples : List (Spec.Tensor Float σ × Spec.Tensor Float τ)
• trainCycle : ℕ → TorchLean.SupervisedSample Float σ τ

def NN.Examples.Models.Operators.Fno1dBurgers.mkEvalData (cfg : BurgersOptions) (data : LoadedData) : IO EvalData

Convert loaded Burgers datasets into the common runtime/evaluation view used by both execution paths.

The fused CUDA path:

• evaluate on fixed deterministic prefixes,
• train by cycling through the finite dataset with seed + step, and
• emit the same train/test MSE metric history.

def NN.Examples.Models.Operators.Fno1dBurgers.pushLossPoint (hist : Runtime.Training.MetricHistory) (step : ℕ) (tag : String) (trainLoss testLoss : Float) : IO Runtime.Training.MetricHistory

Push one train/test MSE point into the metric history and print the tagged report line.

@[reducible, inline]

abbrev NN.Examples.Models.Operators.Fno1dBurgers.FusedCuda.Param : Type

Fused CUDA parameter packet for the real-FFT FNO kernel.

def NN.Examples.Models.Operators.Fno1dBurgers.FusedCuda.meanLoss (ps : Array Param) (samples : List (Spec.Tensor Float σ × Spec.Tensor Float τ)) : IO Float

Mean MSE over a finite evaluation prefix using the fused CUDA FNO implementation.

def NN.Examples.Models.Operators.Fno1dBurgers.FusedCuda.evalLosses (trainEval testEval : List (Spec.Tensor Float σ × Spec.Tensor Float τ)) (ps : Array Param) : IO (Float × Float)

Train/test MSE pair for the current fused CUDA parameters.

def NN.Examples.Models.Operators.Fno1dBurgers.FusedCuda.recordEval (trainEval testEval : List (Spec.Tensor Float σ × Spec.Tensor Float τ)) (hist : Runtime.Training.MetricHistory) (step : ℕ) (ps : Array Param) (tag : String) : IO Runtime.Training.MetricHistory

Append one fused-CUDA evaluation point to the metric history.

def NN.Examples.Models.Operators.Fno1dBurgers.FusedCuda.predict (ps : Array Param) (x : Spec.Tensor Float σ) : IO (Spec.Tensor Float τ)

Predict one Burgers terminal field through the fused CUDA spectral path.

def NN.Examples.Models.Operators.Fno1dBurgers.FusedCuda.trainStep (lr : Float) (ps : Array Param) (adamSt : Runtime.Autograd.Cuda.Fno1dRfftFused.AdamState) (sample : Spec.Tensor Float σ × Spec.Tensor Float τ) : IO (Array Param × Runtime.Autograd.Cuda.Fno1dRfftFused.AdamState)

One fused CUDA Adam update on a single Burgers sample.

def NN.Examples.Models.Operators.Fno1dBurgers.FusedCuda.run (cfg : BurgersOptions) : IO Unit

Run the fused cuFFT/RFFT training path and emit its training and prediction artifacts.

def NN.Examples.Models.Operators.Fno1dBurgers.runPortableDense (opts : Runtime.Autograd.Torch.Options) (cfg : BurgersOptions) : IO Unit

def NN.Examples.Models.Operators.Fno1dBurgers.logRunHeader (opts : Runtime.Autograd.Torch.Options) (cfg : BurgersOptions) : IO Unit

def NN.Examples.Models.Operators.Fno1dBurgers.main (args : List String) : IO UInt32