Parallelized RustBCA python function for 0D compounds with tracking

examples/test_rustbca.py

@@ -12,6 +12,7 @@
 
 def main():
 
+
     #test rotation to and from RustBCA coordinates
 
     #nx, ny, nz is the normal vector (out of surface)
@@ -153,7 +154,7 @@ def main():
     print('Data processing complete.')
     print(f'RustBCA Y: {len(sputtered[:, 0])/number_ions} Yamamura Y: {yamamura_yield}')
     print(f'RustBCA R: {len(reflected[:, 0])/number_ions} Thomas R: {thomas}')
-    print(f'Time per ion: {delta_time/number_ions*1e3} us/{ion["symbol"]}')
+    print(f'Time per ion: {delta_time/number_ions} s/{ion["symbol"]}')
 
     #Next up is the layered target version. I'll add a 50 Angstrom layer of W-H to the top of the target.
 
@@ -170,14 +171,15 @@ def main():
 
     print(f'Running RustBCA for {number_ions} {ion["symbol"]} ions on {target["symbol"]} with hydrogenated layer at {energies_eV[0]/1000.} keV...')
     print(f'This may take several minutes.')
-    #Not the different argument order; when a breaking change is due, this will
+    #Note the different argument order; when a breaking change is due, this will
     #be back-ported to the other bindings as well for consistency.
     output, incident, stopped = compound_bca_list_1D_py(
         ux, uy, uz, energies_eV, [ion['Z']]*number_ions,
         [ion['m']]*number_ions, [ion['Ec']]*number_ions, [ion['Es']]*number_ions, [target['Z'], 1.0], [target['m'], 1.008],
         [target['Ec'], 1.0], [target['Es'], 1.5], [target['Eb'], 0.0], [[target['n']/10**30, target['n']/10**30], [target['n']/10**30, 0.0]], [50.0, 1e6]
     )
 
+
     output = np.array(output)
 
     Z = output[:, 0]
@@ -198,6 +200,80 @@ def main():
     plt.plot([50.0, 50.0], [0.0, np.max(heights)*1.1])
     plt.gca().set_ylim([0.0, np.max(heights)*1.1])
 
+    number_ions = 10000
+
+    #1 keV is above the He on W sputtering threshold of ~150 eV
+    energies_eV = 1000.0*np.ones(number_ions)
+
+    #Working with angles of exactly 0 is problematic due to gimbal lock
+    angle = 0.0001
+
+    #In RustBCA's 0D geometry, +x -> into the surface
+    ux = np.cos(angle*np.pi/180.)*np.ones(number_ions)
+    uy = np.sin(angle*np.pi/180.)*np.ones(number_ions)
+    uz = np.zeros(number_ions)
+
+    Z1 = np.array([ion['Z']]*number_ions)
+    m1 = np.array([ion['m']]*number_ions)
+    Ec1 = np.array([ion['Ec']]*number_ions)
+    Es1 = np.array([ion['Es']]*number_ions)
+
+    print(f'Running tracked RustBCA for {number_ions} {ion["symbol"]} ions on {target["symbol"]} with 1% {ion["symbol"]} at {energies_eV[0]/1000.} keV...')
+    print(f'This may take several minutes.')
+
+    start = time.time()
+    # Note that simple_bca_list_py expects number densities in 1/Angstrom^3
+    output, incident, incident_index = compound_bca_list_tracked_py(
+        energies_eV, ux, uy, uz, Z1, m1, Ec1, Es1,
+        [target['Z'], ion['Z']], [target['m'], ion['m']],
+        [target['Ec'], ion['Ec']], [target['Es'], ion['Es']], [target['n']/10**30, 0.01*target['n']/10**30],
+        [target['Eb'], 0.0])
+    stop = time.time()
+    delta_time = stop - start
+
+    output = np.array(output)
+    Z = output[:, 0]
+    m = output[:, 1]
+    E = output[:, 2]
+    x = output[:, 3]
+    y = output[:, 4]
+    z = output[:, 5]
+    ux = output[:, 6]
+    uy = output[:, 7]
+    uz = output[:, 8]
+
+    #For the python bindings, these conditionals can be used to distinguish
+    #between sputtered, reflected, and implanted particles in the output list
+    sputtered = output[np.logical_and(Z == target['Z'], E > 0), :]
+    reflected = output[np.logical_and(Z == ion['Z'], x < 0), :]
+    implanted = output[np.logical_and(Z == ion['Z'], x > 0), :]
+
+    plt.figure(1)
+    plt.title(f'Sputtered {target["symbol"]} Energy Distribution')
+    plt.xlabel('E [eV]')
+    plt.ylabel('f(E) [A.U.]')
+    plt.hist(sputtered[:, 2], bins=100, density=True, histtype='step')
+
+    plt.figure(2)
+    plt.title(f'Implanted {ion["symbol"]} Depth Distribution')
+    plt.xlabel('x [A]')
+    plt.ylabel('f(x) [A.U.]')
+    plt.hist(implanted[:, 3], bins=100, density=True, histtype='step')
+
+    thomas = thomas_reflection(ion, target, energies_eV[0])
+    yamamura_yield = yamamura(ion, target, energies_eV[0])
+
+    print('Data processing complete.')
+    print(f'RustBCA Y: {len(sputtered[:, 0])/number_ions} Yamamura Y: {yamamura_yield}')
+    print(f'RustBCA R: {len(reflected[:, 0])/number_ions} Thomas R: {thomas}')
+    print(f'Time per ion: {delta_time/number_ions} s/{ion["symbol"]}')
+
+    plt.figure()
+    plt.plot(incident_index)
+    plt.xlabel('Particle number')
+    plt.ylabel('Particle index')
+    plt.legend(['Incident', 'Indicies'])
+
     plt.show()
 
 if __name__ == '__main__':

