def __init__(self):
super(Apply, self).__init__()
self.energy_fn = Energy()
self.grad = C.GradOperation('grad')
self.dt = 0.1
self.temperature = 0.1
self.gamma = 0.1
self.mass = 1.0
self.nu = 1.0 / (self.mass * self.gamma)
self.xi_c = math.sqrt(2.0 * self.temperature * self.dt * self.nu)
self.random = P.StandardNormal(seed=1, seed2=2)
def __init__(self, shape, seed=0, seed2=0):
super(StandardNormalNet, self).__init__()
self.shape = shape
self.seed = seed
self.seed2 = seed2
self.standard_normal = P.StandardNormal(seed, seed2)
os.environ['GLOG_v'] = '3'
# MD configs
dims = 2
dt = 1e-1
temperature = 0.1
gamma = 0.1
mass = 1.0
nu = 1.0 / (mass * gamma)
xi_c = math.sqrt(2.0 * temperature * dt * nu)
dtype = mindspore.float32
# Functions
reduce_sum = P.ReduceSum()
max_op = P.Maximum()
normal = P.StandardNormal(seed=1, seed2=2)
relu = P.ReLU()
@ms_function
def displacement(Ra: Tensor, Rb: Tensor):
dR = Ra - Rb
# periodic
# np.mod(dR + box_size * 0.5, box_size) - box_size * 0.5
return dR
@ms_function
def pairwise_displacement(R: Tensor):
dR = F.expand_dims(R, 1) - F.expand_dims(R, 0)
# periodic
