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