SpectraRust/src/synspec/math/iniset.rs

//! Line selection and frequency grid setup for SYNSPEC.
//!
//! Translated from SYNSPEC `INISET` subroutine (synspec54.f:8074).
//!
//! Selection of lines that may contribute to the opacity,
//! set up auxiliary fields containing line parameters,
//! and set up the set of frequency points.

/// Physical constant: speed of light in Angstrom·Hz
const CNM: f64 = 2.997925e17;
/// Physical constant: speed of light in nm·Hz
const CAS: f64 = 2.997925e18;

/// Parameters for INISET calculation.
pub struct InisetParams<'a> {
    /// Mode flag
    pub imode: i32,
    /// Blanketing flag
    pub iblank: i32,
    /// Frequency window flag (>0 means window mode)
    pub ifwin: i32,
    /// Frequency index for window mode
    pub ifreq: i32,
    /// Starting wavelength (nm)
    pub alam0: f64,
    /// Ending wavelength (nm)
    pub alam1: f64,
    /// Last frequency (Hz)
    pub frlast: f64,
    /// Maximum velocity (cm/s) for window mode
    pub vinf: f64,
    /// Spacing parameter
    pub space0: f64,
    /// Cutoff parameter
    pub cutof0: f64,
    /// Standard temperature (K)
    pub tstd: f64,
    /// Standard Doppler width parameter
    pub dstd: f64,
    /// Central wavelength for spacing (nm)
    pub alamc: f64,
    /// Previous wavelength (nm)
    pub aprev: f64,
    /// Previous wavelength from last set (nm)
    pub alm00: f64,
    /// Maximum frequency (Hz)
    pub frmax: f64,
    /// Standard depth absorption
    pub abstd_idstd: f64,
    /// Relative opacity threshold
    pub relop: f64,
    /// Number of lines in full list
    pub nlin0: i32,
    /// Maximum number of lines in a set
    pub mlin: i32,
    /// Number of frequency points in frequency grid
    pub nfreqs: i32,
    /// Number of molecular line lists
    pub nmlist: i32,
    /// Molecular lines flag
    pub ifmol: i32,
    /// Line frequencies [nlin0]
    pub freq0: &'a [f64],
    /// Line extinction parameters [nlin0]
    pub extin: &'a [f64],
    /// Line profile flags [nlin0]
    pub isprf: &'a [i32],
    /// Line set indices [nlin0]
    pub indlip: &'a [i32],
    /// Last molecular line wavelength per list [nmlist]
    pub alastm: &'a [f64],
    /// Center frequencies for window mode [nfreqs]
    pub freqc: &'a [f64],
    /// Wavelengths for window mode [nfreqs]
    pub wlamc: &'a [f64],
    /// Last index of lines from previous set
    pub illast: i32,
}

/// Result of INISET calculation.
pub struct InisetResult {
    /// Number of selected lines
    pub nlin: i32,
    /// Selected line indices [mlin]
    pub indlin: Vec<i32>,
    /// Number of frequency points
    pub nfreq: i32,
    /// Frequency array [nfreqs]
    pub freq: Vec<f64>,
    /// Weight array [nfreqs]
    pub w: Vec<f64>,
    /// Wavelength array [nfreqs]
    pub wlam: Vec<f64>,
    /// Frequency interpolation coefficient 1 [nfreqs]
    pub frx1: Vec<f64>,
    /// Frequency interpolation coefficient 2 [nfreqs]
    pub frx2: Vec<f64>,
    /// Line center frequency indices [mlin]
    pub ijcntr: Vec<i32>,
    /// Blanketing flag for next iteration
    pub nblank: i32,
    /// Whether molecular lines extend the interval
    pub irlist: i32,
    /// Updated alam0 for next iteration
    pub alam0_next: f64,
    /// Updated alm00 for next iteration
    pub alm00_next: f64,
    /// Updated aprev for next iteration
    pub aprev_next: f64,
    /// Minimum frequency (Hz)
    pub frmin: f64,
    /// Last index of lines from this set
    pub illast_next: i32,
    /// Observed frequencies for window mode [nfreqs]
    pub frqobs: Vec<f64>,
    /// Observed wavelengths for window mode [nfreqs]
    pub wlobs: Vec<f64>,
    /// Planck function for window mode [nfreqs]
    pub bnue: Vec<f64>,
    /// Center frequency indices for window mode [nfreqs]
    pub ijcint: Vec<i32>,
}

