diff --git a/tests/conftest.py b/tests/conftest.py index f6c1182..07c056f 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,21 +1,16 @@ -import tempfile from dataclasses import asdict from pathlib import Path import numpy as np import parsl import pytest -import requests -import yaml from ase import Atoms from ase.build import bulk from ase.calculators.emt import EMT import psiflow from psiflow.data import Dataset, FlowAtoms -from psiflow.models import (AllegroConfig, AllegroModel, MACEConfig, MACEModel, - NequIPConfig, NequIPModel) -from psiflow.reference import EMTReference +from psiflow.models import AllegroConfig, MACEConfig, MACEModel, NequIPConfig def pytest_addoption(parser): @@ -104,7 +99,7 @@ def generate_emt_cu_data(nstates, amplitude): pos = atoms.get_positions() box = atoms.get_cell() atoms_list = [] - for i in range(nstates): + for _ in range(nstates): atoms.set_positions( pos + np.random.uniform(-amplitude, amplitude, size=(len(atoms), 3)) ) diff --git a/tests/test_bias.py b/tests/test_bias.py index c8b37d1..41c6e36 100644 --- a/tests/test_bias.py +++ b/tests/test_bias.py @@ -1,13 +1,8 @@ -import tempfile - import numpy as np import pytest -import requests -import yaml -from ase.build import bulk, make_supercell +from ase.build import make_supercell from psiflow.data import Dataset, NullState -from psiflow.models import NequIPModel from psiflow.walkers import PlumedBias, RandomWalker from psiflow.walkers.bias import ( generate_external_grid, @@ -111,7 +106,7 @@ def test_bias_evaluate(context, dataset): assert np.allclose(volume, values[i, 0]) assert np.allclose(np.zeros(values[:, 1].shape), values[:, 1]) dataset_ = bias.evaluate(dataset, as_dataset=True) - for i, atoms in enumerate(dataset_.as_list().result()): + for atoms in dataset_.as_list().result(): assert np.allclose(atoms.get_volume(), atoms.info["CV1"]) state = dataset_[0].result() state.reset() @@ -142,7 +137,7 @@ def test_bias_external(context, dataset, tmp_path): CV: VOLUME external: EXTERNAL ARG=CV FILE=test_grid """ - bias_function = lambda x: np.exp(-0.01 * (x - 150) ** 2) + bias_function = lambda x: np.exp(-0.01 * (x - 150) ** 2) # noqa: E731 variable = np.linspace(0, 300, 500) grid = generate_external_grid(bias_function, variable, "CV", periodic=False) data = {"EXTERNAL": grid} diff --git a/tests/test_data.py b/tests/test_data.py index fab0803..3b9889c 100644 --- a/tests/test_data.py +++ b/tests/test_data.py @@ -12,24 +12,24 @@ def test_flow_atoms(dataset, tmp_path): atoms = dataset.get(index=0).result().copy() # copy necessary with HTEX! - assert type(atoms) == FlowAtoms + assert type(atoms) is FlowAtoms atoms.reference_status = True atoms_ = atoms.copy() - assert atoms_.reference_status == True + assert atoms_.reference_status atoms_ = FlowAtoms.from_atoms(atoms) - assert atoms_.reference_status == True + assert atoms_.reference_status for i in range(dataset.length().result()): atoms = dataset[i].result() - assert type(atoms) == FlowAtoms - assert atoms.reference_status == True + assert type(atoms) is FlowAtoms + assert atoms.reference_status assert dataset.labeled().length().result() == dataset.length().result() dataset += Dataset([NullState]) assert dataset.length().result() == 1 + dataset.not_null().length().result() - assert atoms.reference_status == True + assert atoms.reference_status atoms.reset() atoms.cell[:] = np.array([[3, 1, 1], [1, 5, 0], [0, -1, 5]]) - assert not "energy" in atoms.info - assert atoms.reference_status == False + assert "energy" not in atoms.info + assert not atoms.reference_status assert tuple(sorted(atoms.elements)) == ("Cu", "H") assert not is_reduced(atoms.cell) atoms.canonical_orientation() @@ -42,7 +42,6 @@ def test_dataset_empty(tmp_path): assert isinstance(dataset.data_future, DataFuture) path_xyz = tmp_path / "test.xyz" dataset.save(path_xyz) # ensure the copy is executed before assert - assert not os.path.isfile(path_xyz) psiflow.wait() assert os.path.isfile(path_xyz) with pytest.raises(ValueError): # cannot save outside cwd @@ -53,8 +52,8 @@ def test_dataset_append(dataset): assert 20 == dataset.length().result() atoms_list = dataset.as_list().result() assert len(atoms_list) == 20 - assert type(atoms_list) == list - assert type(atoms_list[0]) == FlowAtoms + assert type(atoms_list) is list + assert type(atoms_list[0]) is FlowAtoms empty = Dataset([]) # use [] instead of None empty.append(dataset) assert 20 == empty.length().result() @@ -216,7 +215,7 @@ def test_data_elements(dataset): def test_data_reset(dataset): dataset = dataset.reset() - assert not "energy" in dataset[0].result().info + assert "energy" not in dataset[0].result().info def test_nullstate(context): @@ -249,7 +248,7 @@ def test_identifier(dataset): for i in range(data.length().result()): s = data[i].result() if not s == NullState: - assert not "identifier" in s.info + assert "identifier" not in s.info identifier = data.assign_identifiers(10) assert identifier.result() == 10 # none are labeled identifier = dataset.assign_identifiers(10) diff --git a/tests/test_learning.py b/tests/test_learning.py index 6f8a0aa..7e62f0c 100644 --- a/tests/test_learning.py +++ b/tests/test_learning.py @@ -161,7 +161,7 @@ def test_incremental_learning(gpu, tmp_path, mace_config, dataset): walkers, ) assert data.length().result() == len(walkers) # perform 1 iteration - for i, walker in enumerate(walkers): + for walker in walkers: assert not walker.is_reset().result() steps, kappas, centers = walker.bias.get_moving_restraint(variable="CV") assert steps == 10 @@ -172,7 +172,7 @@ def test_temperature_ramp(context): assert apply_temperature_ramp(100, 300, 1, 100) == 300 assert apply_temperature_ramp(100, 500, 3, 550) == 500 T = 100 - for i in range(3): + for _ in range(3): T = apply_temperature_ramp(100, 500, 5, T) assert T == 1 / (1 / 100 - 3 * (1 / 100 - 1 / 500) / 4) assert not T == 500 diff --git a/tests/test_models.py b/tests/test_models.py index 3899dd4..d0577e1 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -1,22 +1,22 @@ import ast -import copy import os -from dataclasses import asdict import numpy as np import pytest import torch -from ase.data import chemical_symbols -from ase.io.extxyz import read_extxyz from parsl.app.futures import DataFuture from parsl.dataflow.futures import AppFuture import psiflow -from psiflow.committee import Committee from psiflow.data import Dataset -from psiflow.execution import ModelEvaluation -from psiflow.models import (AllegroModel, MACEConfig, MACEModel, NequIPConfig, - NequIPModel, load_model) +from psiflow.models import ( + AllegroModel, + MACEConfig, + MACEModel, + NequIPConfig, + NequIPModel, + load_model, +) from psiflow.reference import EMTReference @@ -79,14 +79,14 @@ def test_nequip_save_load(nequip_config, dataset, tmp_path): path_config = tmp_path / "config_after_init.yaml" path_model = tmp_path / "model_undeployed.pth" path_deploy = tmp_path / "model_deployed.pth" - futures = model.save(tmp_path, require_done=True) + model.save(tmp_path, require_done=True) assert os.path.exists(path_config_raw) assert os.path.exists(path_config) assert os.path.exists(path_model) assert os.path.exists(path_deploy) model_ = load_model(tmp_path) - assert type(model_) == NequIPModel + assert type(model_) is NequIPModel assert model_.model_future is not None assert model_.deploy_future is not None e1 = model_.evaluate(dataset.get(indices=[3]))[0].result().info["energy"] @@ -199,14 +199,14 @@ def test_allegro_save_load(allegro_config, dataset, tmp_path): path_config = tmp_path / "config_after_init.yaml" path_model = tmp_path / "model_undeployed.pth" path_deploy = tmp_path / "model_deployed.pth" - futures = model.save(tmp_path, require_done=True) + model.save(tmp_path, require_done=True) assert os.path.exists(path_config_raw) assert os.path.exists(path_config) assert os.path.exists(path_model) assert os.path.exists(path_deploy) model_ = load_model(tmp_path) - assert type(model_) == AllegroModel + assert type(model_) is AllegroModel assert model_.model_future is not None assert model_.deploy_future is not None e1 = model_.evaluate(dataset.get(indices=[3]))[0].result().info["energy"] @@ -303,7 +303,7 @@ def test_mace_save_load(mace_config, dataset, tmp_path): assert os.path.exists(path_deployed) model_ = load_model(tmp_path) - assert type(model_) == MACEModel + assert type(model_) is MACEModel assert model_.model_future is not None assert model_.deploy_future is not None e1 = model_.evaluate(dataset.get(indices=[3]))[0].result().info["energy"] diff --git a/tests/test_reference.py b/tests/test_reference.py index b474381..c112451 100644 --- a/tests/test_reference.py +++ b/tests/test_reference.py @@ -1,22 +1,17 @@ -import os from pathlib import Path import molmod import numpy as np import pytest import requests -from ase import Atoms -from ase.io.extxyz import write_extxyz from ase.units import Pascal -from parsl.app.futures import DataFuture from parsl.dataflow.futures import AppFuture from pymatgen.io.cp2k.inputs import Cp2kInput import psiflow from psiflow.data import Dataset, FlowAtoms, NullState from psiflow.reference import CP2KReference, EMTReference, NWChemReference -from psiflow.reference._cp2k import (insert_atoms_in_input, - insert_filepaths_in_input) +from psiflow.reference._cp2k import insert_filepaths_in_input @pytest.fixture @@ -182,11 +177,11 @@ def test_reference_emt(context, dataset, tmp_path): assert evaluated.length().result() == len(atoms_list) atoms = reference.evaluate(dataset_[5]).result() - assert type(atoms) == FlowAtoms - assert atoms.reference_status == True + assert type(atoms) is FlowAtoms + assert atoms.reference_status atoms = reference.evaluate(dataset_[6]).result() - assert type(atoms) == FlowAtoms - assert atoms.reference_status == False + assert type(atoms) is FlowAtoms + assert not atoms.reference_status def test_cp2k_insert_filepaths(fake_cp2k_input): @@ -242,7 +237,7 @@ def test_cp2k_success(context, cp2k_reference): dataset = Dataset([atoms]) evaluated = cp2k_reference.evaluate(dataset[0]) assert isinstance(evaluated, AppFuture) - assert evaluated.result().reference_status == True + assert evaluated.result().reference_statuss assert Path(evaluated.result().reference_stdout).is_file() assert Path(evaluated.result().reference_stderr).is_file() assert "energy" in evaluated.result().info.keys() @@ -375,7 +370,7 @@ def test_cp2k_failure(context, cp2k_data, tmp_path): ) evaluated = reference.evaluate(atoms) assert isinstance(evaluated, AppFuture) - assert evaluated.result().reference_status == False + assert not evaluated.result().reference_status assert "energy" not in evaluated.result().info.keys() with open(evaluated.result().reference_stdout, "r") as f: log = f.read() @@ -393,7 +388,7 @@ def test_cp2k_timeout(context, cp2k_reference): ) evaluated = cp2k_reference.evaluate(atoms) assert isinstance(evaluated, AppFuture) - assert evaluated.result().reference_status == False + assert not evaluated.result().reference_status assert "energy" not in evaluated.result().info.keys() @@ -438,14 +433,14 @@ def test_nwchem_success(nwchem_reference): dataset = Dataset([atoms]) evaluated = nwchem_reference.evaluate(dataset[0]) assert isinstance(evaluated, AppFuture) - assert evaluated.result().reference_status == True + assert evaluated.result().reference_status assert Path(evaluated.result().reference_stdout).is_file() assert Path(evaluated.result().reference_stderr).is_file() assert "energy" in evaluated.result().info.keys() - assert not "stress" in evaluated.result().info.keys() + assert "stress" not in evaluated.result().info.keys() assert "forces" in evaluated.result().arrays.keys() assert evaluated.result().arrays["forces"][0, 0] < 0 assert evaluated.result().arrays["forces"][1, 0] > 0 - energy_h2 = nwchem_reference.evaluate(dataset) + nwchem_reference.evaluate(dataset) assert nwchem_reference.compute_atomic_energy("H").result() < 0 diff --git a/tests/test_sampling.py b/tests/test_sampling.py index 8b60fbb..4006d77 100644 --- a/tests/test_sampling.py +++ b/tests/test_sampling.py @@ -2,13 +2,12 @@ import parsl from psiflow.committee import Committee -from psiflow.data import Dataset, FlowAtoms +from psiflow.data import FlowAtoms from psiflow.metrics import Metrics, log_dataset from psiflow.models import MACEModel from psiflow.reference import EMTReference from psiflow.sampling import sample_with_committee, sample_with_model -from psiflow.walkers import (BiasedDynamicWalker, DynamicWalker, PlumedBias, - RandomWalker) +from psiflow.walkers import BiasedDynamicWalker, DynamicWalker, PlumedBias, RandomWalker def test_sample_metrics(mace_model, dataset, tmp_path): diff --git a/tests/test_walkers.py b/tests/test_walkers.py index ab07273..d20c4e6 100644 --- a/tests/test_walkers.py +++ b/tests/test_walkers.py @@ -1,22 +1,29 @@ import os -from dataclasses import asdict import numpy as np import pytest -import torch from ase import Atoms from ase.units import kB from parsl.dataflow.futures import AppFuture import psiflow from psiflow.data import Dataset -from psiflow.models import MACEModel, NequIPModel +from psiflow.models import MACEModel from psiflow.utils import copy_app_future -from psiflow.walkers import (BaseWalker, BiasedDynamicWalker, DynamicWalker, - OptimizationWalker, PlumedBias, RandomWalker, - load_walker) -from psiflow.walkers.utils import (get_velocities_at_temperature, - parse_openmm_output, parse_yaff_output) +from psiflow.walkers import ( # noqa: F401 + BaseWalker, + BiasedDynamicWalker, + DynamicWalker, + OptimizationWalker, + PlumedBias, + RandomWalker, + load_walker, +) +from psiflow.walkers.utils import ( + get_velocities_at_temperature, + parse_openmm_output, + parse_yaff_output, +) def test_random_walker_multiply(dataset, tmp_path): @@ -29,13 +36,13 @@ def test_random_walker_multiply(dataset, tmp_path): amplitude_pos=amplitude_pos, amplitude_box=amplitude_box, ) - for i, walker in enumerate(walkers): + for walker in walkers: delta = np.abs(dataset[0].result().positions - walker.state.result().positions) assert np.allclose(delta, 0) delta = np.abs(dataset[0].result().positions - walker.state0.result().positions) assert np.allclose(delta, 0) - data = Dataset([w.propagate(None).state for w in walkers]) - for i, walker in enumerate(walkers): + Dataset([w.propagate(None).state for w in walkers]) + for walker in walkers: delta = np.abs( walker.state0.result().positions - walker.state.result().positions ) @@ -58,7 +65,7 @@ def test_walker_save_load(dataset, mace_model, tmp_path): assert os.path.exists(path_state) assert os.path.exists(path_pars) walker_ = load_walker(tmp_path / "new") - assert type(walker_) == DynamicWalker + assert type(walker_) is DynamicWalker assert np.allclose( walker.state0.result().positions, walker_.state0.result().positions, @@ -415,7 +422,7 @@ def check(walkers): bias = PlumedBias(plumed_input) walkers = BiasedDynamicWalker.multiply(3, dataset, bias=bias, steps=123) assert len(walkers) == 3 - assert type(walkers[0]) == BiasedDynamicWalker + assert type(walkers[0]) is BiasedDynamicWalker assert walkers[0].steps == 123 check(walkers) @@ -460,5 +467,5 @@ def test_moving_restraint_walker(dataset, mace_model): ) walker = BiasedDynamicWalker(dataset[0], bias=bias, steps=30) - metadata = walker.propagate(model=mace_model) + walker.propagate(model=mace_model) assert not walker.is_reset().result()