Native TorchLean FNO1D Burgers

This file is the operator-learning tutorial we want people to read after the basic CNN/MLP examples. The Python helpers do the two jobs Lean should not own here: download/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. We keep the default grid and row counts modest because the first run should answer one question quickly: "is my TorchLean/CUDA path wired correctly?" Once that works, 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

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

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

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

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

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

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

def NN.Examples.Models.Operators.Fno1dBurgers.modelCfg : API.nn.models.Fno1dConfig

@[reducible, inline]

abbrev NN.Examples.Models.Operators.Fno1dBurgers.Vec (n : ℕ) : Shape

@[reducible, inline]

abbrev NN.Examples.Models.Operators.Fno1dBurgers.Mat (m n : ℕ) : Shape

@[reducible, inline]

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

@[reducible, inline]

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

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

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

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

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

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

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

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

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

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

• steps : ℕ
• seed : ℕ
• logEvery : ℕ
• trainRows : ℕ
• testRows : ℕ
• evalRows : ℕ
• lr : Float

@[implicit_reducible]

instance NN.Examples.Models.Operators.Fno1dBurgers.instReprTrainConfig : Repr TrainConfig

def NN.Examples.Models.Operators.Fno1dBurgers.instReprTrainConfig.repr : TrainConfig → ℕ → Std.Format

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

• trainX : System.FilePath
• trainY : System.FilePath
• testX : System.FilePath
• testY : System.FilePath

def NN.Examples.Models.Operators.Fno1dBurgers.DataFiles.paths (files : DataFiles) : List System.FilePath

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

• train : TrainConfig
• files : DataFiles
• plotCsv : System.FilePath
• logJson : System.FilePath

def NN.Examples.Models.Operators.Fno1dBurgers.parseFlags (args : List String) : Except String (RunSpec × List String)

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

def NN.Examples.Models.Operators.Fno1dBurgers.loadDataset {α : Type} [API.Runtime.Scalar α] (xPath yPath : System.FilePath) (n : ℕ) : IO (Runtime.Autograd.Train.Dataset (API.sample.Supervised α σ τ))

def NN.Examples.Models.Operators.Fno1dBurgers.sampleToInputTarget {α : Type} {σ τ : Shape} : API.sample.Supervised α σ τ → Spec.Tensor α σ × Spec.Tensor α τ

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

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

def NN.Examples.Models.Operators.Fno1dBurgers.trainLogNotes (spec : RunSpec) (spectralPath device : String) : Array String

def NN.Examples.Models.Operators.Fno1dBurgers.writeMetricLog (path : System.FilePath) (hist : Runtime.Training.MetricHistory) (spec : RunSpec) (spectralPath device : String) : IO Unit

structure NN.Examples.Models.Operators.Fno1dBurgers.LoadedData (α : Type) : Type

• train : Runtime.Autograd.Train.Dataset (API.sample.Supervised α σ τ)
• test : Runtime.Autograd.Train.Dataset (API.sample.Supervised α σ τ)

def NN.Examples.Models.Operators.Fno1dBurgers.loadData {α : Type} [API.Runtime.Scalar α] (files : DataFiles) (cfg : TrainConfig) : IO (LoadedData α)

@[reducible, inline]

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

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

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

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

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

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)

def NN.Examples.Models.Operators.Fno1dBurgers.FusedCuda.run (spec : RunSpec) : IO Unit

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

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

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