Skip to content

Commit bf6c108

Browse files
committed
Return dict output.
1 parent 3b6a057 commit bf6c108

File tree

2 files changed

+39
-9
lines changed

2 files changed

+39
-9
lines changed

src/pymatgen/core/structure.py

Lines changed: 28 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -3352,8 +3352,8 @@ def get_symmetry_dataset(self, backend: Literal["moyopy"], **kwargs) -> moyopy.M
33523352
def get_symmetry_dataset(self, backend: Literal["spglib"], **kwargs) -> spglib.SpglibDataset: ...
33533353

33543354
def get_symmetry_dataset(
3355-
self, backend: Literal["moyopy", "spglib"] = "spglib", **kwargs
3356-
) -> moyopy.MoyoDataset | spglib.SpglibDataset:
3355+
self, backend: Literal["moyopy", "spglib"] = "spglib", return_raw_dataset=False, **kwargs
3356+
) -> dict | moyopy.MoyoDataset | spglib.SpglibDataset:
33573357
"""Get a symmetry dataset from the structure using either moyopy or spglib backend.
33583358
33593359
If using the spglib backend (default), please cite:
@@ -3365,6 +3365,9 @@ def get_symmetry_dataset(
33653365
Args:
33663366
backend ("moyopy" | "spglib"): Which symmetry analysis backend to use.
33673367
Defaults to "spglib".
3368+
return_raw_dataset (bool): Whether to return the raw Dataset object from the backend. The default is
3369+
False, which returns a dict with a common subset of the data present in both datasets. If you use the
3370+
raw Dataset object, we do not guarantee that the format of the output is not going to change.
33683371
**kwargs: Additional arguments passed to the respective backend's constructor.
33693372
For spglib, these are passed to SpacegroupAnalyzer (e.g. symprec, angle_tolerance).
33703373
For moyopy, these are passed to MoyoDataset constructor.
@@ -3376,6 +3379,9 @@ def get_symmetry_dataset(
33763379
ImportError: If the requested backend is not installed.
33773380
ValueError: If an invalid backend is specified.
33783381
"""
3382+
if backend not in ("moyopy", "spglib"):
3383+
raise ValueError(f"Invalid {backend=}, must be one of moyopy or spglib.")
3384+
33793385
if backend == "moyopy":
33803386
try:
33813387
import moyopy
@@ -3385,16 +3391,31 @@ def get_symmetry_dataset(
33853391

33863392
# Convert structure to MoyoDataset format
33873393
moyo_cell = moyopy.interface.MoyoAdapter.from_structure(self)
3388-
return moyopy.MoyoDataset(cell=moyo_cell, **kwargs)
3394+
dataset = moyopy.MoyoDataset(cell=moyo_cell, **kwargs)
33893395

3390-
if backend == "spglib":
3396+
else:
33913397
from pymatgen.symmetry.analyzer import SpacegroupAnalyzer
33923398

33933399
sga = SpacegroupAnalyzer(self, **kwargs)
3394-
return sga.get_symmetry_dataset()
3400+
dataset = sga.get_symmetry_dataset()
3401+
3402+
if return_raw_dataset:
3403+
return dataset
3404+
3405+
dictdata = {k: getattr(dataset, k) for k in ("hall_number", "number", "site_symmetry_symbols", "wyckoffs")}
3406+
3407+
if backend == "spglib":
3408+
dictdata["international"] = dataset.international
3409+
dictdata["orbits"] = dataset.crystallographic_orbits
3410+
dictdata["std_origin_shift"] = dataset.origin_shift
3411+
else:
3412+
from pymatgen.symmetry.groups import SpaceGroup
3413+
3414+
dictdata["international"] = SpaceGroup.from_int_number(dataset.number).symbol
3415+
dictdata["orbits"] = dataset.orbits
3416+
dictdata["std_origin_shift"] = dataset.std_origin_shift
33953417

3396-
valid_backends = ("moyopy", "spglib")
3397-
raise ValueError(f"Invalid {backend=}, must be one of {valid_backends}")
3418+
return dictdata
33983419

33993420

34003421
class IMolecule(SiteCollection, MSONable):

tests/core/test_structure.py

Lines changed: 11 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -969,15 +969,24 @@ def test_sites_setter(self):
969969
def test_get_symmetry_dataset(self):
970970
"""Test getting symmetry dataset from structure using different backends."""
971971
# Test spglib backend
972-
dataset = self.struct.get_symmetry_dataset(backend="spglib")
972+
for backend in ("spglib", "moyopy"):
973+
dataset = self.struct.get_symmetry_dataset(backend=backend)
974+
assert isinstance(dataset, dict)
975+
assert dataset["number"] == 227
976+
assert dataset["international"] == "Fd-3m"
977+
assert len(dataset["orbits"]) == 2
978+
pytest.approx(dataset["std_origin_shift"], [0, 0, 0])
979+
980+
dataset = self.struct.get_symmetry_dataset(backend="spglib", return_raw_dataset=True)
981+
973982
assert dataset.number == 227 # Fd-3m space group
974983
assert dataset.international == "Fd-3m"
975984
assert len(dataset.rotations) > 0
976985
assert len(dataset.translations) > 0
977986

978987
# Test moyopy backend if available
979988
moyopy = pytest.importorskip("moyopy")
980-
dataset = self.struct.get_symmetry_dataset(backend="moyopy")
989+
dataset = self.struct.get_symmetry_dataset(backend="moyopy", return_raw_dataset=True)
981990
assert isinstance(dataset, moyopy.MoyoDataset)
982991
assert dataset.prim_std_cell.numbers == [14, 14] # Si atomic number is 14
983992

0 commit comments

Comments
 (0)