Calibrate scientific PDE generators

// Uses Philox addressed sampling with explicit seed and draw identities.
// Native conditional-Gaussian closure on an anisotropic screened field.
// Run in a fresh directory: checkpoints use no-overwrite publication.
// This short run demonstrates mechanics; separate 96-fit and 192-fit studies measure quality.
// All values and grid spacings are explicitly nondimensional.
import aner.gen
import aner.nn
import aner.tensor
import aner.math
import aner.parallel
import aner.dev
import aner.random

nn.configure("cpu", "Float64", "preserve")
let count = 320
let values = DynamicArray<Float64>(capacity: count * 35)
let conditions = DynamicArray<Float64>(capacity: count)
var record = 0
while record < count {
    let amplitude = (Float64(record % 5) - 2.0) * 0.35
    conditions.append(amplitude)
    let fluctuation_x = 0.05 * random.normal(mean: 0.0, stddev: 1.0, seed: 903, draw: record)
    let fluctuation_y = 0.03 * random.normal(mean: 0.0, stddev: 1.0, seed: 904, draw: record)
    var row = 0
    while row < 5 {
        var col = 0
        while col < 7 {
            var value = 0.0
            if row > 0 && row < 4 && col > 0 && col < 6 {
                let x_phase = 3.141592653589793 * Float64(col) / 6.0
                let y_phase = 3.141592653589793 * Float64(row) / 4.0
                value = amplitude * math.sin(x_phase) * math.sin(y_phase)
                    + fluctuation_x * math.sin(2.0 * x_phase) * math.sin(y_phase)
                    + fluctuation_y * math.sin(x_phase) * math.sin(2.0 * y_phase)
            }
            values.append(value)
            col = col + 1
        }
        row = row + 1
    }
    record = record + 1
}
let fields = tensor.reshape(values.to_tensor(rows: count, cols: 35), List(count, 5, 7))
let training = gen.field_data(fields, conditions.to_tensor(rows: count, cols: 1), role: "training")
let dx = 1.0 / 3.0
let dy = 0.25
// The two stochastic modes integrate to zero. This is a prescribed affine integral,
// not a conservation law for the source-driven screened equation below.
let mass_scale = dx * dy * (2.0 + math.sqrt(3.0)) * (1.0 + math.sqrt(2.0))
let base = gen.diffusion_model(5, 7, 1, 32, 128, 0.00001, 0.1, 19)
let velocity = base.with_diffusion_velocity_prediction(loss_weighting: "epsilon")
let spatial = velocity.with_diffusion_sobolev_loss(dx: dx, dy: dy, laplacian_weight: 1.0)
let model = spatial.with_mass_constraint(condition_index: 0, scale: mass_scale, offset: 0.0, dx: dx, dy: dy)
let fitted = model.train(training, steps: 500, batch: 32, rate: 0.001, seed: 37)
// Both fits use only bound training records and keep all neural weights fixed.
let covariance_only = fitted.calibrate_diffusion_variance(training, samples_per_timestep: 640, seed: 91)
// Explicit assumption for this linear, Gaussian-forced ensemble: affine conditional Gaussian fields.
// Fit clean training moments without noise integration; exact DDPM transitions preserve those moments.
// This closure bypasses the nonlinear noisy-field NN. The five-argument joint calibrator remains available.
let calibrated = fitted.calibrate_diffusion_gaussian(training)
print(calibrated.loss(training) == fitted.loss(training))
calibrated.save("gaussian-screened.anersci")
let restored = gen.load_scientific("gaussian-screened.anersci")
print(restored.fingerprint() == calibrated.fingerprint())
print(restored.summary())
let heldout_condition = [[0.175]]
let generated = restored.sample_fields(heldout_condition, seed: 11)
let baseline = covariance_only.sample_fields(heldout_condition, seed: 11)
print(nn.item(nn.mse(generated, calibrated.sample_fields(heldout_condition, seed: 11))))
print(nn.item(nn.mse(gen.field_mass(generated, dx, dy), nn.scale(heldout_condition, mass_scale))))
print(gen.field_statistics(generated, dx, dy))

// -div(K grad u) + reaction*u - forcing, with x along columns and y along rows.
// Face tensors can broadcast across a field batch, retaining dtype/device and gradients.
let kx_values = DynamicArray<Float64>(capacity: 30)
let ky_values = DynamicArray<Float64>(capacity: 28)
let reactions = DynamicArray<Float64>(capacity: 35)
let sources = DynamicArray<Float64>(capacity: 35)
var index = 0
while index < 30 {
    kx_values.append(0.7)
    index = index + 1
}
index = 0
while index < 28 {
    ky_values.append(2.0)
    index = index + 1
}
let sx = math.sin(3.141592653589793 / 12.0)
let sy = math.sin(3.141592653589793 / 8.0)
let lambda = 4.0 * 0.7 * sx * sx / (dx * dx) + 4.0 * 2.0 * sy * sy / (dy * dy) + 2.5
var row = 0
while row < 5 {
    var col = 0
    while col < 7 {
        reactions.append(2.5)
        var source = 0.0
        if row > 0 && row < 4 && col > 0 && col < 6 {
            source = lambda * 0.175 * math.sin(3.141592653589793 * Float64(col) / 6.0)
                * math.sin(3.141592653589793 * Float64(row) / 4.0)
        }
        sources.append(source)
        col = col + 1
    }
    row = row + 1
}
let kx = tensor.reshape(kx_values.to_tensor(rows: 1, cols: 30), List(1, 5, 6))
let ky = tensor.reshape(ky_values.to_tensor(rows: 1, cols: 28), List(1, 4, 7))
let reaction = tensor.reshape(reactions.to_tensor(rows: 1, cols: 35), List(1, 5, 7))
let forcing = tensor.reshape(sources.to_tensor(rows: 1, cols: 35), List(1, 5, 7))
let residual = gen.field_elliptic_residual(generated, forcing, kx, ky, reaction, dx, dy)
// Nominal forcing omits stochastic forcing: valid random modes have nonzero residual here.
dev.capture("Conditional Gaussian sample", generated)
dev.capture("Covariance-only sample", baseline)
dev.capture("Screened residual against nominal forcing", residual)
dev.capture("Reverse covariance", calibrated.diffusion_covariance())
// Weight updates invalidate both fitted corrections; resume uses the raw neural objective.
let continued = restored.resume(training, steps: 1)
print(continued.summary())