Model

class apax.nn.models.EnergyDerivativeModel(energy_model: EnergyModel = EnergyModel(     # attributes     representation = GaussianMomentDescriptor(         # attributes         radial_fn = RadialFunction(             # attributes             n_radial = 5             basis_fn = GaussianBasis(                 # attributes                 n_basis = 7                 r_min = 0.5                 r_max = 6.0                 dtype = float32                 spacing = 'linear'             )             n_species = 119             emb_init = 'uniform'             use_embed_norm = True             one_sided_dist = False             dtype = float32         )         n_contr = 8         dtype = float32         apply_mask = True     )     readout = AtomisticReadout(         # attributes         units = [32, 32]         activation_fn=silu         w_init = 'normal'         b_init = 'zeros'         use_ntk = True         n_shallow_ensemble = 0         is_feature_fn = False         dtype = float32     )     scale_shift = PerElementScaleShift(         # attributes         n_species = 119         scale = 1.0         shift = 0.0         dtype = float32     )     property_heads = []     corrections = []     init_box = array([0., 0., 0.])     mask_atoms=True     inference_disp_fn = None ), calc_stress: bool = False, parent: Module | Scope | _Sentinel | None = <flax.linen.module._Sentinel object>, name: str | None = None)

Transforms an EnergyModel into one that also predicts derivatives the total energy. Can calculate forces and stress tensors.

class apax.nn.models.EnergyModel(representation: Module = GaussianMomentDescriptor(     # attributes     radial_fn = RadialFunction(         # attributes         n_radial = 5         basis_fn = GaussianBasis(             # attributes             n_basis = 7             r_min = 0.5             r_max = 6.0             dtype = float32             spacing = 'linear'         )         n_species = 119         emb_init = 'uniform'         use_embed_norm = True         one_sided_dist = False         dtype = float32     )     n_contr = 8     dtype = float32     apply_mask = True ), readout: Module = AtomisticReadout(     # attributes     units = [32, 32]     activation_fn=silu     w_init = 'normal'     b_init = 'zeros'     use_ntk = True     n_shallow_ensemble = 0     is_feature_fn = False     dtype = float32 ), scale_shift: Module = PerElementScaleShift(     # attributes     n_species = 119     scale = 1.0     shift = 0.0     dtype = float32 ), property_heads: list[Module] = <factory>, corrections: list[EmpiricalEnergyTerm] = <factory>, init_box: array = <factory>, mask_atoms: bool = True, inference_disp_fn: Any = None, parent: Module | Scope | _Sentinel | None = <flax.linen.module._Sentinel object>, name: str | None = None)

Model which post processes the output of an atomistic model and adds empirical energy terms.

setup()

Initializes a Module lazily (similar to a lazy __init__).

setup is called once lazily on a module instance when a module is bound, immediately before any other methods like __call__ are invoked, or before a setup-defined attribute on self is accessed.

This can happen in three cases:

1. Immediately when invoking apply(), init() or init_and_output().

2. Once the module is given a name by being assigned to an attribute of another module inside the other module’s setup method (see __setattr__()):

   >>> class MyModule(nn.Module):
   ...   def setup(self):
   ...     submodule = nn.Conv(...)
   
   ...     # Accessing `submodule` attributes does not yet work here.
   
   ...     # The following line invokes `self.__setattr__`, which gives
   ...     # `submodule` the name "conv1".
   ...     self.conv1 = submodule
   
   ...     # Accessing `submodule` attributes or methods is now safe and
   ...     # either causes setup() to be called once.
   

3. Once a module is constructed inside a method wrapped with compact(), immediately before another method is called or setup defined attribute is accessed.

class apax.nn.models.FeatureModel(representation: Module = GaussianMomentDescriptor(     # attributes     radial_fn = RadialFunction(         # attributes         n_radial = 5         basis_fn = GaussianBasis(             # attributes             n_basis = 7             r_min = 0.5             r_max = 6.0             dtype = float32             spacing = 'linear'         )         n_species = 119         emb_init = 'uniform'         use_embed_norm = True         one_sided_dist = False         dtype = float32     )     n_contr = 8     dtype = float32     apply_mask = True ), readout: Module = AtomisticReadout(     # attributes     units = [32, 32]     activation_fn=silu     w_init = 'normal'     b_init = 'zeros'     use_ntk = True     n_shallow_ensemble = 0     is_feature_fn = False     dtype = float32 ), should_average: bool = False, init_box: array = <factory>, inference_disp_fn: Any = None, mask_atoms: bool = True, parent: Module | Scope | _Sentinel | None = <flax.linen.module._Sentinel object>, name: str | None = None)

Model wraps some submodel (e.g. a descriptor) to supply distance computation.

setup()

Initializes a Module lazily (similar to a lazy __init__).

setup is called once lazily on a module instance when a module is bound, immediately before any other methods like __call__ are invoked, or before a setup-defined attribute on self is accessed.

This can happen in three cases:

1. Immediately when invoking apply(), init() or init_and_output().

2. Once the module is given a name by being assigned to an attribute of another module inside the other module’s setup method (see __setattr__()):

   >>> class MyModule(nn.Module):
   ...   def setup(self):
   ...     submodule = nn.Conv(...)
   
   ...     # Accessing `submodule` attributes does not yet work here.
   
   ...     # The following line invokes `self.__setattr__`, which gives
   ...     # `submodule` the name "conv1".
   ...     self.conv1 = submodule
   
   ...     # Accessing `submodule` attributes or methods is now safe and
   ...     # either causes setup() to be called once.
   

3. Once a module is constructed inside a method wrapped with compact(), immediately before another method is called or setup defined attribute is accessed.

class apax.nn.models.ShallowEnsembleModel(energy_model: EnergyModel = EnergyModel(     # attributes     representation = GaussianMomentDescriptor(         # attributes         radial_fn = RadialFunction(             # attributes             n_radial = 5             basis_fn = GaussianBasis(                 # attributes                 n_basis = 7                 r_min = 0.5                 r_max = 6.0                 dtype = float32                 spacing = 'linear'             )             n_species = 119             emb_init = 'uniform'             use_embed_norm = True             one_sided_dist = False             dtype = float32         )         n_contr = 8         dtype = float32         apply_mask = True     )     readout = AtomisticReadout(         # attributes         units = [32, 32]         activation_fn=silu         w_init = 'normal'         b_init = 'zeros'         use_ntk = True         n_shallow_ensemble = 0         is_feature_fn = False         dtype = float32     )     scale_shift = PerElementScaleShift(         # attributes         n_species = 119         scale = 1.0         shift = 0.0         dtype = float32     )     property_heads = []     corrections = []     init_box = array([0., 0., 0.])     mask_atoms=True     inference_disp_fn = None ), calc_stress: bool = False, force_variance: bool = True, chunk_size: int | None = None, parent: Module | Scope | _Sentinel | None = <flax.linen.module._Sentinel object>, name: str | None = None)

Transforms an EnergyModel into one that also predicts derivatives the total energy. Can calculate forces and stress tensors.