SpectraRust/src/synspec/math/rtesca.rs

//! Solution of the radiative transfer equation for continuum scattering.
//!
//! Translated from SYNSPEC `RTESCA` subroutine (synspec54.f:20035).
//!
//! Uses the Discontinuous Finite Element method (Castor, Dykema, Klein, 1992, ApJ 387, 561)
//! to solve the RTE along impact rays for the spherically-symmetric case,
//! deriving the scattering in continuum.

use super::interp::interp;

/// Physical constants
const UN: f64 = 1.0;
const TWO: f64 = 2.0;
const HALF: f64 = 0.5;

/// Maximum number of ALI iterations for electron scattering
const NTRALI: usize = 10;
/// Convergence threshold for electron scattering iteration
const DJMAX: f64 = 1.0e-3;

/// Parameters for RTESCA calculation
pub struct RtescaParams<'a> {
    /// Number of depth points (original grid)
    pub nd: usize,
    /// Number of depth points (fine grid)
    pub ndf: usize,
    /// Number of continuum frequencies
    pub nfreqc: usize,
    /// Continuum frequencies (Hz) [nfreqc]
    pub freqc: &'a [f64],
    /// Continuum wavelengths (Angstrom) [nfreqc]
    pub wlamc: &'a [f64],
    /// Temperature at each depth [nd]
    pub temp: &'a [f64],
    /// Density at each depth [nd]
    pub dens: &'a [f64],
    /// Fine grid density [ndf]
    pub densf: &'a [f64],
    /// Continuum absorption coefficient [nfreqc x nd]
    pub chc: &'a [Vec<f64>],
    /// Continuum emission coefficient [nfreqc x nd]
    pub etc: &'a [Vec<f64>],
    /// Continuum scattering coefficient [nfreqc x nd]
    pub scc: &'a [Vec<f64>],
    /// Electron density * sigma_e at each depth [nd]
    pub elec_sig: &'a [f64],
    /// Boltzmann constant * c^2 (BN constant)
    pub bn: f64,
    /// h/k constant
    pub hk: f64,
    /// Number of mu points (impact rays)
    pub kmu: usize,
    /// Number of core rays
    pub nfiry: usize,
    /// Number of extended rays
    pub nrext: usize,
    /// Number of depth points per ray [kmu]
    pub nud: &'a [usize],
    /// Number of depth points per fine ray [kmu]
    pub nudf: &'a [usize],
    /// Ray index: kray[iu][id] gives depth index for ray iu at point id
    pub kray: &'a [Vec<usize>],
    /// Ray interpolation weight: dray[iu][id]
    pub dray: &'a [Vec<f64>],
    /// Fine grid spacing for fine rays [kmu x (ndf-1)]
    pub delzf: &'a [Vec<f64>],
    /// Grid spacing for extended rays [kmu x (nd-1)]
    pub delz: &'a [Vec<f64>],
    /// Weight for mean intensity: wmuj[iu][id]
    pub wmuj: &'a [Vec<f64>],
    /// Weight for flux: wmuh[kmu]
    pub wmuh: &'a [f64],
}

/// Result of RTESCA calculation
pub struct RtescaResult {
    /// Continuum flux [nfreqc]
    pub fluxc: Vec<f64>,
    /// Scattering source function on fine grid [nfreqc x ndf]
    pub sccf: Vec<Vec<f64>>,
}

