From 81b6af663533c9fb8c3819c8f0445e31860b7093 Mon Sep 17 00:00:00 2001
From: Kamal Choudhary <writetokamal.1989@gmail.com>
Date: Mon, 18 Aug 2025 14:02:16 -0400
Subject: [PATCH 01/23] Update setup.py

---
 setup.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/setup.py b/setup.py
index b9c616c..4059a63 100644
--- a/setup.py
+++ b/setup.py
@@ -21,7 +21,7 @@
         "jarvis-tools>=2021.07.19",
         "torch<=2.2.1",
         "mpmath<=1.3.0",
-        "dgl<=1.1.1",
+        "dgl>=0.6.0",
         "spglib>=2.0.2",
         "scikit-learn>=0.22.2",
         "matplotlib>=3.4.1",

From 35b02192b4c452035b7e76fb4774d82d3b209519 Mon Sep 17 00:00:00 2001
From: knc6 <writetokamal.1989@gmail.com>
Date: Sun, 24 Aug 2025 21:20:13 -0400
Subject: [PATCH 02/23] weights_only=False

---
 alignn/ff/calculators.py | 4 ++++
 alignn/pretrained.py     | 4 +++-
 2 files changed, 7 insertions(+), 1 deletion(-)

diff --git a/alignn/ff/calculators.py b/alignn/ff/calculators.py
index 85f77ff..5f3eb02 100644
--- a/alignn/ff/calculators.py
+++ b/alignn/ff/calculators.py
@@ -262,6 +262,7 @@ def __init__(
                     torch.load(
                         os.path.join(path, model_filename),
                         map_location=self.device,
+                        weights_only=False,
                     )
                 )
             else:
@@ -269,6 +270,7 @@ def __init__(
                     torch.load(
                         os.path.join(path, model_filename),
                         map_location=self.device,
+                        weights_only=False,
                     )["model"]
                 )
             model.to(device)
@@ -445,6 +447,7 @@ def __init__(
             torch.load(
                 os.path.join(ff_path, ff_model_filename),
                 map_location=self.device,
+                weights_only=False,
             )
         )
         ff_model.eval()
@@ -466,6 +469,7 @@ def __init__(
             torch.load(
                 os.path.join(prop_path, prop_model_filename),
                 map_location=self.device,
+                weights_only=False,
             )
         )
         prop_model.eval()
diff --git a/alignn/pretrained.py b/alignn/pretrained.py
index d016e2c..3c10274 100644
--- a/alignn/pretrained.py
+++ b/alignn/pretrained.py
@@ -295,7 +295,9 @@ def get_figshare_model(model_name="jv_formation_energy_peratom_alignn"):
     new_file, filename = tempfile.mkstemp()
     with open(filename, "wb") as f:
         f.write(data)
-    model.load_state_dict(torch.load(filename, map_location=device)["model"])
+    model.load_state_dict(
+        torch.load(filename, map_location=device, weights_only=True)["model"]
+    )
     model.to(device)
     model.eval()
     if os.path.exists(filename):

From 24888918265d0620ad66d1bbca8b29c9dfefdcfa Mon Sep 17 00:00:00 2001
From: knc6 <writetokamal.1989@gmail.com>
Date: Fri, 3 Oct 2025 20:53:07 -0400
Subject: [PATCH 03/23] False weight

---
 alignn/pretrained.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/alignn/pretrained.py b/alignn/pretrained.py
index 3c10274..9243847 100644
--- a/alignn/pretrained.py
+++ b/alignn/pretrained.py
@@ -296,7 +296,7 @@ def get_figshare_model(model_name="jv_formation_energy_peratom_alignn"):
     with open(filename, "wb") as f:
         f.write(data)
     model.load_state_dict(
-        torch.load(filename, map_location=device, weights_only=True)["model"]
+        torch.load(filename, map_location=device, weights_only=False)["model"]
     )
     model.to(device)
     model.eval()

From 1244b02299c50bf0ccfe7485b8bcf4d5e0563a21 Mon Sep 17 00:00:00 2001
From: user <writetokamal.1989@gmail.com>
Date: Sat, 6 Dec 2025 16:18:34 -0500
Subject: [PATCH 04/23] Minor bug fix

---
 alignn/__init__.py       | 2 +-
 alignn/ff/calculators.py | 6 +++++-
 setup.py                 | 2 +-
 3 files changed, 7 insertions(+), 3 deletions(-)

diff --git a/alignn/__init__.py b/alignn/__init__.py
index ca6a785..cdc30cc 100644
--- a/alignn/__init__.py
+++ b/alignn/__init__.py
@@ -1,3 +1,3 @@
 """Version number."""
 
-__version__ = "2025.4.1"
+__version__ = "2025.11.1"
diff --git a/alignn/ff/calculators.py b/alignn/ff/calculators.py
index 5f3eb02..c96bc58 100644
--- a/alignn/ff/calculators.py
+++ b/alignn/ff/calculators.py
@@ -273,7 +273,7 @@ def __init__(
                         weights_only=False,
                     )["model"]
                 )
-            model.to(device)
+            model.to(self.device)
             model.eval()
             self.model = model
         else:
@@ -291,6 +291,10 @@ def calculate(self, atoms, properties=None, system_changes=None):
             atom_features=self.config["atom_features"],
             use_canonize=self.config["use_canonize"],
         )
