diff --git a/cubehandler/__init__.py b/cubehandler/__init__.py index 1c3a0ff..1f9d61b 100644 --- a/cubehandler/__init__.py +++ b/cubehandler/__init__.py @@ -1,9 +1,10 @@ """CubeHandler package.""" -from .cube import Cube +from .cube import Cube, RenderSpec from .version import __version__ __all__ = [ "Cube", + "RenderSpec", "__version__", ] diff --git a/cubehandler/cli/main.py b/cubehandler/cli/main.py index 0649289..0edc411 100644 --- a/cubehandler/cli/main.py +++ b/cubehandler/cli/main.py @@ -1,13 +1,39 @@ -import typer -from pathlib import Path -from ..cube import Cube import enum +import re +from pathlib import Path +from typing import Optional + +import typer import yaml from glob import glob from ..version import __version__ +from ..cube import Cube, RenderSpec, SUPPORTED_IMAGE_FORMATS +from ..version import __version__ + app = typer.Typer(help="Cubehandler: a tool to handle cube files.") +ISO_PATTERN = re.compile( + r"^\s*([+-]?(?:\d+(?:\.\d*)?|\.\d+)(?:[eE][+-]?\d+)?)\s*:\s*(#[0-9A-Fa-f]{6})\s*$" +) +Orientation16 = tuple[ + float, + float, + float, + float, + float, + float, + float, + float, + float, + float, + float, + float, + float, + float, + float, + float, +] class Verbosity(enum.IntEnum): @@ -121,6 +147,85 @@ def run_reduction(inp, out): if verbosity >= Verbosity.INFO: typer.echo(yaml.dump(output, sort_keys=False)) +@app.command(help="Render a cube file.") +def render( + cube_file: str = typer.Argument(..., help="Path to the input cube file."), + orientation: Orientation16 = typer.Option( + ..., + "--orientation", + help="Camera orientation matrix from nglview (16 numbers).", + ), + iso: list[str] = typer.Option( + ..., + "--iso", + help="Isovalue/color pair in the form . Repeat for multiple pairs.", + ), + image_format: str = typer.Option( + ..., + "--format", + help=f"Output image format. Supported: {', '.join(SUPPORTED_IMAGE_FORMATS)}.", + ), + output: Optional[str] = typer.Option( + None, + "--output", + "-o", + help="Output image path. Defaults to _render..", + ), +): + """Render a cube file with nglview-compatible camera orientation.""" + input_path = Path(cube_file) + if not input_path.is_file(): + raise typer.BadParameter( + f"Cube file does not exist: {input_path}", param_hint="cube_file" + ) + + normalized_format = image_format.lower() + if normalized_format not in SUPPORTED_IMAGE_FORMATS: + raise typer.BadParameter( + f"Unsupported format '{image_format}'. " + f"Choose one of: {', '.join(SUPPORTED_IMAGE_FORMATS)}.", + param_hint="--format", + ) + + iso_pairs: list[tuple[float, str]] = [] + for iso_item in iso: + match = ISO_PATTERN.match(iso_item) + if match is None: + raise typer.BadParameter( + f"Invalid --iso value '{iso_item}'. Use .", + param_hint="--iso", + ) + isovalue = float(match.group(1)) + color = match.group(2).upper() + iso_pairs.append((isovalue, color)) + + if output is None: + output_path = Path.cwd() / f"{input_path.stem}_render.{normalized_format}" + else: + output_path = Path(output) + if not output_path.suffix: + output_path = output_path.with_suffix(f".{normalized_format}") + elif output_path.suffix.lower() != f".{normalized_format}": + raise typer.BadParameter( + f"Output extension '{output_path.suffix}' does not match --format '{normalized_format}'.", + param_hint="--output", + ) + + cube = Cube.from_file(input_path) + render_spec = RenderSpec( + orientation16=tuple(float(v) for v in orientation), + iso_pairs=tuple(iso_pairs), + image_format=normalized_format, + output_path=output_path, + ) + + try: + rendered_file = cube.render(render_spec) + except (ImportError, RuntimeError, ValueError) as exc: + raise typer.BadParameter(str(exc)) + + typer.echo(f"Rendered cube file: {rendered_file}") + def parse_pairs(pairs: list[str]) -> list[tuple[Path, float]]: if len(pairs) % 2 != 0: diff --git a/cubehandler/cube.py b/cubehandler/cube.py index c1bb1cc..9f8111c 100644 --- a/cubehandler/cube.py +++ b/cubehandler/cube.py @@ -4,12 +4,40 @@ import io import re +from dataclasses import dataclass +from pathlib import Path +from typing import Sequence import ase import numpy as np from skimage import transform ANG_TO_BOHR = 1.8897259886 +SUPPORTED_IMAGE_FORMATS = ("png", "jpg", "jpeg", "tif", "tiff", "bmp", "pnm", "ps") +_PIL_IMAGE_FORMATS = { + "png": "PNG", + "jpg": "JPEG", + "jpeg": "JPEG", + "tif": "TIFF", + "tiff": "TIFF", + "bmp": "BMP", + "pnm": "PPM", + "ps": "EPS", +} + + +@dataclass(frozen=True) +class RenderSpec: + orientation16: tuple[float, ...] + iso_pairs: tuple[tuple[float, str], ...] + image_format: str + output_path: Path + width: int = 1600 + height: int = 1200 + background: str = "#FFFFFF" + surface_opacity: float = 0.55 + atom_scale: float = 0.8 + bond_radius: float = 0.1 def remove_trailing_zeros(number): @@ -259,6 +287,187 @@ def reduce_data_density_skimage(self, resolution_factor=0.4): new_shape = tuple(int(dim * resolution_factor) for dim in self.data.shape) self.data = transform.resize(self.data, new_shape, anti_aliasing=True) + def render(self, spec: RenderSpec) -> Path: + """Render a cube file using a camera orientation from nglview.""" + import pyvista as pv + from PIL import Image + from ase import data as ase_data + from ase.data import colors as ase_colors + + image_format = spec.image_format.lower() + if image_format not in SUPPORTED_IMAGE_FORMATS: + raise ValueError( + f"Unsupported image format '{spec.image_format}'. " + f"Supported formats: {', '.join(SUPPORTED_IMAGE_FORMATS)}." + ) + if not spec.iso_pairs: + raise ValueError("At least one isovalue/color pair must be provided.") + + output_path = Path(spec.output_path) + output_path.parent.mkdir(parents=True, exist_ok=True) + + rotation, translation, scale = self._decompose_ngl_orientation( + spec.orientation16 + ) + structured_grid = self._build_structured_grid(rotation, translation) + + plotter = pv.Plotter(off_screen=True, window_size=(spec.width, spec.height)) + plotter.set_background(spec.background) + + for isovalue, color in spec.iso_pairs: + contour = structured_grid.contour( + isosurfaces=[isovalue], scalars="cube_data" + ) + if contour.n_points: + plotter.add_mesh( + contour, + color=color, + opacity=spec.surface_opacity, + smooth_shading=True, + ) + + transformed_atom_positions = self._transform_points( + self.ase_atoms.positions, rotation, translation + ) + atom_numbers = self.ase_atoms.get_atomic_numbers() + for number, position in zip(atom_numbers, transformed_atom_positions): + radius = float(ase_data.covalent_radii[number] * spec.atom_scale) + sphere = pv.Sphere( + radius=radius, + center=tuple(position), + theta_resolution=24, + phi_resolution=24, + ) + atom_color = tuple(float(c) for c in ase_colors.cpk_colors[number]) + plotter.add_mesh(sphere, color=atom_color, smooth_shading=True) + + bond_points, _ = self._compute_bonds(self.ase_atoms) + if len(bond_points): + transformed_bond_points = self._transform_points( + bond_points.reshape(-1, 3), rotation, translation + ).reshape(bond_points.shape) + for start, end in transformed_bond_points: + line = pv.Line(start, end) + tube = line.tube(radius=spec.bond_radius) + plotter.add_mesh(tube, color="#A0A0A0", smooth_shading=True) + + plotter.camera.position = (0.0, 0.0, -scale) + plotter.camera.focal_point = (0.0, 0.0, 0.0) + plotter.camera.up = (0.0, 1.0, 0.0) + image_array = plotter.screenshot(return_img=True) + plotter.close() + + image = Image.fromarray(image_array) + pil_format = _PIL_IMAGE_FORMATS[image_format] + if pil_format in {"JPEG", "EPS"} and image.mode in {"RGBA", "LA"}: + image = image.convert("RGB") + image.save(output_path, format=pil_format) + return output_path + + @staticmethod + def _transform_points( + points: np.ndarray, + rotation: np.ndarray, + translation: np.ndarray, + ) -> np.ndarray: + points = np.asarray(points, dtype=float) + return points @ rotation.T + translation + + @staticmethod + def _decompose_ngl_orientation( + orientation16: Sequence[float], + ) -> tuple[np.ndarray, np.ndarray, float]: + orientation_values = np.asarray(orientation16, dtype=float) + if orientation_values.size != 16: + raise ValueError("Camera orientation must contain exactly 16 numbers.") + if not np.all(np.isfinite(orientation_values)): + raise ValueError("Camera orientation must contain only finite numbers.") + + matrix = orientation_values.reshape((4, 4), order="F") + affine = matrix[:3, :3] + translation = matrix[:3, 3] + + column_norms = np.linalg.norm(affine, axis=0) + scale = float(np.mean(column_norms)) + if not np.isfinite(scale) or scale <= 1e-12: + raise ValueError("Could not derive camera scale from orientation matrix.") + + rotation = affine / scale + if np.linalg.matrix_rank(rotation) < 3: + raise ValueError("Camera orientation matrix is singular.") + + u_mat, _, vh_mat = np.linalg.svd(rotation) + rotation = u_mat @ vh_mat + if np.linalg.det(rotation) < 0: + u_mat[:, -1] *= -1 + rotation = u_mat @ vh_mat + + if not np.all(np.isfinite(rotation)): + raise ValueError("Camera orientation produced non-finite rotation values.") + return rotation, translation, scale + + def _build_structured_grid( + self, + rotation: np.ndarray, + translation: np.ndarray, + ): + import pyvista as pv + + nx, ny, nz = (int(v) for v in self.cell_n) + x_idx, y_idx, z_idx = np.meshgrid( + np.arange(nx, dtype=float), + np.arange(ny, dtype=float), + np.arange(nz, dtype=float), + indexing="ij", + ) + origin_ang = self.origin / ANG_TO_BOHR + voxel_vectors_ang = self.dcell_ang + + coords = ( + origin_ang + + x_idx[..., None] * voxel_vectors_ang[0] + + y_idx[..., None] * voxel_vectors_ang[1] + + z_idx[..., None] * voxel_vectors_ang[2] + ) + transformed_coords = self._transform_points( + coords.reshape(-1, 3), rotation, translation + ).reshape(coords.shape) + + grid = pv.StructuredGrid( + transformed_coords[..., 0], + transformed_coords[..., 1], + transformed_coords[..., 2], + ) + grid["cube_data"] = np.asarray(self.data, dtype=float).ravel(order="F") + return grid + + def _compute_bonds(self, structure): + """Create an list of bonds for the structure.""" + + import ase.neighborlist + from ase.data import colors + + if len(structure) <= 1: + return np.empty((0, 2, 3), dtype=float), np.empty((0, 3), dtype=float) + + # The value 1.09 is chosen based on our experience. It is a good compromise between showing too many bonds + # and not showing bonds that should be there. + cutoff = ase.neighborlist.natural_cutoffs(structure, mult=1.09) + + ii, bond_vectors = ase.neighborlist.neighbor_list( + "iD", structure, cutoff, self_interaction=False + ) + # bond start position + v1 = structure.positions[ii] + # middle position + v2 = v1 + bond_vectors * 0.5 + points = np.stack((v1, v2), axis=1) + + # Choose the correct way for computing the cylinder. + numbers = structure.get_atomic_numbers() + + return points, colors.cpk_colors[numbers[ii]] + def rescale_data(self, data): """Rescales the data to be between -1 and 1.""" scaling_factor = max(abs(data.min()), abs(data.max())) diff --git a/pyproject.toml b/pyproject.toml index 753ea16..da0a380 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -20,6 +20,7 @@ classifiers = [ requires-python = ">=3.10" dependencies = [ "ase", + "pyvista>=0.46", "scikit-image", "typer", ]