3d plots of unit cells

CHANGELOG.rst

@@ -12,9 +12,13 @@ Unreleased
 
 Added
 -----
+- The unit cells and atomic positions of ``Phase`` objects can now be visualized
+  using ``Phase.plot_unit_cell``.
 
 Changed
 -------
+- The IPF color keys can take 'x', 'y', and 'z' text strings as direction inputs
+  instead of requiring Vector3d objects.
 
 Removed
 -------

orix/crystal_map/_phase.py

@@ -425,6 +425,61 @@ def expand_asymmetric_unit(self) -> Phase:
 
         return expanded_phase
 
+    def plot_unit_cell(
+        self,
+        return_figure: bool = False,
+        show_xyz: bool = False,
+        show_atoms: bool = True,
+        show_uvw_labels=False,
+        figsize: None | np.array = None,
+    ):
+        """Create a plot of the unit cell for the phase"""
+        # define the vertices and faces of the paralellipid unit cell.
+        import matplotlib.pyplot as plt
+        from mpl_toolkits.mplot3d.art3d import Poly3DCollection
+
+        unit_v = [[x, y, z] for x in [0, 1] for y in [0, 1] for z in [0, 1]]
+        v = np.array(unit_v).dot(self._diffpy_lattice.T)
+        f = [
+            [0, 1, 3, 2],
+            [0, 1, 5, 4],
+            [5, 4, 6, 7],
+            [6, 7, 3, 2],
+            [1, 3, 7, 5],
+            [0, 2, 6, 4],
+        ]
+
+        # set up the plot
+        if figsize is None:
+            fig = plt.figure()
+        else:
+            fig = plt.figure(figsize=figsize)
+        ax = fig.add_subplot(projection="3d")
+        ax.set_aspect("equal")
+        ax.set_proj_type = "ortho"
+        fc = self.color_rgb + (0.1,)
+        ax.add_collection3d(Poly3DCollection(v[f], facecolors=fc, edgecolors=[0, 0, 0]))
+        if show_xyz:
+            ax.plot([0, 1], [0, 0], [0, 0], "grey")
+            ax.plot([0, 0], [0, 1], [0, 0], "grey")
+            ax.plot([0, 0], [0, 0], [0, 1], "grey")
+            ax.text(1, 0, 0, "X")
+            ax.text(0, 1, 0, "Y")
+            ax.text(0, 0, 1, "Z")
+        if show_uvw_labels:
+            ax.text(*v[4] + 0.01, "[100]")
+            ax.text(*v[2] + 0.01, "[010]")
+            ax.text(*v[1] + 0.01, "[001]")
+        if show_atoms and len(self.structure) > 0:
+            for i in range(len(self.structure)):
+                xyz = self.structure[i].xyz.dot(self._diffpy_lattice.T)
+                elem = self.structure[i].element.title()  # titlecase
+                rgb = _jmol_colors.get(elem, [0, 0, 0])
+                ax.scatter(*xyz, facecolor=np.array(rgb) / 255, s=300, edgecolor="k")
+        if return_figure:
+            return fig
+        return
+
 
 def new_structure_matrix_from_alignment(
     old_matrix: np.ndarray,
@@ -493,3 +548,117 @@ def default_lattice(system: VALID_SYSTEMS) -> Lattice:
     else:
         raise ValueError(f"Unknown crystal system {system!r}")
     return lat
+
+
+# The JMOL color dictionary, taken from http://jmol.sourceforge.net/jscolors
+_jmol_colors = {
+    "H": [255, 255, 255],
+    "He": [217, 255, 255],
+    "Li": [204, 128, 255],
+    "Be": [194, 255, 0],
+    "B": [255, 181, 181],
+    "C": [144, 144, 144],
+    "N": [48, 80, 248],
+    "O": [255, 13, 13],
+    "F": [144, 224, 80],
+    "Ne": [179, 227, 245],
+    "Na": [171, 92, 242],
+    "Mg": [138, 255, 0],
+    "Al": [191, 166, 166],
+    "Si": [240, 200, 160],
+    "P": [255, 128, 0],
+    "S": [255, 255, 48],
+    "Cl": [31, 240, 31],
+    "Ar": [128, 209, 227],
+    "K": [143, 64, 212],
+    "Ca": [61, 255, 0],
+    "Sc": [230, 230, 230],
+    "Ti": [191, 194, 199],
+    "V": [166, 166, 171],
+    "Cr": [138, 153, 199],
+    "Mn": [156, 122, 199],
+    "Fe": [224, 102, 51],
+    "Co": [240, 144, 160],
+    "Ni": [80, 208, 80],
+    "Cu": [200, 128, 51],
+    "Zn": [125, 128, 176],
+    "Ga": [194, 143, 143],
+    "Ge": [102, 143, 143],
+    "As": [189, 128, 227],
+    "Se": [255, 161, 0],
+    "Br": [166, 41, 41],
+    "Kr": [92, 184, 209],
+    "Rb": [112, 46, 176],
+    "Sr": [0, 255, 0],
+    "Y": [148, 255, 255],
+    "Zr": [148, 224, 224],
+    "Nb": [115, 194, 201],
+    "Mo": [84, 181, 181],
+    "Tc": [59, 158, 158],
+    "Ru": [36, 143, 143],
+    "Rh": [10, 125, 140],
+    "Pd": [0, 105, 133],
+    "Ag": [192, 192, 192],
+    "Cd": [255, 217, 143],
+    "In": [166, 117, 115],
+    "Sn": [102, 128, 128],
+    "Sb": [158, 99, 181],
+    "Te": [212, 122, 0],
+    "I": [148, 0, 148],
+    "Xe": [66, 158, 176],
+    "Cs": [87, 23, 143],
+    "Ba": [0, 201, 0],
+    "La": [112, 212, 255],
+    "Ce": [255, 255, 199],
+    "Pr": [217, 255, 199],
+    "Nd": [199, 255, 199],
+    "Pm": [163, 255, 199],
+    "Sm": [143, 255, 199],
+    "Eu": [97, 255, 199],
+    "Gd": [69, 255, 199],
+    "Tb": [48, 255, 199],
+    "Dy": [31, 255, 199],
+    "Ho": [0, 255, 156],
+    "Er": [0, 230, 117],
+    "Tm": [0, 212, 82],
+    "Yb": [0, 191, 56],
+    "Lu": [0, 171, 36],
+    "Hf": [77, 194, 255],
+    "Ta": [77, 166, 255],
+    "W": [33, 148, 214],
+    "Re": [38, 125, 171],
+    "Os": [38, 102, 150],
+    "Ir": [23, 84, 135],
+    "Pt": [208, 208, 224],
+    "Au": [255, 209, 35],
+    "Hg": [184, 184, 208],
+    "Tl": [166, 84, 77],
+    "Pb": [87, 89, 97],
+    "Bi": [158, 79, 181],
+    "Po": [171, 92, 0],
+    "At": [117, 79, 69],
+    "Rn": [66, 130, 150],
+    "Fr": [66, 0, 102],
+    "Ra": [0, 125, 0],
+    "Ac": [112, 171, 250],
+    "Th": [0, 186, 255],
+    "Pa": [0, 161, 255],
+    "U": [0, 143, 255],
+    "Np": [0, 128, 255],
+    "Pu": [0, 107, 255],
+    "Am": [84, 92, 242],
+    "Cm": [120, 92, 227],
+    "Bk": [138, 79, 227],
+    "Cf": [161, 54, 212],
+    "Es": [179, 31, 212],
+    "Fm": [179, 31, 186],
+    "Md": [179, 13, 166],
+    "No": [189, 13, 135],
+    "Lr": [199, 0, 102],
+    "Rf": [204, 0, 89],
+    "Db": [209, 0, 79],
+    "Sg": [217, 0, 69],
+    "Bh": [224, 0, 56],
+    "Hs": [230, 0, 46],
+    "Mt": [235, 0, 38],
+}

orix/plot/direction_color_keys/direction_color_key_tsl.py

@@ -136,9 +136,8 @@ def _create_rgba_grid(
         rgba_grid = rgba_grid[::-1]
 
         if return_extent:
-            return rgba_grid, ((x_min, x_max), (y_min, y_max))
-        else:
-            return rgba_grid
+            rgba_grid = [rgba_grid, ((x_min, x_max), (y_min, y_max))]
+        return rgba_grid
 
     def plot(self, return_figure: bool = False) -> Figure | None:
         """Plot the inverse pole figure color key.

orix/plot/orientation_color_keys/ipf_color_key_tsl.py

@@ -42,13 +42,24 @@ def __init__(self, symmetry: Symmetry, direction: Vector3d | None = None) -> Non
 
         Parameters
         ----------
-        symmetry : orix.quaternion.Symmetry
+        symmetry
             (Laue) symmetry of the crystal. If a non-Laue symmetry
             is given, the Laue symmetry of that symmetry will be used.
-        direction : orix.vector.Vector3d, optional
+        direction
             Sample direction. If not given, sample Z direction (out of
             plane) is used.
         """
+        if type(direction) is str:
+            direction = {
+                "x": Vector3d.xvector(),
+                "y": Vector3d.yvector(),
+                "z": Vector3d.zvector(),
+            }.get(direction.lower(), "fail")
+            if type(direction) is str:
+                raise IOError(
+                    "'direction' must be 'x', 'y','z', or a Vector3D instance"
+                )
+
         super().__init__(symmetry.laue, direction=direction)
 
     @property

orix/tests/test_crystal_map/test_phase.py

@@ -17,10 +17,12 @@
 # along with orix. If not, see <http://www.gnu.org/licenses/>.
 #
 
+import diffpy.structure as dst
 from diffpy.structure import Atom, Lattice, Structure, loadStructure
 import numpy as np
 import pytest
 
+import orix.crystal_map as ocm
 from orix.crystal_map import Phase
 from orix.crystal_map._phase import default_lattice, new_structure_matrix_from_alignment
 from orix.quaternion.symmetry import O, Symmetry
@@ -563,7 +565,7 @@ def test_expand_asymmetric_unit(
         # Check atom positions in ORIGINAL lattice alignment
         # Doing the check in orix's alignment makes independently computing expected sites difficult
         s = exp.structure.copy()
-        s.placeInLattice(Lattice(base=phase._diffpy_lattice))
+        s.placeInLattice(dst.Lattice(base=phase._diffpy_lattice))
         # Use set to avoid having to ensure the order is the same
         assert set(tuple(xyz.round(8).tolist()) for xyz in s.xyz) == set(
             expected_atom_positions
@@ -574,7 +576,7 @@ def test_expand_asymmetric_unit(
         assert np.array_equal(base, exp2.structure.lattice.base)
         assert len(exp2.structure) == len(expected_atom_positions)
         s = exp2.structure.copy()
-        s.placeInLattice(Lattice(base=phase._diffpy_lattice))
+        s.placeInLattice(dst.Lattice(base=phase._diffpy_lattice))
         assert set(tuple(xyz.round(8).tolist()) for xyz in s.xyz) == set(
             expected_atom_positions
         )
@@ -954,7 +956,7 @@ def test_expand_asymmetric_unit_from_cif(
         filepath = tmp_path / "tmp.cif"
         with open(filepath, "w") as file:
             file.write(cif_file_content)
-        phase = Phase.from_cif(filepath)
+        phase = ocm.Phase.from_cif(filepath)
         # Asymmetric unit is automatically expanded when read from cif
         assert len(phase.structure) == expected_atom_count
         # Expand just in case
@@ -980,3 +982,39 @@ def test_default_lattice(self):
     def test_default_lattice_raises(self):
         with pytest.raises(ValueError, match="Unknown crystal system 'rhombohedral'"):
             default_lattice("rhombohedral")
+
+    def test_plot_unit_cell(self):
+        import matplotlib.pyplot as plt
+
+        Al2O3_atoms = [
+            dst.Atom("Al", [1 / 3, 2 / 3, 0.815]),
+            dst.Atom("O", [0.361, 1 / 3, 0.583]),
+        ]
+        Al2O3_lattice = dst.Lattice(0.481, 0.481, 1.391, 90, 90, 120)
+        Al2O3_structure = dst.Structure(atoms=Al2O3_atoms, lattice=Al2O3_lattice)
+        Al2O3_phase = ocm.Phase(
+            name="Alumina",
+            space_group=167,
+            structure=Al2O3_structure,
+            color="red",
+        ).expand_asymmetric_unit()
+        Fe_atoms = [
+            dst.Atom("Fe", [0, 0, 0]),
+            dst.Atom("Fe", [1, 0, 0]),
+            dst.Atom("Fe", [0, 1, 0]),
+            dst.Atom("Fe", [0, 0, 1]),
+            dst.Atom("Fe", [0, 1, 1]),
+            dst.Atom("Fe", [1, 0, 1]),
+            dst.Atom("Fe", [1, 1, 0]),
+            dst.Atom("Fe", [1, 1, 1]),
+            dst.Atom("Fe", [1 / 2, 1 / 2, 1 / 2]),
+        ]
+        Fe_lattice = dst.Lattice(1, 1, 1, 90, 90, 90)
+        Fe_structure = dst.Structure(atoms=Fe_atoms, lattice=Fe_lattice)
+        Fe_phase = ocm.Phase(point_group="m3m", structure=Fe_structure)
+        for p in [Al2O3_phase, Fe_phase]:
+            fig1 = p.plot_unit_cell(figsize=[5, 4], return_figure=True)
+            fig2 = p.plot_unit_cell(figsize=np.array([5.1, 4.3]), return_figure=True)
+            p.plot_unit_cell(show_xyz=True, show_atoms=True, show_uvw_labels=True)
+            p.plot_unit_cell(show_xyz=False, show_atoms=False, show_uvw_labels=False)
+            plt.close("all")

orix/tests/test_plot/test_orientation_color_keys.py

@@ -19,9 +19,13 @@
 
 import matplotlib.pyplot as plt
 import numpy as np
+import pytest
 
+import orix.plot as opl
 from orix.plot import EulerColorKey, IPFColorKeyTSL
+import orix.quaternion as oqu
 from orix.quaternion import Orientation, symmetry
+import orix.vector as ove
 from orix.vector import Vector3d
 
 
@@ -34,7 +38,7 @@ def test_orientation2color(self):
             np.radians(((0, 0, 0), (0, 45, 0), (-45, 54.7356, 45))),
             symmetry=pg_oh,
         )
-        ckey_oh = IPFColorKeyTSL(pg_o)
+        ckey_oh = opl.IPFColorKeyTSL(pg_o, direction=ove.Vector3d.zvector())
         assert np.allclose(ckey_oh.symmetry.data, pg_oh.data)
         assert np.allclose(ckey_oh.direction.data, (0, 0, 1))
         assert repr(ckey_oh) == "IPFColorKeyTSL, symmetry: m-3m, direction: [0 0 1]"
@@ -43,6 +47,13 @@ def test_orientation2color(self):
         assert ax_o._symmetry.name == pg_oh.name
         rgb_oh = ckey_oh.orientation2color(ori)
         assert np.allclose(rgb_oh, ((1, 0, 0), (0, 1, 0), (0, 0, 1)), atol=0.1)
+        # using string input gives same result
+        ckey_z = opl.IPFColorKeyTSL(pg_o, direction="z")
+        rgb_z = ckey_z.orientation2color(ori)
+        assert np.allclose(rgb_z, ((1, 0, 0), (0, 1, 0), (0, 0, 1)), atol=0.1)
+        # Using nonsense direction raises informative error
+        with pytest.raises(IOError, match="'direction' must be"):
+            _ = opl.IPFColorKeyTSL(pg_o, direction="bad directions")
 
         # Color [001] and "diagonals" of 2/m IPF red, green and blue
         pg_c2 = symmetry.C2  # 2
@@ -51,7 +62,7 @@ def test_orientation2color(self):
             np.radians(((-90, -90, 0), (0, 90, -55), (0, 90, 55))),
             symmetry=pg_c2h,
         )
-        ckey_c2h = IPFColorKeyTSL(pg_c2, Vector3d.xvector())
+        ckey_c2h = opl.IPFColorKeyTSL(pg_c2, "X")
         assert np.allclose(ckey_c2h.symmetry.data, pg_c2h.data)
         assert np.allclose(ckey_c2h.direction.data, (1, 0, 0))
         assert repr(ckey_c2h) == "IPFColorKeyTSL, symmetry: 2/m, direction: [1 0 0]"