+        print("self.devicee", self.device)
+        print("g", g.device)
+        print("lg", lg.device)
+        print("model", self.model)
 
         if self.config["model"]["alignn_layers"] > 0:
             result = self.model(
diff --git a/setup.py b/setup.py
index 4059a63..d4ab025 100644
--- a/setup.py
+++ b/setup.py
@@ -10,7 +10,7 @@
 
 setuptools.setup(
     name="alignn",
-    version="2025.4.1",
+    version="2025.11.1",
     author="Kamal Choudhary, Brian DeCost",
     author_email="kamal.choudhary@nist.gov",
     description="alignn",

From d760ab9f6990d303d99076514a8642374d90edd2 Mon Sep 17 00:00:00 2001
From: Kamal Choudhary <writetokamal.1989@gmail.com>
Date: Fri, 30 Jan 2026 08:30:22 -0500
Subject: [PATCH 05/23] Update setup.py

---
 setup.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/setup.py b/setup.py
index d4ab025..3bf9207 100644
--- a/setup.py
+++ b/setup.py
@@ -19,7 +19,7 @@
         # "numpy>=1.19.5,<2.0.0",
         "scipy>=1.6.1",
         "jarvis-tools>=2021.07.19",
-        "torch<=2.2.1",
+        "torch>=2.2.1",
         "mpmath<=1.3.0",
         "dgl>=0.6.0",
         "spglib>=2.0.2",

From 8c4d7804f608729e231deed21d3d6d7ab732daac Mon Sep 17 00:00:00 2001
From: Kamal Choudhary <writetokamal.1989@gmail.com>
Date: Fri, 30 Jan 2026 08:35:58 -0500
Subject: [PATCH 06/23] Update setup.py

---
 setup.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/setup.py b/setup.py
index 3bf9207..0e08d66 100644
--- a/setup.py
+++ b/setup.py
@@ -20,6 +20,7 @@
         "scipy>=1.6.1",
         "jarvis-tools>=2021.07.19",
         "torch>=2.2.1",
+        # "torch<=2.2.1",
         "mpmath<=1.3.0",
         "dgl>=0.6.0",
         "spglib>=2.0.2",
@@ -32,7 +33,8 @@
         "flake8>=3.9.1",
         "pycodestyle>=2.7.0",
         "pydocstyle>=6.0.0",
-        "pyparsing>=2.2.1,<3",
+        # "pyparsing>=2.2.1,<3",
+        "pyparsing>=2.2.1",
         "ase",
         "lmdb",
     ],

From 6ce5a3bbc917606b644df28ec3892c327dbf4058 Mon Sep 17 00:00:00 2001
From: user <writetokamal.1989@gmail.com>
Date: Fri, 30 Jan 2026 18:09:06 -0500
Subject: [PATCH 07/23] Pretrained figshare link fix

---
 alignn/ff/all_models_alignn.json          | 116 +++++++++++-----------
 alignn/ff/all_models_alignn_atomwise.json |  52 +++++-----
 alignn/pretrained.py                      | 114 ++++++++++-----------
 3 files changed, 141 insertions(+), 141 deletions(-)

diff --git a/alignn/ff/all_models_alignn.json b/alignn/ff/all_models_alignn.json
index b145a5f..3c7c872 100644
--- a/alignn/ff/all_models_alignn.json
+++ b/alignn/ff/all_models_alignn.json
@@ -1,59 +1,59 @@
 {
-    "jv_formation_energy_peratom_alignn": "https://figshare.com/ndownloader/files/31458679",
-    "jv_optb88vdw_total_energy_alignn": "https://figshare.com/ndownloader/files/31459642",
-    "jv_optb88vdw_bandgap_alignn": "https://figshare.com/ndownloader/files/31459636",
-    "jv_mbj_bandgap_alignn": "https://figshare.com/ndownloader/files/31458694",
-    "jv_spillage_alignn": "https://figshare.com/ndownloader/files/31458736",
-    "jv_slme_alignn": "https://figshare.com/ndownloader/files/31458727",
-    "jv_bulk_modulus_kv_alignn": "https://figshare.com/ndownloader/files/31458649",
-    "jv_shear_modulus_gv_alignn": "https://figshare.com/ndownloader/files/31458724",
-    "jv_n-Seebeck_alignn": "https://figshare.com/ndownloader/files/31458718",
-    "jv_n-powerfact_alignn": "https://figshare.com/ndownloader/files/31458712",
-    "intermat_cbm": "https://figshare.com/ndownloader/files/45392908",
-    "intermat_vbm": "https://figshare.com/ndownloader/files/45392914",
-    "intermat_phi": "https://figshare.com/ndownloader/files/45392911",
-    "jv_magmom_oszicar_alignn": "https://figshare.com/ndownloader/files/31458685",
-    "jv_kpoint_length_unit_alignn": "https://figshare.com/ndownloader/files/31458682",
-    "jv_avg_elec_mass_alignn": "https://figshare.com/ndownloader/files/31458643",
-    "jv_avg_hole_mass_alignn": "https://figshare.com/ndownloader/files/31458646",
-    "jv_epsx_alignn": "https://figshare.com/ndownloader/files/31458667",
-    "jv_mepsx_alignn": "https://figshare.com/ndownloader/files/31458703",
-    "jv_max_efg_alignn": "https://figshare.com/ndownloader/files/31458691",
-    "jv_ehull_alignn": "https://figshare.com/ndownloader/files/31458658",
-    "jv_dfpt_piezo_max_dielectric_alignn": "https://figshare.com/ndownloader/files/31458652",
-    "jv_dfpt_piezo_max_dij_alignn": "https://figshare.com/ndownloader/files/31458655",
-    "jv_exfoliation_energy_alignn": "https://figshare.com/ndownloader/files/31458676",
-    "jv_supercon_tc_alignn": "https://figshare.com/ndownloader/files/38789199",
-    "jv_supercon_edos_alignn": "https://figshare.com/ndownloader/files/39946300",
-    "jv_supercon_debye_alignn": "https://figshare.com/ndownloader/files/39946297",
-    "jv_supercon_a2F_alignn": "https://figshare.com/ndownloader/files/38801886",
-    "mp_e_form_alignn": "https://figshare.com/ndownloader/files/31458811",
-    "mp_gappbe_alignn": "https://figshare.com/ndownloader/files/31458814",
-    "tinnet_O_alignn": "https://figshare.com/ndownloader/files/41962800",
-    "tinnet_N_alignn": "https://figshare.com/ndownloader/files/41962797",
-    "tinnet_OH_alignn": "https://figshare.com/ndownloader/files/41962803",
-    "AGRA_O_alignn": "https://figshare.com/ndownloader/files/41966619",
-    "AGRA_OH_alignn": "https://figshare.com/ndownloader/files/41966610",
-    "AGRA_CHO_alignn": "https://figshare.com/ndownloader/files/41966643",
-    "AGRA_CO_alignn": "https://figshare.com/ndownloader/files/41966634",
-    "AGRA_COOH_alignn": "https://figshare.com/ndownloader/41966646",
-    "qm9_U0_alignn": "https://figshare.com/ndownloader/files/31459054",
-    "qm9_U_alignn": "https://figshare.com/ndownloader/files/31459051",
-    "qm9_alpha_alignn": "https://figshare.com/ndownloader/files/31459027",
-    "qm9_gap_alignn": "https://figshare.com/ndownloader/files/31459036",
-    "qm9_G_alignn": "https://figshare.com/ndownloader/files/31459033",
-    "qm9_HOMO_alignn": "https://figshare.com/ndownloader/files/31459042",
-    "qm9_LUMO_alignn": "https://figshare.com/ndownloader/files/31459045",
-    "qm9_ZPVE_alignn": "https://figshare.com/ndownloader/files/31459057",
-    "hmof_co2_absp_alignn": "https://figshare.com/ndownloader/files/31459198",
-    "hmof_max_co2_adsp_alignn": "https://figshare.com/ndownloader/files/31459207",
-    "hmof_surface_area_m2g_alignn": "https://figshare.com/ndownloader/files/31459222",
-    "hmof_surface_area_m2cm3_alignn": "https://figshare.com/ndownloader/files/31459219",
-    "hmof_pld_alignn": "https://figshare.com/ndownloader/files/31459216",
-    "hmof_lcd_alignn": "https://figshare.com/ndownloader/files/31459201",
-    "hmof_void_fraction_alignn": "https://figshare.com/ndownloader/files/31459228",
-    "ocp2020_all": "https://figshare.com/ndownloader/files/41411025",
-    "ocp2020_100k": "https://figshare.com/ndownloader/files/41967303",
-    "ocp2020_10k": "https://figshare.com/ndownloader/files/41967330",
-    "jv_pdos_alignn": "https://figshare.com/ndownloader/files/36757005"
-}
\ No newline at end of file
+    "jv_formation_energy_peratom_alignn": "https://ndownloader.figshare.com/files/31458679",
+    "jv_optb88vdw_total_energy_alignn": "https://ndownloader.figshare.com/files/31459642",
+    "jv_optb88vdw_bandgap_alignn": "https://ndownloader.figshare.com/files/31459636",
+    "jv_mbj_bandgap_alignn": "https://ndownloader.figshare.com/files/31458694",
+    "jv_spillage_alignn": "https://ndownloader.figshare.com/files/31458736",
+    "jv_slme_alignn": "https://ndownloader.figshare.com/files/31458727",
+    "jv_bulk_modulus_kv_alignn": "https://ndownloader.figshare.com/files/31458649",
+    "jv_shear_modulus_gv_alignn": "https://ndownloader.figshare.com/files/31458724",
+    "jv_n-Seebeck_alignn": "https://ndownloader.figshare.com/files/31458718",
+    "jv_n-powerfact_alignn": "https://ndownloader.figshare.com/files/31458712",
+    "intermat_cbm": "https://ndownloader.figshare.com/files/45392908",
+    "intermat_vbm": "https://ndownloader.figshare.com/files/45392914",
+    "intermat_phi": "https://ndownloader.figshare.com/files/45392911",
+    "jv_magmom_oszicar_alignn": "https://ndownloader.figshare.com/files/31458685",
+    "jv_kpoint_length_unit_alignn": "https://ndownloader.figshare.com/files/31458682",
+    "jv_avg_elec_mass_alignn": "https://ndownloader.figshare.com/files/31458643",
+    "jv_avg_hole_mass_alignn": "https://ndownloader.figshare.com/files/31458646",
+    "jv_epsx_alignn": "https://ndownloader.figshare.com/files/31458667",
+    "jv_mepsx_alignn": "https://ndownloader.figshare.com/files/31458703",
+    "jv_max_efg_alignn": "https://ndownloader.figshare.com/files/31458691",
+    "jv_ehull_alignn": "https://ndownloader.figshare.com/files/31458658",
+    "jv_dfpt_piezo_max_dielectric_alignn": "https://ndownloader.figshare.com/files/31458652",
+    "jv_dfpt_piezo_max_dij_alignn": "https://ndownloader.figshare.com/files/31458655",
+    "jv_exfoliation_energy_alignn": "https://ndownloader.figshare.com/files/31458676",
+    "jv_supercon_tc_alignn": "https://ndownloader.figshare.com/files/38789199",
+    "jv_supercon_edos_alignn": "https://ndownloader.figshare.com/files/39946300",
+    "jv_supercon_debye_alignn": "https://ndownloader.figshare.com/files/39946297",
+    "jv_supercon_a2F_alignn": "https://ndownloader.figshare.com/files/38801886",
+    "mp_e_form_alignn": "https://ndownloader.figshare.com/files/31458811",
+    "mp_gappbe_alignn": "https://ndownloader.figshare.com/files/31458814",
+    "tinnet_O_alignn": "https://ndownloader.figshare.com/files/41962800",
+    "tinnet_N_alignn": "https://ndownloader.figshare.com/files/41962797",
+    "tinnet_OH_alignn": "https://ndownloader.figshare.com/files/41962803",
+    "AGRA_O_alignn": "https://ndownloader.figshare.com/files/41966619",
+    "AGRA_OH_alignn": "https://ndownloader.figshare.com/files/41966610",
+    "AGRA_CHO_alignn": "https://ndownloader.figshare.com/files/41966643",
+    "AGRA_CO_alignn": "https://ndownloader.figshare.com/files/41966634",
+    "AGRA_COOH_alignn": "https://ndownloader.figshare.com/41966646",
+    "qm9_U0_alignn": "https://ndownloader.figshare.com/files/31459054",
+    "qm9_U_alignn": "https://ndownloader.figshare.com/files/31459051",
+    "qm9_alpha_alignn": "https://ndownloader.figshare.com/files/31459027",
+    "qm9_gap_alignn": "https://ndownloader.figshare.com/files/31459036",
+    "qm9_G_alignn": "https://ndownloader.figshare.com/files/31459033",
+    "qm9_HOMO_alignn": "https://ndownloader.figshare.com/files/31459042",
+    "qm9_LUMO_alignn": "https://ndownloader.figshare.com/files/31459045",
+    "qm9_ZPVE_alignn": "https://ndownloader.figshare.com/files/31459057",
+    "hmof_co2_absp_alignn": "https://ndownloader.figshare.com/files/31459198",
+    "hmof_max_co2_adsp_alignn": "https://ndownloader.figshare.com/files/31459207",
+    "hmof_surface_area_m2g_alignn": "https://ndownloader.figshare.com/files/31459222",
+    "hmof_surface_area_m2cm3_alignn": "https://ndownloader.figshare.com/files/31459219",
+    "hmof_pld_alignn": "https://ndownloader.figshare.com/files/31459216",
+    "hmof_lcd_alignn": "https://ndownloader.figshare.com/files/31459201",
+    "hmof_void_fraction_alignn": "https://ndownloader.figshare.com/files/31459228",
+    "ocp2020_all": "https://ndownloader.figshare.com/files/41411025",
+    "ocp2020_100k": "https://ndownloader.figshare.com/files/41967303",
+    "ocp2020_10k": "https://ndownloader.figshare.com/files/41967330",
+    "jv_pdos_alignn": "https://ndownloader.figshare.com/files/36757005"
+}
diff --git a/alignn/ff/all_models_alignn_atomwise.json b/alignn/ff/all_models_alignn_atomwise.json
index 7d42fb8..54a0165 100644
--- a/alignn/ff/all_models_alignn_atomwise.json
+++ b/alignn/ff/all_models_alignn_atomwise.json
@@ -1,28 +1,28 @@
 {
-    "v12.2.2024_dft_3d_307k": "https://figshare.com/ndownloader/files/50904240",
-    "v12.2.2024_mp_1.5mill": "https://figshare.com/ndownloader/files/50904783",
-    "v12.2.2024_mp_187k": "https://figshare.com/ndownloader/files/50904801",
-    "v2024.12.12_dft_3d_multi_prop": "https://figshare.com/ndownloader/files/52025186",
-    "alex_band_gap": "https://figshare.com/ndownloader/files/51993641",
-    "alex_dos_pa": "https://figshare.com/ndownloader/files/51993653",
-    "alex_e_form": "https://figshare.com/ndownloader/files/51993647",
-    "alex_e_hull": "https://figshare.com/ndownloader/files/51993644",
-    "alex_e_total": "https://figshare.com/ndownloader/files/51993650",
-    "alex_mag_per_vol": "https://figshare.com/ndownloader/files/51993659",
-    "alex_vol_pa": "https://figshare.com/ndownloader/files/51993656",
-    "v10.30.2024_dft_3d_307k": "https://figshare.com/ndownloader/files/50634327",
-    "v10.30.2024_mp_168k": "https://figshare.com/ndownloader/files/50634318",
-    "v8.29.2024_dft_3d": "https://figshare.com/ndownloader/files/48889834",
-    "v8.29.2024_mpf": "https://figshare.com/ndownloader/files/48889837",
-    "v5.27.2024": "https://figshare.com/ndownloader/files/47286127",
-    "alignnff_fmult": "https://figshare.com/ndownloader/files/41583585",
-    "alignnff_wt10": "https://figshare.com/ndownloader/files/41583594",
-    "alignnff_fd": "https://figshare.com/ndownloader/files/41583582",
-    "alignnff_wt01": "https://figshare.com/ndownloader/files/41583588",
-    "alignnff_wt1": "https://figshare.com/ndownloader/files/41583591",
-    "fmult_mlearn_only": "https://figshare.com/ndownloader/files/41583597",
-    "aff_Oct23": "https://figshare.com/ndownloader/files/42880573",
-    "revised": "https://figshare.com/ndownloader/files/41583600",
-    "scf_fd_top_10_en_42_fmax_600_wt01": "https://figshare.com/ndownloader/files/41967375",
-    "scf_fd_top_10_en_42_fmax_600_wt10": "https://figshare.com/ndownloader/files/41967372"
+    "v12.2.2024_dft_3d_307k": "https://ndownloader.figshare.com/files/50904240",
+    "v12.2.2024_mp_1.5mill": "https://ndownloader.figshare.com/files/50904783",
+    "v12.2.2024_mp_187k": "https://ndownloader.figshare.com/files/50904801",
+    "v2024.12.12_dft_3d_multi_prop": "https://ndownloader.figshare.com/files/52025186",
+    "alex_band_gap": "https://ndownloader.figshare.com/files/51993641",
+    "alex_dos_pa": "https://ndownloader.figshare.com/files/51993653",
+    "alex_e_form": "https://ndownloader.figshare.com/files/51993647",
+    "alex_e_hull": "https://ndownloader.figshare.com/files/51993644",
+    "alex_e_total": "https://ndownloader.figshare.com/files/51993650",
+    "alex_mag_per_vol": "https://ndownloader.figshare.com/files/51993659",
+    "alex_vol_pa": "https://ndownloader.figshare.com/files/51993656",
+    "v10.30.2024_dft_3d_307k": "https://ndownloader.figshare.com/files/50634327",
+    "v10.30.2024_mp_168k": "https://ndownloader.figshare.com/files/50634318",
+    "v8.29.2024_dft_3d": "https://ndownloader.figshare.com/files/48889834",
+    "v8.29.2024_mpf": "https://ndownloader.figshare.com/files/48889837",
+    "v5.27.2024": "https://ndownloader.figshare.com/files/47286127",
+    "alignnff_fmult": "https://ndownloader.figshare.com/files/41583585",
+    "alignnff_wt10": "https://ndownloader.figshare.com/files/41583594",
+    "alignnff_fd": "https://ndownloader.figshare.com/files/41583582",
+    "alignnff_wt01": "https://ndownloader.figshare.com/files/41583588",
+    "alignnff_wt1": "https://ndownloader.figshare.com/files/41583591",
+    "fmult_mlearn_only": "https://ndownloader.figshare.com/files/41583597",
+    "aff_Oct23": "https://ndownloader.figshare.com/files/42880573",
+    "revised": "https://ndownloader.figshare.com/files/41583600",
+    "scf_fd_top_10_en_42_fmax_600_wt01": "https://ndownloader.figshare.com/files/41967375",
+    "scf_fd_top_10_en_42_fmax_600_wt10": "https://ndownloader.figshare.com/files/41967372"
 }
