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/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

chgnet/utils/vasp_utils.py

@@ -38,7 +38,7 @@ def parse_vasp_dir(
  dict: a dictionary of lists with keys for structure, uncorrected_total_energy,
  energy_per_atom, force, magmom, stress.
  """
- if os.path.isdir(base_dir) is False:
+ if not os.path.isdir(base_dir):
  raise NotADirectoryError(f"{base_dir=} is not a directory")
 
  oszicar_path = zpath(f"{base_dir}/OSZICAR")

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_crystal_graph.py

@@ -1,15 +1,14 @@
 from __future__ import annotations
 
 from time import perf_counter
+from unittest.mock import patch
 
 import numpy as np
 from pymatgen.core import Structure
 
 from chgnet import ROOT
 from chgnet.graph import CrystalGraphConverter
 
-np.random.seed(0)
-
 structure = Structure.from_file(f"{ROOT}/examples/mp-18767-LiMnO2.cif")
 converter = CrystalGraphConverter(atom_graph_cutoff=5, bond_graph_cutoff=3)
 converter_legacy = CrystalGraphConverter(
@@ -127,55 +126,57 @@ def test_crystal_graph_different_cutoff_fast():
 
 
 def test_crystal_graph_perturb_legacy():
- np.random.seed(0)
  structure_perturbed = structure.copy()
- structure_perturbed.perturb(distance=0.1)
+ fixed_rng = np.random.default_rng(0)
+ with patch("numpy.random.default_rng", return_value=fixed_rng):
+ structure_perturbed.perturb(distance=0.1)
 
  start = perf_counter()
  graph = converter_legacy(structure_perturbed)
  print("Legacy test_crystal_graph_perturb time:", perf_counter() - start) # noqa: T201
 
  assert list(graph.atom_frac_coord.shape) == [8, 3]
- assert list(graph.atom_graph.shape) == [410, 2]
- assert (graph.atom_graph[:, 0] == 3).sum().item() == 53
- assert (graph.atom_graph[:, 1] == 3).sum().item() == 53
- assert (graph.atom_graph[:, 1] == 6).sum().item() == 50
-
- assert list(graph.bond_graph.shape) == [688, 5]
- assert (graph.bond_graph[:, 0] == 1).sum().item() == 90
- assert (graph.bond_graph[:, 1] == 36).sum().item() == 17
- assert (graph.bond_graph[:, 3] == 36).sum().item() == 17
- assert (graph.bond_graph[:, 2] == 306).sum().item() == 10
+ assert list(graph.atom_graph.shape) == [420, 2]
+ assert (graph.atom_graph[:, 0] == 3).sum().item() == 54
+ assert (graph.atom_graph[:, 1] == 3).sum().item() == 54
+ assert (graph.atom_graph[:, 1] == 6).sum().item() == 54
+
+ assert list(graph.bond_graph.shape) == [850, 5]
+ assert (graph.bond_graph[:, 0] == 1).sum().item() == 156
+ assert (graph.bond_graph[:, 1] == 36).sum().item() == 18
+ assert (graph.bond_graph[:, 3] == 36).sum().item() == 18
+ assert (graph.bond_graph[:, 2] == 306).sum().item() == 0
  assert (graph.bond_graph[:, 4] == 120).sum().item() == 0
  assert list(graph.lattice.shape) == [3, 3]
- assert list(graph.undirected2directed.shape) == [205]
- assert list(graph.directed2undirected.shape) == [410]
+ assert list(graph.undirected2directed.shape) == [210]
+ assert list(graph.directed2undirected.shape) == [420]
 
 
 def test_crystal_graph_perturb_fast():
- np.random.seed(0)
  structure_perturbed = structure.copy()
- structure_perturbed.perturb(distance=0.1)
+ fixed_rng = np.random.default_rng(0)
+ with patch("numpy.random.default_rng", return_value=fixed_rng):
+ structure_perturbed.perturb(distance=0.1)
 
  start = perf_counter()
  graph = converter_fast(structure_perturbed)
  print("Fast test_crystal_graph_perturb time:", perf_counter() - start) # noqa: T201
 
  assert list(graph.atom_frac_coord.shape) == [8, 3]
- assert list(graph.atom_graph.shape) == [410, 2]
- assert (graph.atom_graph[:, 0] == 3).sum().item() == 53
- assert (graph.atom_graph[:, 1] == 3).sum().item() == 53
- assert (graph.atom_graph[:, 1] == 6).sum().item() == 50
-
- assert list(graph.bond_graph.shape) == [688, 5]
- assert (graph.bond_graph[:, 0] == 1).sum().item() == 90
- assert (graph.bond_graph[:, 1] == 36).sum().item() == 17
- assert (graph.bond_graph[:, 3] == 36).sum().item() == 17
- assert (graph.bond_graph[:, 2] == 306).sum().item() == 10
+ assert list(graph.atom_graph.shape) == [420, 2]
+ assert (graph.atom_graph[:, 0] == 3).sum().item() == 54
+ assert (graph.atom_graph[:, 1] == 3).sum().item() == 54
+ assert (graph.atom_graph[:, 1] == 6).sum().item() == 54
+
+ assert list(graph.bond_graph.shape) == [850, 5]
+ assert (graph.bond_graph[:, 0] == 1).sum().item() == 156
+ assert (graph.bond_graph[:, 1] == 36).sum().item() == 18
+ assert (graph.bond_graph[:, 3] == 36).sum().item() == 18
+ assert (graph.bond_graph[:, 2] == 306).sum().item() == 0
  assert (graph.bond_graph[:, 4] == 120).sum().item() == 0
  assert list(graph.lattice.shape) == [3, 3]
- assert list(graph.undirected2directed.shape) == [205]
- assert list(graph.directed2undirected.shape) == [410]
+ assert list(graph.undirected2directed.shape) == [210]
+ assert list(graph.directed2undirected.shape) == [420]
 
 
 def test_crystal_graph_isotropic_strained_legacy():

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

tests/test_vasp_utils.py

@@ -62,7 +62,7 @@ def test_parse_vasp_dir_without_magmoms(tmp_path: Path):
 
 def test_parse_vasp_dir_no_data():
  # test non-existing directory
- with pytest.raises(FileNotFoundError, match="is not a directory"):
+ with pytest.raises(NotADirectoryError, match="is not a directory"):
  parse_vasp_dir(f"{ROOT}/tests/files/non-existent")
 
  # test existing directory without VASP files