From 5b49f5f0df0a1397a36ee42fd95bd0584c32846b Mon Sep 17 00:00:00 2001 From: njzjz-bot Date: Tue, 28 Jul 2026 19:36:45 +0800 Subject: [PATCH] refactor(lmdb): rename format to deepmd/lmdb Move the LMDB format backend from dpdata/formats/lmdb/ to dpdata/formats/deepmd/lmdb/ and register the canonical format name "deepmd/lmdb". The legacy "lmdb" name is retained as a backward- compatible alias (with from_lmdb/to_lmdb methods) so existing user scripts continue to work. Coding-Agent: opencode opencode-Version: 1.18.8 Model: ustc/glm-5.2 Reasoning-Effort: max --- docs/systems/lmdb.md | 20 +- dpdata/formats/{ => deepmd}/lmdb/__init__.py | 0 dpdata/formats/{ => deepmd}/lmdb/format.py | 16 +- dpdata/plugins/lmdb.py | 4 +- tests/test_lmdb.py | 238 ++++++++++++------- tests/test_lmdb_custom_dtype.py | 54 ++--- 6 files changed, 195 insertions(+), 137 deletions(-) rename dpdata/formats/{ => deepmd}/lmdb/__init__.py (100%) rename dpdata/formats/{ => deepmd}/lmdb/format.py (99%) diff --git a/docs/systems/lmdb.md b/docs/systems/lmdb.md index 67dccc85..63dd0485 100644 --- a/docs/systems/lmdb.md +++ b/docs/systems/lmdb.md @@ -1,6 +1,6 @@ # LMDB Format -The format `lmdb` stores the frames of one or more systems in a single [LMDB](http://www.lmdb.tech/doc/) database, and can be loaded or dumped through {class}`dpdata.System`, {class}`dpdata.LabeledSystem`, and {class}`dpdata.MultiSystems`. The on-disk layout coincides with the LMDB datasets read by the DeePMD-kit data loader. Core fields use the plural on-disk names `coords`, `cells`, `energies`, `forces`, and `virials`, which the DeePMD-kit reader maps to its internal names. Registered additional fields use their `deepmd_name`. +The format `deepmd/lmdb` (with the backward-compatible alias `lmdb`) stores the frames of one or more systems in a single [LMDB](http://www.lmdb.tech/doc/) database, and can be loaded or dumped through {class}`dpdata.System`, {class}`dpdata.LabeledSystem`, and {class}`dpdata.MultiSystems`. The on-disk layout coincides with the LMDB datasets read by the DeePMD-kit data loader. Core fields use the plural on-disk names `coords`, `cells`, `energies`, `forces`, and `virials`, which the DeePMD-kit reader maps to its internal names. Registered additional fields use their `deepmd_name`. In contrast to the directory-based `deepmd/npy` format, every frame is stored as an independent record indexed by a global frame number. Frames within one database may therefore differ in the number of atoms and in chemical composition, which is suited to data sets in which the number of frames per system is small. @@ -29,9 +29,9 @@ A single system or a collection of systems is written to one database. ```python import dpdata -dpdata.LabeledSystem("OUTCAR", fmt="vasp/outcar").to("lmdb", "data.lmdb") +dpdata.LabeledSystem("OUTCAR", fmt="vasp/outcar").to("deepmd/lmdb", "data.lmdb") -dpdata.MultiSystems(*systems).to("lmdb", "data.lmdb") +dpdata.MultiSystems(*systems).to("deepmd/lmdb", "data.lmdb") ``` The element table recorded in `type_map` defaults to the union of the elements present in the data. An explicit table may be supplied through the `type_map` argument, for example the full periodic table. @@ -39,7 +39,7 @@ The element table recorded in `type_map` defaults to the union of the elements p ```python from dpdata.periodic_table import ELEMENTS -dpdata.MultiSystems(*systems).to("lmdb", "data.lmdb", type_map=list(ELEMENTS)) +dpdata.MultiSystems(*systems).to("deepmd/lmdb", "data.lmdb", type_map=list(ELEMENTS)) ``` Frames are committed in batches of 1,000 by default. The `write_batch_size` argument changes the transaction size. If a transaction exceeds `map_size`, the map is enlarged and the same encoded batch is retried before frame counters are advanced. The destination must not exist unless `overwrite=True` is supplied. @@ -50,13 +50,13 @@ Data are first written to a temporary sibling database. After closing it, dpdata The `frame_system_ids` entry of the metadata records, for every frame, the index of the source system it belongs to. DeePMD-kit uses this partition for system-wise sampling, for example through `prob_sys_size`. -When a {class}`dpdata.MultiSystems` is dumped through `to("lmdb", ...)`, its frames are first grouped by chemical formula, and systems that share a formula are merged into a single entry. The resulting `frame_system_ids` therefore reflect the formula grouping rather than the original sources, and the number of systems may be smaller than the number of inputs. +When a {class}`dpdata.MultiSystems` is dumped through `to("deepmd/lmdb", ...)`, its frames are first grouped by chemical formula, and systems that share a formula are merged into a single entry. The resulting `frame_system_ids` therefore reflect the formula grouping rather than the original sources, and the number of systems may be smaller than the number of inputs. -When the original partition must be retained, the function {func}`dpdata.formats.lmdb.dump_systems` writes an ordered sequence of systems without formula merging; each input becomes one system, numbered in iteration order. +When the original partition must be retained, the function {func}`dpdata.formats.deepmd.lmdb.dump_systems` writes an ordered sequence of systems without formula merging; each input becomes one system, numbered in iteration order. ```python import dpdata -from dpdata.formats.lmdb import dump_systems +from dpdata.formats.deepmd.lmdb import dump_systems systems = [dpdata.LabeledSystem(d, fmt="deepmd/npy") for d in directories] dump_systems(systems, "data.lmdb", type_map=["H", "C", "N", "O"]) @@ -80,7 +80,7 @@ Reading an LMDB through dpdata subsequently groups frames by composition and doe ```python import dpdata -ms = dpdata.MultiSystems.from_file("data.lmdb", fmt="lmdb") +ms = dpdata.MultiSystems.from_file("data.lmdb", fmt="deepmd/lmdb") ``` Frames are grouped by composition, and each composition is returned as one system. Atom order is canonicalized by a stable sort on the global atom type; coordinates and all registered atomic fields are permuted consistently. Frames of one composition must have identical field sets and a consistent periodic-boundary condition. By default (`mixed_type=False`) the element table of each resulting system is restricted to the elements that the system contains. When `mixed_type=True`, every system retains the complete element set from the database. Loading through {class}`dpdata.MultiSystems` may normalize the order of `atom_names`, so callers should use element names rather than assume that stored numerical indices are retained. @@ -88,13 +88,13 @@ Frames are grouped by composition, and each composition is returned as one syste The general {class}`dpdata.System` constructor applies its `type_map` argument after format parsing. Consequently, a direct single-system load with an explicit `type_map` retains that complete requested table even when `mixed_type=False`; this is standard `System` behavior. The compact/full distinction above describes the dictionaries yielded to {class}`dpdata.MultiSystems`. ```python -ms = dpdata.MultiSystems.from_file("data.lmdb", fmt="lmdb", mixed_type=True) +ms = dpdata.MultiSystems.from_file("data.lmdb", fmt="deepmd/lmdb", mixed_type=True) ``` A database that holds a single composition may also be read into a {class}`dpdata.LabeledSystem`. If the database contains several compositions, only the first can be represented by a single system and a warning is issued. ```python -ls = dpdata.LabeledSystem("data.lmdb", fmt="lmdb") +ls = dpdata.LabeledSystem("data.lmdb", fmt="deepmd/lmdb") ``` The reader loads all selected frames into memory. The default `max_frames=100000` guard rejects larger data sets before decoding; set it to `None` only when sufficient memory is available. Large training data sets should normally be consumed directly by the DeePMD-kit data loader. diff --git a/dpdata/formats/lmdb/__init__.py b/dpdata/formats/deepmd/lmdb/__init__.py similarity index 100% rename from dpdata/formats/lmdb/__init__.py rename to dpdata/formats/deepmd/lmdb/__init__.py diff --git a/dpdata/formats/lmdb/format.py b/dpdata/formats/deepmd/lmdb/format.py similarity index 99% rename from dpdata/formats/lmdb/format.py rename to dpdata/formats/deepmd/lmdb/format.py index 295de7f5..55d55965 100644 --- a/dpdata/formats/lmdb/format.py +++ b/dpdata/formats/deepmd/lmdb/format.py @@ -1099,21 +1099,21 @@ class LMDBFormat(Format): >>> import dpdata >>> ls = dpdata.LabeledSystem("OUTCAR", fmt="vasp/outcar") - >>> ls.to("lmdb", "data.lmdb") + >>> ls.to("deepmd/lmdb", "data.lmdb") Write many systems into one LMDB, forcing a global type map:: >>> ms = dpdata.MultiSystems(s1, s2, s3) - >>> ms.to("lmdb", "data.lmdb", type_map=["H", "C", "N", "O"]) + >>> ms.to("deepmd/lmdb", "data.lmdb", type_map=["H", "C", "N", "O"]) Read back as standard (per-composition) systems:: - >>> ms = dpdata.MultiSystems.from_file("data.lmdb", fmt="lmdb") + >>> ms = dpdata.MultiSystems.from_file("data.lmdb", fmt="deepmd/lmdb") Read back keeping the full global type map on every system:: >>> ms = dpdata.MultiSystems.from_file( - ... "data.lmdb", fmt="lmdb", mixed_type=True + ... "data.lmdb", fmt="deepmd/lmdb", mixed_type=True ... ) Note that loading through :class:`dpdata.MultiSystems` normalises the @@ -1198,7 +1198,7 @@ def dump_systems( """Write an ordered sequence of systems, one ``frame_system_id`` each. Unlike :meth:`to_multi_systems` (the path used by - ``MultiSystems.to('lmdb', ...)``), the systems are **not** merged by + ``MultiSystems.to('deepmd/lmdb', ...)``), the systems are **not** merged by formula: every element of ``systems`` becomes exactly one source system in the database, numbered ``0, 1, 2, ...`` in iteration order. This preserves the system partition recorded in ``frame_system_ids``, @@ -1452,7 +1452,7 @@ def _first_system(self, file_name, *, require_labeled: bool, **kwargs): warnings.warn( f"LMDB '{file_name}' contains more than one composition; only the " "first is loaded into a single System. Use " - "dpdata.MultiSystems.from_file(..., fmt='lmdb') to load all of them.", + "dpdata.MultiSystems.from_file(..., fmt='deepmd/lmdb') to load all of them.", stacklevel=2, ) return first @@ -2027,7 +2027,7 @@ def dump_systems( Each element of ``systems`` is stored as a distinct source system, numbered ``0, 1, 2, ...`` in iteration order, and recorded in the ``frame_system_ids`` metadata. In contrast to - ``MultiSystems.to('lmdb', ...)``, systems are not merged by formula, so + ``MultiSystems.to('deepmd/lmdb', ...)``, systems are not merged by formula, so the system partition used by DeePMD-kit's ``prob_sys_size`` is kept. Parameters @@ -2057,7 +2057,7 @@ def dump_systems( Examples -------- >>> import dpdata - >>> from dpdata.formats.lmdb import dump_systems + >>> from dpdata.formats.deepmd.lmdb import dump_systems >>> systems = [ ... dpdata.LabeledSystem(d, fmt="deepmd/npy") for d in directories ... ] diff --git a/dpdata/plugins/lmdb.py b/dpdata/plugins/lmdb.py index 5c83dbe2..aa483644 100644 --- a/dpdata/plugins/lmdb.py +++ b/dpdata/plugins/lmdb.py @@ -1,6 +1,8 @@ from __future__ import annotations from dpdata.format import Format -from dpdata.formats.lmdb.format import LMDBFormat +from dpdata.formats.deepmd.lmdb.format import LMDBFormat +# Canonical name; ``lmdb`` is kept as a backward-compatible alias. +Format.register("deepmd/lmdb")(LMDBFormat) Format.register("lmdb")(LMDBFormat) diff --git a/tests/test_lmdb.py b/tests/test_lmdb.py index 5afdcd09..8aa7e378 100644 --- a/tests/test_lmdb.py +++ b/tests/test_lmdb.py @@ -21,7 +21,7 @@ from context import dpdata from dpdata.data_type import Axis, DataError, DataType -from dpdata.formats.lmdb.format import ( +from dpdata.formats.deepmd.lmdb.format import ( LMDBError, LMDBFrameError, LMDBMetadataError, @@ -52,8 +52,8 @@ class TestLMDBLabeledSystem(unittest.TestCase, CompLabeledSys, IsPBC): def setUp(self): self.system_1 = dpdata.LabeledSystem("poscars/OUTCAR.h2o.md", fmt="vasp/outcar") self.lmdb_path = "tmp_labeled.lmdb" - self.system_1.to("lmdb", self.lmdb_path) - self.system_2 = dpdata.LabeledSystem(self.lmdb_path, fmt="lmdb") + self.system_1.to("deepmd/lmdb", self.lmdb_path) + self.system_2 = dpdata.LabeledSystem(self.lmdb_path, fmt="deepmd/lmdb") self.places = 6 self.e_places = 6 self.f_places = 6 @@ -68,8 +68,8 @@ class TestLMDBSystem(unittest.TestCase, CompSys, IsPBC): def setUp(self): self.system_1 = dpdata.System("poscars/POSCAR.h2o.md", fmt="vasp/poscar") self.lmdb_path = "tmp_system.lmdb" - self.system_1.to("lmdb", self.lmdb_path) - self.system_2 = dpdata.System(self.lmdb_path, fmt="lmdb") + self.system_1.to("deepmd/lmdb", self.lmdb_path) + self.system_2 = dpdata.System(self.lmdb_path, fmt="deepmd/lmdb") self.places = 6 self.e_places = 6 self.f_places = 6 @@ -96,8 +96,8 @@ def setUp(self): self.ms_1 = dpdata.MultiSystems(system_1, system_2, system_3) for system in self.ms_1: system.sort_atom_types() - self.ms_1.to("lmdb", self.lmdb_path) - self.ms_2 = dpdata.MultiSystems.from_file(self.lmdb_path, fmt="lmdb") + self.ms_1.to("deepmd/lmdb", self.lmdb_path) + self.ms_2 = dpdata.MultiSystems.from_file(self.lmdb_path, fmt="deepmd/lmdb") self.places = 6 self.e_places = 6 @@ -109,13 +109,51 @@ def tearDown(self): shutil.rmtree(self.lmdb_path) +class TestLMDBLegacyAlias(unittest.TestCase): + """The legacy ``lmdb`` alias and its ``from_lmdb``/``to_lmdb`` methods + must remain equivalent to the canonical ``deepmd/lmdb`` name. + """ + + def setUp(self): + self.system = dpdata.LabeledSystem("poscars/OUTCAR.h2o.md", fmt="vasp/outcar") + self.canonical_path = "tmp_alias_canonical.lmdb" + self.legacy_path = "tmp_alias_legacy.lmdb" + + def tearDown(self): + for path in (self.canonical_path, self.legacy_path): + if os.path.exists(path): + shutil.rmtree(path) + + def test_alias_methods_exist(self): + self.assertTrue(hasattr(self.system, "to_lmdb")) + self.assertTrue(hasattr(self.system, "to_deepmd_lmdb")) + + def test_legacy_alias_round_trip(self): + self.system.to("lmdb", self.legacy_path) + loaded = dpdata.LabeledSystem(self.legacy_path, fmt="deepmd/lmdb") + np.testing.assert_allclose(loaded["energies"], self.system["energies"]) + np.testing.assert_allclose(loaded["forces"], self.system["forces"]) + + def test_canonical_read_via_legacy_alias(self): + self.system.to("deepmd/lmdb", self.canonical_path) + loaded = dpdata.LabeledSystem(self.canonical_path, fmt="lmdb") + np.testing.assert_allclose(loaded["energies"], self.system["energies"]) + np.testing.assert_allclose(loaded["forces"], self.system["forces"]) + + def test_to_lmdb_from_lmdb_methods(self): + self.system.to_lmdb(self.legacy_path) + loaded = dpdata.LabeledSystem().from_lmdb(self.legacy_path) + np.testing.assert_allclose(loaded["energies"], self.system["energies"]) + np.testing.assert_allclose(loaded["forces"], self.system["forces"]) + + class TestLMDBOnDiskFormat(unittest.TestCase): """The on-disk layout must match the DeePMD-kit / reference converters.""" def setUp(self): self.lmdb_path = "tmp_format.lmdb" self.system = dpdata.LabeledSystem("poscars/OUTCAR.h2o.md", fmt="vasp/outcar") - self.system.to("lmdb", self.lmdb_path) + self.system.to("deepmd/lmdb", self.lmdb_path) def tearDown(self): if os.path.exists(self.lmdb_path): @@ -184,7 +222,7 @@ def setUp(self): ) self.ms = dpdata.MultiSystems(s1, s2) self.type_map = ["H", "C", "N", "O"] - self.ms.to("lmdb", self.lmdb_path, type_map=self.type_map) + self.ms.to("deepmd/lmdb", self.lmdb_path, type_map=self.type_map) def tearDown(self): if os.path.exists(self.lmdb_path): @@ -193,7 +231,9 @@ def tearDown(self): def test_mixed_preserves_full_type_map(self): # via MultiSystems the element order is normalized (sorted), but the # full type_map set must be preserved on every system. - ms = dpdata.MultiSystems.from_file(self.lmdb_path, fmt="lmdb", mixed_type=True) + ms = dpdata.MultiSystems.from_file( + self.lmdb_path, fmt="deepmd/lmdb", mixed_type=True + ) for ss in ms: names = [str(n) for n in ss.data["atom_names"]] self.assertEqual(sorted(names), sorted(self.type_map)) @@ -205,15 +245,15 @@ def test_mixed_single_system_preserves_order(self): path = "tmp_mixed_single.lmdb" try: s = dpdata.LabeledSystem("gaussian/methane.gaussianlog", fmt="gaussian/log") - s.to("lmdb", path, type_map=self.type_map) - ls = dpdata.LabeledSystem(path, fmt="lmdb", mixed_type=True) + s.to("deepmd/lmdb", path, type_map=self.type_map) + ls = dpdata.LabeledSystem(path, fmt="deepmd/lmdb", mixed_type=True) self.assertEqual([str(n) for n in ls.data["atom_names"]], self.type_map) finally: if os.path.exists(path): shutil.rmtree(path) def test_standard_compresses_type_map(self): - ms = dpdata.MultiSystems.from_file(self.lmdb_path, fmt="lmdb") + ms = dpdata.MultiSystems.from_file(self.lmdb_path, fmt="deepmd/lmdb") # methane only contains C and H, not N/O for ss in ms: self.assertNotIn("N", ss.data["atom_names"]) @@ -231,7 +271,7 @@ def tearDown(self): def test_explicit_type_map(self): type_map = ["H", "He", "Li", "Be", "B", "C", "N", "O"] - self.system.to("lmdb", self.lmdb_path, type_map=type_map) + self.system.to("deepmd/lmdb", self.lmdb_path, type_map=type_map) with lmdb.open(self.lmdb_path, readonly=True, lock=False) as env: with env.begin() as txn: meta = msgpack.unpackb(txn.get(b"__metadata__"), raw=False) @@ -248,11 +288,11 @@ def test_explicit_type_map(self): self.assertEqual(decoded, expected) def test_missing_element_raises(self): - from dpdata.formats.lmdb.format import LMDBError + from dpdata.formats.deepmd.lmdb.format import LMDBError # water needs O and H; this type_map omits O. with self.assertRaises(LMDBError): - self.system.to("lmdb", self.lmdb_path, type_map=["H", "He"]) + self.system.to("deepmd/lmdb", self.lmdb_path, type_map=["H", "He"]) class TestLMDBReferenceFormatInterop(unittest.TestCase): @@ -305,14 +345,14 @@ def tearDown(self): shutil.rmtree(self.lmdb_path) def test_read_reference(self): - ms = dpdata.MultiSystems.from_file(self.lmdb_path, fmt="lmdb") + ms = dpdata.MultiSystems.from_file(self.lmdb_path, fmt="deepmd/lmdb") self.assertEqual(ms.get_nframes(), 2) self.assertEqual(len(ms), 2) # composition-based grouping: water (3 atoms) + methane (5 atoms) self.assertEqual(sorted(s.get_natoms() for s in ms), [3, 5]) def test_read_reference_values(self): - ms = dpdata.MultiSystems.from_file(self.lmdb_path, fmt="lmdb") + ms = dpdata.MultiSystems.from_file(self.lmdb_path, fmt="deepmd/lmdb") water = next(s for s in ms if s.get_natoms() == 3) # Frames are canonicalized by global type while coordinates and labels # follow the same stable permutation. @@ -405,7 +445,7 @@ def test_type_map_remap_by_name(self): ] _write_raw_lmdb(self.lmdb_path, frames, ["H", "O"]) ms = dpdata.MultiSystems.from_file( - self.lmdb_path, fmt="lmdb", type_map=["O", "N", "C", "H"] + self.lmdb_path, fmt="deepmd/lmdb", type_map=["O", "N", "C", "H"] ) s = ms[0] names = [str(s.data["atom_names"][t]) for t in s.data["atom_types"]] @@ -421,11 +461,11 @@ def test_type_map_missing_element_raises(self): } ] _write_raw_lmdb(self.lmdb_path, frames, ["O", "H"]) - from dpdata.formats.lmdb.format import LMDBError + from dpdata.formats.deepmd.lmdb.format import LMDBError with self.assertRaises(LMDBError): dpdata.MultiSystems.from_file( - self.lmdb_path, fmt="lmdb", type_map=["C", "N"] + self.lmdb_path, fmt="deepmd/lmdb", type_map=["C", "N"] ) def test_inconsistent_keys_raise(self): @@ -449,7 +489,7 @@ def test_inconsistent_keys_raise(self): ] _write_raw_lmdb(self.lmdb_path, frames, ["O", "H"]) with self.assertRaisesRegex(LMDBFrameError, "must contain identical fields"): - dpdata.MultiSystems.from_file(self.lmdb_path, fmt="lmdb") + dpdata.MultiSystems.from_file(self.lmdb_path, fmt="deepmd/lmdb") def test_single_system_multi_composition_warns(self): frames = [ @@ -468,7 +508,7 @@ def test_single_system_multi_composition_warns(self): ] _write_raw_lmdb(self.lmdb_path, frames, ["C", "H"]) with self.assertWarns(UserWarning): - ls = dpdata.LabeledSystem(self.lmdb_path, fmt="lmdb") + ls = dpdata.LabeledSystem(self.lmdb_path, fmt="deepmd/lmdb") # only the first composition is returned self.assertEqual(ls.get_natoms(), 3) @@ -494,7 +534,7 @@ def test_same_composition_different_atom_order_is_canonicalized(self): }, ] _write_raw_lmdb(self.lmdb_path, frames, ["H", "O"]) - ms = dpdata.MultiSystems.from_file(self.lmdb_path, fmt="lmdb") + ms = dpdata.MultiSystems.from_file(self.lmdb_path, fmt="deepmd/lmdb") self.assertEqual(len(ms), 1) system = ms[0] self.assertEqual(system.get_nframes(), 2) @@ -524,7 +564,7 @@ def test_same_composition_mixed_pbc_raises(self): ] _write_raw_lmdb(self.lmdb_path, frames, ["O", "H"]) with self.assertRaisesRegex(LMDBFrameError, "must contain identical fields"): - dpdata.MultiSystems.from_file(self.lmdb_path, fmt="lmdb") + dpdata.MultiSystems.from_file(self.lmdb_path, fmt="deepmd/lmdb") def test_max_frames_guard(self): frames = [ @@ -539,9 +579,11 @@ def test_max_frames_guard(self): ] _write_raw_lmdb(self.lmdb_path, frames, ["H"]) with self.assertRaisesRegex(LMDBError, "exceeding max_frames"): - dpdata.MultiSystems.from_file(self.lmdb_path, fmt="lmdb", max_frames=1) + dpdata.MultiSystems.from_file( + self.lmdb_path, fmt="deepmd/lmdb", max_frames=1 + ) loaded = dpdata.MultiSystems.from_file( - self.lmdb_path, fmt="lmdb", max_frames=None, labeled=False + self.lmdb_path, fmt="deepmd/lmdb", max_frames=None, labeled=False ) self.assertEqual(loaded.get_nframes(), 2) @@ -561,7 +603,7 @@ def test_big_endian_dtype_preserved(self): }, ] _write_raw_lmdb(self.lmdb_path, frames, ["H"]) - system = dpdata.MultiSystems.from_file(self.lmdb_path, fmt="lmdb")[0] + system = dpdata.MultiSystems.from_file(self.lmdb_path, fmt="deepmd/lmdb")[0] self.assertEqual(system["coords"].dtype.str, ">f8") self.assertEqual(system["cells"].dtype.str, ">f8") self.assertEqual(system["energies"].dtype.str, ">f8") @@ -580,7 +622,7 @@ def test_reference_fparam_shape_does_not_collide_with_natoms(self): }, ] _write_raw_lmdb(self.lmdb_path, frames, ["H"]) - dpdata.MultiSystems.from_file(self.lmdb_path, fmt="lmdb", labeled=False) + dpdata.MultiSystems.from_file(self.lmdb_path, fmt="deepmd/lmdb", labeled=False) dtype = next(dt for dt in dpdata.System.DTYPES if dt.name == "fparam") shape = dtype.shape self.assertIsNotNone(shape) @@ -617,7 +659,7 @@ def test_flattened_reference_aparam_is_normalized_and_reordered(self): ] _write_raw_lmdb(self.lmdb_path, frames, ["H", "O"]) system = dpdata.MultiSystems.from_file( - self.lmdb_path, fmt="lmdb", labeled=False + self.lmdb_path, fmt="deepmd/lmdb", labeled=False )[0] np.testing.assert_array_equal(system["aparam"], expected) registered = next( @@ -638,7 +680,7 @@ def test_reference_spin_key_maps_to_dpdata_name(self): ] _write_raw_lmdb(self.lmdb_path, frames, ["H"]) system = dpdata.MultiSystems.from_file( - self.lmdb_path, fmt="lmdb", labeled=False + self.lmdb_path, fmt="deepmd/lmdb", labeled=False )[0] self.assertIn("spins", system.data) self.assertNotIn("spin", system.data) @@ -663,7 +705,9 @@ def test_known_field_shape_hint_cannot_remove_atom_axis(self): }, ) with self.assertRaisesRegex(LMDBMetadataError, "changes its protocol shape"): - dpdata.MultiSystems.from_file(self.lmdb_path, fmt="lmdb", labeled=False) + dpdata.MultiSystems.from_file( + self.lmdb_path, fmt="deepmd/lmdb", labeled=False + ) def test_data_name_hint_must_match_registered_protocol(self): frames = [ @@ -680,7 +724,9 @@ def test_data_name_hint_must_match_registered_protocol(self): metadata_extra={"dp_data_names": {"foo": "spins"}}, ) with self.assertRaises(LMDBMetadataError) as context: - dpdata.MultiSystems.from_file(self.lmdb_path, fmt="lmdb", labeled=False) + dpdata.MultiSystems.from_file( + self.lmdb_path, fmt="deepmd/lmdb", labeled=False + ) self.assertIn("foo", str(context.exception)) self.assertIn("spins", str(context.exception)) @@ -694,7 +740,9 @@ def test_custom_deepmd_core_key_rejected_on_read(self): ] _write_raw_lmdb(self.lmdb_path, frames, ["H"]) with self.assertRaisesRegex(LMDBFrameError, "reserved LMDB key 'coord'"): - dpdata.MultiSystems.from_file(self.lmdb_path, fmt="lmdb", labeled=False) + dpdata.MultiSystems.from_file( + self.lmdb_path, fmt="deepmd/lmdb", labeled=False + ) def test_unknown_atomic_axis_inferred_across_atom_counts(self): frames = [ @@ -711,7 +759,7 @@ def test_unknown_atomic_axis_inferred_across_atom_counts(self): ] _write_raw_lmdb(self.lmdb_path, frames, ["H"]) systems = dpdata.MultiSystems.from_file( - self.lmdb_path, fmt="lmdb", labeled=False + self.lmdb_path, fmt="deepmd/lmdb", labeled=False ) self.assertEqual(systems.get_nframes(), 2) dtype = next(dt for dt in dpdata.System.DTYPES if dt.name == "mystery") @@ -734,7 +782,9 @@ def test_unknown_same_nloc_atom_axis_is_rejected_as_ambiguous(self): ] _write_raw_lmdb(self.lmdb_path, frames, ["H", "O"]) with self.assertRaisesRegex(LMDBFrameError, "ambiguous without dp_data_shapes"): - dpdata.MultiSystems.from_file(self.lmdb_path, fmt="lmdb", labeled=False) + dpdata.MultiSystems.from_file( + self.lmdb_path, fmt="deepmd/lmdb", labeled=False + ) def test_process_global_schema_conflict_raises(self): second_path = "tmp_robust_second.lmdb" @@ -755,9 +805,13 @@ def test_process_global_schema_conflict_raises(self): ] _write_raw_lmdb(self.lmdb_path, first_frames, ["H"]) _write_raw_lmdb(second_path, second_frames, ["H"]) - dpdata.MultiSystems.from_file(self.lmdb_path, fmt="lmdb", labeled=False) + dpdata.MultiSystems.from_file( + self.lmdb_path, fmt="deepmd/lmdb", labeled=False + ) with self.assertRaisesRegex(LMDBError, "process-global definition"): - dpdata.MultiSystems.from_file(second_path, fmt="lmdb", labeled=False) + dpdata.MultiSystems.from_file( + second_path, fmt="deepmd/lmdb", labeled=False + ) finally: if os.path.exists(second_path): shutil.rmtree(second_path) @@ -781,7 +835,9 @@ def test_failed_read_does_not_partially_register_data_types(self): ] _write_raw_lmdb(self.lmdb_path, frames, ["H"]) with self.assertRaisesRegex(LMDBFrameError, "must contain identical fields"): - dpdata.MultiSystems.from_file(self.lmdb_path, fmt="lmdb", labeled=False) + dpdata.MultiSystems.from_file( + self.lmdb_path, fmt="deepmd/lmdb", labeled=False + ) self.assertNotIn("new_scalar", [dt.name for dt in dpdata.System.DTYPES]) def test_caller_construction_failure_rolls_back_registration(self): @@ -794,7 +850,7 @@ def test_caller_construction_failure_rolls_back_registration(self): ] _write_raw_lmdb(self.lmdb_path, frames, ["H"]) with self.assertRaises(DataError): - dpdata.MultiSystems.from_file(self.lmdb_path, fmt="lmdb") + dpdata.MultiSystems.from_file(self.lmdb_path, fmt="deepmd/lmdb") self.assertNotIn( "unlabeled_custom", [dt.name for dt in dpdata.System.DTYPES], @@ -814,7 +870,7 @@ def test_direct_labeled_read_rejects_unlabeled_before_registration(self): ] _write_raw_lmdb(self.lmdb_path, frames, ["H"]) with self.assertRaisesRegex(LMDBFrameError, "required by LabeledSystem"): - dpdata.LabeledSystem(self.lmdb_path, fmt="lmdb") + dpdata.LabeledSystem(self.lmdb_path, fmt="deepmd/lmdb") self.assertNotIn( "unlabeled_custom", [dt.name for dt in dpdata.System.DTYPES], @@ -844,7 +900,7 @@ def _remove_temporary_directories(self): shutil.rmtree(path) def test_system_ids_preserved(self): - from dpdata.formats.lmdb import dump_systems + from dpdata.formats.deepmd.lmdb import dump_systems dump_systems([self.A, self.B, self.C], self.lmdb_path) with lmdb.open(self.lmdb_path, readonly=True, lock=False) as env: @@ -860,7 +916,7 @@ def test_system_ids_preserved(self): def test_contrast_with_multisystems_merge(self): # the default MultiSystems path merges A and C by formula. ms = dpdata.MultiSystems(self.A, self.B, self.C) - ms.to("lmdb", self.lmdb_path) + ms.to("deepmd/lmdb", self.lmdb_path) with lmdb.open(self.lmdb_path, readonly=True, lock=False) as env: with env.begin() as txn: meta = msgpack.unpackb(txn.get(b"__metadata__"), raw=False) @@ -868,14 +924,14 @@ def test_contrast_with_multisystems_merge(self): self.assertEqual(sorted(set(meta["frame_system_ids"])), [0, 1]) def test_roundtrip_total_frames(self): - from dpdata.formats.lmdb import dump_systems + from dpdata.formats.deepmd.lmdb import dump_systems dump_systems([self.A, self.B, self.C], self.lmdb_path) - ms = dpdata.MultiSystems.from_file(self.lmdb_path, fmt="lmdb") + ms = dpdata.MultiSystems.from_file(self.lmdb_path, fmt="deepmd/lmdb") self.assertEqual(ms.get_nframes(), 6) def test_generator_with_type_map(self): - from dpdata.formats.lmdb import dump_systems + from dpdata.formats.deepmd.lmdb import dump_systems def gen(): yield self.A @@ -890,7 +946,7 @@ def gen(): self.assertEqual(meta["frame_system_ids"], [0, 0, 1, 2, 2, 2]) def test_empty_input_rejected(self): - from dpdata.formats.lmdb import dump_systems + from dpdata.formats.deepmd.lmdb import dump_systems with self.assertRaisesRegex(LMDBError, "empty"): dump_systems([], self.lmdb_path) @@ -898,17 +954,17 @@ def test_empty_input_rejected(self): def test_empty_multisystems_rejected(self): with self.assertRaisesRegex(LMDBError, "empty"): - dpdata.MultiSystems().to("lmdb", self.lmdb_path) + dpdata.MultiSystems().to("deepmd/lmdb", self.lmdb_path) self.assertFalse(os.path.exists(self.lmdb_path)) def test_multisystems_input_rejected(self): - from dpdata.formats.lmdb import dump_systems + from dpdata.formats.deepmd.lmdb import dump_systems with self.assertRaisesRegex(TypeError, "original ordered systems"): dump_systems(dpdata.MultiSystems(self.A, self.B), self.lmdb_path) def test_negative_standard_atom_type_rejected(self): - from dpdata.formats.lmdb import dump_systems + from dpdata.formats.deepmd.lmdb import dump_systems data = self.A.data.copy() data["atom_types"] = np.array([-1, 1, 1]) @@ -917,7 +973,7 @@ def test_negative_standard_atom_type_rejected(self): self.assertFalse(os.path.exists(self.lmdb_path)) def test_floating_atom_type_rejected(self): - from dpdata.formats.lmdb import dump_systems + from dpdata.formats.deepmd.lmdb import dump_systems data = self.A.data.copy() data["atom_types"] = np.array([0.0, 1.0, 1.0]) @@ -925,7 +981,7 @@ def test_floating_atom_type_rejected(self): dump_systems([data], self.lmdb_path) def test_uint64_overflow_not_treated_as_virtual_atom(self): - from dpdata.formats.lmdb import dump_systems + from dpdata.formats.deepmd.lmdb import dump_systems data = { "atom_numbs": [1], @@ -941,7 +997,7 @@ def test_uint64_overflow_not_treated_as_virtual_atom(self): dump_systems([data], self.lmdb_path) def test_inconsistent_atom_numbs_rejected(self): - from dpdata.formats.lmdb import dump_systems + from dpdata.formats.deepmd.lmdb import dump_systems data = self.A.data.copy() data["atom_numbs"] = [2, 1] @@ -949,7 +1005,7 @@ def test_inconsistent_atom_numbs_rejected(self): dump_systems([data], self.lmdb_path) def test_nonportable_array_dtypes_rejected(self): - from dpdata.formats.lmdb import dump_systems + from dpdata.formats.deepmd.lmdb import dump_systems for dtype in ( object, @@ -963,7 +1019,7 @@ def test_nonportable_array_dtypes_rejected(self): self.assertFalse(os.path.exists(self.lmdb_path)) def test_mixed_raw_virtual_atoms_are_removed(self): - from dpdata.formats.lmdb import dump_systems + from dpdata.formats.deepmd.lmdb import dump_systems data = { "atom_numbs": [3], @@ -989,7 +1045,7 @@ def test_mixed_raw_virtual_atoms_are_removed(self): self.assertEqual(frame["forces"]["shape"], [2, 3]) def test_ambiguous_mixed_custom_field_rejected(self): - from dpdata.formats.lmdb import dump_systems + from dpdata.formats.deepmd.lmdb import dump_systems data = { "atom_numbs": [3], @@ -1006,19 +1062,19 @@ def test_ambiguous_mixed_custom_field_rejected(self): dump_systems([data], self.lmdb_path) def test_existing_destination_requires_overwrite(self): - from dpdata.formats.lmdb import dump_systems + from dpdata.formats.deepmd.lmdb import dump_systems dump_systems([self.A], self.lmdb_path) with self.assertRaises(FileExistsError): dump_systems([self.B], self.lmdb_path) def test_windows_overwrite_rejected_explicitly(self): - from dpdata.formats.lmdb import dump_systems + from dpdata.formats.deepmd.lmdb import dump_systems dump_systems([self.A], self.lmdb_path) with ( mock.patch( - "dpdata.formats.lmdb.format._IS_WINDOWS", + "dpdata.formats.deepmd.lmdb.format._IS_WINDOWS", True, ), self.assertRaisesRegex(NotImplementedError, "not supported on Windows"), @@ -1031,7 +1087,7 @@ def test_windows_overwrite_rejected_explicitly(self): @unittest.skipIf(os.name == "nt", "LMDB overwrite is POSIX-only") def test_failed_overwrite_preserves_old_database(self): - from dpdata.formats.lmdb import dump_systems + from dpdata.formats.deepmd.lmdb import dump_systems old = _make_labeled_system([0], ["H"], 1) old.data["energies"][:] = 10.0 @@ -1049,7 +1105,7 @@ def test_failed_overwrite_preserves_old_database(self): overwrite=True, write_batch_size=1, ) - loaded = dpdata.LabeledSystem(self.lmdb_path, fmt="lmdb") + loaded = dpdata.LabeledSystem(self.lmdb_path, fmt="deepmd/lmdb") np.testing.assert_array_equal(loaded["energies"], [10.0]) temporary = list( Path(self.lmdb_path).parent.glob(f".{Path(self.lmdb_path).name}.tmp-*") @@ -1058,7 +1114,7 @@ def test_failed_overwrite_preserves_old_database(self): @unittest.skipIf(os.name == "nt", "LMDB overwrite is POSIX-only") def test_successful_overwrite_replaces_database(self): - from dpdata.formats.lmdb import dump_systems + from dpdata.formats.deepmd.lmdb import dump_systems dump_systems([self.A], self.lmdb_path) dump_systems( @@ -1067,7 +1123,7 @@ def test_successful_overwrite_replaces_database(self): overwrite=True, write_batch_size=1, ) - loaded = dpdata.MultiSystems.from_file(self.lmdb_path, fmt="lmdb") + loaded = dpdata.MultiSystems.from_file(self.lmdb_path, fmt="deepmd/lmdb") self.assertEqual(loaded.get_nframes(), 1) self.assertEqual(loaded[0].get_natoms(), 5) with lmdb.open(self.lmdb_path, readonly=True, lock=False) as env: @@ -1076,7 +1132,7 @@ def test_successful_overwrite_replaces_database(self): self.assertIsNone(txn.get(b"000000000001")) def test_duplicate_type_map_rejected(self): - from dpdata.formats.lmdb import dump_systems + from dpdata.formats.deepmd.lmdb import dump_systems with self.assertRaisesRegex(LMDBError, "duplicate"): dump_systems( @@ -1086,7 +1142,7 @@ def test_duplicate_type_map_rejected(self): ) def test_invalid_write_batch_size_rejected_without_artifact(self): - from dpdata.formats.lmdb import dump_systems + from dpdata.formats.deepmd.lmdb import dump_systems with self.assertRaisesRegex(ValueError, "write_batch_size"): dump_systems( @@ -1097,7 +1153,7 @@ def test_invalid_write_batch_size_rejected_without_artifact(self): self.assertFalse(os.path.exists(self.lmdb_path)) def test_map_full_grows_and_retries_batch(self): - from dpdata.formats.lmdb import dump_systems + from dpdata.formats.deepmd.lmdb import dump_systems large = _make_labeled_system([0] * 128, ["H"], 8) dump_systems( @@ -1106,14 +1162,14 @@ def test_map_full_grows_and_retries_batch(self): map_size=4096, write_batch_size=3, ) - loaded = dpdata.LabeledSystem(self.lmdb_path, fmt="lmdb") + loaded = dpdata.LabeledSystem(self.lmdb_path, fmt="deepmd/lmdb") self.assertEqual(loaded.get_nframes(), 8) with lmdb.open(self.lmdb_path, readonly=True, lock=False) as env: self.assertGreater(env.info()["map_size"], 4096) @unittest.skipIf(os.name == "nt", "LMDB overwrite is POSIX-only") def test_active_dpdata_reader_blocks_overwrite_for_path_alias(self): - from dpdata.formats.lmdb import dump_systems + from dpdata.formats.deepmd.lmdb import dump_systems dump_systems([self.A], self.lmdb_path) resolved, _ = _open_read_env(str(Path(self.lmdb_path).resolve())) @@ -1126,12 +1182,12 @@ def test_active_dpdata_reader_blocks_overwrite_for_path_alias(self): ) finally: _close_read_env(resolved) - loaded = dpdata.MultiSystems.from_file(self.lmdb_path, fmt="lmdb") + loaded = dpdata.MultiSystems.from_file(self.lmdb_path, fmt="deepmd/lmdb") self.assertEqual(loaded.get_nframes(), self.A.get_nframes()) @unittest.skipIf(os.name == "nt", "LMDB overwrite is POSIX-only") def test_external_lmdb_reader_blocks_overwrite(self): - from dpdata.formats.lmdb import dump_systems + from dpdata.formats.deepmd.lmdb import dump_systems dump_systems([self.A], self.lmdb_path) external_env = lmdb.open( @@ -1152,12 +1208,12 @@ def test_external_lmdb_reader_blocks_overwrite(self): ) finally: external_env.close() - loaded = dpdata.MultiSystems.from_file(self.lmdb_path, fmt="lmdb") + loaded = dpdata.MultiSystems.from_file(self.lmdb_path, fmt="deepmd/lmdb") self.assertEqual(loaded.get_nframes(), self.A.get_nframes()) @unittest.skipIf(os.name == "nt", "LMDB overwrite is POSIX-only") def test_publish_guard_closes_external_reader_race(self): - from dpdata.formats.lmdb import dump_systems + from dpdata.formats.deepmd.lmdb import dump_systems dump_systems([self.A], self.lmdb_path) replace_entered = threading.Event() @@ -1193,7 +1249,7 @@ def synchronized_replace(source, destination): thread.start() try: with mock.patch( - "dpdata.formats.lmdb.format.os.replace", + "dpdata.formats.deepmd.lmdb.format.os.replace", side_effect=synchronized_replace, ): dump_systems( @@ -1204,13 +1260,13 @@ def synchronized_replace(source, destination): finally: thread.join(timeout=5) self.assertEqual(reader_was_blocked, [True]) - loaded = dpdata.MultiSystems.from_file(self.lmdb_path, fmt="lmdb") + loaded = dpdata.MultiSystems.from_file(self.lmdb_path, fmt="deepmd/lmdb") self.assertEqual(loaded[0].get_natoms(), self.B.get_natoms()) @unittest.skipIf(os.name == "nt", "LMDB overwrite is POSIX-only") def test_staged_validation_failure_preserves_old_database(self): - from dpdata.formats.lmdb import dump_systems - from dpdata.formats.lmdb.format import _LMDBWriter + from dpdata.formats.deepmd.lmdb import dump_systems + from dpdata.formats.deepmd.lmdb.format import _LMDBWriter dump_systems([self.A], self.lmdb_path) with ( @@ -1226,12 +1282,12 @@ def test_staged_validation_failure_preserves_old_database(self): self.lmdb_path, overwrite=True, ) - loaded = dpdata.MultiSystems.from_file(self.lmdb_path, fmt="lmdb") + loaded = dpdata.MultiSystems.from_file(self.lmdb_path, fmt="deepmd/lmdb") self.assertEqual(loaded.get_nframes(), self.A.get_nframes()) @unittest.skipUnless(hasattr(os, "fork"), "requires os.fork") def test_read_cache_resets_after_fork(self): - from dpdata.formats.lmdb import dump_systems + from dpdata.formats.deepmd.lmdb import dump_systems dump_systems([self.A], self.lmdb_path) resolved, _ = _open_read_env(self.lmdb_path) @@ -1276,11 +1332,11 @@ def tearDown(self): def test_missing_metadata(self): lmdb.open(self.lmdb_missing_meta, map_size=1 << 30).close() with self.assertRaises(LMDBMetadataError): - dpdata.MultiSystems.from_file(self.lmdb_missing_meta, fmt="lmdb") + dpdata.MultiSystems.from_file(self.lmdb_missing_meta, fmt="deepmd/lmdb") def test_missing_frame(self): with self.assertRaises(LMDBFrameError): - dpdata.MultiSystems.from_file(self.lmdb_missing_frame, fmt="lmdb") + dpdata.MultiSystems.from_file(self.lmdb_missing_frame, fmt="deepmd/lmdb") def test_non_mapping_metadata_rejected(self): path = "tmp_bad_metadata.lmdb" @@ -1290,7 +1346,7 @@ def test_non_mapping_metadata_rejected(self): txn.put(b"__metadata__", _reference_packb(["not", "a", "mapping"])) env.close() with self.assertRaisesRegex(LMDBMetadataError, "must contain a mapping"): - dpdata.MultiSystems.from_file(path, fmt="lmdb") + dpdata.MultiSystems.from_file(path, fmt="deepmd/lmdb") finally: if os.path.exists(path): shutil.rmtree(path) @@ -1305,7 +1361,7 @@ def test_invalid_msgpack_metadata_wrapped(self): with self.assertRaisesRegex( LMDBMetadataError, "Cannot decode __metadata__" ): - dpdata.MultiSystems.from_file(path, fmt="lmdb") + dpdata.MultiSystems.from_file(path, fmt="deepmd/lmdb") finally: if os.path.exists(path): shutil.rmtree(path) @@ -1329,7 +1385,7 @@ def test_fractional_nframes_rejected(self): with self.assertRaisesRegex( LMDBMetadataError, "nframes must be an integer" ): - dpdata.MultiSystems.from_file(path, fmt="lmdb") + dpdata.MultiSystems.from_file(path, fmt="deepmd/lmdb") finally: if os.path.exists(path): shutil.rmtree(path) @@ -1357,7 +1413,7 @@ def test_non_integer_metadata_vectors_rejected(self): with self.assertRaisesRegex( LMDBMetadataError, "must be an integer" ): - dpdata.MultiSystems.from_file(path, fmt="lmdb") + dpdata.MultiSystems.from_file(path, fmt="deepmd/lmdb") finally: if os.path.exists(path): shutil.rmtree(path) @@ -1389,7 +1445,7 @@ def test_malformed_array_payload_rejected(self): ) env.close() with self.assertRaisesRegex(LMDBFrameError, "Cannot decode array"): - dpdata.MultiSystems.from_file(path, fmt="lmdb") + dpdata.MultiSystems.from_file(path, fmt="deepmd/lmdb") finally: if os.path.exists(path): shutil.rmtree(path) @@ -1398,7 +1454,7 @@ def test_invalid_frame_system_ids_length(self): path = "tmp_invalid_system_ids.lmdb" try: system = _make_labeled_system([0], ["H"], 1) - system.to("lmdb", path) + system.to("deepmd/lmdb", path) env = lmdb.open(path, map_size=1 << 30) with env.begin(write=True) as txn: metadata = msgpack.unpackb(txn.get(b"__metadata__"), raw=False) @@ -1409,7 +1465,7 @@ def test_invalid_frame_system_ids_length(self): ) env.close() with self.assertRaisesRegex(LMDBMetadataError, "frame_system_ids length"): - dpdata.MultiSystems.from_file(path, fmt="lmdb") + dpdata.MultiSystems.from_file(path, fmt="deepmd/lmdb") finally: if os.path.exists(path): shutil.rmtree(path) @@ -1417,13 +1473,13 @@ def test_invalid_frame_system_ids_length(self): def test_existing_external_reader_has_actionable_error(self): path = "tmp_external_reader.lmdb" try: - _make_labeled_system([0], ["H"], 1).to("lmdb", path) + _make_labeled_system([0], ["H"], 1).to("deepmd/lmdb", path) env = lmdb.open(path, readonly=True, lock=False) try: with self.assertRaisesRegex( LMDBError, "another library has already opened" ): - dpdata.MultiSystems.from_file(path, fmt="lmdb") + dpdata.MultiSystems.from_file(path, fmt="deepmd/lmdb") finally: env.close() finally: @@ -1442,12 +1498,12 @@ def tearDown(self): def test_custom_frame_idx_fmt(self): ms = dpdata.MultiSystems(self.system) - ms.to("lmdb", self.lmdb_path, frame_idx_fmt="06d") + ms.to("deepmd/lmdb", self.lmdb_path, frame_idx_fmt="06d") with lmdb.open(self.lmdb_path, readonly=True, lock=False) as env: with env.begin() as txn: self.assertIsNotNone(txn.get(b"000000")) self.assertIsNone(txn.get(b"000000000000")) - loaded = dpdata.LabeledSystem(self.lmdb_path, fmt="lmdb") + loaded = dpdata.LabeledSystem(self.lmdb_path, fmt="deepmd/lmdb") self.assertEqual(len(loaded), len(self.system)) np.testing.assert_allclose(loaded.data["coords"], self.system.data["coords"]) diff --git a/tests/test_lmdb_custom_dtype.py b/tests/test_lmdb_custom_dtype.py index 2eb83b75..fe5da427 100644 --- a/tests/test_lmdb_custom_dtype.py +++ b/tests/test_lmdb_custom_dtype.py @@ -10,7 +10,7 @@ import dpdata from dpdata.data_type import Axis, DataType -from dpdata.formats.lmdb.format import LMDBError, LMDBFrameError +from dpdata.formats.deepmd.lmdb.format import LMDBError, LMDBFrameError class TestLMDBFrameData(unittest.TestCase): @@ -48,21 +48,21 @@ def tearDown(self): shutil.rmtree(self.lmdb_path) def test_frame_data_preservation(self): - self.system.to("lmdb", self.lmdb_path) - loaded = dpdata.LabeledSystem(self.lmdb_path, fmt="lmdb") + self.system.to("deepmd/lmdb", self.lmdb_path) + loaded = dpdata.LabeledSystem(self.lmdb_path, fmt="deepmd/lmdb") np.testing.assert_allclose( loaded.data["frame_data"], self.system.data["frame_data"] ) def test_frame_data_auto_registration(self): - self.system.to("lmdb", self.lmdb_path) + self.system.to("deepmd/lmdb", self.lmdb_path) # simulate a clean session dpdata.System.DTYPES = self.original_system_dtypes dpdata.LabeledSystem.DTYPES = self.original_labeled_system_dtypes self.assertNotIn("frame_data", [dt.name for dt in dpdata.LabeledSystem.DTYPES]) - loaded = dpdata.LabeledSystem(self.lmdb_path, fmt="lmdb") + loaded = dpdata.LabeledSystem(self.lmdb_path, fmt="deepmd/lmdb") self.assertIn("frame_data", [dt.name for dt in dpdata.LabeledSystem.DTYPES]) np.testing.assert_allclose( loaded.data["frame_data"], self.system.data["frame_data"] @@ -111,18 +111,18 @@ def tearDown(self): shutil.rmtree(self.lmdb_path) def test_fparam_aparam_preservation(self): - self.system.to("lmdb", self.lmdb_path) - loaded = dpdata.LabeledSystem(self.lmdb_path, fmt="lmdb") + self.system.to("deepmd/lmdb", self.lmdb_path) + loaded = dpdata.LabeledSystem(self.lmdb_path, fmt="deepmd/lmdb") np.testing.assert_allclose(loaded.data["fparam"], self.system.data["fparam"]) np.testing.assert_allclose(loaded.data["aparam"], self.system.data["aparam"]) def test_fparam_aparam_auto_registration(self): - self.system.to("lmdb", self.lmdb_path) + self.system.to("deepmd/lmdb", self.lmdb_path) dpdata.System.DTYPES = self.original_system_dtypes dpdata.LabeledSystem.DTYPES = self.original_labeled_system_dtypes - loaded = dpdata.LabeledSystem(self.lmdb_path, fmt="lmdb") + loaded = dpdata.LabeledSystem(self.lmdb_path, fmt="deepmd/lmdb") self.assertIn("fparam", [dt.name for dt in dpdata.LabeledSystem.DTYPES]) self.assertIn("aparam", [dt.name for dt in dpdata.LabeledSystem.DTYPES]) np.testing.assert_allclose(loaded.data["fparam"], self.system.data["fparam"]) @@ -132,12 +132,12 @@ def test_symbolic_axis_natoms_preservation(self): # natoms == 3 collides with the trailing coordinate dim; the stored # symbolic shape must still recover (NFRAMES, NATOMS, 3), not # (NFRAMES, NATOMS, NATOMS). - self.system.to("lmdb", self.lmdb_path) + self.system.to("deepmd/lmdb", self.lmdb_path) dpdata.System.DTYPES = self.original_system_dtypes dpdata.LabeledSystem.DTYPES = self.original_labeled_system_dtypes - dpdata.LabeledSystem(self.lmdb_path, fmt="lmdb") + dpdata.LabeledSystem(self.lmdb_path, fmt="deepmd/lmdb") aparam_dt = next( dt for dt in dpdata.LabeledSystem.DTYPES if dt.name == "aparam" ) @@ -198,7 +198,7 @@ def test_deepmd_name_used_on_disk_and_restored(self): dtype=float, ).reshape(self.system.get_nframes(), self.system.get_natoms(), 3) expected = self.system.data["spins"].copy() - self.system.to("lmdb", self.lmdb_path) + self.system.to("deepmd/lmdb", self.lmdb_path) with lmdb.open(self.lmdb_path, readonly=True, lock=False) as env: with env.begin() as txn: @@ -210,7 +210,7 @@ def test_deepmd_name_used_on_disk_and_restored(self): dpdata.System.DTYPES = self.original_system_dtypes dpdata.LabeledSystem.DTYPES = self.original_labeled_system_dtypes - loaded = dpdata.LabeledSystem(self.lmdb_path, fmt="lmdb") + loaded = dpdata.LabeledSystem(self.lmdb_path, fmt="deepmd/lmdb") np.testing.assert_array_equal(loaded["spins"], expected) registered = next( dt for dt in dpdata.LabeledSystem.DTYPES if dt.name == "spins" @@ -240,7 +240,7 @@ def test_duplicate_deepmd_name_rejected(self): LMDBError, "both map to", ): - self.system.to("lmdb", self.lmdb_path) + self.system.to("deepmd/lmdb", self.lmdb_path) def test_deepmd_core_name_collision_rejected(self): malicious = DataType( @@ -260,7 +260,7 @@ def test_deepmd_core_name_collision_rejected(self): 101.0, ) with self.assertRaisesRegex(LMDBError, "reserved LMDB key"): - self.system.to("lmdb", self.lmdb_path) + self.system.to("deepmd/lmdb", self.lmdb_path) def test_additional_protocol_alias_rejected(self): alias = DataType( @@ -279,7 +279,7 @@ def test_additional_protocol_alias_rejected(self): ) ) with self.assertRaisesRegex(LMDBError, "belongs to 'spins'"): - self.system.to("lmdb", self.lmdb_path) + self.system.to("deepmd/lmdb", self.lmdb_path) def test_atom_types_key_collision_rejected(self): malicious = DataType( @@ -294,21 +294,21 @@ def test_atom_types_key_collision_rejected(self): (self.system.get_nframes(), self.system.get_natoms()) ) with self.assertRaisesRegex(LMDBError, "reserved LMDB key"): - self.system.to("lmdb", self.lmdb_path) + self.system.to("deepmd/lmdb", self.lmdb_path) def test_static_field_roundtrip(self): static = DataType("static_data", np.ndarray, shape=(2,), required=False) self._register(static) self.system.data["static_data"] = np.array([1.0, 2.0]) - self.system.to("lmdb", self.lmdb_path) - loaded = dpdata.LabeledSystem(self.lmdb_path, fmt="lmdb") + self.system.to("deepmd/lmdb", self.lmdb_path) + loaded = dpdata.LabeledSystem(self.lmdb_path, fmt="deepmd/lmdb") np.testing.assert_array_equal(loaded["static_data"], [1.0, 2.0]) def test_inconsistent_static_field_raises(self): static = DataType("static_data", np.ndarray, shape=(2,), required=False) self._register(static) self.system.data["static_data"] = np.array([1.0, 2.0]) - self.system.to("lmdb", self.lmdb_path) + self.system.to("deepmd/lmdb", self.lmdb_path) env = lmdb.open(self.lmdb_path, map_size=1 << 30) with env.begin(write=True) as txn: @@ -327,7 +327,7 @@ def test_inconsistent_static_field_raises(self): LMDBFrameError, "Static field", ): - dpdata.LabeledSystem(self.lmdb_path, fmt="lmdb") + dpdata.LabeledSystem(self.lmdb_path, fmt="deepmd/lmdb") def test_nonleading_frame_axis_roundtrip(self): transposed = DataType( @@ -341,8 +341,8 @@ def test_nonleading_frame_axis_roundtrip(self): 2, self.system.get_nframes() ) self.system.data["transposed_frames"] = values - self.system.to("lmdb", self.lmdb_path) - loaded = dpdata.LabeledSystem(self.lmdb_path, fmt="lmdb") + self.system.to("deepmd/lmdb", self.lmdb_path) + loaded = dpdata.LabeledSystem(self.lmdb_path, fmt="deepmd/lmdb") np.testing.assert_array_equal(loaded["transposed_frames"], values) def test_unused_shape_none_data_type_does_not_block_write(self): @@ -353,8 +353,8 @@ def test_unused_shape_none_data_type_does_not_block_write(self): required=False, ) self._register(undefined) - self.system.to("lmdb", self.lmdb_path) - loaded = dpdata.LabeledSystem(self.lmdb_path, fmt="lmdb") + self.system.to("deepmd/lmdb", self.lmdb_path) + loaded = dpdata.LabeledSystem(self.lmdb_path, fmt="deepmd/lmdb") self.assertEqual(loaded.get_nframes(), self.system.get_nframes()) def test_used_shape_none_data_type_rejected(self): @@ -369,7 +369,7 @@ def test_used_shape_none_data_type_rejected(self): (self.system.get_nframes(), 2) ) with self.assertRaisesRegex(LMDBError, "no declared shape"): - self.system.to("lmdb", self.lmdb_path) + self.system.to("deepmd/lmdb", self.lmdb_path) def test_multiple_atom_axes_remove_virtual_atoms(self): hessian = DataType( @@ -398,7 +398,7 @@ def test_multiple_atom_axes_remove_virtual_atoms(self): "energies": np.array([-1.0]), "hessian": np.zeros((1, 3, 3, 3, 3)), } - from dpdata.formats.lmdb import dump_systems + from dpdata.formats.deepmd.lmdb import dump_systems dump_systems([data], self.lmdb_path) with lmdb.open(self.lmdb_path, readonly=True, lock=False) as env: