[WIP] Draft of generalized RustBCA python wrapper(s) (#324)

* Draft of generalized RustBCA python wrapper(s)

* Added tests of rustbca_py to test_different_options.py

* Implemented better python error handling

* Add different interaction potentials to scattering_integrals()

* used macros to streamline both lib.rs geometry types and main.rs geometry types

* Attempt to fix module imports across lib, main

Author: drobnyjt

Cargo.toml

@@ -29,6 +29,7 @@ rcpr = { git = "https://github.com/drobnyjt/rcpr", optional = true}
 ndarray = {version = "0.17.2", features = ["serde"], optional = true}
 parry3d-f64 = {optional = true, version="0.2.0"}
 pyo3 = {version = "0.29.0", optional=true}
+pythonize = {version = "0.29.0", optional=true}
 
 [dev-dependencies]
 float-cmp = "0.10.0"
@@ -45,5 +46,5 @@ cpr_rootfinder = ["rcpr"]
 distributions = ["ndarray"]
 no_list_output = []
 parry3d = ["parry3d-f64"]
-python = ["pyo3"]
+python = ["pyo3", "pythonize"]
 extended_max_z = []

examples/make_input_file_and_run.py

@@ -255,6 +255,17 @@
     'geometry_input': geometry_0D
 }
 
+input_data['options']['name'] = 'rustbca_input_file'
+rustbca_py(input_data, mode)
+s = np.genfromtxt('rustbca_input_filesputtered.output', delimiter=',')
+
+arrays = rustbca_local_py(input_data, mode)
+sputtered = arrays['sputtered']
+
+np.testing.assert_approx_equal(s[0, 2], np.array(arrays['energy'])[sputtered][0])
+
+input_data['options']['name'] = 'input_file'
+
 # Attempt to cleanup line endings
 input_string = dumps(input_data).replace('\r', '')
 with  open('examples/input_file.toml', 'w') as input_file:
@@ -265,6 +276,9 @@
 
 # Read output files - ensure arrays are at least 2D for indexing
 sputtered = np.atleast_2d(np.genfromtxt('input_filesputtered.output', delimiter=','))
+
+np.testing.assert_approx_equal(sputtered[0, 2], s[0, 2])
+
 reflected = np.atleast_2d(np.genfromtxt('input_filereflected.output', delimiter=','))
 implanted = np.atleast_2d(np.genfromtxt('input_filedeposited.output', delimiter=','))
 

examples/test_different_options.py

@@ -51,6 +51,8 @@
 a TOML file using tomlkit.
 
 It runs the input file with cargo run --release and reads the output files.
