Skip to content

GitLab

Switch branch/tag
Find file Normal viewHistoryPermalink
ewald_recp.py 3.14 KB
Newer Older
zhangqha's avatar
BladeDISC DeePMD code  
zhangqha committed
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 
import numpy as np
from typing import Tuple, List

from deepmd.env import tf
from deepmd.common import ClassArg
from deepmd.env import GLOBAL_TF_FLOAT_PRECISION
from deepmd.env import GLOBAL_NP_FLOAT_PRECISION
from deepmd.env import GLOBAL_ENER_FLOAT_PRECISION
from deepmd.env import global_cvt_2_tf_float
from deepmd.env import global_cvt_2_ener_float
from deepmd.env import op_module
from deepmd.env import default_tf_session_config
from deepmd.utils.sess import run_sess

class EwaldRecp () :
    """
    Evaluate the reciprocal part of the Ewald sum
    """
    def __init__(self, 
                 hh,
                 beta):
        """
        Constructor 

        Parameters
        ----------
        hh
                Grid spacing of the reciprocal part of Ewald sum. Unit: A
        beta
                Splitting parameter of the Ewald sum. Unit: A^{-1}
        """
        self.hh = hh
        self.beta = beta
        with tf.Graph().as_default() as graph:
            # place holders
            self.t_nloc       = tf.placeholder(tf.int32, [1], name = "t_nloc")
            self.t_coord      = tf.placeholder(GLOBAL_TF_FLOAT_PRECISION, [None], name='t_coord')
            self.t_charge     = tf.placeholder(GLOBAL_TF_FLOAT_PRECISION, [None], name='t_charge')
            self.t_box        = tf.placeholder(GLOBAL_TF_FLOAT_PRECISION, [None], name='t_box')
            # output            
            self.t_energy, self.t_force, self.t_virial \
                = op_module.ewald_recp(self.t_coord, self.t_charge, self.t_nloc, self.t_box, 
                                       ewald_h = self.hh,
                                       ewald_beta = self.beta)
        self.sess = tf.Session(graph=graph, config=default_tf_session_config)

    def eval(self, 
             coord : np.ndarray, 
             charge : np.ndarray, 
             box : np.ndarray
    ) -> Tuple[np.ndarray, np.ndarray, np.ndarray] :
        """
        Evaluate 
        
        Parameters
        ----------
        coord
                The coordinates of atoms
        charge
                The atomic charge
        box
                The simulation region. PBC is assumed

        Returns
        -------
        e
                The energy 
        f
                The force 
        v
                The virial 
        """
        coord = np.array(coord)
        charge = np.array(charge)
        box = np.array(box)
        nframes = charge.shape[0]
        natoms = charge.shape[1]
        coord = np.reshape(coord, [nframes * 3 * natoms])
        charge = np.reshape(charge, [nframes * natoms])
        box = np.reshape(box, [nframes * 9])

        [energy, force, virial] \
            = run_sess(self.sess, [self.t_energy, self.t_force, self.t_virial], 
                            feed_dict = {
                                self.t_coord:  coord,
                                self.t_charge: charge,
                                self.t_box:    box,
                                self.t_nloc:   [natoms],
                            })

        return energy, force, virial