oxiphysics_gpu/kernels/md_force/
temperaturescalekernel_traits.rs1#[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}