generate_pairwise_interaction(pair_int_kernel, covalent_map, static_args)

This is a calculator generator for pairwise interaction

Input

pair_int_kernel: function type (dr, m, p1i, p1j, p2i, p2j) -> energy : the vectorized kernel function, dr is the distance, m is the topological scaling factor, p1i, p1j, p2i, p2j are pairwise parameters

covalent_map: Na * Na, int: the covalent_map matrix that marks the topological distances between atoms

static_args: dict: a dictionary that stores all static global parameters (such as lmax, kappa, etc)

Output

pair_int: function type (positions, box, pairs, mScales, p1, p2, ...) -> energy The pair interaction calculator. p1, p2 ... involved atomic parameters, the order should be consistent with the order in kernel

Source code in dmff/admp/pairwise.py 
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
def generate_pairwise_interaction(pair_int_kernel, covalent_map, static_args):
    '''
    This is a calculator generator for pairwise interaction 

    Input:
        pair_int_kernel:
            function type (dr, m, p1i, p1j, p2i, p2j) -> energy : the vectorized kernel function, 
            dr is the distance, m is the topological scaling factor, p1i, p1j, p2i, p2j are pairwise parameters

        covalent_map:
            Na * Na, int: the covalent_map matrix that marks the topological distances between atoms

        static_args:
            dict: a dictionary that stores all static global parameters (such as lmax, kappa, etc)

    Output:
        pair_int:
            function type (positions, box, pairs, mScales, p1, p2, ...) -> energy
            The pair interaction calculator. p1, p2 ... involved atomic parameters, the order should be consistent
            with the order in kernel
    '''

    def pair_int(positions, box, pairs, mScales, *atomic_params):
        pairs = regularize_pairs(pairs)

        ri = distribute_v3(positions, pairs[:, 0])
        rj = distribute_v3(positions, pairs[:, 1])
        # ri = positions[pairs[:, 0]]
        # rj = positions[pairs[:, 1]]
        nbonds = covalent_map[pairs[:, 0], pairs[:, 1]]
        mscales = distribute_scalar(mScales, nbonds-1)

        buffer_scales = pair_buffer_scales(pairs)
        mscales = mscales * buffer_scales
        # mscales = mScales[nbonds-1]
        box_inv = jnp.linalg.inv(box)
        dr = ri - rj
        dr = v_pbc_shift(dr, box, box_inv)
        dr = jnp.linalg.norm(dr, axis=1)

        pair_params = []
        for i, param in enumerate(atomic_params):
            pair_params.append(distribute_scalar(param, pairs[:, 0]))
            pair_params.append(distribute_scalar(param, pairs[:, 1]))
            # pair_params.append(param[pairs[:, 0]])
            # pair_params.append(param[pairs[:, 1]])

        energy = jnp.sum(pair_int_kernel(dr, mscales, *pair_params) * buffer_scales)
        return energy

    return pair_int

slater_disp_damping_kernel(dr, m, bi, bj, c6i, c6j, c8i, c8j, c10i, c10j)

Slater-ISA type damping for dispersion: f(x) = -e^{-x} * \sum_{k} x^k/k! x = Br - \frac{2(Br)^2 + 3Br}{(Br)^2 + 3Br + 3} see jctc 12 3851

Source code in dmff/admp/pairwise.py 
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
@vmap
@jit_condition(static_argnums=())
def slater_disp_damping_kernel(dr, m, bi, bj, c6i, c6j, c8i, c8j, c10i, c10j):
    r'''
    Slater-ISA type damping for dispersion:
    f(x) = -e^{-x} * \sum_{k} x^k/k!
    x = Br - \frac{2*(Br)^2 + 3Br}{(Br)^2 + 3*Br + 3}
    see jctc 12 3851
    '''
    b = jnp.sqrt(bi * bj)
    c6 = c6i * c6j
    c8 = c8i * c8j
    c10 = c10i * c10j
    br = b * dr
    br2 = br * br
    x = br - (2*br2 + 3*br) / (br2 + 3*br + 3)
    s6 = 1 + x + x**2/2 + x**3/6 + x**4/24 + x**5/120 + x**6/720
    s8 = s6 + x**7/5040 + x**8/40320
    s10 = s8 + x**9/362880 + x**10/3628800
    exp_x = jnp.exp(-x)
    f6 = exp_x * s6
    f8 = exp_x * s8
    f10 = exp_x * s10
    return (f6*c6/dr**6 + f8*c8/dr**8 + f10*c10/dr**10) * m

slater_sr_kernel(dr, m, ai, aj, bi, bj)

Slater-ISA type short range terms see jctc 12 3851

Source code in dmff/admp/pairwise.py 
159
160
161
162
163
164
165
166
167
168
169
170
171
@vmap
@jit_condition(static_argnums=())
def slater_sr_kernel(dr, m, ai, aj, bi, bj):
    '''
    Slater-ISA type short range terms
    see jctc 12 3851
    '''
    b = jnp.sqrt(bi * bj)
    a = ai * aj
    br = b * dr
    br2 = br * br
    P = 1/3 * br2 + br + 1 
    return a * P * jnp.exp(-br) * m