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 transforms each spatial axis with separate real and imaginary tensors. On CUDA, the command deliberately selects a specialized one-sided real-FFT model whose transforms run through cuFFT. The two paths share the typed field-to-field boundary and training task, but they do not share an identical spectral parameter layout.

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.gridSize : ℕ

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 : ℕ

Spectral mode budget.

The portable full-DFT path uses this as the width of each end band. The fused real-FFT path uses it as the number of stored nonnegative-frequency bins.

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.

@[reducible, inline]

abbrev NN.Examples.Models.Operators.Fno1dBurgers.modelConfig : TorchLean.nn.models.FNO.Config 1

FNO configuration shared by the constructor and sample loaders.

@[reducible, inline]

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

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

@[reducible, inline]

abbrev NN.Examples.Models.Operators.Fno1dBurgers.output : 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.defaultLogPath : 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.Options : Type

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

The record keeps optimizer/log settings, reproducibility, tensor paths, and output artifacts as separate named concerns.

• training : TorchLean.CLI.Training.OptimizerOptions

  Optimizer, step, batching, and logging controls.

• seed : ℕ

  Seed used for initialization and training-row selection.

• data : Support.PairedNpyEvalFlags

  Prepared train/test tensors and evaluation row counts.

• artifacts : Support.CsvArtifactFlags

  Output path for the prediction diagnostic.

@[instance_reducible]

instance NN.Examples.Models.Operators.Fno1dBurgers.instReprOptions : Repr Options

def NN.Examples.Models.Operators.Fno1dBurgers.instReprOptions.repr : Options → ℕ → Std.Format

def NN.Examples.Models.Operators.Fno1dBurgers.dataPaths (config : Options) : Array System.FilePath

All required dataset files for this run.

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

Parse the FNO Burgers command-line options.

def NN.Examples.Models.Operators.Fno1dBurgers.Options.memoryCadence (config : Options) (runtime : TorchLean.Runtime.Config) : ℕ

Effective CUDA-memory-watch cadence for this run.

def NN.Examples.Models.Operators.Fno1dBurgers.Options.logNotes (config : Options) (spectralPath device : String) : Array String

TrainLog note fields for the fused CUDA execution path.

def NN.Examples.Models.Operators.Fno1dBurgers.model : TorchLean.nn.Builder (TorchLean.nn.Sequential input output)

The Fourier neural operator this example trains, at the width, mode count and depth fixed by modelConfig.

def NN.Examples.Models.Operators.Fno1dBurgers.loadSplit (xPath yPath : System.FilePath) (n : ℕ) : IO (TorchLean.Data.SampleStream (TorchLean.Sample.Supervised Float input output))

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

def NN.Examples.Models.Operators.Fno1dBurgers.writePrediction (plotCsv : System.FilePath) (x target prediction : TorchLean.Tensor Float input) : IO Unit

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

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

Two tracked curves, train and test MSE, with the colours the log viewer will use.

def NN.Examples.Models.Operators.Fno1dBurgers.writeLog (destination : Runtime.Training.LogDestination) (history : Runtime.Training.MetricHistory) (config : Options) (spectralPath device : String) : IO Unit

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

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

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

• train : TorchLean.Data.SampleStream (TorchLean.Sample.Supervised Float input output)

  Training split as supervised samples.

• test : TorchLean.Data.SampleStream (TorchLean.Sample.Supervised Float input output)

  Held-out split as supervised samples.

def NN.Examples.Models.Operators.Fno1dBurgers.load (config : Options) : IO Splits

Validate paths and load both Burgers splits.

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

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

• train : TorchLean.Data.SampleStream (TorchLean.Tensor Float input × TorchLean.Tensor Float output)
• test : TorchLean.Data.SampleStream (TorchLean.Tensor Float input × TorchLean.Tensor Float output)
• sample : TorchLean.Sample.Supervised Float input output

  Training sample used to initialize the runtime's compiled evaluation path.

• probe : TorchLean.Tensor Float input × TorchLean.Tensor Float output

  Held-out sample used for the prediction CSV emitted by both execution paths.

• next : ℕ → TorchLean.Sample.Supervised Float input output

def NN.Examples.Models.Operators.Fno1dBurgers.prepare (config : Options) (data : Splits) : IO Evaluation

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

Both execution paths:

• 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.record (history : 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 (parameters : Array Param) (samples : TorchLean.Data.SampleStream (TorchLean.Tensor Float input × TorchLean.Tensor Float output)) : IO Float

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

def NN.Examples.Models.Operators.Fno1dBurgers.FusedCuda.losses (trainEval testEval : TorchLean.Data.SampleStream (TorchLean.Tensor Float input × TorchLean.Tensor Float output)) (parameters : Array Param) : IO (Float × Float)

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

def NN.Examples.Models.Operators.Fno1dBurgers.FusedCuda.record (trainEval testEval : TorchLean.Data.SampleStream (TorchLean.Tensor Float input × TorchLean.Tensor Float output)) (history : Runtime.Training.MetricHistory) (step : ℕ) (parameters : 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.run (config : Options) : IO Unit

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

def NN.Examples.Models.Operators.Fno1dBurgers.runPortable (runtime : TorchLean.Runtime.Config) (config : Options) : IO Unit

Train and evaluate using the portable full-spectrum model.

The transforms have a dense per-axis fallback on CPU. The fused cuFFT path above uses a different one-sided parameter layout; both models can be trained on the same data with the same seed.

def NN.Examples.Models.Operators.Fno1dBurgers.printHeader (runtime : TorchLean.Runtime.Config) (config : Options) : IO Unit

Print the run's configuration: device, execution mode, model geometry, row counts and file paths.

Worth the space, because an FNO run that silently loaded the wrong split looks like a training problem rather than a data problem.

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

Entry point; dispatches to the fused CUDA path when the device supports it and to runPortable otherwise.