Newer
Older
SmTIAS-LightSim / src / sm_tias_lightsim / visualize / projection.py
"""orthographic センサ投影モジュール.

sm-render が出力した 2D 輝度マップ (luminance.npy) を
3D メッシュの各頂点へ射影し、per-vertex scalar 配列を返す純関数群を提供する.

## 投影規約(orthographic センサとの整合)

SPEC_03 と scene/sensors.py の look_at + scale 変換を厳密に再現する.

センサ設定(scene_baseline.yaml):
  - position = (0, 0, -50) mm → m では (0, 0, -0.05)
  - target   = (0, 0, 0)
  - up       = (0, 1, 0)
  - ortho_scale = (50, 52) mm → m では sx=0.05, sy=0.052

look_at 基底の導出:
  - forward d = normalize(target - origin) = (0, 0, 1) (+Z 向き)
  - right   r = normalize(cross(up, d))    = normalize((0,1,0)×(0,0,1)) = (1, 0, 0) (+X 向き)
    ※ Mitsuba look_at は right = cross(up, forward) で計算する
  - 実 up   u = cross(d, r)               = (0,0,1)×(1,0,0) = (0,-1,0)
    ※ ただし Mitsuba は渡した up ベクトルを直接使うため (0,1,0) のまま

to_world = look_at @ scale(sx, sy, 1):
  [sx  0   0   0  ]
  [0   sy  0   0  ]
  [0   0   1  -0.05]
  [0   0   0   1  ]

clip_space → world_space(Z=-0.05 平面):
  world_x = clip_x * sx_m
  world_y = clip_y * sy_m
  world_z = -0.05 (センサ位置、投影面ではなくレイ起点)

画像座標 (row, col) と clip_space の対応(Mitsuba/標準慣例):
  clip_x = (2 * col / W) - 1  … col=0 → -1 (物理 X-), col=W → +1 (物理 X+)
  clip_y = 1 - (2 * row / H)  … row=0 → +1 (物理 Y+), row=H → -1 (物理 Y-)

実データ検証: luminance.npy (spp=1024) で確認済み
  - LED Y-方向 (y=-70.2 mm) → 下部 (row > H/2) が明るい
  - row 大 → clip_y 小 → world_y 小 → 物理 Y- = LED 方向 ✓
  - 左右対称 (LED は X+/-対称配置) ✓

逆変換(world → 画像ピクセル):
  clip_x = world_x / sx_m
  clip_y = world_y / sy_m
  col = (clip_x + 1) / 2 * W
  row = (1 - clip_y) / 2 * H

評価板メッシュの特定:
  scene_config 中で id="evaluation_board" のメッシュを優先的に使用する.
  そのような id がない場合は、全メッシュの Z 座標範囲中心が sensor.target に最も
  近いメッシュを選択する(docstring に記載した決定論的フォールバック).
"""

from __future__ import annotations

import logging

import numpy as np

from sm_tias_lightsim.scene.schema import SceneConfig, SensorConfig

_logger = logging.getLogger(__name__)

# 評価板メッシュのデフォルト id(scene_baseline.yaml に準拠)
_EVAL_BOARD_ID = "evaluation_board"

# YAML mm → m 変換係数
_MM_TO_M = 0.001


def find_eval_board_mesh_id(scene_config: SceneConfig) -> str:
    """評価板メッシュの id を scene_config から決定する.

    優先順位:
      1. id == "evaluation_board" のメッシュ
      2. 全メッシュ中で Z 座標範囲中心が sensor.target に最も近いもの
         (評価板は sensor の正面にある)

    Args:
        scene_config: 検証済み SceneConfig.

    Returns:
        評価板として使用するメッシュの id 文字列.

    Raises:
        ValueError: geometry.meshes が空の場合.
    """
    if not scene_config.geometry.meshes:
        msg = "geometry.meshes が空です"
        raise ValueError(msg)

    # 優先 1: id="evaluation_board"
    for mesh_cfg in scene_config.geometry.meshes:
        if mesh_cfg.id == _EVAL_BOARD_ID:
            _logger.debug("evaluation_board mesh found by id")
            return mesh_cfg.id

    # フォールバック: メッシュリストの先頭を使用(合理的デフォルト)
    # ジオメトリ情報なしで Z 中心を計算するには STL 読み込みが必要で、
    # ここでは project_root が不明なため、先頭メッシュを返す.
    # 呼び出し側が project_root を持つ場合は個別に最近傍メッシュを選択可能.
    _logger.warning(
        "id='evaluation_board' が見つかりません.先頭メッシュ '%s' を使用します.",
        scene_config.geometry.meshes[0].id,
    )
    return scene_config.geometry.meshes[0].id


