Newer
Older
SmTIAS-LightSim / src / sm_tias_lightsim / scene / lights.py
# pyright: reportCallIssue=none
# Mitsuba 3 は variant ベースの動的属性解決を行うため Pylance/Pyright が
# mi.* を object 型と推論し reportCallIssue を誤検知する(coding-style.md の方針).
"""YAML 光源定義 → Mitsuba emitter dict 変換モジュール."""

from __future__ import annotations

from sm_tias_lightsim.runtime.llvm_init import ensure_drjit_libllvm_path
from sm_tias_lightsim.scene.materials import _to_rgb
from sm_tias_lightsim.scene.schema import (
    AreaLight,
    DirectionalLight,
    LightConfig,
    PointLight,
    SpotLight,
)

# mitsuba を import した時点で Dr.Jit が LLVM の初期化を試みる.そのため
# DRJIT_LIBLLVM_PATH の自動探索は import より前に済ませておく必要がある.
# 本モジュールはモジュールレベルで mitsuba を import するので,
# render/runner.py の探索呼び出しでは手遅れになる(環境変数未設定の環境では
# レンダリングが ImportError で失敗し,slow マークのレンダテストが全滅していた).
ensure_drjit_libllvm_path()

import mitsuba as mi  # noqa: E402


def _build_area_light_transform(
    center: tuple[float, float, float],
    half_extents: tuple[float, float],
    direction: tuple[float, float, float],
    position_scale: float,
) -> mi.ScalarTransform4f:
    """AreaLight の中心・半幅・方向から Mitsuba の to_world Transform を組み立てる.

    Mitsuba の rectangle / disk 形状はどちらも XY 平面上の単位図形
    (rectangle は ±1 の正方形,disk は半径 1 の円)で法線が +Z.
    direction ベクトルが面の法線方向を指すように look_at を構成する.

    Args:
        center: 面中心位置(元単位).
        half_extents: X / Y 方向の半幅 (hx, hy)(元単位).
            rectangle は (width/2, height/2),disk は (radius, radius) を渡す.
        direction: 放射方向ベクトル(単位ベクトルでなくても可).
        position_scale: 座標スケール係数(mm→m なら 0.001).

    Returns:
        Mitsuba ScalarTransform4f.
    """
    cx, cy, cz = [c * position_scale for c in center]
    hx, hy = half_extents[0] * position_scale, half_extents[1] * position_scale

    dx, dy, dz = direction
    norm = (dx**2 + dy**2 + dz**2) ** 0.5
    if norm == 0.0:
        msg = "AreaLight direction vector must not be zero."
        raise ValueError(msg)
    dx, dy, dz = dx / norm, dy / norm, dz / norm

    # direction が Y 軸に近い場合は up を Z 軸とし,それ以外は Y 軸を up に使う
    # look_at の origin を center + direction(目線は放射方向の逆から見る)として構成
    if abs(dy) < 0.9:
        up = [0.0, 1.0, 0.0]
    else:
        up = [0.0, 0.0, 1.0]

    # look_at: origin=center,target=center + direction,up=up
    origin = [cx, cy, cz]
    target = [cx + dx, cy + dy, cz + dz]

    transform = mi.ScalarTransform4f().look_at(origin=origin, target=target, up=up)

    # 単位図形(±1 / 半径 1)を半幅でスケールして実寸に合わせる
    scale = mi.ScalarTransform4f().scale([hx, hy, 1.0])

    return transform @ scale


def build_light_dict(light: LightConfig, position_scale: float) -> dict:
    """1 LightConfig を Mitsuba scene dict の emitter エントリに変換する.

    Args:
        light: SceneConfig.lights の 1 要素(PointLight / AreaLight / DirectionalLight).
        position_scale: 座標スケール係数(例:mm→m なら 0.001).
                        位置・サイズに掛けて Mitsuba(m)に揃える.

    Returns:
        Mitsuba scene dict のキーに使える dict.
        - PointLight → {"type": "point", "position": [...], "intensity": {...}}
        - AreaLight → {"type": "rectangle" | "disk", "to_world": ...,
          "emitter": {"type": "area", ...}}
        - DirectionalLight → {"type": "directional", "direction": [...], "irradiance": {...}}

    Raises:
        ValueError: 未対応の type が指定された場合.
    """
    if isinstance(light, PointLight):
        pos = [c * position_scale for c in light.position]
        return {
            "type": "point",
            "position": pos,
            "intensity": {
                "type": "rgb",
                "value": _to_rgb(light.intensity),
            },
        }

    if isinstance(light, AreaLight):
        # 半幅: rectangle は size の半分,disk は radius(スキーマで必須検証済み)
        if light.shape == "disk":
            half_extents = (light.radius, light.radius)
        else:
            half_extents = (light.size[0] / 2.0, light.size[1] / 2.0)
        to_world = _build_area_light_transform(
            center=light.center,
            half_extents=half_extents,
            direction=light.direction,
            position_scale=position_scale,
        )
        return {
            "type": light.shape,
            "to_world": to_world,
            "emitter": {
                "type": "area",
                "radiance": {
                    "type": "rgb",
                    "value": _to_rgb(light.radiance),
                },
            },
        }

    if isinstance(light, DirectionalLight):
        return {
            "type": "directional",
            "direction": list(light.direction),
            "irradiance": {
                "type": "rgb",
                "value": _to_rgb(light.irradiance),
            },
        }

    if isinstance(light, SpotLight):
        # spot emitter は to_world で位置と光軸方向を指定する.
        # Mitsuba の spot は to_world @ (origin=(0,0,0), target=(0,0,1)) を基準とし,
        # to_world で世界座標へ変換するので look_at で組み立てる.
        cx, cy, cz = [c * position_scale for c in light.position]
        dx, dy, dz = light.direction
        norm = (dx**2 + dy**2 + dz**2) ** 0.5
        if norm == 0.0:
            msg = "SpotLight direction vector must not be zero."
            raise ValueError(msg)
        dx, dy, dz = dx / norm, dy / norm, dz / norm
        if abs(dy) < 0.9:
            up = [0.0, 1.0, 0.0]
        else:
            up = [0.0, 0.0, 1.0]
        origin = [cx, cy, cz]
        target = [cx + dx, cy + dy, cz + dz]
        to_world = mi.ScalarTransform4f().look_at(origin=origin, target=target, up=up)
        return {
            "type": "spot",
            "to_world": to_world,
            "cutoff_angle": light.cutoff_angle,
            "beam_width": light.beam_width,
            "intensity": {
                "type": "rgb",
                "value": _to_rgb(light.intensity),
            },
        }

    msg = f"Unsupported light type: {type(light).__name__}"
    raise ValueError(msg)