ewald 0.1.14

Ewald SPME force calculations
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142

//! This module contains code for the short-range component of the Coulomb force.

use std::arch::x86_64::{_CMP_LT_OQ, _mm256_blendv_ps, _mm256_cmp_ps, _mm256_set1_ps};

use lin_alg::f32::{Vec3, Vec3x8, Vec3x16, f32x8, f32x16};
use statrs::function::erf::erfc;

use crate::INV_SQRT_PI;

///  Computes the direct, short-range component. Ideally, use a combined GPU kernel with Lennard Jones,
/// or a SIMD variant, instead of this.  We use this for short-range Coulomb forces on the CPU, as part of SPME.
/// `cutoff_dist` is the distance, in Å, at which we no longer apply any force from this component.
/// α controls the blending of short and long-range forces; 0.35Å for α is a good default for a cutoff of 10Å.
///
/// This assumes diff (and dir) is in order tgt - src.
/// Also returns potential energy. `dir` must be a unit vector.
pub fn force_coulomb_short_range(
    dir: Vec3,
    dist: f32,
    // Included in this form to share between this and Lennard Jones.
    inv_dist: f32,
    q_0: f32,
    q_1: f32,
    cutoff_dist: f32,
    α: f32,
) -> (Vec3, f32) {
    if dist >= cutoff_dist {
        return (Vec3::new_zero(), 0.);
    }

    let α_r = α * dist;
    let erfc_term = erfc(α_r as f64) as f32;
    let charge_term = q_0 * q_1;

    let energy = charge_term * inv_dist * erfc_term;

    let exp_term = (-α_r * α_r).exp();

    let force_mag = charge_term
        * (erfc_term * inv_dist * inv_dist + 2.0 * α * exp_term * INV_SQRT_PI * inv_dist);

    (dir * force_mag, energy)
}

#[cfg(target_arch = "x86_64")]
pub fn force_coulomb_short_range_x8(
    dir: Vec3x8,
    dist: f32x8,
    inv_dist: f32x8,
    q_0: f32x8,
    q_1: f32x8,
    cutoff_dist: f32x8,
    // Alternatively, we could use a normal f32 for this, and splat it in-fn.
    α: f32x8,
) -> (Vec3x8, f32x8) {
    let α_r = α * dist;
    let erfc_term = α_r.erfc();

    let charge_term = q_0 * q_1;

    let energy = charge_term * inv_dist * erfc_term;

    let exp_term = (-α_r * α_r).exp();

    let force_mag = charge_term
        * (erfc_term * inv_dist * inv_dist
            + f32x8::splat(2.) * α * exp_term * f32x8::splat(INV_SQRT_PI) * inv_dist);

    let force = dir * force_mag;

    // This is where we diverge from the syntax of the non-SIMD variant;
    // the outside/inside cutoff.
    // per-lane mask: keep where dist < cutoff_dist, else zero
    unsafe {
        let keep = _mm256_cmp_ps::<{ _CMP_LT_OQ }>(dist.0, cutoff_dist.0);
        let zero = _mm256_set1_ps(0.0);

        let fx = _mm256_blendv_ps(zero, (force.x).0, keep);
        let fy = _mm256_blendv_ps(zero, (force.y).0, keep);
        let fz = _mm256_blendv_ps(zero, (force.z).0, keep);
        let en = _mm256_blendv_ps(zero, energy.0, keep);

        (
            Vec3x8 {
                x: f32x8(fx),
                y: f32x8(fy),
                z: f32x8(fz),
            },
            f32x8(en),
        )
    }
}

#[cfg(target_arch = "x86_64")]
pub fn force_coulomb_short_range_x16(
    dir: Vec3x16,
    dist: f32x16,
    // Included to share between this and Lennard Jones.
    inv_dist: f32x16,
    q_0: f32x16,
    q_1: f32x16,
    cutoff_dist: f32x16,
    // Alternatively, we could use a normal f32 for this, and splat it in-fn.
    α: f32x16,
) -> (Vec3x16, f32x16) {
    let α_r = α * dist;
    let erfc_term = α_r.erfc();

    let charge_term = q_0 * q_1;

    let energy = charge_term * inv_dist * erfc_term;

    let exp_term = (-α_r * α_r).exp();

    let force_mag = charge_term
        * (erfc_term * inv_dist * inv_dist
            + f32x16::splat(2.) * α * exp_term * f32x16::splat(INV_SQRT_PI) * inv_dist);

    let force = dir * force_mag;

    // This is where we diverge from the syntax of the non-SIMD variant;
    // the outside/inside cutoff.
    // per-lane mask: keep where dist < cutoff_dist, else zero
    unsafe {
        use core::arch::x86_64::*;
        let keep: __mmask16 = _mm512_cmp_ps_mask::<{ _CMP_LT_OQ }>(dist.0, cutoff_dist.0);

        let fx = _mm512_maskz_mov_ps(keep, (force.x).0);
        let fy = _mm512_maskz_mov_ps(keep, (force.y).0);
        let fz = _mm512_maskz_mov_ps(keep, (force.z).0);
        let en = _mm512_maskz_mov_ps(keep, energy.0);

        (
            Vec3x16 {
                x: f32x16(fx),
                y: f32x16(fy),
                z: f32x16(fz),
            },
            f32x16(en),
        )
    }
}