def test_merge(): clusters = np.eye(5, dtype=np.int32) groups = [(0, 1, 2), (3, 4)] clusters_ = np.array([ [1, 1, 1, 0, 0], [0, 0, 0, 1, 1], [0, 0, 0, 0, 0], [0, 0, 0, 0, 0], [0, 0, 0, 0, 0], ], dtype=np.int32) assert np.allclose(clusters_, Mapping.merge(clusters, groups))
def test_write_load(tmp_path): harmonic, equivalencies = get_harmonic('uio66_ff') mapping = Mapping(harmonic.atoms.get_masses(), equivalencies) clusters = Mapping.merge( mapping.clusters, [(0, 1, 2, 5), (4, 8, 33, 6, 100)], ) mapping.update_clusters(clusters) mapping.update_identities(validate=True) path_npz = tmp_path / 'mapping.npz' mapping.write(path_npz) loaded = Mapping.load(path_npz) assert np.allclose(loaded.masses, mapping.masses) for i in range(mapping.natoms): assert loaded.atom_types[i] == mapping.atom_types[i] assert np.allclose(loaded.clusters, mapping.clusters)