diff --git a/alignn/pretrained.py b/alignn/pretrained.py
index 9243847..739fd69 100644
--- a/alignn/pretrained.py
+++ b/alignn/pretrained.py
@@ -35,156 +35,156 @@
 # to load as a calculator
 all_models = {
     "jv_formation_energy_peratom_alignn": [
-        "https://figshare.com/ndownloader/files/31458679",
+        "https://ndownloader.figshare.com/files/31458679",
         1,
     ],
     "jv_optb88vdw_total_energy_alignn": [
-        "https://figshare.com/ndownloader/files/31459642",
+        "https://ndownloader.figshare.com/files/31459642",
         1,
     ],
     "jv_optb88vdw_bandgap_alignn": [
-        "https://figshare.com/ndownloader/files/31459636",
+        "https://ndownloader.figshare.com/files/31459636",
         1,
     ],
     "jv_mbj_bandgap_alignn": [
-        "https://figshare.com/ndownloader/files/31458694",
+        "https://ndownloader.figshare.com/files/31458694",
         1,
     ],
     "jv_spillage_alignn": [
-        "https://figshare.com/ndownloader/files/31458736",
+        "https://ndownloader.figshare.com/files/31458736",
         1,
     ],
-    "jv_slme_alignn": ["https://figshare.com/ndownloader/files/31458727", 1],
+    "jv_slme_alignn": ["https://ndownloader.figshare.com/files/31458727", 1],
     "jv_bulk_modulus_kv_alignn": [
-        "https://figshare.com/ndownloader/files/31458649",
+        "https://ndownloader.figshare.com/files/31458649",
         1,
     ],
     "jv_shear_modulus_gv_alignn": [
-        "https://figshare.com/ndownloader/files/31458724",
+        "https://ndownloader.figshare.com/files/31458724",
         1,
     ],
     "jv_n-Seebeck_alignn": [
-        "https://figshare.com/ndownloader/files/31458718",
+        "https://ndownloader.figshare.com/files/31458718",
         1,
     ],
     "jv_n-powerfact_alignn": [
-        "https://figshare.com/ndownloader/files/31458712",
+        "https://ndownloader.figshare.com/files/31458712",
         1,
     ],
     "intermat_cbm": [
-        "https://figshare.com/ndownloader/files/45392908",
+        "https://ndownloader.figshare.com/files/45392908",
         1,
     ],
     "intermat_vbm": [
-        "https://figshare.com/ndownloader/files/45392914",
+        "https://ndownloader.figshare.com/files/45392914",
         1,
     ],
     "intermat_phi": [
-        "https://figshare.com/ndownloader/files/45392911",
+        "https://ndownloader.figshare.com/files/45392911",
         1,
     ],
     "jv_magmom_oszicar_alignn": [
-        "https://figshare.com/ndownloader/files/31458685",
+        "https://ndownloader.figshare.com/files/31458685",
         1,
     ],
     "jv_kpoint_length_unit_alignn": [
-        "https://figshare.com/ndownloader/files/31458682",
+        "https://ndownloader.figshare.com/files/31458682",
         1,
     ],
     "jv_avg_elec_mass_alignn": [
-        "https://figshare.com/ndownloader/files/31458643",
+        "https://ndownloader.figshare.com/files/31458643",
         1,
     ],
     "jv_avg_hole_mass_alignn": [
-        "https://figshare.com/ndownloader/files/31458646",
+        "https://ndownloader.figshare.com/files/31458646",
         1,
     ],
-    "jv_epsx_alignn": ["https://figshare.com/ndownloader/files/31458667", 1],
-    "jv_mepsx_alignn": ["https://figshare.com/ndownloader/files/31458703", 1],
+    "jv_epsx_alignn": ["https://ndownloader.figshare.com/files/31458667", 1],
+    "jv_mepsx_alignn": ["https://ndownloader.figshare.com/files/31458703", 1],
     "jv_max_efg_alignn": [
-        "https://figshare.com/ndownloader/files/31458691",
+        "https://ndownloader.figshare.com/files/31458691",
         1,
     ],
-    "jv_ehull_alignn": ["https://figshare.com/ndownloader/files/31458658", 1],
+    "jv_ehull_alignn": ["https://ndownloader.figshare.com/files/31458658", 1],
     "jv_dfpt_piezo_max_dielectric_alignn": [
-        "https://figshare.com/ndownloader/files/31458652",
+        "https://ndownloader.figshare.com/files/31458652",
         1,
     ],
     "jv_dfpt_piezo_max_dij_alignn": [
-        "https://figshare.com/ndownloader/files/31458655",
+        "https://ndownloader.figshare.com/files/31458655",
         1,
     ],
     "jv_exfoliation_energy_alignn": [
-        "https://figshare.com/ndownloader/files/31458676",
+        "https://ndownloader.figshare.com/files/31458676",
         1,
     ],
     "jv_supercon_tc_alignn": [
-        "https://figshare.com/ndownloader/files/38789199",
+        "https://ndownloader.figshare.com/files/38789199",
         1,
     ],
     "jv_supercon_edos_alignn": [
-        "https://figshare.com/ndownloader/files/39946300",
+        "https://ndownloader.figshare.com/files/39946300",
         1,
     ],
     "jv_supercon_debye_alignn": [
-        "https://figshare.com/ndownloader/files/39946297",
+        "https://ndownloader.figshare.com/files/39946297",
         1,
     ],
     "jv_supercon_a2F_alignn": [
-        "https://figshare.com/ndownloader/files/38801886",
+        "https://ndownloader.figshare.com/files/38801886",
         100,
     ],
     "mp_e_form_alignn": [
-        "https://figshare.com/ndownloader/files/31458811",
+        "https://ndownloader.figshare.com/files/31458811",
         1,
     ],
     "mp_gappbe_alignn": [
-        "https://figshare.com/ndownloader/files/31458814",
+        "https://ndownloader.figshare.com/files/31458814",
         1,
     ],
-    "tinnet_O_alignn": ["https://figshare.com/ndownloader/files/41962800", 1],
-    "tinnet_N_alignn": ["https://figshare.com/ndownloader/files/41962797", 1],
-    "tinnet_OH_alignn": ["https://figshare.com/ndownloader/files/41962803", 1],
-    "AGRA_O_alignn": ["https://figshare.com/ndownloader/files/41966619", 1],
-    "AGRA_OH_alignn": ["https://figshare.com/ndownloader/files/41966610", 1],
-    "AGRA_CHO_alignn": ["https://figshare.com/ndownloader/files/41966643", 1],
-    "AGRA_CO_alignn": ["https://figshare.com/ndownloader/files/41966634", 1],
-    "AGRA_COOH_alignn": ["https://figshare.com/ndownloader/41966646", 1],
-    "qm9_U0_alignn": ["https://figshare.com/ndownloader/files/31459054", 1],
-    "qm9_U_alignn": ["https://figshare.com/ndownloader/files/31459051", 1],
-    "qm9_alpha_alignn": ["https://figshare.com/ndownloader/files/31459027", 1],
-    "qm9_gap_alignn": ["https://figshare.com/ndownloader/files/31459036", 1],
-    "qm9_G_alignn": ["https://figshare.com/ndownloader/files/31459033", 1],
-    "qm9_HOMO_alignn": ["https://figshare.com/ndownloader/files/31459042", 1],
-    "qm9_LUMO_alignn": ["https://figshare.com/ndownloader/files/31459045", 1],
-    "qm9_ZPVE_alignn": ["https://figshare.com/ndownloader/files/31459057", 1],
+    "tinnet_O_alignn": ["https://ndownloader.figshare.com/files/41962800", 1],
+    "tinnet_N_alignn": ["https://ndownloader.figshare.com/files/41962797", 1],
+    "tinnet_OH_alignn": ["https://ndownloader.figshare.com/files/41962803", 1],
+    "AGRA_O_alignn": ["https://ndownloader.figshare.com/files/41966619", 1],
+    "AGRA_OH_alignn": ["https://ndownloader.figshare.com/files/41966610", 1],
+    "AGRA_CHO_alignn": ["https://ndownloader.figshare.com/files/41966643", 1],
+    "AGRA_CO_alignn": ["https://ndownloader.figshare.com/files/41966634", 1],
+    "AGRA_COOH_alignn": ["https://ndownloader.figshare.com/41966646", 1],
+    "qm9_U0_alignn": ["https://ndownloader.figshare.com/files/31459054", 1],
+    "qm9_U_alignn": ["https://ndownloader.figshare.com/files/31459051", 1],
+    "qm9_alpha_alignn": ["https://ndownloader.figshare.com/files/31459027", 1],
+    "qm9_gap_alignn": ["https://ndownloader.figshare.com/files/31459036", 1],
+    "qm9_G_alignn": ["https://ndownloader.figshare.com/files/31459033", 1],
+    "qm9_HOMO_alignn": ["https://ndownloader.figshare.com/files/31459042", 1],
+    "qm9_LUMO_alignn": ["https://ndownloader.figshare.com/files/31459045", 1],
+    "qm9_ZPVE_alignn": ["https://ndownloader.figshare.com/files/31459057", 1],
     "hmof_co2_absp_alignn": [
-        "https://figshare.com/ndownloader/files/31459198",
+        "https://ndownloader.figshare.com/files/31459198",
         5,
     ],
     "hmof_max_co2_adsp_alignn": [
-        "https://figshare.com/ndownloader/files/31459207",
+        "https://ndownloader.figshare.com/files/31459207",
         1,
     ],
     "hmof_surface_area_m2g_alignn": [
-        "https://figshare.com/ndownloader/files/31459222",
+        "https://ndownloader.figshare.com/files/31459222",
         1,
     ],
     "hmof_surface_area_m2cm3_alignn": [
-        "https://figshare.com/ndownloader/files/31459219",
+        "https://ndownloader.figshare.com/files/31459219",
         1,
     ],
-    "hmof_pld_alignn": ["https://figshare.com/ndownloader/files/31459216", 1],
-    "hmof_lcd_alignn": ["https://figshare.com/ndownloader/files/31459201", 1],
+    "hmof_pld_alignn": ["https://ndownloader.figshare.com/files/31459216", 1],
+    "hmof_lcd_alignn": ["https://ndownloader.figshare.com/files/31459201", 1],
     "hmof_void_fraction_alignn": [
-        "https://figshare.com/ndownloader/files/31459228",
+        "https://ndownloader.figshare.com/files/31459228",
         1,
     ],
-    "ocp2020_all": ["https://figshare.com/ndownloader/files/41411025", 1],
-    "ocp2020_100k": ["https://figshare.com/ndownloader/files/41967303", 1],
-    "ocp2020_10k": ["https://figshare.com/ndownloader/files/41967330", 1],
+    "ocp2020_all": ["https://ndownloader.figshare.com/files/41411025", 1],
+    "ocp2020_100k": ["https://ndownloader.figshare.com/files/41967303", 1],
+    "ocp2020_10k": ["https://ndownloader.figshare.com/files/41967330", 1],
     "jv_pdos_alignn": [
-        "https://figshare.com/ndownloader/files/36757005",
+        "https://ndownloader.figshare.com/files/36757005",
         66,
         {"alignn_layers": 6, "gcn_layers": 6},
     ],

From 7d210f55113a231a221d0b70df95cb99e1ea2451 Mon Sep 17 00:00:00 2001
From: Kamal Choudhary <writetokamal.1989@gmail.com>
Date: Fri, 13 Feb 2026 07:51:44 -0500
Subject: [PATCH 08/23] Update ff.py

---
 alignn/ff/ff.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/alignn/ff/ff.py b/alignn/ff/ff.py
index d18eef2..37daa69 100644
--- a/alignn/ff/ff.py
+++ b/alignn/ff/ff.py
@@ -16,7 +16,8 @@
 from ase.optimize.gpmin.gpmin import GPMin
 from ase.optimize.lbfgs import LBFGS, LBFGSLineSearch
 from ase.optimize.mdmin import MDMin
-from ase.constraints import ExpCellFilter
+# from ase.constraints import ExpCellFilter
+from ase.filters import ExpCellFilter
 from ase.eos import EquationOfState
 from ase.units import kJ
 from ase.optimize.sciopt import SciPyFminBFGS, SciPyFminCG

From d54afd2b39cf11a578031e32da6142c5aeb79638 Mon Sep 17 00:00:00 2001
From: user <writetokamal.1989@gmail.com>
Date: Sun, 22 Feb 2026 21:35:47 -0500
Subject: [PATCH 09/23] nal sal

---
 alignn/models/alignn_atomwise.py | 134 ++++++++++++++++++++++++++++++-
 1 file changed, 132 insertions(+), 2 deletions(-)

diff --git a/alignn/models/alignn_atomwise.py b/alignn/models/alignn_atomwise.py
index bbee7c1..92de575 100644
--- a/alignn/models/alignn_atomwise.py
+++ b/alignn/models/alignn_atomwise.py
@@ -10,6 +10,7 @@
 import numpy as np
 from dgl.nn import AvgPooling
 import torch
+from dgl.nn.functional import edge_softmax
 
 # from dgl.nn.functional import edge_softmax
 from typing import Literal
@@ -70,6 +71,10 @@ class ALIGNNAtomWiseConfig(BaseSettings):
     penalty_threshold: float = 1
     additional_output_features: int = 0
     additional_output_weight: float = 0
+    # 🆕 attention controls
+    use_nal: bool = True
+    use_sal: bool = True
+    num_attention_heads: int = 8
 
     class Config:
         """Configure model settings behavior."""
@@ -77,6 +82,99 @@ class Config:
         env_prefix = "jv_model"
 
 
+class NodeAttentionLayer(nn.Module):
+    """
+    N-ALIGNN: Local graph attention over neighbors.
+    Lightweight GAT-style but ALIGNN-compatible.
+    """
+
+    def __init__(self, hidden_dim, num_heads=4):
+        super().__init__()
+        self.hidden_dim = hidden_dim
+        self.num_heads = num_heads
+        self.head_dim = hidden_dim // num_heads
+        assert (
+            hidden_dim % num_heads == 0
+        ), "hidden_dim must be divisible by num_heads"
+        self.q_proj = nn.Linear(hidden_dim, hidden_dim)
+        self.k_proj = nn.Linear(hidden_dim, hidden_dim)
+        self.v_proj = nn.Linear(hidden_dim, hidden_dim)
+
+        self.out_proj = nn.Linear(hidden_dim, hidden_dim)
+        self.norm = nn.LayerNorm(hidden_dim)
+
+    def forward(self, g, x):
+        with g.local_scope():
+
+            Q = self.q_proj(x).view(-1, self.num_heads, self.head_dim)
+            K = self.k_proj(x).view(-1, self.num_heads, self.head_dim)
+            V = self.v_proj(x).view(-1, self.num_heads, self.head_dim)
+
+            g.ndata["Q"] = Q
+            g.ndata["K"] = K
+            g.ndata["V"] = V
+
+            g.apply_edges(fn.u_dot_v("Q", "K", "score"))
+            e = g.edata["score"] / np.sqrt(self.head_dim)
+
+            alpha = edge_softmax(g, e)
+            g.edata["alpha"] = alpha
+
+            g.update_all(
+                fn.u_mul_e("V", "alpha", "m"),
+                fn.sum("m", "h_attn"),
+            )
+
+            h = g.ndata["h_attn"].reshape(-1, self.hidden_dim)
+            h = self.out_proj(h)
+
+            return self.norm(x + h)
+
+
+class SelfAttentionLayer(nn.Module):
+    """
+    T-ALIGNN: Global transformer attention.
+    Works after graph convolutions.
+    """
+
+    def __init__(self, hidden_dim, num_heads=4):
+        super().__init__()
+        self.mha = nn.MultiheadAttention(
+            embed_dim=hidden_dim,
+            num_heads=num_heads,
+            batch_first=True,
+        )
+        self.norm1 = nn.LayerNorm(hidden_dim)
+        self.ffn = nn.Sequential(
+            nn.Linear(hidden_dim, hidden_dim * 4),
+            nn.SiLU(),
+            nn.Linear(hidden_dim * 4, hidden_dim),
+        )
+        self.norm2 = nn.LayerNorm(hidden_dim)
+
+    def forward(self, x, batch_nodes):
+        """
+        x: (N_total, hidden)
+        batch_nodes: list of node counts per graph
+        """
+        outputs = []
+        start = 0
+
+        for n in batch_nodes:
+            h = x[start : start + n].unsqueeze(0)  # (1,n,d)
+
+            attn_out, _ = self.mha(h, h, h)
+            h = self.norm1(h + attn_out)
+
+            ff = self.ffn(h)
+            h = self.norm2(h + ff)
+
+            outputs.append(h.squeeze(0))
+            start += n
+
+        return torch.cat(outputs, dim=0)
+
+
 def cutoff_function_based_edges_old(r, inner_cutoff=4):
     """Apply smooth cutoff to pairwise interactions
 
@@ -309,6 +407,29 @@ def __init__(
             ]
         )
 
+        # 🆕 N-ALIGNN
+        if config.use_nal:
+            self.nal_layers = nn.ModuleList(
+                [
+                    NodeAttentionLayer(
+                        config.hidden_features,
+                        config.num_attention_heads,
+                    )
+                    for _ in range(config.alignn_layers)
+                ]
+            )
+        else:
+            self.nal_layers = None
+
+        # 🆕 T-ALIGNN
+        if config.use_sal:
+            self.sal = SelfAttentionLayer(
+                config.hidden_features,
+                config.num_attention_heads,
+            )
+        else:
+            self.sal = None
+
         self.readout = AvgPooling()
 
         if config.extra_features != 0:
@@ -453,12 +574,21 @@ def forward(
             y = self.edge_embedding(bondlength)
         # y = self.edge_embedding(bondlength)
         # ALIGNN updates: update node, edge, triplet features
-        for alignn_layer in self.alignn_layers:
+
+        for i, alignn_layer in enumerate(self.alignn_layers):
             x, y, z = alignn_layer(g, lg, x, y, z)
 
-        # gated GCN updates: update node, edge features
+            # 🆕 N-ALIGNN
+            if self.nal_layers is not None:
+                x = self.nal_layers[i](g, x)
+
         for gcn_layer in self.gcn_layers:
             x, y = gcn_layer(g, x, y)
+        # 🆕 T-ALIGNN global attention
+        if self.sal is not None:
+            batch_nodes = g.batch_num_nodes()  # .tolist()
+            x = self.sal(x, batch_nodes)
+
         # norm-activation-pool-classify
         out = torch.empty(1)
         additional_out = torch.empty(1)

From 7f99fa1973231f6749ba2c9cb5cd8c1ebebd299a Mon Sep 17 00:00:00 2001
From: knc6 <writetokamal.1989@gmail.com>
Date: Sat, 14 Mar 2026 21:22:54 -0400
Subject: [PATCH 10/23] LogSoftmax and pretrained

---
 alignn/models/alignn_atomwise.py |  4 ++--
 alignn/pretrained.py             | 33 +++++++++++++++++++++++---------
 2 files changed, 26 insertions(+), 11 deletions(-)

diff --git a/alignn/models/alignn_atomwise.py b/alignn/models/alignn_atomwise.py
index bbee7c1..93b42b0 100644
--- a/alignn/models/alignn_atomwise.py
+++ b/alignn/models/alignn_atomwise.py
@@ -344,8 +344,8 @@ def __init__(
             )
         if self.classification:
             self.fc = nn.Linear(config.hidden_features, 1)
-            self.softmax = nn.Sigmoid()
-            # self.softmax = nn.LogSoftmax(dim=1)
+            # self.softmax = nn.Sigmoid()
+            self.softmax = nn.LogSoftmax(dim=1)
         else:
             self.fc = nn.Linear(config.hidden_features, config.output_features)
         self.link = None
diff --git a/alignn/pretrained.py b/alignn/pretrained.py
index 739fd69..947979c 100644
--- a/alignn/pretrained.py
+++ b/alignn/pretrained.py
@@ -312,7 +312,25 @@ def get_prediction(
     max_neighbors=12,
 ):
     """Get model prediction on a single structure."""
-    model = get_figshare_model(model_name)
+    if os.path.isdir(model_name):
+
+        from alignn.models.alignn_atomwise import (
+            ALIGNNAtomWise,
+            ALIGNNAtomWiseConfig,
+        )
+        import torch
+        from jarvis.db.jsonutils import loadjson
+
+        config = loadjson(os.path.join(model_name, "config.json"))
+        model_path = os.path.join(model_name, "best_model.pt")
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        tmp = ALIGNNAtomWiseConfig(**config["model"])
+        model = ALIGNNAtomWise(tmp)
+        model.load_state_dict(torch.load(model_path, map_location=device))
+        model = model.to(device)
+        print(model)
+    else:
+        model = get_figshare_model(model_name)
     # print("Loading completed.")
     g, lg = Graph.atom_dgl_multigraph(
         atoms,
@@ -320,14 +338,11 @@ def get_prediction(
         max_neighbors=max_neighbors,
     )
     lat = torch.tensor(atoms.lattice_mat)
-    out_data = (
-        model([g.to(device), lg.to(device), lat.to(device)])
-        .detach()
-        .cpu()
-        .numpy()
-        .flatten()
-        .tolist()
-    )
+    out_data = model([g.to(device), lg.to(device), lat.to(device)])
+    if isinstance(out_data, dict):
+        out_data = out_data["out"]
+    print("out_data", out_data)
+    out_data = out_data.detach().cpu().numpy().flatten().tolist()
     return out_data
 
 

From 714b566ecddbbea79b4fa1ccf6cc12edb17e712e Mon Sep 17 00:00:00 2001
From: user <writetokamal.1989@gmail.com>
Date: Sat, 14 Mar 2026 22:04:59 -0400
Subject: [PATCH 11/23] new pretrained

---
 alignn/models/alignn_atomwise.py | 134 +------------------------------
 alignn/pretrained.py             |  57 ++++++++-----
 2 files changed, 41 insertions(+), 150 deletions(-)

diff --git a/alignn/models/alignn_atomwise.py b/alignn/models/alignn_atomwise.py
index a080b7c..93b42b0 100644
--- a/alignn/models/alignn_atomwise.py
+++ b/alignn/models/alignn_atomwise.py
@@ -10,7 +10,6 @@
 import numpy as np
 from dgl.nn import AvgPooling
 import torch
-from dgl.nn.functional import edge_softmax
 
 # from dgl.nn.functional import edge_softmax
 from typing import Literal
@@ -71,10 +70,6 @@ class ALIGNNAtomWiseConfig(BaseSettings):
     penalty_threshold: float = 1
     additional_output_features: int = 0
     additional_output_weight: float = 0
-    # 🆕 attention controls
-    use_nal: bool = True
-    use_sal: bool = True
-    num_attention_heads: int = 8
 
     class Config:
         """Configure model settings behavior."""
@@ -82,99 +77,6 @@ class Config:
         env_prefix = "jv_model"
 
 
-class NodeAttentionLayer(nn.Module):
-    """
-    N-ALIGNN: Local graph attention over neighbors.
-    Lightweight GAT-style but ALIGNN-compatible.
-    """
-
-    def __init__(self, hidden_dim, num_heads=4):
-        super().__init__()
-        self.hidden_dim = hidden_dim
-        self.num_heads = num_heads
-        self.head_dim = hidden_dim // num_heads
-        assert (
-            hidden_dim % num_heads == 0
-        ), "hidden_dim must be divisible by num_heads"
-        self.q_proj = nn.Linear(hidden_dim, hidden_dim)
-        self.k_proj = nn.Linear(hidden_dim, hidden_dim)
-        self.v_proj = nn.Linear(hidden_dim, hidden_dim)
-
-        self.out_proj = nn.Linear(hidden_dim, hidden_dim)
-        self.norm = nn.LayerNorm(hidden_dim)
-
-    def forward(self, g, x):
-        with g.local_scope():
-
-            Q = self.q_proj(x).view(-1, self.num_heads, self.head_dim)
-            K = self.k_proj(x).view(-1, self.num_heads, self.head_dim)
-            V = self.v_proj(x).view(-1, self.num_heads, self.head_dim)
-
-            g.ndata["Q"] = Q
-            g.ndata["K"] = K
-            g.ndata["V"] = V
-
-            g.apply_edges(fn.u_dot_v("Q", "K", "score"))
-            e = g.edata["score"] / np.sqrt(self.head_dim)
-
-            alpha = edge_softmax(g, e)
-            g.edata["alpha"] = alpha
-
-            g.update_all(
-                fn.u_mul_e("V", "alpha", "m"),
-                fn.sum("m", "h_attn"),
-            )
-
-            h = g.ndata["h_attn"].reshape(-1, self.hidden_dim)
-            h = self.out_proj(h)
-
-            return self.norm(x + h)
-
-
-class SelfAttentionLayer(nn.Module):
-    """
-    T-ALIGNN: Global transformer attention.
-    Works after graph convolutions.
-    """
-
-    def __init__(self, hidden_dim, num_heads=4):
-        super().__init__()
-        self.mha = nn.MultiheadAttention(
-            embed_dim=hidden_dim,
-            num_heads=num_heads,
-            batch_first=True,
-        )
-        self.norm1 = nn.LayerNorm(hidden_dim)
-        self.ffn = nn.Sequential(
-            nn.Linear(hidden_dim, hidden_dim * 4),
-            nn.SiLU(),
-            nn.Linear(hidden_dim * 4, hidden_dim),
-        )
-        self.norm2 = nn.LayerNorm(hidden_dim)
-
-    def forward(self, x, batch_nodes):
-        """
-        x: (N_total, hidden)
-        batch_nodes: list of node counts per graph
-        """
-        outputs = []
-        start = 0
-
-        for n in batch_nodes:
-            h = x[start : start + n].unsqueeze(0)  # (1,n,d)
-
-            attn_out, _ = self.mha(h, h, h)
-            h = self.norm1(h + attn_out)
-
-            ff = self.ffn(h)
-            h = self.norm2(h + ff)
-
-            outputs.append(h.squeeze(0))
-            start += n
-
-        return torch.cat(outputs, dim=0)
-
-
 def cutoff_function_based_edges_old(r, inner_cutoff=4):
     """Apply smooth cutoff to pairwise interactions
 
@@ -407,29 +309,6 @@ def __init__(
             ]
         )
 
