Changed MAX_CHARGE

valence-labs · Jun 3, 2024 · d52353a · d52353a
1 parent 5e1a3b4
commit d52353a
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 5 deletions.
diff --git a/openqdc/datasets/energies.py b/openqdc/datasets/energies.py
@@ -7,12 +7,11 @@
 from loguru import logger
 
 from openqdc.methods.enums import PotentialMethod
-from openqdc.utils.constants import ATOM_SYMBOLS, ATOMIC_NUMBERS
+from openqdc.utils.constants import ATOM_SYMBOLS, ATOMIC_NUMBERS, MAX_CHARGE_NUMBER
 from openqdc.utils.io import load_pkl, save_pkl
 from openqdc.utils.regressor import Regressor
 
 POSSIBLE_ENERGIES = ["formation", "regression", "null"]
-MAX_CHARGE_NUMBER = 21
 
 
 def dispatch_factory(data, **kwargs) -> "IsolatedEnergyInterface":

diff --git a/openqdc/methods/atom_energies.py b/openqdc/methods/atom_energies.py
@@ -5,7 +5,12 @@
 import numpy as np
 from loguru import logger
 
-from openqdc.utils.constants import ATOMIC_NUMBERS, MAX_ATOMIC_NUMBER, MAX_CHARGE
+from openqdc.utils.constants import (
+    ATOMIC_NUMBERS,
+    MAX_ATOMIC_NUMBER,
+    MAX_CHARGE,
+    MAX_CHARGE_NUMBER,
+)
 
 EF_KEY = Tuple[str, int]
 
@@ -35,7 +40,7 @@ def to_e_matrix(atom_energies: dict) -> np.ndarray:
                         | 2 |    |    |   |    |    |
     """
 
-    matrix = np.zeros((MAX_ATOMIC_NUMBER, MAX_CHARGE * 2 + 1))
+    matrix = np.zeros((MAX_ATOMIC_NUMBER, MAX_CHARGE_NUMBER))
     if len(atom_energies) > 0:
         for key in atom_energies.keys():
             try:

diff --git a/openqdc/utils/constants.py b/openqdc/utils/constants.py
@@ -3,7 +3,8 @@
 import numpy as np
 from rdkit import Chem
 
-MAX_CHARGE: Final[int] = 6
+MAX_CHARGE: Final[int] = 10
+MAX_CHARGE_NUMBER: Final[int] = 2 * MAX_CHARGE + 1
 
 NB_ATOMIC_FEATURES: Final[int] = 5