egsphsp/
lib.rs

1extern crate byteorder;
2extern crate float_cmp;
3extern crate rand;
4
5use std::fmt;
6use std::fs::{remove_file, File, OpenOptions};
7use std::io;
8use std::io::prelude::*;
9use std::io::{BufReader, BufWriter};
10use std::path::Path;
11
12use byteorder::{ByteOrder, LittleEndian};
13use float_cmp::ApproxEqUlps;
14use rand::{rngs::StdRng, seq::SliceRandom, Rng, SeedableRng};
15
16const HEADER_LENGTH: usize = 25;
17const MAX_RECORD_LENGTH: usize = 32;
18const BUFFER_CAPACITY: usize = 1 * 1024 * 1024;
19const MODE_LENGTH: usize = 5;
20const BATCHES: usize = 128; // too high and one hits ulimit (around 1024)
21
22#[derive(Debug, Copy, Clone)]
23pub struct Header {
24    pub mode: [u8; 5],
25    pub total_particles: i32,
26    pub total_photons: i32,
27    pub min_energy: f32,
28    pub max_energy: f32,
29    pub total_particles_in_source: f32,
30    pub record_size: u64,
31    pub using_zlast: bool,
32}
33
34#[derive(Debug, Copy, Clone)]
35pub struct Record {
36    pub latch: u32,
37    total_energy: f32,
38    pub x_cm: f32,
39    pub y_cm: f32,
40    pub x_cos: f32, // TODO verify these are normalized
41    pub y_cos: f32,
42    pub weight: f32, // also carries the sign of the z direction, yikes
43    pub zlast: Option<f32>,
44}
45
46#[derive(Debug)]
47pub struct Transform;
48
49#[derive(Debug)]
50pub enum EGSError {
51    Io(io::Error),
52    BadMode,
53    BadLength,
54    ModeMismatch,
55    HeaderMismatch,
56    RecordMismatch,
57}
58
59pub type EGSResult<T> = Result<T, EGSError>;
60
61impl From<io::Error> for EGSError {
62    fn from(err: io::Error) -> EGSError {
63        EGSError::Io(err)
64    }
65}
66
67impl fmt::Display for EGSError {
68    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
69        match *self {
70            EGSError::Io(ref err) => err.fmt(f),
71            EGSError::BadMode => {
72                write!(
73                    f,
74                    "First 5 bytes of file are invalid, must be MODE0 or MODE2"
75                )
76            }
77            EGSError::BadLength => {
78                write!(
79                    f,
80                    "Number of total particles does notmatch byte length of file"
81                )
82            }
83            EGSError::ModeMismatch => write!(f, "Input file MODE0/MODE2 do not match"),
84            EGSError::HeaderMismatch => write!(f, "Headers are different"),
85            EGSError::RecordMismatch => write!(f, "Records are different"),
86        }
87    }
88}
89
90pub struct PHSPReader {
91    reader: BufReader<File>,
92    pub header: Header,
93    next_record: u64,
94}
95
96pub struct PHSPWriter {
97    writer: BufWriter<File>,
98    pub header: Header,
99}
100
101impl PHSPReader {
102    pub fn from(file: File) -> EGSResult<PHSPReader> {
103        let actual_size = (file.metadata()?).len();
104        let mut reader = BufReader::with_capacity(BUFFER_CAPACITY, file);
105        let mut buffer = [0; HEADER_LENGTH];
106        reader.read_exact(&mut buffer)?;
107        let mut mode = [0; MODE_LENGTH];
108        mode.clone_from_slice(&buffer[0..5]);
109        let header = Header {
110            mode: mode,
111            total_particles: LittleEndian::read_i32(&buffer[5..9]),
112            total_photons: LittleEndian::read_i32(&buffer[9..13]),
113            max_energy: LittleEndian::read_f32(&buffer[13..17]),
114            min_energy: LittleEndian::read_f32(&buffer[17..21]),
115            total_particles_in_source: LittleEndian::read_f32(&buffer[21..25]),
116            using_zlast: &mode == b"MODE2",
117            record_size: if &mode == b"MODE0" {
118                28
119            } else if &mode == b"MODE2" {
120                32
121            } else {
122                return Err(EGSError::BadMode);
123            },
124        };
125        if actual_size != header.expected_size() as u64 {
126            writeln!(
127                &mut std::io::stderr(),
128                "Expected {} bytes in file, not {}",
129                header.expected_size(),
130                actual_size
131            )
132            .unwrap();
133            //return Err(EGSError::BadLength);
134        }
135        reader.consume(header.record_size as usize - HEADER_LENGTH);
136        Ok(PHSPReader {
137            reader: reader,
138            header: header,
139            next_record: 0,
140        })
141    }
142    fn exhausted(&self) -> bool {
143        self.next_record >= self.header.total_particles as u64
144    }
145}
146
147impl Iterator for PHSPReader {
148    type Item = EGSResult<Record>;
149    fn next(&mut self) -> Option<EGSResult<Record>> {
150        if self.next_record >= self.header.total_particles as u64 {
151            return None;
152        }
153        let mut buffer = [0; MAX_RECORD_LENGTH];
154        match self
155            .reader
156            .read_exact(&mut buffer[..self.header.record_size as usize])
157        {
158            Ok(()) => (),
159            Err(err) => return Some(Err(EGSError::Io(err))),
160        };
161        self.next_record += 1;
162        Some(Ok(Record {
163            latch: LittleEndian::read_u32(&buffer[0..4]),
164            total_energy: LittleEndian::read_f32(&buffer[4..8]),
165            x_cm: LittleEndian::read_f32(&buffer[8..12]),
166            y_cm: LittleEndian::read_f32(&buffer[12..16]),
167            x_cos: LittleEndian::read_f32(&buffer[16..20]),
168            y_cos: LittleEndian::read_f32(&buffer[20..24]),
169            weight: LittleEndian::read_f32(&buffer[24..28]),
170            zlast: if self.header.using_zlast {
171                Some(LittleEndian::read_f32(&buffer[28..32]))
172            } else {
173                None
174            },
175        }))
176    }
177}
178
179impl PHSPWriter {
180    pub fn from(file: File, header: &Header) -> EGSResult<PHSPWriter> {
181        let mut writer = BufWriter::with_capacity(BUFFER_CAPACITY, file);
182        let mut buffer = [0; MAX_RECORD_LENGTH];
183        buffer[0..5].clone_from_slice(&header.mode);
184        LittleEndian::write_i32(&mut buffer[5..9], header.total_particles);
185        LittleEndian::write_i32(&mut buffer[9..13], header.total_photons);
186        LittleEndian::write_f32(&mut buffer[13..17], header.max_energy);
187        LittleEndian::write_f32(&mut buffer[17..21], header.min_energy);
188        LittleEndian::write_f32(&mut buffer[21..25], header.total_particles_in_source);
189        writer.write_all(&buffer[..header.record_size as usize])?;
190        Ok(PHSPWriter {
191            header: *header,
192            writer: writer,
193        })
194    }
195
196    pub fn write(&mut self, record: &Record) -> EGSResult<()> {
197        let mut buffer = [0; 32];
198        LittleEndian::write_u32(&mut buffer[0..4], record.latch);
199        LittleEndian::write_f32(&mut buffer[4..8], record.total_energy);
200        LittleEndian::write_f32(&mut buffer[8..12], record.x_cm);
201        LittleEndian::write_f32(&mut buffer[12..16], record.y_cm);
202        LittleEndian::write_f32(&mut buffer[16..20], record.x_cos);
203        LittleEndian::write_f32(&mut buffer[20..24], record.y_cos);
204        LittleEndian::write_f32(&mut buffer[24..28], record.weight);
205        if self.header.using_zlast {
206            LittleEndian::write_f32(&mut buffer[28..32], record.weight);
207        }
208        self.writer
209            .write_all(&buffer[..self.header.record_size as usize])?;
210        Ok(())
211    }
212}
213
214impl Header {
215    fn expected_size(&self) -> usize {
216        (self.total_particles as usize + 1) * self.record_size as usize
217    }
218    pub fn similar_to(&self, other: &Header) -> bool {
219        self.mode == other.mode
220            && self.total_particles == other.total_particles
221            && self.total_photons == other.total_photons
222            && self.max_energy.approx_eq_ulps(&other.max_energy, 10)
223            && self.min_energy.approx_eq_ulps(&other.min_energy, 10)
224            && self
225                .total_particles_in_source
226                .approx_eq_ulps(&other.total_particles_in_source, 2)
227    }
228    fn merge(&mut self, other: &Header) {
229        assert!(&self.mode == &other.mode, "Merge mode mismatch");
230        self.total_particles = self
231            .total_particles
232            .checked_add(other.total_particles)
233            .expect("Too many particles, i32 overflow");
234        self.total_photons += other.total_photons;
235        self.min_energy = self.min_energy.min(other.min_energy);
236        self.max_energy = self.max_energy.max(other.max_energy);
237        self.total_particles_in_source += other.total_particles_in_source;
238    }
239}
240
241impl Record {
242    pub fn similar_to(&self, other: &Record) -> bool {
243        self.latch == other.latch
244            && self.total_energy() - other.total_energy() < 0.01
245            && self.x_cm - other.x_cm < 0.01
246            && self.y_cm - other.y_cm < 0.01
247            && self.x_cos - other.x_cos < 0.01
248            && self.y_cos - other.y_cos < 0.01
249            && self.weight - other.weight < 0.01
250            && self.zlast == other.zlast
251    }
252    pub fn bremsstrahlung_or_annihilation(&self) -> bool {
253        self.latch & 1 != 0
254    }
255    pub fn bit_region(&self) -> u32 {
256        self.latch & 0xfffffe
257    }
258    pub fn region_number(&self) -> u32 {
259        self.latch & 0xf000000
260    }
261    pub fn b29(&self) -> bool {
262        self.latch & (1 << 29) != 0
263    }
264    pub fn charged(&self) -> bool {
265        self.latch & (1 << 30) != 0
266    }
267    pub fn crossed_multiple(&self) -> bool {
268        self.latch & (1 << 30) != 0
269    }
270    pub fn get_weight(&self) -> f32 {
271        self.weight.abs()
272    }
273    pub fn set_weight(&mut self, new_weight: f32) {
274        self.weight = new_weight * self.weight.signum();
275    }
276    pub fn total_energy(&self) -> f32 {
277        self.total_energy.abs()
278    }
279    pub fn z_positive(&self) -> bool {
280        self.weight.is_sign_positive()
281    }
282    pub fn z_cos(&self) -> f32 {
283        (1.0 - (self.x_cos * self.x_cos + self.y_cos * self.y_cos)).sqrt()
284    }
285    pub fn first_scored_by_primary_history(&self) -> bool {
286        return self.total_energy.is_sign_negative();
287    }
288
289    fn translate(&mut self, x: f32, y: f32) {
290        self.x_cm += x;
291        self.y_cm += y;
292    }
293
294    fn transform(&mut self, matrix: &[[f32; 3]; 3]) {
295        let x_cm = self.x_cm;
296        let y_cm = self.y_cm;
297        self.x_cm = matrix[0][0] * x_cm + matrix[0][1] * y_cm + matrix[0][2] * 1.0;
298        self.y_cm = matrix[1][0] * x_cm + matrix[1][1] * y_cm + matrix[1][2] * 1.0;
299        let x_cos = self.x_cos;
300        let y_cos = self.y_cos;
301        self.x_cos = matrix[0][0] * x_cos + matrix[0][1] * y_cos + matrix[0][2] * self.z_cos();
302        self.y_cos = matrix[1][0] * x_cos + matrix[1][1] * y_cos + matrix[1][2] * self.z_cos();
303    }
304}
305
306impl Transform {
307    pub fn reflection(matrix: &mut [[f32; 3]; 3], x_raw: f32, y_raw: f32) {
308        let norm = (x_raw * x_raw + y_raw * y_raw).sqrt();
309        let x = x_raw / norm;
310        let y = y_raw / norm;
311        *matrix = [
312            [x * x - y * y, 2.0 * x * y, 0.0],
313            [2.0 * x * y, y * y - x * x, 0.0],
314            [0.0, 0.0, 1.0],
315        ];
316    }
317    pub fn rotation(matrix: &mut [[f32; 3]; 3], theta: f32) {
318        *matrix = [
319            [theta.cos(), -theta.sin(), 0.0],
320            [theta.sin(), theta.cos(), 0.0],
321            [0.0, 0.0, 1.0],
322        ];
323    }
324}
325
326pub fn randomize(path: &Path, seed: u64) -> EGSResult<()> {
327    let mut rng = StdRng::seed_from_u64(seed);
328    let ifile = File::open(path)?;
329    let mut reader = PHSPReader::from(ifile)?;
330    let header = reader.header;
331    let max_per_batch = reader.header.total_particles as usize / BATCHES + 1;
332    let mut batch_paths = Vec::with_capacity(BATCHES);
333    for i in 0..BATCHES {
334        let mut batch_path = path.to_path_buf();
335        batch_path.set_extension(format!("rand{}", i));
336        batch_paths.push(batch_path);
337    }
338    let mut records = Vec::with_capacity(max_per_batch);
339    for path in batch_paths.iter() {
340        for _ in 0..max_per_batch {
341            match reader.next() {
342                Some(record) => records.push(record.unwrap()),
343                None => (),
344            }
345        }
346        //let mut vec: Vec<Record> = records.collect();
347
348        records.shuffle(&mut rng);
349
350        let header = Header {
351            mode: reader.header.mode,
352            total_particles: records.len() as i32,
353            total_photons: 0,
354            max_energy: 0.0,
355            min_energy: 0.0,
356            total_particles_in_source: 0.0,
357            using_zlast: &reader.header.mode == b"MODE2",
358            record_size: reader.header.record_size,
359        };
360        let ofile = File::create(&path)?;
361        let mut writer = PHSPWriter::from(ofile, &header)?;
362        for record in records.iter() {
363            writer.write(&record)?;
364        }
365        records.clear();
366    }
367    drop(records);
368    let mut readers = Vec::with_capacity(BATCHES);
369    for path in batch_paths.iter() {
370        let ifile = File::open(&path)?;
371        readers.push(PHSPReader::from(ifile)?);
372    }
373
374    let ofile = File::create(path)?;
375    let mut writer = PHSPWriter::from(ofile, &header)?;
376    while readers.len() != 0 {
377        readers.shuffle(&mut rng);
378        for reader in readers.iter_mut() {
379            match reader.next() {
380                Some(record) => writer.write(&record.unwrap())?,
381                None => (),
382            }
383        }
384        readers.retain(|r| !r.exhausted());
385    }
386    for path in batch_paths.iter() {
387        remove_file(path)?;
388    }
389    Ok(())
390}
391
392pub fn combine(input_paths: &[&Path], output_path: &Path, delete: bool) -> EGSResult<()> {
393    assert!(input_paths.len() > 0, "Cannot combine zero files");
394    let reader = PHSPReader::from(File::open(input_paths[0])?)?;
395    let mut final_header = reader.header;
396    for path in input_paths[1..].iter() {
397        let reader = PHSPReader::from(File::open(path)?)?;
398        final_header.merge(&reader.header);
399    }
400    println!("Final header: {:?}", final_header);
401    let ofile = File::create(output_path)?;
402    let mut writer = PHSPWriter::from(ofile, &final_header)?;
403    for path in input_paths.iter() {
404        let reader = PHSPReader::from(File::open(path)?)?;
405        for record in reader {
406            writer.write(&record.unwrap())?
407        }
408        if delete {
409            remove_file(path)?;
410        }
411    }
412    Ok(())
413}
414
415pub fn compare(path1: &Path, path2: &Path) -> EGSResult<()> {
416    let ifile1 = File::open(path1)?;
417    let ifile2 = File::open(path2)?;
418    let reader1 = PHSPReader::from(ifile1)?;
419    let reader2 = PHSPReader::from(ifile2)?;
420    println!("                   First\t\tSecond");
421    println!(
422        "Total particles:   {0: <10}\t\t{1:}",
423        reader1.header.total_particles, reader2.header.total_particles
424    );
425    println!(
426        "Total photons:     {0: <10}\t\t{1}",
427        reader1.header.total_photons, reader2.header.total_photons
428    );
429    println!(
430        "Minimum energy:    {0: <10}\t\t{1}",
431        reader1.header.min_energy, reader2.header.min_energy
432    );
433    println!(
434        "Maximum energy:    {0: <10}\t\t{1}",
435        reader1.header.max_energy, reader2.header.max_energy
436    );
437    println!(
438        "Source particles:  {0: <10}\t\t{1}",
439        reader1.header.total_particles_in_source, reader2.header.total_particles_in_source
440    );
441    if !reader1.header.similar_to(&reader2.header) {
442        println!("Headers different");
443        return Err(EGSError::HeaderMismatch);
444    } else {
445        for (record1, record2) in reader1.zip(reader2) {
446            let r1 = record1.unwrap();
447            let r2 = record2.unwrap();
448            if !r1.similar_to(&r2) {
449                println!("{:?} != {:?}", r1, r2);
450                return Err(EGSError::RecordMismatch);
451            }
452        }
453    }
454    Ok(())
455}
456
457pub fn sample_combine(ipaths: &[&Path], opath: &Path, rate: f64, seed: u64) -> EGSResult<()> {
458    assert!(ipaths.len() > 0, "Cannot combine zero files");
459    let mut rng = StdRng::seed_from_u64(seed);
460    let mut header = Header {
461        mode: *b"MODE0",
462        record_size: 28,
463        using_zlast: false,
464        total_particles: 0,
465        total_photons: 0,
466        min_energy: 1000.0,
467        max_energy: 0.0,
468        total_particles_in_source: 0.0,
469    };
470    let mut writer = PHSPWriter::from(File::create(opath)?, &header)?;
471    for path in ipaths.iter() {
472        let reader = PHSPReader::from(File::open(path)?)?;
473        assert!(!reader.header.using_zlast);
474        println!("Found {} particles", reader.header.total_particles);
475        header.total_particles_in_source += reader.header.total_particles_in_source;
476        let records = reader.filter(|_| rng.random_bool(rate));
477        for record in records.map(|r| r.unwrap()) {
478            header.total_particles = header
479                .total_particles
480                .checked_add(1)
481                .expect("Total particles overflow");
482            if !record.charged() {
483                header.total_photons += 1;
484            }
485            if record.total_energy > 0.0 {
486                header.min_energy = header.min_energy.min(record.total_energy);
487                header.max_energy = header.max_energy.max(record.total_energy);
488            }
489            writer.write(&record)?;
490        }
491        println!("Now have {} particles", header.total_particles);
492    }
493    header.total_particles_in_source *= rate as f32;
494    drop(writer);
495    // write out the header
496    let ofile = OpenOptions::new().write(true).create(true).open(opath)?;
497    PHSPWriter::from(ofile, &header)?;
498    Ok(())
499}
500
501pub fn translate(input_path: &Path, output_path: &Path, x: f32, y: f32) -> EGSResult<()> {
502    let ifile = File::open(input_path)?;
503    let reader = PHSPReader::from(ifile)?;
504    let ofile;
505    if input_path == output_path {
506        println!(
507            "Translating {} in place by ({}, {})",
508            input_path.display(),
509            x,
510            y
511        );
512        ofile = OpenOptions::new()
513            .write(true)
514            .create(true)
515            .open(output_path)?;
516    } else {
517        println!(
518            "Translating {} by ({}, {}) and saving to {}",
519            input_path.display(),
520            x,
521            y,
522            output_path.display()
523        );
524        ofile = File::create(output_path)?;
525    }
526    let mut writer = PHSPWriter::from(ofile, &reader.header)?;
527    let n_particles = reader.header.total_particles;
528    let mut records_translated = 0;
529    for mut record in reader.map(|r| r.unwrap()) {
530        record.translate(x, y);
531        writer.write(&record)?;
532        records_translated += 1;
533    }
534    println!(
535        "Translated {} records, expected {}",
536        records_translated, n_particles
537    );
538    Ok(())
539}
540
541pub fn transform(input_path: &Path, output_path: &Path, matrix: &[[f32; 3]; 3]) -> EGSResult<()> {
542    let ifile = File::open(input_path)?;
543    let reader = PHSPReader::from(ifile)?;
544    let ofile;
545    if input_path == output_path {
546        println!("Transforming {} in place", input_path.display());
547        ofile = OpenOptions::new()
548            .write(true)
549            .create(true)
550            .open(output_path)?;
551    } else {
552        // different path (create/truncate destination)
553        println!(
554            "Transforming {} and saving to {}",
555            input_path.display(),
556            output_path.display()
557        );
558        ofile = File::create(output_path)?;
559    }
560    let mut writer = PHSPWriter::from(ofile, &reader.header)?;
561    let n_particles = reader.header.total_particles;
562    let mut records_transformed = 0;
563    for mut record in reader.map(|r| r.unwrap()) {
564        record.transform(&matrix);
565        writer.write(&record)?;
566        records_transformed += 1;
567    }
568    println!(
569        "Transformed {} records, expected {}",
570        records_transformed, n_particles
571    );
572    Ok(())
573}
574
575pub fn reweight(
576    input_path: &Path,
577    output_path: &Path,
578    f: &dyn Fn(f32) -> f32,
579    number_bins: usize,
580    max_radius: f32,
581) -> EGSResult<()> {
582    let input_file = File::open(input_path)?;
583    let output_file;
584    if input_path == output_path {
585        println!("Reweighting in-place");
586        output_file = OpenOptions::new()
587            .write(true)
588            .create(true)
589            .open(output_path)?;
590    } else {
591        println!("Rewighting and saving to {}", output_path.display());
592        output_file = File::create(output_path)?;
593    }
594    let reader1 = PHSPReader::from(input_file)?;
595    let mut writer1 = PHSPWriter::from(output_file, &reader1.header)?;
596    let bin_size = max_radius / number_bins as f32;
597    println!("Bin size is {}", bin_size);
598    let mut sum_old_weight = 0.0;
599    let mut sum_new_weight = 0.0;
600    for mut record in reader1.map(|r| r.unwrap()) {
601        sum_old_weight += record.weight;
602        let r = (record.x_cm * record.x_cm + record.y_cm * record.y_cm).sqrt();
603        record.weight *= f(r);
604        sum_new_weight += record.weight;
605        writer1.write(&record)?;
606    }
607    drop(writer1);
608    let ifile2 = File::open(input_path)?;
609    let ofile2;
610    if input_path == output_path {
611        ofile2 = OpenOptions::new()
612            .write(true)
613            .create(true)
614            .open(output_path)?;
615    } else {
616        ofile2 = File::create(output_path)?;
617    }
618    let reader2 = PHSPReader::from(ifile2)?;
619    let mut writer2 = PHSPWriter::from(ofile2, &reader2.header)?;
620    let factor = sum_old_weight / sum_new_weight;
621    for mut record in reader2.map(|r| r.unwrap()) {
622        record.weight *= factor;
623        writer2.write(&record)?;
624    }
625    Ok(())
626}
egsphsp/lib.rs

egsphsp/
lib.rs
