Fix calc

Author: knc6

alignn/__init__.py

@@ -1,3 +1,3 @@
 """Version number."""
 
-__version__ = "2025.4.1"
+__version__ = "2025.4.3"

alignn/config.py

@@ -161,6 +161,7 @@ class TrainingConfig(BaseSettings):
 
     # training configuration
     dtype: str = "float32"
+    device: str = "cpu"
     random_seed: Optional[int] = 123
     classification_threshold: Optional[float] = None
     # target_range: Optional[List] = None

alignn/ff/calculators.py

@@ -19,6 +19,7 @@
 import numpy as np
 from tqdm import tqdm
 import torch
+from alignn.config import TrainingConfig
 
 # Reference: https://doi.org/10.1039/D2DD00096B
 
@@ -209,8 +210,8 @@ def __init__(
         if path is None and model is None:
             path = default_path()
         if self.config is None:
-            config = loadjson(os.path.join(path, config_filename))
-            self.config = config
+            self.config = loadjson(os.path.join(path, config_filename))
+            self.config = TrainingConfig(**self.config).model_dump()
         if self.force_mult_natoms:
             self.config["model"]["force_mult_natoms"] = True
 
@@ -262,20 +263,23 @@ def __init__(
                     torch.load(
                         os.path.join(path, model_filename),
                         map_location=self.device,
+                        weights_only=False,
                     )
                 )
             else:
                 model.load_state_dict(
                     torch.load(
                         os.path.join(path, model_filename),
                         map_location=self.device,
+                        weights_only=False,
                     )["model"]
                 )
-            model.to(device)
             model.eval()
+            model.to(device)
             self.model = model
         else:
             model = self.model
+        self.model = self.model.to(self.device)
 
     def calculate(self, atoms, properties=None, system_changes=None):
         """Calculate properties."""
@@ -294,11 +298,13 @@ def calculate(self, atoms, properties=None, system_changes=None):
         )
         if self.config["compute_line_graph"]:
             g, lg = g