src/lib.rs

@@ -89,6 +89,7 @@ pub fn libRustBCA(py: Python, m: &PyModule) -> PyResult<()> {
     m.add_function(wrap_pyfunction!(simple_bca_py, m)?)?;
     m.add_function(wrap_pyfunction!(simple_bca_list_py, m)?)?;
     m.add_function(wrap_pyfunction!(compound_bca_list_py, m)?)?;
+    m.add_function(wrap_pyfunction!(compound_bca_list_tracked_py, m)?)?;
     m.add_function(wrap_pyfunction!(compound_bca_list_1D_py, m)?)?;
     m.add_function(wrap_pyfunction!(sputtering_yield, m)?)?;
     m.add_function(wrap_pyfunction!(reflection_coefficient, m)?)?;
@@ -918,6 +919,136 @@ pub fn compound_bca_list_py(energies: Vec<f64>, ux: Vec<f64>, uy: Vec<f64>, uz:
     (total_output, incident)
 }
 
+#[cfg(feature = "python")]
+///compound_tagged_bca_list_py(ux, uy,  uz, energy, Z1, m1, Ec1, Es1, Z2, m2, Ec2, Es2, n2, Eb2)
+/// runs a BCA simulation for a list of particles and outputs a list of sputtered, reflected, and implanted particles.
+/// Args:
+///    energies (list(f64)): initial ion energies in eV.
+///    ux (list(f64)): initial ion directions x. ux != 0.0 to avoid gimbal lock
+///    uy (list(f64)): initial ion directions y.
+///    uz (list(f64)): initial ion directions z.
+///    Z1 (list(f64)): initial ion atomic numbers.
+///    m1 (list(f64)): initial ion masses in amu.
+///    Ec1 (list(f64)): ion cutoff energies in eV. If ion energy < Ec1, it stops in the material.
+///    Es1 (list(f64)): ion surface binding energies. Assumed planar.
+///    Z2 (list(f64)): target material species atomic numbers.
+///    m2 (list(f64)): target material species masses in amu.
+///    Ec2 (list(f64)): target material species cutoff energies in eV. If recoil energy < Ec2, it stops in the material.
+///    Es2 (list(f64)): target species surface binding energies. Assumed planar.
+///    n2 (list(f64)): target material species atomic number densities in inverse cubic Angstroms.
+///    Eb2 (list(f64)): target material species bulk binding energies in eV.
+/// Returns:
+///    output (NX9 list of f64): each row in the list represents an output particle (implanted,
+///    sputtered, or reflected). Each row consists of:
+///      [Z, m (amu), E (eV), x, y, z, (angstrom), ux, uy, uz]
+///    incident (list(bool)): whether each row of output was an incident ion or originated in the target
+///    incident_index (list(usize)): index of incident particle that caused this particle to be emitted
+#[pyfunction]
+pub fn compound_bca_list_tracked_py(energies: Vec<f64>, ux: Vec<f64>, uy: Vec<f64>, uz: Vec<f64>, Z1: Vec<f64>, m1: Vec<f64>, Ec1: Vec<f64>, Es1: Vec<f64>, Z2: Vec<f64>, m2: Vec<f64>, Ec2: Vec<f64>, Es2: Vec<f64>, n2: Vec<f64>, Eb2: Vec<f64>) -> (Vec<[f64; 9]>, Vec<bool>, Vec<usize>) {
+    let mut total_output = vec![];
+    let mut incident = vec![];
+    let mut incident_index = vec![];
+    let num_species_target = Z2.len();
+    let num_incident_ions = energies.len();
+
+    assert_eq!(ux.len(), num_incident_ions, "Input error: list of x-directions is not the same length as list of incident energies.");
+    assert_eq!(uy.len(), num_incident_ions, "Input error: list of y-directions is not the same length as list of incident energies.");
+    assert_eq!(uz.len(), num_incident_ions, "Input error: list of z-directions is not the same length as list of incident energies.");
+    assert_eq!(Z1.len(), num_incident_ions, "Input error: list of incident atomic numbers is not the same length as list of incident energies.");
+    assert_eq!(m1.len(), num_incident_ions, "Input error: list of incident atomic masses is not the same length as list of incident energies.");
+    assert_eq!(Es1.len(), num_incident_ions, "Input error: list of incident surface binding energies is not the same length as list of incident energies.");
+    assert_eq!(Ec1.len(), num_incident_ions, "Input error: list of incident cutoff energies is not the same length as list of incident energies.");
+
+    assert_eq!(m2.len(), num_species_target, "Input error: list of target atomic masses is not the same length as atomic numbers.");
+    assert_eq!(Ec2.len(), num_species_target, "Input error: list of target cutoff energies is not the same length as atomic numbers.");
+    assert_eq!(Es2.len(), num_species_target, "Input error: list of target surface binding energies is not the same length as atomic numbers.");
+    assert_eq!(Eb2.len(), num_species_target, "Input error: list of target bulk binding energies is not the same length as atomic numbers.");
+    assert_eq!(n2.len(), num_species_target, "Input error: list of target number densities is not the same length as atomic numbers.");
+
+    let options = Options::default_options(true);
+    //options.high_energy_free_flight_paths = true;
+
+    let x = -2.*(n2.iter().sum::<f64>()*10E30).powf(-1./3.);
+    let y = 0.0;
+    let z = 0.0;
+
+    let material_parameters = material::MaterialParameters {
+        energy_unit: "EV".to_string(),
+        mass_unit: "AMU".to_string(),
+        Eb: Eb2,
+        Es: Es2,
+        Ec: Ec2,
+        Ed: vec![0.0; num_species_target],
+        Z: Z2,
+        m: m2,
+        interaction_index: vec![0; num_species_target],
+        surface_binding_model: SurfaceBindingModel::INDIVIDUAL,
+        bulk_binding_model: BulkBindingModel::INDIVIDUAL,
+    };
+
+    let geometry_input = geometry::Mesh0DInput {
+        length_unit: "ANGSTROM".to_string(),
+        densities: n2,
+        electronic_stopping_correction_factor: 1.0
+    };
+
+    let m = material::Material::<Mesh0D>::new(&material_parameters, &geometry_input);
+
+    let mut finished_particles: Vec<particle::Particle> = Vec::new();
+
+    let incident_particles: Vec<particle::Particle> = izip!(energies, ux, uy, uz, Z1, Ec1, Es1, m1)
+        .enumerate()
+        .map(|(index, (energy, ux_, uy_, uz_, Z1_, Ec1_, Es1_, m1_))| {
+            let mut p = particle::Particle::default_incident(
+                m1_,
+                Z1_,
+                energy,
+                Ec1_,
+                Es1_,
+                x,
+                ux_,
+                uy_,
+                uz_
+            );
+            p.tag = index as i32;
+            p
+        }).collect();
+
+        finished_particles.par_extend(
+            incident_particles.into_par_iter()
+            .map(|particle| bca::single_ion_bca(particle, &m, &options))
+            .flatten()
+        );
+
+        for particle in finished_particles {
+            if (particle.left) | (particle.incident) {
+                incident.push(particle.incident);
+                incident_index.push(particle.tag as usize);
+                let mut energy_out;
+                if particle.stopped {
+                    energy_out = 0.;
+                } else {
+                    energy_out = particle.E/EV;
+                }
+                total_output.push(
+                    [
+                        particle.Z,
+                        particle.m/AMU,
+                        energy_out,
+                        particle.pos.x/ANGSTROM,
+                        particle.pos.y/ANGSTROM,
+                        particle.pos.z/ANGSTROM,
+                        particle.dir.x,
+                        particle.dir.y,
+                        particle.dir.z,
+                    ]
+                )
+            }
+        }
+
+    (total_output, incident, incident_index)
+}
+
 #[cfg(feature = "python")]
 ///reflect_single_ion_py(ion, target, vx, vy, vz)
 ///Performs a single BCA ion trajectory in target material with specified incident velocity.