1use std::fmt::Write as _;
9use std::path::Path;
10
11use ndarray::Array1;
12
13use crate::error::{IoError, Result};
14use crate::format::{FortranField, write_fortran_row};
15
16const CHI_DAT_STANDARD_ROW_WIDTH: usize = 4;
17const CHI_DAT_PATH_ROW_WIDTH: usize = 5;
18const CHI_DAT_CKP_ROW_WIDTH: usize = 6;
19const CHI_DAT_ALLOWED_ROW_WIDTHS: &str = "4, 5, or 6";
20
21const CHI_ROW_K: FortranField = FortranField::F {
23 width: 11,
24 precision: 4,
25};
26const CHI_ROW_VALUE: FortranField = FortranField::E {
29 width: 13,
30 precision: 6,
31};
32
33#[derive(Debug, Clone, PartialEq)]
35pub struct ChiDatData {
36 pub header_lines: Vec<String>,
38 pub wave_number: Array1<f64>,
40 pub chi: Array1<f64>,
42 pub magnitude: Array1<f64>,
44 pub phase: Array1<f64>,
46 pub phase_minus_2kr: Option<Array1<f64>>,
48 pub ckp_real: Option<Array1<f64>>,
50 pub ckp_imag: Option<Array1<f64>>,
52}
53
54impl ChiDatData {
55 #[must_use]
57 pub fn point_count(&self) -> usize {
58 self.wave_number.len()
59 }
60
61 #[must_use]
63 pub fn has_path_phase(&self) -> bool {
64 self.phase_minus_2kr.is_some()
65 }
66
67 #[must_use]
69 pub fn has_complex_wave_number(&self) -> bool {
70 self.ckp_real.is_some() && self.ckp_imag.is_some()
71 }
72}
73
74pub fn chi_dat_string(data: &ChiDatData) -> Result<String> {
76 validate_chi_dat(data)?;
77
78 let mut out = String::new();
79 for line in &data.header_lines {
80 writeln!(out, "{line}")?;
81 }
82
83 match (&data.phase_minus_2kr, &data.ckp_real, &data.ckp_imag) {
84 (None, None, None) => {
85 for (((k, chi), magnitude), phase) in data
86 .wave_number
87 .iter()
88 .zip(data.chi.iter())
89 .zip(data.magnitude.iter())
90 .zip(data.phase.iter())
91 {
92 write_chi_row(&mut out, *k, [*chi, *magnitude, *phase])?;
93 }
94 }
95 (Some(phase_minus_2kr), None, None) => {
96 for ((((k, chi), magnitude), phase), path_phase) in data
97 .wave_number
98 .iter()
99 .zip(data.chi.iter())
100 .zip(data.magnitude.iter())
101 .zip(data.phase.iter())
102 .zip(phase_minus_2kr.iter())
103 {
104 write_chi_row(&mut out, *k, [*chi, *magnitude, *phase, *path_phase])?;
105 }
106 }
107 (None, Some(ckp_real), Some(ckp_imag)) => {
108 for (((((k, chi), magnitude), phase), ckp_real), ckp_imag) in data
109 .wave_number
110 .iter()
111 .zip(data.chi.iter())
112 .zip(data.magnitude.iter())
113 .zip(data.phase.iter())
114 .zip(ckp_real.iter())
115 .zip(ckp_imag.iter())
116 {
117 write_chi_row(
118 &mut out,
119 *k,
120 [*chi, *magnitude, *phase, *ckp_real, *ckp_imag],
121 )?;
122 }
123 }
124 _ => {
125 return Err(invalid_chi_dat(
126 "optional columns",
127 "unsupported column combination",
128 ));
129 }
130 }
131
132 Ok(out)
133}
134
135fn write_chi_row<const N: usize>(
136 out: &mut String,
137 wave_number: f64,
138 fields: [f64; N],
139) -> Result<()> {
140 CHI_ROW_K.write(out, wave_number)?;
141 out.push_str(" ");
142 write_fortran_row(
143 out,
144 " ",
145 fields.into_iter().map(|value| (CHI_ROW_VALUE, value)),
146 )?;
147 out.push('\n');
148 Ok(())
149}
150
151pub fn parse_chi_dat(text: &str) -> Result<ChiDatData> {
153 let mut header_lines = Vec::new();
154 let mut row_width = None;
155 let mut wave_number = Vec::new();
156 let mut chi = Vec::new();
157 let mut magnitude = Vec::new();
158 let mut phase = Vec::new();
159 let mut phase_minus_2kr = Vec::new();
160 let mut ckp_real = Vec::new();
161 let mut ckp_imag = Vec::new();
162
163 for (index, raw) in text.lines().enumerate() {
164 let line_number = index + 1;
165 let line = raw.trim_end();
166 let tokens = line.split_whitespace().collect::<Vec<_>>();
167 if tokens.first().is_some_and(|token| is_numeric_token(token)) {
168 let width = tokens.len();
169 if !matches!(
170 width,
171 CHI_DAT_STANDARD_ROW_WIDTH | CHI_DAT_PATH_ROW_WIDTH | CHI_DAT_CKP_ROW_WIDTH
172 ) {
173 return Err(IoError::ChiDatRowWidth {
174 line: line_number,
175 actual: width,
176 expected: CHI_DAT_ALLOWED_ROW_WIDTHS,
177 });
178 }
179 if let Some(expected) = row_width {
180 if width != expected {
181 return Err(IoError::ChiDatRowWidth {
182 line: line_number,
183 actual: width,
184 expected: row_width_label(expected),
185 });
186 }
187 } else {
188 row_width = Some(width);
189 }
190
191 wave_number.push(parse_f64(line_number, "wave number", tokens[0])?);
192 chi.push(parse_f64(line_number, "chi", tokens[1])?);
193 magnitude.push(parse_f64(line_number, "magnitude", tokens[2])?);
194 phase.push(parse_f64(line_number, "phase", tokens[3])?);
195 if width == CHI_DAT_PATH_ROW_WIDTH {
196 phase_minus_2kr.push(parse_f64(line_number, "phase minus 2kr", tokens[4])?);
197 }
198 if width == CHI_DAT_CKP_ROW_WIDTH {
199 ckp_real.push(parse_f64(line_number, "ckp real", tokens[4])?);
200 ckp_imag.push(parse_f64(line_number, "ckp imaginary", tokens[5])?);
201 }
202 } else {
203 header_lines.push(raw.to_string());
204 }
205 }
206
207 let data = ChiDatData {
208 header_lines,
209 wave_number: Array1::from_vec(wave_number),
210 chi: Array1::from_vec(chi),
211 magnitude: Array1::from_vec(magnitude),
212 phase: Array1::from_vec(phase),
213 phase_minus_2kr: (row_width == Some(CHI_DAT_PATH_ROW_WIDTH))
214 .then(|| Array1::from_vec(phase_minus_2kr)),
215 ckp_real: (row_width == Some(CHI_DAT_CKP_ROW_WIDTH)).then(|| Array1::from_vec(ckp_real)),
216 ckp_imag: (row_width == Some(CHI_DAT_CKP_ROW_WIDTH)).then(|| Array1::from_vec(ckp_imag)),
217 };
218 validate_chi_dat(&data)?;
219 Ok(data)
220}
221
222pub fn write_chi_dat(path: impl AsRef<Path>, data: &ChiDatData) -> Result<()> {
224 let path = path.as_ref();
225 std::fs::write(path, chi_dat_string(data)?).map_err(|source| IoError::io(path, source))
226}
227
228pub fn read_chi_dat(path: impl AsRef<Path>) -> Result<ChiDatData> {
230 let path = path.as_ref();
231 let text = std::fs::read_to_string(path).map_err(|source| IoError::io(path, source))?;
232 parse_chi_dat(&text)
233}
234
235pub(crate) fn validate_chi_dat(data: &ChiDatData) -> Result<()> {
236 let point_count = data.point_count();
237 if point_count == 0 {
238 return Err(invalid_chi_dat(
239 "rows",
240 "at least one spectrum row is required",
241 ));
242 }
243 validate_len("chi", data.chi.len(), point_count)?;
244 validate_len("magnitude", data.magnitude.len(), point_count)?;
245 validate_len("phase", data.phase.len(), point_count)?;
246
247 match (&data.phase_minus_2kr, &data.ckp_real, &data.ckp_imag) {
248 (None, None, None) => {}
249 (Some(phase_minus_2kr), None, None) => {
250 validate_len("phase_minus_2kr", phase_minus_2kr.len(), point_count)?;
251 }
252 (None, Some(ckp_real), Some(ckp_imag)) => {
253 validate_len("ckp_real", ckp_real.len(), point_count)?;
254 validate_len("ckp_imag", ckp_imag.len(), point_count)?;
255 }
256 _ => {
257 return Err(invalid_chi_dat(
258 "optional columns",
259 "use either phase_minus_2kr, ckp_real with ckp_imag, or no optional columns",
260 ));
261 }
262 }
263
264 for (row, (((k, chi), magnitude), phase)) in data
265 .wave_number
266 .iter()
267 .zip(data.chi.iter())
268 .zip(data.magnitude.iter())
269 .zip(data.phase.iter())
270 .enumerate()
271 {
272 let row = row + 1;
273 validate_finite_row("wave number", *k, row)?;
274 validate_finite_row("chi", *chi, row)?;
275 validate_finite_row("magnitude", *magnitude, row)?;
276 validate_finite_row("phase", *phase, row)?;
277 }
278 if let Some(phase_minus_2kr) = &data.phase_minus_2kr {
279 for (row, value) in phase_minus_2kr.iter().enumerate() {
280 validate_finite_row("phase minus 2kr", *value, row + 1)?;
281 }
282 }
283 if let Some(ckp_real) = &data.ckp_real {
284 for (row, value) in ckp_real.iter().enumerate() {
285 validate_finite_row("ckp real", *value, row + 1)?;
286 }
287 }
288 if let Some(ckp_imag) = &data.ckp_imag {
289 for (row, value) in ckp_imag.iter().enumerate() {
290 validate_finite_row("ckp imaginary", *value, row + 1)?;
291 }
292 }
293
294 Ok(())
295}
296
297fn validate_len(field: &'static str, actual: usize, expected: usize) -> Result<()> {
298 if actual == expected {
299 Ok(())
300 } else {
301 Err(IoError::ChiDatShape {
302 field,
303 actual,
304 expected,
305 })
306 }
307}
308
309fn parse_f64(line: usize, field: &'static str, token: &str) -> Result<f64> {
310 token
311 .replace(['D', 'd'], "E")
312 .parse::<f64>()
313 .map_err(|_| IoError::ChiDatParse {
314 field,
315 line,
316 token: token.to_string(),
317 })
318}
319
320fn validate_finite_row(field: &'static str, value: f64, row: usize) -> Result<()> {
321 if value.is_finite() {
322 Ok(())
323 } else {
324 Err(IoError::InvalidChiDat {
325 field,
326 message: format!("row {row} value must be finite"),
327 })
328 }
329}
330
331fn invalid_chi_dat(field: &'static str, message: impl Into<String>) -> IoError {
332 IoError::InvalidChiDat {
333 field,
334 message: message.into(),
335 }
336}
337
338fn is_numeric_token(token: &str) -> bool {
339 token.replace(['D', 'd'], "E").parse::<f64>().is_ok()
340}
341
342fn row_width_label(width: usize) -> &'static str {
343 match width {
344 CHI_DAT_STANDARD_ROW_WIDTH => "4",
345 CHI_DAT_PATH_ROW_WIDTH => "5",
346 CHI_DAT_CKP_ROW_WIDTH => "6",
347 _ => CHI_DAT_ALLOWED_ROW_WIDTHS,
348 }
349}
350
351#[cfg(test)]
352mod tests {
353 use super::*;
354
355 #[test]
356 fn parses_feff_chi_reference_shape() -> Result<()> {
357 let data = parse_chi_dat(CHI_DAT)?;
358 assert_eq!(data.point_count(), 3);
359 assert!(!data.has_path_phase());
360 assert!(!data.has_complex_wave_number());
361 assert_eq!(data.wave_number[0], 0.0);
362 assert_eq!(data.chi[1], -1.194138e-1);
363 assert_eq!(data.magnitude[2], 2.750836e-1);
364 assert_eq!(data.phase[0], -2.698164);
365 Ok(())
366 }
367
368 #[test]
369 fn parses_per_path_phase_column() -> Result<()> {
370 let data = parse_chi_dat(CHIP_DAT)?;
371 assert_eq!(data.point_count(), 2);
372 assert!(data.has_path_phase());
373 assert_eq!(
374 data.phase_minus_2kr
375 .as_ref()
376 .ok_or_else(|| invalid_chi_dat("phase_minus_2kr", "missing optional column"))?[1],
377 2.5
378 );
379 Ok(())
380 }
381
382 #[test]
383 fn parses_diagnostic_ckp_columns() -> Result<()> {
384 let data = parse_chi_dat(CHI_CKP_DAT)?;
385 assert_eq!(data.point_count(), 2);
386 assert!(data.has_complex_wave_number());
387 assert_eq!(
388 data.ckp_real
389 .as_ref()
390 .ok_or_else(|| invalid_chi_dat("ckp_real", "missing optional column"))?[0],
391 1.25
392 );
393 assert_eq!(
394 data.ckp_imag
395 .as_ref()
396 .ok_or_else(|| invalid_chi_dat("ckp_imag", "missing optional column"))?[1],
397 -0.0625
398 );
399 Ok(())
400 }
401
402 #[test]
403 fn roundtrips_chi_text() -> Result<()> {
404 let data = parse_chi_dat(CHI_DAT)?;
405 let rendered = chi_dat_string(&data)?;
406 assert_eq!(rendered, CHI_DAT);
407 assert_eq!(parse_chi_dat(&rendered)?, data);
408
409 let chip = parse_chi_dat(CHIP_DAT)?;
410 assert_eq!(chi_dat_string(&chip)?, CHIP_DAT);
411 let ckp = parse_chi_dat(CHI_CKP_DAT)?;
412 assert_eq!(chi_dat_string(&ckp)?, CHI_CKP_DAT);
413 Ok(())
414 }
415
416 #[test]
417 fn rejects_bad_chi_inputs() {
418 assert!(parse_chi_dat("# no data\n").is_err());
419 assert!(parse_chi_dat("1 2 3\n").is_err());
420 assert!(parse_chi_dat("1 2 3 4 5 6 7\n").is_err());
421 assert!(parse_chi_dat("1 2 3 NaN\n").is_err());
422 assert!(parse_chi_dat("1 2 3 4\n2 3 4 5 6\n").is_err());
423 }
424
425 const CHI_DAT: &str = r#"# # Cu FEFF 10.0
426# 0/ 0 paths used
427# -----------------------------------------------------------------------
428# k chi mag phase @#
429 0.0000 -1.159383E-01 2.702278E-01 -2.698164E+00
430 0.0500 -1.194138E-01 2.726708E-01 -2.688285E+00
431 0.1000 -1.229126E-01 2.750836E-01 -2.678386E+00
432"#;
433
434 const CHIP_DAT: &str = r#"# path contribution
435 -----------------------------------------------------------------------
436 k chi mag phase phase-2kr @#
437 0.0000 1.000000E-01 2.000000E-01 1.000000E+00 1.500000E+00
438 0.0500 1.250000E-01 2.250000E-01 2.000000E+00 2.500000E+00
439"#;
440
441 const CHI_CKP_DAT: &str = r#"# diagnostic ckp
442# k chi mag phase @#
443 0.0000 1.000000E-01 2.000000E-01 1.000000E+00 1.250000E+00 -1.250000E-01
444 0.0500 1.250000E-01 2.250000E-01 2.000000E+00 1.500000E+00 -6.250000E-02
445"#;
446
447 mod proptests {
455 use super::*;
456 use crate::format::fortran_exp;
457 use proptest::prelude::*;
458
459 fn snap_fixed(value: f64, precision: usize) -> f64 {
460 format!("{value:.precision$}")
461 .parse::<f64>()
462 .unwrap_or(value)
463 }
464
465 fn snap_exp(value: f64) -> f64 {
466 fortran_exp(value, 13, 6)
467 .trim()
468 .parse::<f64>()
469 .unwrap_or(value)
470 }
471
472 fn wave_number_strategy() -> impl Strategy<Value = f64> {
473 (-999_999_i64..999_999).prop_map(|n| snap_fixed(n as f64 / 10_000.0, 4))
474 }
475
476 fn value_strategy() -> impl Strategy<Value = f64> {
477 (-9.999e6_f64..9.999e6).prop_map(snap_exp)
478 }
479
480 fn row_strategy() -> impl Strategy<Value = (f64, f64, f64, f64)> {
481 (
482 wave_number_strategy(),
483 value_strategy(),
484 value_strategy(),
485 value_strategy(),
486 )
487 }
488
489 proptest! {
490 #[test]
491 fn roundtrips_standard_rows(
492 rows in prop::collection::vec(row_strategy(), 1..6),
493 ) {
494 let data = ChiDatData {
495 header_lines: vec!["# proptest standard header".to_string()],
496 wave_number: Array1::from_iter(rows.iter().map(|row| row.0)),
497 chi: Array1::from_iter(rows.iter().map(|row| row.1)),
498 magnitude: Array1::from_iter(rows.iter().map(|row| row.2)),
499 phase: Array1::from_iter(rows.iter().map(|row| row.3)),
500 phase_minus_2kr: None,
501 ckp_real: None,
502 ckp_imag: None,
503 };
504 let rendered = chi_dat_string(&data)?;
505 let reparsed = parse_chi_dat(&rendered)?;
506 prop_assert_eq!(reparsed, data);
507 }
508
509 #[test]
510 fn roundtrips_path_phase_rows(
511 rows in prop::collection::vec(
512 (row_strategy(), value_strategy()),
513 1..6,
514 ),
515 ) {
516 let data = ChiDatData {
517 header_lines: vec!["# proptest path header".to_string()],
518 wave_number: Array1::from_iter(rows.iter().map(|(row, _)| row.0)),
519 chi: Array1::from_iter(rows.iter().map(|(row, _)| row.1)),
520 magnitude: Array1::from_iter(rows.iter().map(|(row, _)| row.2)),
521 phase: Array1::from_iter(rows.iter().map(|(row, _)| row.3)),
522 phase_minus_2kr: Some(Array1::from_iter(rows.iter().map(|(_, p)| *p))),
523 ckp_real: None,
524 ckp_imag: None,
525 };
526 let rendered = chi_dat_string(&data)?;
527 let reparsed = parse_chi_dat(&rendered)?;
528 prop_assert_eq!(reparsed, data);
529 }
530
531 #[test]
532 fn roundtrips_ckp_rows(
533 rows in prop::collection::vec(
534 (row_strategy(), value_strategy(), value_strategy()),
535 1..6,
536 ),
537 ) {
538 let data = ChiDatData {
539 header_lines: vec!["# proptest ckp header".to_string()],
540 wave_number: Array1::from_iter(rows.iter().map(|(row, ..)| row.0)),
541 chi: Array1::from_iter(rows.iter().map(|(row, ..)| row.1)),
542 magnitude: Array1::from_iter(rows.iter().map(|(row, ..)| row.2)),
543 phase: Array1::from_iter(rows.iter().map(|(row, ..)| row.3)),
544 phase_minus_2kr: None,
545 ckp_real: Some(Array1::from_iter(rows.iter().map(|(_, re, _)| *re))),
546 ckp_imag: Some(Array1::from_iter(rows.iter().map(|(_, _, im)| *im))),
547 };
548 let rendered = chi_dat_string(&data)?;
549 let reparsed = parse_chi_dat(&rendered)?;
550 prop_assert_eq!(reparsed, data);
551 }
552 }
553 }
554}