- bugfix reduce_elements functionality when elements are not ordered lexicographically (#19)

- add _configure_keras_backend to the top-most module
- bump version to 0.5.5

Co-authored-by: Yury Lysogorskiy <[email protected]>

Author: yury-lysogorskiy

docs/gracemaker/faq.md

@@ -1,3 +1,22 @@
+## Resolving the `TypeError: 'NoneType' object is not callable` in TensorFlow Callbacks
+
+If you encounter a `TypeError: 'NoneType' object is not callable` error, typically after the first epoch, the traceback will look similar to this:
+```python
+...
+    if self.monitor_op(current, self.best):
+       ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+TypeError: 'NoneType' object is not callable
+```
+This issue often occurs due to a change in how TensorFlow/Keras handles callbacks in newer versions.
+To resolve this, ensure that you set the following environment variable before running `gracemaker`:
+```bash
+export TF_USE_LEGACY_KERAS=1
+```
+
+Setting this variable forces the use of the legacy Keras backend, which resolves compatibility conflicts with certain callback implementations.
+
+---
+
 ## How to Continue a Current Fit?
 
 - Run `gracemaker -r` in the folder of the original fit to restart from the previous best-test-loss checkpoint.  

docs/gracemaker/inputfile.md

@@ -7,7 +7,7 @@ seed: 42
 cutoff: 6
 
 # cutoff_dict: {Mo: 4, MoNb: 3, W: 5, Ta*: 7 } ## Defining cutoff for each bond type separately, used by certain models
-## possible defaults: DEFAULT_CUTOFF_1L, DEFAULT_CUTOFF_2L
+## possible defaults: DEFAULT_CUTOFF_1L, DEFAULT_CUTOFF_2L, CUTOFF_2L
 
 ######################
 ##       DATA       ##
@@ -23,6 +23,11 @@ data:
   # save_dataset: False # default is True
   # stress_units: eV/A3 # eV/A3 (default) or GPa or kbar or -kbar
   # max_workers: 6 # for parallel data builder
+  
+  ## Extra input/reference DataBuilder/s required for model
+  # extra_components: {
+  #   MagMomDataBuilder: {},
+  # }
 
 ######################
 ##    POTENTIAL     ##
@@ -86,12 +91,13 @@ fit:
   maxiter: 500 # Max number of optimization epochs
   optimizer: Adam
     # Optimization with Adam: good for large number of parameters, first-order method
-  opt_params: { learning_rate: 0.01, use_ema: True, ema_momentum: 0.99,  weight_decay: 1.e-20, clipnorm: 1.0}
+  opt_params: { learning_rate: 0.008, use_ema: True, ema_momentum: 0.99,  weight_decay: 1.e-20, clipnorm: 1.0}
   # reset_optimizer: True  # reset optimizer state, after being loaded from checkpoint
   # reset_epoch_and_step: False # reset epoch and step internal counters (stored in checkpoint)
   scheduler: cosine_decay # scheduler for learning-rate reduction during training 
     # available options are: reduce_on_plateau, cosine_decay, linear_decay, exponential_decay
-  scheduler_params: {"warmup_epochs": 2, "cold_learning_rate": 0.1, "minimal_learning_rate": 0.05}
+  scheduler_params: {"minimal_learning_rate": 0.0001}
+  #scheduler_params: {"warmup_epochs": 2, "cold_learning_rate": 0.1, "minimal_learning_rate": 0.05}
     # If :warmup_epochs: > 0, begin optimization with :cold_learning_rate: and reach :opt_params::learning_rate:
     # within :warmup_epochs: (can be < 1). Else, begin optimization with :opt_params::learning_rate: and decay down to
     # minimum_learning_rate within :maxiter: epochs

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "tensorpotential"
-version = "0.5.4"  
+version = "0.5.5"
 authors = [
     { name = "Anton Bochkarev", email = "[email protected]" },
     { name = "Yury Lysogorskiy", email = "[email protected]" },

tensorpotential/__init__.py

@@ -1,3 +1,42 @@
+import os
+import sys
+import warnings
+
+
+def _configure_keras_backend(verbose=True):
+    """
+    Sets TF_USE_LEGACY_KERAS=1 and informs the user.
+    Must be run before 'import tensorflow'.
+    """
+    target_val = "1"
+    env_key = "TF_USE_LEGACY_KERAS"
+
+    existing_val = os.environ.get(env_key)
+
+    # CRITICAL CHECK: Is TensorFlow already loaded?
+    if 'tensorflow' in sys.modules and existing_val!=target_val:
+        warnings.warn(
+            f"TensorFlow was imported before {__name__} could set {env_key}={target_val}. "
+            "The flag may be ignored. Please import this package first or continue at your own risk.",
+            RuntimeWarning,
+            stacklevel=2
+        )
+        return
+
+    if existing_val is None or existing_val=='':
+        # It is missing, set it and inform.
+        os.environ[env_key] = target_val
+        if verbose:
+            msg = f"[{__name__}] Info: Environment variable {env_key} is automatically set to '{target_val}'."
+            print(msg)
+    elif existing_val not in [target_val, 'true']:
+        if verbose:
+            msg = f"[{__name__}] Warning: Environment variable {env_key} is already set to '{existing_val}', but tensorpotential requires '{target_val}'. Do it at your own risk"
+            print(msg)
+
+# Run immediately on import
+_configure_keras_backend(verbose=True)
+
 from tensorpotential.tensorpot import TensorPotential
 from tensorpotential.tpmodel import TPModel
 from tensorpotential.loss import LossFunction, L2Loss

tensorpotential/calculator/asecalculator.py

@@ -382,9 +382,38 @@ def __init__(
         )
         self.data_builders = [self.geom_data_builder]
         if constants.ATOMIC_MAGMOM in self.data_keys:
-            from tensorpotential.experimental.mag.databuilder import MagMomDataBuilder
+            try:
+                from tensorpotential.experimental.mag.databuilder import (
+                    MagMomDataBuilder,
+                )
+            except ModuleNotFoundError:
+                raise ImportError(
+                    "TensorPotential.experimental.mag.databuilder not found"
+                )
 
             self.data_builders.append(MagMomDataBuilder())
+        if constants.ATOMIC_POS in self.data_keys:
+            try:
+                from tensorpotential.experimental.gen_tensor.databuilder import (
+                    PositionsDataBuilder,
+                )
+            except ModuleNotFoundError:
+                raise ImportError(
+                    "TensorPotential.experimental.gen_tensor.databuilder not found"
+                )
+
+            self.data_builders.append(PositionsDataBuilder(cutoff=self.cutoff))
+        if constants.CELL_VECTORS in self.data_keys:
+            try:
+                from tensorpotential.experimental.gen_tensor.databuilder import (
+                    CellDataBuilder,
+                )
+            except ModuleNotFoundError:
+                raise ImportError(
+                    "TensorPotential.experimental.gen_tensor.databuilder not found"
+                )
+
+            self.data_builders.append(CellDataBuilder(cutoff=self.cutoff))
 
         self.padding_manager = PaddingManager(
             data_builders=self.data_builders,

tensorpotential/cli/train.py

@@ -25,6 +25,7 @@
     LRSchedulerFactory,
     CustomReduceLROnPlateau,
 )
+from tensorpotential.utils import NumpyEncoder
 
 LEGACY_SCHEDULER_PARAMS = "learning_rate_reduction"
 SCHEDULER_PARAMS = "scheduler_params"
@@ -41,7 +42,7 @@ def dump_metrics(filename, metrics):
     directory = os.path.dirname(filename)
     if not os.path.exists(directory):
         os.makedirs(directory, exist_ok=True)
-    json_repr = json.dumps(metrics)
+    json_repr = json.dumps(metrics, cls=NumpyEncoder)
     with open(filename, "at") as f:
         print("-", json_repr, file=f)
 

tensorpotential/constants.py

@@ -16,6 +16,10 @@
 N_STRUCTURES_BATCH_REAL: Final[str] = "batch_total_num_structures"
 N_STRUCTURES_BATCH_TOTAL: Final[str] = "n_struct_total"
 
+CELL_VECTORS: Final[str] = "cell_vectors"
+
+ATOMIC_POS: Final[str] = "atomic_positions"
+
 ATOMIC_MU_I: Final[str] = "atomic_mu_i"
 ATOMIC_MAGMOM: Final[str] = "atomic_magmom"
 ATOMS_TO_STRUCTURE_MAP: Final[str] = "map_atoms_to_structure"

tensorpotential/data/databuilder.py

@@ -21,7 +21,8 @@
 
 from tensorpotential import constants
 from tensorpotential.data.process_df import ENERGY_CORRECTED_COL, FORCES_COL, STRESS_COL
-from tensorpotential.utils import process_cutoff_dict
+from tensorpotential.utils import process_cutoff_dict, enforce_pbc
+
 
 # from ase.neighborlist import neighbor_list as nl
 
@@ -45,19 +46,6 @@ def f():
     thread.start()
 
 
-def enforce_pbc(atoms, cutoff):
-    """Enforce periodic boundary conditions for a given cutoff."""
-    pos = atoms.get_positions()
-    if (atoms.get_pbc() == 0).all():
-        max_d = np.max(np.linalg.norm(pos - pos[0], axis=1))
-        cell = np.eye(3) * ((max_d + cutoff) * 2)
-        atoms.set_cell(cell)
-        atoms.center()
-    atoms.set_pbc(True)
-
-    return atoms
-
-
 def transparent_iterator(iterator, *arg, **kwarg):
     return iterator
 

tensorpotential/instructions/compute.py

@@ -976,12 +976,21 @@ def frwrd(self, input_data: dict, training=False):
             )
 
     def get_index_to_select(self, elements_to_select):
+        # 1. Create a dictionary mapping {symbol: index} for fast O(1) lookup.
+        # We decode the bytes to string here to match the format of elements_to_select.
+        symbol_to_index_map = {
+            sym.decode(): idx
+            for idx, sym in zip(self.element_map_index.numpy(), self.element_map_symbols.numpy())
+        }
+
+        # 2. Iterate through the input list (elements_to_select) to preserve its order.
         index_to_select = []
-        for ei, es in zip(
-            self.element_map_index.numpy(), self.element_map_symbols.numpy()
-        ):
-            if es.decode() in elements_to_select:
-                index_to_select.append(ei)
+        for element in elements_to_select:
+            # Only append if the element exists in our map
+            if element in symbol_to_index_map:
+                index_to_select.append(symbol_to_index_map[element])
+            else:
+                raise ValueError(f"Element {element} not found in the map ({symbol_to_index_map}).")
 
         return tf.constant(index_to_select, dtype=tf.int32)
 

tensorpotential/scripts/grace_predict.py

@@ -23,6 +23,8 @@
 
 tqdm.pandas()
 
+NUMBER_ASSERT_ERRORS_SHOWN = 3
+
 
 def set_magmom(at, magmoms):
     magmoms = np.array(magmoms)
@@ -38,15 +40,25 @@ def set_magmom(at, magmoms):
     return at
 
 
-def predict(row, calc):
+def predict(row, calc, raise_errors):
     at = row["ase_atoms"].copy()
     if "mag_mom" in row:
         at = set_magmom(at, row["mag_mom"])
     at.calc = calc
-    e = at.get_potential_energy()
-    f = at.get_forces()
-    s = at.get_stress()
-    return {"energy": e, "forces": f, "stress": s}
+    try:
+        e = at.get_potential_energy()
+        f = at.get_forces()
+        s = at.get_stress()
+        return {"energy": e, "forces": f, "stress": s}
+    except AssertionError as e:
+        if raise_errors:
+            raise e
+        global NUMBER_ASSERT_ERRORS_SHOWN
+        if NUMBER_ASSERT_ERRORS_SHOWN > 0:
+            print("Error: ", e)
+            NUMBER_ASSERT_ERRORS_SHOWN -= 1
+            print("No more errors will be shown.")
+        return {}
 
 
 def main(args=None):
@@ -79,11 +91,21 @@ def main(args=None):
         dest="output",
     )
 
