nucl-parquet 0.13.6

//! EADL atomic relaxation data (fluorescence X-rays and Auger electrons).
//!
//! After a photoelectric interaction creates a vacancy in an inner shell,
//! the atom relaxes via radiative (X-ray) or non-radiative (Auger) transitions.
//! This module provides the transition probabilities and energies needed to
//! sample the resulting secondary particles.

use std::collections::HashMap;
use std::fs;
use std::path::Path;

use arrow::array::{Float64Array, Int32Array};
use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder;

/// Type of atomic relaxation transition.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum TransitionType {
    /// Radiative transition — emits a characteristic X-ray photon.
    Radiative,
    /// Auger transition — emits an electron.
    Auger,
}

/// A single atomic relaxation transition.
#[derive(Debug, Clone)]
pub struct Transition {
    /// Shell where the vacancy was created (e.g., "K", "L1").
    pub vacancy_shell: String,
    /// Shell that fills the vacancy (e.g., "L3").
    pub filling_shell: String,
    /// Radiative (X-ray) or Auger (electron).
    pub transition_type: TransitionType,
    /// Transition energy in keV.
    pub energy_kev: f64,
    /// Transition probability (fractional, sums to ~1 per shell).
    pub probability: f64,
    /// Binding energy of the vacancy shell in keV.
    pub edge_kev: f64,
}

/// EADL atomic relaxation database.
///
/// Thread-safe: `Send + Sync`.
pub struct RelaxationDb {
    /// Z -> list of transitions (sorted by vacancy_shell, then probability desc)
    transitions: HashMap<u8, Vec<Transition>>,
}

impl RelaxationDb {
    /// Load EADL data from the nucl-parquet `meta/` directory.
    ///
    /// Reads `meta/eadl/*.parquet`.
    pub fn open(meta_dir: impl AsRef<Path>) -> crate::Result<Self> {
        let dir = meta_dir.as_ref().join("eadl");
        let mut transitions: HashMap<u8, Vec<Transition>> = HashMap::new();

        if !dir.exists() {
            return Ok(Self { transitions });
        }

        for entry in fs::read_dir(&dir)? {
            let entry = entry?;
            let path = entry.path();
            if path.extension().and_then(|e| e.to_str()) != Some("parquet") {
                continue;
            }

            let file = fs::File::open(&path)?;
            if let Some((z, trans_list)) = Self::parse_one(file)? {
                transitions.insert(z, trans_list);
            }
        }

        Ok(Self { transitions })
    }

    /// Construct from a single element's EADL Parquet bytes.
    ///
    /// The element Z is determined from the `Z` column in the data.
    pub fn from_bytes(data: &[u8]) -> crate::Result<Self> {
        let bytes = bytes::Bytes::from(data.to_vec());
        let mut transitions: HashMap<u8, Vec<Transition>> = HashMap::new();
        if let Some((z, trans_list)) = Self::parse_one(bytes)? {
            transitions.insert(z, trans_list);
        }
        Ok(Self { transitions })
    }

    /// Parse a single parquet source into (Z, transitions).
    fn parse_one(
        reader_source: impl parquet::file::reader::ChunkReader + 'static,
    ) -> crate::Result<Option<(u8, Vec<Transition>)>> {
        let reader = ParquetRecordBatchReaderBuilder::try_new(reader_source)?.build()?;

        let mut z_val: Option<u8> = None;
        let mut trans_list = Vec::new();

        for batch in reader {
            let batch = batch?;

            let z_col = batch
                .column_by_name("Z")
                .and_then(|c| c.as_any().downcast_ref::<Int32Array>());
            let vac_col_ref = batch.column_by_name("vacancy_shell");
            let vac_values = vac_col_ref.and_then(|c| crate::interp::as_string_array(c));
            let fill_col_ref = batch.column_by_name("filling_shell");
            let fill_values = fill_col_ref.and_then(|c| crate::interp::as_string_array(c));
            let type_col_ref = batch.column_by_name("transition_type");
            let type_values = type_col_ref.and_then(|c| crate::interp::as_string_array(c));
            let energy_col = batch
                .column_by_name("energy_keV")
                .and_then(|c| c.as_any().downcast_ref::<Float64Array>());
            let prob_col = batch
                .column_by_name("probability")
                .and_then(|c| c.as_any().downcast_ref::<Float64Array>());
            let edge_col = batch
                .column_by_name("edge_keV")
                .and_then(|c| c.as_any().downcast_ref::<Float64Array>());

            if let (Some(z), Some(vac), Some(fill), Some(tt), Some(e), Some(p), Some(edge)) = (
                z_col,
                vac_values,
                fill_values,
                type_values,
                energy_col,
                prob_col,
                edge_col,
            ) {
                for i in 0..batch.num_rows() {
                    if z_val.is_none() {
                        z_val = Some(z.value(i) as u8);
                    }
                    trans_list.push(Transition {
                        vacancy_shell: vac[i].unwrap_or("").to_string(),
                        filling_shell: fill[i].unwrap_or("").to_string(),
                        transition_type: match tt[i].unwrap_or("") {
                            "radiative" => TransitionType::Radiative,
                            _ => TransitionType::Auger,
                        },
                        energy_kev: e.value(i),
                        probability: p.value(i),
                        edge_kev: edge.value(i),
                    });
                }
            }
        }

        match z_val {
            Some(z) => {
                // Sort by shell, then probability descending
                trans_list.sort_by(|a, b| {
                    a.vacancy_shell
                        .cmp(&b.vacancy_shell)
                        .then(b.probability.partial_cmp(&a.probability).unwrap())
                });
                Ok(Some((z, trans_list)))
            }
            None => Ok(None),
        }
    }

