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(
203                &mut buffer[28..32],
204                record.zlast.expect("MODE2 record missing zlast"),
205            );
206        }
207        self.writer
208            .write_all(&buffer[..self.header.record_size as usize])?;
209        Ok(())
210    }
211}
212
213impl Header {
214    fn expected_size(&self) -> usize {
215        (self.total_particles as usize + 1) * self.record_size as usize
216    }
217    pub fn similar_to(&self, other: &Header) -> bool {
218        self.mode == other.mode
219            && self.total_particles == other.total_particles
220            && self.total_photons == other.total_photons
221            && self.max_energy.approx_eq_ulps(&other.max_energy, 10)
222            && self.min_energy.approx_eq_ulps(&other.min_energy, 10)
223            && self
224                .total_particles_in_source
225                .approx_eq_ulps(&other.total_particles_in_source, 2)
226    }
227    fn merge(&mut self, other: &Header) {
228        assert!(self.mode == other.mode, "Merge mode mismatch");
229        self.total_particles = self
230            .total_particles
231            .checked_add(other.total_particles)
232            .expect("Too many particles, i32 overflow");
233        self.total_photons += other.total_photons;
234        self.min_energy = self.min_energy.min(other.min_energy);
235        self.max_energy = self.max_energy.max(other.max_energy);
236        self.total_particles_in_source += other.total_particles_in_source;
237    }
238}
239
240impl Record {
241    pub fn similar_to(&self, other: &Record) -> bool {
242        self.latch == other.latch
243            && (self.total_energy() - other.total_energy()).abs() < 0.01
244            && (self.x_cm - other.x_cm).abs() < 0.01
245            && (self.y_cm - other.y_cm).abs() < 0.01
246            && (self.x_cos - other.x_cos).abs() < 0.01
247            && (self.y_cos - other.y_cos).abs() < 0.01
248            && (self.weight - other.weight).abs() < 0.01
249            && self.zlast == other.zlast
250    }
251    pub fn bremsstrahlung_or_annihilation(&self) -> bool {
252        self.latch & 1 != 0
253    }
254    pub fn bit_region(&self) -> u32 {
255        self.latch & 0xfffffe
256    }
257    pub fn region_number(&self) -> u32 {
258        self.latch & 0xf000000
259    }
260    pub fn b29(&self) -> bool {
261        self.latch & (1 << 29) != 0
262    }
263    pub fn charged(&self) -> bool {
264        self.latch & (1 << 30) != 0
265    }
266    pub fn crossed_multiple(&self) -> bool {
267        self.latch & (1 << 31) != 0
268    }
269    pub fn get_weight(&self) -> f32 {
270        self.weight.abs()
271    }
272    pub fn set_weight(&mut self, new_weight: f32) {
273        self.weight = new_weight * self.weight.signum();
274    }
275    pub fn total_energy(&self) -> f32 {
276        self.total_energy.abs()
277    }
278    pub fn z_positive(&self) -> bool {
279        self.weight.is_sign_positive()
280    }
281    pub fn z_cos(&self) -> f32 {
282        (1.0 - (self.x_cos * self.x_cos + self.y_cos * self.y_cos)).sqrt()
283    }
284    pub fn first_scored_by_primary_history(&self) -> bool {
285        self.total_energy.is_sign_negative()
286    }
287
288    fn translate(&mut self, x: f32, y: f32) {
289        self.x_cm += x;
290        self.y_cm += y;
291    }
292
293    fn transform(&mut self, matrix: &[[f32; 3]; 3]) {
294        let x_cm = self.x_cm;
295        let y_cm = self.y_cm;
296        self.x_cm = matrix[0][0] * x_cm + matrix[0][1] * y_cm + matrix[0][2] * 1.0;
297        self.y_cm = matrix[1][0] * x_cm + matrix[1][1] * y_cm + matrix[1][2] * 1.0;
298        let x_cos = self.x_cos;
299        let y_cos = self.y_cos;
300        let z_cos = self.z_cos();
301        self.x_cos = matrix[0][0] * x_cos + matrix[0][1] * y_cos + matrix[0][2] * z_cos;
302        self.y_cos = matrix[1][0] * x_cos + matrix[1][1] * y_cos + matrix[1][2] * 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            if let Some(record) = reader.next() { records.push(record.unwrap()) }
342        }
343        //let mut vec: Vec<Record> = records.collect();
344
345        records.shuffle(&mut rng);
346
347        let header = Header {
348            mode: reader.header.mode,
349            total_particles: records.len() as i32,
350            total_photons: 0,
351            max_energy: 0.0,
352            min_energy: 0.0,
353            total_particles_in_source: 0.0,
354            using_zlast: &reader.header.mode == b"MODE2",
355            record_size: reader.header.record_size,
356        };
357        let ofile = File::create(path)?;
358        let mut writer = PHSPWriter::from(ofile, &header)?;
359        for record in records.iter() {
360            writer.write(record)?;
361        }
362        records.clear();
363    }
364    drop(records);
365    let mut readers = Vec::with_capacity(BATCHES);
366    for path in batch_paths.iter() {
367        let ifile = File::open(path)?;
368        readers.push(PHSPReader::from(ifile)?);
369    }
370
371    let ofile = File::create(path)?;
372    let mut writer = PHSPWriter::from(ofile, &header)?;
373    while !readers.is_empty() {
374        readers.shuffle(&mut rng);
375        for reader in readers.iter_mut() {
376            if let Some(record) = reader.next() { writer.write(&record.unwrap())? }
377        }
378        readers.retain(|r| !r.exhausted());
379    }
380    for path in batch_paths.iter() {
381        remove_file(path)?;
382    }
383    Ok(())
384}
385
386pub fn combine(input_paths: &[&Path], output_path: &Path, delete: bool) -> EGSResult<()> {
387    assert!(!input_paths.is_empty(), "Cannot combine zero files");
388    let reader = PHSPReader::from(File::open(input_paths[0])?)?;
389    let mut final_header = reader.header;
390    for path in input_paths[1..].iter() {
391        let reader = PHSPReader::from(File::open(path)?)?;
392        final_header.merge(&reader.header);
393    }
394    println!("Final header: {:?}", final_header);
395    let ofile = File::create(output_path)?;
396    let mut writer = PHSPWriter::from(ofile, &final_header)?;
397    for path in input_paths.iter() {
398        let reader = PHSPReader::from(File::open(path)?)?;
399        for record in reader {
400            writer.write(&record.unwrap())?
401        }
402        if delete {
403            remove_file(path)?;
404        }
405    }
406    Ok(())
407}
408
409pub fn compare(path1: &Path, path2: &Path) -> EGSResult<()> {
410    let ifile1 = File::open(path1)?;
411    let ifile2 = File::open(path2)?;
412    let reader1 = PHSPReader::from(ifile1)?;
413    let reader2 = PHSPReader::from(ifile2)?;
414    println!("                   First\t\tSecond");
415    println!(
416        "Total particles:   {0: <10}\t\t{1:}",
417        reader1.header.total_particles, reader2.header.total_particles
418    );
419    println!(
420        "Total photons:     {0: <10}\t\t{1}",
421        reader1.header.total_photons, reader2.header.total_photons
422    );
423    println!(
424        "Minimum energy:    {0: <10}\t\t{1}",
425        reader1.header.min_energy, reader2.header.min_energy
426    );
427    println!(
428        "Maximum energy:    {0: <10}\t\t{1}",
429        reader1.header.max_energy, reader2.header.max_energy
430    );
431    println!(
432        "Source particles:  {0: <10}\t\t{1}",
433        reader1.header.total_particles_in_source, reader2.header.total_particles_in_source
434    );
435    if !reader1.header.similar_to(&reader2.header) {
436        println!("Headers different");
437        return Err(EGSError::HeaderMismatch);
438    } else {
439        for (record1, record2) in reader1.zip(reader2) {
440            let r1 = record1.unwrap();
441            let r2 = record2.unwrap();
442            if !r1.similar_to(&r2) {
443                println!("{:?} != {:?}", r1, r2);
444                return Err(EGSError::RecordMismatch);
445            }
446        }
447    }
448    Ok(())
449}
450
451pub fn sample_combine(ipaths: &[&Path], opath: &Path, rate: f64, seed: u64) -> EGSResult<()> {
452    assert!(!ipaths.is_empty(), "Cannot combine zero files");
453    let mut rng = StdRng::seed_from_u64(seed);
454    let mut header = Header {
455        mode: *b"MODE0",
456        record_size: 28,
457        using_zlast: false,
458        total_particles: 0,
459        total_photons: 0,
460        min_energy: 1000.0,
461        max_energy: 0.0,
462        total_particles_in_source: 0.0,
463    };
464    let mut writer = PHSPWriter::from(File::create(opath)?, &header)?;
465    for path in ipaths.iter() {
466        let reader = PHSPReader::from(File::open(path)?)?;
467        assert!(!reader.header.using_zlast);
468        println!("Found {} particles", reader.header.total_particles);
469        header.total_particles_in_source += reader.header.total_particles_in_source;
470        let records = reader.filter(|_| rng.random_bool(rate));
471        for record in records.map(|r| r.unwrap()) {
472            header.total_particles = header
473                .total_particles
474                .checked_add(1)
475                .expect("Total particles overflow");
476            if !record.charged() {
477                header.total_photons += 1;
478            }
479            let energy = record.total_energy();
480            header.min_energy = header.min_energy.min(energy);
481            header.max_energy = header.max_energy.max(energy);
482            writer.write(&record)?;
483        }
484        println!("Now have {} particles", header.total_particles);
485    }
486    header.total_particles_in_source *= rate as f32;
487    drop(writer);
488    // write out the header
489    let ofile = OpenOptions::new()
490        .write(true)
491        .create(true)
492        .truncate(false)
493        .open(opath)?;
494    PHSPWriter::from(ofile, &header)?;
495    Ok(())
496}
497
498pub fn translate(input_path: &Path, output_path: &Path, x: f32, y: f32) -> EGSResult<()> {
499    let ifile = File::open(input_path)?;
500    let reader = PHSPReader::from(ifile)?;
501    let ofile = if input_path == output_path {
502        println!(
503            "Translating {} in place by ({}, {})",
504            input_path.display(),
505            x,
506            y
507        );
508        OpenOptions::new()
509            .write(true)
510            .create(true)
511            .truncate(false)
512            .open(output_path)?
513    } else {
514        println!(
515            "Translating {} by ({}, {}) and saving to {}",
516            input_path.display(),
517            x,
518            y,
519            output_path.display()
520        );
521        File::create(output_path)?
522    };
523    let mut writer = PHSPWriter::from(ofile, &reader.header)?;
524    let n_particles = reader.header.total_particles;
525    let mut records_translated = 0;
526    for mut record in reader.map(|r| r.unwrap()) {
527        record.translate(x, y);
528        writer.write(&record)?;
529        records_translated += 1;
530    }
531    println!(
532        "Translated {} records, expected {}",
533        records_translated, n_particles
534    );
535    Ok(())
536}
537
538pub fn transform(input_path: &Path, output_path: &Path, matrix: &[[f32; 3]; 3]) -> EGSResult<()> {
539    let ifile = File::open(input_path)?;
540    let reader = PHSPReader::from(ifile)?;
541    let ofile = if input_path == output_path {
542        println!("Transforming {} in place", input_path.display());
543        OpenOptions::new()
544            .write(true)
545            .create(true)
546            .truncate(false)
547            .open(output_path)?
548    } else {
549        println!(
550            "Transforming {} and saving to {}",
551            input_path.display(),
552            output_path.display()
553        );
554        File::create(output_path)?
555    };
556    let mut writer = PHSPWriter::from(ofile, &reader.header)?;
557    let n_particles = reader.header.total_particles;
558    let mut records_transformed = 0;
559    for mut record in reader.map(|r| r.unwrap()) {
560        record.transform(matrix);
561        writer.write(&record)?;
562        records_transformed += 1;
563    }
564    println!(
565        "Transformed {} records, expected {}",
566        records_transformed, n_particles
567    );
568    Ok(())
569}
570
571pub fn reweight(
572    input_path: &Path,
573    output_path: &Path,
574    f: &dyn Fn(f32) -> f32,
575    _number_bins: usize,
576    _max_radius: f32,
577) -> EGSResult<()> {
578    if input_path == output_path {
579        println!("Reweighting in-place");
580    } else {
581        println!("Reweighting and saving to {}", output_path.display());
582    }
583
584    let reader1 = PHSPReader::from(File::open(input_path)?)?;
585    let mut sum_old_weight = 0.0_f32;
586    let mut sum_new_weight = 0.0_f32;
587    for record in reader1.map(|r| r.unwrap()) {
588        sum_old_weight += record.weight;
589        let r = (record.x_cm * record.x_cm + record.y_cm * record.y_cm).sqrt();
590        sum_new_weight += record.weight * f(r);
591    }
592
593    let reader2 = PHSPReader::from(File::open(input_path)?)?;
594    let output_file = if input_path == output_path {
595        OpenOptions::new()
596            .write(true)
597            .create(true)
598            .truncate(false)
599            .open(output_path)?
600    } else {
601        File::create(output_path)?
602    };
603    let mut writer = PHSPWriter::from(output_file, &reader2.header)?;
604    let factor = sum_old_weight / sum_new_weight;
605    for mut record in reader2.map(|r| r.unwrap()) {
606        let r = (record.x_cm * record.x_cm + record.y_cm * record.y_cm).sqrt();
607        record.weight *= f(r) * factor;
608        writer.write(&record)?;
609    }
610    Ok(())
611}
612
613#[cfg(test)]
614mod tests {
615    use super::*;
616    use std::sync::atomic::{AtomicU64, Ordering};
617
618    static COUNTER: AtomicU64 = AtomicU64::new(0);
619
620    fn tmp_path(label: &str) -> std::path::PathBuf {
621        let n = COUNTER.fetch_add(1, Ordering::SeqCst);
622        let pid = std::process::id();
623        std::env::temp_dir().join(format!("beamdpr_test_{}_{}_{}.egsphsp1", label, pid, n))
624    }
625
626    fn make_record(latch: u32, energy: f32, x: f32, y: f32, zlast: Option<f32>) -> Record {
627        Record {
628            latch,
629            total_energy: energy,
630            x_cm: x,
631            y_cm: y,
632            x_cos: 0.1,
633            y_cos: 0.2,
634            weight: 1.0,
635            zlast,
636        }
637    }
638
639    fn write_phsp(path: &Path, header: &Header, records: &[Record]) {
640        let f = File::create(path).unwrap();
641        let mut writer = PHSPWriter::from(f, header).unwrap();
642        for r in records {
643            writer.write(r).unwrap();
644        }
645    }
646
647    #[test]
648    fn reweight_applies_radial_function_and_normalizes() {
649        let input = tmp_path("reweight_in");
650        let output = tmp_path("reweight_out");
651        let header = Header {
652            mode: *b"MODE0",
653            total_particles: 3,
654            total_photons: 3,
655            min_energy: 1.0,
656            max_energy: 1.0,
657            total_particles_in_source: 10.0,
658            record_size: 28,
659            using_zlast: false,
660        };
661        let mut records = vec![
662            make_record(0, 1.0, 0.0, 0.0, None),
663            make_record(0, 1.0, 1.0, 0.0, None),
664            make_record(0, 1.0, 2.0, 0.0, None),
665        ];
666        for r in records.iter_mut() {
667            r.weight = 2.0;
668            r.x_cos = 0.0;
669            r.y_cos = 0.0;
670        }
671        write_phsp(&input, &header, &records);
672
673        reweight(&input, &output, &|r| r + 1.0, 10, 5.0).unwrap();
674
675        let reader = PHSPReader::from(File::open(&output).unwrap()).unwrap();
676        let out: Vec<Record> = reader.map(|r| r.unwrap()).collect();
677        let _ = remove_file(&input);
678        let _ = remove_file(&output);
679
680        // sum_old = 6, sum_new = 2*1 + 2*2 + 2*3 = 12, factor = 0.5
681        // expected = original_weight * f(r) * factor = 2 * (r+1) * 0.5 = r + 1
682        let expected = [1.0_f32, 2.0, 3.0];
683        for (i, r) in out.iter().enumerate() {
684            assert!(
685                (r.weight - expected[i]).abs() < 1e-4,
686                "record {}: expected weight {}, got {}",
687                i,
688                expected[i],
689                r.weight
690            );
691        }
692    }
693
694    #[test]
695    fn sample_combine_uses_abs_energy_for_min_max() {
696        let input = tmp_path("sample_in");
697        let output = tmp_path("sample_out");
698        let header = Header {
699            mode: *b"MODE0",
700            total_particles: 3,
701            total_photons: 3,
702            min_energy: 0.5,
703            max_energy: 3.0,
704            total_particles_in_source: 10.0,
705            record_size: 28,
706            using_zlast: false,
707        };
708        let records = vec![
709            make_record(0, 0.5, 0.0, 0.0, None),
710            make_record(0, -3.0, 0.0, 0.0, None),
711            make_record(0, 1.5, 0.0, 0.0, None),
712        ];
713        write_phsp(&input, &header, &records);
714
715        sample_combine(&[&input], &output, 1.0, 0).unwrap();
716
717        let reader = PHSPReader::from(File::open(&output).unwrap()).unwrap();
718        let got = reader.header;
719        let _ = remove_file(&input);
720        let _ = remove_file(&output);
721
722        assert_eq!(got.total_particles, 3);
723        assert!(
724            (got.max_energy - 3.0).abs() < 1e-5,
725            "max_energy should be 3.0 (abs of -3.0), got {}",
726            got.max_energy
727        );
728        assert!(
729            (got.min_energy - 0.5).abs() < 1e-5,
730            "min_energy should be 0.5, got {}",
731            got.min_energy
732        );
733    }
734
735    #[test]
736    fn crossed_multiple_is_independent_of_charged() {
737        let mut r = make_record(0, 1.0, 0.0, 0.0, None);
738        r.latch = 1 << 30;
739        assert!(r.charged());
740        assert!(
741            !r.crossed_multiple(),
742            "crossed_multiple should be bit 31, distinct from charged (bit 30)"
743        );
744
745        r.latch = 1 << 31;
746        assert!(!r.charged());
747        assert!(r.crossed_multiple(), "bit 31 should mean crossed_multiple");
748    }
749
750    #[test]
751    fn transform_uses_original_z_cos_for_y_row() {
752        let mut record = make_record(0, 1.0, 0.0, 0.0, None);
753        record.x_cos = 0.6;
754        record.y_cos = 0.0;
755        let original_z_cos = record.z_cos();
756        let matrix = [
757            [0.5, 0.0, 0.0],
758            [0.0, 1.0, 1.0],
759            [0.0, 0.0, 1.0],
760        ];
761        record.transform(&matrix);
762        let expected_y_cos = 1.0 * original_z_cos;
763        assert!(
764            (record.y_cos - expected_y_cos).abs() < 1e-5,
765            "expected y_cos {} (from original z_cos), got {}",
766            expected_y_cos,
767            record.y_cos
768        );
769    }
770
771    #[test]
772    fn similar_to_detects_negative_difference() {
773        let a = make_record(0, 1.0, 0.0, 0.0, None);
774        let mut b = make_record(0, 1.0, 0.0, 0.0, None);
775        b.x_cm = 100.0;
776        assert!(
777            !a.similar_to(&b),
778            "records with x_cm differing by 100 should not be similar"
779        );
780    }
781
782    #[test]
783    fn mode2_writer_preserves_zlast() {
784        let path = tmp_path("mode2_zlast");
785        let header = Header {
786            mode: *b"MODE2",
787            total_particles: 2,
788            total_photons: 1,
789            min_energy: 0.5,
790            max_energy: 2.0,
791            total_particles_in_source: 100.0,
792            record_size: 32,
793            using_zlast: true,
794        };
795        let r1 = make_record(0, 1.0, 0.0, 0.0, Some(7.25));
796        let r2 = make_record(1 << 30, 2.0, 1.0, 1.0, Some(-3.5));
797        write_phsp(&path, &header, &[r1, r2]);
798
799        let reader = PHSPReader::from(File::open(&path).unwrap()).unwrap();
800        let records: Vec<Record> = reader.map(|r| r.unwrap()).collect();
801        let _ = remove_file(&path);
802
803        assert_eq!(records.len(), 2);
804        assert_eq!(records[0].zlast, Some(7.25), "first zlast not preserved");
805        assert_eq!(records[1].zlast, Some(-3.5), "second zlast not preserved");
806    }
807}
egsphsp/lib.rs

egsphsp/
lib.rs
