Added example code blocks

Author: e3p7y

.DS_Store

@@ -21,7 +21,11 @@ where $P$ are the physical parameters, $\overline{T}$ is the averaged temperatur
 
 ## Generating Data for Training
 
-To train the neural network, we can generate data using the function $\overline{wT} = cos(sin(T^3)) + sin(cos(T^2))$ with $N$ spatial points discretized by a finite difference method, with the time domain $t \in [0,1.5]$ and Neumann zero-flux boundary conditions, meaning $\frac{\partial \overline{T}}{\partial z} = 0$ at the edges.
+To train the neural network, we can generate data using the function: 
+
+$$\overline{wT} = {cos(sin(T^3)) + sin(cos(T^2))}$$
+
+with $N$ spatial points discretized by a finite difference method, with the time domain $t \in [0,1.5]$ and Neumann zero-flux boundary conditions, meaning $\frac{\partial \overline{T}}{\partial z} = 0$ at the edges.
 
 
 ## Training Neural Network

docs/src/examples/pde/boussinesq.md

@@ -0,0 +1,122 @@
+using OrdinaryDiffEq, SciMLSensitivity
+using Lux, Optimization, OptimizationOptimisers, ComponentArrays
+using Random, Statistics, LinearAlgebra
+
+# ==========================================
+# 1. Problem Setup & Data Generation
+# ==========================================
+const N = 32                  # Number of spatial points
+const L = 1.0f0               # Length of vertical domain
+const dx = L / N              # Grid spacing
+const κ = 0.01f0              # Thermal diffusivity (kappa)
+const tspan = (0.0f0, 1.5f0)  # Time domain
+
+# "True" closure term: wT = cos(sin(T^3)) + sin(cos(T^2))
+function true_flux_closure(T)
+    cos(sin(T^3)) + sin(cos(T^2))
+end
+
+# Discretized PDE (Method of Lines)
+function boussinesq_pde!(du, u, p, t, flux_model)
+    # Calculate Flux Term (wT)
+    # flux_model can be the true function or the Neural Network
+    wT = flux_model(u, p)
+    
+    # Finite Difference Scheme
+    # Equation: ∂T/∂t = κ * ∂²T/∂z² - ∂(wT)/∂z
+    
+    # Interior Points (2 to N-1)
+    for i in 2:N-1
+        # Diffusion: Central Difference
+        diffusion = κ * (u[i+1] - 2u[i] + u[i-1]) / (dx^2)
+        
+        # Advection/Flux: Central Difference
+        advection = (wT[i+1] - wT[i-1]) / (2dx)
+        
+        du[i] = diffusion - advection
+    end
+    
+    # Boundary Conditions: Neumann Zero-Flux (∂T/∂z = 0)
+    # Implies T[0] = T[2] and T[N+1] = T[N-1] (Mirror ghost points)
+    # At boundaries, if T is flat, wT(T) is flat, so ∂(wT)/∂z = 0
+    
+    # Left Boundary (z=0)
+    du[1] = κ * (2(u[2] - u[1])) / (dx^2) 
+    
+    # Right Boundary (z=L)
+    du[N] = κ * (2(u[N-1] - u[N])) / (dx^2)
+end
+
+# Initial Condition: A simple curve (e.g., Gaussian or linear gradient)
+u0 = Float32[exp(-(x-0.5)^2 / 0.1) for x in range(0, L, length=N)]
+
+# wrapper for the "True" solver
+function pde_true!(du, u, p, t)
+    # Helper to map function over vector u
+    wrapper(u_in, p_in) = true_flux_closure.(u_in)
+    boussinesq_pde!(du, u, p, t, wrapper)
+end
+
+# Generate Training Data
+prob_true = ODEProblem(pde_true!, u0, tspan)
+sol_true = solve(prob_true, Tsit5(), saveat=0.1)
+training_data = Array(sol_true)
+
+# ==========================================
+# 2. Define the Universal Differential Equation
+# ==========================================
+# Neural Network to approximate wT = U(T)
+# Input: 1 (Temperature), Output: 1 (Flux)
+rng = Random.default_rng()
+nn = Lux.Chain(
+    Lux.Dense(1, 8, tanh),
+    Lux.Dense(8, 8, tanh),
+    Lux.Dense(8, 1)
+)
+p_nn, st = Lux.setup(rng, nn)
+
+# Wrapper for NN to fit the pde function signature
+function nn_closure(u, p)
+    # Lux expects (Features, Batch). Reshape u to (1, N)
+    u_in = reshape(u, 1, :)
+    pred = nn(u_in, p, st)[1]
+    return reshape(pred, :) # Return as vector
+end
+
+function pde_ude!(du, u, p, t)
+    boussinesq_pde!(du, u, p, t, nn_closure)
+end
+
+prob_ude = ODEProblem(pde_ude!, u0, tspan, p_nn)
+
+# ==========================================
+# 3. Training Loop
+# ==========================================
+function loss(p)
+    # Solve UDE with current NN parameters
+    # Use InterpolatingAdjoint for memory efficiency with stiff PDEs
+    sol_pred = solve(prob_ude, Tsit5(), p=p, saveat=0.1, 
+                     sensealg=InterpolatingAdjoint(autojacvec=ReverseDiffVJP(true)))
+    
+    # Calculate MSE against training data
+    if sol_pred.retcode != :Success
+        return Inf
+    end
+    return mean(abs2, Array(sol_pred) .- training_data)
+end
+
+# Define Optimization Problem
+adtype = Optimization.AutoZygote()
+optf = Optimization.OptimizationFunction((x, p) -> loss(x), adtype)
+optprob = Optimization.OptimizationProblem(optf, ComponentVector(p_nn))
+
+# Training Stage 1: ADAM with lr=1e-2 for 200 iterations
+println("Starting Training Phase 1...")
+res1 = Optimization.solve(optprob, OptimizationOptimisers.Adam(0.01), maxiters=200)
+
+# Training Stage 2: ADAM with lr=1e-3 for 1000 iterations
+println("Starting Training Phase 2...")
+optprob2 = Optimization.OptimizationProblem(optf, res1.u)
+res2 = Optimization.solve(optprob2, OptimizationOptimisers.Adam(0.001), maxiters=1000)
+
+println("Final Loss: ", loss(res2.u))