Skip to content

Commit

Permalink
Return dict output.
Browse files Browse the repository at this point in the history
  • Loading branch information
shyuep committed Jan 28, 2025
1 parent 3b6a057 commit bf6c108
Show file tree
Hide file tree
Showing 2 changed files with 39 additions and 9 deletions.
35 changes: 28 additions & 7 deletions src/pymatgen/core/structure.py
Original file line number Diff line number Diff line change
Expand Up @@ -3352,8 +3352,8 @@ def get_symmetry_dataset(self, backend: Literal["moyopy"], **kwargs) -> moyopy.M
def get_symmetry_dataset(self, backend: Literal["spglib"], **kwargs) -> spglib.SpglibDataset: ...

def get_symmetry_dataset(
self, backend: Literal["moyopy", "spglib"] = "spglib", **kwargs
) -> moyopy.MoyoDataset | spglib.SpglibDataset:
self, backend: Literal["moyopy", "spglib"] = "spglib", return_raw_dataset=False, **kwargs
) -> dict | moyopy.MoyoDataset | spglib.SpglibDataset:
"""Get a symmetry dataset from the structure using either moyopy or spglib backend.
If using the spglib backend (default), please cite:
Expand All @@ -3365,6 +3365,9 @@ def get_symmetry_dataset(
Args:
backend ("moyopy" | "spglib"): Which symmetry analysis backend to use.
Defaults to "spglib".
return_raw_dataset (bool): Whether to return the raw Dataset object from the backend. The default is
False, which returns a dict with a common subset of the data present in both datasets. If you use the
raw Dataset object, we do not guarantee that the format of the output is not going to change.
**kwargs: Additional arguments passed to the respective backend's constructor.
For spglib, these are passed to SpacegroupAnalyzer (e.g. symprec, angle_tolerance).
For moyopy, these are passed to MoyoDataset constructor.
Expand All @@ -3376,6 +3379,9 @@ def get_symmetry_dataset(
ImportError: If the requested backend is not installed.
ValueError: If an invalid backend is specified.
"""
if backend not in ("moyopy", "spglib"):
raise ValueError(f"Invalid {backend=}, must be one of moyopy or spglib.")

if backend == "moyopy":
try:
import moyopy
Expand All @@ -3385,16 +3391,31 @@ def get_symmetry_dataset(

# Convert structure to MoyoDataset format
moyo_cell = moyopy.interface.MoyoAdapter.from_structure(self)
return moyopy.MoyoDataset(cell=moyo_cell, **kwargs)
dataset = moyopy.MoyoDataset(cell=moyo_cell, **kwargs)

if backend == "spglib":
else:
from pymatgen.symmetry.analyzer import SpacegroupAnalyzer

sga = SpacegroupAnalyzer(self, **kwargs)
return sga.get_symmetry_dataset()
dataset = sga.get_symmetry_dataset()

if return_raw_dataset:
return dataset

dictdata = {k: getattr(dataset, k) for k in ("hall_number", "number", "site_symmetry_symbols", "wyckoffs")}

if backend == "spglib":
dictdata["international"] = dataset.international
dictdata["orbits"] = dataset.crystallographic_orbits
dictdata["std_origin_shift"] = dataset.origin_shift
else:
from pymatgen.symmetry.groups import SpaceGroup

dictdata["international"] = SpaceGroup.from_int_number(dataset.number).symbol
dictdata["orbits"] = dataset.orbits
dictdata["std_origin_shift"] = dataset.std_origin_shift

valid_backends = ("moyopy", "spglib")
raise ValueError(f"Invalid {backend=}, must be one of {valid_backends}")
return dictdata


class IMolecule(SiteCollection, MSONable):
Expand Down
13 changes: 11 additions & 2 deletions tests/core/test_structure.py
Original file line number Diff line number Diff line change
Expand Up @@ -969,15 +969,24 @@ def test_sites_setter(self):
def test_get_symmetry_dataset(self):
"""Test getting symmetry dataset from structure using different backends."""
# Test spglib backend
dataset = self.struct.get_symmetry_dataset(backend="spglib")
for backend in ("spglib", "moyopy"):
dataset = self.struct.get_symmetry_dataset(backend=backend)
assert isinstance(dataset, dict)
assert dataset["number"] == 227
assert dataset["international"] == "Fd-3m"
assert len(dataset["orbits"]) == 2
pytest.approx(dataset["std_origin_shift"], [0, 0, 0])

dataset = self.struct.get_symmetry_dataset(backend="spglib", return_raw_dataset=True)

assert dataset.number == 227 # Fd-3m space group
assert dataset.international == "Fd-3m"
assert len(dataset.rotations) > 0
assert len(dataset.translations) > 0

# Test moyopy backend if available
moyopy = pytest.importorskip("moyopy")
dataset = self.struct.get_symmetry_dataset(backend="moyopy")
dataset = self.struct.get_symmetry_dataset(backend="moyopy", return_raw_dataset=True)
assert isinstance(dataset, moyopy.MoyoDataset)
assert dataset.prim_std_cell.numbers == [14, 14] # Si atomic number is 14

Expand Down

0 comments on commit bf6c108

Please sign in to comment.