Skip to main content

oxiphysics_gpu/kernels/md_force/
virialstresstensorkernel_traits.rs

1//! # VirialStressTensorKernel - Trait Implementations
2//!
3//! This module contains trait implementations for `VirialStressTensorKernel`.
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::VirialStressTensorKernel;
13
14impl ComputeKernel for VirialStressTensorKernel {
15    fn name(&self) -> &str {
16        "VirialStressTensorKernel"
17    }
18    fn execute(&self, inputs: &[&[f64]], outputs: &mut [Vec<f64>], work_size: usize) {
19        if inputs.len() < 2 || outputs.is_empty() {
20            return;
21        }
22        let pos = inputs[0];
23        let epsilon = inputs[1][0];
24        let sigma = inputs[1][1];
25        let cutoff = inputs[1][2];
26        let n = work_size;
27        let cutoff2 = cutoff * cutoff;
28        let mut w = [0.0f64; 6];
29        for i in 0..n {
30            for j in (i + 1)..n {
31                let dx = pos[i * 3] - pos[j * 3];
32                let dy = pos[i * 3 + 1] - pos[j * 3 + 1];
33                let dz = pos[i * 3 + 2] - pos[j * 3 + 2];
34                let r2 = dx * dx + dy * dy + dz * dz;
35                if r2 >= cutoff2 || r2 < 1e-30 {
36                    continue;
37                }
38                let r2_inv = 1.0 / r2;
39                let sr2 = sigma * sigma * r2_inv;
40                let sr6 = sr2 * sr2 * sr2;
41                let sr12 = sr6 * sr6;
42                let f_over_r2 = 24.0 * epsilon * (2.0 * sr12 - sr6) * r2_inv;
43                let r_vec = [dx, dy, dz];
44                let c = [(0usize, 0usize), (1, 1), (2, 2), (0, 1), (0, 2), (1, 2)];
45                for (ci, &(a, b)) in c.iter().enumerate() {
46                    w[ci] -= r_vec[a] * f_over_r2 * r_vec[b];
47                }
48            }
49        }
50        outputs[0] = w.to_vec();
51    }
52}