+
+It also runs with rustbca_py and ensures the output files are identical.
 '''
 
 def run_test(
@@ -275,6 +277,19 @@ def run_test(
     reflected = np.atleast_2d(np.genfromtxt(f'input_file_{index}reflected.output', delimiter=','))
     implanted = np.atleast_2d(np.genfromtxt(f'input_file_{index}deposited.output', delimiter=','))
 
+    input_data['options']['name'] = f'python_{index}'
+    if run_sim:
+        rustbca_py(input_data, geometry_mode=mode)
+
+    # Read output files - ensure arrays are at least 2D for indexing
+    sputtered_py = np.atleast_2d(np.genfromtxt(f'python_{index}sputtered.output', delimiter=','))
+    reflected_py = np.atleast_2d(np.genfromtxt(f'python_{index}reflected.output', delimiter=','))
+    implanted_py = np.atleast_2d(np.genfromtxt(f'python_{index}deposited.output', delimiter=','))
+
+    np.testing.assert_allclose(sputtered, sputtered_py)
+    np.testing.assert_allclose(reflected, reflected_py)
+    np.testing.assert_allclose(implanted, implanted_py)
+
     return sputtered, reflected, implanted
 
 num_bins = 75

examples/test_scattering_integrals.py

@@ -10,7 +10,7 @@
 from formulas import *
 
 energies = np.logspace(0, 4, 4)
-impact_parameters = np.logspace(-3, 3, 100)
+impact_parameters = np.logspace(-3, 2, 100)
 
 ion = helium
 target = boron
@@ -24,23 +24,28 @@
 
 linestyles = ['-', '--', ':', '-.']
 
-for linestyle, energy in zip(linestyles, energies):
-    gm = np.zeros(100)
-    gl = np.zeros(100)
-    mw = np.zeros(100)
-    magic = np.zeros(100)
-    for index, p in enumerate(impact_parameters):
-        gm[index], gl[index], mw[index], magic[index] = scattering_integrals(Za, Zb, Ma, Mb, energy, p)
-    plt.semilogx(impact_parameters, gm, label=f'Gauss-Mehler, E={np.round(energy/1000)} keV', linestyle=linestyle)
-    plt.semilogx(impact_parameters, gl, label=f'Gauss-Legendre, E={np.round(energy/1000)} keV', linestyle=linestyle)
-    plt.semilogx(impact_parameters, mw, label=f'Mendenhall-Weller, E={np.round(energy/1000)} keV', linestyle=linestyle)
-    plt.semilogx(impact_parameters, magic, label=f'MAGIC, E={np.round(energy/1000)} keV', linestyle=linestyle)
-    plt.gca().set_prop_cycle(None)
-
-    np.testing.assert_allclose(gm, gl, atol=5e-3) # 0.5% seems reasonable? max is ~0.3%
-    np.testing.assert_allclose(gm, mw, atol=5e-3)
-    np.testing.assert_allclose(mw, gl, atol=5e-3)
-    plt.legend()
-    plt.xlabel('p [A]')
-    plt.ylabel('theta [rad]')
+for potential in ["KR_C", "MOLIERE", "ZBL"]:
+    plt.figure()
+    plt.title(f'Scattering Angles for {potential}')
+    for linestyle, energy in zip(linestyles, energies):
+        gm = np.zeros(100)
+        gl = np.zeros(100)
+        mw = np.zeros(100)
+        magic = np.zeros(100)
+        for index, p in enumerate(impact_parameters):
+            gm[index], gl[index], mw[index], magic[index] = scattering_integrals(Za, Zb, Ma, Mb, energy, p, interaction_potential=potential)
+
+        plt.semilogx(impact_parameters, gm, label=f'Gauss-Mehler, E={np.round(energy/1000, 3)} keV', linestyle=linestyle)
+        plt.semilogx(impact_parameters, gl, label=f'Gauss-Legendre, E={np.round(energy/1000, 3)} keV', linestyle=linestyle)
+        plt.semilogx(impact_parameters, mw, label=f'Mendenhall-Weller, E={np.round(energy/1000, 3)} keV', linestyle=linestyle)
+        plt.semilogx(impact_parameters, magic, label=f'MAGIC, E={np.round(energy/1000, 3)} keV', linestyle=linestyle)
+        plt.gca().set_prop_cycle(None)
+
+        np.testing.assert_allclose(gm, gl, atol=5e-3) # 0.5% seems reasonable? max is ~0.3%
+        np.testing.assert_allclose(gm, mw, atol=5e-3)
+        np.testing.assert_allclose(mw, gl, atol=5e-3)
+        plt.legend()
+        plt.xlabel('p [A]')
+        plt.ylabel('theta [rad]')
+
 if show_plots: plt.show()

src/input.rs

@@ -381,9 +381,7 @@ impl Options {
     }
 }
 
-pub fn input<T: Geometry>(input_file: String) -> (Vec<particle::ParticleInput>, material::Material<T>, Options, OutputUnits)
-where <T as Geometry>::InputFileFormat: Deserialize<'static> + 'static {
-
+pub fn read_input_file<T: Geometry>(input_file: String) -> <T as Geometry>::InputFileFormat {
     //Read input file, convert to string, and open with toml
     let mut input_toml = String::new();
     let mut file = OpenOptions::new()
@@ -394,7 +392,19 @@ where <T as Geometry>::InputFileFormat: Deserialize<'static> + 'static {
         .unwrap_or_else(|_| panic!("Input errror: could not open input file {}.", &input_file));
     file.read_to_string(&mut input_toml).context("Could not convert TOML file to string.").unwrap();
 
-    let input: <T as Geometry>::InputFileFormat = InputFile::new(&input_toml);
+    InputFile::new(&input_toml)
+}
+
+pub fn input<T: Geometry>(input_file: String) -> (Vec<particle::ParticleInput>, material::Material<T>, Options, OutputUnits)
+where <T as Geometry>::InputFileFormat: Deserialize<'static> + 'static {
+
+    let input: <T as Geometry>::InputFileFormat = read_input_file::<T>(input_file);
+
+    process_input_file(input)
+
+}
+
+pub fn process_input_file<T: Geometry>(input: <T as Geometry>::InputFileFormat) -> (Vec<particle::ParticleInput>, material::Material<T>, Options, OutputUnits) {
 
     //Unpack toml information into structs
     let options = (*input.get_options()).clone();

src/lib.rs

@@ -9,10 +9,10 @@ use std::alloc::{dealloc, Layout};
 use std::mem::align_of;
 
 //Parallelization - currently only used in python library functions
+//#[cfg(feature = "python")]
+//use rayon::ThreadPoolBuilder;
 #[cfg(feature = "python")]
-use rayon::prelude::*;
-#[cfg(feature = "python")]
-use rayon::*;
+use rayon::iter::{IndexedParallelIterator, ParallelExtend, IntoParallelIterator, ParallelIterator};
 
 //Error handling crate
 use anyhow::{Result, Context, anyhow};
@@ -53,6 +53,10 @@ use std::f64::consts::SQRT_2;
 use pyo3::prelude::*;
 #[cfg(feature = "python")]
 use pyo3::types::*;
+#[cfg(feature = "python")]
+use pythonize::*;
+#[cfg(feature = "python")]
+use pyo3::exceptions::{PyValueError, PyRuntimeError};
 
 //Load internal modules
 pub mod material;
@@ -68,9 +72,7 @@ pub mod consts;
 pub mod structs;
 pub mod sphere;
 pub mod math;
-
-#[cfg(feature = "parry3d")]
-pub mod parry;
+pub mod physics;
 
 pub use crate::enums::*;
 pub use crate::consts::*;
@@ -81,6 +83,10 @@ pub use crate::geometry::{Geometry, GeometryElement, Mesh0D, Mesh1D, Mesh2D};
 pub use crate::sphere::{Sphere, SphereInput, InputSphere};
 pub use crate::math::*;
 pub use crate::material::*;
+pub use crate::physics::*;
+
+#[cfg(feature = "parry3d")]
+pub mod parry;
 
 #[cfg(feature = "parry3d")]
 pub use crate::parry::{ParryBall, ParryBallInput, InputParryBall, ParryTriMesh, ParryTriMeshInput, InputParryTriMesh};
@@ -139,6 +145,12 @@ mod libRustBCA {
 
     #[pymodule_export]
     use super::scattering_integrals;
+
+    #[pymodule_export]
+    use super::rustbca_py;
+
+    #[pymodule_export]
+    use super::rustbca_local_py;
 }
 
 #[derive(Debug)]
@@ -2131,7 +2143,6 @@ pub fn compound_reflection_coefficient<'py>(ion: &Bound<'py, PyDict>, targets: V
 
                 let mut residue = residue.lock().unwrap();
                 *residue = *residue + residue_part;
-
             }
         }
     });
@@ -2161,20 +2172,118 @@ fn moller_knuth_two_sum(a: f64, b: f64) -> (f64, f64) {
     let r = delta_a + delta_b;
     (s, r)
 }
+
 #[cfg(feature = "python")]
 #[pyfunction]
-#[pyo3(signature = (Za, Zb, Ma, Mb, E0, p, n_gl_points=100))]
-fn scattering_integrals(Za: f64, Zb: f64, Ma: f64, Mb: f64, E0: f64, p: f64, n_gl_points: usize) -> (f64, f64, f64, f64) {
+#[pyo3(signature = (Za, Zb, Ma, Mb, E0, p, n_gl_points=100, interaction_potential="KR_C"))]
+fn scattering_integrals(Za: f64, Zb: f64, Ma: f64, Mb: f64, E0: f64, p: f64, n_gl_points: usize, interaction_potential: &str) -> PyResult<(f64, f64, f64, f64)> {
     let E0 = E0*EV;
     let p = p*ANGSTROM;
 
-    let x0_newton = bca::newton_rootfinder(Za, Zb, Ma, Mb, E0, p, InteractionPotential::KR_C, 1000, 1E-12).unwrap();
+    let potential = match interaction_potential {
+        "KR_C" => InteractionPotential::KR_C,
+        "LENZ_JENSEN" => InteractionPotential::LENZ_JENSEN,
+        "MOLIERE" => InteractionPotential::MOLIERE,
+        "ZBL" => InteractionPotential::ZBL,
+        _ => return Err(PyValueError::new_err(format!("Unimplemented interaction potential {}; try 'KR_C'", interaction_potential)))
+    };
+
+    let x0_newton = bca::newton_rootfinder(Za, Zb, Ma, Mb, E0, p, potential, 1000, 1E-12).map_err(
+        |error| PyRuntimeError::new_err(format!("Rootfinder failed to find distance of closest approach; check input values."))
+    )?;
 
     //Compute center of mass deflection angle with each algorithm
-    let theta_gm = bca::gauss_mehler(Za, Zb, Ma, Mb, E0, p, x0_newton, InteractionPotential::KR_C, n_gl_points);
-    let theta_gl = bca::gauss_legendre(Za, Zb, Ma, Mb, E0, p, x0_newton, InteractionPotential::KR_C);
-    let theta_mw = bca::mendenhall_weller(Za, Zb, Ma, Mb, E0, p, x0_newton, InteractionPotential::KR_C);
-    let theta_magic = bca::magic(Za, Zb, Ma, Mb, E0, p, x0_newton, InteractionPotential::KR_C);
+    let theta_gm = bca::gauss_mehler(Za, Zb, Ma, Mb, E0, p, x0_newton, potential, n_gl_points);
+    let theta_gl = bca::gauss_legendre(Za, Zb, Ma, Mb, E0, p, x0_newton, potential);
+    let theta_mw = bca::mendenhall_weller(Za, Zb, Ma, Mb, E0, p, x0_newton, potential);
+    let theta_magic = bca::magic(Za, Zb, Ma, Mb, E0, p, x0_newton, potential);
+
+    Ok((theta_gm, theta_gl, theta_mw, theta_magic))
+}
+#[cfg(feature = "python")]
+macro_rules! geometry_typed_loops {
+    ($geometry_type:ty, $input:expr, $python:expr) => {
+        {
+            let input: <$geometry_type as geometry::Geometry>::InputFileFormat = depythonize(&$input).unwrap();
+            let (particle_input_array, material, options, output_units) = input::process_input_file(input);
+            let pool = rayon::ThreadPoolBuilder::new().num_threads(options.num_threads).build().unwrap();
+            pool.install( ||
+                physics::physics_loop::<$geometry_type>(particle_input_array, material, options, output_units)
+            );
+            Ok(())
+        }
+    }
+}
+
+#[cfg(feature = "python")]
+#[pyfunction]
+#[pyo3(signature=(input, geometry_mode="1D"))]
+fn rustbca_py<'py>(python: Python<'py>, input: &Bound<'py, PyDict>, geometry_mode: &str) -> PyResult<()> {
+    match geometry_mode {
+        "0D" => geometry_typed_loops!(Mesh0D, input, python),
+        "1D" => geometry_typed_loops!(Mesh1D, input, python),
+        "2D" => geometry_typed_loops!(Mesh2D, input, python),
+        "HOMOGENEOUS2D" => geometry_typed_loops!(Mesh2D, input, python),
+        "SPHERE" => geometry_typed_loops!(Sphere, input, python),
+        #[cfg(feature="parry3d")]
+        "BALL" => geometry_typed_loops!(ParryBall, input, python),
+        #[cfg(feature="parry3d")]
+        "TRIMESH" => geometry_typed_loops!(ParryTriMesh, input, python),
+       _ => Err(PyValueError::new_err(format!("Input Error: Unimplemented geometry mode {}; try '1D'", geometry_mode)))
+    }
+}
 
-    (theta_gm, theta_gl, theta_mw, theta_magic)
+/*
+Notes on macros - this is the first I have written, so I'm taking notes here as I go.
+macro_rules! makes a macro - here, the macro is called geometry_types_silent_loops
+macros pattern match an argument and replace it with anything you want
+I want it to take a tuple of a string (e.g., "1D") and a type (e.g., Mesh1D)
+and plop those into corresponding match arms.
+The first line tells the macro to expect an argument with that pattern.
+arguments are $<name>:<designator>. Designators:
+block
+expr is used for expressions
+ident is used for variable/function names
+item
+literal is used for literal constants
+pat (pattern)
+path
+stmt (statement)
+tt (token tree)
+ty (type)
+vis (visibility qualifier)
+*/
+#[cfg(feature = "python")]
+macro_rules! geometry_typed_silent_loops {
+    ($geometry_type:ty, $input:expr, $python:expr) => {
+        {
+            let input: <$geometry_type as geometry::Geometry>::InputFileFormat = depythonize(&$input).unwrap();
+            let (particle_input_array, material, options, output_units) = input::process_input_file(input);
+            let pool = rayon::ThreadPoolBuilder::new().num_threads(options.num_threads).build().unwrap();
+            let finished_particles = pool.install( ||
+                physics::silent_physics_loop::<$geometry_type>(particle_input_array, material, options, output_units.clone())
+            );
+            let finished_particles_container = physics::process_finished_particles_to_arrays(finished_particles, output_units);
+            Ok(pythonize($python, &finished_particles_container)?)
+        }
+    }
+}
+
+#[cfg(feature = "python")]
+#[pyfunction]
+#[pyo3(signature=(input, geometry_mode="1D"))]
+fn rustbca_local_py<'py>(python: Python<'py>, input: &Bound<'py, PyDict>, geometry_mode: &str) -> PyResult<Bound<'py, PyAny>> {
+
+    match geometry_mode {
+        "0D" => geometry_typed_silent_loops!(Mesh0D, input, python),
+        "1D" => geometry_typed_silent_loops!(Mesh1D, input, python),
+        "2D" => geometry_typed_silent_loops!(Mesh2D, input, python),
+        "HOMOGENEOUS2D" => geometry_typed_silent_loops!(Mesh2D, input, python),
+        "SPHERE" => geometry_typed_silent_loops!(Sphere, input, python),
+        #[cfg(feature="parry3d")]
+        "BALL" => geometry_typed_silent_loops!(ParryBall, input, python),
+        #[cfg(feature="parry3d")]
+        "TRIMESH" => geometry_typed_silent_loops!(ParryTriMesh, input, python),
+        _ => Err(PyValueError::new_err(format!("Input Error: Unimplemented geometry mode {}; try '1D'", geometry_mode)))
+    }
 }