Improve POV-Ray support and add CLI rendering interface

crystal_toolkit/cli/cli.py

@@ -0,0 +1,54 @@
+from pathlib import Path
+
+import rich_click as click
+from tqdm import tqdm
+
+
+@click.group()
+def cli():
+    pass
+
+
+@cli.command()
+@click.argument("input", required=True)
+def render(input):
+    from pymatgen.analysis.local_env import CrystalNN
+    from pymatgen.core.structure import Structure
+
+    from crystal_toolkit.core.scene import Scene
+    from crystal_toolkit.helpers.povray.renderer import POVRayRenderer
+
+    input_path = Path(input)
+    if input_path.is_file():
+        paths = [input_path]  # load CIF
+    else:
+        paths = list(input_path.glob("*.cif"))
+
+    r = POVRayRenderer()
+
+    structures = {}
+    for path in tqdm(paths, desc="Reading structures"):
+        try:
+            structures[path] = Structure.from_file(path)
+        except Exception as exc:
+            print(f"Failed to parse {path}: {exc}")
+
+    def _get_scene(struct: Structure) -> Scene:
+        # opinionated defaults, would be better to be customizable
+        nn = CrystalNN()
+        sg = nn.get_bonded_structure(struct)
+        return sg.get_scene(explicitly_calculate_polyhedra_hull=True)
+
+    scenes = {}
+    for path, structure in tqdm(structures.items(), desc="Preparing scenes"):
+        try:
+            scenes[path] = _get_scene(structure)
+        except Exception as exc:
+            print(f"Failed to parse {path}: {exc}")
+
+    for path, scene in tqdm(scenes.items(), desc="Rendering scenes"):
+        r.write_scene_to_file(scene, filename=f"{path.stem}.png")
+
+
+if __name__ == "__main__":
+    cli()

crystal_toolkit/helpers/__init__.py

@@ -0,0 +1 @@
+from crystal_toolkit.helpers.povray.renderer import POVRayRenderer

crystal_toolkit/helpers/povray/renderer.py

