Skip to content

recip

torch_admp.recip

Reciprocal space operations for torch-admp.

This module provides functions for reciprocal space calculations used in Particle Mesh Ewald (PME) and other reciprocal space methods, including B-spline interpolation, charge spreading, and k-point setup.

Q_m_peratom(charges: torch.Tensor, sph_harms: torch.Tensor, pme_order: int)

Computes . See eq. (49) of https://doi.org/10.1021/ct5007983

Inputs: Q: N_a * (l+1)2 matrix containing global frame multipole moments up to lmax, sph_harms: N_a, 216, (l+1)2 lmax: int: maximal L

Output: Q_m_pera: N_a * 216 matrix, values of theta evaluated on a 6 * 6 block about the atoms

Source code in torch_admp/recip.py 
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
def Q_m_peratom(
    charges: torch.Tensor,
    sph_harms: torch.Tensor,
    pme_order: int,
):
    """
    Computes <R_t|Q>. See eq. (49) of https://doi.org/10.1021/ct5007983

    Inputs:
        Q:
            N_a * (l+1)**2 matrix containing global frame multipole moments up to lmax,
        sph_harms:
            N_a, 216, (l+1)**2
        lmax:
            int: maximal L

    Output:
        Q_m_pera:
            N_a * 216 matrix, values of theta evaluated on a 6 * 6 block about the atoms
    """
    n_mesh = int(pme_order**3)
    N_a = sph_harms.shape[0]
    Q_dbf = torch.atleast_2d(charges)[:, 0:1]
    Q_m_pera = torch.sum(Q_dbf[:, None, :] * sph_harms, dim=2)
    assert Q_m_pera.shape == (N_a, n_mesh)
    return Q_m_pera

Q_mesh_on_m(Q_mesh_pera: torch.Tensor, m_u0: torch.Tensor, t_kmesh: torch.Tensor, shifts: torch.Tensor)

Reduce the local Q_m_peratom into the global mesh

Input: Q_mesh_pera, m_u0, N

Output: Q_mesh: Nx * Ny * Nz matrix

Source code in torch_admp/recip.py 
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
def Q_mesh_on_m(
    Q_mesh_pera: torch.Tensor,
    m_u0: torch.Tensor,
    t_kmesh: torch.Tensor,
    shifts: torch.Tensor,
):
    """
    Reduce the local Q_m_peratom into the global mesh

    Input:
        Q_mesh_pera, m_u0, N

    Output:
        Q_mesh:
            Nx * Ny * Nz matrix
    """
    indices_arr = torch.fmod(
        m_u0[:, None, :] + shifts + t_kmesh[None, None, :] * 10, t_kmesh[None, None, :]
    )
    Q_mesh = torch.zeros(
        int(t_kmesh[0].item()) * int(t_kmesh[1].item()) * int(t_kmesh[2].item()),
        device=t_kmesh.device,
        dtype=Q_mesh_pera.dtype,
    )
    indices_0 = indices_arr[:, :, 0].flatten()
    indices_1 = indices_arr[:, :, 1].flatten()
    indices_2 = indices_arr[:, :, 2].flatten()
    flat_indices = (
        indices_0 * int(t_kmesh[1].item()) * int(t_kmesh[2].item())
        + indices_1 * int(t_kmesh[2].item())
        + indices_2
    )
    Q_mesh.index_add_(0, flat_indices, Q_mesh_pera.view(-1))
    Q_mesh = Q_mesh.reshape(
        int(t_kmesh[0].item()),
        int(t_kmesh[1].item()),
        int(t_kmesh[2].item()),
    )

    return Q_mesh

bspline(u: torch.Tensor)

Computes the cardinal B-spline function

Source code in torch_admp/recip.py 
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
143
def bspline(u: torch.Tensor):
    """
    Computes the cardinal B-spline function
    """
    return torch.where(
        (u >= 0.0) & (u < 1.0),
        u**5 / 120,
        torch.where(
            u < 2.0,
            u**5 / 120 - (u - 1) ** 5 / 20,
            torch.where(
                u < 3.0,
                u**5 / 120 + (u - 2) ** 5 / 8 - (u - 1) ** 5 / 20,
                torch.where(
                    u < 4.0,
                    u**5 / 120
                    - (u - 3) ** 5 / 6
                    + (u - 2) ** 5 / 8
                    - (u - 1) ** 5 / 20,
                    torch.where(
                        u < 5.0,
                        u**5 / 24
                        - u**4
                        + 19 * u**3 / 2
                        - 89 * u**2 / 2
                        + 409 * u / 4
                        - 1829 / 20,
                        torch.where(
                            u < 6.0,
                            -(u**5) / 120
                            + u**4 / 4
                            - 3 * u**3
                            + 18 * u**2
                            - 54 * u
                            + 324 / 5,
                            torch.zeros_like(u),
                        ),
                    ),
                ),
            ),
        ),
    )

