def __init__ (self,
descrpt : tf.Tensor,
neuron : List[int] = [120,120,120],
resnet_dt : bool = True,
sel_type : List[int] = None,
fit_diag : bool = True,
scale : List[float] = None,
diag_shift : List[float] = None,
seed : int = None,
activation_function : str = 'tanh',
precision : str = 'default'
) -> None:
"""
Constructor
"""
if not isinstance(descrpt, DescrptSeA) :
raise RuntimeError('GlobalPolarFittingSeA only supports DescrptSeA')
self.ntypes = descrpt.get_ntypes()
self.dim_descrpt = descrpt.get_dim_out()
self.polar_fitting = PolarFittingSeA(descrpt,
neuron,
resnet_dt,
sel_type,
fit_diag,
scale,
diag_shift,
seed,
activation_function,
precision)
def __init__ (self,
descrpt : tf.Tensor,
neuron : List[int] = [120,120,120],
resnet_dt : bool = True,
sel_type : List[int] = None,
seed : int = None,
activation_function : str = 'tanh',
precision : str = 'default',
uniform_seed: bool = False
) -> None:
"""
Constructor
"""
if not isinstance(descrpt, DescrptSeA) :
raise RuntimeError('DipoleFittingSeA only supports DescrptSeA')
self.ntypes = descrpt.get_ntypes()
self.dim_descrpt = descrpt.get_dim_out()
# args = ClassArg()\
# .add('neuron', list, default = [120,120,120], alias = 'n_neuron')\
# .add('resnet_dt', bool, default = True)\
# .add('sel_type', [list,int], default = [ii for ii in range(self.ntypes)], alias = 'dipole_type')\
# .add('seed', int)\
# .add("activation_function", str, default = "tanh")\
# .add('precision', str, default = "default")
# class_data = args.parse(jdata)
self.n_neuron = neuron
self.resnet_dt = resnet_dt
self.sel_type = sel_type
if self.sel_type is None:
self.sel_type = [ii for ii in range(self.ntypes)]
self.seed = seed
self.uniform_seed = uniform_seed
self.seed_shift = one_layer_rand_seed_shift()
self.fitting_activation_fn = get_activation_func(activation_function)
self.fitting_precision = get_precision(precision)
self.dim_rot_mat_1 = descrpt.get_dim_rot_mat_1()
self.dim_rot_mat = self.dim_rot_mat_1 * 3
self.useBN = False
self.fitting_net_variables = None
self.mixed_prec = None
def embed_atom_type(
ntypes: int,
natoms: tf.Tensor,
type_embedding: tf.Tensor,
):
"""
Make the embedded type for the atoms in system.
The atoms are assumed to be sorted according to the type,
thus their types are described by a `tf.Tensor` natoms, see explanation below.
Parameters
----------
ntypes:
Number of types.
natoms:
The number of atoms. This tensor has the length of Ntypes + 2
natoms[0]: number of local atoms
natoms[1]: total number of atoms held by this processor
natoms[i]: 2 <= i < Ntypes+2, number of type i atoms
type_embedding:
The type embedding.
It has the shape of [ntypes, embedding_dim]
Returns
-------
atom_embedding
The embedded type of each atom.
It has the shape of [numb_atoms, embedding_dim]
"""
te_out_dim = type_embedding.get_shape().as_list()[-1]
atype = []
for ii in range(ntypes):
atype.append(tf.tile([ii], [natoms[2 + ii]]))
atype = tf.concat(atype, axis=0)
atm_embed = tf.nn.embedding_lookup(
type_embedding, tf.cast(atype, dtype=tf.int32)) #(nf*natom)*nchnl
atm_embed = tf.reshape(atm_embed, [-1, te_out_dim])
return atm_embed
def __init__(self,
descrpt: tf.Tensor,
neuron: List[int] = [120, 120, 120],
resnet_dt: bool = True,
numb_fparam: int = 0,
numb_aparam: int = 0,
rcond: float = 1e-3,
tot_ener_zero: bool = False,
trainable: List[bool] = None,
seed: int = None,
atom_ener: List[float] = [],
activation_function: str = 'tanh',
precision: str = 'default',
uniform_seed: bool = False) -> None:
"""
Constructor
"""
# model param
self.ntypes = descrpt.get_ntypes()
self.dim_descrpt = descrpt.get_dim_out()
# args = ()\
# .add('numb_fparam', int, default = 0)\
# .add('numb_aparam', int, default = 0)\
# .add('neuron', list, default = [120,120,120], alias = 'n_neuron')\
# .add('resnet_dt', bool, default = True)\
# .add('rcond', float, default = 1e-3) \
# .add('tot_ener_zero', bool, default = False) \
# .add('seed', int) \
# .add('atom_ener', list, default = [])\
# .add("activation_function", str, default = "tanh")\
# .add("precision", str, default = "default")\
# .add("trainable", [list, bool], default = True)
self.numb_fparam = numb_fparam
self.numb_aparam = numb_aparam
self.n_neuron = neuron
self.resnet_dt = resnet_dt
self.rcond = rcond
self.seed = seed
self.uniform_seed = uniform_seed
self.seed_shift = one_layer_rand_seed_shift()
self.tot_ener_zero = tot_ener_zero
self.fitting_activation_fn = get_activation_func(activation_function)
self.fitting_precision = get_precision(precision)
self.trainable = trainable
if self.trainable is None:
self.trainable = [True for ii in range(len(self.n_neuron) + 1)]
if type(self.trainable) is bool:
self.trainable = [self.trainable] * (len(self.n_neuron) + 1)
assert (len(self.trainable) == len(self.n_neuron) +
1), 'length of trainable should be that of n_neuron + 1'
self.atom_ener = []
self.atom_ener_v = atom_ener
for at, ae in enumerate(atom_ener):
if ae is not None:
self.atom_ener.append(
tf.constant(ae,
self.fitting_precision,
name="atom_%d_ener" % at))
else:
self.atom_ener.append(None)
self.useBN = False
self.bias_atom_e = np.zeros(self.ntypes, dtype=np.float64)
# data requirement
if self.numb_fparam > 0:
add_data_requirement('fparam',
self.numb_fparam,
atomic=False,
must=True,
high_prec=False)
self.fparam_avg = None
self.fparam_std = None
self.fparam_inv_std = None
if self.numb_aparam > 0:
add_data_requirement('aparam',
self.numb_aparam,
atomic=True,
must=True,
high_prec=False)
self.aparam_avg = None
self.aparam_std = None
self.aparam_inv_std = None
self.fitting_net_variables = None
self.mixed_prec = None
def __init__ (self,
descrpt : tf.Tensor,
neuron : List[int] = [120,120,120],
resnet_dt : bool = True,
sel_type : List[int] = None,
fit_diag : bool = True,
scale : List[float] = None,
shift_diag : bool = True, # YWolfeee: will support the user to decide whether to use this function
#diag_shift : List[float] = None, YWolfeee: will not support the user to assign a shift
seed : int = None,
activation_function : str = 'tanh',
precision : str = 'default',
uniform_seed: bool = False
) -> None:
"""
Constructor
Parameters
----------
descrpt : tf.Tensor
The descrptor
neuron : List[int]
Number of neurons in each hidden layer of the fitting net
resnet_dt : bool
Time-step `dt` in the resnet construction:
y = x + dt * \phi (Wx + b)
sel_type : List[int]
The atom types selected to have an atomic polarizability prediction. If is None, all atoms are selected.
fit_diag : bool
Fit the diagonal part of the rotational invariant polarizability matrix, which will be converted to normal polarizability matrix by contracting with the rotation matrix.
scale : List[float]
The output of the fitting net (polarizability matrix) for type i atom will be scaled by scale[i]
diag_shift : List[float]
The diagonal part of the polarizability matrix of type i will be shifted by diag_shift[i]. The shift operation is carried out after scale.
seed : int
Random seed for initializing the network parameters.
activation_function : str
The activation function in the embedding net. Supported options are {0}
precision : str
The precision of the embedding net parameters. Supported options are {1}
uniform_seed
Only for the purpose of backward compatibility, retrieves the old behavior of using the random seed
"""
if not isinstance(descrpt, DescrptSeA) :
raise RuntimeError('PolarFittingSeA only supports DescrptSeA')
self.ntypes = descrpt.get_ntypes()
self.dim_descrpt = descrpt.get_dim_out()
# args = ClassArg()\
# .add('neuron', list, default = [120,120,120], alias = 'n_neuron')\
# .add('resnet_dt', bool, default = True)\
# .add('fit_diag', bool, default = True)\
# .add('diag_shift', [list,float], default = [0.0 for ii in range(self.ntypes)])\
# .add('scale', [list,float], default = [1.0 for ii in range(self.ntypes)])\
# .add('sel_type', [list,int], default = [ii for ii in range(self.ntypes)], alias = 'pol_type')\
# .add('seed', int)\
# .add("activation_function", str , default = "tanh")\
# .add('precision', str, default = "default")
# class_data = args.parse(jdata)
self.n_neuron = neuron
self.resnet_dt = resnet_dt
self.sel_type = sel_type
self.fit_diag = fit_diag
self.seed = seed
self.uniform_seed = uniform_seed
self.seed_shift = one_layer_rand_seed_shift()
#self.diag_shift = diag_shift
self.shift_diag = shift_diag
self.scale = scale
self.fitting_activation_fn = get_activation_func(activation_function)
self.fitting_precision = get_precision(precision)
if self.sel_type is None:
self.sel_type = [ii for ii in range(self.ntypes)]
if self.scale is None:
self.scale = [1.0 for ii in range(self.ntypes)]
#if self.diag_shift is None:
# self.diag_shift = [0.0 for ii in range(self.ntypes)]
if type(self.sel_type) is not list:
self.sel_type = [self.sel_type]
self.constant_matrix = np.zeros(len(self.sel_type)) # len(sel_type) x 1, store the average diagonal value
#if type(self.diag_shift) is not list:
# self.diag_shift = [self.diag_shift]
if type(self.scale) is not list:
self.scale = [self.scale]
self.dim_rot_mat_1 = descrpt.get_dim_rot_mat_1()
self.dim_rot_mat = self.dim_rot_mat_1 * 3
self.useBN = False
self.fitting_net_variables = None
self.mixed_prec = None