def _eval_fv(self, coords, cells, atom_types, ext_f):
# reshape the inputs
cells = np.reshape(cells, [-1, 9])
nframes = cells.shape[0]
coords = np.reshape(coords, [nframes, -1])
natoms = coords.shape[1] // 3
# sort inputs
coords, atom_types, imap, sel_at, sel_imap = self.sort_input(
coords, atom_types, sel_atoms=self.get_sel_type())
# make natoms_vec and default_mesh
natoms_vec = self.make_natoms_vec(atom_types)
assert (natoms_vec[0] == natoms)
default_mesh = make_default_mesh(cells)
# evaluate
tensor = []
feed_dict_test = {}
feed_dict_test[self.t_natoms] = natoms_vec
feed_dict_test[self.t_type] = np.tile(atom_types,
[nframes, 1]).reshape([-1])
feed_dict_test[self.t_coord] = coords.reshape([-1])
feed_dict_test[self.t_box] = cells.reshape([-1])
feed_dict_test[self.t_mesh] = default_mesh.reshape([-1])
feed_dict_test[self.t_ef] = ext_f.reshape([-1])
# print(run_sess(self.sess, tf.shape(self.t_tensor), feed_dict = feed_dict_test))
fout, vout, avout \
= run_sess(self.sess, [self.force, self.virial, self.av],
feed_dict = feed_dict_test)
# print('fout: ', fout.shape, fout)
fout = self.reverse_map(np.reshape(fout, [nframes, -1, 3]), imap)
fout = np.reshape(fout, [nframes, -1])
return fout, vout, avout
def eval(self, coords, cells, atom_types, atomic=True):
# standarize the shape of inputs
coords = np.array(coords)
cells = np.array(cells)
atom_types = np.array(atom_types, dtype=int)
# reshape the inputs
cells = np.reshape(cells, [-1, 9])
nframes = cells.shape[0]
coords = np.reshape(coords, [nframes, -1])
natoms = coords.shape[1] // 3
# sort inputs
coords, atom_types, imap, sel_at, sel_imap = self.sort_input(
coords, atom_types, sel_atoms=self.get_sel_type())
# make natoms_vec and default_mesh
natoms_vec = self.make_natoms_vec(atom_types)
assert (natoms_vec[0] == natoms)
# evaluate
tensor = []
feed_dict_test = {}
feed_dict_test[self.t_natoms] = natoms_vec
feed_dict_test[self.t_type] = atom_types
t_out = [self.t_tensor]
for ii in range(nframes):
feed_dict_test[self.t_coord] = np.reshape(coords[ii:ii + 1, :],
[-1])
feed_dict_test[self.t_box] = np.reshape(cells[ii:ii + 1, :], [-1])
feed_dict_test[self.t_mesh] = make_default_mesh(cells[ii:ii +
1, :])
v_out = self.sess.run(t_out, feed_dict=feed_dict_test)
tensor.append(v_out[0])
# reverse map of the outputs
if atomic:
tensor = np.array(tensor)
tensor = self.reverse_map(
np.reshape(tensor, [nframes, -1, self.variable_dof]), sel_imap)
tensor = np.reshape(
tensor, [nframes, len(sel_at), self.variable_dof])
else:
tensor = np.reshape(tensor, [nframes, self.variable_dof])
return tensor
def _eval_inner(self,
coords,
cells,
atom_types,
fparam=None,
aparam=None,
atomic=False,
efield=None):
# standarize the shape of inputs
atom_types = np.array(atom_types, dtype=int).reshape([-1])
natoms = atom_types.size
coords = np.reshape(np.array(coords), [-1, natoms * 3])
nframes = coords.shape[0]
if cells is None:
pbc = False
# make cells to work around the requirement of pbc
cells = np.tile(np.eye(3), [nframes, 1]).reshape([nframes, 9])
else:
pbc = True
cells = np.array(cells).reshape([nframes, 9])
if self.has_fparam:
assert (fparam is not None)
fparam = np.array(fparam)
if self.has_aparam:
assert (aparam is not None)
aparam = np.array(aparam)
if self.has_efield:
assert (
efield is not None
), "you are using a model with external field, parameter efield should be provided"
efield = np.array(efield)
# reshape the inputs
if self.has_fparam:
fdim = self.get_dim_fparam()
if fparam.size == nframes * fdim:
fparam = np.reshape(fparam, [nframes, fdim])
elif fparam.size == fdim:
fparam = np.tile(fparam.reshape([-1]), [nframes, 1])
else:
raise RuntimeError(
'got wrong size of frame param, should be either %d x %d or %d'
% (nframes, fdim, fdim))
if self.has_aparam:
fdim = self.get_dim_aparam()
if aparam.size == nframes * natoms * fdim:
aparam = np.reshape(aparam, [nframes, natoms * fdim])
elif aparam.size == natoms * fdim:
aparam = np.tile(aparam.reshape([-1]), [nframes, 1])
elif aparam.size == fdim:
aparam = np.tile(aparam.reshape([-1]), [nframes, natoms])
else:
raise RuntimeError(
'got wrong size of frame param, should be either %d x %d x %d or %d x %d or %d'
% (nframes, natoms, fdim, natoms, fdim, fdim))
# sort inputs
coords, atom_types, imap = self.sort_input(coords, atom_types)
if self.has_efield:
efield = np.reshape(efield, [nframes, natoms, 3])
efield = efield[:, imap, :]
efield = np.reshape(efield, [nframes, natoms * 3])
# make natoms_vec and default_mesh
natoms_vec = self.make_natoms_vec(atom_types)
assert (natoms_vec[0] == natoms)
# evaluate
feed_dict_test = {}
feed_dict_test[self.t_natoms] = natoms_vec
feed_dict_test[self.t_type] = np.tile(atom_types,
[nframes, 1]).reshape([-1])
t_out = [self.t_energy, self.t_force, self.t_virial]
if atomic:
t_out += [self.t_ae, self.t_av]
feed_dict_test[self.t_coord] = np.reshape(coords, [-1])
feed_dict_test[self.t_box] = np.reshape(cells, [-1])
if self.has_efield:
feed_dict_test[self.t_efield] = np.reshape(efield, [-1])
if pbc:
feed_dict_test[self.t_mesh] = make_default_mesh(cells)
else:
feed_dict_test[self.t_mesh] = np.array([], dtype=np.int32)
if self.has_fparam:
feed_dict_test[self.t_fparam] = np.reshape(fparam, [-1])
if self.has_aparam:
feed_dict_test[self.t_aparam] = np.reshape(aparam, [-1])
v_out = run_sess(self.sess, t_out, feed_dict=feed_dict_test)
energy = v_out[0]
force = v_out[1]
virial = v_out[2]
if atomic:
ae = v_out[3]
av = v_out[4]
# reverse map of the outputs
force = self.reverse_map(np.reshape(force, [nframes, -1, 3]), imap)
if atomic:
ae = self.reverse_map(np.reshape(ae, [nframes, -1, 1]), imap)
av = self.reverse_map(np.reshape(av, [nframes, -1, 9]), imap)
energy = np.reshape(energy, [nframes, 1])
force = np.reshape(force, [nframes, natoms, 3])
virial = np.reshape(virial, [nframes, 9])
if atomic:
ae = np.reshape(ae, [nframes, natoms, 1])
av = np.reshape(av, [nframes, natoms, 9])
return energy, force, virial, ae, av
else:
return energy, force, virial
def eval_inner(self,
coords,
cells,
atom_types,
fparam = None,
aparam = None,
atomic = False) :
# standarize the shape of inputs
atom_types = np.array(atom_types, dtype = int).reshape([-1])
natoms = atom_types.size
coords = np.reshape(np.array(coords), [-1, natoms * 3])
nframes = coords.shape[0]
if cells is None:
pbc = False
# make cells to work around the requirement of pbc
cells = np.tile(np.eye(3), [nframes, 1]).reshape([nframes, 9])
else:
pbc = True
cells = np.array(cells).reshape([nframes, 9])
if self.has_fparam :
assert(fparam is not None)
fparam = np.array(fparam)
if self.has_aparam :
assert(aparam is not None)
aparam = np.array(aparam)
# reshape the inputs
if self.has_fparam :
fdim = self.get_dim_fparam()
if fparam.size == nframes * fdim :
fparam = np.reshape(fparam, [nframes, fdim])
elif fparam.size == fdim :
fparam = np.tile(fparam.reshape([-1]), [nframes, 1])
else :
raise RuntimeError('got wrong size of frame param, should be either %d x %d or %d' % (nframes, fdim, fdim))
if self.has_aparam :
fdim = self.get_dim_aparam()
if aparam.size == nframes * natoms * fdim:
aparam = np.reshape(aparam, [nframes, natoms * fdim])
elif aparam.size == natoms * fdim :
aparam = np.tile(aparam.reshape([-1]), [nframes, 1])
elif aparam.size == fdim :
aparam = np.tile(aparam.reshape([-1]), [nframes, natoms])
else :
raise RuntimeError('got wrong size of frame param, should be either %d x %d x %d or %d x %d or %d' % (nframes, natoms, fdim, natoms, fdim, fdim))
# sort inputs
coords, atom_types, imap = self.sort_input(coords, atom_types)
# make natoms_vec and default_mesh
natoms_vec = self.make_natoms_vec(atom_types)
assert(natoms_vec[0] == natoms)
# evaluate
energy = []
force = []
virial = []
ae = []
av = []
feed_dict_test = {}
feed_dict_test[self.t_natoms] = natoms_vec
feed_dict_test[self.t_type ] = atom_types
t_out = [self.t_energy,
self.t_force,
self.t_virial]
if atomic :
t_out += [self.t_ae,
self.t_av]
for ii in range(nframes) :
feed_dict_test[self.t_coord] = np.reshape(coords[ii:ii+1, :], [-1])
feed_dict_test[self.t_box ] = np.reshape(cells [ii:ii+1, :], [-1])
if pbc:
feed_dict_test[self.t_mesh ] = make_default_mesh(cells[ii:ii+1, :])
else:
feed_dict_test[self.t_mesh ] = np.array([], dtype = np.int32)
if self.has_fparam:
feed_dict_test[self.t_fparam] = np.reshape(fparam[ii:ii+1, :], [-1])
if self.has_aparam:
feed_dict_test[self.t_aparam] = np.reshape(aparam[ii:ii+1, :], [-1])
v_out = self.sess.run (t_out, feed_dict = feed_dict_test)
energy.append(v_out[0])
force .append(v_out[1])
virial.append(v_out[2])
if atomic:
ae.append(v_out[3])
av.append(v_out[4])
# reverse map of the outputs
force = self.reverse_map(np.reshape(force, [nframes,-1,3]), imap)
if atomic :
ae = self.reverse_map(np.reshape(ae, [nframes,-1,1]), imap)
av = self.reverse_map(np.reshape(av, [nframes,-1,9]), imap)
energy = np.reshape(energy, [nframes, 1])
force = np.reshape(force, [nframes, natoms, 3])
virial = np.reshape(virial, [nframes, 9])
if atomic:
ae = np.reshape(ae, [nframes, natoms, 1])
av = np.reshape(av, [nframes, natoms, 9])
return energy, force, virial, ae, av
else :
return energy, force, virial
def eval_full(self,
coords: np.ndarray,
cells: np.ndarray,
atom_types: List[int],
atomic: bool = False,
fparam: Optional[np.array] = None,
aparam: Optional[np.array] = None,
efield: Optional[np.array] = None) -> Tuple[np.ndarray, ...]:
"""Evaluate the model with interface similar to the energy model.
Will return global tensor, component-wise force and virial
and optionally atomic tensor and atomic virial.
Parameters
----------
coords
The coordinates of atoms.
The array should be of size nframes x natoms x 3
cells
The cell of the region.
If None then non-PBC is assumed, otherwise using PBC.
The array should be of size nframes x 9
atom_types
The atom types
The list should contain natoms ints
atomic
Whether to calculate atomic tensor and virial
fparam
Not used in this model
aparam
Not used in this model
efield
Not used in this model
Returns
-------
tensor
The global tensor.
shape: [nframes x nout]
force
The component-wise force (negative derivative) on each atom.
shape: [nframes x nout x natoms x 3]
virial
The component-wise virial of the tensor.
shape: [nframes x nout x 9]
atom_tensor
The atomic tensor. Only returned when atomic == True
shape: [nframes x natoms x nout]
atom_virial
The atomic virial. Only returned when atomic == True
shape: [nframes x nout x natoms x 9]
"""
assert self._support_gfv, \
f"do not support eval_full with old tensor model"
# standarize the shape of inputs
atom_types = np.array(atom_types, dtype=int).reshape([-1])
natoms = atom_types.size
coords = np.reshape(np.array(coords), [-1, natoms * 3])
nframes = coords.shape[0]
if cells is None:
pbc = False
cells = np.tile(np.eye(3), [nframes, 1]).reshape([nframes, 9])
else:
pbc = True
cells = np.array(cells).reshape([nframes, 9])
nout = self.output_dim
# sort inputs
coords, atom_types, imap, sel_at, sel_imap = self.sort_input(
coords, atom_types, sel_atoms=self.get_sel_type())
# make natoms_vec and default_mesh
natoms_vec = self.make_natoms_vec(atom_types)
assert (natoms_vec[0] == natoms)
# evaluate
feed_dict_test = {}
feed_dict_test[self.t_natoms] = natoms_vec
feed_dict_test[self.t_type] = np.tile(atom_types,
[nframes, 1]).reshape([-1])
feed_dict_test[self.t_coord] = np.reshape(coords, [-1])
feed_dict_test[self.t_box] = np.reshape(cells, [-1])
if pbc:
feed_dict_test[self.t_mesh] = make_default_mesh(cells)
else:
feed_dict_test[self.t_mesh] = np.array([], dtype=np.int32)
t_out = [self.t_global_tensor, self.t_force, self.t_virial]
if atomic:
t_out += [self.t_tensor, self.t_atom_virial]
v_out = self.sess.run(t_out, feed_dict=feed_dict_test)
gt = v_out[0] # global tensor
force = v_out[1]
virial = v_out[2]
if atomic:
at = v_out[3] # atom tensor
av = v_out[4] # atom virial
# please note here the shape are wrong!
force = self.reverse_map(
np.reshape(force, [nframes * nout, natoms, 3]), imap)
if atomic:
at = self.reverse_map(np.reshape(
at, [nframes, len(sel_at), nout]), sel_imap)
av = self.reverse_map(np.reshape(av, [nframes * nout, natoms, 9]),
imap)
# make sure the shapes are correct here
gt = np.reshape(gt, [nframes, nout])
force = np.reshape(force, [nframes, nout, natoms, 3])
virial = np.reshape(virial, [nframes, nout, 9])
if atomic:
at = np.reshape(at, [nframes, len(sel_at), self.output_dim])
av = np.reshape(av, [nframes, nout, natoms, 9])
return gt, force, virial, at, av
else:
return gt, force, virial
def eval(self,
coords: np.ndarray,
cells: np.ndarray,
atom_types: List[int],
atomic: bool = True,
fparam: Optional[np.ndarray] = None,
aparam: Optional[np.ndarray] = None,
efield: Optional[np.ndarray] = None) -> np.ndarray:
"""Evaluate the model.
Parameters
----------
coords
The coordinates of atoms.
The array should be of size nframes x natoms x 3
cells
The cell of the region.
If None then non-PBC is assumed, otherwise using PBC.
The array should be of size nframes x 9
atom_types
The atom types
The list should contain natoms ints
atomic
If True (default), return the atomic tensor
Otherwise return the global tensor
fparam
Not used in this model
aparam
Not used in this model
efield
Not used in this model
Returns
-------
tensor
The returned tensor
If atomic == False then of size nframes x output_dim
else of size nframes x natoms x output_dim
"""
# standarize the shape of inputs
atom_types = np.array(atom_types, dtype=int).reshape([-1])
natoms = atom_types.size
coords = np.reshape(np.array(coords), [-1, natoms * 3])
nframes = coords.shape[0]
if cells is None:
pbc = False
cells = np.tile(np.eye(3), [nframes, 1]).reshape([nframes, 9])
else:
pbc = True
cells = np.array(cells).reshape([nframes, 9])
# sort inputs
coords, atom_types, imap, sel_at, sel_imap = self.sort_input(
coords, atom_types, sel_atoms=self.get_sel_type())
# make natoms_vec and default_mesh
natoms_vec = self.make_natoms_vec(atom_types)
assert (natoms_vec[0] == natoms)
# evaluate
feed_dict_test = {}
feed_dict_test[self.t_natoms] = natoms_vec
feed_dict_test[self.t_type] = np.tile(atom_types,
[nframes, 1]).reshape([-1])
feed_dict_test[self.t_coord] = np.reshape(coords, [-1])
feed_dict_test[self.t_box] = np.reshape(cells, [-1])
if pbc:
feed_dict_test[self.t_mesh] = make_default_mesh(cells)
else:
feed_dict_test[self.t_mesh] = np.array([], dtype=np.int32)
if atomic:
assert "global" not in self.model_type, \
f"cannot do atomic evaluation with model type {self.model_type}"
t_out = [self.t_tensor]
else:
assert self._support_gfv or "global" in self.model_type, \
f"do not support global tensor evaluation with old {self.model_type} model"
t_out = [
self.t_global_tensor if self._support_gfv else self.t_tensor
]
v_out = self.sess.run(t_out, feed_dict=feed_dict_test)
tensor = v_out[0]
# reverse map of the outputs
if atomic:
tensor = np.array(tensor)
tensor = self.reverse_map(
np.reshape(tensor, [nframes, -1, self.output_dim]), sel_imap)
tensor = np.reshape(
tensor, [nframes, len(sel_at), self.output_dim])
else:
tensor = np.reshape(tensor, [nframes, self.output_dim])
return tensor
def _prepare_feed_dict(
self,
coords,
cells,
atom_types,
fparam=None,
aparam=None,
atomic=False,
efield=None
):
# standarize the shape of inputs
natoms, nframes = self._get_natoms_and_nframes(coords, atom_types)
atom_types = np.array(atom_types, dtype = int).reshape([-1])
coords = np.reshape(np.array(coords), [-1, natoms * 3])
if cells is None:
pbc = False
# make cells to work around the requirement of pbc
cells = np.tile(np.eye(3), [nframes, 1]).reshape([nframes, 9])
else:
pbc = True
cells = np.array(cells).reshape([nframes, 9])
if self.has_fparam :
assert(fparam is not None)
fparam = np.array(fparam)
if self.has_aparam :
assert(aparam is not None)
aparam = np.array(aparam)
if self.has_efield :
assert(efield is not None), "you are using a model with external field, parameter efield should be provided"
efield = np.array(efield)
# reshape the inputs
if self.has_fparam :
fdim = self.get_dim_fparam()
if fparam.size == nframes * fdim :
fparam = np.reshape(fparam, [nframes, fdim])
elif fparam.size == fdim :
fparam = np.tile(fparam.reshape([-1]), [nframes, 1])
else :
raise RuntimeError('got wrong size of frame param, should be either %d x %d or %d' % (nframes, fdim, fdim))
if self.has_aparam :
fdim = self.get_dim_aparam()
if aparam.size == nframes * natoms * fdim:
aparam = np.reshape(aparam, [nframes, natoms * fdim])
elif aparam.size == natoms * fdim :
aparam = np.tile(aparam.reshape([-1]), [nframes, 1])
elif aparam.size == fdim :
aparam = np.tile(aparam.reshape([-1]), [nframes, natoms])
else :
raise RuntimeError('got wrong size of frame param, should be either %d x %d x %d or %d x %d or %d' % (nframes, natoms, fdim, natoms, fdim, fdim))
# sort inputs
coords, atom_types, imap = self.sort_input(coords, atom_types)
if self.has_efield:
efield = np.reshape(efield, [nframes, natoms, 3])
efield = efield[:,imap,:]
efield = np.reshape(efield, [nframes, natoms*3])
# make natoms_vec and default_mesh
natoms_vec = self.make_natoms_vec(atom_types)
assert(natoms_vec[0] == natoms)
# evaluate
feed_dict_test = {}
feed_dict_test[self.t_natoms] = natoms_vec
feed_dict_test[self.t_type ] = np.tile(atom_types, [nframes, 1]).reshape([-1])
feed_dict_test[self.t_coord] = np.reshape(coords, [-1])
feed_dict_test[self.t_box ] = np.reshape(cells , [-1])
if self.has_efield:
feed_dict_test[self.t_efield]= np.reshape(efield, [-1])
if pbc:
feed_dict_test[self.t_mesh ] = make_default_mesh(cells)
else:
feed_dict_test[self.t_mesh ] = np.array([], dtype = np.int32)
if self.has_fparam:
feed_dict_test[self.t_fparam] = np.reshape(fparam, [-1])
if self.has_aparam:
feed_dict_test[self.t_aparam] = np.reshape(aparam, [-1])
return feed_dict_test, imap