diff --git a/crystal_toolkit/cli/__init__.py b/crystal_toolkit/cli/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/crystal_toolkit/cli/cli.py b/crystal_toolkit/cli/cli.py new file mode 100644 index 00000000..b273c9f4 --- /dev/null +++ b/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() diff --git a/crystal_toolkit/helpers/__init__.py b/crystal_toolkit/helpers/__init__.py index e69de29b..17d4efcb 100644 --- a/crystal_toolkit/helpers/__init__.py +++ b/crystal_toolkit/helpers/__init__.py @@ -0,0 +1 @@ +from crystal_toolkit.helpers.povray.renderer import POVRayRenderer diff --git a/crystal_toolkit/helpers/povray/__init__.py b/crystal_toolkit/helpers/povray/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/crystal_toolkit/helpers/povray/renderer.py b/crystal_toolkit/helpers/povray/renderer.py new file mode 100644 index 00000000..983cd313 --- /dev/null +++ b/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> +}} +""" diff --git a/crystal_toolkit/helpers/povray/templates/cylinder.pov b/crystal_toolkit/helpers/povray/templates/cylinder.pov new file mode 100644 index 00000000..dff4f1ce --- /dev/null +++ b/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 %} diff --git a/crystal_toolkit/helpers/povray/templates/header.pov b/crystal_toolkit/helpers/povray/templates/header.pov new file mode 100644 index 00000000..78afb5b4 --- /dev/null +++ b/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 diff --git a/crystal_toolkit/helpers/povray/templates/lights.pov b/crystal_toolkit/helpers/povray/templates/lights.pov new file mode 100644 index 00000000..77b2b8cf --- /dev/null +++ b/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 *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 + 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 + 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 + media_interaction off + media_attenuation off +} diff --git a/crystal_toolkit/helpers/povray/templates/line.pov b/crystal_toolkit/helpers/povray/templates/line.pov new file mode 100644 index 00000000..34fa30ad --- /dev/null +++ b/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 %} diff --git a/crystal_toolkit/helpers/povray/templates/render.ini b/crystal_toolkit/helpers/povray/templates/render.ini new file mode 100644 index 00000000..c4155537 --- /dev/null +++ b/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 diff --git a/crystal_toolkit/helpers/povray/templates/sphere.pov b/crystal_toolkit/helpers/povray/templates/sphere.pov new file mode 100644 index 00000000..8b2cfbf3 --- /dev/null +++ b/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 %} diff --git a/crystal_toolkit/helpers/povray_renderer.py b/crystal_toolkit/helpers/povray_renderer.py deleted file mode 100644 index 8eb86dc2..00000000 --- a/crystal_toolkit/helpers/povray_renderer.py +++ /dev/null @@ -1,246 +0,0 @@ -"""Export wrapper for POV-Ray. - -For creating publication quality plots. -""" - -from __future__ import annotations - -from jinja2 import Environment - -HEAD = """ -#version 3.7 ; -global_settings { assumed_gamma 1.8 - ambient_light rgb<1, 1, 1> -} -background { rgb 0. } // Set the background to black - -/* -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 - -""" - -CAMERA = """ -/* -Define the camera and the view of the atoms -*/ - -camera { - orthographic - location - look_at - sky <0, 0, 1> -} - -""" - -LIGHTS = """ -/* -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 *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 - 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 - 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 - media_interaction off - media_attenuation off -} - -""" - -TEMP_SPHERE = """ -// Draw atoms in the scene - -{% for val in positions -%} -Atom(<{{val}}>, {{radius}}, {{color}}, plastic_atom_finish) -{% endfor %} -""" - -TEMP_CYLINDER = """ -// 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 %} -""" - -TEMP_LINE = """ -// 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 %} -""" - - -def pov_write_data(input_scene_comp, fstream): - """Parse a primitive display object in crystaltoolkit and print it to POV-Ray - input_scene_comp fstream. - """ - vect = "{:.4f},{:.4f},{:.4f}" - - if input_scene_comp["type"] == "spheres": - # Render atoms - positions = input_scene_comp["positions"] - positions = [vect.format(*pos) for pos in positions] - color = input_scene_comp["color"].replace("#", "") - color = tuple(int(color[i : i + 2], 16) / 255.0 for i in (0, 2, 4)) - color = f"rgb<{vect.format(*color)}>" - - fstream.write( - Environment() - .from_string(TEMP_SPHERE) - .render( - positions=positions, - radius=input_scene_comp["radius"], - color=color, - ) - ) - - if input_scene_comp["type"] == "cylinders": - # Render bonds between atoms - posPairs = [ - [vect.format(*ipos), vect.format(*fpos)] - for ipos, fpos in input_scene_comp["positionPairs"] - ] - color = input_scene_comp["color"].replace("#", "") - color = tuple(int(color[i : i + 2], 16) / 255.0 for i in (0, 2, 4)) - color = f"rgb<{vect.format(*color)}>" - fstream.write( - Environment() - .from_string(TEMP_CYLINDER) - .render(posPairs=posPairs, color=color) - ) - - if input_scene_comp["type"] == "lines": - # Render the cell - pos1, pos2 = ( - input_scene_comp["positions"][0::2], - input_scene_comp["positions"][1::2], - ) - cylCaps = {tuple(pos) for pos in input_scene_comp["positions"]} - cylCaps = [vect.format(*pos) for pos in cylCaps] - posPairs = [ - [vect.format(*ipos), vect.format(*fpos)] for ipos, fpos in zip(pos1, pos2) - ] - fstream.write( - Environment() - .from_string(TEMP_LINE) - .render(posPairs=posPairs, cylCaps=cylCaps) - ) - - -def filter_data(scene_data, fstream): - """Recursively traverse the scene_data dictionary to find objects to draw.""" - if "type" in scene_data: - pov_write_data(scene_data, fstream) - else: - for itr in scene_data["contents"]: - filter_data(itr, fstream) - - -def write_pov_file(smc, file_name): - """Args: - smc (StructureMoleculeComponent): Object containing the scene data. - file_name (str): name of the file to write to. - """ - with open(file_name, "w") as fstream: - fstream.write(HEAD) - fstream.write(CAMERA) - fstream.write(LIGHTS) - filter_data(smc.initial_scene_data, fstream) - - render_settings = get_render_settings() - with open(file_name, "w") as file: - file.write(render_settings) - - -def get_render_settings(file_name): - """Creates a POV-Ray render.ini file.""" - image_name = f"{file_name[:-4]}.png" - - return f""" -Input_File_Name = {file_name} -Output_File_Name = {image_name} -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=5 -Declare=k=4 -# -- Set the look_at position -Declare=ii=0 -Declare=jj=0 -Declare=kk=0 -""" diff --git a/crystal_toolkit/settings.py b/crystal_toolkit/settings.py index 6ed95ef4..50aa346d 100644 --- a/crystal_toolkit/settings.py +++ b/crystal_toolkit/settings.py @@ -1,6 +1,7 @@ from __future__ import annotations from pathlib import Path +from shutil import which from typing import Literal, Optional from pydantic import Field, HttpUrl, RedisDsn @@ -102,6 +103,13 @@ class Settings(BaseSettings): help="Default radius for displaying atoms when uniform radii are chosen.", ) + # Renderer settings. These control settings for any additional renderers like POV-Ray and Asymptote. + + POVRAY_PATH: Optional[str] = Field( + default=which("povray") or "/opt/homebrew/bin/povray", + help="Path to POV-Ray binary. Tested with 3.7.0.10.unofficial via `brew install povray` on macOS.", + ) + # Materials Project API settings. # TODO: These should be deprecated in favor of setti API_KEY: Optional[str] = Field(default="", help="Materials Project API key.") diff --git a/pyproject.toml b/pyproject.toml index e7b1609d..cc71c86a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -48,6 +48,9 @@ test = ["playwright", "pytest", "pytest-playwright"] repo = "https://github.com/materialsproject/crystaltoolkit" docs = "https://docs.crystaltoolkit.org" +[project.scripts] +ctk = "crystal_toolkit.cli.cli:cli" + [tool.setuptools.packages.find] exclude = ["docs_rst"]