+    parser.add_argument(
+        "-e",
+        "--raise-errors",
+        help="Whether to NOT ignore errors and stop the program.",
+        action="store_true",
+        default=False,
+        dest="raise_errors",
+    )
+
     args_parse = parser.parse_args(args)
 
     model_path = os.path.abspath(args_parse.model_path)
     dataset_file = args_parse.dataset_file
     output_file = args_parse.output
+    raise_errors = args_parse.raise_errors
 
     logger.info(f"Loading model from: {model_path}")
     calc = TPCalculator(
@@ -98,11 +120,10 @@ def main(args=None):
 
     logger.info(f"Starting prediction")
 
-    df["prediction"] = df.progress_apply(predict, axis=1, args=(calc,))
-    df["energy_predicted"] = df["prediction"].map(lambda x: x["energy"])
-    df["forces_predicted"] = df["prediction"].map(lambda x: x["forces"])
-    df["stress_predicted"] = df["prediction"].map(lambda x: x["stress"])
-    # df = df.drop(columns=["ase_atoms", "prediction"])
+    df["prediction"] = df.progress_apply(predict, axis=1, args=(calc, raise_errors))
+    df["energy_predicted"] = df["prediction"].map(lambda x: x.get("energy"))
+    df["forces_predicted"] = df["prediction"].map(lambda x: x.get("forces"))
+    df["stress_predicted"] = df["prediction"].map(lambda x: x.get("stress"))
 
     logger.info(f"Saving dataset to {output_file}")
     df.drop(columns=["ase_atoms", "prediction"]).to_pickle(