egsphsp/
lib.rs

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

egsphsp/
lib.rs
