Merge pull request #263 from AI-multimodal/dev/mc-fix-238

Resolve save multiplicities

lightshow/_tests/test_metadata.py

@@ -4,8 +4,8 @@
 from lightshow.parameters.feff import FEFFParameters
 
 
-def test_multiplicipty_writing(database_from_file, tmp_path):
-    target = Path(tmp_path) / Path("multi") / Path("destination")
+def test_multiplicity_writing(database_from_file, tmp_path):
+    target = Path(tmp_path) / "t"
     target.mkdir(exist_ok=True, parents=True)
     R = 10.0
     feff_parameters = FEFFParameters(
@@ -32,13 +32,12 @@ def test_multiplicipty_writing(database_from_file, tmp_path):
         pbar=False,
     )
     for k, v in database_from_file.metadata.items():
-        metadata_fn = target / Path(k) / Path("metadata.json")
-        with open(metadata_fn) as f:
-            d_metadata = json.load(f)
+        fname = target / k / Path("multiplicity.json")
+        with open(fname) as f:
+            mult = json.load(f)
         prim_meta = v["primitive"]
         for i_site, n_multi in zip(
             prim_meta["sites"], prim_meta["multiplicities"]
         ):
             i_site = str(i_site)
-            assert "multiplicities" in d_metadata
-            assert d_metadata["multiplicities"][i_site] == n_multi
+            assert mult[i_site] == n_multi

lightshow/database.py

@@ -283,8 +283,8 @@ def database_status(self):
     def __init__(
         self,
         structures,
-        metadata=dict(),
-        supercells=dict(),
+        metadata=None,
+        supercells=None,
         supercell_cutoff=None,
         inequivalent_sites_initialized=False,
         supercells_initialized=False,
@@ -293,8 +293,14 @@ def __init__(
         the classmethods to initialize this object."""
 
         self._structures = structures
-        self._metadata = metadata
-        self._supercells = supercells
+        if metadata is None:
+            self._metadata = {}
+        else:
+            self._metadata = metadata
+        if supercells is None:
+            self._supercells = {}
+        else:
+            self._supercells = supercells
         self._supercell_cutoff = supercell_cutoff
         self._inequivalent_sites_initialized = inequivalent_sites_initialized
         self._supercells_initialized = supercells_initialized
@@ -360,32 +366,21 @@ def _write_unit_cells(self, root, pbar=False):
             fname = Path(root) / key / "POSCAR"
             structure.to(fmt="POSCAR", filename=str(fname))
 
-    def _write_origin_paths(self, root, pbar=False):
-        """A helper method for writing important metadata for each of the
-        structures if the data was loaded from disk.
-
-        Parameters
-        ----------
-        root : os.PathLike
-        pbar : bool, optional
-        """
+    def _write_multiplicity(self, root, pbar=False):
+        """Helper method for writing the site multiplicity information of
+        the structure."""
 
         for key, metadata in tqdm(self._metadata.items(), disable=not pbar):
-            if "origin" not in metadata.keys():
-                continue
-            fname = Path(root) / key / "metadata.json"
-            origin = str(Path(metadata["origin"]).resolve())
-            new_metadata = {"origin": origin}
+            fname = Path(root) / key / "multiplicity.json"
             multiplicities = {
                 i_site: n_multi
                 for i_site, n_multi in zip(
                     metadata["primitive"]["sites"],
                     metadata["primitive"]["multiplicities"],
                 )
             }
-            new_metadata["multiplicities"] = multiplicities
             with open(fname, "w") as outfile:
-                json.dump(new_metadata, outfile, indent=4, sort_keys=True)
+                json.dump(multiplicities, outfile, indent=4, sort_keys=True)
 
     def write(
         self,
@@ -395,6 +390,7 @@ def write(
         pbar=True,
         copy_script=None,
         write_unit_cells=True,
+        write_multiplicity=True,
     ):
         """The core method of the :class:`.Database` class. This method will
         write all input files specified in the ``options`` parameter to disk.
@@ -452,6 +448,9 @@ def write(
         write_unit_cells : bool, optional
             If True, writes the unit cells in the materials directory in
             POSCAR format. Very useful!
+        write_multiplicity : bool, optional
+            If True, writes the multiplicities of each atom in the unit cell
+            to a multiplicities.json file.
         """
 
         self._setup_preliminary_attributes()
@@ -536,8 +535,8 @@ def write(
 
         if write_unit_cells:
             self._write_unit_cells(root, pbar=pbar)
-
-        self._write_origin_paths(root, pbar=pbar)
+        if write_multiplicity:
+            self._write_multiplicity(root, pbar=pbar)
 
         # Save a metadata file (not a serialized version of this class) to
         # disk along with the input files