-        # 🆕 N-ALIGNN
-        if config.use_nal:
-            self.nal_layers = nn.ModuleList(
-                [
-                    NodeAttentionLayer(
-                        config.hidden_features,
-                        config.num_attention_heads,
-                    )
-                    for _ in range(config.alignn_layers)
-                ]
-            )
-        else:
-            self.nal_layers = None
-
-        # 🆕 T-ALIGNN
-        if config.use_sal:
-            self.sal = SelfAttentionLayer(
-                config.hidden_features,
-                config.num_attention_heads,
-            )
-        else:
-            self.sal = None
-
         self.readout = AvgPooling()
 
         if config.extra_features != 0:
@@ -574,21 +453,12 @@ def forward(
             y = self.edge_embedding(bondlength)
         # y = self.edge_embedding(bondlength)
         # ALIGNN updates: update node, edge, triplet features
-
-        for i, alignn_layer in enumerate(self.alignn_layers):
+        for alignn_layer in self.alignn_layers:
             x, y, z = alignn_layer(g, lg, x, y, z)
 
-            # 🆕 N-ALIGNN
-            if self.nal_layers is not None:
-                x = self.nal_layers[i](g, x)
-
+        # gated GCN updates: update node, edge features
         for gcn_layer in self.gcn_layers:
             x, y = gcn_layer(g, x, y)
-        # 🆕 T-ALIGNN global attention
-        if self.sal is not None:
-            batch_nodes = g.batch_num_nodes()  # .tolist()
-            x = self.sal(x, batch_nodes)
-
         # norm-activation-pool-classify
         out = torch.empty(1)
         additional_out = torch.empty(1)
diff --git a/alignn/pretrained.py b/alignn/pretrained.py
index 947979c..e523552 100644
--- a/alignn/pretrained.py
+++ b/alignn/pretrained.py
@@ -14,10 +14,14 @@
 import argparse
 from jarvis.core.atoms import Atoms
 from alignn.graphs import Graph
-from jarvis.db.jsonutils import dumpjson
+from jarvis.db.jsonutils import loadjson, dumpjson
 import pandas as pd
 from alignn.dataset import get_torch_dataset
 import numpy as np
+from alignn.models.alignn_atomwise import (
+    ALIGNNAtomWise,
+    ALIGNNAtomWiseConfig,
+)
 
 # from jarvis.core.graphs import Graph
 
@@ -188,6 +192,10 @@
         66,
         {"alignn_layers": 6, "gcn_layers": 6},
     ],
