oxiphysics_gpu/kernels/md_force/
angleforcekernel_traits.rs1#[allow(unused_imports)]
12use super::functions::*;
13use crate::compute::ComputeKernel;
14
15use super::types::AngleForceKernel;
16
17#[allow(clippy::needless_range_loop)]
18impl ComputeKernel for AngleForceKernel {
19 fn name(&self) -> &str {
20 "AngleForceKernel"
21 }
22 fn execute(&self, inputs: &[&[f64]], outputs: &mut [Vec<f64>], work_size: usize) {
23 if inputs.len() < 2 || outputs.len() < 2 {
24 return;
25 }
26 let pos = inputs[0];
27 let angle_data = inputs[1];
28 let n = work_size;
29 let num_angles = angle_data.len() / 5;
30 let mut forces = vec![0.0f64; n * 3];
31 let mut energies = vec![0.0f64; num_angles];
32 for b in 0..num_angles {
33 let i = angle_data[b * 5] as usize;
34 let j = angle_data[b * 5 + 1] as usize;
35 let k = angle_data[b * 5 + 2] as usize;
36 let k_theta = angle_data[b * 5 + 3];
37 let theta0 = angle_data[b * 5 + 4];
38 if i >= n || j >= n || k >= n {
39 continue;
40 }
41 let rji = [
42 pos[i * 3] - pos[j * 3],
43 pos[i * 3 + 1] - pos[j * 3 + 1],
44 pos[i * 3 + 2] - pos[j * 3 + 2],
45 ];
46 let rjk = [
47 pos[k * 3] - pos[j * 3],
48 pos[k * 3 + 1] - pos[j * 3 + 1],
49 pos[k * 3 + 2] - pos[j * 3 + 2],
50 ];
51 let len_ji = (rji[0] * rji[0] + rji[1] * rji[1] + rji[2] * rji[2]).sqrt();
52 let len_jk = (rjk[0] * rjk[0] + rjk[1] * rjk[1] + rjk[2] * rjk[2]).sqrt();
53 if len_ji < 1e-30 || len_jk < 1e-30 {
54 continue;
55 }
56 let cos_theta = ((rji[0] * rjk[0] + rji[1] * rjk[1] + rji[2] * rjk[2])
57 / (len_ji * len_jk))
58 .clamp(-1.0, 1.0);
59 let theta = cos_theta.acos();
60 let delta = theta - theta0;
61 energies[b] = 0.5 * k_theta * delta * delta;
62 let sin_theta = (1.0 - cos_theta * cos_theta).sqrt().max(1e-12);
63 let d_prefactor = -k_theta * delta / sin_theta;
64 for dim in 0..3 {
65 let d_cos_d_ri =
66 rjk[dim] / (len_ji * len_jk) - cos_theta * rji[dim] / (len_ji * len_ji);
67 let d_cos_d_rk =
68 rji[dim] / (len_ji * len_jk) - cos_theta * rjk[dim] / (len_jk * len_jk);
69 let fi = d_prefactor * d_cos_d_ri;
70 let fk = d_prefactor * d_cos_d_rk;
71 forces[i * 3 + dim] += fi;
72 forces[k * 3 + dim] += fk;
73 forces[j * 3 + dim] -= fi + fk;
74 }
75 }
76 outputs[0] = forces;
77 outputs[1] = energies;
78 }
79}