+            # print('self.model',self.model.device)
+            # print("delf.device",self.device)
             result = self.model(
                 (
                     g.to(self.device),
                     lg.to(self.device),
-                    torch.tensor(atoms.cell)
+                    torch.tensor(np.array(atoms.cell))
                     .type(torch.get_default_dtype())
                     .to(self.device),
                 )
@@ -454,6 +460,7 @@ def __init__(
             torch.load(
                 os.path.join(ff_path, ff_model_filename),
                 map_location=self.device,
+                weights_only=False,
             )
         )
         ff_model.eval()
@@ -475,10 +482,13 @@ def __init__(
             torch.load(
                 os.path.join(prop_path, prop_model_filename),
                 map_location=self.device,
+                weights_only=False,
             )
         )
         prop_model.eval()
         self.prop_model = prop_model
+        self.ff_model = self.ff_model.to(self.device)
+        self.prop_model = self.prop_model.to(self.device)
 
     def calculate(
         self,

alignn/graphs.py

@@ -452,14 +452,15 @@ def get_line_graph(
         # print(g)
         g.ndata["cart_coords"] = compute_cartesian_coordinates(g, lat)
         g.ndata["cart_coords"].requires_grad_(True)
-        r, bondlength = compute_pair_vector_and_distance(g)
-        dst_pos = g.ndata["cart_coords"][g.edges()[1]] + g.edata["images"]
-        src_pos = g.ndata["cart_coords"][g.edges()[0]]
-        bond_vec = dst_pos - src_pos
-        bond_dist = torch.norm(bond_vec, dim=1)
-        pos = g.ndata["cart_coords"]
-        g.edata["bond_dist"] = bond_dist
-        g.edata["r"] = bond_vec
+        # r, bondlength = compute_pair_vector_and_distance(g)
+        # dst_pos = g.ndata["cart_coords"][g.edges()[1]] + g.edata["images"]
+        # src_pos = g.ndata["cart_coords"][g.edges()[0]]
+        # bond_vec = dst_pos - src_pos
+        # bond_dist = torch.norm(bond_vec, dim=1)
+        # pos = g.ndata["cart_coords"]
+        # g.edata["bond_dist"] = bond_dist
+        # g.edata["r"] = bond_vec
+
         src, dst = g.edges()  # shape: [E], [E]
         pos = g.ndata["cart_coords"]
         src1 = src.unsqueeze(1)  # [E,1]
@@ -485,7 +486,7 @@ def get_line_graph(
         eid1, eid2 = angle_mask.nonzero(as_tuple=True)
         # Create the line graph
         lg = dgl.graph((eid1, eid2), num_nodes=len(src))
-        lg.ndata["r"] = bond_vec
+        lg.ndata["r"] = g.edata["r"]  # bond_vec
         lg.apply_edges(compute_bond_cosines)
         return lg
 

alignn/models/alignn.py

@@ -38,6 +38,8 @@ class ALIGNNConfig(BaseSettings):
     classification: bool = False
     num_classes: int = 2
     extra_features: int = 0
+    inner_cutoff: int = 3
+    lighten_edges: int = False
 
     class Config:
         """Configure model settings behavior."""

alignn/models/alignn_atomwise.py

@@ -374,10 +374,11 @@ def forward(
         """
         if len(g) == 3:
             g, lg, lat = g
-            lg = lg.local_var()
+            # lg = lg.local_var()
             # print('lg',lg)
             # z = self.angle_embedding(lg.edata.pop("h"))
             z = self.angle_embedding(lg.edata["h"])
+            # lg = lg.local_var()
         else:
             g, lat = g
             if len(self.alignn_layers) > 0:

alignn/models/utils.py

@@ -31,6 +31,7 @@ def __init__(
             # SchNet-style
             # set lengthscales relative to granularity of RBF expansion
             self.lengthscale = np.diff(self.centers).mean()
+
             self.gamma = 1 / self.lengthscale
 
         else:

alignn/tests/test_alignn_ff.py

@@ -20,6 +20,8 @@
 from jarvis.io.vasp.inputs import Poscar
 from alignn.ff.ff import get_figshare_model_prop, get_figshare_model_ff
 import os
+from elastic import get_elementary_deformations, get_elastic_tensor
+import elastic, ase
 
 # JVASP-25139
 pos = """Rb8
@@ -154,6 +156,24 @@ def test_jdft_mbj_gap():
     val = atoms.get_potential_energy()  # gap
 
 
+def test_elastic():
+    model_path = get_figshare_model_ff(
+        model_name="v12.2.2024_dft_3d_307k"
+    )  # default_path()
+    calc = AlignnAtomwiseCalculator(path=model_path)
+    atoms = Poscar.from_string(pos).atoms
+    ase_atoms = atoms.ase_converter()
+    ase_atoms.calc = calc
+    systems = get_elementary_deformations(ase_atoms)
+    cij_order = elastic.elastic.get_cij_order(ase_atoms)
+    Cij, Bij = get_elastic_tensor(ase_atoms, systems)
+    for i, j in zip(cij_order, Cij):
+        # print(i, j / ase.units.GPa)
+        assert j / ase.units.GPa > 0
+
+
+# test_elastic()
+# test_ev()
 # test_jdft_mbj_gap()
 # test_alexandria_gap()
 # test_alexandria_gap()

environment.yml

@@ -236,6 +236,7 @@ dependencies:
       - cryptography==42.0.5
       - dgl==2.1.0
       - dill==0.3.8
+      - elastic==5.2.5.3
       - filelock==3.13.1
       - flake8==7.0.0
       - fsspec==2024.3.1

setup.py

@@ -10,7 +10,7 @@
 
 setuptools.setup(
     name="alignn",
-    version="2025.4.1",
+    version="2025.4.3",
     author="Kamal Choudhary, Brian DeCost",
     author_email="kamal.choudhary@nist.gov",
     description="alignn",