oxiphysics_gpu/kernels/md_force/
temperaturescalekernel_traits.rs1use 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}