egsphsp/
lib.rs

1extern crate float_cmp;
2extern crate byteorder;
3extern crate rand;
4extern crate cpu_time;
5
6use std::error::Error;
7use std::fs::{File, OpenOptions, remove_file};
8use std::io::{BufReader, BufWriter};
9use std::io::prelude::*;
10use std::path::Path;
11use std::str;
12use std::io;
13use std::fmt;
14
15use cpu_time::ProcessTime;
16use std::time::Duration;
17use byteorder::{ByteOrder, LittleEndian};
18use rand::{SeedableRng, StdRng, Rng};
19use float_cmp::ApproxEqUlps;
20
21const HEADER_LENGTH: usize = 25;
22const MAX_RECORD_LENGTH: usize = 32;
23const BUFFER_CAPACITY: usize = 1 * 1024 * 1024;
24const MODE_LENGTH: usize = 5;
25
26#[derive(Debug, Copy, Clone)]
27pub struct Header {
28    pub mode: [u8; 5],
29    pub total_particles: i32,
30    pub total_photons: i32,
31    pub min_energy: f32,
32    pub max_energy: f32,
33    pub total_particles_in_source: f32,
34    pub record_size: u64,
35    pub using_zlast: bool,
36}
37
38#[derive(Debug, Copy, Clone)]
39pub struct Record {
40    pub latch: u32,
41    total_energy: f32,
42    pub x_cm: f32,
43    pub y_cm: f32,
44    pub x_cos: f32,
45    pub y_cos: f32,
46    pub weight: f32, // also carries the sign of the z direction, yikes
47    pub zlast: Option<f32>,
48}
49
50#[derive(Debug)]
51pub struct Transform;
52
53#[derive(Debug)]
54pub enum EGSError {
55    Io(io::Error),
56    BadMode,
57    BadLength,
58    ModeMismatch,
59    HeaderMismatch,
60    RecordMismatch,
61}
62
63pub type EGSResult<T> = Result<T, EGSError>;
64
65
66impl From<io::Error> for EGSError {
67    fn from(err: io::Error) -> EGSError {
68        EGSError::Io(err)
69    }
70}
71
72impl fmt::Display for EGSError {
73    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
74        match *self {
75            EGSError::Io(ref err) => err.fmt(f),
76            EGSError::BadMode => {
77                write!(f,
78                       "First 5 bytes of file are invalid, must be MODE0 or MODE2")
79            }
80            EGSError::BadLength => {
81                write!(f,
82                       "Number of total particles does notmatch byte length of file")
83            }
84            EGSError::ModeMismatch => write!(f, "Input file MODE0/MODE2 do not match"),
85            EGSError::HeaderMismatch => write!(f, "Headers are different"),
86            EGSError::RecordMismatch => write!(f, "Records are different"),
87        }
88    }
89}
90
91impl Error for EGSError {
92    fn description(&self) -> &str {
93        match *self {
94            EGSError::Io(ref err) => err.description(),
95            EGSError::BadMode => "invalid mode",
96            EGSError::BadLength => "bad file length",
97            EGSError::ModeMismatch => "mode mismatch",
98            EGSError::HeaderMismatch => "header mismatch",
99            EGSError::RecordMismatch => "record mismatch",
100        }
101    }
102
103    fn cause(&self) -> Option<&dyn Error> {
104        match *self {
105            EGSError::Io(ref err) => Some(err),
106            EGSError::BadMode => None,
107            EGSError::BadLength => None,
108            EGSError::ModeMismatch => None,
109            EGSError::HeaderMismatch => None,
110            EGSError::RecordMismatch => None,
111        }
112    }
113}
114
115pub struct PHSPReader {
116    reader: BufReader<File>,
117    pub header: Header,
118    next_record: u64,
119}
120
121pub struct PHSPWriter {
122    writer: BufWriter<File>,
123    pub header: Header,
124}
125
126
127impl PHSPReader {
128    pub fn from(file: File) -> EGSResult<PHSPReader> {
129        let actual_size = file.metadata()?.len();
130        let mut reader = BufReader::with_capacity(BUFFER_CAPACITY, file);
131        let mut buffer = [0; HEADER_LENGTH];
132        reader.read_exact(&mut buffer)?;
133        let mut mode = [0; MODE_LENGTH];
134        mode.clone_from_slice(&buffer[0..5]);
135        let header = Header {
136            mode: mode,
137            total_particles: LittleEndian::read_i32(&buffer[5..9]),
138            total_photons: LittleEndian::read_i32(&buffer[9..13]),
139            max_energy: LittleEndian::read_f32(&buffer[13..17]),
140            min_energy: LittleEndian::read_f32(&buffer[17..21]),
141            total_particles_in_source: LittleEndian::read_f32(&buffer[21..25]),
142            using_zlast: &mode == b"MODE2",
143            record_size: if &mode == b"MODE0" {
144                28
145            } else if &mode == b"MODE2" {
146                32
147            } else {
148                return Err(EGSError::BadMode);
149            },
150        };
151        if actual_size != header.expected_size() as u64 {
152            writeln!(&mut std::io::stderr(),
153                     "Expected {} bytes in file, not {}",
154                     header.expected_size(),
155                     actual_size)
156                .unwrap();
157            //return Err(EGSError::BadLength);
158        }
159        reader.consume(header.record_size as usize - HEADER_LENGTH);
160        Ok(PHSPReader {
161            reader: reader,
162            header: header,
163            next_record: 0,
164        })
165    }
166}
167
168impl Iterator for PHSPReader {
169    type Item = EGSResult<Record>;
170    fn next(&mut self) -> Option<EGSResult<Record>> {
171        if self.next_record >= self.header.total_particles as u64 {
172            return None;
173        }
174        let mut buffer = [0; MAX_RECORD_LENGTH];
175        match self.reader.read_exact(&mut buffer[..self.header.record_size as usize]) {
176            Ok(()) => (),
177            Err(err) => return Some(Err(EGSError::Io(err))),
178        };
179        self.next_record += 1;
180        Some(Ok(Record {
181            latch: LittleEndian::read_u32(&buffer[0..4]),
182            total_energy: LittleEndian::read_f32(&buffer[4..8]),
183            x_cm: LittleEndian::read_f32(&buffer[8..12]),
184            y_cm: LittleEndian::read_f32(&buffer[12..16]),
185            x_cos: LittleEndian::read_f32(&buffer[16..20]),
186            y_cos: LittleEndian::read_f32(&buffer[20..24]),
187            weight: LittleEndian::read_f32(&buffer[24..28]),
188            zlast: if self.header.using_zlast {
189                Some(LittleEndian::read_f32(&buffer[28..32]))
190            } else {
191                None
192            },
193        }))
194    }
195}
196
197impl PHSPWriter {
198    pub fn from(file: File, header: &Header) -> EGSResult<PHSPWriter> {
199        let mut writer = BufWriter::with_capacity(BUFFER_CAPACITY, file);
200        let mut buffer = [0; MAX_RECORD_LENGTH];
201        buffer[0..5].clone_from_slice(&header.mode);
202        LittleEndian::write_i32(&mut buffer[5..9], header.total_particles);
203        LittleEndian::write_i32(&mut buffer[9..13], header.total_photons);
204        LittleEndian::write_f32(&mut buffer[13..17], header.max_energy);
205        LittleEndian::write_f32(&mut buffer[17..21], header.min_energy);
206        LittleEndian::write_f32(&mut buffer[21..25], header.total_particles_in_source);
207        writer.write_all(&buffer[..header.record_size as usize])?;
208        Ok(PHSPWriter {
209            header: *header,
210            writer: writer,
211        })
212    }
213
214    pub fn write(&mut self, record: &Record) -> EGSResult<()> {
215        let mut buffer = [0; 32];
216        LittleEndian::write_u32(&mut buffer[0..4], record.latch);
217        LittleEndian::write_f32(&mut buffer[4..8], record.total_energy);
218        LittleEndian::write_f32(&mut buffer[8..12], record.x_cm);
219        LittleEndian::write_f32(&mut buffer[12..16], record.y_cm);
220        LittleEndian::write_f32(&mut buffer[16..20], record.x_cos);
221        LittleEndian::write_f32(&mut buffer[20..24], record.y_cos);
222        LittleEndian::write_f32(&mut buffer[24..28], record.weight);
223        if self.header.using_zlast {
224            LittleEndian::write_f32(&mut buffer[28..32], record.weight);
225        }
226        self.writer.write_all(&buffer[..self.header.record_size as usize])?;
227        Ok(())
228    }
229}
230
231impl Header {
232    fn expected_size(&self) -> usize {
233        (self.total_particles as usize + 1) * self.record_size as usize
234    }
235    pub fn similar_to(&self, other: &Header) -> bool {
236        self.mode == other.mode && self.total_particles == other.total_particles &&
237        self.total_photons == other.total_photons &&
238        self.max_energy.approx_eq_ulps(&other.max_energy, 10) &&
239        self.min_energy.approx_eq_ulps(&other.min_energy, 10) &&
240        self.total_particles_in_source.approx_eq_ulps(&other.total_particles_in_source, 2)
241    }
242    fn merge(&mut self, other: &Header) {
243        assert!(&self.mode == &other.mode, "Merge mode mismatch");
244        self.total_particles = self.total_particles
245            .checked_add(other.total_particles)
246            .expect("Too many particles, i32 overflow");
247        self.total_photons += other.total_photons;
248        self.min_energy = self.min_energy.min(other.min_energy);
249        self.max_energy = self.max_energy.max(other.max_energy);
250        self.total_particles_in_source += other.total_particles_in_source;
251    }
252}
253
254
255impl Record {
256    pub fn similar_to(&self, other: &Record) -> bool {
257        self.latch == other.latch && self.total_energy() - other.total_energy() < 0.01 &&
258        self.x_cm - other.x_cm < 0.01 && self.y_cm - other.y_cm < 0.01 &&
259        self.x_cos - other.x_cos < 0.01 && self.y_cos - other.y_cos < 0.01 &&
260        self.weight - other.weight < 0.01 && self.zlast == other.zlast
261    }
262    pub fn bremsstrahlung_or_annihilation(&self) -> bool {
263        self.latch & 1 != 0
264    }
265    pub fn bit_region(&self) -> u32 {
266        self.latch & 0xfffffe
267    }
268    pub fn region_number(&self) -> u32 {
269        self.latch & 0xf000000
270    }
271    pub fn b29(&self) -> bool {
272        self.latch & (1 << 29) != 0
273    }
274    pub fn charged(&self) -> bool {
275        self.latch & (1 << 30) != 0
276    }
277    pub fn crossed_multiple(&self) -> bool {
278        self.latch & (1 << 30) != 0
279    }
280    pub fn get_weight(&self) -> f32 {
281        self.weight.abs()
282    }
283    pub fn set_weight(&mut self, new_weight: f32) {
284        self.weight = new_weight * self.weight.signum();
285    }
286    pub fn total_energy(&self) -> f32 {
287        self.total_energy.abs()
288    }
289    pub fn z_positive(&self) -> bool {
290        self.weight.is_sign_positive()
291    }
292    pub fn z_cos(&self) -> f32 {
293        (1.0 - (self.x_cos * self.x_cos + self.y_cos * self.y_cos)).sqrt()
294    }
295    pub fn first_scored_by_primary_history(&self) -> bool {
296        return self.total_energy.is_sign_negative();
297    }
298
299    fn transform(&mut self, matrix: &[[f32; 3]; 3]) {
300        let x_cm = self.x_cm;
301        let y_cm = self.y_cm;
302        self.x_cm = matrix[0][0] * x_cm + matrix[0][1] * y_cm + matrix[0][2] * 1.0;
303        self.y_cm = matrix[1][0] * x_cm + matrix[1][1] * y_cm + matrix[1][2] * 1.0;
304        let x_cos = self.x_cos;
305        let y_cos = self.y_cos;
306        self.x_cos = matrix[0][0] * x_cos + matrix[0][1] * y_cos + matrix[0][2] * self.z_cos();
307        self.y_cos = matrix[1][0] * x_cos + matrix[1][1] * y_cos + matrix[1][2] * self.z_cos();
308    }
309}
310
311impl Transform {
312    pub fn rotation(matrix: &mut [[f32; 3]; 3], theta: f32) {
313        *matrix =
314            [[theta.cos(), -theta.sin(), 0.0], [theta.sin(), theta.cos(), 0.0], [0.0, 0.0, 1.0]];
315    }
316}
317
318
319
320pub fn combine(input_paths: &[&Path], output_path: &Path, delete: bool) -> EGSResult<()> {
321    assert!(input_paths.len() > 0, "Cannot combine zero files");
322    let start = ProcessTime::now();
323    let reader = PHSPReader::from(File::open(input_paths[0])?)?;
324    let mut final_header = reader.header;
325    for path in input_paths[1..].iter() {
326        let reader = PHSPReader::from(File::open(path)?)?;
327        final_header.merge(&reader.header);
328    }
329    println!("");
330    println!("Final header: {:?}", final_header);
331    println!("");
332    let ofile = File::create(output_path)?;
333    let mut writer = PHSPWriter::from(ofile, &final_header)?;
334    for path in input_paths.iter() {
335        let reader = PHSPReader::from(File::open(path)?)?;
336        for record in reader {
337            writer.write(&record.unwrap())?
338        }
339        if delete {
340            remove_file(path)?;
341        }
342    }
343    let cpu_time: Duration = start.elapsed();
344    println!("CPU time: {:?}", cpu_time);
345    Ok(())
346}
347
348pub fn sample(ipaths: &[&Path], opath: &Path, rate: u32, seed: &[usize]) -> EGSResult<()> {
349    assert!(ipaths.len() > 0, "Cannot combine zero files");
350    let mut rng: StdRng = SeedableRng::from_seed(seed);
351    let mut header = Header {
352        mode: *b"MODE0",
353        record_size: 28,
354        using_zlast: false,
355        total_particles: 0,
356        total_photons: 0,
357        min_energy: 1000.0,
358        max_energy: 0.0,
359        total_particles_in_source: 0.0,
360    };
361    let mut writer = PHSPWriter::from(File::create(opath)?, &header)?;
362    for path in ipaths.iter() {
363        let reader = PHSPReader::from(File::open(path)?)?;
364        assert!(!reader.header.using_zlast);
365        println!("Found {} particles", reader.header.total_particles);
366        header.total_particles_in_source += reader.header.total_particles_in_source;
367        let records = reader.filter(|_| rng.gen_weighted_bool(rate));
368        for record in records.map(|r| r.unwrap()) {
369            header.total_particles =
370                header.total_particles.checked_add(1).expect("Total particles overflow");
371            if !record.charged() {
372                header.total_photons += 1;
373            }
374            if record.total_energy > 0.0 {
375                header.min_energy = header.min_energy.min(record.total_energy);
376                header.max_energy = header.max_energy.max(record.total_energy);
377            }
378            writer.write(&record)?;
379        }
380        println!("Now have {} particles", header.total_particles);
381    }
382    header.total_particles_in_source /= rate as f32;
383    drop(writer);
384    // write out the header
385    let ofile = OpenOptions::new().write(true).create(true).open(opath)?;
386    PHSPWriter::from(ofile, &header)?;
387    Ok(())
388}
389
390pub fn transform(input_path: &Path, output_path: &Path, matrix: &[[f32; 3]; 3]) -> EGSResult<()> {
391    let ifile = File::open(input_path)?;
392    let reader = PHSPReader::from(ifile)?;
393    let ofile;
394    if input_path == output_path {
395        println!("Transforming {} in place", input_path.display());
396        ofile = OpenOptions::new().write(true).create(true).open(output_path)?;
397    } else {
398        // different path (create/truncate destination)
399        println!("Transforming {} and saving to {}",
400                 input_path.display(),
401                 output_path.display());
402        ofile = File::create(output_path)?;
403    }
404    let mut writer = PHSPWriter::from(ofile, &reader.header)?;
405    let n_particles = reader.header.total_particles;
406    let mut records_transformed = 0;
407    for mut record in reader.map(|r| r.unwrap()) {
408        record.transform(&matrix);
409        writer.write(&record)?;
410        records_transformed += 1;
411    }
412    println!("Transformed {} records, expected {}",
413             records_transformed,
414             n_particles);
415    Ok(())
416}
egsphsp/lib.rs

egsphsp/
lib.rs
