From 26318276bf65a928ad769a5c65e09b065b240260 Mon Sep 17 00:00:00 2001 From: Alin M Elena Date: Sat, 28 Mar 2026 15:12:05 +0000 Subject: [PATCH 01/16] add gpu workflow and tests --- .github/workflows/bluesky.yml | 58 +++++++ tests/data/mlip_fine_tune.yml | 1 + tests/data/mlip_train.yml | 1 + tests/test_descriptors.py | 48 ++++-- tests/test_descriptors_cli.py | 65 +++++++- tests/test_elasticity.py | 18 +- tests/test_elasticity_cli.py | 18 +- tests/test_eos.py | 48 +++++- tests/test_eos_cli.py | 81 +++++++-- tests/test_geom_opt.py | 103 ++++++++++-- tests/test_geomopt_cli.py | 173 ++++++++++++++++--- tests/test_md.py | 296 ++++++++++++++++++++++++--------- tests/test_md_cli.py | 125 ++++++++++++-- tests/test_mlip_calculators.py | 73 ++++---- tests/test_neb.py | 65 ++++++-- tests/test_neb_cli.py | 61 ++++++- tests/test_phonons.py | 50 +++++- tests/test_phonons_cli.py | 123 ++++++++++++-- tests/test_single_point.py | 136 +++++++++++---- tests/test_singlepoint_cli.py | 108 ++++++++++-- tests/test_train_cli.py | 84 ++++++++-- 21 files changed, 1419 insertions(+), 316 deletions(-) create mode 100644 .github/workflows/bluesky.yml diff --git a/.github/workflows/bluesky.yml b/.github/workflows/bluesky.yml new file mode 100644 index 00000000..b839e50f --- /dev/null +++ b/.github/workflows/bluesky.yml @@ -0,0 +1,58 @@ +name: gpu-ci + +on: [push, pull_request] + +jobs: + + tests: + runs-on: [self-hosted, gpu] + if: github.repository == 'stfc/janus-core' + timeout-minutes: 60 + strategy: + matrix: + python-version: ["3.10","3.11","3.12"] + + steps: + - uses: actions/checkout@v4 + + - name: Install uv + uses: astral-sh/setup-uv@v6 + with: + python-version: ${{ matrix.python-version }} + + - name: Install "all" dependencies + run: uv sync --extra chgnet --extra dpa3 --extra d3 --extra grace --extra mace --extra orb --extra upet --extra plumed + + - name: Install PLUMED + uses: Iximiel/install-plumed@v1 + id: plumed + continue-on-error: true + + - name: Set environment variable based on plumed success + run: | + if [ "${{ steps.plumed.outcome }}" = "success" ]; then + echo "PLUMED_KERNEL=${{ steps.plumed.outputs.plumed_prefix }}/lib/libplumedKernel.dylib" >> $GITHUB_ENV + fi + + - name: Download extra models + run: | + python3 tests/models/extra_models.py tests/models/extra + + - name: Run test suite + env: + # show timings of tests + PYTEST_ADDOPTS: "--durations=0" + run: uv run --no-sync pytest -k cuda + + - name: Install updated e3nn dependencies + run: | + uv sync --extra mattersim --extra fairchem --extra sevennet --extra nequip --extra d3 + uv pip install --reinstall pynvml + uv pip install "fairchem-core[torch-extras]" --no-build-isolation + + - name: Run test suite for updated e3nn dependencies + env: + # show timings of tests + PYTEST_ADDOPTS: "--durations=0" + HF_TOKEN: ${{ secrets.HF_TOKEN }} + run: uv run --no-sync pytest tests/test_{mlip_calculators,single_point}.py -k cuda diff --git a/tests/data/mlip_fine_tune.yml b/tests/data/mlip_fine_tune.yml index e058b31d..4edb0c21 100644 --- a/tests/data/mlip_fine_tune.yml +++ b/tests/data/mlip_fine_tune.yml @@ -32,3 +32,4 @@ keep_checkpoints: False save_cpu: True weight_decay: 1e-8 eval_interval: 2 +plot: False diff --git a/tests/data/mlip_train.yml b/tests/data/mlip_train.yml index bd4f5f05..75c3bbbf 100644 --- a/tests/data/mlip_train.yml +++ b/tests/data/mlip_train.yml @@ -44,3 +44,4 @@ clip_grad: 100 keep_checkpoints: False keep_isolated_atoms: True save_cpu: True +plot: False diff --git a/tests/test_descriptors.py b/tests/test_descriptors.py index 11ca59ed..e853cadf 100644 --- a/tests/test_descriptors.py +++ b/tests/test_descriptors.py @@ -7,6 +7,7 @@ from ase import Atoms from ase.io import read import pytest +import torch from janus_core.calculations.descriptors import Descriptors from janus_core.calculations.single_point import SinglePoint @@ -17,11 +18,15 @@ MODEL_PATH = Path(__file__).parent / "models" / "mace_mp_small.model" -def test_calc_descriptors(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_calc_descriptors(tmp_path, device): """Test calculating equation of state from ASE atoms object.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + struct = read(DATA_PATH / "NaCl.cif") log_file = tmp_path / "descriptors.log" - struct.calc = choose_calculator(arch="mace_mp", model=MODEL_PATH) + struct.calc = choose_calculator(arch="mace_mp", model=MODEL_PATH, device=device) descriptors = Descriptors( struct, @@ -41,13 +46,18 @@ def test_calc_descriptors(tmp_path): ) -def test_calc_per_element(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_calc_per_element(tmp_path, device): """Test calculating descriptors for each element from SinglePoint object.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + log_file = tmp_path / "descriptors.log" single_point = SinglePoint( struct=DATA_PATH / "NaCl.cif", arch="mace", model=MODEL_PATH, + device=device, ) descriptors = Descriptors( @@ -68,14 +78,19 @@ def test_calc_per_element(tmp_path): assert atoms.info["mace_Na_descriptor"] == pytest.approx(-0.0020374985791535563) -def test_logging(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_logging(tmp_path, device): """Test attaching logger to Descriptors and emissions are saved to info.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + log_file = tmp_path / "descriptors.log" single_point = SinglePoint( struct=DATA_PATH / "NaCl.cif", arch="mace_mp", model=MODEL_PATH, + device=device, ) descriptors = Descriptors( @@ -92,12 +107,17 @@ def test_logging(tmp_path): assert single_point.struct.info["emissions"] > 0 -def test_dispersion(): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_dispersion(device): """Test using mace_mp with dispersion.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + single_point = SinglePoint( struct=DATA_PATH / "NaCl.cif", arch="mace_mp", calc_kwargs={"dispersion": False}, + device=device, ) descriptors = Descriptors( @@ -110,6 +130,7 @@ def test_dispersion(): struct=DATA_PATH / "NaCl.cif", arch="mace_mp", calc_kwargs={"dispersion": True}, + device=device, ) descriptors_disp = Descriptors( @@ -118,17 +139,24 @@ def test_dispersion(): ) descriptors_disp.run() - assert ( - descriptors_disp.struct.info["mace_mp_d3_descriptor"] - == descriptors.struct.info["mace_mp_descriptor"] + assert descriptors_disp.struct.info["mace_mp_d3_descriptor"] == pytest.approx( + descriptors.struct.info["mace_mp_descriptor"] ) -def test_not_implemented_error(): - """Test correct error raised if descriptors not implemented.""" +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_not_implemented_error(device): + """Test correct implemented error raised if descriptors not installed.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + + from tests.utils import skip_extras + + skip_extras("chgnet") single_point = SinglePoint( struct=DATA_PATH / "NaCl.cif", arch="chgnet", + device=device, ) with pytest.raises(NotImplementedError): Descriptors( diff --git a/tests/test_descriptors_cli.py b/tests/test_descriptors_cli.py index 8af2e382..e425ff2c 100644 --- a/tests/test_descriptors_cli.py +++ b/tests/test_descriptors_cli.py @@ -6,6 +6,7 @@ from ase.io import read import pytest +import torch from typer.testing import CliRunner import yaml @@ -31,8 +32,12 @@ def test_help(): assert "Usage: janus descriptors [OPTIONS]" in strip_ansi_codes(result.stdout) -def test_descriptors(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_descriptors(tmp_path, device): """Test calculating MLIP descriptors.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + with chdir(tmp_path): results_dir = Path("janus_results") out_path = results_dir / "NaCl-descriptors.extxyz" @@ -47,6 +52,8 @@ def test_descriptors(tmp_path): DATA_PATH / "NaCl.cif", "--arch", "mace_mp", + "--device", + device, ], ) assert result.exit_code == 0 @@ -100,8 +107,12 @@ def test_descriptors(tmp_path): clear_log_handlers() -def test_calc_per_element(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_calc_per_element(tmp_path, device): """Test calculating MLIP descriptors for each element.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + out_path = tmp_path / "test" / "NaCl-descriptors.extxyz" log_path = tmp_path / "test.log" summary_path = tmp_path / "summary.yml" @@ -114,6 +125,8 @@ def test_calc_per_element(tmp_path): DATA_PATH / "NaCl.cif", "--arch", "mace_mp", + "--device", + device, "--out", out_path, "--calc-per-element", @@ -142,8 +155,12 @@ def test_calc_per_element(tmp_path): ) -def test_invariant(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_invariant(tmp_path, device): """Test setting invariant_only to false.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + out_path = tmp_path / "NaCl-descriptors.extxyz" log_path = tmp_path / "test.log" summary_path = tmp_path / "summary.yml" @@ -156,6 +173,8 @@ def test_invariant(tmp_path): DATA_PATH / "NaCl.cif", "--arch", "mace_mp", + "--device", + device, "--out", out_path, "--no-invariants-only", @@ -184,8 +203,12 @@ def test_invariant(tmp_path): ) -def test_traj(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_traj(tmp_path, device): """Test calculating descriptors for a trajectory.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + out_path = tmp_path / "benzene-descriptors.extxyz" log_path = tmp_path / "test.log" summary_path = tmp_path / "summary.yml" @@ -198,6 +221,8 @@ def test_traj(tmp_path): DATA_PATH / "benzene-traj.xyz", "--arch", "mace_mp", + "--device", + device, "--out", out_path, "--read-kwargs", @@ -217,8 +242,12 @@ def test_traj(tmp_path): assert "mace_mp_descriptor" in atoms[1].info -def test_per_atom(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_per_atom(tmp_path, device): """Test calculating descriptors for each atom.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + out_path = tmp_path / "NaCl-descriptors.extxyz" log_path = tmp_path / "test.log" summary_path = tmp_path / "summary.yml" @@ -231,6 +260,8 @@ def test_per_atom(tmp_path): DATA_PATH / "NaCl.cif", "--arch", "mace_mp", + "--device", + device, "--out", out_path, "--calc-per-atom", @@ -261,8 +292,12 @@ def test_per_atom(tmp_path): ) -def test_no_carbon(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_no_carbon(tmp_path, device): """Test disabling carbon tracking.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + out_path = tmp_path / "test" / "NaCl-descriptors.extxyz" log_path = tmp_path / "test.log" summary_path = tmp_path / "summary.yml" @@ -275,6 +310,8 @@ def test_no_carbon(tmp_path): DATA_PATH / "NaCl.cif", "--arch", "mace_mp", + "--device", + device, "--out", out_path, "--log", @@ -292,8 +329,12 @@ def test_no_carbon(tmp_path): assert "emissions" not in descriptors_summary -def test_file_prefix(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_file_prefix(tmp_path, device): """Test file prefix creates directories and affects all files.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + file_prefix = tmp_path / "test/test" result = runner.invoke( app, @@ -303,6 +344,8 @@ def test_file_prefix(tmp_path): DATA_PATH / "NaCl.cif", "--arch", "mace_mp", + "--device", + device, "--file-prefix", file_prefix, ], @@ -317,8 +360,12 @@ def test_file_prefix(tmp_path): } -def test_model(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_model(tmp_path, device): """Test model passed correctly.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + file_prefix = tmp_path / "NaCl" results_path = tmp_path / "NaCl-descriptors.extxyz" log_path = tmp_path / "test.log" @@ -331,6 +378,8 @@ def test_model(tmp_path): DATA_PATH / "NaCl.cif", "--arch", "mace_mp", + "--device", + device, "--model", MACE_PATH, "--log", diff --git a/tests/test_elasticity.py b/tests/test_elasticity.py index bd97a44e..6b15ee69 100644 --- a/tests/test_elasticity.py +++ b/tests/test_elasticity.py @@ -7,7 +7,9 @@ from ase.build import bulk from ase.io import read import numpy as np +import pytest from pytest import approx +import torch from janus_core.calculations.elasticity import Elasticity from janus_core.helpers.mlip_calculators import choose_calculator @@ -17,15 +19,19 @@ MODEL_PATH = Path(__file__).parent / "models" / "mace_mp_small.model" -def test_calc_elasticity(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_calc_elasticity(tmp_path, device): """Test calculating elasticity for Aluminium.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + elasticity_path = tmp_path / "Al-elastic_tensor.dat" log_file = tmp_path / "Al-elasticity.log" generated_path = tmp_path / "Al-elasticity-generated.extxyz" minimized_path = tmp_path / "Al-elasticity-opt.extxyz" struct = bulk("Al", crystalstructure="fcc") - struct.calc = choose_calculator(arch="mace_mp", model=MODEL_PATH) + struct.calc = choose_calculator(arch="mace_mp", model=MODEL_PATH, device=device) elasticity = Elasticity( struct, @@ -86,15 +92,19 @@ def test_calc_elasticity(tmp_path): assert struct.info["strain"] == approx(strain.voigt) -def test_no_optimize_no_write_voigt(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_no_optimize_no_write_voigt(tmp_path, device): """Test calculating elasticity for Aluminium without optimization.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + elasticity_path = tmp_path / "Al-elastic_tensor.dat" log_file = tmp_path / "Al-elasticity.log" generated_path = tmp_path / "Al-elasticity-generated.extxyz" minimized_path = tmp_path / "Al-elasticity-opt.extxyz" struct = bulk("Al", crystalstructure="fcc") - struct.calc = choose_calculator(arch="mace_mp", model=MODEL_PATH) + struct.calc = choose_calculator(arch="mace_mp", model=MODEL_PATH, device=device) elasticity = Elasticity( struct, diff --git a/tests/test_elasticity_cli.py b/tests/test_elasticity_cli.py index 2bdb02d0..800f84f8 100644 --- a/tests/test_elasticity_cli.py +++ b/tests/test_elasticity_cli.py @@ -6,7 +6,9 @@ from ase.io import read import numpy as np +import pytest from pytest import approx +import torch from typer.testing import CliRunner import yaml @@ -26,8 +28,12 @@ def test_help(): assert "Usage: janus elasticity [OPTIONS]" in strip_ansi_codes(result.stdout) -def test_elasticity_opt_none(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_elasticity_opt_none(tmp_path, device): """Test calculating the ElasticTensor from the CLI without optimisation.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + with chdir(tmp_path): results_dir = Path("./janus_results") elasticity_path = results_dir / "NaCl-elastic_tensor.dat" @@ -44,6 +50,8 @@ def test_elasticity_opt_none(tmp_path): DATA_PATH / "NaCl.cif", "--arch", "mace_mp", + "--device", + device, "--n-strains", "2", "--no-minimize", @@ -95,8 +103,12 @@ def test_elasticity_opt_none(tmp_path): assert elasticity_summary["config"]["n_strains"] == 2 -def test_elasticity_opt_all(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_elasticity_opt_all(tmp_path, device): """Test calculating the ElasticTensor from the command line.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + with chdir(tmp_path): results_dir = Path("./janus_results") elasticity_path = results_dir / "NaCl-elastic_tensor.dat" @@ -113,6 +125,8 @@ def test_elasticity_opt_all(tmp_path): DATA_PATH / "NaCl.cif", "--arch", "mace_mp", + "--device", + device, "--n-strains", "2", "--minimize-all", diff --git a/tests/test_eos.py b/tests/test_eos.py index 4d1102ed..604a8534 100644 --- a/tests/test_eos.py +++ b/tests/test_eos.py @@ -8,6 +8,7 @@ from ase.eos import EquationOfState from ase.io import read import pytest +import torch from janus_core.calculations.eos import EoS from janus_core.calculations.single_point import SinglePoint @@ -18,14 +19,19 @@ MODEL_PATH = Path(__file__).parent / "models" / "mace_mp_small.model" -def test_calc_eos(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_calc_eos(tmp_path, device): """Test calculating equation of state from ASE atoms object.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + struct = read(DATA_PATH / "NaCl.cif") log_file = tmp_path / "eos.log" - struct.calc = choose_calculator(arch="mace_mp", model=MODEL_PATH) + struct.calc = choose_calculator(arch="mace_mp", model=MODEL_PATH, device=device) eos = EoS( struct, + device=device, file_prefix=tmp_path / "NaCl", log_kwargs={"filename": log_file}, ) @@ -39,16 +45,22 @@ def test_calc_eos(tmp_path): ) -def test_no_optimize(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_no_optimize(tmp_path, device): """Test not optimizing structure before calculation.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + log_file = tmp_path / "eos.log" single_point = SinglePoint( struct=DATA_PATH / "NaCl.cif", arch="mace", model=MODEL_PATH, + device=device, ) eos = EoS( single_point.struct, + device=device, minimize=False, file_prefix=tmp_path / "NaCl", log_kwargs={"filename": log_file}, @@ -64,10 +76,17 @@ def test_no_optimize(tmp_path): @pytest.mark.parametrize( "arch, device", - [("chgnet", "cpu"), ("sevennet", "cpu")], + [ + ("chgnet", "cpu"), + ("chgnet", "cuda"), + ("sevennet", "cpu"), + ("sevennet", "cuda"), + ], ) def test_extras(arch, device, tmp_path): """Test extra potentials.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") skip_extras(arch) eos_fit_path = tmp_path / "NaCl-eos-fit.dat" @@ -97,12 +116,16 @@ def test_extras(arch, device, tmp_path): pytest.skip() -def test_invalid_struct(): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_invalid_struct(device): """Test setting invalid structure.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") single_point = SinglePoint( struct=DATA_PATH / "benzene-traj.xyz", arch="mace_mp", model=MODEL_PATH, + device=device, ) with pytest.raises(NotImplementedError): @@ -115,18 +138,24 @@ def test_invalid_struct(): ) -def test_logging(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_logging(tmp_path, device): """Test attaching logger to EoS and emissions are saved to info.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + log_file = tmp_path / "eos.log" single_point = SinglePoint( struct=DATA_PATH / "NaCl.cif", arch="mace_mp", model=MODEL_PATH, + device=device, ) eos = EoS( single_point.struct, + device=device, file_prefix=tmp_path / "NaCl", log_kwargs={"filename": log_file}, ) @@ -139,14 +168,19 @@ def test_logging(tmp_path): assert single_point.struct.info["emissions"] > 0 -def test_plot(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_plot(tmp_path, device): """Test plotting equation of state.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + plot_file = tmp_path / "plot.svg" eos = EoS( struct=DATA_PATH / "NaCl.cif", arch="mace_mp", model=MODEL_PATH, + device=device, plot_to_file=True, plot_kwargs={"filename": plot_file}, file_prefix=tmp_path / "NaCl", diff --git a/tests/test_eos_cli.py b/tests/test_eos_cli.py index 11abdba6..5c2429cd 100644 --- a/tests/test_eos_cli.py +++ b/tests/test_eos_cli.py @@ -6,6 +6,7 @@ from ase.io import read import pytest +import torch from typer.testing import CliRunner import yaml @@ -31,8 +32,12 @@ def test_help(): assert "Usage: janus eos [OPTIONS]" in strip_ansi_codes(result.stdout) -def test_eos(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_eos(tmp_path, device): """Test calculating the equation of state.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + with chdir(tmp_path): results_dir = Path("./janus_results") eos_raw_path = results_dir / "NaCl-eos-raw.dat" @@ -48,6 +53,8 @@ def test_eos(tmp_path): DATA_PATH / "NaCl.cif", "--arch", "mace_mp", + "--device", + device, ], ) assert result.exit_code == 0 @@ -111,8 +118,12 @@ def test_eos(tmp_path): clear_log_handlers() -def test_setting_lattice(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_setting_lattice(tmp_path, device): """Test setting the lattice constants.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + file_prefix = tmp_path / "example" eos_raw_path = tmp_path / "example-eos-raw.dat" @@ -124,6 +135,8 @@ def test_setting_lattice(tmp_path): DATA_PATH / "NaCl.cif", "--arch", "mace_mp", + "--device", + device, "--min-volume", 0.8, "--max-volume", @@ -171,8 +184,12 @@ def test_invalid_lattice(option, value, tmp_path): assert isinstance(result.exception, ValueError) -def test_minimising_all(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_minimising_all(tmp_path, device): """Test minimising structures with different lattice constants.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + file_prefix = tmp_path / "NaCl" log_path = tmp_path / "NaCl-eos-log.yml" @@ -184,6 +201,8 @@ def test_minimising_all(tmp_path): DATA_PATH / "NaCl.cif", "--arch", "mace_mp", + "--device", + device, "--n-volumes", 4, "--minimize-all", @@ -204,8 +223,12 @@ def test_minimising_all(tmp_path): ) -def test_writing_structs(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_writing_structs(tmp_path, device): """Test writing out generated structures.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + file_prefix = tmp_path / "test" / "example" raw_path = tmp_path / "test" / "example-eos-raw.dat" fit_path = tmp_path / "test" / "example-eos-fit.dat" @@ -221,6 +244,8 @@ def test_writing_structs(tmp_path): DATA_PATH / "NaCl.cif", "--arch", "mace_mp", + "--device", + device, "--n-volumes", 4, "--file-prefix", @@ -275,9 +300,13 @@ def test_error_write_geomopt(tmp_path): assert isinstance(result.exception, ValueError) +@pytest.mark.parametrize("device", ["cpu", "cuda"]) @pytest.mark.parametrize("read_kwargs", ["{'index': 1}", "{}"]) -def test_valid_traj_input(read_kwargs, tmp_path): +def test_valid_traj_input(read_kwargs, tmp_path, device): """Test valid trajectory input structure handled.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + file_prefix = tmp_path / "traj" eos_raw_path = tmp_path / "traj-eos-raw.dat" eos_fit_path = tmp_path / "traj-eos-fit.dat" @@ -290,6 +319,8 @@ def test_valid_traj_input(read_kwargs, tmp_path): DATA_PATH / "NaCl-traj.xyz", "--arch", "mace_mp", + "--device", + device, "--read-kwargs", read_kwargs, "--file-prefix", @@ -323,8 +354,12 @@ def test_invalid_traj_input(tmp_path): assert isinstance(result.exception, ValueError) -def test_plot(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_plot(tmp_path, device): """Test plotting equation of state.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + file_prefix = tmp_path / "NaCl" raw_path = tmp_path / "NaCl-eos-raw.dat" fit_path = tmp_path / "NaCl-eos-fit.dat" @@ -340,6 +375,8 @@ def test_plot(tmp_path): DATA_PATH / "NaCl.cif", "--arch", "mace_mp", + "--device", + device, "--n-volumes", 4, "--plot-to-file", @@ -364,8 +401,12 @@ def test_plot(tmp_path): check_output_files(eos_summary, output_files) -def test_no_carbon(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_no_carbon(tmp_path, device): """Test disabling carbon tracking.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + file_prefix = tmp_path / "NaCl" summary_path = tmp_path / "NaCl-eos-summary.yml" @@ -377,6 +418,8 @@ def test_no_carbon(tmp_path): DATA_PATH / "NaCl.cif", "--arch", "mace_mp", + "--device", + device, "--n-volumes", 4, "--no-tracker", @@ -392,8 +435,12 @@ def test_no_carbon(tmp_path): assert "emissions" not in eos_summary -def test_model(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_model(tmp_path, device): """Test model passed correctly.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + file_prefix = tmp_path / "NaCl" generated_path = tmp_path / "NaCl-generated.extxyz" log_path = tmp_path / "test.log" @@ -406,6 +453,8 @@ def test_model(tmp_path): DATA_PATH / "NaCl.cif", "--arch", "mace_mp", + "--device", + device, "--model", MACE_PATH, "--n-volumes", @@ -446,8 +495,12 @@ def test_missing_arch(tmp_path): assert "Missing option" in result.stderr -def test_info(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_info(tmp_path, device): """Test info written to generated structures.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + file_prefix = tmp_path / "NaCl" generated_path = tmp_path / "NaCl-generated.extxyz" @@ -459,6 +512,8 @@ def test_info(tmp_path): DATA_PATH / "NaCl.cif", "--arch", "mace_mp", + "--device", + device, "--n-volumes", 4, "--file-prefix", @@ -480,8 +535,12 @@ def test_info(tmp_path): assert struct.info["config_type"] == "eos" -def test_info_min(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_info_min(tmp_path, device): """Test info written to generated structures after minimisation.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + file_prefix = tmp_path / "NaCl" generated_path = tmp_path / "NaCl-generated.extxyz" @@ -493,6 +552,8 @@ def test_info_min(tmp_path): DATA_PATH / "NaCl.cif", "--arch", "mace_mp", + "--device", + device, "--n-volumes", 4, "--file-prefix", diff --git a/tests/test_geom_opt.py b/tests/test_geom_opt.py index 67fa314a..ebe8b2fe 100644 --- a/tests/test_geom_opt.py +++ b/tests/test_geom_opt.py @@ -7,6 +7,7 @@ from ase.filters import FrechetCellFilter, UnitCellFilter from ase.io import read import pytest +import torch from janus_core.calculations.geom_opt import GeomOpt from janus_core.calculations.single_point import SinglePoint @@ -35,31 +36,41 @@ ] +@pytest.mark.parametrize("device", ["cpu", "cuda"]) @pytest.mark.parametrize("arch, struct, expected, kwargs", test_data) -def test_optimize(arch, struct, expected, kwargs): +def test_optimize(arch, struct, expected, kwargs, device): """Test optimizing geometry using MACE.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + single_point = SinglePoint( struct=DATA_PATH / struct, arch=arch, model=MODEL_PATH, + device=device, ) optimizer = GeomOpt(single_point.struct, **kwargs) optimizer.run() assert single_point.struct.get_potential_energy() == pytest.approx( - expected, rel=1e-8 + expected, rel=1e-6 ) -def test_saving_struct(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_saving_struct(tmp_path, device): """Test saving optimized structure.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + results_path = tmp_path / "NaCl.extxyz" single_point = SinglePoint( struct=DATA_PATH / "NaCl.cif", arch="mace", model=MODEL_PATH, + device=device, properties="energy", ) @@ -76,12 +87,17 @@ def test_saving_struct(tmp_path): assert opt_struct.info["mace_energy"] < init_energy -def test_saving_traj(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_saving_traj(tmp_path, device): """Test saving optimization trajectory output.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + optimizer = GeomOpt( struct=DATA_PATH / "NaCl.cif", arch="mace", model=MODEL_PATH, + device=device, write_traj=True, traj_kwargs={"filename": tmp_path / "NaCl.traj"}, ) @@ -125,18 +141,24 @@ def test_traj_without_write(tmp_path): ) -def test_hydrostatic_strain(): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_hydrostatic_strain(device): """Test setting hydrostatic strain for filter.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + single_point_1 = SinglePoint( struct=DATA_PATH / "NaCl-deformed.cif", arch="mace", model=MODEL_PATH, + device=device, ) single_point_2 = SinglePoint( struct=DATA_PATH / "NaCl-deformed.cif", arch="mace", model=MODEL_PATH, + device=device, ) optimizer_1 = GeomOpt( @@ -169,10 +191,14 @@ def test_hydrostatic_strain(): assert single_point_2.struct.cell.cellpar() == pytest.approx(expected_2) -def test_set_calc(): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_set_calc(device): """Test setting the calculator without SinglePoint.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + struct = read(DATA_PATH / "NaCl.cif") - struct.calc = choose_calculator(arch="mace_mp", model=MODEL_PATH) + struct.calc = choose_calculator(arch="mace_mp", model=MODEL_PATH, device=device) init_energy = struct.get_potential_energy() optimizer = GeomOpt(struct) @@ -180,12 +206,17 @@ def test_set_calc(): assert struct.get_potential_energy() < init_energy -def test_converge_warning(): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_converge_warning(device): """Test warning raised if not converged.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + single_point = SinglePoint( struct=DATA_PATH / "NaCl-deformed.cif", arch="mace_mp", model=MODEL_PATH, + device=device, ) optimizer = GeomOpt(single_point.struct, steps=1) @@ -193,12 +224,17 @@ def test_converge_warning(): optimizer.run() -def test_restart(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_restart(tmp_path, device): """Test restarting geometry optimization.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + single_point = SinglePoint( struct=DATA_PATH / "NaCl-deformed.cif", arch="mace_mp", model=MODEL_PATH, + device=device, properties="energy", ) @@ -227,12 +263,17 @@ def test_restart(tmp_path): assert final_energy < intermediate_energy -def test_space_group(): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_space_group(device): """Test spacegroup of the structure.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + single_point = SinglePoint( struct=DATA_PATH / "NaCl-sg.cif", arch="mace_mp", model=MODEL_PATH, + device=device, ) optimizer = GeomOpt(single_point.struct, fmax=0.001) @@ -242,14 +283,19 @@ def test_space_group(): assert single_point.struct.info["final_spacegroup"] == "Fm-3m (225)" -def test_str_optimizer(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_str_optimizer(tmp_path, device): """Test setting optimizer function with string.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + log_file = tmp_path / "opt.log" single_point = SinglePoint( struct=DATA_PATH / "NaCl-sg.cif", arch="mace_mp", model=MODEL_PATH, + device=device, ) optimizer = GeomOpt( @@ -281,14 +327,19 @@ def test_invalid_str_optimizer(): ) -def test_str_filter(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_str_filter(tmp_path, device): """Test setting filter with string.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + log_file = tmp_path / "opt.log" single_point = SinglePoint( struct=DATA_PATH / "NaCl-sg.cif", arch="mace_mp", model=MODEL_PATH, + device=device, ) optimizer = GeomOpt( @@ -339,13 +390,18 @@ def test_invalid_struct(): ) -def test_logging(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_logging(tmp_path, device): """Test attaching logger to GeomOpt and emissions are saved to info.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + log_file = tmp_path / "geomopt.log" single_point = SinglePoint( struct=DATA_PATH / "NaCl.cif", arch="mace_mp", model=MODEL_PATH, + device=device, ) assert "emissions" not in single_point.struct.info @@ -361,12 +417,17 @@ def test_logging(tmp_path): assert single_point.struct.info["emissions"] > 0 -def test_write_xyz(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_write_xyz(tmp_path, device): """Test writing a non-extended xyz file.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + optimizer = GeomOpt( struct=DATA_PATH / "NaCl-deformed.cif", arch="mace_mp", model=MODEL_PATH, + device=device, fmax=0.1, ) optimizer.run() @@ -381,8 +442,12 @@ def test_missing_arch(struct): GeomOpt(struct=struct) -def test_traj_new_dir(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_traj_new_dir(tmp_path, device): """Test writing trajectory extxyz in new directory via file_prefix.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + new_dir = tmp_path / "test" / "test" traj_path = new_dir / "NaCl-traj.extxyz" @@ -390,6 +455,7 @@ def test_traj_new_dir(tmp_path): struct=DATA_PATH / "NaCl.cif", arch="mace_mp", model=MODEL_PATH, + device=device, write_traj=True, file_prefix=new_dir / "NaCl", ) @@ -399,8 +465,12 @@ def test_traj_new_dir(tmp_path): assert len(traj) == 3 -def test_traj_kwargs_new_dir(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_traj_kwargs_new_dir(tmp_path, device): """Test writing trajectory in new directory via traj_kwargs.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + new_dir = tmp_path / "test" / "test" traj_path = new_dir / "NaCl-traj.traj" @@ -408,6 +478,7 @@ def test_traj_kwargs_new_dir(tmp_path): struct=DATA_PATH / "NaCl.cif", arch="mace_mp", model=MODEL_PATH, + device=device, write_traj=True, file_prefix=tmp_path / "NaCl", traj_kwargs={"filename": traj_path}, diff --git a/tests/test_geomopt_cli.py b/tests/test_geomopt_cli.py index 1ef21e4b..08e9904e 100644 --- a/tests/test_geomopt_cli.py +++ b/tests/test_geomopt_cli.py @@ -7,6 +7,7 @@ from ase import Atoms from ase.io import read import pytest +import torch from typer.testing import CliRunner import yaml @@ -33,8 +34,12 @@ def test_help(): assert "Usage: janus geomopt [OPTIONS]" in strip_ansi_codes(result.stdout) -def test_geomopt(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_geomopt(tmp_path, device): """Test geomopt calculation.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + with chdir(tmp_path): results_dir = Path("janus_results") results_path = results_dir / "NaCl-opt.extxyz" @@ -49,6 +54,8 @@ def test_geomopt(tmp_path): DATA_PATH / "NaCl.cif", "--arch", "mace_mp", + "--device", + device, "--fmax", "0.2", ], @@ -63,8 +70,12 @@ def test_geomopt(tmp_path): clear_log_handlers() -def test_log(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_log(tmp_path, device): """Test log correctly written for geomopt.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + results_path = tmp_path / "NaCl-opt.extxyz" log_path = tmp_path / "test.log" summary_path = tmp_path / "summary.yml" @@ -77,6 +88,8 @@ def test_log(tmp_path): DATA_PATH / "NaCl.cif", "--arch", "mace_mp", + "--device", + device, "--out", results_path, "--log", @@ -99,8 +112,12 @@ def test_log(tmp_path): ) -def test_traj(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_traj(tmp_path, device): """Test trajectory correctly written for geomopt.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + results_path = tmp_path / "NaCl-opt.extxyz" traj_path = tmp_path / "test.extxyz" log_path = tmp_path / "test.log" @@ -114,6 +131,8 @@ def test_traj(tmp_path): DATA_PATH / "NaCl.cif", "--arch", "mace_mp", + "--device", + device, "--out", results_path, "--write-traj", @@ -146,8 +165,12 @@ def test_traj(tmp_path): check_output_files(geomopt_summary, output_files) -def test_opt_fully(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_opt_fully(tmp_path, device): """Test passing --opt-cell-fully without --opt-cell-lengths.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + results_path = tmp_path / "NaCl-opt.extxyz" log_path = tmp_path / "test.log" summary_path = tmp_path / "summary.yml" @@ -160,6 +183,8 @@ def test_opt_fully(tmp_path): DATA_PATH / "NaCl-deformed.cif", "--arch", "mace_mp", + "--device", + device, "--out", results_path, "--opt-cell-fully", @@ -188,8 +213,12 @@ def test_opt_fully(tmp_path): assert atoms.cell.cellpar() == pytest.approx(expected) -def test_opt_fully_and_vectors(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_opt_fully_and_vectors(tmp_path, device): """Test passing --opt-cell-fully with --opt-cell-lengths.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + results_path = tmp_path / "NaCl-opt.extxyz" log_path = tmp_path / "test.log" summary_path = tmp_path / "summary.yml" @@ -202,6 +231,8 @@ def test_opt_fully_and_vectors(tmp_path): DATA_PATH / "NaCl-deformed.cif", "--arch", "mace_mp", + "--device", + device, "--opt-cell-fully", "--opt-cell-lengths", "--out", @@ -228,8 +259,12 @@ def test_opt_fully_and_vectors(tmp_path): assert atoms.cell.cellpar() == pytest.approx(expected) -def test_vectors_not_opt_fully(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_vectors_not_opt_fully(tmp_path, device): """Test passing --opt-cell-lengths without --opt-cell-fully.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + results_path = tmp_path / "NaCl-opt.extxyz" log_path = tmp_path / "test.log" summary_path = tmp_path / "summary.yml" @@ -242,6 +277,8 @@ def test_vectors_not_opt_fully(tmp_path): DATA_PATH / "NaCl.cif", "--arch", "mace_mp", + "--device", + device, "--out", results_path, "--opt-cell-lengths", @@ -259,9 +296,13 @@ def test_vectors_not_opt_fully(tmp_path): test_data = ["--opt-cell-lengths", "--opt-cell-fully"] +@pytest.mark.parametrize("device", ["cpu", "cuda"]) @pytest.mark.parametrize("option", test_data) -def test_scalar_pressure(option, tmp_path): +def test_scalar_pressure(option, tmp_path, device): """Test passing --pressure with --opt-cell-lengths.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + results_path = tmp_path / "NaCl-opt.extxyz" log_path = tmp_path / "test.log" summary_path = tmp_path / "summary.yml" @@ -274,6 +315,8 @@ def test_scalar_pressure(option, tmp_path): DATA_PATH / "NaCl-deformed.cif", "--arch", "mace_mp", + "--device", + device, "--out", results_path, option, @@ -317,8 +360,12 @@ def test_opt_kwargs_traj(tmp_path): assert isinstance(result.exception, ValueError) -def test_restart(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_restart(tmp_path, device): """Test restarting geometry optimization.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + data_path = DATA_PATH / "NaCl-deformed.cif" restart_path = tmp_path / "NaCl-res.pkl" results_path = tmp_path / "NaCl-opt.extxyz" @@ -335,6 +382,8 @@ def test_restart(tmp_path): data_path, "--arch", "mace_mp", + "--device", + device, "--out", results_path, "--minimize-kwargs", @@ -377,8 +426,12 @@ def test_restart(tmp_path): assert final_energy < intermediate_energy -def test_summary(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_summary(tmp_path, device): """Test summary file can be read correctly.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + results_path = tmp_path / "NaCl-results.extxyz" log_path = tmp_path / "test.log" summary_path = tmp_path / "summary.yml" @@ -391,6 +444,8 @@ def test_summary(tmp_path): DATA_PATH / "NaCl.cif", "--arch", "mace_mp", + "--device", + device, "--out", results_path, "--log", @@ -429,8 +484,12 @@ def test_summary(tmp_path): check_output_files(geomopt_summary, output_files) -def test_config(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_config(tmp_path, device): """Test passing a config file with opt_kwargs.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + results_path = tmp_path / "NaCl-results.extxyz" log_path = tmp_path / "test.log" summary_path = tmp_path / "summary.yml" @@ -443,6 +502,8 @@ def test_config(tmp_path): DATA_PATH / "NaCl.cif", "--arch", "mace_mp", + "--device", + device, "--out", results_path, "--log", @@ -481,8 +542,12 @@ def test_invalid_config(): assert isinstance(result.exception, ValueError) -def test_const_volume(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_const_volume(tmp_path, device): """Test setting constant volume with --opt-cell-fully.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + results_path = tmp_path / "NaCl-opt.extxyz" log_path = tmp_path / "test.log" summary_path = tmp_path / "summary.yml" @@ -497,6 +562,8 @@ def test_const_volume(tmp_path): DATA_PATH / "NaCl-deformed.cif", "--arch", "mace_mp", + "--device", + device, "--out", results_path, "--opt-cell-fully", @@ -512,8 +579,12 @@ def test_const_volume(tmp_path): assert_log_contains(log_path, includes=["constant_volume: True"]) -def test_optimizer_str(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_optimizer_str(tmp_path, device): """Test setting optimizer function.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + results_path = tmp_path / "NaCl-opt.extxyz" log_path = tmp_path / "test.log" summary_path = tmp_path / "summary.yml" @@ -526,6 +597,8 @@ def test_optimizer_str(tmp_path): DATA_PATH / "NaCl-deformed.cif", "--arch", "mace_mp", + "--device", + device, "--out", results_path, "--optimizer", @@ -543,8 +616,12 @@ def test_optimizer_str(tmp_path): ) -def test_filter_str(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_filter_str(tmp_path, device): """Test setting filter.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + results_path = tmp_path / "NaCl-opt.extxyz" log_path = tmp_path / "test.log" summary_path = tmp_path / "summary.yml" @@ -557,6 +634,8 @@ def test_filter_str(tmp_path): DATA_PATH / "NaCl-deformed.cif", "--arch", "mace_mp", + "--device", + device, "--out", results_path, "--opt-cell-fully", @@ -603,9 +682,13 @@ def test_filter_str_error(tmp_path): assert isinstance(result.exception, ValueError) +@pytest.mark.parametrize("device", ["cpu", "cuda"]) @pytest.mark.parametrize("read_kwargs", ["{'index': 1}", "{}"]) -def test_valid_traj_input(read_kwargs, tmp_path): +def test_valid_traj_input(read_kwargs, tmp_path, device): """Test valid trajectory input structure handled.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + results_path = tmp_path / "benezene-traj.extxyz" log_path = tmp_path / "test.log" summary_path = tmp_path / "summary.yml" @@ -618,6 +701,8 @@ def test_valid_traj_input(read_kwargs, tmp_path): DATA_PATH / "benzene-traj.xyz", "--arch", "mace_mp", + "--device", + device, "--out", results_path, "--read-kwargs", @@ -661,8 +746,12 @@ def test_invalid_traj_input(tmp_path): assert isinstance(result.exception, ValueError) -def test_reuse_output(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_reuse_output(tmp_path, device): """Test using geomopt output as new input.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + results_path_1 = tmp_path / "test" / "NaCl-opt-1.extxyz" results_path_2 = tmp_path / "test" / "NaCl-opt-2.extxyz" log_path = tmp_path / "test.log" @@ -676,6 +765,8 @@ def test_reuse_output(tmp_path): DATA_PATH / "NaCl-deformed.cif", "--arch", "mace_mp", + "--device", + device, "--fmax", 0.1, "--out", @@ -699,6 +790,8 @@ def test_reuse_output(tmp_path): results_path_1, "--arch", "mace_mp", + "--device", + device, "--fmax", 0.01, "--out", @@ -716,8 +809,12 @@ def test_reuse_output(tmp_path): assert results_1.positions != pytest.approx(results_2.positions) -def test_symmetrize(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_symmetrize(tmp_path, device): """Test symmetrizing final structure.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + results_path_1 = tmp_path / "test" / "NaCl-opt-1.extxyz" results_path_2 = tmp_path / "test" / "NaCl-opt-2.extxyz" log_path = tmp_path / "test.log" @@ -731,6 +828,8 @@ def test_symmetrize(tmp_path): DATA_PATH / "NaCl-deformed.cif", "--arch", "mace_mp", + "--device", + device, "--fmax", 0.001, "--out", @@ -752,6 +851,8 @@ def test_symmetrize(tmp_path): results_path_1, "--arch", "mace_mp", + "--device", + device, "--fmax", 0.001, "--out", @@ -778,8 +879,12 @@ def test_symmetrize(tmp_path): assert results_2.cell.cellpar() == pytest.approx(expected) -def test_no_carbon(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_no_carbon(tmp_path, device): """Test disabling carbon tracking.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + results_path = tmp_path / "NaCl-results.extxyz" log_path = tmp_path / "test.log" summary_path = tmp_path / "summary.yml" @@ -792,6 +897,8 @@ def test_no_carbon(tmp_path): DATA_PATH / "NaCl.cif", "--arch", "mace_mp", + "--device", + device, "--out", results_path, "--log", @@ -809,8 +916,12 @@ def test_no_carbon(tmp_path): assert "emissions" not in geomopt_summary -def test_units(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_units(tmp_path, device): """Test correct units are saved.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + results_path = tmp_path / "NaCl-opt.extxyz" log_path = tmp_path / "test.log" summary_path = tmp_path / "summary.yml" @@ -823,6 +934,8 @@ def test_units(tmp_path): DATA_PATH / "NaCl.cif", "--arch", "mace_mp", + "--device", + device, "--out", results_path, "--log", @@ -840,8 +953,12 @@ def test_units(tmp_path): assert atoms.info["units"][prop] == units -def test_file_prefix(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_file_prefix(tmp_path, device): """Test file prefix creates directories and affects all files.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + file_prefix = tmp_path / "test/test" result = runner.invoke( app, @@ -851,6 +968,8 @@ def test_file_prefix(tmp_path): DATA_PATH / "NaCl.cif", "--arch", "mace_mp", + "--device", + device, "--file-prefix", file_prefix, "--write-traj", @@ -890,8 +1009,12 @@ def test_traj_kwargs_no_write(tmp_path): assert "trajectory writing not enabled" in result.exception.args[0] -def test_model(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_model(tmp_path, device): """Test model passed correctly.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + file_prefix = tmp_path / "NaCl" results_path = tmp_path / "NaCl-opt.extxyz" log_path = tmp_path / "test.log" @@ -904,6 +1027,8 @@ def test_model(tmp_path): DATA_PATH / "NaCl.cif", "--arch", "mace_mp", + "--device", + device, "--model", MACE_PATH, "--log", @@ -938,8 +1063,12 @@ def test_missing_arch(tmp_path): assert "Missing option" in result.stderr -def test_info(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_info(tmp_path, device): """Test info written to output structure.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + file_prefix = tmp_path / "NaCl" results_path = tmp_path / "NaCl-opt.extxyz" @@ -951,6 +1080,8 @@ def test_info(tmp_path): DATA_PATH / "NaCl.cif", "--arch", "mace_mp", + "--device", + device, "--file-prefix", file_prefix, "--no-tracker", diff --git a/tests/test_md.py b/tests/test_md.py index 23290d7c..17c8dbfb 100644 --- a/tests/test_md.py +++ b/tests/test_md.py @@ -9,6 +9,7 @@ import ase.md.nose_hoover_chain import numpy as np import pytest +import torch from janus_core.calculations.md import NPH, NPT, NVE, NVT, NVT_CSVR, NVT_NH from janus_core.calculations.single_point import SinglePoint @@ -58,13 +59,18 @@ ) +@pytest.mark.parametrize("device", ["cpu", "cuda"]) @pytest.mark.parametrize("ensemble, expected", test_data) -def test_init(ensemble, expected): +def test_init(ensemble, expected, device): """Test initialising ensembles.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + single_point = SinglePoint( struct=DATA_PATH / "NaCl.cif", arch="mace", model=MODEL_PATH, + device=device, ) dyn = ensemble( struct=single_point.struct, @@ -72,7 +78,8 @@ def test_init(ensemble, expected): assert dyn.ensemble == expected -def test_deprecation_npt_mtk(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_deprecation_npt_mtk(tmp_path, device): """Test FutureWarning raise for ensemble ntp-mtk.""" with chdir(tmp_path): with pytest.warns( @@ -115,52 +122,12 @@ def test_deprecation_npt_mtk(tmp_path): assert len(lines) == 6 -def test_npt(tmp_path): - """Test NPT molecular dynamics.""" - with chdir(tmp_path): - results_dir = Path("janus_results") - restart_path_1 = results_dir / "Cl4Na4-npt-T300.0-p1.0-res-2.extxyz" - restart_path_2 = results_dir / "Cl4Na4-npt-T300.0-p1.0-res-4.extxyz" - restart_final = results_dir / "Cl4Na4-npt-T300.0-p1.0-final.extxyz" - traj_path = results_dir / "Cl4Na4-npt-T300.0-p1.0-traj.extxyz" - stats_path = results_dir / "Cl4Na4-npt-T300.0-p1.0-stats.dat" - - single_point = SinglePoint( - struct=DATA_PATH / "NaCl.cif", - arch="mace", - model=MODEL_PATH, - ) - npt = NPT( - struct=single_point.struct, - pressure=1.0, - temp=300.0, - steps=4, - traj_every=1, - restart_every=2, - stats_every=1, - ) - - npt.run() - restart_atoms_1 = read(restart_path_1) - assert isinstance(restart_atoms_1, Atoms) - restart_atoms_2 = read(restart_path_2) - assert isinstance(restart_atoms_2, Atoms) - restart_atoms_final = read(restart_final) - assert isinstance(restart_atoms_final, Atoms) - traj = read(traj_path, index=":") - assert all(isinstance(image, Atoms) for image in traj) - # Includes step 0 - assert len(traj) == 5 - - with open(stats_path, encoding="utf8") as stats_file: - lines = stats_file.readlines() - assert "Target_P [GPa]" in lines[0] and "Target_T [K]" in lines[0] - # Includes step 0 - assert len(lines) == 6 - - -def test_nvt_nh(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_nvt_nh(tmp_path, device): """Test NVT-Nosé–Hoover molecular dynamics.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + with chdir(tmp_path): results_dir = Path("janus_results") restart_path = results_dir / "Cl4Na4-nvt-nh-T300.0-res-3.extxyz" @@ -202,8 +169,12 @@ def test_nvt_nh(tmp_path): assert len(lines) == 5 -def test_nve(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_nve(tmp_path, device): """Test NVE molecular dynamics.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + file_prefix = tmp_path / "Cl4Na4-nve-T300.0" traj_path = tmp_path / "Cl4Na4-nve-T300.0-traj.extxyz" stats_path = tmp_path / "Cl4Na4-nve-T300.0-stats.dat" @@ -234,8 +205,12 @@ def test_nve(tmp_path): assert len(lines) == 6 -def test_nph(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_nph(tmp_path, device): """Test NPH molecular dynamics.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + with chdir(tmp_path): results_dir = Path("janus_results") restart_path = results_dir / "Cl4Na4-nph-T300.0-p0.0-res-2.extxyz" @@ -277,8 +252,12 @@ def test_nph(tmp_path): assert len(lines) == 3 -def test_nvt_csvr(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_nvt_csvr(tmp_path, device): """Test NVT CSVR molecular dynamics.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + with chdir(tmp_path): results_dir = Path("janus_results") restart_path_1 = results_dir / "NaCl-nvt-csvr-T300.0-res-2.extxyz" @@ -291,6 +270,7 @@ def test_nvt_csvr(tmp_path): struct=DATA_PATH / "NaCl.cif", arch="mace", model=MODEL_PATH, + device=device, temp=300.0, steps=4, traj_every=1, @@ -316,10 +296,14 @@ def test_nvt_csvr(tmp_path): assert len(lines) == 6 +@pytest.mark.parametrize("device", ["cpu", "cuda"]) @pytest.mark.skipif(MTK_IMPORT_FAILED, reason="Requires updated version of ASE") @pytest.mark.parametrize("mtk_flavour", ["iso", "aniso"]) -def test_npt_mtk(tmp_path, mtk_flavour): +def test_npt_mtk(tmp_path, device,mtk_flavour): """Test NPT MTK molecular dynamics.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + with chdir(tmp_path): results_dir = Path("janus_results") stem = f"NaCl-npt-mtk-{mtk_flavour}-T300.0-p0.0001" @@ -333,6 +317,7 @@ def test_npt_mtk(tmp_path, mtk_flavour): struct=DATA_PATH / "NaCl.cif", arch="mace", model=MODEL_PATH, + device=device, temp=300.0, pressure=0.0001, steps=4, @@ -365,8 +350,12 @@ def test_npt_mtk(tmp_path, mtk_flavour): assert len(lines) == 6 -def test_restart(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_restart(tmp_path, device): """Test restarting molecular dynamics simulation.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + file_prefix = tmp_path / "Cl4Na4-nvt-T300.0" traj_path = tmp_path / "Cl4Na4-nvt-T300.0-traj.extxyz" stats_path = tmp_path / "Cl4Na4-nvt-T300.0-stats.dat" @@ -399,6 +388,7 @@ def test_restart(tmp_path): restart=True, restart_auto=False, file_prefix=file_prefix, + device=device, ) nvt_restart.run() assert nvt_restart.offset == 4 @@ -415,8 +405,12 @@ def test_restart(tmp_path): assert len(traj) == 9 -def test_restart_with_d3(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_restart_with_d3(tmp_path, device): """Test restarting molecular dynamics simulation (with D3).""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + file_prefix = tmp_path / "Cl4Na4-nvt-T300.0" traj_path = tmp_path / "Cl4Na4-nvt-T300.0-traj.extxyz" stats_path = tmp_path / "Cl4Na4-nvt-T300.0-stats.dat" @@ -425,6 +419,7 @@ def test_restart_with_d3(tmp_path): struct=DATA_PATH / "NaCl.cif", arch="mace", model=MODEL_PATH, + device=device, calc_kwargs={"dispersion": True}, ) nvt = NVT( @@ -451,6 +446,7 @@ def test_restart_with_d3(tmp_path): restart=True, restart_auto=False, file_prefix=file_prefix, + device=device, calc_kwargs={"dispersion": True}, ) nvt_restart.run() @@ -468,8 +464,12 @@ def test_restart_with_d3(tmp_path): assert len(traj) == 9 -def test_minimize(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_minimize(tmp_path, device): """Test geometry optimzation before dynamics.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + file_prefix = tmp_path / "Cl4Na4-nvt-T300.0" single_point = SinglePoint( @@ -487,6 +487,7 @@ def test_minimize(tmp_path): stats_every=1, minimize=True, file_prefix=file_prefix, + device=device, ) nvt.run() @@ -494,8 +495,12 @@ def test_minimize(tmp_path): assert final_energy < init_energy -def test_reset_velocities(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_reset_velocities(tmp_path, device): """Test rescaling velocities before dynamics.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + file_prefix = tmp_path / "Cl4Na4-nvt-T300.0" single_point = SinglePoint( @@ -527,8 +532,12 @@ def test_reset_velocities(tmp_path): assert final_momentum == pytest.approx(0, abs=1e-10) -def test_no_reset_velocities(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_no_reset_velocities(tmp_path, device): """Test not resetting velocities before dynamics.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + struct_file = DATA_PATH / "lj-traj.xyz" file_prefix = tmp_path / "md" @@ -548,6 +557,7 @@ def test_no_reset_velocities(tmp_path): stats_every=1, rescale_velocities=False, file_prefix=file_prefix, + device=device, ) nvt.run() @@ -558,8 +568,12 @@ def test_no_reset_velocities(tmp_path): assert final_velocities == pytest.approx(init_velocities) -def test_reset_velocities_zero_kinetic_energy(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_reset_velocities_zero_kinetic_energy(tmp_path, device): """Test velocities are set for zero kinetic energy.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + file_prefix = tmp_path / "Cl4Na4-nvt-T300.0" log_file = tmp_path / "nvt.log" @@ -582,6 +596,7 @@ def test_reset_velocities_zero_kinetic_energy(tmp_path): stats_every=1, rescale_velocities=False, file_prefix=file_prefix, + device=device, log_kwargs={"filename": log_file}, ) nvt.run() @@ -603,8 +618,12 @@ def test_reset_velocities_zero_kinetic_energy(tmp_path): assert final_velocities != pytest.approx(init_velocities) -def test_remove_rot(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_remove_rot(tmp_path, device): """Test removing rotation before dynamics.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + file_prefix = tmp_path / "Cl4Na4-nvt-T300.0" single_point = SinglePoint( @@ -631,6 +650,7 @@ def test_remove_rot(tmp_path): remove_rot=True, equil_steps=1, file_prefix=file_prefix, + device=device, ) nvt.run() @@ -645,8 +665,12 @@ def test_remove_rot(tmp_path): assert final_ang_mom == pytest.approx(0) -def test_traj_start(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_traj_start(tmp_path, device): """Test starting trajectory after n steps.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + file_prefix = tmp_path / "Cl4Na4-nvt-T300.0" traj_path = tmp_path / "Cl4Na4-nvt-T300.0-traj.extxyz" @@ -687,8 +711,12 @@ def test_traj_start(tmp_path): assert len(traj) == 2 -def test_minimize_every(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_minimize_every(tmp_path, device): """Test setting minimize_every.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + file_prefix = tmp_path / "Cl4Na4-nvt-T300.0" log_file = tmp_path / "nvt.log" @@ -718,8 +746,12 @@ def test_minimize_every(tmp_path): ) -def test_rescale_every(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_rescale_every(tmp_path, device): """Test setting minimize_every.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + file_prefix = tmp_path / "Cl4Na4-nvt-T300.0" log_file = tmp_path / "nvt.log" @@ -760,8 +792,12 @@ def test_rescale_every(tmp_path): ) -def test_rotate_restart(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_rotate_restart(tmp_path, device): """Test setting rotate_restart.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + file_prefix = tmp_path / "NaCl-nvt" restart_path_1 = tmp_path / "NaCl-nvt-res-1.extxyz" restart_path_2 = tmp_path / "NaCl-nvt-res-2.extxyz" @@ -781,6 +817,7 @@ def test_rotate_restart(tmp_path): restart_every=1, restarts_to_keep=2, file_prefix=file_prefix, + device=device, ) nvt.run() @@ -789,14 +826,18 @@ def test_rotate_restart(tmp_path): assert restart_path_3.exists() -def test_atoms_struct(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_atoms_struct(tmp_path, device): """Test passing a structure with an attached calculator.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + file_prefix = tmp_path / "Cl4Na4-nvt-T300.0" traj_path = tmp_path / "Cl4Na4-nvt-T300.0-traj.extxyz" stats_path = tmp_path / "Cl4Na4-nvt-T300.0-stats.dat" struct = read(DATA_PATH / "NaCl.cif") - struct.calc = choose_calculator(arch="mace_mp", model=MODEL_PATH) + struct.calc = choose_calculator(arch="mace_mp", model=MODEL_PATH, device=device) nvt = NVT( struct=struct, @@ -805,6 +846,7 @@ def test_atoms_struct(tmp_path): stats_every=1, traj_every=1, file_prefix=file_prefix, + device=device, ) nvt.run() @@ -818,9 +860,13 @@ def test_atoms_struct(tmp_path): assert len(lines) == 6 +@pytest.mark.parametrize("device", ["cpu", "cuda"]) @pytest.mark.parametrize("ensemble, tag", test_data) -def test_stats(tmp_path, ensemble, tag): +def test_stats(tmp_path, ensemble, tag, device): """Test stats file has correct structure and entries for all ensembles.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + file_prefix = tmp_path / tag / "NaCl" single_point = SinglePoint( struct=DATA_PATH / "NaCl.cif", @@ -832,6 +878,7 @@ def test_stats(tmp_path, ensemble, tag): steps=2, stats_every=1, file_prefix=file_prefix, + device=device, ) md.run() @@ -845,9 +892,13 @@ def test_stats(tmp_path, ensemble, tag): assert stat_data.units[etot_index] == "eV" +@pytest.mark.parametrize("device", ["cpu", "cuda"]) @pytest.mark.parametrize("ensemble", ensembles_with_thermostat) -def test_heating(tmp_path, capsys, ensemble): +def test_heating(tmp_path, capsys, ensemble, device): """Test heating with no MD.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + file_prefix = tmp_path / "NaCl-heating" final_file = tmp_path / "NaCl-heating-final.extxyz" log_file = tmp_path / "NaCl.log" @@ -888,9 +939,13 @@ def test_heating(tmp_path, capsys, ensemble): assert "2/2" in capsys.readouterr().out +@pytest.mark.parametrize("device", ["cpu", "cuda"]) @pytest.mark.parametrize("ensemble", ensembles_without_thermostat) -def test_no_thermostat_heating(tmp_path, ensemble): +def test_no_thermostat_heating(tmp_path, ensemble, device): """Test that temperature ramp with no thermostat throws an error.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + file_prefix = tmp_path / "NaCl-heating" final_file = tmp_path / "NaCl-heating-final.extxyz" log_file = tmp_path / "NaCl.log" @@ -913,14 +968,19 @@ def test_no_thermostat_heating(tmp_path, ensemble): temp_step=20, temp_time=0.5, log_kwargs={"filename": log_file}, + device=device, ) md.run() assert not final_file.exists() +@pytest.mark.parametrize("device", ["cpu", "cuda"]) @pytest.mark.parametrize("ensemble", ensembles_with_thermostat) -def test_noramp_heating(tmp_path, ensemble): +def test_noramp_heating(tmp_path, ensemble, device): """Test ValueError is thrown for invalid temperature ramp.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + file_prefix = tmp_path / "NaCl-heating" single_point = SinglePoint( @@ -936,12 +996,17 @@ def test_noramp_heating(tmp_path, ensemble): temp_start=10, temp_end=10, temp_step=20, + device=device, ) +@pytest.mark.parametrize("device", ["cpu", "cuda"]) @pytest.mark.parametrize("ensemble", ensembles_with_thermostat) -def test_heating_md(tmp_path, capsys, ensemble): +def test_heating_md(tmp_path, capsys, ensemble, device): """Test heating followed by MD.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + file_prefix = tmp_path / "NaCl-heating" stats_path = tmp_path / "NaCl-heating-stats.dat" log_file = tmp_path / "NaCl.log" @@ -992,13 +1057,18 @@ def test_heating_md(tmp_path, capsys, ensemble): assert "9/9" in capsys.readouterr().out -def test_heating_restart(tmp_path, capsys): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_heating_restart(tmp_path, capsys, device): """Test restarting during temperature ramp.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + file_prefix = tmp_path / "NaCl-heating" md = NVT( struct=DATA_PATH / "NaCl.cif", arch="mace_mp", + device=device, temp=10.0, steps=3, traj_every=100, @@ -1013,6 +1083,7 @@ def test_heating_restart(tmp_path, capsys): md_restart = NVT( struct=DATA_PATH / "NaCl.cif", arch="mace_mp", + device=device, temp=30.0, steps=2, traj_every=100, @@ -1038,8 +1109,12 @@ def test_heating_restart(tmp_path, capsys): assert "8/8" in capsys.readouterr().out -def test_heating_files(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_heating_files(tmp_path, device): """Test default heating file names.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + with chdir(tmp_path): results_dir = Path("./janus_results") traj_heating_path = results_dir / "Cl4Na4-nvt-T10-T20-traj.extxyz" @@ -1050,6 +1125,7 @@ def test_heating_files(tmp_path): struct=DATA_PATH / "NaCl.cif", arch="mace", model=MODEL_PATH, + device=device, ) nvt = NVT( struct=single_point.struct, @@ -1061,6 +1137,7 @@ def test_heating_files(tmp_path): temp_end=20, temp_step=10, temp_time=2, + device=device, ) nvt.run() @@ -1079,8 +1156,12 @@ def test_heating_files(tmp_path): assert stats.data[2, 16] == 20.0 -def test_heating_md_files(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_heating_md_files(tmp_path, device): """Test default heating files when also running md.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + with chdir(tmp_path): results_dir = Path("janus_results") traj_heating_path = results_dir / "Cl4Na4-nvt-T10-T20-T25.0-traj.extxyz" @@ -1091,6 +1172,7 @@ def test_heating_md_files(tmp_path): struct=DATA_PATH / "NaCl.cif", arch="mace", model=MODEL_PATH, + device=device, ) nvt = NVT( struct=single_point.struct, @@ -1102,6 +1184,7 @@ def test_heating_md_files(tmp_path): temp_end=20, temp_step=10, temp_time=2, + device=device, ) nvt.run() @@ -1120,14 +1203,19 @@ def test_heating_md_files(tmp_path): assert stats.data[3, 16] == 25 -def test_ramp_negative(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_ramp_negative(tmp_path, device): """Test ValueError is thrown for negative values in temperature ramp.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + file_prefix = tmp_path / "NaCl-heating" single_point = SinglePoint( struct=DATA_PATH / "NaCl.cif", arch="mace", model=MODEL_PATH, + device=device, ) with pytest.raises(ValueError): @@ -1138,6 +1226,7 @@ def test_ramp_negative(tmp_path): temp_end=10, temp_step=20, temp_time=10, + device=device, ) with pytest.raises(ValueError): NVT( @@ -1147,11 +1236,16 @@ def test_ramp_negative(tmp_path): temp_end=-5, temp_step=20, temp_time=10, + device=device, ) -def test_cooling(tmp_path, capsys): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_cooling(tmp_path, capsys, device): """Test cooling with no MD.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + file_prefix = tmp_path / "NaCl-cooling" final_path = tmp_path / "NaCl-cooling-final.extxyz" stats_cooling_path = tmp_path / "NaCl-cooling-stats.dat" @@ -1176,6 +1270,7 @@ def test_cooling(tmp_path, capsys): temp_time=1, log_kwargs={"filename": log_file}, enable_progress_bar=True, + device=device, ) nvt.run() assert_log_contains( @@ -1200,14 +1295,18 @@ def test_cooling(tmp_path, capsys): assert "2/2" in capsys.readouterr().out -def test_heating_too_short(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_heating_too_short(tmp_path, device): """Test that heating with steps shorter than 1 timestep fails.""" - file_prefix = tmp_path / "NaCl-heating" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + file_prefix = tmp_path / "NaCl-heating" single_point = SinglePoint( struct=DATA_PATH / "NaCl.cif", arch="mace", model=MODEL_PATH, + device=device, ) with pytest.raises(ValueError, match="Temperature ramp step time"): @@ -1222,12 +1321,17 @@ def test_heating_too_short(tmp_path): temp_end=20.0, temp_step=20, temp_time=0.5, + device=device, ) nvt.run() -def test_ensemble_kwargs(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_ensemble_kwargs(tmp_path, device): """Test setting integrator kwargs.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + log_file = tmp_path / "nvt.log" single_point = SinglePoint( @@ -1240,6 +1344,7 @@ def test_ensemble_kwargs(tmp_path): struct=single_point.struct, ensemble_kwargs={"mask": (0, 1, 0)}, log_kwargs={"filename": log_file}, + device=device, ) expected_mask = [[False, False, False], [False, True, False], [False, False, False]] @@ -1264,8 +1369,12 @@ def test_invalid_struct(): ) -def test_logging(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_logging(tmp_path, device): """Test attaching logger to NVT and emissions are saved to info.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + log_file = tmp_path / "md.log" file_prefix = tmp_path / "md" @@ -1273,6 +1382,7 @@ def test_logging(tmp_path): struct=DATA_PATH / "NaCl.cif", arch="mace_mp", model=MODEL_PATH, + device=device, ) assert "emissions" not in single_point.struct.info @@ -1291,8 +1401,12 @@ def test_logging(tmp_path): assert single_point.struct.info["emissions"] > 0 -def test_auto_restart(tmp_path, capsys): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_auto_restart(tmp_path, capsys, device): """Test auto restarting simulation.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + with chdir(tmp_path): # Include T300.0 to test Path.stem vs Path.name final_path = "md-T300.0-final.extxyz" @@ -1317,6 +1431,7 @@ def test_auto_restart(tmp_path, capsys): final_file=final_path, restart_every=4, timestep=100, + device=device, ) nvt.run() @@ -1350,6 +1465,7 @@ def test_auto_restart(tmp_path, capsys): final_file=final_path, log_kwargs={"filename": log_file}, enable_progress_bar=True, + device=device, ) assert_log_contains(log_file, includes="Auto restart successful") @@ -1382,8 +1498,12 @@ def test_auto_restart(tmp_path, capsys): assert "7/7" in capsys.readouterr().out -def test_auto_restart_restart_stem(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_auto_restart_restart_stem(tmp_path, device): """Test auto restarting simulation with restart stem defined.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + file_prefix = tmp_path / "npt" traj_path = tmp_path / "npt-traj.extxyz" stats_path = tmp_path / "npt-stats.dat" @@ -1405,6 +1525,7 @@ def test_auto_restart_restart_stem(tmp_path): restart_stem=str(restart_stem), restart_every=4, timestep=10, + device=device, ) npt.run() @@ -1437,6 +1558,7 @@ def test_auto_restart_restart_stem(tmp_path): file_prefix=file_prefix, traj_every=1, log_kwargs={"filename": log_file}, + device=device, ) assert_log_contains(log_file, includes="Auto restart successful") @@ -1464,8 +1586,12 @@ def test_auto_restart_restart_stem(tmp_path): assert len(final_traj) == 10 -def test_set_info(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_set_info(tmp_path, device): """Test info is set at correct frequency.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + file_prefix = tmp_path / "npt" traj_path = tmp_path / "npt-traj.extxyz" @@ -1483,6 +1609,7 @@ def test_set_info(tmp_path): file_prefix=file_prefix, seed=2024, traj_every=10, + device=device, ) npt.run() @@ -1491,9 +1618,13 @@ def test_set_info(tmp_path): assert final_struct.info["density"] == pytest.approx(2.120952627887493) +@pytest.mark.parametrize("device", ["cpu", "cuda"]) @pytest.mark.parametrize("ensemble, tag", test_data) -def test_progress_bar_complete(tmp_path, capsys, ensemble, tag): +def test_progress_bar_complete(tmp_path, capsys, ensemble, tag, device): """Test progress bar completes in all ensembles.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + file_prefix = tmp_path / f"Cl4Na4-{tag}-T300.0" md = ensemble( @@ -1503,6 +1634,7 @@ def test_progress_bar_complete(tmp_path, capsys, ensemble, tag): steps=2, file_prefix=file_prefix, enable_progress_bar=True, + device=device, ) md.run() diff --git a/tests/test_md_cli.py b/tests/test_md_cli.py index ac839606..439e76e4 100644 --- a/tests/test_md_cli.py +++ b/tests/test_md_cli.py @@ -9,6 +9,7 @@ import ase.md.nose_hoover_chain import numpy as np import pytest +import torch from typer.testing import CliRunner import yaml @@ -62,9 +63,12 @@ def test_md_help(): ] +@pytest.mark.parametrize("device", ["cpu", "cuda"]) @pytest.mark.parametrize("ensemble", test_data) -def test_md(ensemble, tmp_path): +def test_md(ensemble, tmp_path, device): """Test all MD simulations are able to run.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") # Expected default file prefix for each ensemble file_prefix = { "nvt": "NaCl-nvt-T300.0-", @@ -120,6 +124,8 @@ def test_md(ensemble, tmp_path): " 'vaf_x': {'a': 'Velocity', 'points': 10," "'a_kwargs': {'components': ['x']}, 'b_kwargs': '.'}}" ), + "--device", + device, ], ) @@ -170,8 +176,11 @@ def test_md(ensemble, tmp_path): clear_log_handlers() -def test_log(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_log(tmp_path, device): """Test log correctly written for MD.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") file_prefix = tmp_path / "NaCl" stats_path = tmp_path / "NaCl-stats.dat" log_path = tmp_path / "NaCl-md-log.yml" @@ -192,6 +201,8 @@ def test_log(tmp_path): 1, "--file-prefix", file_prefix, + "--device", + device, ], ) assert result.exit_code == 0 @@ -219,8 +230,11 @@ def test_log(tmp_path): assert final_temp == pytest.approx(final_temp) -def test_seed(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_seed(tmp_path, device): """Test seed enables reproducable results for NVT.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") file_prefix = tmp_path / "nvt-T300" stats_path = tmp_path / "nvt-T300-stats.dat" @@ -291,8 +305,11 @@ def test_seed(tmp_path): assert stats_1 == stats_2 -def test_summary(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_summary(tmp_path, device): """Test summary file can be read correctly.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") file_prefix = tmp_path / "nve" stats_path = tmp_path / "nve-stats.dat" traj_path = tmp_path / "nve-traj.extxyz" @@ -354,8 +371,11 @@ def test_summary(tmp_path): check_output_files(summary=summary, output_files=output_files) -def test_config(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_config(tmp_path, device): """Test passing a config file.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") file_prefix = tmp_path / "nvt" log_path = tmp_path / "nvt-md-log.yml" summary_path = tmp_path / "nvt-md-summary.yml" @@ -377,6 +397,8 @@ def test_config(tmp_path): "--minimize", "--config", DATA_PATH / "md_config.yml", + "--device", + device, ], ) assert result.exit_code == 0 @@ -399,9 +421,12 @@ def test_config(tmp_path): assert_log_contains(log_path, includes=["hydrostatic_strain: True"]) +@pytest.mark.parametrize("device", ["cpu", "cuda"]) @pytest.mark.parametrize("ensemble", test_data) -def test_heating(tmp_path, ensemble): +def test_heating(tmp_path, ensemble, device): """Test heating before MD.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") file_prefix = tmp_path / "nvt-T300" result = runner.invoke( @@ -428,6 +453,8 @@ def test_heating(tmp_path, ensemble): 10, "--temp-time", 2, + "--device", + device, ], ) if ensemble in ("nve", "nph"): @@ -456,8 +483,11 @@ def test_invalid_config(): assert isinstance(result.exception, ValueError) -def test_ensemble_kwargs(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_ensemble_kwargs(tmp_path, device): """Test passing ensemble-kwargs to NPT.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") struct_path = DATA_PATH / "NaCl.cif" file_prefix = tmp_path / "test" / "md" final_path = tmp_path / "test" / "md-final.extxyz" @@ -483,6 +513,8 @@ def test_ensemble_kwargs(tmp_path): ensemble_kwargs, "--stats-every", 1, + "--device", + device, ], ) @@ -535,8 +567,11 @@ def test_invalid_ensemble_kwargs(tmp_path): assert isinstance(result.exception, TypeError) -def test_final_name(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_final_name(tmp_path, device): """Test specifying the final file name.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") file_prefix = tmp_path / "npt" stats_path = tmp_path / "npt-stats.dat" traj_path = tmp_path / "npt-traj.extxyz" @@ -562,6 +597,8 @@ def test_final_name(tmp_path): file_prefix, "--final-file", final_path, + "--device", + device, ], ) assert result.exit_code == 0 @@ -570,8 +607,11 @@ def test_final_name(tmp_path): assert final_path.exists() -def test_write_kwargs(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_write_kwargs(tmp_path, device): """Test passing write-kwargs.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") struct_path = DATA_PATH / "NaCl.cif" file_prefix = tmp_path / "md" final_path = tmp_path / "md-final.extxyz" @@ -599,6 +639,8 @@ def test_write_kwargs(tmp_path): write_kwargs, "--traj-every", 1, + "--device", + device, ], ) @@ -619,9 +661,12 @@ def test_write_kwargs(tmp_path): assert "mace_mp_energy" in final_atoms.info +@pytest.mark.parametrize("device", ["cpu", "cuda"]) @pytest.mark.parametrize("read_kwargs", ["{'index': 1}", "{}"]) -def test_valid_traj_input(read_kwargs, tmp_path): +def test_valid_traj_input(read_kwargs, tmp_path, device): """Test valid trajectory input structure handled.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") file_prefix = tmp_path / "traj" final_path = tmp_path / "traj-final.extxyz" @@ -641,6 +686,8 @@ def test_valid_traj_input(read_kwargs, tmp_path): 2, "--read-kwargs", read_kwargs, + "--device", + device, ], ) assert result.exit_code == 0 @@ -674,8 +721,11 @@ def test_invalid_traj_input(tmp_path): assert isinstance(result.exception, ValueError) -def test_minimize_kwargs_filename(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_minimize_kwargs_filename(tmp_path, device): """Test passing filename via minimize kwargs to MD.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") file_prefix = tmp_path / "test" / "md" opt_path = tmp_path / "test" / "test.extxyz" traj_path = tmp_path / "test" / "md-traj.extxyz" @@ -705,6 +755,8 @@ def test_minimize_kwargs_filename(tmp_path): "--minimize", "--minimize-kwargs", f"{{'write_kwargs': {{'filename': '{opt_path.as_posix()}'}}}}", + "--device", + device, ], ) assert result.exit_code == 0 @@ -736,8 +788,11 @@ def test_minimize_kwargs_filename(tmp_path): check_output_files(summary=summary, output_files=output_files) -def test_minimize_kwargs_write_results(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_minimize_kwargs_write_results(tmp_path, device): """Test passing write_results via minimize kwargs to MD.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") file_prefix = tmp_path / "test" / "md" opt_path = tmp_path / "test" / "md-opt.extxyz" traj_path = tmp_path / "test" / "md-traj.extxyz" @@ -765,6 +820,8 @@ def test_minimize_kwargs_write_results(tmp_path): "--minimize", "--minimize-kwargs", "{'write_results': True}", + "--device", + device, ], ) assert result.exit_code == 0 @@ -778,8 +835,11 @@ def test_minimize_kwargs_write_results(tmp_path): assert isinstance(atoms, Atoms) -def test_auto_restart(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_auto_restart(tmp_path, device): """Test auto restart with file_prefix.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") file_prefix = tmp_path / "md" traj_path = tmp_path / "md-traj.extxyz" stats_path = tmp_path / "md-stats.dat" @@ -813,6 +873,8 @@ def test_auto_restart(tmp_path): log_path, "--summary", summary_path, + "--device", + device, ], ) assert result.exit_code == 0 @@ -855,6 +917,8 @@ def test_auto_restart(tmp_path): log_path, "--summary", summary_path, + "--device", + device, ], ) assert result.exit_code == 0 @@ -876,8 +940,11 @@ def test_auto_restart(tmp_path): assert "7/7" in result.stdout -def test_no_carbon(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_no_carbon(tmp_path, device): """Test disabling carbon tracking.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") file_prefix = tmp_path / "nvt" summary_path = tmp_path / "nvt-md-summary.yml" @@ -896,6 +963,8 @@ def test_no_carbon(tmp_path): file_prefix, "--steps", 1, + "--device", + device, ], ) @@ -910,8 +979,11 @@ def test_no_carbon(tmp_path): @pytest.mark.parametrize("ensemble", ("nvt", "npt", "nvt-csvr")) @pytest.mark.parametrize("output_every", (1, 2)) @pytest.mark.parametrize("heating", (True, False)) -def test_consistent_stats_traj(tmp_path, ensemble, output_every, heating): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_consistent_stats_traj(tmp_path, ensemble, output_every, heating, device): """Test data saved to statistics is consistent with trajectory info.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") file_prefix = tmp_path / ensemble stats_path = tmp_path / f"{ensemble}-stats.dat" traj_path = tmp_path / f"{ensemble}-traj.extxyz" @@ -933,6 +1005,8 @@ def test_consistent_stats_traj(tmp_path, ensemble, output_every, heating): output_every, "--traj-every", output_every, + "--device", + device, ] if heating: @@ -970,8 +1044,11 @@ def test_consistent_stats_traj(tmp_path, ensemble, output_every, heating): assert target_temps == pytest.approx(stats_target_temps, rel=1e5) -def test_no_progress(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_no_progress(tmp_path, device): """Test disabling progress bar.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") file_prefix = tmp_path / "nvt" result = runner.invoke( @@ -990,6 +1067,8 @@ def test_no_progress(tmp_path): "--steps", 2, "--no-progress-bar", + "--device", + device, ], ) @@ -998,8 +1077,11 @@ def test_no_progress(tmp_path): assert "2/2" not in result.output -def test_model(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_model(tmp_path, device): """Test model passed correctly.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") file_prefix = tmp_path / "NaCl" results_path = tmp_path / "NaCl-final.extxyz" log_path = tmp_path / "test.log" @@ -1023,6 +1105,8 @@ def test_model(tmp_path): "--file-prefix", file_prefix, "--no-tracker", + "--device", + device, ], ) assert result.exit_code == 0 @@ -1052,8 +1136,11 @@ def test_missing_arch(tmp_path): assert "Missing option" in result.stderr -def test_info(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_info(tmp_path, device): """Test info written to output structures.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") file_prefix = tmp_path / "NaCl" traj_path = tmp_path / "NaCl-traj.extxyz" final_path = tmp_path / "NaCl-final.extxyz" @@ -1078,6 +1165,8 @@ def test_info(tmp_path): "--file-prefix", file_prefix, "--no-tracker", + "--device", + device, ], ) assert result.exit_code == 0 diff --git a/tests/test_mlip_calculators.py b/tests/test_mlip_calculators.py index 58ac3000..244c6da0 100644 --- a/tests/test_mlip_calculators.py +++ b/tests/test_mlip_calculators.py @@ -7,6 +7,7 @@ from zipfile import BadZipFile import pytest +import torch from janus_core.helpers.mlip_calculators import add_dispersion, choose_calculator from tests.utils import skip_extras @@ -63,48 +64,52 @@ ) +@pytest.mark.parametrize("device", ["cpu", "cuda"]) @pytest.mark.parametrize( - "arch, device, kwargs", + "arch, kwargs", [ - ("chgnet", "cpu", {}), - ("chgnet", "cpu", {"model": "0.2.0"}), - ("chgnet", "cpu", {"model": CHGNET_PATH}), - ("chgnet", "cpu", {"model": CHGNET_MODEL}), - ("dpa3", "cpu", {"model": DPA3_PATH}), - ("grace", "cpu", {}), - ("grace", "cpu", {"model": "GRACE-1L-MP-r6"}), - ("mace", "cpu", {"model": MACE_MP_PATH}), - ("mace", "cpu", {"model_paths": MACE_MP_PATH}), - ("mace_mp", "cpu", {}), - ("mace_mp", "cpu", {"model": "small"}), - ("mace_mp", "cpu", {"model": MACE_MP_PATH}), - ("mace_off", "cpu", {}), - ("mace_off", "cpu", {"model": "medium"}), - ("mace_off", "cpu", {"model": MACE_OFF_PATH}), - ("mace_omol", "cpu", {}), - ("mace_omol", "cpu", {"model": "extra_large"}), - ("mattersim", "cpu", {}), - ("mattersim", "cpu", {"model": "mattersim-v1.0.0-1m"}), - ("nequip", "cpu", {"model": NEQUIP_PATH}), - ("orb", "cpu", {}), - ("orb", "cpu", {"model": ORB_MODEL}), - ("upet", "cpu", {}), - ("upet", "cpu", {"model": PET_MAD_CHECKPOINT}), - ("upet", "cpu", {"checkpoint_path": PET_MAD_CHECKPOINT}), - ("sevennet", "cpu", {"model": SEVENNET_PATH}), - ("sevennet", "cpu", {"path": SEVENNET_PATH}), - ("sevennet", "cpu", {}), - ("sevennet", "cpu", {"model": "sevennet-0"}), - ("fairchem", "cpu", {"model": UMA_LABEL}), - ("fairchem", "cpu", {"model_name": UMA_LABEL}), - ("fairchem", "cpu", {"model": UMA_PREDICT_UNIT}), - ("fairchem", "cpu", {"predict_unit": UMA_PREDICT_UNIT}), + ("chgnet", {}), + ("chgnet", {"model": "0.2.0"}), + ("chgnet", {"model": CHGNET_PATH}), + ("chgnet", {"model": CHGNET_MODEL}), + ("dpa3", {"model": DPA3_PATH}), + ("grace", {}), + ("grace", {"model": "GRACE-1L-MP-r6"}), + ("mace", {"model": MACE_MP_PATH}), + ("mace", {"model_paths": MACE_MP_PATH}), + ("mace_mp", {}), + ("mace_mp", {"model": "small"}), + ("mace_mp", {"model": MACE_MP_PATH}), + ("mace_off", {}), + ("mace_off", {"model": "medium"}), + ("mace_off", {"model": MACE_OFF_PATH}), + ("mace_omol", {}), + ("mace_omol", {"model": "extra_large"}), + ("mattersim", {}), + ("mattersim", {"model": "mattersim-v1.0.0-1m"}), + ("nequip", {"model": NEQUIP_PATH}), + ("orb", {}), + ("orb", {"model": ORB_MODEL}), + ("upet", {}), + ("upet", {"model": PET_MAD_CHECKPOINT}), + ("upet", {"checkpoint_path": PET_MAD_CHECKPOINT}), + ("sevennet", {"model": SEVENNET_PATH}), + ("sevennet", {"path": SEVENNET_PATH}), + ("sevennet", {}), + ("sevennet", {"model": "sevennet-0"}), + ("fairchem", {"model": UMA_LABEL}), + ("fairchem", {"model_name": UMA_LABEL}), + ("fairchem", {"model": UMA_PREDICT_UNIT}), + ("fairchem", {"predict_unit": UMA_PREDICT_UNIT}), ], ) def test_mlips(arch, device, kwargs): """Test calculators can be configured.""" skip_extras(arch) + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + try: calculator = choose_calculator(arch=arch, device=device, **kwargs) assert calculator.parameters["version"] is not None diff --git a/tests/test_neb.py b/tests/test_neb.py index c5c09f3f..59b0b97b 100644 --- a/tests/test_neb.py +++ b/tests/test_neb.py @@ -7,6 +7,7 @@ from ase.io import read, write import pytest +import torch from janus_core.calculations.neb import NEB from janus_core.calculations.single_point import SinglePoint @@ -38,8 +39,12 @@ def LFPO_end_b(LFPO): return struct -def test_neb(tmp_path, LFPO_start_b, LFPO_end_b): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_neb(tmp_path, LFPO_start_b, LFPO_end_b, device): """Test NEB.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + # Write initial and final structures init_struct = tmp_path / "init_struct.cif" final_struct = tmp_path / "final_struct.cif" @@ -54,6 +59,7 @@ def test_neb(tmp_path, LFPO_start_b, LFPO_end_b): n_images=5, neb_kwargs={"method": "aseneb"}, file_prefix=tmp_path / "LFPO", + device=device, ) neb.run() @@ -63,18 +69,24 @@ def test_neb(tmp_path, LFPO_start_b, LFPO_end_b): assert neb.results["delta_E"] == pytest.approx(-3.0149328722473e-07) -def test_neb_pymatgen(tmp_path, LFPO_start_b, LFPO_end_b): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_neb_pymatgen(tmp_path, LFPO_start_b, LFPO_end_b, device): """Test pymatgen interpolation.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + file_prefix = tmp_path / "LFPO" single_point_start = SinglePoint( struct=LFPO_start_b, arch="mace", model=MODEL_PATH, + device=device, ) single_point_end = SinglePoint( struct=LFPO_end_b, arch="mace_mp", model=MODEL_PATH, + device=device, ) neb = NEB( @@ -87,6 +99,7 @@ def test_neb_pymatgen(tmp_path, LFPO_start_b, LFPO_end_b): fmax=4, neb_kwargs={"method": "aseneb"}, file_prefix=file_prefix, + device=device, ) neb.run() @@ -97,13 +110,19 @@ def test_neb_pymatgen(tmp_path, LFPO_start_b, LFPO_end_b): assert neb.results["max_force"] == pytest.approx(2.533305796378263) -def test_set_calc(tmp_path, LFPO_start_b, LFPO_end_b): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_set_calc(tmp_path, LFPO_start_b, LFPO_end_b, device): """Test setting the calculators explicitly.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + file_prefix = tmp_path / "LFPO" start_struct = LFPO_start_b end_struct = LFPO_end_b - start_struct.calc = choose_calculator(arch="mace_mp", model=MODEL_PATH) - end_struct.calc = choose_calculator(arch="mace_mp", model=MODEL_PATH) + start_struct.calc = choose_calculator( + arch="mace_mp", model=MODEL_PATH, device=device + ) + end_struct.calc = choose_calculator(arch="mace_mp", model=MODEL_PATH, device=device) neb = NEB( init_struct=start_struct, @@ -114,6 +133,7 @@ def test_set_calc(tmp_path, LFPO_start_b, LFPO_end_b): fmax=4, neb_kwargs={"method": "aseneb"}, file_prefix=file_prefix, + device=device, ) neb.run() @@ -124,8 +144,12 @@ def test_set_calc(tmp_path, LFPO_start_b, LFPO_end_b): assert neb.results["delta_E"] == pytest.approx(-3.0149328722473e-07) -def test_neb_functions(tmp_path, LFPO_start_b, LFPO_end_b): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_neb_functions(tmp_path, LFPO_start_b, LFPO_end_b, device): """Test individual NEB functions.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + file_prefix = tmp_path / "LFPO" neb = NEB( @@ -137,6 +161,7 @@ def test_neb_functions(tmp_path, LFPO_start_b, LFPO_end_b): interpolator="ase", neb_kwargs={"method": "aseneb"}, file_prefix=file_prefix, + device=device, ) neb.interpolate() neb.optimize() @@ -148,8 +173,12 @@ def test_neb_functions(tmp_path, LFPO_start_b, LFPO_end_b): assert neb.results["delta_E"] == pytest.approx(-3.0149328722473e-07) -def test_neb_plot(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_neb_plot(tmp_path, device): """Test plotting NEB before running NEBTools.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + file_prefix = tmp_path / "LFPO" neb = NEB( @@ -159,6 +188,7 @@ def test_neb_plot(tmp_path): steps=2, neb_kwargs={"method": "aseneb"}, file_prefix=file_prefix, + device=device, ) neb.optimize() neb.plot() @@ -170,8 +200,12 @@ def test_neb_plot(tmp_path): assert neb.results["max_force"] == pytest.approx(1.5425684122118983) -def test_converge_warning(tmp_path, LFPO_start_b, LFPO_end_b): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_converge_warning(tmp_path, LFPO_start_b, LFPO_end_b, device): """Test warning raised if NEB does not converge.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + neb = NEB( init_struct=LFPO_start_b, final_struct=LFPO_end_b, @@ -182,14 +216,19 @@ def test_converge_warning(tmp_path, LFPO_start_b, LFPO_end_b): fmax=0.1, neb_kwargs={"method": "aseneb"}, file_prefix=tmp_path / "LFPO", + device=device, ) with pytest.warns(UserWarning, match="NEB optimization has not converged"): neb.run() assert not neb.converged -def test_restart(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_restart(tmp_path, device): """Test restarting NEB optimisation.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + neb = NEB( neb_structs=DATA_PATH / "LiFePO4-neb-band.xyz", arch="mace", @@ -198,6 +237,7 @@ def test_restart(tmp_path): file_prefix=tmp_path / "LFPO", neb_kwargs={"method": "aseneb"}, fmax=1.3, + device=device, ) neb.optimize() neb.run_nebtools() @@ -211,8 +251,12 @@ def test_restart(tmp_path): assert final_force < init_force -def test_restart_update_climb(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_restart_update_climb(tmp_path, device): """Test updating NEB climb setting when continuing optimization.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + results = {} for label in ("climb", "no-climb"): @@ -224,6 +268,7 @@ def test_restart_update_climb(tmp_path): fmax=1.3, neb_kwargs={"method": "aseneb"}, file_prefix=tmp_path / "LFPO", + device=device, ) neb.run() neb.neb.climb = label == "climb" diff --git a/tests/test_neb_cli.py b/tests/test_neb_cli.py index 18e63b63..12bac4bc 100644 --- a/tests/test_neb_cli.py +++ b/tests/test_neb_cli.py @@ -7,6 +7,7 @@ from ase import Atoms from ase.io import read import pytest +import torch from typer.testing import CliRunner import yaml @@ -32,8 +33,12 @@ def test_help(): assert "Usage: janus neb [OPTIONS]" in strip_ansi_codes(result.stdout) -def test_neb(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_neb(tmp_path, device): """Test calculating force constants and band structure.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + with chdir(tmp_path): results_dir = Path("./janus_results") results_path = results_dir / "LiFePO4_start-neb-results.dat" @@ -60,6 +65,8 @@ def test_neb(tmp_path): 5, "--plot-band", "--write-band", + "--device", + device, ], ) assert result.exit_code == 0 @@ -113,8 +120,12 @@ def test_neb(tmp_path): clear_log_handlers() -def test_minimize(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_minimize(tmp_path, device): """Test minimizing structures before interpolation and optimization.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + file_prefix = tmp_path / "LFPO" results_path = tmp_path / "LFPO-neb-results.dat" min_init_path = tmp_path / "LFPO-init-opt.extxyz" @@ -144,6 +155,8 @@ def test_minimize(tmp_path): "--no-tracker", "--file-prefix", file_prefix, + "--device", + device, ], ) assert result.exit_code == 0 @@ -185,8 +198,12 @@ def test_minimize(tmp_path): ) -def test_bands(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_bands(tmp_path, device): """Test using band that has already been generated.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + file_prefix = tmp_path / "LFPO" log_path = tmp_path / "LFPO-neb-log.yml" @@ -202,6 +219,8 @@ def test_bands(tmp_path): 2, "--file-prefix", file_prefix, + "--device", + device, ], ) assert result.exit_code == 0 @@ -285,8 +304,12 @@ def test_invalid_n_images(tmp_path): assert isinstance(result.exception, ValueError) -def test_neb_class(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_neb_class(tmp_path, device): """Test passing neb_class and invalid neb_kwargs.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + file_prefix = tmp_path / "LFPO" log_path = tmp_path / "LFPO-neb-log.yml" @@ -313,6 +336,8 @@ def test_neb_class(tmp_path): "--no-tracker", "--file-prefix", file_prefix, + "--device", + device, ], ) assert result.exit_code == 1 @@ -320,8 +345,12 @@ def test_neb_class(tmp_path): assert_log_contains(log_path, includes="Using NEB class: DyNEB") -def test_interpolator(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_interpolator(tmp_path, device): """Test passing interpolator_kwargs.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + file_prefix = tmp_path / "LFPO" results_path = tmp_path / "LFPO-neb-results.dat" @@ -346,6 +375,8 @@ def test_interpolator(tmp_path): "--no-tracker", "--file-prefix", file_prefix, + "--device", + device, ], ) assert result.exit_code == 0 @@ -359,8 +390,12 @@ def test_interpolator(tmp_path): ) -def test_optimzer(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_optimzer(tmp_path, device): """Test passing optimizer and optimizer_kwargs.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + file_prefix = tmp_path / "LFPO" results_path = tmp_path / "LFPO-neb-results.dat" log_path = tmp_path / "LFPO-neb-log.yml" @@ -386,6 +421,8 @@ def test_optimzer(tmp_path): "--no-tracker", "--file-prefix", file_prefix, + "--device", + device, ], ) assert result.exit_code == 0 @@ -462,8 +499,12 @@ def test_invalid_opt(tmp_path): assert isinstance(result.exception, AttributeError) -def test_model(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_model(tmp_path, device): """Test model passed correctly.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + file_prefix = tmp_path / "NaCl" results_path = tmp_path / "NaCl-neb-band.extxyz" log_path = tmp_path / "test.log" @@ -523,8 +564,12 @@ def test_missing_arch(tmp_path): assert "Missing option" in result.stderr -def test_info(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_info(tmp_path, device): """Test info written to output structures.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + file_prefix = tmp_path / "LiFePO4" neb_path = tmp_path / "LiFePO4-neb-band.extxyz" diff --git a/tests/test_phonons.py b/tests/test_phonons.py index c774e6df..afdb554b 100644 --- a/tests/test_phonons.py +++ b/tests/test_phonons.py @@ -7,6 +7,7 @@ from ase.io import read from h5py import File import pytest +import torch from janus_core.calculations.phonons import Phonons from janus_core.calculations.single_point import SinglePoint @@ -17,40 +18,55 @@ MODEL_PATH = Path(__file__).parent / "models" / "mace_mp_small.model" -def test_init(): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_init(device): """Test initialising Phonons.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + single_point = SinglePoint( struct=DATA_PATH / "NaCl.cif", arch="mace", model=MODEL_PATH, + device=device, ) phonons = Phonons(struct=single_point.struct) assert str(phonons.file_prefix) == str(Path("janus_results") / "Cl4Na4") -def test_calc_phonons(): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_calc_phonons(device): """Test calculating phonons from ASE atoms object.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + struct = read(DATA_PATH / "NaCl.cif") - struct.calc = choose_calculator(arch="mace_mp", model=MODEL_PATH) + struct.calc = choose_calculator(arch="mace_mp", model=MODEL_PATH, device=device) phonons = Phonons( struct=struct, + device=device, ) phonons.calc_force_constants(write_force_consts=False) assert "phonon" in phonons.results -def test_force_consts_compression(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_force_consts_compression(tmp_path, device): """Test compression of force constants.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + log_file = tmp_path / "phonons.log" force_constants = tmp_path / "NaCl-force_constants.hdf5" struct = read(DATA_PATH / "NaCl.cif") - struct.calc = choose_calculator(arch="mace_mp", model=MODEL_PATH) + struct.calc = choose_calculator(arch="mace_mp", model=MODEL_PATH, device=device) phonons = Phonons( struct=struct, + device=device, file_prefix=tmp_path / "NaCl", log_kwargs={"filename": log_file}, hdf5=True, @@ -64,8 +80,12 @@ def test_force_consts_compression(tmp_path): assert h5f["force_constants"].compression == "gzip" -def test_optimize(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_optimize(tmp_path, device): """Test optimizing structure before calculation.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + log_file = tmp_path / "phonons.log" opt_file = tmp_path / "NaCl-opt.extxyz" @@ -73,9 +93,11 @@ def test_optimize(tmp_path): struct=DATA_PATH / "NaCl.cif", arch="mace", model=MODEL_PATH, + device=device, ) phonons = Phonons( struct=single_point.struct, + device=device, log_kwargs={"filename": log_file}, minimize=True, minimize_kwargs={"write_kwargs": {"filename": opt_file}}, @@ -107,18 +129,24 @@ def test_invalid_struct(): ) -def test_logging(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_logging(tmp_path, device): """Test attaching logger to Phonons and emissions are saved to info.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + log_file = tmp_path / "phonons.log" single_point = SinglePoint( struct=DATA_PATH / "NaCl.cif", arch="mace_mp", model=MODEL_PATH, + device=device, ) phonons = Phonons( struct=single_point.struct, + device=device, log_kwargs={"filename": log_file}, write_results=False, ) @@ -131,14 +159,19 @@ def test_logging(tmp_path): assert single_point.struct.info["emissions"] > 0 -def test_symmetrize(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_symmetrize(tmp_path, device): """Test symmetrize.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + file_prefix = tmp_path / "NaCl" phonons_1 = Phonons( struct=DATA_PATH / "NaCl-deformed.cif", arch="mace_mp", model=MODEL_PATH, + device=device, write_results=False, minimize=True, minimize_kwargs={"fmax": 0.001}, @@ -151,6 +184,7 @@ def test_symmetrize(tmp_path): struct=DATA_PATH / "NaCl-deformed.cif", arch="mace_mp", model=MODEL_PATH, + device=device, write_results=False, minimize=True, minimize_kwargs={"fmax": 0.001}, diff --git a/tests/test_phonons_cli.py b/tests/test_phonons_cli.py index e9b5ca56..75527c20 100644 --- a/tests/test_phonons_cli.py +++ b/tests/test_phonons_cli.py @@ -7,6 +7,7 @@ from ase.io import read from h5py import File as HDF5Open import pytest +import torch from typer.testing import CliRunner import yaml @@ -32,8 +33,12 @@ def test_help(): assert "Usage: janus phonons [OPTIONS]" in strip_ansi_codes(result.stdout) -def test_phonons(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_phonons(tmp_path, device): """Test calculating force constants and band structure.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + with chdir(tmp_path): results_dir = Path("janus_results") phonopy_path = results_dir / "NaCl-phonopy.yml" @@ -51,6 +56,8 @@ def test_phonons(tmp_path): "mace_mp", "--no-hdf5", "--bands", + "--device", + device, ], ) assert result.exit_code == 0 @@ -104,8 +111,12 @@ def test_phonons(tmp_path): clear_log_handlers() -def test_bands_simple(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_bands_simple(tmp_path, device): """Test calculating force constants and reduced bands information.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + file_prefix = tmp_path / "NaCl" autoband_results = tmp_path / "NaCl-auto_bands.yml" summary_path = tmp_path / "NaCl-phonons-summary.yml" @@ -125,6 +136,8 @@ def test_bands_simple(tmp_path): "--no-hdf5", "--file-prefix", file_prefix, + "--device", + device, ], ) assert result.exit_code == 0 @@ -146,8 +159,12 @@ def test_bands_simple(tmp_path): assert phonon_summary["config"]["bands"] -def test_qpoints_simple(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_qpoints_simple(tmp_path, device): """Test qpoints mode.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + with chdir(tmp_path): results_dir = Path("janus_results") qpoints_results = results_dir / "NaCl-qpoints.hdf5" @@ -167,6 +184,8 @@ def test_qpoints_simple(tmp_path): "--qpoints", "--write-full", "--hdf5", + "--device", + device, ], ) assert result.exit_code == 0 @@ -191,9 +210,13 @@ def test_qpoints_simple(tmp_path): assert phonon_summary["config"]["qpoints"] +@pytest.mark.parametrize("device", ["cpu", "cuda"]) @pytest.mark.parametrize("compression", [None, "gzip", "lzf"]) -def test_hdf5(tmp_path, compression): +def test_hdf5(tmp_path, compression, device): """Test saving force constants and bands to HDF5 in new directory.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + file_prefix = tmp_path / "test" / "NaCl" phonon_results = tmp_path / "test" / "NaCl-phonopy.yml" hdf5_results = tmp_path / "test" / "NaCl-force_constants.hdf5" @@ -218,6 +241,8 @@ def test_hdf5(tmp_path, compression): file_prefix, "--hdf5", *compression_kwargs, + "--device", + device, ], ) assert result.exit_code == 0 @@ -263,8 +288,12 @@ def test_hdf5(tmp_path, compression): assert h5f["force_constants"].compression == compression -def test_thermal_props(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_thermal_props(tmp_path, device): """Test calculating thermal properties.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + file_prefix = tmp_path / "test" / "NaCl" thermal_results = tmp_path / "test" / "NaCl-thermal.yml" summary_path = tmp_path / "test" / "NaCl-phonons-summary.yml" @@ -281,6 +310,8 @@ def test_thermal_props(tmp_path): "--no-hdf5", "--file-prefix", file_prefix, + "--device", + device, ], ) assert result.exit_code == 0 @@ -296,8 +327,12 @@ def test_thermal_props(tmp_path): check_output_files(summary=phonon_summary, output_files=output_files) -def test_dos(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_dos(tmp_path, device): """Test calculating the DOS.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + file_prefix = tmp_path / "test" / "NaCl" dos_results = tmp_path / "test" / "NaCl-dos.dat" @@ -315,6 +350,8 @@ def test_dos(tmp_path): "--no-hdf5", "--file-prefix", file_prefix, + "--device", + device, ], ) assert result.exit_code == 0 @@ -324,8 +361,12 @@ def test_dos(tmp_path): assert lines[-1].split()[0] == "0.0000000000" -def test_pdos(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_pdos(tmp_path, device): """Test calculating the PDOS.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + file_prefix = tmp_path / "test" / "NaCl" pdos_results = tmp_path / "test" / "NaCl-pdos.dat" @@ -343,6 +384,8 @@ def test_pdos(tmp_path): "--no-hdf5", "--file-prefix", file_prefix, + "--device", + device, ], ) assert result.exit_code == 0 @@ -353,8 +396,12 @@ def test_pdos(tmp_path): assert lines[-1].split()[0] == "0.0000000000" -def test_plot(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_plot(tmp_path, device): """Test for plotting routines.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + file_prefix = tmp_path / "NaCl" phonon_results = tmp_path / "NaCl-phonopy.yml" bands_path = tmp_path / "NaCl-auto_bands.yml" @@ -388,6 +435,8 @@ def test_plot(tmp_path): "--no-write-full", "--file-prefix", file_prefix, + "--device", + device, ], ) assert result.exit_code == 0 @@ -435,9 +484,13 @@ def test_plot(tmp_path): ] +@pytest.mark.parametrize("device", ["cpu", "cuda"]) @pytest.mark.parametrize("supercell,supercell_matrix", test_data) -def test_supercell(supercell, supercell_matrix, tmp_path): +def test_supercell(supercell, supercell_matrix, tmp_path, device): """Test setting the supercell.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + file_prefix = tmp_path / "NaCl" param_file = tmp_path / "NaCl-phonopy.yml" @@ -492,8 +545,12 @@ def test_invalid_supercell(supercell, tmp_path): assert result.exit_code == 1 or result.exit_code == 2 -def test_minimize_kwargs(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_minimize_kwargs(tmp_path, device): """Test setting optimizer function and writing optimized structure.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + file_prefix = tmp_path / "test" opt_path = tmp_path / "test-opt.extxyz" log_path = tmp_path / "test-phonons-log.yml" @@ -525,8 +582,12 @@ def test_minimize_kwargs(tmp_path): assert opt_path.exists() -def test_minimize_filename(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_minimize_filename(tmp_path, device): """Test minimize filename overwrites default.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + file_prefix = tmp_path / "test" / "test" opt_path = tmp_path / "test" / "geomopt-opt.extxyz" summary_path = tmp_path / "test" / "test-phonons-summary.yml" @@ -563,9 +624,13 @@ def test_minimize_filename(tmp_path): check_output_files(summary=phonon_summary, output_files=output_files) +@pytest.mark.parametrize("device", ["cpu", "cuda"]) @pytest.mark.parametrize("read_kwargs", ["{'index': 0}", "{}"]) -def test_valid_traj_input(read_kwargs, tmp_path): +def test_valid_traj_input(read_kwargs, tmp_path, device): """Test valid trajectory input structure handled.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + file_prefix = tmp_path / "traj" phonon_results = tmp_path / "traj-phonopy.yml" @@ -612,8 +677,12 @@ def test_invalid_traj_input(tmp_path): assert isinstance(result.exception, ValueError) -def test_no_carbon(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_no_carbon(tmp_path, device): """Test disabling carbon tracking.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + file_prefix = tmp_path / "NaCl" summary_path = tmp_path / "NaCl-phonons-summary.yml" @@ -629,6 +698,8 @@ def test_no_carbon(tmp_path): "--no-tracker", "--file-prefix", file_prefix, + "--device", + device, ], ) assert result.exit_code == 0 @@ -639,8 +710,12 @@ def test_no_carbon(tmp_path): assert "emissions" not in phonon_summary -def test_displacement_kwargs(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_displacement_kwargs(tmp_path, device): """Test displacement_kwargs can be set.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + file_prefix_1 = tmp_path / "NaCl_1" file_prefix_2 = tmp_path / "NaCl_2" displacement_file_1 = tmp_path / "NaCl_1-phonopy.yml" @@ -659,6 +734,8 @@ def test_displacement_kwargs(tmp_path): "{'is_plusminus': True}", "--file-prefix", file_prefix_1, + "--device", + device, ], ) assert result.exit_code == 0 @@ -676,6 +753,8 @@ def test_displacement_kwargs(tmp_path): "{'is_plusminus': False}", "--file-prefix", file_prefix_2, + "--device", + device, ], ) assert result.exit_code == 0 @@ -694,8 +773,12 @@ def test_displacement_kwargs(tmp_path): assert n_displacements_2 == 2 -def test_paths(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_paths(tmp_path, device): """Test passing qpoint file.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + file_prefix = tmp_path / "NaCl" qpoint_file = DATA_PATH / "paths.yml" band_results = tmp_path / "NaCl-bands.yml" @@ -714,6 +797,8 @@ def test_paths(tmp_path): qpoint_file, "--file-prefix", file_prefix, + "--device", + device, ], ) assert result.exit_code == 0 @@ -725,8 +810,12 @@ def test_paths(tmp_path): assert bands["npath"] == 1 -def test_model(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_model(tmp_path, device): """Test model passed correctly.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + file_prefix = tmp_path / "NaCl" results_path = tmp_path / "NaCl-opt.extxyz" log_path = tmp_path / "test.log" @@ -748,6 +837,8 @@ def test_model(tmp_path): "--no-tracker", "--no-hdf5", "--minimize", + "--device", + device, ], ) assert result.exit_code == 0 diff --git a/tests/test_single_point.py b/tests/test_single_point.py index 50518db9..a834137b 100644 --- a/tests/test_single_point.py +++ b/tests/test_single_point.py @@ -11,6 +11,7 @@ from ase.io import read from numpy import isfinite import pytest +import torch from janus_core.calculations.single_point import SinglePoint from tests.utils import chdir, read_atoms, skip_extras @@ -42,19 +43,26 @@ ] +@pytest.mark.parametrize("device", ["cpu", "cuda"]) @pytest.mark.parametrize( "struct, expected, properties, prop_key, calc_kwargs, idx", test_data ) -def test_potential_energy(struct, expected, properties, prop_key, calc_kwargs, idx): +def test_potential_energy( + struct, expected, properties, prop_key, calc_kwargs, idx, device +): """Test single point energy using MACE calculators.""" skip_extras("mace") + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + single_point = SinglePoint( struct=DATA_PATH / struct, arch="mace", model=MACE_PATH, calc_kwargs=calc_kwargs, properties=properties, + device=device, ) results = single_point.run()[prop_key] @@ -72,50 +80,46 @@ def test_potential_energy(struct, expected, properties, prop_key, calc_kwargs, i assert results_2 == pytest.approx(expected) +@pytest.mark.parametrize("device", ["cpu", "cuda"]) @pytest.mark.parametrize( - "arch, device, expected_energy, struct, kwargs", + "arch, expected_energy, struct, kwargs", [ - ("chgnet", "cpu", -29.331436157226562, "NaCl.cif", {}), - ("dpa3", "cpu", -27.053507387638092, "NaCl.cif", {"model": DPA3_PATH}), + ("chgnet", -29.331436157226562, "NaCl.cif", {}), + ("dpa3", -27.053507387638092, "NaCl.cif", {"model": DPA3_PATH}), ( "fairchem", - "cpu", -27.10070295, "NaCl.cif", {"model": UMA_LABEL}, ), - ("grace", "cpu", -27.081155042373453, "NaCl.cif", {}), - ("mace_off", "cpu", -2081.1209264240006, "H2O.cif", {}), - ("mace_omol", "cpu", -2079.8650795528843, "H2O.cif", {}), - ("mattersim", "cpu", -27.06208038330078, "NaCl.cif", {}), + ("grace", -27.081155042373453, "NaCl.cif", {}), + ("mace_off", -2081.1209264240006, "H2O.cif", {}), + ("mace_omol", -2079.8650795528843, "H2O.cif", {}), + ("mattersim", -27.06208038330078, "NaCl.cif", {}), ( "nequip", - "cpu", -169815.1282456301, "toluene.xyz", {"model": NEQUIP_PATH}, ), - ("orb", "cpu", -27.08186149597168, "NaCl.cif", {}), - ("orb", "cpu", -27.089094161987305, "NaCl.cif", {"model": "orb-v2"}), - ("upet", "cpu", -30.168052673339844, "NaCl.cif", {}), + ("orb", -27.08186149597168, "NaCl.cif", {}), + ("orb", -27.089094161987305, "NaCl.cif", {"model": "orb-v2"}), + ("upet", -30.168052673339844, "NaCl.cif", {}), ( "upet", - "cpu", -27.47624969482422, "NaCl.cif", {"model": PET_MAD_CHECKPOINT}, ), ( "sevennet", - "cpu", -27.061979293823242, "NaCl.cif", {"model": SEVENNET_PATH}, ), - ("sevennet", "cpu", -27.061979293823242, "NaCl.cif", {}), + ("sevennet", -27.061979293823242, "NaCl.cif", {}), ( "sevennet", - "cpu", -27.061979293823242, "NaCl.cif", {"model": "SevenNet-0_11July2024"}, @@ -126,6 +130,9 @@ def test_extras(arch, device, expected_energy, struct, kwargs): """Test single point energy using extra MLIP calculators.""" skip_extras(arch) + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + try: single_point = SinglePoint( struct=DATA_PATH / struct, @@ -150,14 +157,19 @@ def test_extras(arch, device, expected_energy, struct, kwargs): raise err -def test_single_point_none(): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_single_point_none(device): """Test single point stress using MACE calculator.""" skip_extras("mace") + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + single_point = SinglePoint( struct=DATA_PATH / "NaCl.cif", arch="mace", model=MACE_PATH, + device=device, ) results = single_point.run() @@ -165,14 +177,19 @@ def test_single_point_none(): assert prop in results -def test_single_point_clean(): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_single_point_clean(device): """Test single point stress using MACE calculator.""" skip_extras("mace") + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + single_point = SinglePoint( struct=DATA_PATH / "H2O.cif", arch="mace", model=MACE_PATH, + device=device, ) results = single_point.run() @@ -181,15 +198,20 @@ def test_single_point_clean(): assert "mace_stress" not in results -def test_single_point_traj(): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_single_point_traj(device): """Test single point stress using MACE calculator.""" skip_extras("mace") + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + single_point = SinglePoint( struct=DATA_PATH / "benzene-traj.xyz", arch="mace", model=MACE_PATH, properties="energy", + device=device, ) assert len(single_point.struct) == 2 @@ -204,10 +226,14 @@ def test_single_point_traj(): ) -def test_single_point_write(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_single_point_write(tmp_path, device): """Test writing singlepoint results.""" skip_extras("mace") + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + with chdir(tmp_path): data_path = DATA_PATH / "NaCl.cif" results_dir = Path("janus_results") @@ -218,6 +244,7 @@ def test_single_point_write(tmp_path): arch="mace", model=MACE_PATH, write_results=True, + device=device, ) assert "mace_forces" not in single_point.struct.arrays @@ -241,10 +268,14 @@ def test_single_point_write(tmp_path): ) -def test_single_point_write_kwargs(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_single_point_write_kwargs(tmp_path, device): """Test passing write_kwargs to singlepoint results.""" skip_extras("mace") + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + data_path = DATA_PATH / "NaCl.cif" results_path = tmp_path / "NaCl.extxyz" @@ -254,6 +285,7 @@ def test_single_point_write_kwargs(tmp_path): model=MACE_PATH, write_results=True, write_kwargs={"filename": results_path}, + device=device, ) assert "mace_forces" not in single_point.struct.arrays @@ -263,10 +295,14 @@ def test_single_point_write_kwargs(tmp_path): assert "mace_forces" in atoms.arrays -def test_single_point_molecule(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_single_point_molecule(tmp_path, device): """Test singlepoint results for isolated molecule.""" skip_extras("mace") + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + data_path = DATA_PATH / "H2O.cif" results_path = tmp_path / "H2O.extxyz" single_point = SinglePoint( @@ -274,6 +310,7 @@ def test_single_point_molecule(tmp_path): arch="mace", model=MACE_PATH, properties="energy", + device=device, ) assert isfinite(single_point.run()["energy"]).all() @@ -305,16 +342,21 @@ def test_invalid_prop(): ) -def test_atoms(): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_atoms(device): """Test passing ASE Atoms structure.""" skip_extras("mace") + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + struct = read(DATA_PATH / "NaCl.cif") single_point = SinglePoint( struct=struct, arch="mace", model=MACE_PATH, properties="energy", + device=device, ) assert single_point.run()["energy"] < 0 @@ -328,10 +370,14 @@ def test_no_atoms_or_path(): ) -def test_invalidate_calc(): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_invalidate_calc(device): """Test setting invalidate_calc via write_kwargs.""" skip_extras("mace") + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + struct = DATA_PATH / "NaCl.cif" single_point = SinglePoint( @@ -339,6 +385,7 @@ def test_invalidate_calc(): arch="mace", model=MACE_PATH, write_kwargs={"invalidate_calc": False}, + device=device, ) single_point.run() @@ -349,10 +396,14 @@ def test_invalidate_calc(): assert "energy" not in single_point.struct.calc.results -def test_logging(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_logging(tmp_path, device): """Test attaching logger to SinglePoint and emissions are saved to info.""" skip_extras("mace") + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + log_file = tmp_path / "sp.log" single_point = SinglePoint( @@ -361,6 +412,7 @@ def test_logging(tmp_path): model=MACE_PATH, properties="energy", log_kwargs={"filename": log_file}, + device=device, ) assert "emissions" not in single_point.struct.info @@ -372,15 +424,20 @@ def test_logging(tmp_path): assert single_point.struct.info["emissions"] > 0 -def test_hessian(): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_hessian(device): """Test Hessian.""" skip_extras("mace") + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + sp = SinglePoint( model=MACE_PATH, struct=DATA_PATH / "NaCl.cif", arch="mace_mp", properties="hessian", + device=device, ) results = sp.run() assert "hessian" in results @@ -388,15 +445,20 @@ def test_hessian(): assert "mace_mp_hessian" in sp.struct.info -def test_hessian_traj(): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_hessian_traj(device): """Test calculating Hessian for trajectory.""" skip_extras("mace") + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + sp = SinglePoint( model=MACE_PATH, struct=DATA_PATH / "benzene-traj.xyz", arch="mace_mp", properties="hessian", + device=device, ) results = sp.run() assert "hessian" in results @@ -452,6 +514,7 @@ def test_missing_arch(struct): SinglePoint(struct=struct) +@pytest.mark.parametrize("device", ["cpu", "cuda"]) @pytest.mark.parametrize( "arch, kwargs, pred", [ @@ -466,11 +529,14 @@ def test_missing_arch(struct): ("sevennet", {}, -0.08281749), ], ) -def test_dispersion(arch, kwargs, pred): +def test_dispersion(arch, kwargs, pred, device): """Test dispersion correction.""" skip_extras(arch) pytest.importorskip("torch_dftd") + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + try: data_path = DATA_PATH / "benzene.xyz" sp_no_d3 = SinglePoint( @@ -478,6 +544,7 @@ def test_dispersion(arch, kwargs, pred): arch=arch, properties="energy", calc_kwargs={"dispersion": False}, + device=device, ) assert not isinstance(sp_no_d3.struct.calc, SumCalculator) no_d3_results = sp_no_d3.run() @@ -487,6 +554,7 @@ def test_dispersion(arch, kwargs, pred): arch=arch, properties="energy", calc_kwargs={"dispersion": True, "dispersion_kwargs": {**kwargs}}, + device=device, ) assert isinstance(sp_d3.struct.calc, SumCalculator) d3_results = sp_d3.run() @@ -500,11 +568,15 @@ def test_dispersion(arch, kwargs, pred): raise err -def test_mace_mp_dispersion(): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_mace_mp_dispersion(device): """Test mace_mp dispersion correction matches default.""" skip_extras("mace_mp") pytest.importorskip("torch_dftd") + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + from mace.calculators import mace_mp data_path = DATA_PATH / "benzene.xyz" @@ -514,6 +586,7 @@ def test_mace_mp_dispersion(): arch="mace_mp", properties="energy", calc_kwargs={"dispersion": False}, + device=device, ).run()["energy"] d3_energy = SinglePoint( @@ -521,10 +594,11 @@ def test_mace_mp_dispersion(): arch="mace_mp", properties="energy", calc_kwargs={"dispersion": True}, + device=device, ).run()["energy"] struct = read(data_path) - struct.calc = mace_mp(model="small", dispersion=True) + struct.calc = mace_mp(model="small", dispersion=True, device=device) mace_d3_energy = SinglePoint( struct=struct, diff --git a/tests/test_singlepoint_cli.py b/tests/test_singlepoint_cli.py index 0ac37110..ce72f206 100644 --- a/tests/test_singlepoint_cli.py +++ b/tests/test_singlepoint_cli.py @@ -6,6 +6,8 @@ from ase import Atoms from ase.io import read +import pytest +import torch from typer.testing import CliRunner import yaml @@ -32,8 +34,12 @@ def test_singlepoint_help(): assert "Usage: janus singlepoint [OPTIONS]" in strip_ansi_codes(result.stdout) -def test_singlepoint(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_singlepoint(tmp_path, device): """Test singlepoint calculation.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + with chdir(tmp_path): results_dir = Path("janus_results") results_path = results_dir / "NaCl-results.extxyz" @@ -48,6 +54,8 @@ def test_singlepoint(tmp_path): DATA_PATH / "NaCl.cif", "--arch", "mace_mp", + "--device", + device, ], ) assert result.exit_code == 0 @@ -77,8 +85,12 @@ def test_singlepoint(tmp_path): clear_log_handlers() -def test_properties(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_properties(tmp_path, device): """Test properties for singlepoint calculation in a new directory.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + results_path_1 = tmp_path / "test" / "H2O-energy-results.extxyz" results_path_2 = tmp_path / "test" / "H2O-stress-results.extxyz" log_path = tmp_path / "test.log" @@ -101,6 +113,8 @@ def test_properties(tmp_path): log_path, "--summary", summary_path, + "--device", + device, ], ) assert result.exit_code == 0 @@ -125,6 +139,8 @@ def test_properties(tmp_path): log_path, "--summary", summary_path, + "--device", + device, ], ) assert result.exit_code == 0 @@ -134,8 +150,12 @@ def test_properties(tmp_path): assert "mace_mp_energy" not in atoms.info -def test_read_kwargs(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_read_kwargs(tmp_path, device): """Test setting read_kwargs for singlepoint calculation.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + results_path = tmp_path / "benzene-traj-results.extxyz" log_path = tmp_path / "test.log" summary_path = tmp_path / "summary.yml" @@ -158,6 +178,8 @@ def test_read_kwargs(tmp_path): log_path, "--summary", summary_path, + "--device", + device, ], ) assert result.exit_code == 0 @@ -166,8 +188,12 @@ def test_read_kwargs(tmp_path): assert isinstance(atoms, list) -def test_calc_kwargs(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_calc_kwargs(tmp_path, device): """Test setting calc_kwargs for singlepoint calculation.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + results_path = tmp_path / "NaCl-results.extxyz" log_path = tmp_path / "test.log" summary_path = tmp_path / "summary.yml" @@ -190,14 +216,20 @@ def test_calc_kwargs(tmp_path): log_path, "--summary", summary_path, + "--device", + device, ], ) assert result.exit_code == 0 assert "Using float32 for MACECalculator" in result.stdout -def test_log(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_log(tmp_path, device): """Test log correctly written for singlepoint.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + results_path = tmp_path / "NaCl-results.extxyz" log_path = tmp_path / "test.log" summary_path = tmp_path / "summary.yml" @@ -218,6 +250,8 @@ def test_log(tmp_path): log_path, "--summary", summary_path, + "--device", + device, ], ) assert result.exit_code == 0 @@ -225,8 +259,12 @@ def test_log(tmp_path): assert_log_contains(log_path, includes=["Starting single point calculation"]) -def test_summary(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_summary(tmp_path, device): """Test summary file can be read correctly.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + results_path = tmp_path / "benzene-traj-results.extxyz" log_path = tmp_path / "test.log" summary_path = tmp_path / "summary.yml" @@ -247,6 +285,8 @@ def test_summary(tmp_path): log_path, "--summary", summary_path, + "--device", + device, ], ) @@ -281,8 +321,12 @@ def test_summary(tmp_path): check_output_files(summary=sp_summary, output_files=output_files) -def test_config(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_config(tmp_path, device): """Test passing a config file with read kwargs, and values to be overwritten.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + results_path = tmp_path / "benzene-traj-results.extxyz" log_path = tmp_path / "test.log" summary_path = tmp_path / "summary.yml" @@ -303,6 +347,8 @@ def test_config(tmp_path): summary_path, "--config", DATA_PATH / "singlepoint_config.yml", + "--device", + device, ], ) assert result.exit_code == 0 @@ -335,8 +381,12 @@ def test_invalid_config(): assert isinstance(result.exception, ValueError) -def test_write_kwargs(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_write_kwargs(tmp_path, device): """Test setting invalidate_calc and write_results via write_kwargs.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + results_path = tmp_path / "NaCl-results.extxyz" log_path = tmp_path / "test.log" summary_path = tmp_path / "summary.yml" @@ -357,6 +407,8 @@ def test_write_kwargs(tmp_path): log_path, "--summary", summary_path, + "--device", + device, ], ) assert result.exit_code == 0 @@ -367,8 +419,12 @@ def test_write_kwargs(tmp_path): assert "forces" in atoms.calc.results -def test_write_cif(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_write_cif(tmp_path, device): """Test writing out a cif file.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + results_path = tmp_path / "NaCl-results.cif" log_path = tmp_path / "test.log" summary_path = tmp_path / "summary.yml" @@ -389,6 +445,8 @@ def test_write_cif(tmp_path): log_path, "--summary", summary_path, + "--device", + device, ], ) assert result.exit_code == 0 @@ -396,8 +454,12 @@ def test_write_cif(tmp_path): assert isinstance(atoms, Atoms) -def test_hessian(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_hessian(tmp_path, device): """Test Hessian calculation.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + results_path = tmp_path / "NaCl-results.extxyz" log_path = tmp_path / "test.log" summary_path = tmp_path / "summary.yml" @@ -423,6 +485,8 @@ def test_hessian(tmp_path): log_path, "--summary", summary_path, + "--device", + device, ], ) assert result.exit_code == 0 @@ -435,8 +499,12 @@ def test_hessian(tmp_path): assert atoms.info["units"]["hessian"] == "ev/Ang^2" -def test_no_carbon(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_no_carbon(tmp_path, device): """Test disabling carbon tracking.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + results_path = tmp_path / "NaCl-results.extxyz" log_path = tmp_path / "test.log" summary_path = tmp_path / "summary.yml" @@ -458,6 +526,8 @@ def test_no_carbon(tmp_path): "--no-tracker", "--summary", summary_path, + "--device", + device, ], ) assert result.exit_code == 0 @@ -468,8 +538,12 @@ def test_no_carbon(tmp_path): assert "emissions" not in sp_summary -def test_file_prefix(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_file_prefix(tmp_path, device): """Test file prefix creates directories and affects all files.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + file_prefix = tmp_path / "test/test" result = runner.invoke( app, @@ -481,6 +555,8 @@ def test_file_prefix(tmp_path): "mace_mp", "--file-prefix", file_prefix, + "--device", + device, ], ) assert result.exit_code == 0 @@ -493,8 +569,12 @@ def test_file_prefix(tmp_path): } -def test_model(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_model(tmp_path, device): """Test model passed correctly.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") + file_prefix = tmp_path / "NaCl" results_path = tmp_path / "NaCl-results.extxyz" log_path = tmp_path / "test.log" @@ -514,6 +594,8 @@ def test_model(tmp_path): "--file-prefix", file_prefix, "--no-tracker", + "--device", + device, ], ) assert result.exit_code == 0 diff --git a/tests/test_train_cli.py b/tests/test_train_cli.py index 572c1e5d..1e199074 100644 --- a/tests/test_train_cli.py +++ b/tests/test_train_cli.py @@ -6,7 +6,7 @@ from ase.io import read, write import pytest -from pytest import skip +import torch from typer.testing import CliRunner import yaml @@ -30,7 +30,9 @@ runner = CliRunner() -def write_tmp_config_mace(config_path: Path, tmp_path: Path) -> Path: +def write_tmp_config_mace( + config_path: Path, tmp_path: Path, device: str = "cpu" +) -> Path: """ Fix paths in config files and write corrected config to tmp_path for mace. @@ -60,6 +62,8 @@ def write_tmp_config_mace(config_path: Path, tmp_path: Path) -> Path: if file in config and (MODEL_PATH / Path(config[file]).name).exists(): config[file] = str(MODEL_PATH / Path(config[file]).name) + config["device"] = device + # Write out temporary config with corrected paths tmp_config = tmp_path / "config.yml" with open(tmp_config, "w", encoding="utf8") as file: @@ -73,6 +77,7 @@ def write_tmp_config_nequip( tmp_path: Path, fine_tune: bool = False, model_type: str = "package", + device: str = "cpu", ) -> Path: """ Fix paths in config files and write corrected config to tmp_path for nequip. @@ -120,6 +125,10 @@ def write_tmp_config_nequip( if (MODEL_PATH / pth).is_file(): model_dict[f"{model_type}_path"] = str(MODEL_PATH / pth) + if "trainer" not in config: + config["trainer"] = {} + config["trainer"]["accelerator"] = "gpu" if device == "cuda" else "cpu" + # Write out temporary config with corrected paths tmp_config = tmp_path / "config.yaml" with open(tmp_config, "w", encoding="utf8") as file: @@ -188,8 +197,11 @@ def test_help(): assert "Usage: janus train [OPTIONS]" in strip_ansi_codes(result.stdout) -def test_train(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_train(tmp_path, device): """Test MLIP training.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") skip_extras("mace") with chdir(tmp_path): @@ -202,7 +214,9 @@ def test_train(tmp_path): log_path = results_dir / "train-log.yml" summary_path = results_dir / "train-summary.yml" - config = write_tmp_config_mace(DATA_PATH / "mlip_train.yml", Path.cwd()) + config = write_tmp_config_mace( + DATA_PATH / "mlip_train.yml", Path.cwd(), device=device + ) result = runner.invoke( app, @@ -249,15 +263,20 @@ def test_train(tmp_path): clear_log_handlers() -def test_train_with_foundation(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_train_with_foundation(tmp_path, device): """Test MLIP training raises error with foundation_model in config.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") skip_extras("mace") results_dir = tmp_path / "janus_results" log_path = results_dir / "test.log" summary_path = results_dir / "summary.yml" - config = write_tmp_config_mace(DATA_PATH / "mlip_train_invalid.yml", tmp_path) + config = write_tmp_config_mace( + DATA_PATH / "mlip_train_invalid.yml", tmp_path, device=device + ) result = runner.invoke( app, @@ -276,8 +295,11 @@ def test_train_with_foundation(tmp_path): assert isinstance(result.exception, ValueError) -def test_fine_tune(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_fine_tune(tmp_path, device): """Test MLIP fine-tuning.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") skip_extras("mace") with chdir(tmp_path): @@ -291,7 +313,13 @@ def test_fine_tune(tmp_path): summary_path = tmp_path / "summary.yml" logs_path = results_dir / "logs" +<<<<<<< HEAD config = write_tmp_config_mace(DATA_PATH / "mlip_fine_tune.yml", Path.cwd()) +======= + config = write_tmp_config_mace( + DATA_PATH / "mlip_fine_tune.yml", Path(), device=device + ) +>>>>>>> d7be6e9 (add gpu workflow and tests) result = runner.invoke( app, @@ -318,8 +346,11 @@ def test_fine_tune(tmp_path): clear_log_handlers() -def test_fine_tune_no_foundation(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_fine_tune_no_foundation(tmp_path, device): """Test MLIP fine-tuning raises errors without foundation_model.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") skip_extras("mace") log_path = tmp_path / "test.log" @@ -345,8 +376,11 @@ def test_fine_tune_no_foundation(tmp_path): assert isinstance(result.exception, ValueError) -def test_fine_tune_invalid_foundation(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_fine_tune_invalid_foundation(tmp_path, device): """Test MLIP fine-tuning raises errors with invalid foundation_model.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") skip_extras("mace") log_path = tmp_path / "test.log" @@ -371,8 +405,11 @@ def test_fine_tune_invalid_foundation(tmp_path): assert isinstance(result.exception, ValueError) -def test_no_carbon(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_no_carbon(tmp_path, device): """Test disabling carbon tracking.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") skip_extras("mace") with chdir(tmp_path): @@ -386,7 +423,9 @@ def test_no_carbon(tmp_path): log_path = tmp_path / "test.log" summary_path = tmp_path / "summary.yml" - config = write_tmp_config_mace(DATA_PATH / "mlip_train.yml", Path.cwd()) + config = write_tmp_config_mace( + DATA_PATH / "mlip_train.yml", Path.cwd(), device=device + ) result = runner.invoke( app, @@ -419,8 +458,11 @@ def test_no_carbon(tmp_path): clear_log_handlers() -def test_nequip_train(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_nequip_train(tmp_path, device): """Test training with nequip.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") skip_extras("nequip") with chdir(tmp_path): @@ -436,7 +478,7 @@ def test_nequip_train(tmp_path): train_log_path = results_dir / "train_log" metrics_path = results_dir / "train_log/version_0/metrics.csv" - config_path = write_tmp_config_nequip(config_path, tmp_path) + config_path = write_tmp_config_nequip(config_path, tmp_path, device=device) result = runner.invoke( app, @@ -467,14 +509,17 @@ def test_nequip_train(tmp_path): assert header[:3] == ["epoch", "lr-Adam", "step"] -def test_nequip_train_invalid_config_suffix(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_nequip_train_invalid_config_suffix(tmp_path, device): """Test training with nequip.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") skip_extras("nequip") with chdir(tmp_path): config_path = DATA_PATH / "mlip_train.yml" - config_path = write_tmp_config_mace(config_path, tmp_path) + config_path = write_tmp_config_mace(config_path, tmp_path, device=device) result = runner.invoke( app, @@ -488,8 +533,11 @@ def test_nequip_train_invalid_config_suffix(tmp_path): not NEQUIP_EXTRA_MODEL_PATH.exists(), reason=f"Extra model: {NEQUIP_EXTRA_MODEL_PATH} not downloaded.", ) -def test_nequip_fine_tune_foundation(tmp_path): +@pytest.mark.parametrize("device", ["cpu", "cuda"]) +def test_nequip_fine_tune_foundation(tmp_path, device): """Test fine-tuning with a nequip foundation model.""" + if device == "cuda" and not torch.cuda.is_available(): + pytest.skip("CUDA not available") skip_extras("nequip") with chdir(tmp_path): @@ -503,7 +551,7 @@ def test_nequip_fine_tune_foundation(tmp_path): metrics_path = results_dir / "train_log/version_0/metrics.csv" config_path = write_tmp_config_nequip( - DATA_PATH / "nequip_fine_tune.yaml", tmp_path, True + DATA_PATH / "nequip_fine_tune.yaml", tmp_path, True, device=device ) result = runner.invoke( @@ -591,7 +639,7 @@ def test_sevennet_fine_tune_foundation(tmp_path): skip_extras("sevennet") if not SEVENNET_EXTRA_MODEL_PATH.exists(): - skip(f"Extra model: {SEVENNET_EXTRA_MODEL_PATH} not downloaded.") + pytest.skip(f"Extra model: {SEVENNET_EXTRA_MODEL_PATH} not downloaded.") with chdir(tmp_path): log_path = tmp_path / "test.log" From c56fe6a4f58fc228efa52c7147656b2bd7363787 Mon Sep 17 00:00:00 2001 From: Alin M Elena Date: Sun, 29 Mar 2026 08:34:26 +0100 Subject: [PATCH 02/16] run only cpu on cpu --- .github/workflows/ci.yml | 4 ++-- .github/workflows/mac.yml | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index d46171dd..08fc7d6a 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -46,7 +46,7 @@ jobs: env: # show timings of tests PYTEST_ADDOPTS: "--durations=0" - run: uv run --no-sync pytest --cov janus_core --cov-append . + run: uv run --no-sync pytest --cov janus_core --cov-append . -k "not cuda" - name: Install updated e3nn dependencies run: | @@ -60,7 +60,7 @@ jobs: # show timings of tests PYTEST_ADDOPTS: "--durations=0" HF_TOKEN: ${{ secrets.HF_TOKEN }} - run: uv run --no-sync pytest tests/test_{mlip_calculators,single_point}.py + run: uv run --no-sync pytest tests/test_{mlip_calculators,single_point}.py -k "not cuda" - name: Report coverage to Coveralls uses: coverallsapp/github-action@v2 diff --git a/.github/workflows/mac.yml b/.github/workflows/mac.yml index cbc70f94..d62e8958 100644 --- a/.github/workflows/mac.yml +++ b/.github/workflows/mac.yml @@ -42,7 +42,7 @@ jobs: env: # show timings of tests PYTEST_ADDOPTS: "--durations=0" - run: uv run --no-sync pytest + run: uv run --no-sync pytest -k "not cuda" - name: Install updated e3nn dependencies run: | @@ -55,4 +55,4 @@ jobs: # show timings of tests PYTEST_ADDOPTS: "--durations=0" HF_TOKEN: ${{ secrets.HF_TOKEN }} - run: uv run --no-sync pytest tests/test_{mlip_calculators,single_point}.py + run: uv run --no-sync pytest tests/test_{mlip_calculators,single_point}.py -k "not cuda" From ed82efbad35468c26ce5facd43ce1c4133b4fecd Mon Sep 17 00:00:00 2001 From: Alin M Elena Date: Mon, 30 Mar 2026 07:01:18 +0100 Subject: [PATCH 03/16] disable tracker on gpu --- .github/workflows/bluesky.yml | 2 ++ 1 file changed, 2 insertions(+) diff --git a/.github/workflows/bluesky.yml b/.github/workflows/bluesky.yml index b839e50f..581e3207 100644 --- a/.github/workflows/bluesky.yml +++ b/.github/workflows/bluesky.yml @@ -8,6 +8,8 @@ jobs: runs-on: [self-hosted, gpu] if: github.repository == 'stfc/janus-core' timeout-minutes: 60 + env: + CODECARBON_DISABLED: 'true' strategy: matrix: python-version: ["3.10","3.11","3.12"] From c0a008e4e284cb845b67e0a47aa6ad0be52f0314 Mon Sep 17 00:00:00 2001 From: Alin M Elena Date: Fri, 10 Apr 2026 15:55:14 +0100 Subject: [PATCH 04/16] turn codecarbon error to a warning --- janus_core/helpers/log.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/janus_core/helpers/log.py b/janus_core/helpers/log.py index dbcc9cc7..02756b49 100644 --- a/janus_core/helpers/log.py +++ b/janus_core/helpers/log.py @@ -206,7 +206,7 @@ def config_tracker( carbon_logger.removeHandler(carbon_logger.handlers[0]) if not hasattr(tracker, "_emissions"): - raise ValueError( + raise Warning( "Carbon tracker has not been configured correctly. Please try " "reconfiguring, or disable the tracker." ) From 5dcae04b9b0c0a29188b51753df60dd10857621e Mon Sep 17 00:00:00 2001 From: Alin M Elena Date: Fri, 10 Apr 2026 15:55:41 +0100 Subject: [PATCH 05/16] turn codecarbon error to a warning --- janus_core/helpers/log.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/janus_core/helpers/log.py b/janus_core/helpers/log.py index 02756b49..fd56e7a1 100644 --- a/janus_core/helpers/log.py +++ b/janus_core/helpers/log.py @@ -187,6 +187,7 @@ def config_tracker( OfflineEmissionsTracker | None Configured offline codecarbon tracker, if logger is specified. """ + tracker = None if janus_logger and track_carbon: carbon_logger = LoggerOutput(janus_logger) tracker = OfflineEmissionsTracker( @@ -211,7 +212,4 @@ def config_tracker( "reconfiguring, or disable the tracker." ) - else: - tracker = None - return tracker From a50d421c0b6b5a66ac0300455009599e8cd18ee8 Mon Sep 17 00:00:00 2001 From: Alin M Elena Date: Fri, 10 Apr 2026 16:04:42 +0100 Subject: [PATCH 06/16] fix precommit --- tests/test_md.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_md.py b/tests/test_md.py index 17c8dbfb..8aff0aa7 100644 --- a/tests/test_md.py +++ b/tests/test_md.py @@ -299,7 +299,7 @@ def test_nvt_csvr(tmp_path, device): @pytest.mark.parametrize("device", ["cpu", "cuda"]) @pytest.mark.skipif(MTK_IMPORT_FAILED, reason="Requires updated version of ASE") @pytest.mark.parametrize("mtk_flavour", ["iso", "aniso"]) -def test_npt_mtk(tmp_path, device,mtk_flavour): +def test_npt_mtk(tmp_path, device, mtk_flavour): """Test NPT MTK molecular dynamics.""" if device == "cuda" and not torch.cuda.is_available(): pytest.skip("CUDA not available") From 5b1857ed6ce647120a3c664bbc6e6b7c91b95338 Mon Sep 17 00:00:00 2001 From: Alin M Elena Date: Fri, 10 Apr 2026 16:25:47 +0100 Subject: [PATCH 07/16] change to warn --- janus_core/helpers/log.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/janus_core/helpers/log.py b/janus_core/helpers/log.py index fd56e7a1..66bfe445 100644 --- a/janus_core/helpers/log.py +++ b/janus_core/helpers/log.py @@ -5,6 +5,7 @@ import json import logging from typing import Literal +from warnings import warn from codecarbon import OfflineEmissionsTracker from codecarbon.output import LoggerOutput @@ -207,9 +208,10 @@ def config_tracker( carbon_logger.removeHandler(carbon_logger.handlers[0]) if not hasattr(tracker, "_emissions"): - raise Warning( + raise warn( "Carbon tracker has not been configured correctly. Please try " - "reconfiguring, or disable the tracker." + "reconfiguring, or disable the tracker.", + stacklevel=2, ) return tracker From d4196f363efeb687766902d6124d669d2555fd17 Mon Sep 17 00:00:00 2001 From: Alin M Elena Date: Fri, 10 Apr 2026 16:30:13 +0100 Subject: [PATCH 08/16] fix borked merge --- tests/test_train_cli.py | 4 ---- 1 file changed, 4 deletions(-) diff --git a/tests/test_train_cli.py b/tests/test_train_cli.py index 1e199074..da3a7d48 100644 --- a/tests/test_train_cli.py +++ b/tests/test_train_cli.py @@ -313,13 +313,9 @@ def test_fine_tune(tmp_path, device): summary_path = tmp_path / "summary.yml" logs_path = results_dir / "logs" -<<<<<<< HEAD - config = write_tmp_config_mace(DATA_PATH / "mlip_fine_tune.yml", Path.cwd()) -======= config = write_tmp_config_mace( DATA_PATH / "mlip_fine_tune.yml", Path(), device=device ) ->>>>>>> d7be6e9 (add gpu workflow and tests) result = runner.invoke( app, From 1514c9bd7e21d77439b7c64e078770d0620712c1 Mon Sep 17 00:00:00 2001 From: Alin M Elena Date: Fri, 10 Apr 2026 16:54:49 +0100 Subject: [PATCH 09/16] fix various --- janus_core/helpers/log.py | 2 +- tests/test_log.py | 2 +- tests/test_single_point.py | 2 ++ 3 files changed, 4 insertions(+), 2 deletions(-) diff --git a/janus_core/helpers/log.py b/janus_core/helpers/log.py index 66bfe445..ffe557e8 100644 --- a/janus_core/helpers/log.py +++ b/janus_core/helpers/log.py @@ -208,7 +208,7 @@ def config_tracker( carbon_logger.removeHandler(carbon_logger.handlers[0]) if not hasattr(tracker, "_emissions"): - raise warn( + warn( "Carbon tracker has not been configured correctly. Please try " "reconfiguring, or disable the tracker.", stacklevel=2, diff --git a/tests/test_log.py b/tests/test_log.py index b8d0bfbd..ec086ffd 100644 --- a/tests/test_log.py +++ b/tests/test_log.py @@ -48,5 +48,5 @@ def test_tracker_error(tmp_path): log_path = tmp_path / "test.log" logger = config_logger(name=__name__, filename=log_path) - with pytest.raises(ValueError): + with pytest.raises(Warning): config_tracker(janus_logger=logger, country_2letter_iso_code=123) diff --git a/tests/test_single_point.py b/tests/test_single_point.py index a834137b..6c18e88a 100644 --- a/tests/test_single_point.py +++ b/tests/test_single_point.py @@ -584,6 +584,7 @@ def test_mace_mp_dispersion(device): no_d3_energy = SinglePoint( struct=data_path, arch="mace_mp", + model="small", properties="energy", calc_kwargs={"dispersion": False}, device=device, @@ -592,6 +593,7 @@ def test_mace_mp_dispersion(device): d3_energy = SinglePoint( struct=data_path, arch="mace_mp", + model="small", properties="energy", calc_kwargs={"dispersion": True}, device=device, From 033844c5c0bc230fecb4643aa04dd9c4d7ca9453 Mon Sep 17 00:00:00 2001 From: ElliottKasoar <45317199+ElliottKasoar@users.noreply.github.com> Date: Fri, 10 Apr 2026 17:02:58 +0100 Subject: [PATCH 10/16] Fix testing tracker warning --- tests/test_log.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_log.py b/tests/test_log.py index ec086ffd..018ac48e 100644 --- a/tests/test_log.py +++ b/tests/test_log.py @@ -48,5 +48,5 @@ def test_tracker_error(tmp_path): log_path = tmp_path / "test.log" logger = config_logger(name=__name__, filename=log_path) - with pytest.raises(Warning): + with pytest.warns(UserWarning): config_tracker(janus_logger=logger, country_2letter_iso_code=123) From 5ade16e1d895b2c7efb6bdca5be4c76e1e12c828 Mon Sep 17 00:00:00 2001 From: ElliottKasoar <45317199+ElliottKasoar@users.noreply.github.com> Date: Fri, 10 Apr 2026 17:38:24 +0100 Subject: [PATCH 11/16] Set tracker to None if fails to set up --- janus_core/helpers/log.py | 1 + 1 file changed, 1 insertion(+) diff --git a/janus_core/helpers/log.py b/janus_core/helpers/log.py index ffe557e8..095afd20 100644 --- a/janus_core/helpers/log.py +++ b/janus_core/helpers/log.py @@ -213,5 +213,6 @@ def config_tracker( "reconfiguring, or disable the tracker.", stacklevel=2, ) + tracker = None return tracker From 082917056e190de15bb587c78f3b0e6394ef368d Mon Sep 17 00:00:00 2001 From: alin m elena Date: Fri, 10 Apr 2026 20:02:16 +0100 Subject: [PATCH 12/16] add mccabe complexity flag --- pyproject.toml | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/pyproject.toml b/pyproject.toml index c1389b8a..250c543a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -168,10 +168,15 @@ select = [ "N", # isort "UP", + # mccabe complexity + "C901", ] [tool.ruff.lint.per-file-ignores] 'containers/*.py' = ['D100', 'E501'] +[tool.ruff.lint.mccabe] +max-complexity = 10 + [tool.ruff.lint.isort] force-sort-within-sections = true required-imports = ["from __future__ import annotations"] From c2fa0180912f22cacb9f52d470575e76da478b20 Mon Sep 17 00:00:00 2001 From: alin m elena Date: Fri, 10 Apr 2026 20:06:56 +0100 Subject: [PATCH 13/16] add mccabe complexity flag 5 --- pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 250c543a..dfd39e99 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -175,7 +175,7 @@ select = [ 'containers/*.py' = ['D100', 'E501'] [tool.ruff.lint.mccabe] -max-complexity = 10 +max-complexity = 5 [tool.ruff.lint.isort] force-sort-within-sections = true From c0260f7bff97f2dae4b1eb3509902326abda2fab Mon Sep 17 00:00:00 2001 From: Alin Marin Elena Date: Fri, 10 Apr 2026 21:40:20 +0100 Subject: [PATCH 14/16] Update .github/workflows/bluesky.yml Co-authored-by: Elliott Kasoar <45317199+ElliottKasoar@users.noreply.github.com> --- .github/workflows/bluesky.yml | 2 -- 1 file changed, 2 deletions(-) diff --git a/.github/workflows/bluesky.yml b/.github/workflows/bluesky.yml index 581e3207..b839e50f 100644 --- a/.github/workflows/bluesky.yml +++ b/.github/workflows/bluesky.yml @@ -8,8 +8,6 @@ jobs: runs-on: [self-hosted, gpu] if: github.repository == 'stfc/janus-core' timeout-minutes: 60 - env: - CODECARBON_DISABLED: 'true' strategy: matrix: python-version: ["3.10","3.11","3.12"] From aaacc49a6f3aa5a1aae4430d165404de47d2e144 Mon Sep 17 00:00:00 2001 From: Alin Marin Elena Date: Fri, 10 Apr 2026 21:40:38 +0100 Subject: [PATCH 15/16] Update .github/workflows/bluesky.yml Co-authored-by: Elliott Kasoar <45317199+ElliottKasoar@users.noreply.github.com> --- .github/workflows/bluesky.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/bluesky.yml b/.github/workflows/bluesky.yml index b839e50f..da3cbf18 100644 --- a/.github/workflows/bluesky.yml +++ b/.github/workflows/bluesky.yml @@ -16,7 +16,7 @@ jobs: - uses: actions/checkout@v4 - name: Install uv - uses: astral-sh/setup-uv@v6 + uses: astral-sh/setup-uv@v7 with: python-version: ${{ matrix.python-version }} From 421e22a72715f8ad27f797f3bebc65ebc83f8a4b Mon Sep 17 00:00:00 2001 From: Alin Marin Elena Date: Fri, 10 Apr 2026 21:40:52 +0100 Subject: [PATCH 16/16] Update .github/workflows/bluesky.yml Co-authored-by: Elliott Kasoar <45317199+ElliottKasoar@users.noreply.github.com> --- .github/workflows/bluesky.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/bluesky.yml b/.github/workflows/bluesky.yml index da3cbf18..d57d2def 100644 --- a/.github/workflows/bluesky.yml +++ b/.github/workflows/bluesky.yml @@ -13,7 +13,7 @@ jobs: python-version: ["3.10","3.11","3.12"] steps: - - uses: actions/checkout@v4 + - uses: actions/checkout@v6 - name: Install uv uses: astral-sh/setup-uv@v7