get_recip_grid_vectors(box_inv: torch.Tensor, t_kmesh: torch.Tensor)

Compute reciprocal lattice vectors of grids

PARAMETER DESCRIPTION
box_inv

(3 * 3)-matrix for inv cell vectors inv_box = torch.linalg.inv(box)

TYPE: Tensor

t_kmesh

(3,)-shaped tensor [kx, ky, kz]

TYPE: Tensor

RETURNS DESCRIPTION
recip_grid_vectors

(3 * 3)-matrix for reciprocal lattice vectors of grids

TYPE: Tensor

Source code in torch_admp/recip.py 
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
def get_recip_grid_vectors(
    box_inv: torch.Tensor,
    t_kmesh: torch.Tensor,
):
    """
    Compute reciprocal lattice vectors of grids

    Parameters
    ----------
    box_inv : torch.Tensor
        (3 * 3)-matrix for inv cell vectors
        inv_box = torch.linalg.inv(box)
    t_kmesh : torch.Tensor
        (3,)-shaped tensor [kx, ky, kz]

    Returns
    -------
    recip_grid_vectors: torch.Tensor
        (3 * 3)-matrix for reciprocal lattice vectors of grids
    """
    recip_grid_vectors = (t_kmesh.reshape(1, 3) * box_inv).transpose(0, 1)
    return recip_grid_vectors

setup_kpts(box_inv, kpts_int)

This function sets up the k-points used for reciprocal space calculations

Input: box_inv: 3 * 3, three axis arranged in rows kpts_int: n_k * 3 matrix

Output: kpts: 4 * K, K=K1K2K3, contains kx, ky, kz, k^2 for each kpoint

Source code in torch_admp/recip.py 
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
def setup_kpts(box_inv, kpts_int):
    """
    This function sets up the k-points used for reciprocal space calculations

    Input:
        box_inv:
            3 * 3, three axis arranged in rows
        kpts_int:
            n_k * 3 matrix

    Output:
        kpts:
            4 * K, K=K1*K2*K3, contains kx, ky, kz, k^2 for each kpoint
    """
    # in this array, a*, b*, c* (without 2*pi) are arranged in column
    # K * 3, coordinate in reciprocal space
    kpts = 2 * torch.pi * torch.matmul(kpts_int.double(), box_inv)
    ksq = torch.sum(kpts**2, dim=1)
    # 4 * K
    kpts = torch.hstack((kpts, ksq[:, None])).transpose(0, 1)
    return kpts

setup_kpts_integer(t_kmesh: torch.Tensor)

Set up integer k-points for reciprocal space calculations.

PARAMETER DESCRIPTION
t_kmesh

Mesh dimensions [Kx, Ky, Kz]

TYPE: Tensor

RETURNS DESCRIPTION
Tensor