@@ -0,0 +1,179 @@
+"""Export wrapper for POV-Ray.
+
+For creating publication quality plots.
+"""
+
+from __future__ import annotations
+
+import os
+import shutil
+import subprocess
+from pathlib import Path
+from tempfile import TemporaryDirectory
+from typing import ClassVar
+from warnings import warn
+
+import numpy as np
+from jinja2 import Environment  # TODO: add to requirements
+from matplotlib.colors import to_hex
+
+from crystal_toolkit.core.scene import Cylinders, Lines, Primitive, Scene, Spheres
+from crystal_toolkit.settings import MODULE_PATH, SETTINGS
+
+
+class POVRayRenderer:
+    """A class to interface with the POV-Ray command line tool (ray tracer)."""
+
+    _TEMPLATES: ClassVar[dict[str, str]] = {
+        path.stem: path.read_text()
+        for path in (MODULE_PATH / "helpers" / "povray" / "templates").glob("*")
+    }
+    _ENV: ClassVar[Environment] = Environment()
+
+    @staticmethod
+    def write_scene_to_file(scene: Scene, filename: str | Path):
+        """Render a Scene to a PNG file using POV-Ray."""
+        current_dir = Path.cwd()
+
+        with TemporaryDirectory() as temp_dir:
+            os.chdir(temp_dir)
+
+            POVRayRenderer.write_povray_input_scene_and_settings(
+                scene, image_filename="crystal_toolkit_scene.png"
+            )
+            POVRayRenderer.call_povray()
+
+            shutil.copy("crystal_toolkit_scene.png", current_dir / filename)
+
+        os.chdir(current_dir)
+
+    @staticmethod
+    def call_povray(
+        povray_args: tuple[str] = ("render.ini",),
+        povray_path: str = SETTINGS.POVRAY_PATH,
+    ):
+        """
+        Run POV-Ray. Prefer `render_scene` method unless advanced user.
+        """
+
+        povray_args = [povray_path, *povray_args]
+        result = subprocess.run(
+            povray_args, capture_output=True, text=True, check=False
+        )
+
+        if result.returncode != 0:
+            raise RuntimeError(
+                f"{povray_path} exit code: {result.returncode}."
+                f"Please check your POV-Ray installation."
+                f"\nStdout:\n\n{result.stdout}\n\nStderr:\n\n{result.stderr}"
+            )
+
+    @staticmethod
+    def write_povray_input_scene_and_settings(
+        scene,
+        scene_filename="crystal_toolkit_scene.pov",
+        settings_filename="render.ini",
+        image_filename="crystal_toolkit_scene.png",
+    ):
+        """
+        Prefer `render_scene` method unless advanced user.
+        """
+
+        with open(scene_filename, "w") as f:
+            scene_str = POVRayRenderer.scene_to_povray(scene)
+
+            f.write(POVRayRenderer._TEMPLATES["header"])
+            f.write(POVRayRenderer._get_camera_for_scene(scene))
+            f.write(POVRayRenderer._TEMPLATES["lights"])
+            f.write(scene_str)
+
+        render_settings = POVRayRenderer._ENV.from_string(
+            POVRayRenderer._TEMPLATES["render"]
+        ).render(filename=scene_filename, image_filename=image_filename)
+        with open(settings_filename, "w") as f:
+            f.write(render_settings)
+
+    @staticmethod
+    def scene_to_povray(scene: Scene) -> str:
+        povray_str = ""
+
+        for item in scene.contents:
+            if isinstance(item, Primitive):
+                povray_str += POVRayRenderer.primitive_to_povray(obj=item)
+
+            elif isinstance(item, Scene):
+                povray_str += POVRayRenderer.scene_to_povray(scene=item)
+
+        return povray_str
+
+    @staticmethod
+    def primitive_to_povray(obj: Primitive) -> str:
+        vect = "{:.4f},{:.4f},{:.4f}"
+
+        if isinstance(obj, Spheres):
+            positions = obj.positions
+            positions = [vect.format(*pos) for pos in positions]
+            color = POVRayRenderer._format_color_to_povray(obj.color)
+
+            return POVRayRenderer._ENV.from_string(
+                POVRayRenderer._TEMPLATES["sphere"]
+            ).render(positions=positions, radius=obj.radius, color=color)
+
+        elif isinstance(obj, Cylinders):
+            position_pairs = [
+                [vect.format(*ipos), vect.format(*fpos)]
+                for ipos, fpos in obj.positionPairs
+            ]
+            color = POVRayRenderer._format_color_to_povray(obj.color)
+            return POVRayRenderer._ENV.from_string(
+                POVRayRenderer._TEMPLATES["cylinder"]
+            ).render(posPairs=position_pairs, color=color)
+
+        elif isinstance(obj, Lines):
+            pos1, pos2 = (
+                obj.positions[0::2],
+                obj.positions[1::2],
+            )
+            cylCaps = {tuple(pos) for pos in obj.positions}
+            cylCaps = [vect.format(*pos) for pos in cylCaps]
+            position_pairs = [
+                [vect.format(*ipos), vect.format(*fpos)]
+                for ipos, fpos in zip(pos1, pos2)
+            ]
+            return POVRayRenderer._ENV.from_string(
+                POVRayRenderer._TEMPLATES["line"]
+            ).render(posPairs=position_pairs, cylCaps=cylCaps)
+
+        elif isinstance(obj, Primitive):
+            warn(
+                f"Skipping {type(obj)}, not yet implemented. Submit PR to add support."
+            )
+
+        return ""
+
+    @staticmethod
+    def _format_color_to_povray(color: str) -> str:
+        """Convert a matplotlib-compatible color string to a POV-Ray color string."""
+        vect = "{:.4f},{:.4f},{:.4f}"
+        color = to_hex(color)
+        color = color.replace("#", "")
+        color = tuple(int(color[i : i + 2], 16) / 255.0 for i in (0, 2, 4))
+        return f"rgb<{vect.format(*color)}>"
+
+    @staticmethod
+    def _get_camera_for_scene(scene: Scene) -> str:
+        """Creates a camera in POV-Ray format for a given scene with respect to its bounding box."""
+
+        bounding_box = scene.bounding_box  # format is [min_corner, max_corner]
+        center = (np.array(bounding_box[0]) + bounding_box[1]) / 2
+        size = np.array(bounding_box[1]) - bounding_box[0]
+        camera_pos = center + np.array([0, 0, 1.2 * size[2]])
+
+        return f"""
+camera {{
+   orthographic
+   location <{camera_pos[0]:.4f}, {camera_pos[1]:.4f}, {camera_pos[2]:.4f}>
+   look_at <{center[0]:.4f}, {center[1]:.4f}, {center[2]:.4f}>
+   sky <0, 0, 1>
+}}
+"""

crystal_toolkit/helpers/povray/templates/cylinder.pov

