From 19de8fffce687c5fc81c942b66a8af3d1dd0f456 Mon Sep 17 00:00:00 2001
From: Janosh Riebesell <janosh.riebesell@gmail.com>
Date: Mon, 23 Oct 2023 07:18:46 -0700
Subject: [PATCH] v0.3.1

---
 chgnet/model/composition_model.py | 4 ++--
 chgnet/trainer/trainer.py         | 6 +++---
 pyproject.toml                    | 4 ++--
 tests/test_model.py               | 2 +-
 4 files changed, 8 insertions(+), 8 deletions(-)

diff --git a/chgnet/model/composition_model.py b/chgnet/model/composition_model.py
index f67cbb71..01e1bb77 100644
--- a/chgnet/model/composition_model.py
+++ b/chgnet/model/composition_model.py
@@ -70,7 +70,7 @@ def _assemble_graphs(self, graphs: list[CrystalGraph]):
             assembled batch_graph that contains all information for model.
         """
         composition_feas = []
-        for _graph_idx, graph in enumerate(graphs):
+        for graph in graphs:
             composition_fea = torch.bincount(
                 graph.atomic_number - 1, minlength=self.max_num_elements
             )
@@ -201,7 +201,7 @@ def initialize_from(self, dataset: str):
         """Initialize pre-fitted weights from a dataset."""
         if dataset in ["MPtrj", "MPtrj_e"]:
             self.initialize_from_MPtrj()
-        elif dataset in ["MPF"]:
+        elif dataset == "MPF":
             self.initialize_from_MPF()
         else:
             raise NotImplementedError(f"{dataset=} not supported yet")
diff --git a/chgnet/trainer/trainer.py b/chgnet/trainer/trainer.py
index 38f409ae..466baa0c 100644
--- a/chgnet/trainer/trainer.py
+++ b/chgnet/trainer/trainer.py
@@ -151,7 +151,7 @@ def __init__(
                 eta_min=decay_fraction * learning_rate,
             )
             self.scheduler_type = "cos"
-        elif scheduler in ["CosRestartLR"]:
+        elif scheduler == "CosRestartLR":
             scheduler_params = kwargs.pop(
                 "scheduler_params", {"decay_fraction": 1e-2, "T_0": 10, "T_mult": 2}
             )
@@ -471,7 +471,7 @@ def get_best_model(self):
         if self.best_model is None:
             raise RuntimeError("the model needs to be trained first")
         MAE = min(self.training_history["e"]["val"])
-        print(f"Best model has val {MAE = :.4}")
+        print(f"Best model has val {MAE =:.4}")
         return self.best_model
 
     @property
@@ -616,7 +616,7 @@ def __init__(
             self.criterion = nn.MSELoss()
         elif criterion in ["MAE", "mae", "l1"]:
             self.criterion = nn.L1Loss()
-        elif criterion in ["Huber"]:
+        elif criterion == "Huber":
             self.criterion = nn.HuberLoss(delta=delta)
         else:
             raise NotImplementedError
diff --git a/pyproject.toml b/pyproject.toml
index c6caf92a..8113797b 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "chgnet"
-version = "0.3.0"
+version = "0.3.1"
 description = "Pretrained Universal Neural Network Potential for Charge-informed Atomistic Modeling"
 authors = [{ name = "Bowen Deng", email = "bowendeng@berkeley.edu" }]
 requires-python = ">=3.9"
@@ -45,7 +45,7 @@ find = { include = ["chgnet*"], exclude = ["tests", "tests*"] }
 
 [tool.setuptools.package-data]
 "chgnet" = ["*.json"]
-"chgnet.pretrained" = ["**/*"]
+"chgnet.pretrained" = ["*", "**/*"]
 
 [tool.ruff]
 target-version = "py39"
diff --git a/tests/test_model.py b/tests/test_model.py
index 6e7964a1..1536eb44 100644
--- a/tests/test_model.py
+++ b/tests/test_model.py
@@ -188,7 +188,7 @@ def test_predict_supercell() -> None:
         np.repeat(pristine_prediction["f"], 4, axis=0), rel=1e-4, abs=1e-4
     )
 
-    assert out["s"] == pytest.approx(pristine_prediction["s"], rel=1e-4, abs=1e-4)
+    assert out["s"] == pytest.approx(pristine_prediction["s"], rel=1e-3, abs=1e-3)
 
     assert out["site_energies"] == pytest.approx(
         np.repeat(pristine_prediction["site_energies"], 4), rel=1e-4, abs=1e-4