n_k * 3 matrix of integer k-points, where n_k = Kx * Ky * Kz
Source code in torch_admp/recip.py 
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
def setup_kpts_integer(
    t_kmesh: torch.Tensor,
):
    """
    Set up integer k-points for reciprocal space calculations.

    Parameters
    ----------
    t_kmesh : torch.Tensor
        Mesh dimensions [Kx, Ky, Kz]

    Returns
    -------
    torch.Tensor
        n_k * 3 matrix of integer k-points, where n_k = Kx * Ky * Kz
    """
    kx, ky, kz = [
        torch.roll(
            torch.arange(
                -(int(t_kmesh[i].item()) - 1) // 2,
                (int(t_kmesh[i].item()) + 1) // 2,
                device=t_kmesh.device,
            ),
            shifts=[-(int(t_kmesh[i].item()) - 1) // 2],
        )
        for i in range(3)
    ]
    kpts_int = torch.hstack(
        [ki.flatten()[:, None] for ki in torch.meshgrid(kx, ky, kz, indexing="ij")]
    )
    return kpts_int

sph_harmonics_GO(u0: torch.Tensor, shifts: torch.Tensor, pme_order: int)

Find out the value of spherical harmonics GRADIENT OPERATORS, assume the order is: 00, 10, 11c, 11s, 20, 21c, 21s, 22c, 22s, ... Currently supports lmax <= 2

PARAMETER DESCRIPTION
u0

(N_a * 3)-matrix containing all positions

TYPE: Tensor

recip_grid_vectors

(3 * 3)-matrix for reciprocal lattice vectors of grids

TYPE: Tensor

RETURNS DESCRIPTION
harmonics

a Na * (63) * (l+1)^2 matrix, STGO operated on theta, evaluated at 666 integer points about reference points m_u0

TYPE: Tensor

Source code in torch_admp/recip.py 
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
def sph_harmonics_GO(
    u0: torch.Tensor,
    shifts: torch.Tensor,
    pme_order: int,
):
    """
    Find out the value of spherical harmonics GRADIENT OPERATORS, assume the order is:
    00, 10, 11c, 11s, 20, 21c, 21s, 22c, 22s, ...
    Currently supports lmax <= 2

    Parameters
    ----------
    u0 : torch.Tensor
        (N_a * 3)-matrix containing all positions
    recip_grid_vectors : torch.Tensor
        (3 * 3)-matrix for reciprocal lattice vectors of grids

    Returns
    -------
    harmonics: torch.Tensor
        a Na * (6**3) * (l+1)^2 matrix, STGO operated on theta,
        evaluated at 6*6*6 integer points about reference points m_u0
    """
    n_mesh = int(pme_order**3)
    N_a = u0.shape[0]

    # mesh points around each site
    u = u0[:, None, :] + shifts
    u_reshape = torch.reshape(u, (N_a * n_mesh, 3))
    # bspline may have little different value
    M_u = bspline(u_reshape)
    theta = torch.prod(M_u, dim=-1)
    return theta.reshape(N_a, n_mesh, 1)

spread_charges(positions: torch.Tensor, box_inv: torch.Tensor, charges: torch.Tensor, t_kmesh: torch.Tensor, shifts: torch.Tensor, pme_order: int)

This is the high level wrapper function, in charge of spreading the charges/multipoles on grid

PARAMETER DESCRIPTION
positions

Na * 3: positions of atoms

TYPE: Tensor

box

3 * 3: cell vectors

TYPE: Tensor

charges

Na * (lmax+1)**2: the multipole of each atomic site in global frame

TYPE: Tensor

RETURNS DESCRIPTION
Tensor

Nx * Ny * Nz: the meshed multipoles
Output

Q_mesh: K1 * K2 * K3: the meshed multipoles
Source code in torch_admp/recip.py 
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
def spread_charges(
    positions: torch.Tensor,
    box_inv: torch.Tensor,
    charges: torch.Tensor,
    t_kmesh: torch.Tensor,
    shifts: torch.Tensor,
    pme_order: int,
):
    """
    This is the high level wrapper function, in charge of spreading the charges/multipoles on grid

    Parameters
    ----------
    positions : torch.Tensor
        Na * 3: positions of atoms
    box : torch.Tensor
        3 * 3: cell vectors
    charges : torch.Tensor
        Na * (lmax+1)**2: the multipole of each atomic site in global frame

    Returns
    -------
    torch.Tensor
        Nx * Ny * Nz: the meshed multipoles
    Output:
        Q_mesh:
            K1 * K2 * K3: the meshed multipoles

    """
    recip_grid_vectors = get_recip_grid_vectors(box_inv, t_kmesh)
    # For each atom, find the reference mesh point, and u position of the site
    m_u0, u0 = u_reference(positions, recip_grid_vectors, pme_order)
    # find out the STGO values of each grid point
    sph_harms = sph_harmonics_GO(u0, shifts, pme_order)
    # find out the local meshed values for each site
    Q_mesh_pera = Q_m_peratom(charges, sph_harms, pme_order)
    return Q_mesh_on_m(Q_mesh_pera, m_u0, t_kmesh, shifts)

u_reference(positions: torch.Tensor, recip_grid_vectors: torch.Tensor, pme_order: int)

Each atom is meshed to dispersion_ORDER**3 points on the m-meshgrid. This function computes the xyz-index of the reference point, which is the point on the meshgrid just above atomic coordinates, and the corresponding values of xyz fractional displacements from real coordinate to the reference point.

PARAMETER DESCRIPTION
positions

Na * 3: positions of atoms

TYPE: Tensor

recip_grid_vectors

(3 * 3)-matrix for reciprocal lattice vectors of grids

TYPE: Tensor

RETURNS DESCRIPTION
m_u0

N_a * 3 matrix, positions of the reference points of R_a on the m-meshgrid

TYPE: Tensor

u0

N_a * 3 matrix, (R_a - R_m)*a_star values

TYPE: Tensor

Source code in torch_admp/recip.py 
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
def u_reference(
    positions: torch.Tensor,
    recip_grid_vectors: torch.Tensor,
    pme_order: int,
):
    """
    Each atom is meshed to dispersion_ORDER**3 points on the m-meshgrid.
    This function computes the xyz-index of the reference point, which is the point on the meshgrid just above atomic coordinates,
    and the corresponding values of xyz fractional displacements from real coordinate to the reference point.

    Parameters
    ----------
    positions : torch.Tensor
        Na * 3: positions of atoms
    recip_grid_vectors : torch.Tensor
        (3 * 3)-matrix for reciprocal lattice vectors of grids

    Returns
    -------
    m_u0: torch.Tensor
        N_a * 3 matrix, positions of the reference points of R_a on the m-meshgrid
    u0: torch.Tensor
        N_a * 3 matrix, (R_a - R_m)*a_star values
    """
    R_in_m_basis = torch.einsum("ij,kj->ki", recip_grid_vectors, positions)
    m_u0 = torch.ceil(R_in_m_basis).to(torch.int)
    u0 = (m_u0 - R_in_m_basis) + pme_order / 2
    return m_u0, u0