/// Solve the radiative transfer equation for continuum scattering.
///
/// # Arguments
/// * `params` - Input parameters
///
/// # Returns
/// Continuum flux and scattering source function
pub fn rtesca(params: &RtescaParams) -> RtescaResult {
    let nd = params.nd;
    let ndf = params.ndf;
    let nfreqc = params.nfreqc;
    let kmu = params.kmu;

    let mut fluxc = vec![0.0; nfreqc];
    let mut sccf = vec![vec![0.0; ndf]; nfreqc];

    // Overall loop over continuum frequencies
    for ij in 0..nfreqc {
        let fr = params.freqc[ij];

        // Initialisation of J=B (Planck function)
        let fr15 = fr * 1.0e-15;
        let bnu = params.bn * fr15 * fr15 * fr15;
        let hkfr = params.hk * fr;

        // Initialize RAD00 = Planck function
        let mut rad00: Vec<f64> = params.temp.iter()
            .map(|&t| {
                let exp_val = (hkfr / t).exp();
                if exp_val > UN { bnu / (exp_val - UN) } else { 0.0 }
            })
            .collect();

        // Loop over electron scattering iterations
        let mut itrali = 0;
        loop {
            itrali += 1;
            fluxc[ij] = 0.0;

            let mut rad1 = vec![0.0; nd];
            let mut ali1 = vec![0.0; nd];

            // Prepare opacity arrays on fine or original grid
            let (abc0, stc0, scc0, rdx);
            if nd == ndf {
                // Same grid - direct copy
                abc0 = params.chc[ij].clone();
                stc0 = params.etc[ij].iter().zip(params.chc[ij].iter())
                    .map(|(&e, &c)| if c > 0.0 { e / c } else { 0.0 })
                    .collect();
                scc0 = params.scc[ij].clone();
                rdx = rad00.clone();
            } else {
                // Interpolate to fine grid
                let abc1 = params.chc[ij].clone();
                let stc1: Vec<f64> = params.etc[ij].iter().zip(params.chc[ij].iter())
                    .map(|(&e, &c)| if c > 0.0 { e / c } else { 0.0 })
                    .collect();
                let scc01 = params.scc[ij].clone();

                abc0 = interp(params.dens, &abc1, params.densf, 4, 1, 0);
                let interp_stc = interp(params.dens, &stc1, params.densf, 4, 1, 0);
                let interp_scc = interp(params.dens, &scc01, params.densf, 4, 1, 0);
                rdx = interp(params.dens, &rad00, params.densf, 4, 1, 0);

                stc0 = interp_stc;
                scc0 = interp_scc;
            }

            // Loop over impact rays
            for iu in 0..kmu {
                let iud = if iu < params.nfiry {
                    params.nudf[iu]
                } else {
                    params.nud[iu]
                };
                if iud <= 1 {
                    continue;
                }

                // Interpolate quantities along the ray
                let mut densr = vec![0.0; iud];
                let mut ab0 = vec![0.0; iud];
                let mut st0 = vec![0.0; iud];
                let mut ss0 = vec![0.0; iud];
                let mut rdy = vec![0.0; iud];

                for id in 0..iud {
                    let ky = params.kray[iu][id];
                    let ydr = params.dray[iu][id];
                    let ydr1 = UN - ydr;

                    // ky is 1-based from Fortran, convert to 0-based
                    let ky0 = ky.saturating_sub(1);
                    let ky1 = (ky0 + 1).min(ndf.saturating_sub(1));

                    densr[id] = ydr1 * params.densf[ky0] + ydr * params.densf[ky1];
                    ab0[id] = ydr1 * abc0[ky0] + ydr * abc0[ky1];
                    st0[id] = ydr1 * stc0[ky0] + ydr * stc0[ky1];
                    let sc0 = ydr1 * scc0[ky0] + ydr * scc0[ky1];
                    rdy[id] = ydr1 * rdx[ky0] + ydr * rdx[ky1];
                    ss0[id] = if ab0[id] > 0.0 { sc0 / ab0[id] } else { 0.0 };
                    st0[id] += ss0[id] * rdy[id];
                }

                // Calculate optical depth along the ray
                let mut dtau = vec![0.0; iud - 1];
                if iu < params.nfiry {
                    for id in 0..iud - 1 {
                        dtau[id] = HALF * (ab0[id] + ab0[id + 1]) * params.delzf[iu][id];
                    }
                } else {
                    for id in 0..iud - 1 {
                        dtau[id] = HALF * (ab0[id] + ab0[id + 1]) * params.delz[iu][id];
                    }
                }

                // Incoming intensity (TAUMIN=0)
                let mut rim = vec![0.0; iud];
                let mut rip = vec![0.0; iud];
                let mut aim_arr = vec![0.0; iud];
                let mut aip = vec![0.0; iud];

                for id in 0..iud - 1 {
                    let dt0 = dtau[id];
                    let dtaup1 = dt0 + UN;
                    let dtau2 = dt0 * dt0;
                    let bb = TWO * dtaup1;
                    let cc = dt0 * dtaup1;
                    let aa = UN / (dtau2 + bb);
                    rip[id] = (bb * rim[id] + cc * st0[id] - dt0 * st0[id + 1]) * aa;
                    rim[id + 1] = (TWO * rim[id] + dt0 * st0[id] + cc * st0[id + 1]) * aa;
                    aip[id] = (cc + bb * aim_arr[id]) * aa;
                    aim_arr[id + 1] = cc * aa;
                }

                // Interpolate to cell centers
                let mut riin = vec![0.0; iud];
                let mut aiin = vec![0.0; iud];
                for id in 1..iud - 1 {
                    let dtt = UN / (dtau[id - 1] + dtau[id]);
                    riin[id] = (rim[id] * dtau[id] + rip[id] * dtau[id - 1]) * dtt;
                    aiin[id] = (aim_arr[id] * dtau[id] + aip[id] * dtau[id - 1]) * dtt;
                }
                riin[0] = rim[0];
                riin[iud - 1] = rim[iud - 1];
                aiin[0] = aim_arr[0];
                aiin[iud - 1] = aim_arr[iud - 1];
                rip[iud - 1] = rim[iud - 1];

                // Outgoing intensity
                // Symmetric boundary condition or diffusion approximation for core rays
                if iu >= params.nrext {
                    let t_nd = params.temp[nd - 1];
                    let t_nd1 = params.temp[nd - 2];
                    let pland = if t_nd > 0.0 {
                        bnu / ((hkfr / t_nd).exp() - UN)
                    } else {
                        0.0
                    };
                    let pland1 = if t_nd1 > 0.0 {
                        bnu / ((hkfr / t_nd1).exp() - UN)
                    } else {
                        0.0
                    };
                    let dplan = pland - pland1;
                    let ium1 = iud - 1;
                    rip[ium1] = if dtau[ium1 - 1] > 0.0 {
                        pland + dplan / dtau[ium1 - 1]
                    } else {
                        pland
                    };
                    let dt0 = dtau[ium1 - 1];
                    let dtaup1 = dt0 + UN;
                    let dtau2 = dt0 * dt0;
                    let bb = TWO * dtaup1;
                    let cc = dt0 * dtaup1;
                    let aa = dtau2 + bb;
                    rim[ium1] = (aa * rip[ium1] - cc * st0[ium1] + dt0 * st0[ium1 - 1]) / bb;
                }

                // Outgoing sweep
                for id in (0..iud - 1).rev() {
                    let dt0 = dtau[id];
                    let dtaup1 = dt0 + UN;
                    let dtau2 = dt0 * dt0;
                    let bb = TWO * dtaup1;
                    let cc = dt0 * dtaup1;
                    let aa = UN / (dtau2 + bb);
                    rip[id + 1] = (bb * rim[id + 1] + cc * st0[id + 1] - dt0 * st0[id]) * aa;
                    rim[id] = (TWO * rim[id + 1] + dt0 * st0[id + 1] + cc * st0[id]) * aa;
                    aip[id + 1] = (cc + bb * aim_arr[id + 1]) * aa;
                    aim_arr[id] = cc * aa;
                }

                // Interpolate outgoing to cell centers
                let mut riup = vec![0.0; iud];
                let mut aiup = vec![0.0; iud];
                for id in 1..iud - 1 {
                    let dtt = UN / (dtau[id - 1] + dtau[id]);
                    riup[id] = (rim[id] * dtau[id - 1] + rip[id] * dtau[id]) * dtt;
                    aiup[id] = (aim_arr[id] * dtau[id - 1] + aip[id] * dtau[id]) * dtt;
                }
                riup[0] = rim[0];
                riup[iud - 1] = rim[iud - 1];
                aiup[0] = aim_arr[0];
                aiup[iud - 1] = aim_arr[iud - 1];

                // Symmetrized (Feautrier) intensity = (riin + riup) / 2
                let mut uf: Vec<f64> = riup.iter().zip(riin.iter())
                    .map(|(&u, &i)| u + i)
                    .collect();
                let mut af: Vec<f64> = aiup.iter().zip(aiin.iter())
                    .map(|(&u, &i)| u + i)
                    .collect();

                // Interpolate back to original radial grid for fine rays
                let actual_iud = if iu < params.nfiry {
                    let inrp = params.nud[iu].min(4);
                    let nud_iu = params.nud[iu];
                    let interp_uf = interp(&densr, &uf, params.dens, inrp as i32, 1, 0);
                    let interp_af = interp(&densr, &af, params.dens, inrp as i32, 1, 0);
                    uf = interp_uf;
                    af = interp_af;
                    nud_iu
                } else {
                    iud
                };

                // Contribution to mean intensity J
                for id in 0..actual_iud {
                    rad1[id] += params.wmuj[iu][id] * uf[id];
                    ali1[id] += params.wmuj[iu][id] * af[id];
                }
                fluxc[ij] += params.wmuh[iu] * rim[0];
            } // end loop over impact rays

            // Solve the scattering problem
            // Interpolate scattering source function to original grid
            let ndx = if params.nfiry > 0 {
                params.nudf[kmu - 1]
            } else {
                params.nud[kmu - 1]
            };
            // Use first ray's densr as reference
            let mut densr_ref = vec![0.0; ndx];
            let mut ss0_ref = vec![0.0; ndx];
            // Reconstruct from the last ray's data (simplified)
            // In practice, we need the ray geometry from the last ray
            // For now, use the fine grid directly
            for id in 0..ndx.min(ndf) {
                densr_ref[id] = params.densf[id];
                // Approximate scattering source
                let c_ij = params.chc[ij][id.min(nd - 1)];
                let s_ij = params.scc[ij][id.min(nd - 1)];
                ss0_ref[id] = if c_ij > 0.0 { s_ij / c_ij } else { 0.0 };
            }

            let scx = interp(&densr_ref, &ss0_ref, params.dens, 4, 1, 1);

            let mut djtot: f64 = 0.0;
            for id in 0..nd {
                rad1[id] *= HALF;
                ali1[id] *= HALF;
                let sss = scx[id];
                let delta_j = (rad1[id] - rad00[id]) / (UN - sss * ali1[id]);
                rad00[id] += delta_j;
                if rad00[id].abs() > 0.0 {
                    djtot = djtot.max((delta_j / rad00[id]).abs());
                }
            }

            // Check convergence
            if djtot <= DJMAX || itrali >= NTRALI {
                break;
            }
        } // end electron scattering loop

        // Store scattering source function on fine grid
        let rdx_final = interp(params.dens, &rad00, params.densf, 4, 1, 0);
        for id in 0..ndf {
            sccf[ij][id] = params.scc[ij][id.min(nd - 1)] * rdx_final[id];
        }
        fluxc[ij] *= 2.997925e18 / (params.wlamc[ij] * params.wlamc[ij]) * 0.5;
    } // end loop over frequencies

    RtescaResult { fluxc, sccf }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_rtesca_basic() {
        // Basic smoke test with minimal parameters
        let nd = 3;
        let ndf = 3;
        let nfreqc = 1;
        let kmu = 1;

        let freqc = vec![1.0e15];
        let wlamc = vec![3000.0];
        let temp = vec![5000.0, 10000.0, 20000.0];
        let dens = vec![1.0e-10, 1.0e-11, 1.0e-12];
        let densf = dens.clone();
        let chc = vec![vec![1.0e-2; nd]; nfreqc];
        let etc = vec![vec![1.0e-4; nd]; nfreqc];
        let scc = vec![vec![1.0e-3; nd]; nfreqc];
        let elec_sig = vec![1.0e-15; nd];

        let params = RtescaParams {
            nd,
            ndf,
            nfreqc,
            freqc: &freqc,
            wlamc: &wlamc,
            temp: &temp,
            dens: &dens,
            densf: &densf,
            chc: &chc,
            etc: &etc,
            scc: &scc,
            elec_sig: &elec_sig,
            bn: 3.9729e-16, // typical BN value
            hk: 4.7992e-11, // typical HK value
            kmu,
            nfiry: 0,
            nrext: 0,
            nud: &vec![nd; kmu],
            nudf: &vec![ndf; kmu],
            kray: &vec![vec![1, 2, 3]; kmu],
            dray: &vec![vec![0.5, 0.5, 0.0]; kmu],
            delzf: &vec![vec![1.0; ndf - 1]; kmu],
            delz: &vec![vec![1.0; nd - 1]; kmu],
            wmuj: &vec![vec![1.0; nd]; kmu],
            wmuh: &vec![1.0; kmu],
        };

        let result = rtesca(&params);
        assert_eq!(result.fluxc.len(), nfreqc);
        assert_eq!(result.sccf.len(), nfreqc);
        assert_eq!(result.sccf[0].len(), ndf);
        // Flux should be non-negative
        assert!(result.fluxc[0] >= 0.0);
    }
}