diff --git a/psiflow/committee.py b/psiflow/committee.py index 1d33c3d..235a9c8 100644 --- a/psiflow/committee.py +++ b/psiflow/committee.py @@ -37,12 +37,12 @@ def _compute_disagreements( disagreements = np.zeros(lengths[0]) if metric == "mean_force": for i in range(lengths[0]): - forces = np.zeros((len(inputs), len(data[0][0]), 3)) + forces = np.zeros((len(inputs), len(data[0][i]), 3)) for j in range(len(inputs)): if data[j][i] == NullState: assert j == 0 # nullstates do not depend on j break - forces[j] = data[j][i].arrays["forces"] + forces[j, :] = data[j][i].arrays["forces"] SE = (forces - np.mean(forces, axis=0, keepdims=True)) ** 2 RMSE = np.sqrt(np.mean(SE)) disagreements[i] = RMSE diff --git a/tests/conftest.py b/tests/conftest.py index bbb0c15..4458d22 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -6,7 +6,7 @@ import pytest import yaml from ase import Atoms -from ase.build import bulk +from ase.build import bulk, make_supercell from ase.calculators.emt import EMT import psiflow @@ -94,8 +94,10 @@ def mace_config(): return asdict(mace_config) -def generate_emt_cu_data(nstates, amplitude): - atoms = bulk("Cu", "fcc", a=3.6, cubic=True) +def generate_emt_cu_data(nstates, amplitude, supercell=None): + if supercell is None: + supercell = np.eye(3) + atoms = make_supercell(bulk("Cu", "fcc", a=3.6, cubic=True), supercell) atoms.calc = EMT() pos = atoms.get_positions() box = atoms.get_cell() @@ -120,6 +122,7 @@ def generate_emt_cu_data(nstates, amplitude): @pytest.fixture def dataset(context): data = generate_emt_cu_data(20, 0.2) + data += generate_emt_cu_data(5, 0.15, supercell=np.diag([1, 2, 1])) data_ = [FlowAtoms.from_atoms(atoms) for atoms in data] for atoms in data_: atoms.reference_status = True diff --git a/tests/test_data.py b/tests/test_data.py index 3b9889c..0b3c692 100644 --- a/tests/test_data.py +++ b/tests/test_data.py @@ -49,21 +49,21 @@ def test_dataset_empty(tmp_path): def test_dataset_append(dataset): - assert 20 == dataset.length().result() + l = dataset.length().result() # noqa: E741 atoms_list = dataset.as_list().result() - assert len(atoms_list) == 20 + assert len(atoms_list) == l 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() + assert l == empty.length().result() dataset.append(dataset) - assert 40 == dataset.length().result() + assert 2 * l == dataset.length().result() added = dataset + dataset - assert added.length().result() == 80 - assert dataset.length().result() == 40 # may not changed + assert added.length().result() == 4 * l + assert dataset.length().result() == 2 * l # must not have changed dataset += dataset - assert dataset.length().result() == 80 # may not changed + assert dataset.length().result() == 4 * l # must not have changed # test canonical transformation transformed = dataset.canonical_orientation() @@ -90,10 +90,12 @@ def test_dataset_slice(dataset): dataset_ = dataset.shuffle() equal = np.array([False] * dataset.length().result()) for i in range(dataset.length().result()): - equal[i] = np.allclose( - dataset_[i].result().get_positions(), - dataset[i].result().get_positions(), - ) + pos0 = dataset_[i].result().get_positions() + pos1 = dataset[i].result().get_positions() + if pos0.shape == pos1.shape: + equal[i] = np.allclose(pos0, pos1) + else: + equal[i] = False assert not np.all(equal) @@ -104,7 +106,7 @@ def test_dataset_from_xyz(tmp_path, dataset): loaded = Dataset.load(path_xyz) data = read(path_xyz, index=":") - for i in range(20): + for i in range(dataset.length().result()): assert np.allclose( dataset[i].result().get_positions(), loaded[i].result().get_positions(), @@ -266,9 +268,9 @@ def test_data_offset(dataset): } data = dataset.subtract_offset(**atomic_energies) data_ = data.add_offset(**atomic_energies) - natoms = len(dataset[0].result()) - offset = (natoms - 1) * atomic_energies["Cu"] + atomic_energies["H"] for i in range(dataset.length().result()): + natoms = len(dataset[i].result()) + offset = (natoms - 1) * atomic_energies["Cu"] + atomic_energies["H"] assert np.allclose( data[i].result().info["energy"], dataset[i].result().info["energy"] - offset, diff --git a/tests/test_learning.py b/tests/test_learning.py index 4b3e77f..f5dd6d3 100644 --- a/tests/test_learning.py +++ b/tests/test_learning.py @@ -72,8 +72,8 @@ def test_sequential_learning(gpu, tmp_path, mace_config, dataset): data = learning.run(model, reference, walkers, dataset) psiflow.wait() assert data.labeled().length().result() == len(walkers) + dataset.length().result() - assert learning.identifier.result() == 25 - assert data.assign_identifiers().result() == 25 + assert learning.identifier.result() == 5 + dataset.length().result() + assert data.assign_identifiers().result() == 5 + dataset.length().result() data = learning.run(model, reference, walkers) assert data.length().result() == 0 # iteration 0 already performed diff --git a/tests/test_models.py b/tests/test_models.py index 836a1d7..324e03a 100644 --- a/tests/test_models.py +++ b/tests/test_models.py @@ -103,47 +103,6 @@ def test_nequip_seed(nequip_config): assert model.seed == 112 -def test_nequip_offset(nequip_config, dataset): - config = NequIPConfig(**nequip_config) - model = NequIPModel(config) - model.initialize(dataset[:2]) - with pytest.raises(AssertionError): - model.add_atomic_energy("H", 1) # cannot change this after initialization - assert not model.do_offset - errors = Dataset.get_errors( - dataset, - model.evaluate(dataset), - properties=["energy"], - ) - assert np.mean(errors.result()) < 1e3 # in meV/atom - - reference = EMTReference() - atomic_energies = { - "H": 3e2, - "Cu": reference.compute_atomic_energy("Cu", box_size=6), # future - } - errors_ = Dataset.get_errors( - dataset.subtract_offset(**atomic_energies), - model.evaluate(dataset), - ) - assert np.mean(errors_.result()) > 1e3 # in meV/atom - model.reset() - for element, energy in atomic_energies.items(): - model.add_atomic_energy(element, energy) - assert model.do_offset - model.initialize(dataset[:2]) - errors_same = Dataset.get_errors( - dataset, - model.evaluate(dataset), - properties=["energy"], - ) - assert np.allclose( - errors_same.result(), - errors.result(), - atol=1e-1, - ) - - def test_allegro_init(allegro_config, dataset): model = AllegroModel(allegro_config) model.seed = 1 @@ -325,7 +284,7 @@ def test_mace_seed(mace_config): def test_mace_offset(mace_config, dataset, tmp_path): config = MACEConfig(**mace_config) model = MACEModel(config) - model.initialize(dataset[:2]) + model.initialize(dataset) with pytest.raises(AssertionError): model.add_atomic_energy("H", 1) # cannot change this after initialization assert not model.do_offset @@ -350,7 +309,7 @@ def test_mace_offset(mace_config, dataset, tmp_path): for element, energy in atomic_energies.items(): model.add_atomic_energy(element, energy) assert model.do_offset - model.initialize(dataset[:2]) + model.initialize(dataset) errors_same = Dataset.get_errors( dataset, model.evaluate(dataset), @@ -359,7 +318,7 @@ def test_mace_offset(mace_config, dataset, tmp_path): assert np.allclose( errors_same.result(), errors.result(), - atol=1e-1, + atol=0.1, ) model.save(tmp_path) psiflow.wait()