Retain Structure.properties in structure_from_abivars()/structure_to_abivars() round trip (#3552)

* retain Structure.properties in structure_from_abivars()/structure_to_abivars() round trip

* test_structure.py add self.struct.properties["foo"] = "bar" for test_to_from_abivars()

Author: janosh

pymatgen/core/structure.py

@@ -2632,12 +2632,12 @@ def as_dataframe(self):
         return df
 
     @classmethod
-    def from_dict(cls, d: dict[str, Any], fmt: Literal["abivars"] | None = None) -> Structure:
+    def from_dict(cls, dct: dict[str, Any], fmt: Literal["abivars"] | None = None) -> Structure:
         """Reconstitute a Structure object from a dict representation of Structure
         created using as_dict().
 
         Args:
-            d (dict): Dict representation of structure.
+            dct (dict): Dict representation of structure.
             fmt ('abivars' | None): Use structure_from_abivars() to parse the dict. Defaults to None.
 
         Returns:
@@ -2646,12 +2646,12 @@ def from_dict(cls, d: dict[str, Any], fmt: Literal["abivars"] | None = None) ->
         if fmt == "abivars":
             from pymatgen.io.abinit.abiobjects import structure_from_abivars
 
-            return structure_from_abivars(cls=cls, **d)
+            return structure_from_abivars(cls=cls, **dct)
 
-        lattice = Lattice.from_dict(d["lattice"])
-        sites = [PeriodicSite.from_dict(sd, lattice) for sd in d["sites"]]
-        charge = d.get("charge")
-        return cls.from_sites(sites, charge=charge, properties=d.get("properties"))
+        lattice = Lattice.from_dict(dct["lattice"])
+        sites = [PeriodicSite.from_dict(sd, lattice) for sd in dct["sites"]]
+        charge = dct.get("charge")
+        return cls.from_sites(sites, charge=charge, properties=dct.get("properties"))
 
     def to(self, filename: str | Path = "", fmt: str = "", **kwargs) -> str:
         """Outputs the structure to a file or string.

pymatgen/io/abinit/abiobjects.py

@@ -95,12 +95,12 @@ def lattice_from_abivars(cls=None, *args, **kwargs):
     raise ValueError(f"Don't know how to construct a Lattice from dict:\n{pformat(kwargs)}")
 
 
-def structure_from_abivars(cls=None, *args, **kwargs):
+def structure_from_abivars(cls=None, *args, **kwargs) -> Structure:
     """
     Build a Structure object from a dictionary with ABINIT variables.
 
     Args:
-        cls: Structure class to be instantiated. pymatgen.core.structure.Structure if cls is None
+        cls: Structure class to be instantiated. Defaults to Structure.
 
     Example:
         al_structure = structure_from_abivars(
@@ -115,30 +115,28 @@ def structure_from_abivars(cls=None, *args, **kwargs):
     `xred` can be replaced with `xcart` or `xangst`.
     """
     kwargs.update(dict(*args))
-    d = kwargs
-
-    cls = Structure if cls is None else cls
+    cls = cls or Structure
 
     # lattice = Lattice.from_dict(d, fmt="abivars")
-    lattice = lattice_from_abivars(**d)
-    coords, coords_are_cartesian = d.get("xred"), False
+    lattice = lattice_from_abivars(**kwargs)
+    coords, coords_are_cartesian = kwargs.get("xred"), False
 
     if coords is None:
-        coords = d.get("xcart")
+        coords = kwargs.get("xcart")
         if coords is not None:
-            if "xangst" in d:
+            if "xangst" in kwargs:
                 raise ValueError("xangst and xcart are mutually exclusive")
             coords = ArrayWithUnit(coords, "bohr").to("ang")
         else:
-            coords = d.get("xangst")
+            coords = kwargs.get("xangst")
         coords_are_cartesian = True
 
     if coords is None:
-        raise ValueError(f"Cannot extract coordinates from:\n {d}")
+        raise ValueError(f"Cannot extract coordinates from:\n {kwargs}")
 
     coords = np.reshape(coords, (-1, 3))
 
-    znucl_type, typat = d["znucl"], d["typat"]
+    znucl_type, typat = kwargs["znucl"], kwargs["typat"]
 
     if not isinstance(znucl_type, Iterable):
         znucl_type = [znucl_type]
@@ -160,6 +158,7 @@ def structure_from_abivars(cls=None, *args, **kwargs):
         validate_proximity=False,
         to_unit_cell=False,
         coords_are_cartesian=coords_are_cartesian,
+        properties=kwargs.get("properties"),
     )
 
 
