fix: use lmdb_get to access already loaded data without re-reading

mattergen/evaluation/reference/reference_dataset_serializer.py

@@ -136,22 +136,20 @@ def __init__(self, lmdb_path: Path, cleanup_dir: bool = False):
         """
         self.env = lmdb_open(lmdb_path, readonly=True)
         self.num_entries_by_chemsys_reduced_formulas = (
-            self._build_num_entries_by_chemsys_reduced_formulas(lmdb_path)
+            self._build_num_entries_by_chemsys_reduced_formulas()
         )
         self.total_num_entries = sum(
             sum(d.values()) for d in self.num_entries_by_chemsys_reduced_formulas.values()
         )
         # close the LMDB environment when this object is garbage collected
         weakref.finalize(self, self._cleanup, self.env, cleanup_dir)
 
-    def _build_num_entries_by_chemsys_reduced_formulas(
-        self, lmdb_path: Path
-    ) -> dict[str, dict[str, int]]:
-        chemical_systems = lmdb_read_metadata(lmdb_path, "chemical_systems")
+    def _build_num_entries_by_chemsys_reduced_formulas(self) -> dict[str, dict[str, int]]:        
         result: defaultdict[str, dict[str, int]] = defaultdict(dict)
         with self.env.begin() as txn:
+            chemical_systems = lmdb_get(txn, "chemical_systems")
             for chemsys in chemical_systems:
-                reduced_formulas = lmdb_read_metadata(lmdb_path, f"{chemsys}.reduced_formulas")
+                reduced_formulas = lmdb_get(txn, f"{chemsys}.reduced_formulas")
                 for reduced_formula in reduced_formulas:
                     result[chemsys][reduced_formula] = lmdb_get(
                         txn, f"{chemsys}.{reduced_formula}.length"

mattergen/tests/test_reference_dataset_serializer.py

@@ -0,0 +1,32 @@
+from pathlib import Path
+
+from pymatgen.core import Lattice, Structure
+from pymatgen.entries.computed_entries import ComputedStructureEntry
+
+from mattergen.evaluation.reference.reference_dataset import ReferenceDataset
+from mattergen.evaluation.reference.reference_dataset_serializer import LMDBGZSerializer
+
+
+def test_deserialize_does_not_reopen_same_lmdb_while_environment_is_active(
+    tmp_path: Path,
+) -> None:
+    serializer = LMDBGZSerializer()
+    dataset_path = tmp_path / "reference.lmdb.gz"
+    entry = ComputedStructureEntry(
+        structure=Structure(
+            lattice=Lattice.cubic(3.5),
+            species=["Fe", "O"],
+            coords=[[0, 0, 0], [0.5, 0.5, 0.5]],
+        ),
+        energy=0.0,
+    )
+    serializer.serialize(
+        ReferenceDataset.from_entries("reference", [entry]),
+        dataset_path,
+    )
+
+    reference = serializer.deserialize(dataset_path)
+    assert reference.name == "reference"
+    assert reference.impl.chemical_systems == ("Fe-O",)
+    assert len(reference) == 1
+    reference.impl.cleanup(cleanup_dir=True)