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.

Why we use the real-split FNO path in this executable:

• NN.FNO1D.model is the mathematically clean complex-domain implementation.
• The eager CUDA backend stores float32 buffers, not complex buffers.
• On CUDA this run uses the fused spectralConv1dRfft autograd primitive, which represents Fourier weights by real/imaginary float32 buffers and executes the real FFT path through cuFFT.
• On CPU it falls back to the dense DFT implementation. That is slower, but it is the useful reference path when someone wants to inspect the math without CUDA in the way.

The training task follows the standard FNO Burgers setup: learn the operator u₀(x) ↦ 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. Kept small so the eager reference path remains usable.

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.Fno1dConfig

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

@[reducible, inline]

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

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

@[reducible, inline]

abbrev NN.Examples.Models.Operators.Fno1dBurgers.τ : 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.API.Common.SeededModelTrainFlags, NN.API.Common.PairedNpyEvalFlags, NN.API.Common.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 : Runtime.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

@[implicit_reducible]

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

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

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 : TorchLean.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 shared by the fused CUDA and portable dense execution paths.

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 (TorchLean.Data.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 : TorchLean.Training.MetricHistory

def NN.Examples.Models.Operators.Fno1dBurgers.writeMetricLog (dest : TorchLean.Training.LogDestination) (hist : TorchLean.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 : TorchLean.Data.Dataset (TorchLean.SupervisedSample Float σ τ)

  Training split as supervised samples.

• test : TorchLean.Data.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.

Both the fused CUDA path and the portable dense 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 : TorchLean.Training.MetricHistory) (step : ℕ) (tag : String) (trainLoss testLoss : Float) : IO TorchLean.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 : TorchLean.Training.MetricHistory) (step : ℕ) (ps : Array Param) (tag : String) : IO TorchLean.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 the same artifacts as the dense path.

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

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

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