    /// Get all transitions for element Z.
    pub fn transitions(&self, z: u8) -> &[Transition] {
        self.transitions
            .get(&z)
            .map(|v| v.as_slice())
            .unwrap_or(&[])
    }

    /// Get transitions for a specific vacancy shell (e.g., "K").
    pub fn shell_transitions(&self, z: u8, shell: &str) -> Vec<&Transition> {
        self.transitions(z)
            .iter()
            .filter(|t| t.vacancy_shell == shell)
            .collect()
    }

    /// Get only radiative (X-ray) transitions for element Z and shell.
    pub fn radiative_transitions(&self, z: u8, shell: &str) -> Vec<&Transition> {
        self.transitions(z)
            .iter()
            .filter(|t| t.vacancy_shell == shell && t.transition_type == TransitionType::Radiative)
            .collect()
    }

    /// Fluorescence yield for a shell = sum of radiative transition probabilities.
    pub fn fluorescence_yield(&self, z: u8, shell: &str) -> f64 {
        self.transitions(z)
            .iter()
            .filter(|t| t.vacancy_shell == shell && t.transition_type == TransitionType::Radiative)
            .map(|t| t.probability)
            .sum()
    }

    /// Check if data is loaded for element Z.
    pub fn has_element(&self, z: u8) -> bool {
        self.transitions.contains_key(&z)
    }
}

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

    fn data_meta_dir() -> std::path::PathBuf {
        std::path::Path::new(env!("CARGO_MANIFEST_DIR"))
            .join("..")
            .join("..")
            .join("..")
            .join("data")
            .join("meta")
    }

    #[test]
    #[ignore = "requires nucl-parquet data files"]
    fn from_bytes_matches_open() {
        let db_file = RelaxationDb::open(data_meta_dir()).unwrap();
        assert!(db_file.has_element(29));
        // Read the first valid EADL file and load via from_bytes
        let eadl_dir = data_meta_dir().join("eadl");
        let first_file = std::fs::read_dir(&eadl_dir)
            .unwrap()
            .filter_map(|e| e.ok())
            .find(|e| {
                e.path().extension().and_then(|x| x.to_str()) == Some("parquet")
                    && !e
                        .path()
                        .file_name()
                        .and_then(|n| n.to_str())
                        .is_some_and(|n| n.starts_with("._"))
            })
            .expect("at least one EADL file");
        let data = std::fs::read(first_file.path()).unwrap();
        let db_bytes = RelaxationDb::from_bytes(&data).unwrap();
        // Find what Z the bytes DB loaded and verify transitions match the file DB
        for z in 1..=100u8 {
            if db_bytes.has_element(z) {
                let t_file = db_file.transitions(z);
                let t_bytes = db_bytes.transitions(z);
                assert_eq!(
                    t_file.len(),
                    t_bytes.len(),
                    "Z={z} transition count mismatch"
                );
                break;
            }
        }
    }

    #[test]
    #[ignore = "requires nucl-parquet data files"]
    fn load_and_query_cu() {
        let db = RelaxationDb::open("../../meta").unwrap();
        assert!(db.has_element(29));

        let k_trans = db.shell_transitions(29, "K");
        assert!(!k_trans.is_empty());

        // K fluorescence yield for Cu should be ~0.44
        let fy = db.fluorescence_yield(29, "K");
        assert!(fy > 0.3 && fy < 0.6, "Cu K fluorescence yield: {fy}");

        // K-alpha line should be ~8 keV
        let radiative = db.radiative_transitions(29, "K");
        let strongest = radiative
            .iter()
            .max_by(|a, b| a.probability.partial_cmp(&b.probability).unwrap());
        assert!(strongest.is_some());
        let ka = strongest.unwrap();
        assert!(
            ka.energy_kev > 7.0 && ka.energy_kev < 9.0,
            "Cu K-alpha: {} keV",
            ka.energy_kev
        );
    }
}