rustytorch/crates/specialized/rtx-fsi/tests/fsi2_harness/overset_march.rs

//! P5 (§5.12): the coupled march on the overset fluid — the harness's
//! rigid phase, release, and per-step subiterated coupling (predictor,
//! IQN-ILS / Aitken passes each re-marching the fluid from the step's
//! snapshot, the acceptance rule of `march.rs`), without the embedded
//! march's rescue machinery (refuted, retired). Returns the harness's
//! `MarchResult` plus the composite's own counters.

use std::cell::RefCell;
use std::io::Write as _;

use nalgebra::Vector3;
use rtx_fea::analysis::{
    AnalysisConfig, ConvergenceCriteria, DynamicState, NonlinearDynamicAnalysis,
};
use rtx_fea::materials::{LinearElastic, MaterialDatabase};
use rtx_fea::mesh::{MaterialId, NodeId};
use rtx_fsi::{IqnIls, Subiterated};

use super::march::MarchResult;
use super::overset::OversetFluid;
use super::{clamp_left, env_or, median, mid_amp, BenchmarkCase};

/// The overset march's knobs (`RTX_<PREFIX>_*`).
#[derive(Debug, Clone)]
pub struct OversetMarchConfig {
    pub ny: usize,
    pub flag_nx: usize,
    pub t_release: f64,
    pub t_end: f64,
    pub subcycle: usize,
    pub tol_floor: f64,
    pub rtol: f64,
    pub stall_accept: f64,
    pub max_subiterations: usize,
    pub coupler: String,
    pub reuse: usize,
    pub initial_relaxation: f64,
    pub c1_interface: bool,
    pub predictor: String,
    /// Winslow sweeps per patch regeneration.
    pub sweeps: usize,
    /// Schwarz rounds per corrector (the P5 budget: 3).
    pub max_rounds: usize,
    pub csv_path: Option<String>,
    pub trace_steps: usize,
}

impl OversetMarchConfig {
    pub fn from_env(prefix: &str, d: OversetMarchConfig) -> OversetMarchConfig {
        let num = |k: &str, v: f64| env_or(&format!("RTX_{prefix}_{k}"), v);
        let text =
            |k: &str, v: &str| std::env::var(format!("RTX_{prefix}_{k}")).unwrap_or(v.into());
        OversetMarchConfig {
            ny: num("NY", d.ny as f64) as usize,
            flag_nx: num("FLAG_NX", d.flag_nx as f64) as usize,
            t_release: num("T_RELEASE", d.t_release),
            t_end: num("T_END", d.t_end),
            subcycle: num("SUBCYCLE", d.subcycle as f64) as usize,
            tol_floor: num("TOL_FLOOR", d.tol_floor),
            rtol: num("RTOL", d.rtol),
            stall_accept: num("STALL_ACCEPT", d.stall_accept),
            max_subiterations: num("MAX_SUBIT", d.max_subiterations as f64) as usize,
            coupler: text("COUPLER", &d.coupler),
            reuse: num("REUSE", d.reuse as f64) as usize,
            initial_relaxation: num("OMEGA0", d.initial_relaxation),
            c1_interface: num("C1", if d.c1_interface { 1.0 } else { 0.0 }) > 0.5,
            predictor: text("PREDICTOR", &d.predictor),
            sweeps: num("SWEEPS", d.sweeps as f64) as usize,
            max_rounds: num("MAX_ROUNDS", d.max_rounds as f64) as usize,
            csv_path: std::env::var(format!("RTX_{prefix}_CSV"))
                .ok()
                .or(d.csv_path),
            trace_steps: num("TRACE", d.trace_steps as f64) as usize,
        }
    }
}

/// The march's outcome: the harness's statistics, the composite's
/// counters, and the death (if any) instead of a panic.
pub struct OversetMarchResult {
    pub result: MarchResult,
    pub death: Option<(usize, f64, String)>,
    pub rounds_mean: f64,
    pub reclassified_mean: f64,
    pub fresh_mean: f64,
    pub regen_count: usize,
    pub regen_seconds: f64,
    pub fluid_seconds: f64,
    pub structure_seconds: f64,
    pub faces_used: usize,
}

pub fn run_march_overset(case: BenchmarkCase, config: &OversetMarchConfig) -> OversetMarchResult {
    let cfg = config.clone();
    let mut fluid = OversetFluid::build_case(case, cfg.ny, cfg.flag_nx, cfg.sweeps, cfg.max_rounds)
        .expect("overset fluid");
    let dt_fluid = fluid.dt_fluid;
    let dt = dt_fluid * cfg.subcycle as f64;
    let zero_d = vec![0.0; 2 * fluid.interface.wetted.len()];

    // Phase 1: rigid flag to t_release.
    let start = std::time::Instant::now();
    let rigid_steps = (cfg.t_release / dt_fluid).round() as usize;
    for _ in 0..rigid_steps {
        fluid.step().expect("rigid fluid step");
    }
    let (rigid_drag, rigid_lift) = fluid.measure_force();
    println!(
        "  {} OVERSET rigid phase: {rigid_steps} steps (dt {dt_fluid:.3e}) to t = {:.2} s in {:.0} s wall; wall drag {rigid_drag:.2} (rigid-flag reference {:.1}; the overset's own CFD2 at ny = 41: 137.8), lift {rigid_lift:.2}; rounds mean {:.2}",
        case.name,
        cfg.t_release,
        start.elapsed().as_secs_f64(),
        case.rigid_drag_reference,
        fluid.rounds_total.get() as f64 / fluid.correctors_total.get().max(1) as f64
    );

    // The flag: nonlinear Newmark stepper at the coupled dt (as march.rs).
    let mut db = MaterialDatabase::new();
    db.add_material(
        MaterialId(0),
        LinearElastic::new(case.e_s, case.nu_s).with_density(case.rho_s),
        None,
    );
    let analysis = NonlinearDynamicAnalysis::new(
        fluid.mesh.clone(),
        db,
        clamp_left(&fluid.mesh),
        dt,
        1,
        AnalysisConfig::default(),
    )
    .with_total_lagrangian()
    .with_convergence_criteria(ConvergenceCriteria {
        max_iterations: 60,
        ..ConvergenceCriteria::default()
    });
    let flag = RefCell::new(analysis.stepper().unwrap());
    let wetted_dofs: Vec<[usize; 2]> = fluid
        .interface
        .wetted
        .iter()
        .map(|&id| {
            let dofs = flag.borrow().node_dofs(id);
            [dofs[0], dofs[1]]
        })
        .collect();
    let a_dofs = flag.borrow().node_dofs(fluid.a_node);
    let extract = |state: &DynamicState| -> Vec<f64> {
        let mut d = vec![0.0; 2 * wetted_dofs.len()];
        for (k, dofs) in wetted_dofs.iter().enumerate() {
            d[2 * k] = state.displacement[dofs[0]];
            d[2 * k + 1] = state.displacement[dofs[1]];
        }
        d
    };
    let extract_velocity = |state: &DynamicState| -> Vec<f64> {
        let mut v = vec![0.0; 2 * wetted_dofs.len()];
        for (k, dofs) in wetted_dofs.iter().enumerate() {
            v[2 * k] = state.velocity[dofs[0]];
            v[2 * k + 1] = state.velocity[dofs[1]];
        }
        v
    };

    // Phase 2: release under the current load.
    let (nodal0, conservation0, faces0) = fluid.sample_load(&zero_d);
    flag.borrow_mut().set_nodal_forces(&nodal0);
    let mut flag_state = flag.borrow_mut().rest_state().unwrap();
    let mut committed_nodal = nodal0;
    let mut worst_conservation = conservation0;
    println!(
        "  release: {faces0} wall faces transferred (conservation defect {conservation0:.2e}); initial tip acceleration |a| = {:.3e}",
        (a_dofs.iter().map(|&k| flag_state.acceleration[k].powi(2)).sum::<f64>()).sqrt()
    );

    let fluid = RefCell::new(fluid);
    let mut iqn = (cfg.coupler == "iqn").then(|| {
        IqnIls::new(cfg.max_subiterations, 1.0)
            .unwrap()
            .with_reuse(cfg.reuse)
            .with_initial_relaxation(cfg.initial_relaxation)
            .unwrap()
    });

    let coupled_steps = ((cfg.t_end - cfg.t_release) / dt).round() as usize;
    let mut times = Vec::with_capacity(coupled_steps);
    let mut ux_series = Vec::with_capacity(coupled_steps);
    let mut uy_series = Vec::with_capacity(coupled_steps);
    let mut force_times = Vec::new();
    let mut drag_series = Vec::new();
    let mut lift_series = Vec::new();
    let (mut interval_drag, mut interval_lift) = (Vec::new(), Vec::new());
    let mut total_subiterations = 0usize;
    let mut max_subiterations = 0usize;
    let mut stalled_steps = 0usize;
    let mut retried_steps = 0usize;
    let mut worst_stall = 0.0_f64;
    let mut faces_used = faces0;
    let mut death: Option<(usize, f64, String)> = None;
    let mut csv = cfg.csv_path.as_ref().map(|p| {
        let mut f = std::fs::File::create(p).expect("csv");
        writeln!(f, "t,ux,uy,drag,lift").unwrap();
        f
    });
    let (t_fluid, t_structure) = (std::cell::Cell::new(0.0_f64), std::cell::Cell::new(0.0_f64));
    let phase_start = std::time::Instant::now();

    for step in 0..coupled_steps {
        let d_n = extract(&flag_state);
        let v_n: Option<Vec<f64>> = cfg.c1_interface.then(|| extract_velocity(&flag_state));
        let d_predicted = if cfg.predictor == "kinematic" {
            let v = extract_velocity(&flag_state);
            d_n.iter().zip(&v).map(|(d, v)| d + dt * v).collect()
        } else {
            flag.borrow_mut().set_nodal_forces(&committed_nodal);
            let (predicted, _) = flag.borrow_mut().step(&flag_state).unwrap();
            extract(&predicted)
        };
        let saved = fluid.borrow().snapshot();
        type PassResult = (DynamicState, Vec<(NodeId, Vector3<f64>)>, f64, usize);
        let latest: RefCell<Option<PassResult>> = RefCell::new(None);
        let pass = |d_candidate: &[f64]| -> Vec<f64> {
            let fs = std::time::Instant::now();
            let mut fl = fluid.borrow_mut();
            fl.restore(&saved);
            fl.advance_subcycled(&d_n, d_candidate, cfg.subcycle, v_n.as_deref())
                .expect("fluid pass");
            let (nodal, conservation, faces) = fl.sample_load(d_candidate);
            t_fluid.set(t_fluid.get() + fs.elapsed().as_secs_f64());
            let ss = std::time::Instant::now();
            let mut flag_ref = flag.borrow_mut();
            flag_ref.set_nodal_forces(&nodal);
            let (candidate_state, _) = flag_ref.step(&flag_state).unwrap();
            t_structure.set(t_structure.get() + ss.elapsed().as_secs_f64());
            let d_new = extract(&candidate_state);
            if step < cfg.trace_steps {
                let residual: f64 = d_new
                    .iter()
                    .zip(d_candidate)
                    .map(|(a, b)| (a - b) * (a - b))
                    .sum::<f64>()
                    .sqrt();
                let load: f64 = nodal.iter().map(|(_, f)| f.norm()).sum();
                println!(
                    "    step {step} pass: |d_new − d_candidate| = {residual:.3e}, |d_new| = {:.3e}, nodal load {load:.2}, {faces} faces",
                    d_new.iter().map(|v| v * v).sum::<f64>().sqrt()
                );
            }
            *latest.borrow_mut() = Some((candidate_state, nodal, conservation, faces));
            d_new
        };

        let increment: f64 = d_predicted
            .iter()
            .zip(&d_n)
            .map(|(a, b)| (a - b) * (a - b))
            .sum::<f64>()
            .sqrt();
        let tol_step = cfg.tol_floor.max(cfg.rtol * increment);
        let retry_at = (5.0 * tol_step).max(0.1 * increment);
        let acceptable = (cfg.stall_accept * tol_step).max(0.1 * increment);
        let mut outcome = if let Some(iqn) = iqn.as_mut() {
            iqn.set_tolerance(tol_step).unwrap();
            iqn.solve(&d_predicted, pass)
        } else {
            Subiterated::aitken(cfg.max_subiterations, tol_step)
                .unwrap()
                .solve(&d_predicted, pass)
        };
        if let (Err(e), Some(iqn_ref)) = (&outcome, iqn.as_mut()) {
            let recoverable = matches!(
                e,
                rtx_fsi::FsiError::CouplingNotConverged { residual, .. }
                | rtx_fsi::FsiError::CouplingDiverged { residual, .. }
                if *residual >= retry_at
            );
            if recoverable {
                iqn_ref.reset_history();
                retried_steps += 1;
                outcome = iqn_ref.solve(&d_predicted, pass);
            }
        }
        let t_now = cfg.t_release + (step + 1) as f64 * dt;
        match outcome {
            Ok(c) => {
                total_subiterations += c.iterations;
                max_subiterations = max_subiterations.max(c.iterations);
            }
            Err(
                rtx_fsi::FsiError::CouplingNotConverged {
                    iterations,
                    residual,
                    ..
                }
                | rtx_fsi::FsiError::CouplingDiverged {
                    iterations,
                    residual,
                },
            ) if residual < acceptable => {
                stalled_steps += 1;
                worst_stall = worst_stall.max(residual);
                total_subiterations += iterations;
                max_subiterations = max_subiterations.max(iterations);
            }
            Err(e) => {
                println!(
                    "  {} OVERSET DEATH at step {step} t = {t_now:.4}: {e:?} (increment {increment:.3e}, tol {tol_step:.3e}, acceptable {acceptable:.3e})",
                    case.name
                );
                death = Some((step, t_now, format!("{e:?}")));
                break;
            }
        }
        let (new_state, nodal, conservation, faces) = latest.borrow_mut().take().expect("pass ran");
        flag_state = new_state;
        committed_nodal = nodal;
        worst_conservation = worst_conservation.max(conservation);
        faces_used = faces;
        let ux = flag_state.displacement[a_dofs[0]];
        let uy = flag_state.displacement[a_dofs[1]];
        times.push(t_now);
        ux_series.push(ux);
        uy_series.push(uy);
        let (drag_now, lift_now) = fluid.borrow().measure_force();
        interval_drag.push(drag_now);
        interval_lift.push(lift_now);
        if (step + 1) % 10 == 0 {
            let drag = median(&mut interval_drag);
            let lift = median(&mut interval_lift);
            interval_drag.clear();
            interval_lift.clear();
            force_times.push(t_now);
            drag_series.push(drag);
            lift_series.push(lift);
            if let Some(f) = csv.as_mut() {
                writeln!(f, "{t_now:.6},{ux:.6e},{uy:.6e},{drag:.6e},{lift:.6e}").unwrap();
            }
        } else if let Some(f) = csv.as_mut() {
            writeln!(f, "{t_now:.6},{ux:.6e},{uy:.6e},,").unwrap();
        }
        if (step + 1) % 500 == 0 {
            let window = &uy_series[uy_series.len().saturating_sub(500)..];
            let (w_mid, w_amp) = mid_amp(window);
            let fl = fluid.borrow();
            println!(
                "  t = {t_now:.3} s ({} steps): uy(A) = {uy:.3e} (window mid {w_mid:.3e} amp {w_amp:.3e}), drag {drag_now:.1} lift {lift_now:.1}, {:.1} subit/step, rounds mean {:.2}, reclassified/step {:.1}, regen {:.0} s of {:.0} s fluid, {:.0} s wall",
                step + 1,
                total_subiterations as f64 / (step + 1) as f64,
                fl.rounds_total.get() as f64 / fl.correctors_total.get().max(1) as f64,
                fl.reclassified_total.get() as f64 / (rigid_steps + (step + 1) * cfg.subcycle * 4).max(1) as f64,
                fl.regen_seconds.get(),
                t_fluid.get(),
                phase_start.elapsed().as_secs_f64()
            );
        }
    }
    let fl = fluid.borrow();
    let final_state_finite = flag_state.displacement.iter().all(|v| v.is_finite());
    let steps_done = times.len();
    OversetMarchResult {
        result: MarchResult {
            dt,
            coupled_steps: steps_done,
            times,
            ux: ux_series,
            uy: uy_series,
            force_times,
            drag: drag_series,
            lift: lift_series,
            rigid_drag,
            rigid_lift,
            mean_subiterations: total_subiterations as f64 / steps_done.max(1) as f64,
            max_subiterations,
            stalled_steps,
            retried_steps,
            worst_stall,
            worst_conservation,
            skipped: 0,
            spiked: 0,
            newton_rescues: flag.borrow().rescue_counts(),
            coupling_rescues: 0,
            coupling_rescue_failures: 0,
            rescue_records: Vec::new(),
            final_state_finite,
            elapsed: start.elapsed().as_secs_f64(),
        },
        death,
        rounds_mean: fl.rounds_total.get() as f64 / fl.correctors_total.get().max(1) as f64,
        reclassified_mean: fl.reclassified_total.get() as f64
            / (rigid_steps + steps_done * cfg.subcycle).max(1) as f64,
        fresh_mean: fl.fresh_total.get() as f64
            / (rigid_steps + steps_done * cfg.subcycle).max(1) as f64,
        regen_count: fl.regen_count.get(),
        regen_seconds: fl.regen_seconds.get(),
        fluid_seconds: t_fluid.get(),
        structure_seconds: t_structure.get(),
        faces_used,
    }
}