Merge pull request #181 from usnistgov/redw

Faster training with compute_line_graph False/lg_on_fly

Author: knc6

alignn/__init__.py

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

alignn/ff/calculators.py

@@ -213,6 +213,7 @@ def __init__(
             self.config = config
         if self.force_mult_natoms:
             self.config["model"]["force_mult_natoms"] = True
+
         if self.include_stress:
             self.implemented_properties = ["energy", "forces", "stress"]
             if (
@@ -226,6 +227,11 @@ def __init__(
 
         else:
             self.implemented_properties = ["energy", "forces"]
+        if (
+            "calculate_gradient" in self.config["model"]
+            and self.config["model"]["calculate_gradient"]
+        ):
+            self.trained_stress = True
 
         if (
             batch_stress is not None
@@ -296,24 +302,57 @@ def calculate(self, atoms, properties=None, system_changes=None):
             )
         else:
             result = self.model(
-                (g.to(self.device, torch.tensor(atoms.cell).to(self.device)))
+                (g.to(self.device), torch.tensor(atoms.cell).to(self.device))
             )
+        # print("result",result)
         if "atomwise" in self.config["model"]["name"]:
             forces = forces = (
                 result["grad"].detach().cpu().numpy() * self.force_multiplier
             )
         else:
             forces = np.zeros((3, 3))
-        if "atomwise" in self.config["model"]["name"] and self.trained_stress:
-            stress = (
-                full_3x3_to_voigt_6_stress(
-                    result["stresses"][:3].reshape(3, 3).detach().cpu().numpy()
+        # print("self.trained_stress",self.trained_stress)
+        # if self.trained_stress:
+        #    # if "atomwise" in self.config["model"]["name"]
+        #    # and self.trained_stress:
+        #    stress = (
+        #        full_3x3_to_voigt_6_stress(
+        #            result["stresses"][:3].r
+        # eshape(3, 3).detach().cpu().numpy()
+        #        )
+        #        * self.stress_wt
+        #        / 160.21766208
+        #    )
+        # else:
+        #    stress = np.zeros((3, 3))
+        if (
+            "calculate_gradient" in self.config["model"]
+            and self.config["model"]["calculate_gradient"]
+        ):
+            try:
+                stress = (
+                    full_3x3_to_voigt_6_stress(
+                        result["stresses"][:3]
+                        .reshape(3, 3)
+                        .detach()
+                        .cpu()
+                        .numpy()
+                    )
+                    * self.stress_wt
+                    / 160.21766208
                 )
-                * self.stress_wt
-                / 160.21766208
-            )
+            except Exception:
+                stress = np.zeros((3, 3))
+                pass
         else:
             stress = np.zeros((3, 3))
+        # stress = (
+        #    full_3x3_to_voigt_6_stress(
+        #        result["stresses"][:3].reshape(3, 3).detach().cpu().numpy()
+        #    )
+        #    * self.stress_wt
+        #    / 160.21766208
+        # )
         if "atomwise" in self.config["model"]["name"]:
             energy = result["out"].detach().cpu().numpy()
         else:
@@ -325,6 +364,7 @@ def calculate(self, atoms, properties=None, system_changes=None):
         if self.force_mult_batchsize:
             forces *= self.config["batch_size"]
 
+        # print("stress cal",stress)
         self.results = {
             "energy": energy,
             "forces": forces,
@@ -449,6 +489,7 @@ def calculate(
             atom_features=self.ff_config["atom_features"],
             use_canonize=self.ff_config["use_canonize"],
         )
+        # print("config",self.ff_config)
         result_ff = self.ff_model(
             (
                 g.to(self.device),
@@ -459,6 +500,7 @@ def calculate(
             )
         )
         forces = forces = result_ff["grad"].detach().cpu().numpy()
+
         stress = (
             full_3x3_to_voigt_6_stress(
                 result_ff["stresses"][:3].reshape(3, 3).detach().cpu().numpy()

alignn/models/alignn_atomwise.py

@@ -371,20 +371,32 @@ def forward(
         z: angle features (lg.edata)
         """
         if len(self.alignn_layers) > 0:
-            g, lg, lat = g
-            lg = lg.local_var()
-            # print('lg',lg)
-            # angle features (fixed)
-            z = self.angle_embedding(lg.edata.pop("h"))
+            if len(g) == 3:
+                g, lg, lat = g
+                lg = lg.local_var()
+                # z = self.angle_embedding(lg.edata.pop("h"))
+                z = self.angle_embedding(lg.edata["h"])
+            else:
+                g, lat = 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)
+                lg = g.line_graph(shared=True)
+                lg.ndata["r"] = r
+                lg.apply_edges(compute_bond_cosines)
+                # print('lg',lg)
+                # angle features (fixed)
+        else:
+            g, lat = g
         if self.config.extra_features != 0:
             features = g.ndata["extra_features"]
             # print('features',features,features.shape)
             features = self.extra_feature_embedding(features)
-        g = g.local_var()
+        # g = g.local_var()
         result = {}
-
         # initial node features: atom feature network...
-        x = g.ndata.pop("atom_features")
+        x = g.ndata["atom_features"]
+        # x = g.ndata.pop("atom_features")
         # print('x1',x,x.shape)
 
         x = self.atom_embedding(x)
@@ -416,7 +428,8 @@ def forward(
 
             lg.ndata["r"] = r  # overwrites precomputed r values
             lg.apply_edges(compute_bond_cosines)  # overwrites precomputed h
-            z = self.angle_embedding(lg.edata.pop("h"))
+            z = self.angle_embedding(lg.edata["h"])
+            # z = self.angle_embedding(lg.edata.pop("h"))
 
         # r = g.edata["r"].clone().detach().requires_grad_(True)
         if self.config.use_cutoff_function:
@@ -623,6 +636,7 @@ def forward(
                         stress = self.config.stress_multiplier * torch.stack(
                             stresses
                         )
+                        # print("stress",stress)
                     # print("stress2", stress, stress.shape)
                     # virial = (
                     #    160.21766208

alignn/tests/test_alignn_ff.py

@@ -154,6 +154,9 @@ def test_jdft_mbj_gap():
     val = atoms.get_potential_energy()  # gap
 
 
+# test_jdft_mbj_gap()
+# test_alexandria_gap()
+# test_alexandria_gap()
 # print('test_graph_builder')
 # test_graph_builder()
 # print('test_ev')

setup.py

@@ -10,18 +10,18 @@
 
 setuptools.setup(
     name="alignn",
-    version="2024.12.12",
+    version="2025.4.1",
     author="Kamal Choudhary, Brian DeCost",
     author_email="kamal.choudhary@nist.gov",
     description="alignn",
     install_requires=[
-        "numpy>=1.19.5",
+        "numpy<2.0",
         # "numpy>=1.19.5,<2.0.0",
         "scipy>=1.6.1",
         "jarvis-tools>=2021.07.19",
-        "torch>=2.0.0",
+        "torch<=2.2.1",
         "mpmath<=1.3.0",
-        "dgl>=0.6.0",
+        "dgl<=1.1.1",
         "spglib>=2.0.2",
         "scikit-learn>=0.22.2",
         "matplotlib>=3.4.1",