Skip to main content

oxiphysics_gpu/kernels/md_force/
angleforcekernel_traits.rs

1//! # AngleForceKernel - Trait Implementations
2//!
3//! This module contains trait implementations for `AngleForceKernel`.
4//!
5//! ## Implemented Traits
6//!
7//! - `ComputeKernel`
8//!
9//! 🤖 Generated with [SplitRS](https://github.com/cool-japan/splitrs)
10use crate::compute::ComputeKernel;
11
12use super::types::AngleForceKernel;
13
14impl ComputeKernel for AngleForceKernel {
15    fn name(&self) -> &str {
16        "AngleForceKernel"
17    }
18    fn execute(&self, inputs: &[&[f64]], outputs: &mut [Vec<f64>], work_size: usize) {
19        if inputs.len() < 2 || outputs.len() < 2 {
20            return;
21        }
22        let pos = inputs[0];
23        let angle_data = inputs[1];
24        let n = work_size;
25        let num_angles = angle_data.len() / 5;
26        let mut forces = vec![0.0f64; n * 3];
27        let mut energies = vec![0.0f64; num_angles];
28        for b in 0..num_angles {
29            let i = angle_data[b * 5] as usize;
30            let j = angle_data[b * 5 + 1] as usize;
31            let k = angle_data[b * 5 + 2] as usize;
32            let k_theta = angle_data[b * 5 + 3];
33            let theta0 = angle_data[b * 5 + 4];
34            if i >= n || j >= n || k >= n {
35                continue;
36            }
37            let rji = [
38                pos[i * 3] - pos[j * 3],
39                pos[i * 3 + 1] - pos[j * 3 + 1],
40                pos[i * 3 + 2] - pos[j * 3 + 2],
41            ];
42            let rjk = [
43                pos[k * 3] - pos[j * 3],
44                pos[k * 3 + 1] - pos[j * 3 + 1],
45                pos[k * 3 + 2] - pos[j * 3 + 2],
46            ];
47            let len_ji = (rji[0] * rji[0] + rji[1] * rji[1] + rji[2] * rji[2]).sqrt();
48            let len_jk = (rjk[0] * rjk[0] + rjk[1] * rjk[1] + rjk[2] * rjk[2]).sqrt();
49            if len_ji < 1e-30 || len_jk < 1e-30 {
50                continue;
51            }
52            let cos_theta = ((rji[0] * rjk[0] + rji[1] * rjk[1] + rji[2] * rjk[2])
53                / (len_ji * len_jk))
54                .clamp(-1.0, 1.0);
55            let theta = cos_theta.acos();
56            let delta = theta - theta0;
57            energies[b] = 0.5 * k_theta * delta * delta;
58            let sin_theta = (1.0 - cos_theta * cos_theta).sqrt().max(1e-12);
59            let d_prefactor = -k_theta * delta / sin_theta;
60            for dim in 0..3 {
61                let d_cos_d_ri =
62                    rjk[dim] / (len_ji * len_jk) - cos_theta * rji[dim] / (len_ji * len_ji);
63                let d_cos_d_rk =
64                    rji[dim] / (len_ji * len_jk) - cos_theta * rjk[dim] / (len_jk * len_jk);
65                let fi = d_prefactor * d_cos_d_ri;
66                let fk = d_prefactor * d_cos_d_rk;
67                forces[i * 3 + dim] += fi;
68                forces[k * 3 + dim] += fk;
69                forces[j * 3 + dim] -= fi + fk;
70            }
71        }
72        outputs[0] = forces;
73        outputs[1] = energies;
74    }
75}