ruff fixes

chgnet/graph/converter.py

@@ -248,7 +248,7 @@ def _create_graph_fast(
         """
         center_index = np.ascontiguousarray(center_index)
         neighbor_index = np.ascontiguousarray(neighbor_index)
-        image = np.ascontiguousarray(image, dtype=np.int_)
+        image = np.ascontiguousarray(image, dtype=np.int64)
         distance = np.ascontiguousarray(distance)
         gc_saved = gc.get_threshold()
         gc.set_threshold(0)

chgnet/model/composition_model.py

@@ -146,7 +146,7 @@ def fit(
             if isinstance(structure, Structure):
                 atomic_number = torch.tensor(
                     [site.specie.Z for site in structure],
-                    dtype=int,
+                    dtype=np.int64,
                     requires_grad=False,
                 )
             else:

chgnet/model/model.py

@@ -672,7 +672,7 @@ def from_dict(cls, dct: dict, **kwargs) -> Self:
     @classmethod
     def from_file(cls, path: str, **kwargs) -> Self:
         """Build a CHGNet from a saved file."""
-        state = torch.load(path, map_location=torch.device("cpu"))
+        state = torch.load(path, map_location=torch.device("cpu"), weights_only=False)
         return cls.from_dict(state["model"], **kwargs)
 
     @classmethod

pyproject.toml

@@ -9,9 +9,9 @@ license = { text = "Modified BSD" }
 dependencies = [
     "ase>=3.23.0",
     "cython>=3",
-    "numpy>=1.26,<2",
+    "numpy>=1.26",
     "nvidia-ml-py3>=7.352.0",
-    "pymatgen>=2023.10.11",
+    "pymatgen>=2024.9.10",
     "torch>=1.11.0",
     "typing-extensions>=4.12",
 ]
@@ -29,7 +29,7 @@ classifiers = [
 ]
 
 [project.optional-dependencies]
-test = ["pytest-cov>=4", "pytest>=8"]
+test = ["pytest-cov>=4", "pytest>=8", "wandb>=0.17"]
 # needed to run interactive example notebooks
 examples = ["crystal-toolkit>=2023.11.3", "pandas>=2.2"]
 docs = ["lazydocs>=0.4"]
@@ -48,7 +48,12 @@ find = { include = ["chgnet*"], exclude = ["tests", "tests*"] }
 "chgnet.pretrained" = ["*", "**/*"]
 
 [build-system]
-requires = ["Cython", "setuptools>=65,<70", "wheel"]
+requires = [
+    "Cython",
+    "setuptools>=65",
+    "wheel",
+    "numpy>=2.0.0",
+]
 build-backend = "setuptools.build_meta"
 
 [tool.ruff]

tests/test_trainer.py

@@ -153,19 +153,19 @@ def test_wandb_init(mock_wandb):
     )
 
 
-def test_wandb_log_frequency(mock_wandb):
+def test_wandb_log_frequency(tmp_path, mock_wandb):
     trainer = Trainer(model=chgnet, wandb_path="test-project/test-run", epochs=1)
 
     # Test epoch logging
-    trainer.train(train_loader, val_loader, wandb_log_freq="epoch", save_dir="")
+    trainer.train(train_loader, val_loader, wandb_log_freq="epoch", save_dir=tmp_path)
     assert (
         mock_wandb.log.call_count == 2 * trainer.epochs
     ), "Expected one train and one val log per epoch"
 
     mock_wandb.log.reset_mock()
 
     # Test batch logging
-    trainer.train(train_loader, val_loader, wandb_log_freq="batch", save_dir="")
+    trainer.train(train_loader, val_loader, wandb_log_freq="batch", save_dir=tmp_path)
     expected_batch_calls = trainer.epochs * len(train_loader)
     assert (
         mock_wandb.log.call_count > expected_batch_calls
@@ -183,5 +183,5 @@ def test_wandb_log_frequency(mock_wandb):
 
     # Test no logging when wandb_path is not provided
     trainer_no_wandb = Trainer(model=chgnet, epochs=1)
-    trainer_no_wandb.train(train_loader, val_loader)
+    trainer_no_wandb.train(train_loader, val_loader, save_dir=tmp_path)
     mock_wandb.log.assert_not_called()