Skip to main content

rustpix_io/
writer.rs

1//! File writers for processed data.
2//!
3
4use crate::Result;
5use rustpix_core::neutron::{Neutron, NeutronBatch};
6use std::fs::File;
7use std::io::{BufWriter, Write};
8use std::path::Path;
9
10/// Writer for processed data output.
11///
12/// Writes processed neutron data to files in various formats.
13pub struct DataFileWriter {
14    writer: BufWriter<File>,
15}
16
17impl DataFileWriter {
18    /// Creates a new file writer.
19    ///
20    /// # Errors
21    /// Returns an error if the file cannot be created.
22    pub fn create<P: AsRef<Path>>(path: P) -> Result<Self> {
23        let file = File::create(path)?;
24        let writer = BufWriter::new(file);
25        Ok(Self { writer })
26    }
27
28    /// Writes neutrons as CSV.
29    ///
30    /// # Errors
31    /// Returns an error if writing to the underlying file fails.
32    pub fn write_neutrons_csv(&mut self, neutrons: &[Neutron]) -> Result<()> {
33        writeln!(self.writer, "x,y,tof,tot,n_hits,chip_id")?;
34
35        for n in neutrons {
36            writeln!(
37                self.writer,
38                "{},{},{},{},{},{}",
39                n.x, n.y, n.tof, n.tot, n.n_hits, n.chip_id
40            )?;
41        }
42
43        self.writer.flush()?;
44        Ok(())
45    }
46
47    /// Writes neutrons as binary data.
48    ///
49    /// Format per neutron: `f64` (x) + `f64` (y) + `u32` (tof) + `u16` (tot)
50    /// + `u16` (`n_hits`) + `u8` (`chip_id`) + 3 reserved bytes.
51    ///
52    /// Total: 28 bytes per neutron.
53    ///
54    /// # Errors
55    /// Returns an error if writing to the underlying file fails.
56    pub fn write_neutrons_binary(&mut self, neutrons: &[Neutron]) -> Result<()> {
57        for n in neutrons {
58            self.writer.write_all(&n.x.to_le_bytes())?;
59            self.writer.write_all(&n.y.to_le_bytes())?;
60            self.writer.write_all(&n.tof.to_le_bytes())?;
61            self.writer.write_all(&n.tot.to_le_bytes())?;
62            self.writer.write_all(&n.n_hits.to_le_bytes())?;
63            self.writer.write_all(&[n.chip_id])?;
64            self.writer.write_all(&[0u8; 3])?; // Reserved/padding
65        }
66
67        self.writer.flush()?;
68        Ok(())
69    }
70
71    /// Writes neutron batch as CSV.
72    ///
73    /// # Errors
74    /// Returns an error if writing to the underlying file fails.
75    pub fn write_neutron_batch_csv(
76        &mut self,
77        batch: &NeutronBatch,
78        include_header: bool,
79    ) -> Result<()> {
80        if include_header {
81            writeln!(self.writer, "x,y,tof,tot,n_hits,chip_id")?;
82        }
83
84        for i in 0..batch.len() {
85            writeln!(
86                self.writer,
87                "{},{},{},{},{},{}",
88                batch.x[i],
89                batch.y[i],
90                batch.tof[i],
91                batch.tot[i],
92                batch.n_hits[i],
93                batch.chip_id[i]
94            )?;
95        }
96
97        self.writer.flush()?;
98        Ok(())
99    }
100
101    /// Writes neutron batch as binary data.
102    ///
103    /// # Errors
104    /// Returns an error if writing to the underlying file fails.
105    pub fn write_neutron_batch_binary(&mut self, batch: &NeutronBatch) -> Result<()> {
106        for i in 0..batch.len() {
107            self.writer.write_all(&batch.x[i].to_le_bytes())?;
108            self.writer.write_all(&batch.y[i].to_le_bytes())?;
109            self.writer.write_all(&batch.tof[i].to_le_bytes())?;
110            self.writer.write_all(&batch.tot[i].to_le_bytes())?;
111            self.writer.write_all(&batch.n_hits[i].to_le_bytes())?;
112            self.writer.write_all(&[batch.chip_id[i]])?;
113            self.writer.write_all(&[0u8; 3])?; // Reserved/padding
114        }
115
116        self.writer.flush()?;
117        Ok(())
118    }
119
120    /// Flushes the writer.
121    ///
122    /// # Errors
123    /// Returns an error if the underlying writer fails to flush.
124    pub fn flush(&mut self) -> Result<()> {
125        self.writer.flush()?;
126        Ok(())
127    }
128}
129
130#[cfg(test)]
131mod tests {
132    use super::*;
133    use tempfile::NamedTempFile;
134
135    #[test]
136    fn test_write_neutrons_csv() {
137        let file = NamedTempFile::new().unwrap();
138        let mut writer = DataFileWriter::create(file.path()).unwrap();
139
140        let neutrons = vec![
141            Neutron::new(1.5, 2.5, 1000, 100, 5, 0),
142            Neutron::new(10.3, 20.7, 2000, 200, 8, 1),
143        ];
144
145        writer.write_neutrons_csv(&neutrons).unwrap();
146
147        let content = std::fs::read_to_string(file.path()).unwrap();
148        assert!(content.contains("x,y,tof,tot,n_hits,chip_id"));
149        assert!(content.contains("1.5,2.5,1000,100,5,0"));
150        assert!(content.contains("10.3,20.7,2000,200,8,1"));
151    }
152
153    #[test]
154    fn test_write_neutrons_binary() {
155        let file = NamedTempFile::new().unwrap();
156        let mut writer = DataFileWriter::create(file.path()).unwrap();
157
158        let neutrons = vec![Neutron::new(1.5, 2.5, 1000, 100, 5, 0)];
159
160        writer.write_neutrons_binary(&neutrons).unwrap();
161
162        let data = std::fs::read(file.path()).unwrap();
163        // 8 (f64) + 8 (f64) + 4 (u32) + 2 (u16) + 2 (u16) + 1 (u8) + 3 (reserved) = 28 bytes
164        assert_eq!(data.len(), 28);
165    }
166}