Conserve quantities in generated diffusion

// Uses Philox addressed sampling with explicit seed and draw identities.
// Native constrained DDPM workflow. Run in a fresh directory for checkpoint creation.
// This short run demonstrates mechanics;
// the two-seed held-out experiment supplies calibration/physical quality evidence.
// Values and spacing 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 values = DynamicArray<Float64>(capacity: 80 * 25)
let conditions = DynamicArray<Float64>(capacity: 80)
var record = 0
while record < 80 {
    let amplitude = (Float64(record % 5) - 2.0) * 0.35
    conditions.append(amplitude)
    let fluctuation = 0.05 * random.normal(mean: 0.0, stddev: 1.0, seed: 903, draw: record)
    var row = 0
    while row < 5 {
        var col = 0
        while col < 5 {
            var value = 0.0
            if row > 0 && row < 4 && col > 0 && col < 4 {
                let x = Float64(col) / 4.0
                let y = Float64(row) / 4.0
                value = amplitude * math.sin(3.141592653589793 * x) * math.sin(3.141592653589793 * y)
                    + fluctuation * math.sin(6.283185307179586 * x) * math.sin(3.141592653589793 * y)
            }
            values.append(value)
            col = col + 1
        }
        row = row + 1
    }
    record = record + 1
}
let fields = tensor.reshape(values.to_tensor(rows: 80, cols: 25), List(80, 5, 5))
let training = gen.field_data(fields, conditions.to_tensor(rows: 80, cols: 1), role: "training")
let dx = 0.25
let dy = 0.25
// First sine mode has integral dx*dy*(1+sqrt(2))^2; second mode integrates zero.
let mass_scale = dx * dy * (1.0 + math.sqrt(2.0)) * (1.0 + math.sqrt(2.0))
let base = gen.diffusion_model(5, 5, 1, 32, 128, 0.00001, 0.1, 19)
let model = base.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)
// Calibration freezes the weights and uses only the bound training records.
let calibrated = fitted.calibrate_diffusion_variance(training, samples_per_timestep: 256, seed: 91)
calibrated.save("conserved-calibrated.anersci")
let restored = gen.load_scientific("conserved-calibrated.anersci")
print(restored.fingerprint() == calibrated.fingerprint())
print(restored.summary())
let heldout_condition = [[0.175]]
let generated = restored.sample_fields(heldout_condition, seed: 11)
let target = nn.scale(heldout_condition, mass_scale)
print(nn.item(gen.field_mass(generated, dx, dy)))
print(nn.item(target))
print(nn.item(nn.mse(gen.field_mass(generated, dx, dy), target)))
print(gen.field_statistics(generated, dx, dy))
dev.capture("Conserved sample", generated)
dev.capture("Reverse covariance", calibrated.diffusion_covariance())
// Resume changes weights, so the old calibration is cleared and must be refitted.
let continued = restored.resume(training, steps: 1)
print(continued.summary())