@@ -199,11 +198,10 @@ def structure_to_abivars(structure, enforce_znucl=None, enforce_typat=None, **kw
     """
     if not structure.is_ordered:
         raise ValueError(
-            """\
-Received disordered structure with partial occupancies that cannot be converted into an Abinit input.
-Please use OrderDisorderedStructureTransformation or EnumerateStructureTransformation
-to build an appropriate supercell from partial occupancies or, alternatively, use the Rigid Band Model
-or the Virtual Crystal Approximation."""
+            "Received disordered structure with partial occupancies that cannot be converted into an "
+            "Abinit input. Please use OrderDisorderedStructureTransformation or EnumerateStructureTransformation "
+            "to build an appropriate supercell from partial occupancies or, alternatively, use the Rigid Band Model "
+            "or the Virtual Crystal Approximation."
         )
 
     n_atoms = len(structure)
@@ -227,7 +225,6 @@ def structure_to_abivars(structure, enforce_znucl=None, enforce_typat=None, **kw
     if not enforce_order:
         types_of_specie = species_by_znucl(structure)
 
-        # [ntypat] list
         znucl_type = [specie.number for specie in types_of_specie]
         typat = np.zeros(n_atoms, int)
         for atm_idx, site in enumerate(structure):
@@ -236,49 +233,50 @@ def structure_to_abivars(structure, enforce_znucl=None, enforce_typat=None, **kw
         znucl_type = enforce_znucl
         typat = enforce_typat
 
-    rprim = ArrayWithUnit(structure.lattice.matrix, "ang").to("bohr")
-    angdeg = structure.lattice.angles
-    xred = np.reshape([site.frac_coords for site in structure], (-1, 3))
+    r_prim = ArrayWithUnit(structure.lattice.matrix, "ang").to("bohr")
+    ang_deg = structure.lattice.angles
+    x_red = np.reshape([site.frac_coords for site in structure], (-1, 3))
 
     # Set small values to zero. This usually happens when the CIF file
     # does not give structure parameters with enough digits.
-    rprim = np.where(np.abs(rprim) > 1e-8, rprim, 0.0)
-    xred = np.where(np.abs(xred) > 1e-8, xred, 0.0)
+    r_prim = np.where(np.abs(r_prim) > 1e-8, r_prim, 0.0)
+    x_red = np.where(np.abs(x_red) > 1e-8, x_red, 0.0)
 
     # Info on atoms.
     dct = {
         "natom": n_atoms,
         "ntypat": n_types_atom,
         "typat": typat,
         "znucl": znucl_type,
-        "xred": xred,
+        "xred": x_red,
+        "properties": structure.properties,
     }
 
     # Add info on the lattice.
     # Should we use (rprim, acell) or (angdeg, acell) to specify the lattice?
-    geomode = kwargs.pop("geomode", "rprim")
-    if geomode == "automatic":
-        geomode = "rprim"
+    geo_mode = kwargs.pop("geomode", "rprim")
+    if geo_mode == "automatic":
+        geo_mode = "rprim"
         if structure.lattice.is_hexagonal:  # or structure.lattice.is_rhombohedral
-            geomode = "angdeg"
-            angdeg = structure.lattice.angles
+            geo_mode = "angdeg"
+            ang_deg = structure.lattice.angles
             # Here one could polish a bit the numerical values if they are not exact.
             # Note that in pmg the angles are 12, 20, 01 while in Abinit 12, 02, 01
             # One should make sure that the orientation is preserved (see Curtarolo's settings)
 
-    if geomode == "rprim":
+    if geo_mode == "rprim":
         dct.update(
             acell=3 * [1.0],
-            rprim=rprim,
+            rprim=r_prim,
         )
 
-    elif geomode == "angdeg":
+    elif geo_mode == "angdeg":
         dct.update(
             acell=ArrayWithUnit(structure.lattice.abc, "ang").to("bohr"),
-            angdeg=angdeg,
+            angdeg=ang_deg,
         )
     else:
-        raise ValueError(f"Wrong value for {geomode=}")
+        raise ValueError(f"Wrong value for {geo_mode=}")
 
     return dct
 

pymatgen/io/abinit/variable.py

@@ -5,7 +5,7 @@
 import collections
 import collections.abc
 import string
-from collections.abc import Iterable
+from collections.abc import Sequence
 
 import numpy as np
 
@@ -23,7 +23,7 @@
 class InputVariable:
     """An Abinit input variable."""
 
-    def __init__(self, name, value, units="", valperline=3):
+    def __init__(self, name: str, value, units: str = "", valperline: int = 3) -> None:
         """
         Args:
             name: Name of the variable.
@@ -35,12 +35,11 @@ def __init__(self, name, value, units="", valperline=3):
         self.value = value
         self._units = units
 
-        # Maximum number of values per line.
-        self.valperline = valperline
+        self.valperline = valperline  # Maximum number of values per line.
         if name == "bdgw":
             self.valperline = 2
 
-        if isinstance(self.value, Iterable) and isinstance(self.value[-1], str) and self.value[-1] in _UNITS:
+        if isinstance(self.value, Sequence) and isinstance(self.value[-1], str) and self.value[-1] in _UNITS:
             self.value = list(self.value)
             self._units = self.value.pop(-1)
 

tests/core/test_structure.py

@@ -203,37 +203,34 @@ def test_as_dict(self):
         struct = IStructure(self.lattice, [{si: 0.5, mn: 0.5}, {si: 0.5}], coords)
         assert "lattice" in struct.as_dict()
         assert "sites" in struct.as_dict()
-        d = self.propertied_structure.as_dict()
-        assert d["sites"][0]["properties"]["magmom"] == 5
+        dct = self.propertied_structure.as_dict()
+        assert dct["sites"][0]["properties"]["magmom"] == 5
         coords = [[0, 0, 0], [0.75, 0.5, 0.75]]
         struct = IStructure(
             self.lattice,
-            [
-                {Species("O", -2, spin=3): 1.0},
-                {Species("Mg", 2, spin=2): 0.8},
-            ],
+            [{Species("O", -2, spin=3): 1.0}, {Species("Mg", 2, spin=2): 0.8}],
             coords,
             site_properties={"magmom": [5, -5]},
             properties={"general_property": "test"},
         )
-        d = struct.as_dict()
-        assert d["sites"][0]["properties"]["magmom"] == 5
-        assert d["sites"][0]["species"][0]["spin"] == 3
-        assert d["properties"]["general_property"] == "test"
+        dct = struct.as_dict()
+        assert dct["sites"][0]["properties"]["magmom"] == 5
+        assert dct["sites"][0]["species"][0]["spin"] == 3
+        assert dct["properties"]["general_property"] == "test"
 
-        d = struct.as_dict(0)
-        assert "volume" not in d["lattice"]
-        assert "xyz" not in d["sites"][0]
+        dct = struct.as_dict(0)
+        assert "volume" not in dct["lattice"]
+        assert "xyz" not in dct["sites"][0]
 
     def test_from_dict(self):
-        d = self.propertied_structure.as_dict()
-        struct = IStructure.from_dict(d)
+        dct = self.propertied_structure.as_dict()
+        struct = IStructure.from_dict(dct)
         assert struct[0].magmom == 5
-        d = self.propertied_structure.as_dict(0)
-        s2 = IStructure.from_dict(d)
+        dct = self.propertied_structure.as_dict(0)
+        s2 = IStructure.from_dict(dct)
         assert struct == s2
 
-        d = {
+        dct = {
             "lattice": {
                 "a": 3.8401979337,
                 "volume": 40.044794644251596,
@@ -282,7 +279,7 @@ def test_from_dict(self):
             ],
             "properties": {"test_property": "test"},
         }
-        struct = IStructure.from_dict(d)
+        struct = IStructure.from_dict(dct)
         assert struct[0].magmom == 5
         assert struct[0].specie.spin == 3
         assert struct.properties["test_property"] == "test"
@@ -818,7 +815,7 @@ def test_get_dist_matrix(self):
         ans = [[0.0, 2.3516318], [2.3516318, 0.0]]
         assert_allclose(self.struct.distance_matrix, ans)
 
-    def test_to_from_file_string(self):
+    def test_to_from_file_and_string(self):
         for fmt in ["cif", "json", "poscar", "cssr"]:
             struct = self.struct.to(fmt=fmt)
             assert struct is not None
@@ -855,8 +852,8 @@ def test_to_from_file_string(self):
         with pytest.raises(ValueError, match="Invalid format='badformat'"):
             self.struct.to(fmt="badformat")
 
-        self.struct.to(filename="POSCAR.testing.gz")
-        struct = Structure.from_file("POSCAR.testing.gz")
+        self.struct.to(filename=(gz_json_path := "POSCAR.testing.gz"))
+        struct = Structure.from_file(gz_json_path)
         assert struct == self.struct
 
         # test CIF file with unicode error
@@ -897,6 +894,7 @@ def setUp(self):
         coords = [[0, 0, 0], [0.75, 0.5, 0.75]]
         lattice = Lattice([[3.8401979337, 0, 0], [1.9200989668, 3.3257101909, 0], [0, -2.2171384943, 3.1355090603]])
         self.struct = Structure(lattice, ["Si", "Si"], coords)
+        self.struct.properties["foo"] = "bar"
         self.cu_structure = Structure(lattice, ["Cu", "Cu"], coords)
         self.disordered = Structure.from_spacegroup("Im-3m", Lattice.cubic(3), [Composition("Fe0.5Mn0.5")], [[0, 0, 0]])
         self.labeled_structure = Structure(lattice, ["Si", "Si"], coords, labels=["Si1", "Si2"])
@@ -1271,13 +1269,14 @@ def test_as_from_dict(self):
         assert isinstance(s1, Structure)
 
     def test_default_dict_attrs(self):
-        d = self.struct.as_dict()
-        assert d["charge"] == 0
+        dct = self.struct.as_dict()
+        assert dct["charge"] == 0
+        assert dct["properties"] == {"foo": "bar"}
 
     def test_to_from_abivars(self):
         """Test as_dict, from_dict with fmt == abivars."""
-        d = self.struct.as_dict(fmt="abivars")
-        s2 = Structure.from_dict(d, fmt="abivars")
+        dct = self.struct.as_dict(fmt="abivars")
+        s2 = Structure.from_dict(dct, fmt="abivars")
         assert s2 == self.struct
         assert isinstance(s2, Structure)
 
@@ -1292,8 +1291,8 @@ def test_to_from_file_string(self):
             assert isinstance(ss, Structure)
 
         # to/from file
-        self.struct.to(filename="POSCAR.testing")
-        assert os.path.isfile("POSCAR.testing")
+        self.struct.to(filename=(poscar_path := "POSCAR.testing"))
+        assert os.path.isfile(poscar_path)
 
         for ext in (".json", ".json.gz", ".json.bz2", ".json.xz", ".json.lzma"):
             self.struct.to(filename=f"json-struct{ext}")

tests/io/abinit/test_inputs.py

@@ -256,7 +256,6 @@ def test_api(self):
 
 class TestShiftMode(PymatgenTest):
     def test_shiftmode(self):
-        """Testing shiftmode."""
         gamma = ShiftMode.GammaCentered
         assert ShiftMode.from_object("G") == gamma
         assert ShiftMode.from_object(gamma) == gamma