Added the Crystal.groupby method which allows to group atoms by some …

…measure of site symmetry (#12)

CHANGELOG.rst

@@ -4,6 +4,11 @@ What's new
 
 .. currentmodule:: crystals
 
+Release 1.6.0
+-------------
+
+* Added the :meth:`Crystal.groupby` method to group unit cell atoms by site-symmetry (#12).
+
 Release 1.5.0
 -------------
 
@@ -15,7 +20,7 @@ Release 1.5.0
 Release 1.4.1
 -------------
 
-* Fixed an issue with the `tag` attribute of `Atom`s not being propagated properly (#9).
+* Fixed an issue with the `tag` attribute of `Atom` not being propagated properly (#9).
 
 Release 1.4.0
 -------------

crystals/__init__.py

@@ -5,7 +5,7 @@
 __author__ = "Laurent P. René de Cotret"
 __email__ = "laurent.renedecotret@mail.mcgill.ca"
 __license__ = "GPLv3"
-__version__ = "1.5.0"
+__version__ = "1.6.0"
 
 from .atom import Atom
 from .atom import Element

crystals/crystal.py

@@ -1,8 +1,9 @@
 # -*- coding: utf-8 -*-
+from collections import namedtuple
 from copy import deepcopy
 from enum import Enum, unique
 from functools import lru_cache
-from itertools import chain, combinations, islice, product
+from itertools import chain, combinations, islice, product, groupby, starmap
 from operator import attrgetter
 from os import PathLike
 from pathlib import Path
@@ -37,6 +38,14 @@
 CIF_ENTRIES = frozenset((Path(__file__).parent / "cifs").glob("*.cif"))
 LONGEST_CHEMICAL_SYMBOL = max(len(s) for s in chemical_symbols)
 
+ACCEPTABLE_SITE_SYMMETRIES = frozenset(
+    {
+        "wyckoffs",
+        "crystallographic_orbits",
+        "equivalent_atoms",
+    }
+)
+
 is_atom = lambda a: isinstance(a, Atom)
 is_structure = lambda s: isinstance(s, AtomicStructure)
 
@@ -489,25 +498,25 @@ def symmetry(
         info : dict
             Dictionary of space-group information. The following keys are available:
 
-            * ``'international_symbol'``: International Tables of Crystallography
-              space-group symbol (short);
+                * ``'international_symbol'``: International Tables of Crystallography
+                  space-group symbol (short);
 
-            * ``'international_full'``: International Tables of
-              Crystallography space-group full symbol;
+                * ``'international_full'``: International Tables of
+                  Crystallography space-group full symbol;
 
-            * ``'hall_symbol'`` : Hall symbol;
+                * ``'hall_symbol'`` : Hall symbol;
 
-            * ``'hm_symbol'`` : Hermann-Mauguin symbol;
+                * ``'hm_symbol'`` : Hermann-Mauguin symbol;
 
-            *``'centering'``: Centering-type ("P", "F", etc.);
+                * ``'centering'``: Centering-type ("P", "F", etc.);
 
-            * ``'pointgroup'`` : International Tables of
-              Crystallography point-group;
+                * ``'pointgroup'`` : International Tables of
+                  Crystallography point-group;
 
-            * ``'international_number'`` : International Tables of
-              Crystallography space-group number (between 1 and 230);
+                * ``'international_number'`` : International Tables of
+                  Crystallography space-group number (between 1 and 230);
 
-            * ``'hall_number'`` : Hall number (between 1 and 531).
+                * ``'hall_number'`` : Hall number (between 1 and 531).
 
         Raises
         ------
@@ -681,6 +690,87 @@ def centering(self) -> "CenteringType":
         """Centering type of this crystals."""
         return self.symmetry()["centering"]
 
+    # TODO: interpolate the docstring to specify the allowed values of `by`.
+    #       f-strings cannot be used as docstrings unfortunately.
+    def groupby(
+        self,
+        by: str,
+        symprec: float = 1e-2,
+        angle_tolerance: float = -1.0,
+    ) -> Dict[Any, Iterable[Atom]]:
+        """
+        Group unit cell atoms by some measure of site symmetry, for example Wyckoff letters
+        or crystallographic orbits.
+
+        .. versionadded:: 1.6.0
+
+        Parameters
+        ----------
+        by : str
+            Measure of site symmetry by which to group atoms. Can be one of "wyckoffs", "crystallographic_orbits",
+            or "equivalent_atoms". See the documentation of `spglib` for a description of these quantities.
+        symprec : float, optional
+            Symmetry-search distance tolerance in Cartesian coordinates [Angstroms].
+        angle_tolerance: float, optional
+            Symmetry-search tolerance in degrees. If the value is negative (default),
+            an internally optimized routine is used to judge symmetry.
+
+        Returns
+        -------
+        groups: dict
+            Dictionary where each key is a distinct site symmetry (e.g. a distinct orbit), and values are
+            iterable of :class:`Atom` which share the same site symmetry.
+
+        Raises
+        ------
+        RuntimeError
+            If symmetry-determination has not succeeded.
+        ValueError
+            If `by` is not one of the supported values.
+
+        Notes
+        -----
+        Note that crystals generated from the Protein Data Bank are often incomplete;
+        in such cases the site symmetry information will be incorrect.
+        """
+        if by not in ACCEPTABLE_SITE_SYMMETRIES:
+            raise ValueError(
+                f"Cannot group atoms by site-symmetry '{by}'. Expected one of: {ACCEPTABLE_SITE_SYMMETRIES}."
+            )
+
+        # Iterating over the atoms of a `Crystal` has an undefined order. As such,
+        # we need to first 'freeze' the iteration order, and create a spglib-compatible
+        # cell *with this order*. That's why we can't re-use `Crystal._spglib_cell`.
+        sortkey = lambda atm: np.linalg.norm(atm.coords_fractional)
+        atoms = sorted(self.unitcell, key=sortkey)
+
+        unitcell_array = np.stack([np.asarray(atm) for atm in atoms])
+        spglib_cell = (
+            np.array(self.lattice_vectors),
+            unitcell_array[:, 1:],
+            unitcell_array[:, 0],
+        )
+
+        symmetry_dataset = get_symmetry_dataset(
+            cell=spglib_cell, symprec=symprec, angle_tolerance=angle_tolerance
+        )
+
+        mksite = namedtuple("Site", ("atom", *ACCEPTABLE_SITE_SYMMETRIES))
+
+        sites = starmap(
+            mksite,
+            zip(
+                atoms,
+                *(symmetry_dataset[measure] for measure in ACCEPTABLE_SITE_SYMMETRIES),
+            ),
+        )
+
+        by = attrgetter(by)
+        return {
+            k: sorted((site.atom for site in v), key=sortkey)
+            for k, v in groupby(iterable=sorted(sites, key=by), key=by)
+        }
+
     def indexed_by(self, lattice: Union[Lattice, np.ndarray]) -> "Crystal":
         """
         Return a crystal structure, indexed by another lattice/crystal structure.

crystals/tests/test_crystal.py

@@ -10,7 +10,11 @@
 import numpy as np
 from crystals import Atom, AtomicStructure, CenteringType, Crystal
 from crystals.affine import translation_matrix
-from crystals.crystal import symmetry_expansion, symmetry_reduction
+from crystals.crystal import (
+    symmetry_expansion,
+    symmetry_reduction,
+    ACCEPTABLE_SITE_SYMMETRIES,
+)
 import pytest
 
 np.random.seed(23)
@@ -225,6 +229,25 @@ def test_supercell_preserve_attributes(tag, symbol, occupancy):
         assert supercell_atom.symbol == atom.symbol
 
 
+# The two tested site symmetry measures aren't *always* equivalent, but often are
+@pytest.mark.parametrize(
+    "site_symmetry_measure", ["crystallographic_orbits", "equivalent_atoms"]
+)
+@pytest.mark.parametrize(
+    "crystal, num_distinct_sites", (["C", 2], ["vo2-m1", 3], ["vo2-rutile", 2])
+)
+def test_groupby(site_symmetry_measure, crystal, num_distinct_sites):
+    cryst = Crystal.from_database(crystal)
+    gps = cryst.groupby(site_symmetry_measure)
+    assert len(gps) == num_distinct_sites
+
+
+def test_groupby_unknown_symmetry():
+    cryst = Crystal.from_database("diamond")
+    with pytest.raises(ValueError):
+        cryst.groupby("this should not work")
+
+
 def test_indexed_by_trivial_reindexing():
     """Test re-indexing a crystal by itself"""
     c1 = Crystal.from_database("Pu-gamma")

docs/guides/crystals.rst

@@ -402,6 +402,24 @@ and the translation is the right-most column. Example with iteration:
 
 Symmetry operations in reciprocal space are also made available via :meth:`Crystal.reciprocal_symmetry_operations`.
 
+
+Grouping atoms by site-symmetry
+-------------------------------
+
+For various reasons, it might be useful to know which of the atoms in a :class:`Crystal` are related by some site symmetry, for example Wyckoff letters
+or crystallographic orbits. This can be done with the :meth:`Crystal.groupby` method, which groups atoms according to their site symmetry. For example:
+
+    >>> graphite = Crystal.from_database('C')
+    >>> groups = graphite.groupby(by="crystallographic_orbits")
+    >>> for orbit, atoms in groups.items():
+    ...     print(f"Orbit {orbit}: {atoms}")
+    ... 
+    Orbit 0: [< Atom C  @ (0.00, 0.00, 0.25) >, < Atom C  @ (0.00, 0.00, 0.75) >]
+    Orbit 2: [< Atom C  @ (0.33, 0.67, 0.25) >, < Atom C  @ (0.67, 0.33, 0.75) >]
+
+Supported site-symmetry measures are currently `"crystallographic_orbits"`, `"wyckoffs"`, and `"equivalent_atoms"`. See 
+the `spglib <http://atztogo.github.io/spglib/>`_ documentation for a description of these measures of symmetry.
+
 Cell refinements and manipulations
 ----------------------------------