Flexible property prediction heads

apax/bal/api.py

@@ -16,6 +16,7 @@
     check_for_ensemble,
     restore_parameters,
 )
+from apax.utils.transform import make_energy_only_model
 
 
 def create_feature_fn(
@@ -150,12 +151,13 @@ def kernel_selection(
     _, init_box = dataset.init_input()
 
     Builder = config.model.get_builder()
-    builder = Builder(config.model.get_dict(), n_species=119)
+    builder = Builder(config.model.model_dump(), n_species=119)
 
-    model = builder.build_energy_model(apply_mask=True, init_box=init_box)
+    energy_model = builder.build_energy_model(apply_mask=True, init_box=init_box)
+    energy_model = make_energy_only_model(energy_model.apply)
 
     feature_fn = create_feature_fn(
-        model, params, base_feature_map, feature_transforms, is_ensemble
+        energy_model, params, base_feature_map, feature_transforms, is_ensemble
     )
     g = compute_features(feature_fn, dataset)
     km = kernel.KernelMatrix(g, n_train)

apax/bal/feature_maps.py

@@ -65,7 +65,7 @@ def inner(ll_params):
                     inputs["box"],
                     inputs["offsets"],
                 )
-                out = model.apply(full_params, R, Z, idx, box, offsets)
+                out = model(full_params, R, Z, idx, box, offsets)
                 # take mean in case of shallow ensemble
                 # no effect for single model
                 out = jnp.mean(out)
@@ -108,7 +108,7 @@ def apply(self, model: EnergyModel) -> FeatureMap:
         def ll_grad(params, inputs):
             ll_params, remaining_params = extract_feature_params(params, self.layer_name)
 
-            energy_fn = lambda *inputs: jnp.mean(model.apply(*inputs))
+            energy_fn = lambda *inputs: jnp.mean(model(*inputs))
             force_fn = jax.grad(energy_fn, 1)
 
             def inner(ll_params):
@@ -184,7 +184,7 @@ def inner(params):
                     inputs["box"],
                     inputs["offsets"],
                 )
-                out = model.apply(params, R, Z, idx, box, offsets)
+                out = model(params, R, Z, idx, box, offsets)
                 # take mean in case of shallow ensemble
                 # no effect for single model
                 out = jnp.mean(out)
@@ -214,7 +214,7 @@ class IdentityFeatures(FeatureTransformation, extra="forbid"):
     name: Literal["identity"]
 
     def apply(self, model: EnergyModel) -> FeatureMap:
