MPI functionality for the Photon struct

Cargo.toml

@@ -14,6 +14,7 @@ exclude = ["docs/", "input/", "output/"]
 
 [features]
 default = []
+mpi = ["dep:mpi", "dep:memoffset"]
 
 [dependencies]
 # ARCTK Dependencies
@@ -49,7 +50,11 @@ lidrs = "0.2.*"
 
 # Formats for the File I/O Library. 
 json5 = "0.4.*"
-netcdf = "0.8.1"
+netcdf = "^0.8.1"
+
+# Optional MPI dependencies
+mpi = { version = "0.7.0", optional = true }
+memoffset = { version = "0.9.0", optional = true }
 
 [dev-dependencies]
 tempfile = "3.2.*"
@@ -64,4 +69,9 @@ rustdoc-args = [ "--html-in-header", "./src/docs-header.html" ]
 [lib]
 
 [[bin]]
-name = "mcrt"
+name = "mcrt"
+
+[[bin]]
+name = "mpi_tests"
+path = "src/bin/tests/mpi_tests.rs"
+required-features = ["mpi"]

src/bin/tests/mpi_tests.rs

@@ -0,0 +1,68 @@
+//! Simple integration tests for MPI functionality
+
+use mpi::{
+    point_to_point::send_receive_replace_into,
+    traits::*,
+};
+use Aetherus::{
+    math::{Formula, Probability, Point3, Dir3},
+    geom::{Emitter, Ray},
+    phys::{Light, Material, Photon, PhotonBuf},
+};
+use rand;
+
+
+/// Main program function
+fn main() {
+
+    // Init MPI communicator
+    let comm = mpi::initialize().unwrap();
+    let world = comm.world();
+    let size = world.size();
+    let rank = world.rank();
+
+    if size != 2 {
+        panic!("Test case only works with 2 MPI ranks");
+    }
+
+    // Init photon on home rank
+    if rank == 0 {
+
+        // Create light source
+        let mut rng = rand::thread_rng();
+        let ray = Ray::new(Point3::new(0.0, 0.0, 0.0), Dir3::new(1.0, 0.0, 0.0));
+        let emitter = Emitter::new_beam(ray.clone());
+        let mat = get_air_material();
+        let light = Light::new(1.0, emitter, Probability::new_point(1.0), &mat);
+
+        // Now emit a photon
+        let photon = light.emit(&mut rng, 1.0);
+
+        // Send photon across MPI rank
+        let photon_buffer = PhotonBuf::new(&photon);
+        world.process_at_rank(1).send(&photon_buffer);
+
+    } else {
+        let msg = world.process_at_rank(0).receive::<PhotonBuf>().0;
+
+        let phot_return = msg.as_photon();
+
+        assert_eq!(phot_return.ray().pos(), &Point3::new(0.0, 0.0, 0.0));
+        assert_eq!(phot_return.ray().dir(), &Dir3::new(1.0, 0.0, 0.0));
+        assert_eq!(phot_return.weight(), 1.0);
+        assert_eq!(phot_return.wavelength(), 1.0);
+        assert_eq!(phot_return.power(), 1.0);
+
+    }
+
+}
+
+fn get_air_material() -> Material {
+    Material::new(
+        Formula::Constant { c: 1.0 }, 
+        Formula::Constant { c: 1.0e-6 }, 
+        None, 
+        None, 
+        Formula::Constant { c: 0.1 }
+    )
+}

src/phys/photon.rs

@@ -1,6 +1,14 @@
 //! Photon particle.
+use crate::{access, clone, geom::Ray, math::{Dir3, Point3}};
 
-use crate::{access, clone, geom::Ray};
+#[cfg(feature = "mpi")]
+use mpi::{
+    datatype::{UncommittedUserDatatype, UserDatatype},
+    traits::*,
+    Address,
+};
+#[cfg(feature = "mpi")]
+use memoffset::{offset_of};
 
 /// Photon.
 #[derive(Clone)]
@@ -42,3 +50,116 @@ impl Photon {
         self.weight = 0.0;
     }
 }
