ipython tests

Author: jmmshn

tests/asy_test/core/__init__.py

@@ -0,0 +1,35 @@
+import plotly.express as px
+from monty.json import MSONable
+
+from crystal_toolkit.core.jupyter import patch_msonable
+from crystal_toolkit.core.scene import Scene
+
+
+def test_patch_msonable():
+    patch_msonable()
+
+    class GetSceneClass(MSONable):
+        def get_scene(self):
+            return Scene(name="test_scene")
+
+    class GetPlotClass(MSONable):
+        def get_plot(self):
+            """Dummy plotly object"""
+            return px.scatter(x=[1, 2, 3], y=[1, 2, 3])
+
+    class AsDictClass(MSONable):
+        def __init__(self, a: int) -> None:
+            self.a = a
+
+    # The output of _ipython_display_ is None
+    # However, the logic for the creating the different output
+    # dictionaries should be executed so the following tests
+    # are still valuable.
+    as_dict_class = AsDictClass(1)
+    assert as_dict_class._ipython_display_() is None
+
+    get_scene_class = GetSceneClass()
+    assert get_scene_class._ipython_display_() is None
+
+    get_plot_class = GetPlotClass()
+    assert get_plot_class._ipython_display_() is None

tests/asy_test/core/test_jupyter.py

@@ -0,0 +1,175 @@
+from __future__ import annotations
+
+from pymatgen.core import Lattice, Structure
+
+from crystal_toolkit.core.legend import Legend
+
+
+class TestLegend:
+    def setup_method(self, method):
+        self.struct = Structure(
+            Lattice.cubic(5),
+            ["H", "O", "In"],
+            [[0, 0, 0], [0.5, 0.5, 0.5], [0.5, 0, 0]],
+            site_properties={
+                "example_site_prop": [5, 0, -3],
+                "example_categorical_site_prop": ["4a", "4a", "8b"],
+            },
+        )
+
+        self.site0 = self.struct[0]
+        self.sp0 = next(iter(self.site0.species))
+
+        self.site1 = self.struct[1]
+        self.sp1 = next(iter(self.site1.species))
+
+        self.site2 = self.struct[2]
+        self.sp2 = next(iter(self.site2.species))
+
+        self.struct_disordered = Structure(
+            Lattice.cubic(5),
+            ["H", "O", {"In": 0.5, "Al": 0.5}],
+            [[0, 0, 0], [0.5, 0.5, 0.5], [0.5, 0, 0]],
+            site_properties={"example_site_prop": [5, 0, -3]},
+        )
+
+        self.site_d = self.struct_disordered[2]
+        self.site_d_sp0 = next(iter(self.site_d.species))
+        self.site_d_sp1 = list(self.site_d.species)[1]
+
+        self.struct_manual = Structure(
+            Lattice.cubic(5),
+            ["H", "O2-", "In"],
+            [[0, 0, 0], [0.5, 0.5, 0.5], [0.5, 0, 0]],
+            site_properties={"display_color": [[255, 0, 0], "blue", "#00ff00"]},
+        )
+
+    def test_get_color(self):
+        # test default
+
+        legend = Legend(self.struct, color_scheme="VESTA")
+
+        color = legend.get_color(self.sp0)
+        assert color == "#ffcccc"
+
+        # element-based color schemes shouldn't change if you supply a site
+        color = legend.get_color(self.sp0, site=self.site0)
+        assert color == "#ffcccc"
+
+        color = legend.get_color(self.sp2)
+        assert color == "#a67573"
+
+        assert legend.get_legend()["colors"] == {
+            "#a67573": "In",
+            "#fe0300": "O",
+            "#ffcccc": "H",
+        }
+
+        # test alternate
+
+        legend = Legend(self.struct, color_scheme="Jmol")
+
+        color = legend.get_color(self.sp0)
+        assert color == "#ffffff"
+
+        assert legend.get_legend()["colors"] == {
+            "#a67573": "In",
+            "#ff0d0d": "O",
+            "#ffffff": "H",
+        }
+
+        # test coloring by site properties
+
+        legend = Legend(self.struct, color_scheme="example_site_prop")
+
+        color = legend.get_color(self.sp0, site=self.site0)
+        assert color == "#b30326"
+
+        color = legend.get_color(self.sp1, site=self.site1)
+        assert color == "#000000"
+
+        color = legend.get_color(self.sp2, site=self.site2)
+        assert color == "#7b9ef8"
+
+        assert legend.get_legend()["colors"] == {
+            "#7b9ef8": "-3.00",
+            "#b30326": "5.00",
+            "#000000": "0.00",
+        }
+
+        # test accessible
+
+        legend = Legend(self.struct, color_scheme="accessible")
+
+        color = legend.get_color(self.sp0, site=self.site0)
+        assert color == "#ffffff"
+
+        color = legend.get_color(self.sp1, site=self.site1)
+        assert color == "#d55e00"
+
+        color = legend.get_color(self.sp2, site=self.site2)
+        assert color == "#cc79a7"
+
+        assert legend.get_legend()["colors"] == {
+            "#cc79a7": "In",
+            "#d55e00": "O",
+            "#ffffff": "H",
+        }
+
+        # test disordered
+
+        legend = Legend(self.struct_disordered)
+
+        color = legend.get_color(self.site_d_sp0, site=self.site_d)
+        assert color == "#a67573"
+
+        color = legend.get_color(self.site_d_sp1, site=self.site_d)
+        assert color == "#bfa6a6"
+
+        assert legend.get_legend()["colors"] == {
+            "#a67573": "In",
+            "#bfa6a6": "Al",
+            "#ff0d0d": "O",
+            "#ffffff": "H",
+        }
+
+        # test categorical
+
+        legend = Legend(self.struct, color_scheme="example_categorical_site_prop")
+
+        assert legend.get_legend()["colors"] == {"#377eb8": "8b", "#e41a1c": "4a"}
+
+        # test pre-defined
+
+        legend = Legend(self.struct_manual)
+
+        assert legend.get_legend()["colors"] == {
+            "#0000ff": "O²⁻",
+            "#00ff00": "In",
+            "#ff0000": "H",
+        }
+
+    def test_get_radius(self):
+        legend = Legend(self.struct, radius_scheme="uniform")
+
+        assert legend.get_radius(sp=self.sp0) == 0.5
+
+        legend = Legend(self.struct, radius_scheme="covalent")
+
+        assert legend.get_radius(sp=self.sp1) == 0.66
+
+        legend = Legend(self.struct, radius_scheme="specified_or_average_ionic")
+
+        assert legend.get_radius(sp=self.sp2) == 0.94
+
+    def test_msonable(self):
+        legend = Legend(self.struct)
+        legend_dict = legend.as_dict()
+        legend_from_dict = Legend.from_dict(legend_dict)
+
+        assert legend.get_legend() == legend_from_dict.get_legend()
+
+    def test_get_tiling(self):
+        scene = self.struct.get_scene()
+        assert hasattr(scene, "lattice")
+        assert scene.lattice == [[5.0, 0, 0], [0, 5.0, 0], [0, 0, 5.0]]