-        return model.apply
+        return model
 
 
 FeatureMapOptions = TypeAdapter(

apax/cli/apax_app.py

@@ -251,7 +251,7 @@ def visualize_model(
 
     R, Z, idx, box, offsets = make_minimal_input()
     Builder = config.model.get_builder()
-    builder = Builder(config.model.get_dict(), n_species=10)
+    builder = Builder(config.model.model_dump(), n_species=10)
     model = builder.build_energy_model()
     print(model.tabulate(jax.random.PRNGKey(0), R, Z, idx, box, offsets))
 

apax/cli/templates/train_config_full.yaml

@@ -18,8 +18,6 @@ data:
     processing: cached
     shuffle_buffer_size: 1000
 
-  additional_properties_info: {}
-
   n_train: 1000
   n_valid: 100
 

apax/config/model_config.py

@@ -114,7 +114,28 @@ class ExponentialRepulsion(Correction, extra="forbid"):
     r_max: NonNegativeFloat = 1.5
 
 
-EmpiricalCorrection = Union[ZBLRepulsion, ExponentialRepulsion]
+class LatentEwald(Correction, extra="forbid"):
+    name: Literal["latent_ewald"]
+    kgrid: list
+    sigma: float = 1.0
+
+
+EmpiricalCorrection = Union[ZBLRepulsion, ExponentialRepulsion, LatentEwald]
+
+
+class PropertyHead(BaseModel, extra="forbid"):
+    """ """
+
+    name: str
+    aggregation: str = "none"
+    mode: str = "l0"
+
+    nn: List[PositiveInt] = [128, 128]
+    n_shallow_members: int = 0
+    w_init: Literal["normal", "lecun"] = "lecun"
+    b_init: Literal["normal", "zeros"] = "zeros"
+    use_ntk: bool = False
+    dtype: Literal["fp32", "fp64"] = "fp32"
 
 
 class BaseModelConfig(BaseModel, extra="forbid"):
@@ -156,6 +177,8 @@ class BaseModelConfig(BaseModel, extra="forbid"):
 
     ensemble: Optional[EnsembleConfig] = None
 
+    property_heads: list[PropertyHead] = []
+
     # corrections
     empirical_corrections: list[EmpiricalCorrection] = []
 
@@ -165,17 +188,6 @@ class BaseModelConfig(BaseModel, extra="forbid"):
     readout_dtype: Literal["fp32", "fp64"] = "fp32"
     scale_shift_dtype: Literal["fp32", "fp64"] = "fp64"
 
-    def get_dict(self):
-        import jax.numpy as jnp
-
-        model_dict = self.model_dump()
-        prec_dict = {"fp32": jnp.float32, "fp64": jnp.float64}
-        model_dict["descriptor_dtype"] = prec_dict[model_dict["descriptor_dtype"]]
-        model_dict["readout_dtype"] = prec_dict[model_dict["readout_dtype"]]
-        model_dict["scale_shift_dtype"] = prec_dict[model_dict["scale_shift_dtype"]]
-
-        return model_dict
-
 
 class GMNNConfig(BaseModelConfig, extra="forbid"):
     """

apax/config/train_config.py

@@ -116,8 +116,6 @@ class DataConfig(BaseModel, extra="forbid"):
         | Number of validation examples to be evaluated at once.
     shuffle_buffer_size : int, default = 1000
         | Size of the `tf.data` shuffle buffer.
-    additional_properties_info : dict, optional
-        | dict of property name, shape (ragged or fixed) pairs. Currently unused.
     energy_regularisation :
         | Magnitude of the regularization in the per-element energy regression.
     pos_unit : str, default = "Ang"
@@ -141,7 +139,6 @@ class DataConfig(BaseModel, extra="forbid"):
     n_valid: PositiveInt = 100
     batch_size: PositiveInt = 32
     valid_batch_size: PositiveInt = 100
-    additional_properties_info: dict[str, str] = {}
 
     shift_method: str = "per_element_regression_shift"
     shift_options: dict = {"energy_regularisation": 1.0}

apax/data/input_pipeline.py

@@ -90,6 +90,7 @@ def __init__(
         n_jit_steps=1,
         pos_unit: str = "Ang",
         energy_unit: str = "eV",
+        additional_properties: list[tuple] = [],
         pre_shuffle=False,
         shuffle_buffer_size=1000,
         ignore_labels=False,
@@ -102,6 +103,7 @@ def __init__(
         self.n_data = len(atoms_list)
         self.batch_size = self.validate_batch_size(bs)
         self.pos_unit = pos_unit
+        self.additional_properties = additional_properties
 
         if pre_shuffle:
             shuffle(atoms_list)
@@ -112,7 +114,9 @@ def __init__(
         self.max_atoms = max_atoms
         self.max_nbrs = max_nbrs
         if atoms_list[0].calc and not ignore_labels:
-            self.labels = atoms_to_labels(atoms_list, pos_unit, energy_unit)
+            self.labels = atoms_to_labels(
+                atoms_list, pos_unit, energy_unit, additional_properties
+            )
         else:
             self.labels = None
 
@@ -161,8 +165,16 @@ def prepare_data(self, i):
             labels["forces"] = np.pad(
                 labels["forces"], ((0, zeros_to_add), (0, 0)), "constant"
             )
+
+        for prop in self.additional_properties:
+            name, shape = prop
+            if shape[0] == "natoms":
+                pad_shape = [(0, zeros_to_add)] + [(0, 0)] * (len(shape) - 1)
+                labels[name] = np.pad(labels[name], pad_shape, "constant")
+
         inputs = {k: tf.constant(v) for k, v in inputs.items()}
         labels = {k: tf.constant(v) for k, v in labels.items()}
+
         return (inputs, labels)
 
     def enqueue(self, num_elements):
@@ -202,6 +214,14 @@ def make_signature(self) -> tf.TensorSpec:
             label_signature["stress"] = tf.TensorSpec(
                 (3, 3), dtype=tf.float64, name="stress"
             )
+
+        for prop in self.additional_properties:
+            name, shape = prop
+            if shape[0] == "natoms":
+                shape[0] = self.max_atoms
+
+            sig = tf.TensorSpec(tuple(shape), dtype=tf.float64, name=name)
+            label_signature[name] = sig
         signature = (input_signature, label_signature)
         return signature
 
@@ -377,14 +397,21 @@ def round_up_to_multiple(value, multiple):
 
 class BatchProcessor:
     def __init__(
-        self, cutoff, atom_padding: int, nl_padding: int, forces=True, stress=False
+        self,
+        cutoff,
+        atom_padding: int,
+        nl_padding: int,
+        forces=True,
+        stress=False,
+        additional_properties=[],
     ) -> None:
         self.cutoff = cutoff
         self.atom_padding = atom_padding
         self.nl_padding = nl_padding
 
         self.forces = forces
         self.stress = stress
+        self.additional_properties = additional_properties
 
     def __call__(self, samples: list[dict]):
         n_samples = len(samples)
@@ -401,7 +428,12 @@ def __call__(self, samples: list[dict]):
         labels = {
             "energy": np.zeros(n_samples, dtype=np.float64),
         }
-
+        for prop in self.additional_properties:
+            name, shape = prop
+            if shape[0] == "natoms":
+                shape = [max_atoms] + shape[1:]
+            shape = [n_samples] + shape
+            labels[name] = np.zeros(shape, dtype=np.float64)
         if self.forces:
             labels["forces"] = np.zeros((n_samples, max_atoms, 3), dtype=np.float64)
         if self.stress:
@@ -425,6 +457,13 @@ def __call__(self, samples: list[dict]):
             if self.stress:
                 labels["stress"][i] = lab["stress"]
 
+            for prop in self.additional_properties:
+                name, shape = prop
+                if shape[0] == "natoms":
+                    labels[name][i, : inp["n_atoms"]] = lab[name]
+                else:
+                    labels[name][i] = lab[name]
+
         max_nbrs = np.max([idx.shape[1] for idx in idxs])
         max_nbrs = round_up_to_multiple(max_nbrs, self.nl_padding)
 
@@ -470,6 +509,7 @@ def __init__(
         nl_padding: int = 2000,
         pos_unit: str = "Ang",
         energy_unit: str = "eV",
+        additional_properties=[],
         pre_shuffle=False,
     ) -> None:
         self.cutoff = cutoff
@@ -478,6 +518,7 @@ def __init__(
         self.n_data = len(atoms_list)
         self.batch_size = self.validate_batch_size(bs)
         self.pos_unit = pos_unit
+        self.additional_properties = additional_properties
 
         if num_workers:
             self.num_workers = num_workers
@@ -488,7 +529,9 @@ def __init__(
 
         # Transform atoms into inputs and labels
         self.inputs = atoms_to_inputs(atoms_list, pos_unit)
-        self.labels = atoms_to_labels(atoms_list, pos_unit, energy_unit)
+        self.labels = atoms_to_labels(
+            atoms_list, pos_unit, energy_unit, additional_properties
+        )
         label_keys = self.labels.keys()
 
         self.data = list(
@@ -500,7 +543,7 @@ def __init__(
         forces = "forces" in label_keys
         stress = "stress" in label_keys
         self.prepare_batch = BatchProcessor(
-            cutoff, atom_padding, nl_padding, forces, stress
+            cutoff, atom_padding, nl_padding, forces, stress, additional_properties
         )
 
         self.count = 0

apax/layers/descriptor/basis_functions.py

@@ -6,6 +6,7 @@
 import numpy as np
 
 from apax.layers.initializers import uniform_range
+from apax.utils.convert import str_to_dtype
 
 
 class GaussianBasis(nn.Module):
@@ -15,6 +16,8 @@ class GaussianBasis(nn.Module):
     dtype: Any = jnp.float32
 
     def setup(self):
+        dtype = str_to_dtype(self.dtype)
+
         self.betta = self.n_basis**2 / self.r_max**2
         self.rad_norm = (2.0 * self.betta / np.pi) ** 0.25
         shifts = self.r_min + (self.r_max - self.r_min) / self.n_basis * np.arange(
@@ -23,7 +26,7 @@ def setup(self):
 
         # shape: 1 x n_basis
         shifts = einops.repeat(shifts, "n_basis -> 1 n_basis")
-        self.shifts = jnp.asarray(shifts, dtype=self.dtype)
+        self.shifts = jnp.asarray(shifts, dtype=dtype)
 
     def __call__(self, dr):
         dr = einops.repeat(dr, "neighbors -> neighbors 1")
@@ -47,7 +50,8 @@ class BesselBasis(nn.Module):
     dtype: Any = jnp.float32
 
     def setup(self):
-        self.n = jnp.arange(self.n_basis, dtype=self.dtype)
+        dtype = str_to_dtype(self.dtype)
+        self.n = jnp.arange(self.n_basis, dtype=dtype)
 
     def __call__(self, dr):
         dr = einops.repeat(dr, "neighbors -> neighbors 1")
@@ -69,10 +73,9 @@ class RadialFunction(nn.Module):
     dtype: Any = jnp.float32
 
     def setup(self):
+        dtype = str_to_dtype(self.dtype)
         self.r_max = self.basis_fn.r_max
-        self.embed_norm = jnp.array(
-            1.0 / np.sqrt(self.basis_fn.n_basis), dtype=self.dtype
-        )
+        self.embed_norm = jnp.array(1.0 / np.sqrt(self.basis_fn.n_basis), dtype=dtype)
         if self.one_sided_dist:
             lower_bound = 0.0
         else:
@@ -81,7 +84,7 @@ def setup(self):
         if self.emb_init is not None:
             self._n_radial = self.n_radial
             if self.emb_init == "uniform":
-                emb_initializer = uniform_range(lower_bound, 1.0, dtype=self.dtype)
+                emb_initializer = uniform_range(lower_bound, 1.0, dtype=dtype)
                 self.embeddings = self.param(
                     "atomic_type_embedding",
                     emb_initializer,
@@ -91,7 +94,7 @@ def setup(self):
                         self.n_radial,
                         self.basis_fn.n_basis,
                     ),
-                    self.dtype,
+                    dtype,
                 )
             else:
                 raise ValueError(
@@ -102,7 +105,8 @@ def setup(self):
             self._n_radial = self.basis_fn.n_basis
 
     def __call__(self, dr, Z_i, Z_j):
-        dr = dr.astype(self.dtype)
+        dtype = str_to_dtype(self.dtype)
+        dr = dr.astype(dtype)
         # basis shape: neighbors x n_basis
         basis = self.basis_fn(dr)
 
@@ -128,6 +132,6 @@ def __call__(self, dr, Z_i, Z_j):
 
         radial_function = radial_function * cutoff
 
-        assert radial_function.dtype == self.dtype
+        assert radial_function.dtype == dtype
 
         return radial_function