1use crate::Result;
5use rustpix_core::neutron::{Neutron, NeutronBatch};
6use std::fs::File;
7use std::io::{BufWriter, Write};
8use std::path::Path;
9
10pub struct DataFileWriter {
14 writer: BufWriter<File>,
15}
16
17impl DataFileWriter {
18 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 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 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])?; }
66
67 self.writer.flush()?;
68 Ok(())
69 }
70
71 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 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])?; }
115
116 self.writer.flush()?;
117 Ok(())
118 }
119
120 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 assert_eq!(data.len(), 28);
165 }
166}