v0.3.1

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")

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

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 = "[email protected]" }]
 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"

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