+
+/// Photon reconstructed into raw data for MPI buffer.
+#[cfg(feature = "mpi")]
+#[derive(Clone)]
+pub struct PhotonBuf {
+    /// Ray of travel broken down to component arrays
+    pub ray_pos: [f64; 3],
+    pub ray_dir: [f64; 3],
+    /// Statistical weight.
+    pub weight: f64,
+    /// Wavelength (m).
+    pub wavelength: f64,
+    /// Power (J/s).
+    pub power: f64,
+}
+
+#[cfg(feature = "mpi")]
+impl PhotonBuf {
+
+    /// Construct a new instance.
+    #[inline]
+    #[must_use]
+    pub fn new(photon: &Photon) -> Self {
+        Self {
+            ray_pos: [photon.ray().pos().x(), photon.ray().pos().y(), photon.ray().pos().z()],
+            ray_dir: [photon.ray().dir().x(), photon.ray().dir().y(), photon.ray().dir().z()],
+            weight: photon.weight(),
+            wavelength: photon.wavelength(),
+            power: photon.power(),
+        }
+    }
+
+    /// Convert photon buffer back to Photon struct
+    #[inline]
+    pub fn as_photon(self) -> Photon {
+        let ray = Ray::new(
+            Point3::new(self.ray_pos[0], self.ray_pos[1], self.ray_pos[2]),
+            Dir3::new(self.ray_dir[0], self.ray_dir[1], self.ray_dir[2]));
+        return Photon::new(ray, self.wavelength, self.power);
+    }
+
+}
+
+#[cfg(feature = "mpi")]
+unsafe impl Equivalence for PhotonBuf {
+    type Out = UserDatatype;
+    fn equivalent_datatype() -> Self::Out {
+        UserDatatype::structured(
+            &[1, 1, 1, 1, 1],
+            &[
+                offset_of!(PhotonBuf, ray_pos) as Address,
+                offset_of!(PhotonBuf, ray_dir) as Address,
+                offset_of!(PhotonBuf, weight) as Address,
+                offset_of!(PhotonBuf, wavelength) as Address,
+                offset_of!(PhotonBuf, power) as Address,
+            ],
+            &[
+                UncommittedUserDatatype::contiguous(3, &f64::equivalent_datatype()).as_ref(),
+                UncommittedUserDatatype::contiguous(3, &f64::equivalent_datatype()).as_ref(),
+                f64::equivalent_datatype().into(),
+                f64::equivalent_datatype().into(),
+                f64::equivalent_datatype().into(),
+            ],
+        )
+    }
+}
+
+
+#[cfg(test)]
+mod tests {
+    use super::Photon;
+    use crate::geom::Ray;
+    use crate::math::{Dir3, Point3};
+    use assert_approx_eq::assert_approx_eq;
+    use std::f64;
+
+    #[cfg(feature = "mpi")]
+    use super::PhotonBuf;
+
+    /// Check that the creation and accessing code is working correctly.
+    #[test]
+    #[cfg(feature = "mpi")]
+    fn buf_init_test() {
+        let ray = Ray::new(Point3::new(0.0, 0.0, 0.0), Dir3::new(1.0, 1.0, 1.0));
+        let phot = Photon::new(ray, 500.0, 10.0);
+
+        let phot_buf  = PhotonBuf::new(&phot);
+        // Check that we get the correct
+        assert_eq!(phot_buf.ray_pos, [0.0, 0.0, 0.0]);
+        //assert_eq!(phot_buf.ray_dir, [1.0, 1.0, 1.0]);
+        assert_eq!(phot_buf.weight, 1.0);
+        assert_eq!(phot_buf.wavelength, 500.0);
+        assert_eq!(phot_buf.power, 10.0);
+    }
+
+    /// Check that arrays destruct correctly
+    #[test]
+    #[cfg(feature = "mpi")]
+    fn buf_as_photon_test() {
+        let ray = Ray::new(Point3::new(0.0, 0.0, 0.0), Dir3::new(1.0, 1.0, 1.0));
+        let phot = Photon::new(ray, 500.0, 10.0);
+        let phot_buf = PhotonBuf::new(&phot);
+
+        let phot_return = phot_buf.as_photon();
+
+        assert_eq!(phot.ray.pos(), phot_return.ray.pos());
+        assert_eq!(phot.ray.dir(), phot_return.ray.dir());
+        assert_eq!(phot.weight(), phot_return.weight());
+        assert_eq!(phot.wavelength(), phot_return.wavelength());
+        assert_eq!(phot.power(), phot_return.power());
+    }
+
+}