@@ -0,0 +1,7 @@
+// Draw bonds between atoms in the scene
+
+#declare bond_texture = texture { pigment { {{color}} } finish { plastic_atom_finish } };
+
+{% for ipos, fpos in posPairs -%}
+cylinder { <{{ipos}}>, <{{fpos}}>, 0.1 texture { bond_texture } no_shadow }
+{% endfor %}

crystal_toolkit/helpers/povray/templates/header.pov

@@ -0,0 +1,27 @@
+#version 3.7 ;
+global_settings { assumed_gamma 1.8
+                  ambient_light rgb<1, 1, 1>
+}
+background { colour srgbt <0.0, 0.0, 0.0, 1.0> } // Set the background to transparent
+
+/*
+Create an Atom object along with some textures.
+The arguments are: Atom( position, radius, color, finish )
+*/
+
+#declare plastic_atom_finish = finish {
+                                specular 0.2
+                                roughness 0.001
+                                ambient 0.075
+                                diffuse 0.55
+                                brilliance 1.5
+                                conserve_energy
+                              }
+
+#macro Atom (P1, R1, C1, F1)
+  #local T = texture {
+                       pigment { C1 }
+                       finish { F1 }
+                     }
+  sphere { P1, R1 texture {T} no_shadow }
+#end

crystal_toolkit/helpers/povray/templates/lights.pov

@@ -0,0 +1,46 @@
+/*
+Define light sources to illuminate the atoms. For visualizing mediam
+media_interaction and media_attenuation are set to "off" so voxel
+data is rendered to be transparent. Lights are automatically oriented
+with respect to the camera position.
+*/
+
+// Overhead light source
+light_source {
+    <0, 0, 10>
+    color rgb <1,1,1>*0.5
+    parallel
+    point_at <ii, jj, kk>*0.5
+    media_interaction off
+    media_attenuation off
+}
+
+// Rear (forward-facing) light source
+light_source {
+    < (i-ii), (j-jj), (k-kk)>*4
+    color rgb <1,1,1> * 0.5
+    parallel
+    point_at <ii, jj, kk>
+    media_interaction off
+    media_attenuation off
+}
+
+// Left light source
+light_source {
+    <( (i-ii)*cos(60*pi/180) - (j-jj)*sin(60*pi/180) ), ( (i-ii)*sin(60*pi/180) + (j-jj)*cos(60*pi/180) ), k>
+    color rgb <1,1,1>*0.5
+    parallel
+    point_at <ii, jj, kk>
+    media_interaction off
+    media_attenuation off
+}
+
+// Right light source
+light_source {
+    <( (i-ii)*cos(-60*pi/180) - (j-jj)*sin(-60*pi/180) ), ( (i-ii)*sin(-60*pi/180) + (j-jj)*cos(-60*pi/180) ), k>
+    color rgb <1,1,1>*0.5
+    parallel
+    point_at <ii, jj, kk>
+    media_interaction off
+    media_attenuation off
+}

crystal_toolkit/helpers/povray/templates/line.pov

@@ -0,0 +1,11 @@
+// Draw the edges of the supercell in the scene
+
+#declare bbox = texture { pigment { rgb <1,1,1> } }
+
+{% for ipos, fpos in posPairs -%}
+cylinder {<{{ipos}}>, <{{fpos}}>, 0.02 texture {bbox} no_shadow}
+{% endfor %}
+
+{% for val in cylCaps -%}
+sphere {<{{val}}>, 0.02 texture {bbox} no_shadow}
+{% endfor %}

crystal_toolkit/helpers/povray/templates/render.ini

@@ -0,0 +1,22 @@
+Input_File_Name = {{filename}}
+Output_File_Name = {{image_filename}}
+Output_Alpha=On
+Display = 1
+# -- Option to switch on the density
+Declare=render_density=0     # 0 = off, 1 = on
+Quality = 9
+Height = 1200
+Width = 1600
+# -- Uncomment below for higher quality rendering
+Antialias = On
+Antialias_Threshold = 0.01
+Antialias_Depth = 4
+Jitter_Amount = 1.0
+# -- Set the camera position
+Declare=i=8
+Declare=j=8
+Declare=k=8
+# -- Set the look_at position
+Declare=ii=0
+Declare=jj=0
+Declare=kk=0

crystal_toolkit/helpers/povray/templates/sphere.pov

@@ -0,0 +1,5 @@
+// Draw Spheres
+
+{% for val in positions -%}
+Atom(<{{val}}>, {{radius}}, {{color}}, plastic_atom_finish)
+{% endfor %}