def __init__(
self,
maxl,
max_sh,
num_cg_levels,
num_channels,
num_species,
cutoff_type,
hard_cut_rad,
soft_cut_rad,
soft_cut_width,
weight_init,
level_gain,
charge_power,
basis_set,
charge_scale,
bag_scale,
device=None,
dtype=None,
cg_dict=None,
) -> None:
# Parameters
level_gain = expand_var_list(level_gain, num_cg_levels)
hard_cut_rad = expand_var_list(hard_cut_rad, num_cg_levels)
soft_cut_rad = expand_var_list(soft_cut_rad, num_cg_levels)
soft_cut_width = expand_var_list(soft_cut_width, num_cg_levels)
maxl = expand_var_list(maxl, num_cg_levels)
max_sh = expand_var_list(max_sh, num_cg_levels)
num_channels = expand_var_list(num_channels, num_cg_levels + 1)
super().__init__(maxl=max(maxl + max_sh),
device=device,
dtype=dtype,
cg_dict=cg_dict)
self.num_cg_levels = num_cg_levels
self.num_channels = num_channels
self.charge_power = charge_power
self.charge_scale = charge_scale
self.bag_scale = bag_scale
self.num_species = num_species
# Set up spherical harmonics
self.sph_harms = SphericalHarmonicsRel(maxl=max(max_sh),
conj=True,
device=self.device,
dtype=self.dtype,
cg_dict=self.cg_dict)
# Set up position functions, now independent of spherical harmonics
self.rad_funcs = RadialFilters(
max_sh=max_sh,
basis_set=basis_set,
num_channels_out=num_channels,
num_levels=num_cg_levels,
device=self.device,
dtype=self.dtype,
)
tau_pos = self.rad_funcs.tau
num_scalars_in = self.num_species * (self.charge_power +
1) + self.num_species
num_scalars_out = num_channels[0]
self.input_func_atom = InputLinear(num_scalars_in,
num_scalars_out,
device=self.device,
dtype=self.dtype)
self.input_func_edge = NoLayer()
tau_in_atom = self.input_func_atom.tau
tau_in_edge = self.input_func_edge.tau
self.cormorant_cg = CormorantCG(maxl=maxl,
max_sh=max_sh,
tau_in_atom=tau_in_atom,
tau_in_edge=tau_in_edge,
tau_pos=tau_pos,
num_cg_levels=num_cg_levels,
num_channels=num_channels,
level_gain=level_gain,
weight_init=weight_init,
cutoff_type=cutoff_type,
hard_cut_rad=hard_cut_rad,
soft_cut_rad=soft_cut_rad,
soft_cut_width=soft_cut_width,
cat=True,
gaussian_mask=False,
device=self.device,
dtype=self.dtype,
cg_dict=self.cg_dict)
def __init__(self, maxl, max_sh, num_cg_levels, num_channels, num_species,
cutoff_type, hard_cut_rad, soft_cut_rad, soft_cut_width,
weight_init, level_gain, charge_power, basis_set,
charge_scale, gaussian_mask,
top, input, num_mpnn_layers, activation='leakyrelu',
device=None, dtype=None, cg_dict=None):
logging.info('Initializing network!')
level_gain = expand_var_list(level_gain, num_cg_levels)
hard_cut_rad = expand_var_list(hard_cut_rad, num_cg_levels)
soft_cut_rad = expand_var_list(soft_cut_rad, num_cg_levels)
soft_cut_width = expand_var_list(soft_cut_width, num_cg_levels)
maxl = expand_var_list(maxl, num_cg_levels)
max_sh = expand_var_list(max_sh, num_cg_levels)
num_channels = expand_var_list(num_channels, num_cg_levels+1)
logging.info('hard_cut_rad: {}'.format(hard_cut_rad))
logging.info('soft_cut_rad: {}'.format(soft_cut_rad))
logging.info('soft_cut_width: {}'.format(soft_cut_width))
logging.info('maxl: {}'.format(maxl))
logging.info('max_sh: {}'.format(max_sh))
logging.info('num_channels: {}'.format(num_channels))
super().__init__(maxl=max(maxl+max_sh), device=device, dtype=dtype, cg_dict=cg_dict)
device, dtype, cg_dict = self.device, self.dtype, self.cg_dict
print('CGDICT', cg_dict.maxl)
self.num_cg_levels = num_cg_levels
self.num_channels = num_channels
self.charge_power = charge_power
self.charge_scale = charge_scale
self.num_species = num_species
# Set up spherical harmonics
self.sph_harms = SphericalHarmonicsRel(max(max_sh), conj=True,
device=device, dtype=dtype, cg_dict=cg_dict)
# Set up position functions, now independent of spherical harmonics
self.rad_funcs = RadialFilters(max_sh, basis_set, num_channels, num_cg_levels,
device=self.device, dtype=self.dtype)
tau_pos = self.rad_funcs.tau
num_scalars_in = self.num_species * (self.charge_power + 1)
num_scalars_out = num_channels[0]
self.input_func_atom = InputMPNN(num_scalars_in, num_scalars_out, num_mpnn_layers,
soft_cut_rad[0], soft_cut_width[0], hard_cut_rad[0],
activation=activation, device=self.device, dtype=self.dtype)
self.input_func_edge = NoLayer()
tau_in_atom = self.input_func_atom.tau
tau_in_edge = self.input_func_edge.tau
self.cormorant_cg = ENN(maxl, max_sh, tau_in_atom, tau_in_edge,
tau_pos, num_cg_levels, num_channels, level_gain, weight_init,
cutoff_type, hard_cut_rad, soft_cut_rad, soft_cut_width,
cat=True, gaussian_mask=False,
device=self.device, dtype=self.dtype, cg_dict=self.cg_dict)
tau_cg_levels_atom = self.cormorant_cg.tau_levels_atom
tau_cg_levels_edge = self.cormorant_cg.tau_levels_edge
self.get_scalars_atom = GetScalarsAtom(tau_cg_levels_atom,
device=self.device, dtype=self.dtype)
self.get_scalars_edge = NoLayer()
num_scalars_atom = self.get_scalars_atom.num_scalars
num_scalars_edge = self.get_scalars_edge.num_scalars
self.output_layer_atom = OutputPMLP(num_scalars_atom, activation=activation,
device=self.device, dtype=self.dtype)
self.output_layer_edge = NoLayer()
logging.info('Model initialized. Number of parameters: {}'.format(
sum([p.nelement() for p in self.parameters()])))
def __init__(self, maxl, max_sh, num_cg_levels, num_channels, num_species,
cutoff_type, hard_cut_rad, soft_cut_rad, soft_cut_width,
weight_init, level_gain, charge_power, basis_set,
charge_scale, gaussian_mask, cgprod_bounded=True,
cg_pow_normalization='none', cg_agg_normalization='none',
device=None, dtype=None, cg_dict=None):
logging.info('Initializing network!')
level_gain = expand_var_list(level_gain, num_cg_levels)
hard_cut_rad = expand_var_list(hard_cut_rad, num_cg_levels)
soft_cut_rad = expand_var_list(soft_cut_rad, num_cg_levels)
soft_cut_width = expand_var_list(soft_cut_width, num_cg_levels)
maxl = expand_var_list(maxl, num_cg_levels)
max_sh = expand_var_list(max_sh, num_cg_levels)
num_channels = expand_var_list(num_channels, num_cg_levels+1)
logging.info('hard_cut_rad: {}'.format(hard_cut_rad))
logging.info('soft_cut_rad: {}'.format(soft_cut_rad))
logging.info('soft_cut_width: {}'.format(soft_cut_width))
logging.info('maxl: {}'.format(maxl))
logging.info('max_sh: {}'.format(max_sh))
logging.info('num_channels: {}'.format(num_channels))
super().__init__(maxl=max(maxl+max_sh), device=device, dtype=dtype, cg_dict=cg_dict)
device, dtype, cg_dict = self.device, self.dtype, self.cg_dict
self.num_cg_levels = num_cg_levels
self.num_channels = num_channels
self.charge_power = charge_power
self.charge_scale = charge_scale
self.num_species = num_species
# Set up spherical harmonics
self.sph_harms = SphericalHarmonicsRel(max(max_sh), conj=True,
device=device, dtype=dtype, cg_dict=cg_dict)
# Set up position functions, now independent of spherical harmonics
self.rad_funcs = RadialFilters(max_sh, basis_set, num_channels, num_cg_levels,
device=self.device, dtype=self.dtype)
tau_pos = self.rad_funcs.tau
# Set up input layers
num_scalars_in = self.num_species * (self.charge_power + 1)
num_scalars_out = num_channels[0]
self.input_func_atom = InputLinear(num_scalars_in, num_scalars_out,
device=self.device, dtype=self.dtype)
self.input_func_edge = NoLayer()
# Set up the central Clebsch-Gordan network
tau_in_atom = self.input_func_atom.tau
tau_in_edge = self.input_func_edge.tau
self.cormorant_cg = ENN(maxl, max_sh, tau_in_atom, tau_in_edge,
tau_pos, num_cg_levels, num_channels, level_gain, weight_init,
cutoff_type, hard_cut_rad, soft_cut_rad, soft_cut_width,
cgprod_bounded=cgprod_bounded,
cg_pow_normalization=cg_pow_normalization,
cg_agg_normalization=cg_agg_normalization,
device=self.device, dtype=self.dtype, cg_dict=self.cg_dict)
# Get atom and edge scalars
tau_cg_levels_atom = self.cormorant_cg.tau_levels_atom
tau_cg_levels_edge = self.cormorant_cg.tau_levels_edge
self.get_scalars_atom = GetScalarsAtom(tau_cg_levels_atom,
device=self.device, dtype=self.dtype)
self.get_scalars_edge = NoLayer()
# Set up the output networks
num_scalars_atom = self.get_scalars_atom.num_scalars
num_scalars_edge = self.get_scalars_edge.num_scalars
self.output_layer_atom = OutputLinear(num_scalars_atom, bias=True,
device=self.device, dtype=self.dtype)
self.output_layer_edge = NoLayer()
logging.info('Model initialized. Number of parameters: {}'.format(
sum([p.nelement() for p in self.parameters()])))