/// Select lines and set up frequency grid.
///
/// Translates SYNSPEC INISET (synspec54.f:8074).
///
/// # Arguments
/// * `params` - Input parameters
///
/// # Returns
/// Frequency grid and line selection results
pub fn iniset(params: &InisetParams) -> InisetResult {
    let nfreqs = params.nfreqs as usize;
    let mlin = params.mlin as usize;

    // Initialize output arrays
    let mut freq = vec![0.0; nfreqs + 1];
    let mut w = vec![0.0; nfreqs + 1];
    let mut wlam = vec![0.0; nfreqs + 1];
    let mut frx1 = vec![0.0; nfreqs + 1];
    let mut frx2 = vec![0.0; nfreqs + 1];
    let mut indlin = vec![0i32; mlin + 1];
    let mut ijcntr = vec![0i32; mlin + 1];
    let mut frqobs = vec![0.0; nfreqs + 1];
    let mut wlobs = vec![0.0; nfreqs + 1];
    let mut bnue = vec![0.0; nfreqs + 1];
    let mut ijcint = vec![0i32; nfreqs + 1];

    let mut nlin: i32 = 0;
    let mut nblank = params.iblank + 1;
    let mut irlist = 0;

    // Calculate minimum frequency from starting wavelength
    let mut frmin = CNM / params.alam0;
    let mut frm = frmin;

    // Determine starting frequency index
    let ij0: usize = if params.ifwin <= 0 { 3 } else { 1 };
    let mut ij = ij0;
    freq[ij0] = frm;

    // Calculate spacing
    let mut space = params.space0;
    if params.alamc > 0.0 {
        space = params.space0 * params.alam0 / params.alamc;
    }
    if params.space0 < 0.0 {
        space = -params.space0;
    }

    // IMODE=2 special case
    if params.imode == 2 {
        let nfrp = (params.nfreqs + 1) as usize;
        let w0 = space;
        // Jump to frequency point setup (label 105)
        let fract = freq[ij];
        let mut alact = CNM / fract;
        for _k in 0..nfrp {
            alact += w0;
            ij += 1;
            if ij > nfreqs {
                break;
            }
            freq[ij] = CNM / alact;
            if ij > 1 {
                w[ij] += (freq[ij - 1] - freq[ij]) * 0.5;
                w[ij - 1] += (freq[ij - 1] - freq[ij]) * 0.5;
            }
        }
    } else {
        // Main line selection loop
        let mut il0: i32 = 0;
        let mut iprset: i32 = 0;
        let mut ireadp = if params.iblank <= 1
            || params.imode == 1
            || params.imode == -1
        {
            0
        } else {
            1
        };

        // Calculate cutoff and Doppler parameters
        let (cutoff, dopstd, distan, spac, dista0, _astd, _avab) =
            if params.ifwin <= 0 {
                let cutoff = params.cutof0;
                let dopstd = 1e7 / params.alam0 * params.dstd;
                let distan = 0.15 * dopstd;
                let spac = 3e16 / params.alam0 / params.alam0 * space;
                let dista0 = 0.14 * spac;
                (cutoff, dopstd, distan, spac, dista0, 1.0, params.abstd_idstd * params.relop)
            } else {
                (params.cutof0, 0.0, 0.0, space, 0.0, 1.0, 0.0)
            };

        if params.iblank >= 2 && params.imode == -1 {
            il0 = params.illast;
        }

        // Main loop over lines
        loop {
            // Set up line index
            if ireadp == 1 {
                iprset += 1;
                let idx = iprset as usize;
                if idx <= params.indlip.len() {
                    il0 = params.indlip[idx - 1];
                }
                if il0 as usize <= params.freq0.len()
                    && params.freq0[il0 as usize - 1] < frmin
                {
                    ireadp = 0;
                    il0 = if iprset > 1 {
                        params.indlip[(iprset - 2) as usize] + 1
                    } else {
                        1
                    };
                }
            } else {
                il0 += 1;
            }

            if il0 > params.nlin0 {
                break;
            }

            let fr0 = if il0 as usize <= params.freq0.len() {
                params.freq0[il0 as usize - 1]
            } else {
                break;
            };
            let alam = CNM / fr0;

            // Window mode spacing adjustment
            let (cutoff, _dopstd, _distan, spac, dista0) = if params.ifwin > 0 {
                let mut space_adj = space;
                if params.alamc > 0.0 {
                    space_adj = params.space0 * alam / params.alamc;
                }
                if params.space0 < 0.0 {
                    space_adj = -params.space0;
                }
                let cutoff = params.cutof0 * alam / params.alamc;
                let dopstd = 1e7 / alam * params.dstd;
                let distan = 0.15 * dopstd;
                let spac = if params.ifreq % 10 > 0 {
                    3e16 / alam / alam * space_adj
                } else {
                    space_adj
                };
                let dista0 = 0.14 * spac;
                (cutoff, dopstd, distan, spac, dista0)
            } else {
                (cutoff, dopstd, distan, spac, dista0)
            };

            // IMODE=1: adjust starting wavelength
            if params.imode == 1 && nlin == 0 && ij == 3
                && alam >= params.alam0 + 2.0 * cutoff {
                    // Update alam0 and frmin
                    let alam0_new = alam - cutoff + 0.0001;
                    frmin = CNM / alam0_new;
                    frm = frmin;
                    ij = ij0;
                    freq[ij0] = frm;
                }

            // First selection: wavelength range
            if alam < params.alam0 - cutoff {
                continue;
            }
            if ij < (params.nfreqs + 1) as usize {
                // Continue to second selection
            } else if alam > params.alam1 + cutoff {
                break;
            }

            // Second selection: line strengths
            let mut _istr = 0;
            if params.imode >= 1 {
                _istr = 1;
            } else {
                let ext = if il0 as usize <= params.extin.len() {
                    params.extin[il0 as usize - 1]
                } else {
                    0.0
                };
                let frli0_new = fr0 - ext - spac;

                let frmiv = if params.ifwin > 0 {
                    frmin * (1.0 + params.vinf / 2.997925e10)
                } else {
                    frmin
                };
                if alam < params.alam0 && fr0 - frmiv > ext + spac {
                    continue;
                }
                _istr = 1;

                let frmav = if params.ifwin > 0 {
                    params.frmax * (1.0 - params.vinf / 2.997925e10)
                } else {
                    params.frmax
                };
                if ij >= (params.nfreqs + 1) as usize && frmav - fr0 > ext + spac {
                    continue;
                }

                let _ = frli0_new; // Used for FRLI0 update
            }

            // Select line
            nlin += 1;
            if nlin > params.mlin {
                break; // Too many lines
            }
            indlin[nlin as usize] = il0;
            let _alamcu = alam + cutoff;

            // Frequency points and weights
            if ij >= (params.nfreqs + 1) as usize {
                continue;
            }
            if fr0 > frmin {
                continue;
            }

            let delt = (frm - fr0).abs();
            if delt < dista0 && params.imode != 1 {
                continue;
            }

            let dfrel = CNM * (1.0 / fr0 - 1.0 / frm) / space;
            let mut nfrp = (dfrel as i32) + 1;
            if nfrp <= 2 {
                nfrp = 2;
            }
            let w0 = CNM * (1.0 / fr0 - 1.0 / frm) / nfrp as f64;
            frm = fr0;

            // Generate frequency points
            let mut fract = freq[ij];
            let mut alact = CNM / fract;

            for _k in 0..nfrp {
                fract -= w0;
                alact += w0;

                if params.imode < 1 && nfrp != 2 {
                    let frli0_check = fr0 - spac;
                    if fract < frli0_check && fract > fr0 + spac {
                        continue;
                    }
                }

                ij += 1;
                if ij > nfreqs {
                    break;
                }
                freq[ij] = CNM / alact;
                if ij > 1 {
                    w[ij] += (freq[ij - 1] - freq[ij]) * 0.5;
                    w[ij - 1] += (freq[ij - 1] - freq[ij]) * 0.5;
                }
            }

            if ij <= nfreqs {
                ijcntr[nlin as usize] = ij as i32;
            }
        }

        // Truncate interval if needed
        let ijmx = if params.ifwin > 0 { ij } else { 2 };
        if freq[ijmx] < params.frlast {
            freq[ijmx] = params.frlast;
            if params.ifwin <= 0 && ij > 1 {
                w[1] = 0.5 * (freq[1] - freq[2]);
                w[2] = w[1];
            }
            // Find IJMAX
            let mut ijmax = ij.min(nfreqs);
            for k in ij0..=ij.min(nfreqs) {
                if freq[k] < params.frlast {
                    ijmax = k;
                }
            }
            let nfreq_new = ijmax + 1;
            if nfreq_new <= nfreqs {
                freq[nfreq_new] = params.frlast;
                if nfreq_new > 1 {
                    w[nfreq_new] = 0.5 * (freq[nfreq_new - 1] - freq[nfreq_new]);
                }
                if nfreq_new > 2 {
                    w[nfreq_new - 1] = w[nfreq_new]
                        + 0.5 * (freq[nfreq_new - 2] - freq[nfreq_new - 1]);
                }
            }
        }
    }

    // Calculate frequency interpolation coefficients
    let nfreq_actual = if params.imode != -1 {
        if params.ifwin <= 0 {
            let xx = if freq.len() > 2 { freq[2] - freq[1] } else { 1.0 };
            for k in 1..=nfreqs {
                if freq[k] != 0.0 {
                    wlam[k] = CAS / freq[k];
                }
                if xx != 0.0 {
                    frx1[k] = (freq[k] - freq[1]) / xx;
                    frx2[k] = (freq[2] - freq[k]) / xx;
                }
            }
        } else {
            for k in 1..=nfreqs {
                if freq[k] != 0.0 {
                    wlam[k] = CAS / freq[k];
                    frqobs[k] = freq[k];
                    wlobs[k] = wlam[k];
                    let fr = freq[k];
                    bnue[k] = 1.47450e-47 * fr * fr * fr; // BN * fr^3

                    // Find center frequency index
                    let mut ijc = 1;
                    for ijc_inner in 1..params.freqc.len() {
                        if wlam[k] <= params.wlamc[ijc_inner - 1] {
                            ijc = ijc_inner;
                            break;
                        }
                    }
                    ijcint[k] = (ijc as i32 - 1).max(1);
                    let ijci = ijcint[k] as usize;
                    if ijci + 1 < params.freqc.len()
                        && params.freqc[ijci] != params.freqc[ijci + 1]
                    {
                        frx1[k] = (freq[k] - params.freqc[ijci + 1])
                            / (params.freqc[ijci] - params.freqc[ijci + 1]);
                    }
                }
            }
        }
        nfreqs as i32
    } else {
        nfreqs as i32
    };

    // Calculate frequency indices of line centers
    if params.imode != -1 && nlin > 0 {
        let xx = if freq.len() > 2 { freq[2] - freq[1] } else { 1.0 };
        if xx != 0.0 {
            let dfrcon = (nfreq_actual - ij0 as i32) as f64;
            let dfrcon = -dfrcon / xx;
            for il in 1..=nlin as usize {
                let il_idx = indlin[il] as usize;
                if il_idx > 0 && il_idx <= params.freq0.len() {
                    let fr0 = params.freq0[il_idx - 1];
                    let xjc = 3.0 + dfrcon * (freq[1] - fr0);
                    let mut ijc = xjc as i32;
                    if ijc > ij0 as i32 && ijc < nfreq_actual {
                        // Find closest frequency
                        if fr0 < freq[ijc as usize] {
                            let mut ijc0 = ijc;
                            let mut dfr0 = freq[ijc0 as usize] - fr0;
                            loop {
                                ijc0 += 1;
                                if ijc0 as usize >= freq.len() {
                                    break;
                                }
                                let dfr = (freq[ijc0 as usize] - fr0).abs();
                                if dfr < dfr0 {
                                    ijc = ijc0;
                                    dfr0 = dfr;
                                } else {
                                    break;
                                }
                            }
                        } else if fr0 > freq[ijc as usize] {
                            let mut ijc0 = ijc;
                            let mut dfr0 = fr0 - freq[ijc0 as usize];
                            loop {
                                ijc0 -= 1;
                                if ijc0 < 1 {
                                    break;
                                }
                                let dfr = (freq[ijc0 as usize] - fr0).abs();
                                if dfr < dfr0 {
                                    ijc = ijc0;
                                    dfr0 = dfr;
                                } else {
                                    break;
                                }
                            }
                        }
                        ijcntr[il] = ijc;
                    }
                }
            }
        }
    }

    // Update blanketing flag
    let nfreq_out = if nfreq_actual > 0 && (nfreq_actual as usize) <= nfreqs {
        nfreq_actual
    } else {
        nfreqs as i32
    };

    if freq.len() > nfreq_out as usize && freq[nfreq_out as usize] <= params.frlast {
        nblank = params.iblank;
    }

    // Molecular line correction
    if params.nmlist > 0 && params.ifmol > 0 {
        for ilist in 0..params.nmlist as usize {
            if ilist < params.alastm.len()
                && params.alastm[ilist] > 0.0
                && params.alastm[ilist] <= params.alam1
            {
                nblank = params.iblank;
                irlist = 1;
            }
        }
    }

    // Update illast
    let illast_next = if nlin > 0 { indlin[nlin as usize] } else { 0 };

    InisetResult {
        nlin,
        indlin,
        nfreq: nfreq_out,
        freq: freq.clone(),
        w,
        wlam,
        frx1,
        frx2,
        ijcntr,
        nblank,
        irlist,
        alam0_next: params.alam1,
        alm00_next: if nfreq_out > 0 && (nfreq_out as usize) <= nfreqs {
            CNM / freq[nfreq_out as usize]
        } else {
            0.0
        },
        aprev_next: params.alam0,
        frmin,
        illast_next,
        frqobs,
        wlobs,
        bnue,
        ijcint,
    }
}

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

    #[test]
    fn test_iniset_basic() {
        let freq0 = vec![1e15, 1.1e15, 1.2e15];
        let extin = vec![1e10, 1e10, 1e10];
        let isprf = vec![0, 0, 0];
        let indlip = vec![1, 2, 3];
        let alastm = vec![];
        let freqc = vec![1e15, 1.2e15];
        let wlamc = vec![2997.925, 2498.271];

        let params = InisetParams {
            imode: 0,
            iblank: 0,
            ifwin: 0,
            ifreq: 0,
            alam0: 200.0,
            alam1: 300.0,
            frlast: 1e15,
            vinf: 0.0,
            space0: 0.5,
            cutof0: 100.0,
            tstd: 10000.0,
            dstd: 2.0,
            alamc: 0.0,
            aprev: 0.0,
            alm00: 0.0,
            frmax: 1.5e15,
            abstd_idstd: 1.0,
            relop: 0.01,
            nlin0: 3,
            mlin: 100,
            nfreqs: 100,
            nmlist: 0,
            ifmol: 0,
            freq0: &freq0,
            extin: &extin,
            isprf: &isprf,
            indlip: &indlip,
            alastm: &alastm,
            freqc: &freqc,
            wlamc: &wlamc,
            illast: 0,
        };

        let result = iniset(&params);
        assert!(result.nlin >= 0);
        assert!(result.nfreq > 0);
    }
}