+    "jv_raman_alignn": [
+        "https://ndownloader.figshare.com/files/62805487",
+        200,
+    ],
 }
 
 
@@ -283,6 +291,7 @@ def get_figshare_model(model_name="jv_formation_energy_peratom_alignn"):
     print("Path", os.path.abspath(path))
     print("Config", os.path.abspath(cfg))
     config = json.loads(zipfile.ZipFile(path).read(cfg))
+    print("config", config, type(config))
     # print("Loading the zipfile...", zipfile.ZipFile(path).namelist())
     data = zipfile.ZipFile(path).read(tmp)
     # model = ALIGNN(
@@ -290,18 +299,34 @@ def get_figshare_model(model_name="jv_formation_energy_peratom_alignn"):
     #        name="alignn", output_features=output_features, **config_params
     #    )
     # )
-    model = ALIGNN(ALIGNNConfig(**config["model"]))
+    print(config, type(config))
+    if config["model"]["name"] == "alignn":
+        model = ALIGNN(ALIGNNConfig(**config["model"]))
+        new_file, filename = tempfile.mkstemp()
+        with open(filename, "wb") as f:
+            f.write(data)
+        model.load_state_dict(
+            torch.load(filename, map_location=device, weights_only=False)[
+                "model"
+            ]
+        )
+        model.to(device)
+        model.eval()
+        if os.path.exists(filename):
+            os.remove(filename)
+    if config["model"]["name"] == "alignn_atomwise":
+        model = ALIGNNAtomWise(ALIGNNAtomWiseConfig(**config["model"]))
+        new_file, filename = tempfile.mkstemp()
+        with open(filename, "wb") as f:
+            f.write(data)
+        model.load_state_dict(
+            torch.load(filename, map_location=device, weights_only=False)
+        )
+        model.to(device)
+        model.eval()
+        if os.path.exists(filename):
+            os.remove(filename)
 
-    new_file, filename = tempfile.mkstemp()
-    with open(filename, "wb") as f:
-        f.write(data)
-    model.load_state_dict(
-        torch.load(filename, map_location=device, weights_only=False)["model"]
-    )
-    model.to(device)
-    model.eval()
-    if os.path.exists(filename):
-        os.remove(filename)
     return model
 
 
@@ -314,12 +339,8 @@ def get_prediction(
     """Get model prediction on a single structure."""
     if os.path.isdir(model_name):
 
-        from alignn.models.alignn_atomwise import (
-            ALIGNNAtomWise,
-            ALIGNNAtomWiseConfig,
-        )
-        import torch
-        from jarvis.db.jsonutils import loadjson
+        # import torch
+        # from jarvis.db.jsonutils import loadjson
 
         config = loadjson(os.path.join(model_name, "config.json"))
         model_path = os.path.join(model_name, "best_model.pt")

From 643874415e9252066f44e5fc7fc38f7c61d13aeb Mon Sep 17 00:00:00 2001
From: knc6 <writetokamal.1989@gmail.com>
Date: Fri, 3 Apr 2026 07:10:57 -0400
Subject: [PATCH 12/23] Raman calc update

---
 alignn/ff/calculators.py | 8 ++++----
 alignn/pretrained.py     | 3 ++-
 2 files changed, 6 insertions(+), 5 deletions(-)

diff --git a/alignn/ff/calculators.py b/alignn/ff/calculators.py
index c96bc58..a47d5fe 100644
--- a/alignn/ff/calculators.py
+++ b/alignn/ff/calculators.py
@@ -291,10 +291,10 @@ def calculate(self, atoms, properties=None, system_changes=None):
             atom_features=self.config["atom_features"],
             use_canonize=self.config["use_canonize"],
         )
-        print("self.devicee", self.device)
-        print("g", g.device)
-        print("lg", lg.device)
-        print("model", self.model)
+        # print("self.devicee", self.device)
+        # print("g", g.device)
+        # print("lg", lg.device)
+        # print("model", self.model)
 
         if self.config["model"]["alignn_layers"] > 0:
             result = self.model(
diff --git a/alignn/pretrained.py b/alignn/pretrained.py
index e523552..0cbd275 100644
--- a/alignn/pretrained.py
+++ b/alignn/pretrained.py
@@ -193,7 +193,8 @@
         {"alignn_layers": 6, "gcn_layers": 6},
     ],
     "jv_raman_alignn": [
-        "https://ndownloader.figshare.com/files/62805487",
+        "https://figshare.com/ndownloader/files/63429102",
+        # "https://ndownloader.figshare.com/files/62805487",
         200,
     ],
 }

From 855f4a5b59a56fbd933f7b274b54cb1c9a45d4bb Mon Sep 17 00:00:00 2001
From: knc6 <writetokamal.1989@gmail.com>
Date: Fri, 3 Apr 2026 07:45:52 -0400
Subject: [PATCH 13/23] Old raman4x4

---
 alignn/pretrained.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/alignn/pretrained.py b/alignn/pretrained.py
index 0cbd275..e523552 100644
--- a/alignn/pretrained.py
+++ b/alignn/pretrained.py
@@ -193,8 +193,7 @@
         {"alignn_layers": 6, "gcn_layers": 6},
     ],
     "jv_raman_alignn": [
-        "https://figshare.com/ndownloader/files/63429102",
-        # "https://ndownloader.figshare.com/files/62805487",
+        "https://ndownloader.figshare.com/files/62805487",
         200,
     ],
 }

From 3160e72eb2ff088efd2613aeee6b31a09d5133a7 Mon Sep 17 00:00:00 2001
From: user <writetokamal.1989@gmail.com>
Date: Sun, 5 Apr 2026 08:58:32 -0400
Subject: [PATCH 14/23] Alignn version update

---
 alignn/__init__.py |  2 +-
 setup.py           | 11 ++++++-----
 2 files changed, 7 insertions(+), 6 deletions(-)