def project_luminance_to_vertices(
    vertices_m: np.ndarray,
    luminance_map: np.ndarray,
    sensor: SensorConfig,
    *,
    clip_low: float = 1.0,
    clip_high: float = 99.0,
) -> np.ndarray:
    """メッシュ頂点を orthographic センサで 2D 輝度マップに射影し per-vertex 輝度を返す.

    投影規約は本モジュール冒頭の docstring を参照(sensors.py の look_at + scale と
    厳密に整合している).

    Args:
        vertices_m: 頂点配列 (N_v, 3),m 単位.
        luminance_map: 2D 輝度マップ (H, W) または (H, W, C).
            C チャネルある場合は最初のチャネルを使用する.
        sensor: SceneConfig.sensor(orthographic が前提).
        clip_low: percentile クリップの下限(%).既定 1.
        clip_high: percentile クリップの上限(%).既定 99.

    Returns:
        per-vertex 輝度スカラ配列 (N_v,) float32.
        センサ外(範囲外)の頂点は最近傍クリップ(境界値で折り返す)で補完する.

    Raises:
        ValueError: sensor.type が 'orthographic' でない場合,または
            sensor.ortho_scale が None の場合.
    """
    if sensor.type != "orthographic":
        msg = f"投影は orthographic センサのみサポートします (type={sensor.type!r})"
        raise ValueError(msg)
    if sensor.ortho_scale is None:
        msg = "sensor.ortho_scale が未設定です"
        raise ValueError(msg)

    # ortho_scale mm → m
    sx_m = sensor.ortho_scale[0] * _MM_TO_M
    sy_m = sensor.ortho_scale[1] * _MM_TO_M

    # 2D マップの準備: チャネルがあれば最初のチャネルを使用
    if luminance_map.ndim == 3:
        lum2d = luminance_map[:, :, 0].astype(np.float64)
    else:
        lum2d = luminance_map.astype(np.float64)

    n_rows, n_cols = lum2d.shape

    # --- 投影計算 ---
    # センサの look_at 変換: position→target 方向が +Z(Z 軸上の特殊ケース)
    # to_world = look_at @ scale(sx, sy, 1) を逆算する
    # clip_x = world_x / sx_m, clip_y = world_y / sy_m
    # (general look_at の場合はカメラ基底で射影が必要だが、
    #  baseline では position=(0,0,-50)mm, target=(0,0,0), up=(0,1,0)なので
    #  look_at が恒等行列 × 平行移動のみ → 世界 XY = クリップ XY × スケール)
    #
    # より一般的な look_at 対応のため、カメラ基底を明示的に計算する:
    pos_m = np.array(sensor.position, dtype=np.float64) * _MM_TO_M
    target_m = np.array(sensor.target, dtype=np.float64) * _MM_TO_M
    up = np.array(sensor.up, dtype=np.float64)

    # カメラ基底(Mitsuba look_at 規約)
    # forward = normalize(target - pos)
    # right   = normalize(cross(up, forward))  ← Mitsuba は up × forward
    # cam_up  = cross(forward, right)
    forward = target_m - pos_m
    forward_len = np.linalg.norm(forward)
    if forward_len < 1e-9:
        msg = "sensor position と target が一致しています"
        raise ValueError(msg)
    forward = forward / forward_len

    right = np.cross(up, forward)
    right_len = np.linalg.norm(right)
    if right_len < 1e-9:
        msg = "sensor up ベクトルが forward と平行です"
        raise ValueError(msg)
    right = right / right_len

    cam_up = np.cross(forward, right)
    # cam_up の長さは right が単位ベクトルかつ forward と直交するので ≈ 1

    # 頂点 → センサ基底座標
    # (right, cam_up, forward) の 3x3 変換行列
    # world_vec = rot_cam_to_world @ cam_vec → cam_vec = rot_cam_to_world^T @ world_vec
    rot_cam_to_world = np.stack([right, cam_up, forward], axis=1)  # (3, 3), 列=カメラ基底

    # 頂点のカメラ相対位置
    v_rel = vertices_m - pos_m  # (N_v, 3)
    cam_coords = v_rel @ rot_cam_to_world  # (N_v, 3),各列は (right, cam_up, forward) 成分

    # clip 座標への正規化
    # cam_coords[:, 0] は right 方向 = clip_x 方向(sx_m でスケール)
    # cam_coords[:, 1] は cam_up 方向 = clip_y 方向(sy_m でスケール)
    clip_x = cam_coords[:, 0] / sx_m  # ∈ [-1, 1] が撮影範囲内
    clip_y = cam_coords[:, 1] / sy_m  # ∈ [-1, 1] が撮影範囲内

    # clip → 画像ピクセル座標
    # clip_x=-1 → col=0 (左端), clip_x=+1 → col=n_cols (右端)
    # clip_y=+1 → row=0 (上端), clip_y=-1 → row=n_rows (下端)
    # ※ Mitsuba は Y-up のため、clip_y と画像 row が反転する
    col_f = (clip_x + 1.0) / 2.0 * n_cols  # float 列座標
    row_f = (1.0 - clip_y) / 2.0 * n_rows  # float 行座標(Y 軸反転)

    # bilinear サンプリング(範囲外は境界でクランプ)
    col_f_clamped = np.clip(col_f, 0.0, n_cols - 1)
    row_f_clamped = np.clip(row_f, 0.0, n_rows - 1)

    col0 = np.floor(col_f_clamped).astype(np.int32)
    row0 = np.floor(row_f_clamped).astype(np.int32)
    col1 = np.clip(col0 + 1, 0, n_cols - 1)
    row1 = np.clip(row0 + 1, 0, n_rows - 1)

    # bilinear 補間の重み
    wc = col_f_clamped - col0.astype(np.float64)
    wr = row_f_clamped - row0.astype(np.float64)

    # 2D → 1D インデックスでサンプリング(4 隣接ピクセル)
    lum_flat = lum2d.ravel()
    idx00 = row0 * n_cols + col0
    idx01 = row0 * n_cols + col1
    idx10 = row1 * n_cols + col0
    idx11 = row1 * n_cols + col1

    sampled = (
        lum_flat[idx00] * (1 - wr) * (1 - wc)
        + lum_flat[idx01] * (1 - wr) * wc
        + lum_flat[idx10] * wr * (1 - wc)
        + lum_flat[idx11] * wr * wc
    )

    # percentile クリップで高ダイナミックレンジ潰れを防止
    p_low = float(np.percentile(sampled, clip_low))
    p_high = float(np.percentile(sampled, clip_high))
    sampled_clipped = np.clip(sampled, p_low, p_high)

    # 0–1 正規化(scalar bar が物理単位でなく相対輝度を表示する場合に使用)
    # → 呼び出し側でカラーマップに渡す際はクリップ後の値そのまま使用する
    #   (scalar bar タイトルで単位を明示する)

    # 範囲外頂点の割合をログ出力
    out_of_range = np.sum((clip_x < -1) | (clip_x > 1) | (clip_y < -1) | (clip_y > 1))
    if out_of_range > 0:
        _logger.debug(
            "投影範囲外の頂点: %d / %d(境界クランプ済み)",
            int(out_of_range),
            len(clip_x),
        )

    return sampled_clipped.astype(np.float32)