nereids_endf/
retrieval.rs

1//! ENDF file download and local caching.
2//!
3//! Downloads ENDF files from the IAEA Nuclear Data Services and caches them
4//! locally for offline use. Follows the URL patterns established by PLEIADES.
5//!
6//! ## PLEIADES Reference
7//! - `pleiades/nuclear/manager.py` — URL construction, cache directory layout
8//! - `pleiades/nuclear/models.py` — library enum, filename patterns
9
10use nereids_core::elements;
11use nereids_core::types::Isotope;
12use std::fs;
13use std::io::Read;
14use std::path::{Path, PathBuf};
15
16/// ENDF evaluated nuclear data libraries.
17#[derive(Debug, Clone, Copy, PartialEq, Eq)]
18pub enum EndfLibrary {
19    /// ENDF/B-VIII.0 (default, well-tested).
20    EndfB8_0,
21    /// ENDF/B-VIII.1 (latest release, Aug 2024).
22    EndfB8_1,
23    /// JEFF-3.3 (European library).
24    Jeff3_3,
25    /// JENDL-5 (Japanese library).
26    Jendl5,
27}
28
29impl EndfLibrary {
30    /// URL path component for this library.
31    fn url_path(&self) -> &'static str {
32        match self {
33            Self::EndfB8_0 => "ENDF-B-VIII.0/n",
34            Self::EndfB8_1 => "ENDF-B-VIII.1/n",
35            Self::Jeff3_3 => "JEFF-3.3/n",
36            Self::Jendl5 => "JENDL-5/n",
37        }
38    }
39
40    /// Cache directory name.
41    fn cache_dir_name(&self) -> &'static str {
42        match self {
43            Self::EndfB8_0 => "ENDF-B-VIII.0",
44            Self::EndfB8_1 => "ENDF-B-VIII.1",
45            Self::Jeff3_3 => "JEFF-3.3",
46            Self::Jendl5 => "JENDL-5",
47        }
48    }
49
50    /// Construct the ZIP filename for a given isotope.
51    ///
52    /// IAEA uses two naming conventions (MAT always 4-digit zero-padded):
53    /// - VIII.0, JEFF-3.3: MAT-first `n_{mat:04}_{z}-{Sym}-{a}.zip` (Z unpadded)
54    /// - VIII.1, JENDL-5: Z-first   `n_{z:03}-{Sym}-{a}_{mat:04}.zip` (Z 3-digit)
55    fn zip_filename(&self, isotope: &Isotope, mat: u32) -> String {
56        let sym = elements::element_symbol(isotope.z()).unwrap_or("X");
57        let z = isotope.z();
58        let a = isotope.a();
59        match self {
60            Self::EndfB8_0 | Self::Jeff3_3 => {
61                format!("n_{mat:04}_{z}-{sym}-{a}.zip")
62            }
63            Self::EndfB8_1 | Self::Jendl5 => {
64                format!("n_{z:03}-{sym}-{a}_{mat:04}.zip")
65            }
66        }
67    }
68}
69
70/// ENDF file retrieval manager with local caching.
71pub struct EndfRetriever {
72    /// Root cache directory.
73    cache_root: PathBuf,
74    /// Base URL for IAEA downloads.
75    base_url: String,
76}
77
78impl EndfRetriever {
79    /// Create a new retriever with default cache location (~/.cache/nereids/endf/).
80    pub fn new() -> Self {
81        let cache_root = dirs::cache_dir()
82            .unwrap_or_else(|| PathBuf::from(".cache"))
83            .join("nereids")
84            .join("endf");
85        Self {
86            cache_root,
87            base_url: "https://www-nds.iaea.org/public/download-endf".to_string(),
88        }
89    }
90
91    /// Create a retriever with a custom cache directory.
92    pub fn with_cache_dir(cache_dir: impl Into<PathBuf>) -> Self {
93        Self {
94            cache_root: cache_dir.into(),
95            base_url: "https://www-nds.iaea.org/public/download-endf".to_string(),
96        }
97    }
98
99    /// Get the cache directory for a specific library.
100    fn cache_dir(&self, library: EndfLibrary) -> PathBuf {
101        self.cache_root.join(library.cache_dir_name())
102    }
103
104    /// Get the cached ENDF file path for an isotope.
105    fn cache_file_path(&self, isotope: &Isotope, library: EndfLibrary) -> PathBuf {
106        let sym = elements::element_symbol(isotope.z()).unwrap_or("X");
107        let filename = format!("{}-{}.endf", sym, isotope.a());
108        self.cache_dir(library).join(filename)
109    }
110
111    /// Retrieve the ENDF file for an isotope, using cache if available.
112    ///
113    /// Returns the path to the cached ENDF file and its contents as a string.
114    ///
115    /// # Arguments
116    /// * `isotope` — The isotope to retrieve data for.
117    /// * `library` — The ENDF library to use.
118    /// * `mat` — The ENDF MAT (material) number.
119    pub fn get_endf_file(
120        &self,
121        isotope: &Isotope,
122        library: EndfLibrary,
123        mat: u32,
124    ) -> Result<(PathBuf, String), EndfRetrievalError> {
125        let cache_path = self.cache_file_path(isotope, library);
126
127        // Check cache first.
128        if cache_path.exists() {
129            let contents = fs::read_to_string(&cache_path)?;
130            return Ok((cache_path, contents));
131        }
132
133        // Download from IAEA.
134        let contents = self.download_endf(isotope, library, mat)?;
135
136        // Cache the file.
137        if let Some(parent) = cache_path.parent() {
138            fs::create_dir_all(parent)?;
139        }
140        fs::write(&cache_path, &contents)?;
141
142        Ok((cache_path, contents))
143    }
144
145    /// Download ENDF file from IAEA and extract from ZIP archive.
146    fn download_endf(
147        &self,
148        isotope: &Isotope,
149        library: EndfLibrary,
150        mat: u32,
151    ) -> Result<String, EndfRetrievalError> {
152        let zip_filename = library.zip_filename(isotope, mat);
153        let url = format!("{}/{}/{}", self.base_url, library.url_path(), zip_filename);
154
155        let response = reqwest::blocking::get(&url).map_err(|e| {
156            EndfRetrievalError::NetworkError(format!("Failed to connect to {}: {}", url, e))
157        })?;
158
159        let status = response.status();
160        if status == reqwest::StatusCode::NOT_FOUND {
161            return Err(EndfRetrievalError::NotInLibrary {
162                isotope: format!(
163                    "{}-{}",
164                    nereids_core::elements::element_symbol(isotope.z()).unwrap_or("?"),
165                    isotope.a()
166                ),
167                library: library.cache_dir_name().to_string(),
168            });
169        }
170        if !status.is_success() {
171            return Err(EndfRetrievalError::NetworkError(format!(
172                "HTTP {} for {}",
173                status, url
174            )));
175        }
176
177        let bytes = response.bytes().map_err(|e| {
178            EndfRetrievalError::NetworkError(format!("Failed to read response body: {}", e))
179        })?;
180
181        // Extract ENDF file from ZIP archive.
182        self.extract_endf_from_zip(&bytes)
183    }
184
185    /// Extract the ENDF data file from a ZIP archive.
186    ///
187    /// Looks for files ending in .endf, .dat, or .txt within the archive.
188    fn extract_endf_from_zip(&self, zip_bytes: &[u8]) -> Result<String, EndfRetrievalError> {
189        let cursor = std::io::Cursor::new(zip_bytes);
190        let mut archive = zip::ZipArchive::new(cursor)
191            .map_err(|e| EndfRetrievalError::Parse(format!("Invalid ZIP archive: {}", e)))?;
192
193        // Find the ENDF data file within the archive.
194        for i in 0..archive.len() {
195            let mut file = archive.by_index(i).map_err(|e| {
196                EndfRetrievalError::Parse(format!("Failed to read ZIP entry: {}", e))
197            })?;
198
199            let name = file.name().to_lowercase();
200            if name.ends_with(".endf") || name.ends_with(".dat") || name.ends_with(".txt") {
201                let mut contents = String::new();
202                file.read_to_string(&mut contents).map_err(|e| {
203                    EndfRetrievalError::Parse(format!("Failed to read ENDF content: {}", e))
204                })?;
205                return Ok(contents);
206            }
207        }
208
209        // If no obvious extension, try the first file.
210        if !archive.is_empty() {
211            let mut file = archive.by_index(0).map_err(|e| {
212                EndfRetrievalError::Parse(format!("Failed to read ZIP entry: {}", e))
213            })?;
214            let mut contents = String::new();
215            file.read_to_string(&mut contents).map_err(|e| {
216                EndfRetrievalError::Parse(format!("Failed to read ENDF content: {}", e))
217            })?;
218            return Ok(contents);
219        }
220
221        Err(EndfRetrievalError::Parse(
222            "No ENDF data file found in ZIP archive".to_string(),
223        ))
224    }
225
226    /// Load an ENDF file from a local path (no download).
227    pub fn load_local(path: &Path) -> Result<String, EndfRetrievalError> {
228        fs::read_to_string(path).map_err(EndfRetrievalError::from)
229    }
230
231    /// Clear the cache for a specific library, or all if `None`.
232    pub fn clear_cache(&self, library: Option<EndfLibrary>) -> Result<(), EndfRetrievalError> {
233        match library {
234            Some(lib) => {
235                let dir = self.cache_dir(lib);
236                if dir.exists() {
237                    fs::remove_dir_all(&dir)?;
238                }
239            }
240            None => {
241                if self.cache_root.exists() {
242                    fs::remove_dir_all(&self.cache_root)?;
243                }
244            }
245        }
246        Ok(())
247    }
248}
249
250impl Default for EndfRetriever {
251    fn default() -> Self {
252        Self::new()
253    }
254}
255
256/// Look up the ENDF MAT number for a ground-state isotope.
257///
258/// Covers 535 isotopes from the ENDF/B-VIII.0 neutrons sublibrary
259/// (via the `endf-mat` crate). This replaces the previous hand-coded
260/// 47-entry table and fixes several incorrect MAT values (Cd-113,
261/// Hf-177/178, W-182/183/184/186 were off by one isotope offset).
262pub fn mat_number(isotope: &Isotope) -> Option<u32> {
263    endf_mat::mat_number(isotope.z(), isotope.a())
264}
265
266/// Errors from ENDF retrieval operations.
267#[derive(Debug, thiserror::Error)]
268pub enum EndfRetrievalError {
269    /// Transport-level failure (connection refused, DNS error, non-404 HTTP error, etc.).
270    #[error("Network error: {0}")]
271    NetworkError(String),
272
273    /// The isotope exists in ENDF/B-VIII.0 but is not available in the requested library.
274    #[error("{isotope} is not available in the {library} library")]
275    NotInLibrary { isotope: String, library: String },
276
277    #[error("Parse error: {0}")]
278    Parse(String),
279
280    #[error("I/O error: {0}")]
281    Io(#[from] std::io::Error),
282
283    #[error("Isotope not found in MAT database: {0}")]
284    UnknownIsotope(String),
285}