Skip to main content

oxiphysics_gpu/kernels/md_force/
temperaturescalekernel_traits.rs

1//! # TemperatureScaleKernel - Trait Implementations
2//!
3//! This module contains trait implementations for `TemperatureScaleKernel`.
4//!
5//! ## Implemented Traits
6//!
7//! - `ComputeKernel`
8//!
9//! 🤖 Generated with [SplitRS](https://github.com/cool-japan/splitrs)
10
11#[allow(unused_imports)]
12use super::functions::*;
13use crate::compute::ComputeKernel;
14
15use super::types::TemperatureScaleKernel;
16
17impl ComputeKernel for TemperatureScaleKernel {
18    fn name(&self) -> &str {
19        "TemperatureScaleKernel"
20    }
21    fn execute(&self, inputs: &[&[f64]], outputs: &mut [Vec<f64>], work_size: usize) {
22        if inputs.len() < 3 || outputs.len() < 2 {
23            return;
24        }
25        let vel_flat = inputs[0];
26        let masses = inputs[1];
27        let params = inputs[2];
28        if params.len() < 2 {
29            return;
30        }
31        let t_target = params[0];
32        let kb = params[1];
33        let n = work_size;
34        let ke2: f64 = (0..n)
35            .map(|i| {
36                let m = if i < masses.len() { masses[i] } else { 1.0 };
37                let vx = vel_flat[i * 3];
38                let vy = vel_flat[i * 3 + 1];
39                let vz = vel_flat[i * 3 + 2];
40                m * (vx * vx + vy * vy + vz * vz)
41            })
42            .sum();
43        let n_dof = (3 * n) as f64;
44        let t_before = if kb > 1e-30 { ke2 / (n_dof * kb) } else { 0.0 };
45        let scale = if t_before > 1e-30 && t_target >= 0.0 {
46            (t_target / t_before).sqrt()
47        } else {
48            1.0
49        };
50        let mut new_vel = vel_flat.to_vec();
51        for i in 0..n {
52            new_vel[i * 3] *= scale;
53            new_vel[i * 3 + 1] *= scale;
54            new_vel[i * 3 + 2] *= scale;
55        }
56        let ke2_after: f64 = (0..n)
57            .map(|i| {
58                let m = if i < masses.len() { masses[i] } else { 1.0 };
59                let vx = new_vel[i * 3];
60                let vy = new_vel[i * 3 + 1];
61                let vz = new_vel[i * 3 + 2];
62                m * (vx * vx + vy * vy + vz * vz)
63            })
64            .sum();
65        let t_after = if kb > 1e-30 {
66            ke2_after / (n_dof * kb)
67        } else {
68            0.0
69        };
70        outputs[0] = new_vel;
71        outputs[1] = vec![t_before, t_after];
72    }
73}