From bbf4020226ffaf0f30068bb6674e44c79c2f1e0f Mon Sep 17 00:00:00 2001
From: knc6 <writetokamal.1989@gmail.com>
Date: Sun, 14 Apr 2024 08:59:27 -0400
Subject: [PATCH] Use canonize default and pretrained with best_model.pt

---
 alignn/__init__.py   | 2 +-
 alignn/graphs.py     | 3 ++-
 alignn/pretrained.py | 8 +++++---
 setup.py             | 2 +-
 4 files changed, 9 insertions(+), 6 deletions(-)

diff --git a/alignn/__init__.py b/alignn/__init__.py
index 2683b5af..83674b0f 100644
--- a/alignn/__init__.py
+++ b/alignn/__init__.py
@@ -1,3 +1,3 @@
 """Version number."""
 
-__version__ = "2024.3.24"
+__version__ = "2024.4.10"
diff --git a/alignn/graphs.py b/alignn/graphs.py
index 25ca9aaf..74a10b7f 100644
--- a/alignn/graphs.py
+++ b/alignn/graphs.py
@@ -364,7 +364,8 @@ def atom_dgl_multigraph(
         max_attempts=3,
         id: Optional[str] = None,
         compute_line_graph: bool = True,
-        use_canonize: bool = False,
+        use_canonize: bool = True,
+        # use_canonize: bool = False,
         use_lattice_prop: bool = False,
         cutoff_extra=3.5,
     ):
diff --git a/alignn/pretrained.py b/alignn/pretrained.py
index 87cae37b..02c7994a 100644
--- a/alignn/pretrained.py
+++ b/alignn/pretrained.py
@@ -12,14 +12,14 @@
 import torch
 import sys
 import json
-
-# from jarvis.db.jsonutils import loadjson
 import argparse
 from jarvis.core.atoms import Atoms
-from jarvis.core.graphs import Graph
+from alignn.graphs import Graph
 from jarvis.db.jsonutils import dumpjson
 import pandas as pd
 
+# from jarvis.core.graphs import Graph
+
 tqdm.pandas()
 
 """
@@ -257,6 +257,8 @@ def get_figshare_model(model_name="jv_formation_energy_peratom_alignn"):
             chks.append(i)
         if "config.json" in i:
             cfg = i
+        if "best_model.pt" in i:
+            tmp = i
 
     print("Using chk file", tmp, "from ", chks)
     print("Path", os.path.abspath(path))
diff --git a/setup.py b/setup.py
index a6ca5f14..e075929c 100644
--- a/setup.py
+++ b/setup.py
@@ -10,7 +10,7 @@
 
 setuptools.setup(
     name="alignn",
-    version="2024.3.24",
+    version="2024.4.10",
     author="Kamal Choudhary, Brian DeCost",
     author_email="kamal.choudhary@nist.gov",
     description="alignn",