diff --git a/alignn/__init__.py b/alignn/__init__.py
index cdc30cc..e396d61 100644
--- a/alignn/__init__.py
+++ b/alignn/__init__.py
@@ -1,3 +1,3 @@
 """Version number."""
 
-__version__ = "2025.11.1"
+__version__ = "2025.12.1"
diff --git a/setup.py b/setup.py
index 0e08d66..8acfb73 100644
--- a/setup.py
+++ b/setup.py
@@ -1,6 +1,6 @@
 """ALIGNN: Atomistic LIne Graph Neural Network.
 
-https://jarvis.nist.gov.
+https://atomgpt.org/alignn
 """
 
 import setuptools
@@ -10,19 +10,20 @@
 
 setuptools.setup(
     name="alignn",
-    version="2025.11.1",
+    version="2025.12.1",
     author="Kamal Choudhary, Brian DeCost",
     author_email="kamal.choudhary@nist.gov",
     description="alignn",
     install_requires=[
-        "numpy<2.0",
+        # "numpy<2.0",
         # "numpy>=1.19.5,<2.0.0",
+        "numpy>=1.19.5",
         "scipy>=1.6.1",
         "jarvis-tools>=2021.07.19",
         "torch>=2.2.1",
         # "torch<=2.2.1",
         "mpmath<=1.3.0",
-        "dgl>=0.6.0",
+        # "dgl>=0.6.0",
         "spglib>=2.0.2",
         "scikit-learn>=0.22.2",
         "matplotlib>=3.4.1",
@@ -45,7 +46,7 @@
     ],
     long_description=long_description,
     long_description_content_type="text/markdown",
-    url="https://github.com/usnistgov/alignn",
+    url="https://github.com/atomgptlab/alignn",
     packages=setuptools.find_packages(),
     package_data={
         "alignn.ff": [

From b4c5ddc3c65b38fbd070dba171474dd1c982901f Mon Sep 17 00:00:00 2001
From: user <writetokamal.1989@gmail.com>
Date: Sun, 5 Apr 2026 09:17:03 -0400
Subject: [PATCH 15/23] ase filters

---
 environment.yml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/environment.yml b/environment.yml
index 886addb..115f917 100644
--- a/environment.yml
+++ b/environment.yml
@@ -229,7 +229,7 @@ dependencies:
   - pip:
       - annotated-types==0.6.0
       - appdirs==1.4.4
-      - ase==3.22.1
+      - ase
       - autograd==1.6.2
       - bokeh==2.4.2
       - cif2cell==2.1.0

From 23254f3516ebd4aeae5e145ce2ce15c9b20e1ef0 Mon Sep 17 00:00:00 2001
From: user <writetokamal.1989@gmail.com>
Date: Sun, 5 Apr 2026 09:37:20 -0400
Subject: [PATCH 16/23] bug fix

---
 alignn/ff/ff.py                  | 13 +++++++++----
 alignn/models/alignn_atomwise.py |  2 ++
 alignn/pretrained.py             |  2 +-
 3 files changed, 12 insertions(+), 5 deletions(-)

diff --git a/alignn/ff/ff.py b/alignn/ff/ff.py
index 37daa69..846a515 100644
--- a/alignn/ff/ff.py
+++ b/alignn/ff/ff.py
@@ -16,6 +16,7 @@
 from ase.optimize.gpmin.gpmin import GPMin
 from ase.optimize.lbfgs import LBFGS, LBFGSLineSearch
 from ase.optimize.mdmin import MDMin
+
 # from ase.constraints import ExpCellFilter
 from ase.filters import ExpCellFilter
 from ase.eos import EquationOfState
@@ -462,7 +463,8 @@ def run_nvt_langevin(
             self.timestep,
             temperature_K=temperature_K,
             friction=friction,
-            communicator=self.communicator,
+            comm=self.communicator,
+            # communicator=self.communicator,
         )
         # Create monitors for logfile and a trajectory file
         # logfile = os.path.join(".", "%s.log" % filename)
@@ -495,7 +497,8 @@ def run_nvt_andersen(
             self.timestep,
             temperature_K=temperature_K,
             andersen_prob=andersen_prob,
-            communicator=self.communicator,
+            comm=self.communicator,
+            # communicator=self.communicator,
         )
         # Create monitors for logfile and a trajectory file
         # logfile = os.path.join(".", "%s.log" % filename)
@@ -530,7 +533,8 @@ def run_nvt_berendsen(
             self.timestep,
             temperature_K=temperature_K,
             taut=taut,
-            communicator=self.communicator,
+            comm=self.communicator,
+            # communicator=self.communicator,
         )
         # Create monitors for logfile and a trajectory file
         # logfile = os.path.join(".", "%s.log" % filename)
@@ -569,7 +573,8 @@ def run_npt_berendsen(
             taup=taup,
             pressure=pressure,
             compressibility=compressibility,
-            communicator=self.communicator,
+            comm=self.communicator,
+            # communicator=self.communicator,
         )
         # Create monitors for logfile and a trajectory file
         # logfile = os.path.join(".", "%s.log" % filename)
diff --git a/alignn/models/alignn_atomwise.py b/alignn/models/alignn_atomwise.py
index 93b42b0..40a3dd4 100644
--- a/alignn/models/alignn_atomwise.py
+++ b/alignn/models/alignn_atomwise.py
@@ -650,6 +650,8 @@ def forward(
 
         if self.classification:
             # out = torch.max(out,dim=1)
+            if out.dim() == 1:
+                out = out.unsqueeze(0)
             out = self.softmax(out)
         result["out"] = out
         result["additional"] = additional_out
diff --git a/alignn/pretrained.py b/alignn/pretrained.py
index e523552..7bc8fe6 100644
--- a/alignn/pretrained.py
+++ b/alignn/pretrained.py
@@ -337,6 +337,7 @@ def get_prediction(
     max_neighbors=12,
 ):
     """Get model prediction on a single structure."""
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
     if os.path.isdir(model_name):
 
         # import torch
@@ -344,7 +345,6 @@ def get_prediction(
 
         config = loadjson(os.path.join(model_name, "config.json"))
         model_path = os.path.join(model_name, "best_model.pt")
-        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
         tmp = ALIGNNAtomWiseConfig(**config["model"])
         model = ALIGNNAtomWise(tmp)
         model.load_state_dict(torch.load(model_path, map_location=device))

From 76c56ebe8f3e83a9144bb08250062a77a2f1eac6 Mon Sep 17 00:00:00 2001
From: user <writetokamal.1989@gmail.com>
Date: Sun, 5 Apr 2026 11:46:46 -0400
Subject: [PATCH 17/23] Pytest fix, new ASE

---
 alignn/ff/calculators.py         |  2 +
 alignn/ff/ff.py                  | 64 +++++++++++++++++++++++++-------
 alignn/models/alignn_atomwise.py |  3 +-
 3 files changed, 55 insertions(+), 14 deletions(-)

diff --git a/alignn/ff/calculators.py b/alignn/ff/calculators.py
index a47d5fe..eee7d0e 100644
--- a/alignn/ff/calculators.py
+++ b/alignn/ff/calculators.py
@@ -455,6 +455,7 @@ def __init__(
             )
         )
         ff_model.eval()
+        ff_model.to(self.device)
         self.ff_model = ff_model
         if prop_path is None and prop_model is None:
             prop_path = get_figshare_model_ff(
@@ -477,6 +478,7 @@ def __init__(
             )
         )
         prop_model.eval()
+        prop_model.to(self.device)
         self.prop_model = prop_model
 
     def calculate(
diff --git a/alignn/ff/ff.py b/alignn/ff/ff.py
index 846a515..9f2bc0a 100644
--- a/alignn/ff/ff.py
+++ b/alignn/ff/ff.py
@@ -458,13 +458,24 @@ def run_nvt_langevin(
             self.set_momentum_maxwell_boltzmann(
                 temperature_K=initial_temperature_K
             )
+        # self.dyn = Langevin(
+        #    self.atoms,
+        #    self.timestep,
+        #    temperature_K=temperature_K,
+        #    friction=friction,
+        #    comm=self.communicator,
+        #    # communicator=self.communicator,
+        # )
+        langevin_kwargs = dict(
+            temperature_K=temperature_K,
+            friction=friction,
+        )
+        if self.communicator is not None:
+            langevin_kwargs["comm"] = self.communicator
         self.dyn = Langevin(
             self.atoms,
             self.timestep,
-            temperature_K=temperature_K,
-            friction=friction,
-            comm=self.communicator,
-            # communicator=self.communicator,
+            **langevin_kwargs,
         )
         # Create monitors for logfile and a trajectory file
         # logfile = os.path.join(".", "%s.log" % filename)
@@ -492,13 +503,25 @@ def run_nvt_andersen(
             self.set_momentum_maxwell_boltzmann(
                 temperature_K=initial_temperature_K
             )
+        # self.dyn = Andersen(
+        #    self.atoms,
+        #    self.timestep,
+        #    temperature_K=temperature_K,
+        #    andersen_prob=andersen_prob,
+        #    comm=self.communicator,
+        #    # communicator=self.communicator,
+        # )
+
+        andersen_kwargs = dict(
+            temperature_K=temperature_K,
+            andersen_prob=andersen_prob,
+        )
+        if self.communicator is not None:
+            andersen_kwargs["comm"] = self.communicator
         self.dyn = Andersen(
             self.atoms,
             self.timestep,
-            temperature_K=temperature_K,
-            andersen_prob=andersen_prob,
-            comm=self.communicator,
-            # communicator=self.communicator,
+            **andersen_kwargs,
         )
         # Create monitors for logfile and a trajectory file
         # logfile = os.path.join(".", "%s.log" % filename)
@@ -565,16 +588,31 @@ def run_npt_berendsen(
             self.set_momentum_maxwell_boltzmann(
                 temperature_K=initial_temperature_K
             )
-        self.dyn = NPTBerendsen(
-            self.atoms,
-            self.timestep,
+        # self.dyn = NPTBerendsen(
+        #    self.atoms,
+        #    self.timestep,
+        #    temperature_K=temperature_K,
+        #    taut=taut,
+        #    taup=taup,
+        #    pressure=pressure,
+        #    compressibility=compressibility,
+        #    comm=self.communicator,
+        #    # communicator=self.communicator,
+        # )
+
+        npt_kwargs = dict(
             temperature_K=temperature_K,
             taut=taut,
             taup=taup,
             pressure=pressure,
             compressibility=compressibility,
-            comm=self.communicator,
-            # communicator=self.communicator,
+        )
+        if self.communicator is not None:
+            npt_kwargs["comm"] = self.communicator
+        self.dyn = NPTBerendsen(
+            self.atoms,
+            self.timestep,
+            **npt_kwargs,
         )
         # Create monitors for logfile and a trajectory file
         # logfile = os.path.join(".", "%s.log" % filename)
diff --git a/alignn/models/alignn_atomwise.py b/alignn/models/alignn_atomwise.py
index 40a3dd4..aff0243 100644
--- a/alignn/models/alignn_atomwise.py
+++ b/alignn/models/alignn_atomwise.py
@@ -343,7 +343,8 @@ def __init__(
                 config.hidden_features, config.additional_output_features
             )
         if self.classification:
-            self.fc = nn.Linear(config.hidden_features, 1)
+            # self.fc = nn.Linear(config.hidden_features, 1)
+            self.fc = nn.Linear(config.hidden_features, 2)
             # self.softmax = nn.Sigmoid()
             self.softmax = nn.LogSoftmax(dim=1)
         else:

From b635d65ad43f1498d1357e06e29a15d57d2d725c Mon Sep 17 00:00:00 2001
From: user <writetokamal.1989@gmail.com>
Date: Sun, 5 Apr 2026 17:40:30 -0400
Subject: [PATCH 18/23] Store directory update

---
 README.md                | 44 ++++++++++++++++++++--------------------
 alignn/__init__.py       |  2 +-
 alignn/ff/calculators.py |  4 +++-
 alignn/ff/ff.py          |  8 +++++---
 alignn/pretrained.py     |  5 +++--
 setup.py                 |  2 +-
 6 files changed, 35 insertions(+), 30 deletions(-)

diff --git a/README.md b/README.md
index 0230fb4..847cd90 100644
--- a/README.md
+++ b/README.md
@@ -1,9 +1,9 @@
-![alt text](https://github.com/usnistgov/alignn/actions/workflows/main.yml/badge.svg)
-[![codecov](https://codecov.io/gh/usnistgov/alignn/branch/main/graph/badge.svg?token=S5X4OYC80V)](https://codecov.io/gh/usnistgov/alignn)
+![alt text](https://github.com/atomgptlab/alignn/actions/workflows/main.yml/badge.svg)
+[![codecov](https://codecov.io/gh/atomgptlab/alignn/branch/main/graph/badge.svg?token=S5X4OYC80V)](https://codecov.io/gh/atomgptlab/alignn)
 [![PyPI version](https://badge.fury.io/py/alignn.svg)](https://badge.fury.io/py/alignn)
-![GitHub tag (latest by date)](https://img.shields.io/github/v/tag/usnistgov/alignn)
-![GitHub code size in bytes](https://img.shields.io/github/languages/code-size/usnistgov/alignn)
-![GitHub commit activity](https://img.shields.io/github/commit-activity/y/usnistgov/alignn)
+![GitHub tag (latest by date)](https://img.shields.io/github/v/tag/atomgptlab/alignn)
+![GitHub code size in bytes](https://img.shields.io/github/languages/code-size/atomgptlab/alignn)
+![GitHub commit activity](https://img.shields.io/github/commit-activity/y/atomgptlab/alignn)
 [![Downloads](https://pepy.tech/badge/alignn)](https://pepy.tech/project/alignn)
 <!--
 [![PWC](https://img.shields.io/endpoint.svg?url=https://paperswithcode.com/badge/atomistic-line-graph-neural-network-for/formation-energy-on-materials-project)](https://paperswithcode.com/sota/formation-energy-on-materials-project?p=atomistic-line-graph-neural-network-for)
@@ -35,7 +35,7 @@ Atomisitic line graph neural network-based FF (ALIGNN-FF) (https://pubs.rsc.org/
 
 
 
-![ALIGNN layer schematic](https://github.com/usnistgov/alignn/blob/develop/alignn/tex/schematic_lg.jpg)
+![ALIGNN layer schematic](https://github.com/atomgptlab/alignn/blob/develop/alignn/tex/schematic_lg.jpg)
 
 <a name="install"></a>
 Installation
@@ -72,7 +72,7 @@ You can laso install a development version of alignn by cloning the repository a
 conda create --name my_alignn python=3.10 -y
 conda activate my_alignn
 conda install dgl=2.1.0 pytorch torchvision torchaudio pytorch-cuda -c pytorch -c nvidia
-git clone https://github.com/usnistgov/alignn
+git clone https://github.com/atomgptlab/alignn
 cd alignn
 python -m pip install -e .
 ```
@@ -122,9 +122,9 @@ Users can keep their structure files in `POSCAR`, `.cif`, `.xyz` or `.pdb` files
 
 In this directory, `id_prop.csv`, the filenames, and correponding target values are kept in `comma separated values (csv) format`.
 
-Here is an example of training OptB88vdw bandgaps of 50 materials from JARVIS-DFT database. The example is created using the [generate_sample_data_reg.py](https://github.com/usnistgov/alignn/blob/main/alignn/examples/sample_data/scripts/generate_sample_data_reg.py) script. Users can modify the script for more than 50 data, or make their own dataset in this format. For list of available datasets see [Databases](https://jarvis-tools.readthedocs.io/en/master/databases.html).
+Here is an example of training OptB88vdw bandgaps of 50 materials from JARVIS-DFT database. The example is created using the [generate_sample_data_reg.py](https://github.com/atomgptlab/alignn/blob/main/alignn/examples/sample_data/scripts/generate_sample_data_reg.py) script. Users can modify the script for more than 50 data, or make their own dataset in this format. For list of available datasets see [Databases](https://jarvis-tools.readthedocs.io/en/master/databases.html).
 
-The dataset in split in 80:10:10 as training-validation-test set (controlled by `train_ratio, val_ratio, test_ratio`) . To change the split proportion and other parameters, change the [config_example.json](https://github.com/usnistgov/alignn/blob/main/alignn/examples/sample_data/config_example.json) file. If, users want to train on certain sets and val/test on another dataset, set `n_train`, `n_val`, `n_test` manually in the `config_example.json` and also set `keep_data_order` as True there so that random shuffle is disabled.
+The dataset in split in 80:10:10 as training-validation-test set (controlled by `train_ratio, val_ratio, test_ratio`) . To change the split proportion and other parameters, change the [config_example.json](https://github.com/atomgptlab/alignn/blob/main/alignn/examples/sample_data/config_example.json) file. If, users want to train on certain sets and val/test on another dataset, set `n_train`, `n_val`, `n_test` manually in the `config_example.json` and also set `keep_data_order` as True there so that random shuffle is disabled.
 
 A brief help guide (`-h`) can be obtained as follows.
 
@@ -169,7 +169,7 @@ To finetune model, use `--restart_model_path` tag as well in the above with the
 train_alignn.py --root_dir "alignn/examples/sample_data_ff" --restart_model_path "temp/best_model.pt" --config "alignn/examples/sample_data_ff/config_example_atomwise.json" --output_dir="temp1"
 ```
 
-Starting version v2024.10.30, we also allow global training for multi-output along with energy (graph wise output), forces (atomwise gradients), charges/magnetic moments etc. (atomwise but non-gradients) properties with or without additional fingerprints/features in graph. See examples [here](https://github.com/usnistgov/alignn/tree/main/alignn/examples). 
+Starting version v2024.10.30, we also allow global training for multi-output along with energy (graph wise output), forces (atomwise gradients), charges/magnetic moments etc. (atomwise but non-gradients) properties with or without additional fingerprints/features in graph. See examples [here](https://github.com/atomgptlab/alignn/tree/main/alignn/examples). 
 
 Multi-GPU training is allowed with `DistributedDataParallel` with `torchrun` command. This feature is not thoroughly tested yet. 
 Example:
@@ -179,7 +179,7 @@ torchrun --nproc_per_node=4 train_alignn.py --root_dir DataDir --config config.j
 ```
 For multi-GPU training make sure you have correct SLURM/PBS script setup correctly such as `#SBATCH -n 4, #SBATCH -N 1, #SBATCH --gres=gpu:4` etc.
 
-High-throughput like training: Users can also try training using multiple example scripts to run multiple dataset (such as JARVIS-DFT, Materials project, QM9_JCTC etc.). Look into the [alignn/scripts/train_*.py](https://github.com/usnistgov/alignn/tree/main/alignn/scripts) folder. This is done primarily to make the trainings more automated rather than making folder/ csv files etc.
+High-throughput like training: Users can also try training using multiple example scripts to run multiple dataset (such as JARVIS-DFT, Materials project, QM9_JCTC etc.). Look into the [alignn/scripts/train_*.py](https://github.com/atomgptlab/alignn/tree/main/alignn/scripts) folder. This is done primarily to make the trainings more automated rather than making folder/ csv files etc.
 These scripts automatically download datasets from [Databases in jarvis-tools](https://jarvis-tools.readthedocs.io/en/master/databases.html) and train several models. Make sure you specify your specific queuing system details in the scripts.
 
 
@@ -193,7 +193,7 @@ Using pre-trained models
 
 All the trained models are distributed on [Figshare](https://figshare.com/projects/ALIGNN_models/126478.
 
-The [pretrained.py script](https://github.com/usnistgov/alignn/blob/develop/alignn/pretrained.py) can be applied to use them. These models can be used to directly make predictions.
+The [pretrained.py script](https://github.com/atomgptlab/alignn/blob/develop/alignn/pretrained.py) can be applied to use them. These models can be used to directly make predictions.
 
 A brief help section (`-h`) is shown using:
 
@@ -226,18 +226,18 @@ Several supporting scripts for stucture optimization, equation of states, phonon
 Web-apps
 ------------
 
-A basic web-app is for direct-prediction available at [JARVIS-ALIGNN app](https://jarvis.nist.gov/jalignn/). Given atomistic structure in POSCAR format it predict formation energy, total energy per atom and bandgap using data trained on JARVIS-DFT dataset.
+A basic web-app is for direct-prediction available at [AtomGPT.org ALIGNN app](https://atomgpt.org/alignn). Given atomistic structure in POSCAR format it predict formation energy, total energy per atom and bandgap using data trained on JARVIS-DFT dataset.
 
-Similarly, a web-app for [ALIGNN-FF](https://jarvis.nist.gov/jalignnff/) for structure optimization is also available.
+Similarly, a web-app for [ALIGNN-FF](https://atomgpt.org/alignn_ff_dynamics) for structure optimization is also available.
 
-![JARVIS-ALIGNN](https://github.com/usnistgov/alignn/blob/develop/alignn/tex/jalignn.PNG)
+![JARVIS-ALIGNN](https://github.com/atomgptlab/alignn/blob/develop/alignn/tex/jalignn.PNG)
 
 
 
 <a name="alignnff"></a>
 ALIGNN-FF ASE Calculaor
 -------------------------
-![ALIGNN FF](https://github.com/usnistgov/alignn/blob/develop/alignn/tex/animation.gif)
+![ALIGNN FF](https://github.com/atomgptlab/alignn/blob/develop/alignn/tex/animation.gif)
 
 
 [ASE calculator](https://wiki.fysik.dtu.dk/ase/ase/calculators/calculators.html) provides interface to various codes. An example for ALIGNN-FF is give below. Note that there are multiple pretrained ALIGNN-FF models available, here we use the deafult_path model. As more accurate models are developed, they will be made available as well:
@@ -374,7 +374,7 @@ The results from models other than ALIGNN are reported as given in corresponding
 
 ### 4) On QM9 dataset
 
-Note the [issue](https://github.com/usnistgov/alignn/issues/54) related to QM9 dataset. The results from models other than ALIGNN are reported as given in corresponding papers, not necessarily reproduced by us. These models were trained with same parameters as solid-state databases but for 1000 epochs.
+Note the [issue](https://github.com/atomgptlab/alignn/issues/54) related to QM9 dataset. The results from models other than ALIGNN are reported as given in corresponding papers, not necessarily reproduced by us. These models were trained with same parameters as solid-state databases but for 1000 epochs.
 
 | Target | Units | SchNet | MEGNet  | DimeNet++ | ALIGNN |
 |:------:|-------|--------|---------|-----------|--------|
@@ -440,7 +440,7 @@ Useful notes (based on some of the queries we received)
 2) While comnventional '.cif' and '.pdb' files can be read using jarvis-tools, for complex files you might have to install `cif2cell` and `pytraj` respectively i.e.`pip install cif2cell==2.0.0a3` and `conda install -c ambermd pytraj`.
 3) Make sure you use `batch_size` as 32 or 64 for large datasets, and not 2 as given in the example config file, else it will take much longer to train, and performnce might drop a lot.
 4) Note that `train_alignn.py` and `pretrained.py` in alignn folder are actually python executable scripts. So, even if you don't provide absolute path of these scripts, they should work.
-5) Learn about the issue with QM9 results here: https://github.com/usnistgov/alignn/issues/54
+5) Learn about the issue with QM9 results here: https://github.com/atomgptlab/alignn/issues/54
 6) Make sure you have `pandas` version as >1.2.3.
 7) Starting March 2024, pytroch-ignite dependency will be removed to enable conda-forge build.
 
@@ -459,7 +459,7 @@ References
 8) [Unified graph neural network force-field for the periodic table](https://pubs.rsc.org/en/content/articlehtml/2023/dd/d2dd00096b)
 9) [Large Scale Benchmark of Materials Design Methods](https://www.nature.com/articles/s41524-024-01259-w)
 10) [Prediction of Magnetic Properties in van der Waals Magnets using Graph Neural Networks](https://doi.org/10.1103/PhysRevMaterials.8.114002)
-11) [CHIPS-FF: Benchmarking universal force-fields](https://github.com/usnistgov/chipsff)
+11) [CHIPS-FF: Benchmarking universal force-fields](https://github.com/atomgptlab/chipsff)
 
 
 Please see detailed publications list [here](https://jarvis-tools.readthedocs.io/en/master/publications.html).
@@ -468,13 +468,13 @@ Please see detailed publications list [here](https://jarvis-tools.readthedocs.io
 How to contribute
 -----------------
 
-For detailed instructions, please see [Contribution instructions](https://github.com/usnistgov/jarvis/blob/master/Contribution.rst)
+For detailed instructions, please see [Contribution instructions](https://github.com/atomgptlab/jarvis/blob/master/Contribution.rst)
 
 <a name="corres"></a>
 Correspondence
 --------------------
 
-Please report bugs as Github issues (https://github.com/usnistgov/alignn/issues) or email to kamal.choudhary@nist.gov.
+Please report bugs as Github issues (https://github.com/atomgptlab/alignn/issues) or email to drkamal@jhu.edu .
 
 <a name="fund"></a>
 Funding support
@@ -487,4 +487,4 @@ NIST-CHIPS (https://www.nist.gov/chips)
 Code of conduct
 --------------------
 
-Please see [Code of conduct](https://github.com/usnistgov/jarvis/blob/master/CODE_OF_CONDUCT.md)
+Please see [Code of conduct](https://github.com/atomgptlab/jarvis/blob/master/CODE_OF_CONDUCT.md)
diff --git a/alignn/__init__.py b/alignn/__init__.py
index e396d61..0c35900 100644
--- a/alignn/__init__.py
+++ b/alignn/__init__.py
@@ -1,3 +1,3 @@
 """Version number."""
 
-__version__ = "2025.12.1"
+__version__ = "2026.4.1"
diff --git a/alignn/ff/calculators.py b/alignn/ff/calculators.py
index eee7d0e..ebc1e6c 100644
--- a/alignn/ff/calculators.py
+++ b/alignn/ff/calculators.py
@@ -19,6 +19,7 @@
 import numpy as np
 from tqdm import tqdm
 import torch
+from jarvis.core.utils import get_cache_dir
 
 # Reference: https://doi.org/10.1039/D2DD00096B
 
@@ -44,7 +45,8 @@ def get_figshare_model_ff(
     all_models_ff = get_all_models()
     # https://doi.org/10.6084/m9.figshare.23695695
     if dir_path is None:
-        dir_path = str(os.path.join(os.path.dirname(__file__), model_name))
+        # dir_path = str(os.path.join(os.path.dirname(__file__), model_name))
+        dir_path = os.path.join(get_cache_dir("alignn_ff"), model_name)
     # cwd=os.getcwd()
     dir_path = os.path.abspath(dir_path)
     if not os.path.exists(dir_path):
diff --git a/alignn/ff/ff.py b/alignn/ff/ff.py
index 9f2bc0a..b065e4f 100644
--- a/alignn/ff/ff.py
+++ b/alignn/ff/ff.py
@@ -47,7 +47,7 @@
 from matplotlib.gridspec import GridSpec
 from sklearn.metrics import mean_absolute_error
 from tqdm import tqdm
-
+from jarvis.core.utils import get_cache_dir
 # import torch
 from alignn.ff.calculators import (
     AlignnAtomwiseCalculator,
@@ -85,7 +85,8 @@ def get_figshare_model_ff(
     all_models_ff = get_all_models()
     # https://doi.org/10.6084/m9.figshare.23695695
     if dir_path is None:
-        dir_path = str(os.path.join(os.path.dirname(__file__), model_name))
+        # dir_path = str(os.path.join(os.path.dirname(__file__), model_name))
+        dir_path = os.path.join(get_cache_dir("alignn_ff"), model_name)
     # cwd=os.getcwd()
     dir_path = os.path.abspath(dir_path)
     if not os.path.exists(dir_path):
@@ -145,7 +146,8 @@ def get_figshare_model_prop(
     all_models_prop = get_all_models_prop()
     # https://doi.org/10.6084/m9.figshare.23695695
     if dir_path is None:
-        dir_path = str(os.path.join(os.path.dirname(__file__), model_name))
+        # dir_path = str(os.path.join(os.path.dirname(__file__), model_name))
+        dir_path = os.path.join(get_cache_dir("alignn_models"), model_name)
     # cwd=os.getcwd()
     dir_path = os.path.abspath(dir_path)
     if not os.path.exists(dir_path):
diff --git a/alignn/pretrained.py b/alignn/pretrained.py
index 7bc8fe6..c05a9d5 100644
--- a/alignn/pretrained.py
+++ b/alignn/pretrained.py
@@ -22,7 +22,7 @@
     ALIGNNAtomWise,
     ALIGNNAtomWiseConfig,
 )
-
+from jarvis.core.utils import get_cache_dir
 # from jarvis.core.graphs import Graph
 
 tqdm.pandas()
@@ -260,7 +260,8 @@ def get_figshare_model(model_name="jv_formation_energy_peratom_alignn"):
     # else:
     #    config_params = {}
     zfile = model_name + ".zip"
-    path = str(os.path.join(os.path.dirname(__file__), zfile))
+    # path = str(os.path.join(os.path.dirname(__file__), zfile))
+    path = os.path.join(get_cache_dir("alignn_models"), zfile)
     if not os.path.isfile(path):
         response = requests.get(url, stream=True)
         total_size_in_bytes = int(response.headers.get("content-length", 0))
diff --git a/setup.py b/setup.py
index 8acfb73..ab9bb21 100644
--- a/setup.py
+++ b/setup.py
@@ -10,7 +10,7 @@
 
 setuptools.setup(
     name="alignn",
-    version="2025.12.1",
+    version="2026.4.1",
     author="Kamal Choudhary, Brian DeCost",
     author_email="kamal.choudhary@nist.gov",
     description="alignn",

From e1e7252343460341821cae53236b0404a8aa9246 Mon Sep 17 00:00:00 2001
From: user <writetokamal.1989@gmail.com>
Date: Sun, 5 Apr 2026 17:58:45 -0400
Subject: [PATCH 19/23] database update

---
 README.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/README.md b/README.md
index 847cd90..59edf33 100644
--- a/README.md
+++ b/README.md
@@ -226,7 +226,7 @@ Several supporting scripts for stucture optimization, equation of states, phonon
 Web-apps
 ------------
 
-A basic web-app is for direct-prediction available at [AtomGPT.org ALIGNN app](https://atomgpt.org/alignn). Given atomistic structure in POSCAR format it predict formation energy, total energy per atom and bandgap using data trained on JARVIS-DFT dataset.
+A basic web-app is for direct-prediction available at [AtomGPT.org ALIGNN app](https://atomgpt.org/alignn). Given atomistic structure in POSCAR format it predict formation energy, total energy per atom, bandgap and other props using data trained on JARVIS-DFT dataset.
 
 Similarly, a web-app for [ALIGNN-FF](https://atomgpt.org/alignn_ff_dynamics) for structure optimization is also available.
 

From d0f216bc5f325718ed61f0307c3886299092d445 Mon Sep 17 00:00:00 2001
From: user <writetokamal.1989@gmail.com>
Date: Sun, 5 Apr 2026 18:09:29 -0400
Subject: [PATCH 20/23] database update

---
 environment.yml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/environment.yml b/environment.yml
index 115f917..6b42391 100644
--- a/environment.yml
+++ b/environment.yml
@@ -60,7 +60,7 @@ dependencies:
   - idna=3.6=pyhd8ed1ab_0
   - importlib-metadata=7.1.0=pyha770c72_0
   - iniconfig=2.0.0=pyhd8ed1ab_0
-  - jarvis-tools=2024.3.24=pyhd8ed1ab_0
+  - jarvis-tools
   - jinja2=3.1.3=pyhd8ed1ab_0
   - joblib=1.3.2=pyhd8ed1ab_0
   - jsonpatch=1.33=pyhd8ed1ab_0

From e9fb8e25d2f371599c4bd34711db8c51b49c37a6 Mon Sep 17 00:00:00 2001
From: user <writetokamal.1989@gmail.com>
Date: Sun, 5 Apr 2026 18:44:04 -0400
Subject: [PATCH 21/23] DB update

---
 README.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/README.md b/README.md
index 59edf33..0d07c5d 100644
--- a/README.md
+++ b/README.md
@@ -226,7 +226,7 @@ Several supporting scripts for stucture optimization, equation of states, phonon
 Web-apps
 ------------
 
-A basic web-app is for direct-prediction available at [AtomGPT.org ALIGNN app](https://atomgpt.org/alignn). Given atomistic structure in POSCAR format it predict formation energy, total energy per atom, bandgap and other props using data trained on JARVIS-DFT dataset.
+A basic web-app is for direct-prediction available at [AtomGPT.org ALIGNN app](https://atomgpt.org/alignn). Given atomistic structure in POSCAR format it predict formation energy, total energy per atom, bandgap and other properties using data trained on JARVIS-DFT dataset.
 
 Similarly, a web-app for [ALIGNN-FF](https://atomgpt.org/alignn_ff_dynamics) for structure optimization is also available.
 

From dccb9b4ce73f625005eb7a692936866076042152 Mon Sep 17 00:00:00 2001
From: Kamal Choudhary <writetokamal.1989@gmail.com>
Date: Sun, 5 Apr 2026 17:24:56 -0700
Subject: [PATCH 22/23] Update yml

---
 environment.yml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/environment.yml b/environment.yml
index 6b42391..f3f686f 100644
--- a/environment.yml
+++ b/environment.yml
@@ -60,7 +60,6 @@ dependencies:
   - idna=3.6=pyhd8ed1ab_0
   - importlib-metadata=7.1.0=pyha770c72_0
   - iniconfig=2.0.0=pyhd8ed1ab_0
-  - jarvis-tools
   - jinja2=3.1.3=pyhd8ed1ab_0
   - joblib=1.3.2=pyhd8ed1ab_0
   - jsonpatch=1.33=pyhd8ed1ab_0
@@ -240,6 +239,7 @@ dependencies:
       - flake8==7.0.0
       - fsspec==2024.3.1
       - future==1.0.0
+      - jarvis-tools
       - lightgbm==3.3.1
       - mccabe==0.7.0
       - mkdocs-material==9.5.14

From 6786071d9e456ea5e27b6025c9698dac5d8e3e7e Mon Sep 17 00:00:00 2001
From: user <writetokamal.1989@gmail.com>
Date: Sun, 5 Apr 2026 21:22:09 -0400
Subject: [PATCH 23/23] workflows

---
 .github/workflows/main.yml | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml
index 01b3c92..c12b9f3 100644
--- a/.github/workflows/main.yml
+++ b/.github/workflows/main.yml
@@ -35,7 +35,8 @@ jobs:
             find . -type f > before_test_files.txt
             conda env create -f environment.yml
             conda activate my_alignn
-            conda install pytest coverage codecov pandas numpy matplotlib phonopy scikit-learn jarvis-tools --quiet
+            conda install pytest coverage codecov pandas numpy matplotlib phonopy scikit-learn  --quiet
+            #conda install pytest coverage codecov pandas numpy matplotlib phonopy scikit-learn jarvis-tools --quiet
             export DGLBACKEND=pytorch
             export CUDA_VISIBLE_DEVICES="-1"
             pip install phonopy flake8 pytest pycodestyle pydocstyle codecov pytest-cov coverage