Newer
Older
SmTIAS-LightSim / src / sm_tias_lightsim / scene / schema.py
"""YAML シーン定義スキーマ(Pydantic モデル).

SPEC_02_シーン定義仕様 に準拠したモデル群と,YAML ファイルローダを提供する.
"""

from __future__ import annotations

from pathlib import Path
from typing import Annotated, Literal

import yaml
from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator

# ---------------------------------------------------------------------------
# 共通基底・型エイリアス
# ---------------------------------------------------------------------------


# 全モデルで frozen=True を共有するため基底クラスに集約する
class _FrozenBase(BaseModel):
    model_config = ConfigDict(frozen=True)


# 3D ベクトル(mm 座標・方向ベクトル・RGB 輝度)の共通型
Vec3 = tuple[float, float, float]


# ---------------------------------------------------------------------------
# Film / Sensor
# ---------------------------------------------------------------------------


class FilmConfig(_FrozenBase):
    """レンダリングフィルム(出力解像度・ピクセルフィルタ)設定."""

    width: int = Field(..., ge=1, description="フィルム幅 [px]")
    height: int = Field(..., ge=1, description="フィルム高さ [px]")
    rfilter: Literal["gaussian", "box", "tent"] = Field(
        default="gaussian", description="ピクセル再構成フィルタ"
    )


class SensorConfig(_FrozenBase):
    """カメラセンサ設定.現在は perspective のみ対応."""

    type: Literal["perspective"]
    position: Vec3 = Field(..., description="カメラ位置 [mm]")
    target: Vec3 = Field(..., description="注視点 [mm]")
    up: Vec3 = Field(..., description="上方向ベクトル")
    fov_deg: float = Field(..., description="水平方向 FOV [deg]")
    film: FilmConfig

    @field_validator("fov_deg")
    @classmethod
    def _validate_fov(cls, v: float) -> float:
        """FOV が 0 < fov_deg < 180 の範囲内であることを検証する."""
        if not (0.0 < v < 180.0):
            msg = f"fov_deg must be in range (0, 180), got {v}"
            raise ValueError(msg)
        return v


# ---------------------------------------------------------------------------
# Render
# ---------------------------------------------------------------------------


class RenderConfig(_FrozenBase):
    """レンダリングパラメータ設定."""

    variant: str = Field(default="llvm_ad_rgb", description="Mitsuba variant 文字列")
    integrator: Literal["path", "direct", "volpath"] = Field(
        default="path", description="積分器の種類"
    )
    max_depth: int = Field(default=8, ge=1, description="最大光線追跡深度")
    spp: int = Field(default=256, ge=1, description="ピクセルあたりのサンプル数")


# ---------------------------------------------------------------------------
# Materials(discriminated union)
# ---------------------------------------------------------------------------


class DiffuseMaterial(_FrozenBase):
    """Lambertian 拡散反射材料."""

    type: Literal["diffuse"]
    reflectance: float | Vec3 = Field(..., description="反射率(スカラまたは RGB タプル)")


class _IorBase(_FrozenBase):
    """屈折率フィールドを持つ誘電体材料の共通基底."""

    int_ior: float = Field(..., description="内側の屈折率")
    ext_ior: float = Field(..., description="外側の屈折率")


class DielectricMaterial(_IorBase):
    """誘電体材料(厚みあり)."""

    type: Literal["dielectric"]


class ThinDielectricMaterial(_IorBase):
    """薄膜誘電体材料(ディフューザー薄板向け)."""

    type: Literal["thindielectric"]


# type フィールドで識別される discriminated union
MaterialConfig = Annotated[
    DiffuseMaterial | DielectricMaterial | ThinDielectricMaterial,
    Field(discriminator="type"),
]


# ---------------------------------------------------------------------------
# Lights(discriminated union)
# ---------------------------------------------------------------------------


class PointLight(_FrozenBase):
    """点光源(LED 単体モデル等)."""

    id: str = Field(..., description="シーン内ユニーク ID")
    type: Literal["point"]
    position: Vec3 = Field(..., description="位置 [mm]")
    intensity: float | Vec3 = Field(..., description="強度(スカラまたは RGB タプル)")
    profile: Literal["lambertian", "isotropic"] = Field(..., description="配光プロファイル")


class AreaLight(_FrozenBase):
    """面光源(矩形)."""

    id: str = Field(..., description="シーン内ユニーク ID")
    type: Literal["area"]
    shape: Literal["rectangle"] = Field(..., description="面形状")
    center: Vec3 = Field(..., description="中心位置 [mm]")
    size: tuple[float, float] = Field(..., description="サイズ [mm, mm]")
    direction: Vec3 = Field(..., description="放射方向ベクトル")
    radiance: float | Vec3 = Field(..., description="放射輝度(スカラまたは RGB タプル)")


class DirectionalLight(_FrozenBase):
    """平行光源(無限遠光源)."""

    id: str = Field(..., description="シーン内ユニーク ID")
    type: Literal["directional"]
    direction: Vec3 = Field(..., description="光の進行方向ベクトル")
    irradiance: float | Vec3 = Field(..., description="照度(スカラまたは RGB タプル)")


# type フィールドで識別される discriminated union
LightConfig = Annotated[
    PointLight | AreaLight | DirectionalLight,
    Field(discriminator="type"),
]


# ---------------------------------------------------------------------------
# Geometry
# ---------------------------------------------------------------------------


class MeshConfig(_FrozenBase):
    """メッシュ(STL)定義."""

    id: str = Field(..., description="シーン内ユニーク ID")
    path: str = Field(..., description="プロジェクトルートからの相対パス(存在チェックは行わない)")
    unit: Literal["mm", "m"] = Field(..., description="モデルの単位系")
    material: str = Field(..., description="materials の dict キー参照")


class GeometryConfig(_FrozenBase):
    """ジオメトリ設定(メッシュ一覧)."""

    meshes: list[MeshConfig] = Field(..., min_length=1, description="メッシュ定義リスト")


# ---------------------------------------------------------------------------
# Evaluation
# ---------------------------------------------------------------------------


class RoiConfig(_FrozenBase):
    """評価 ROI 設定(image_pixel 型のみ対応).

    scene_face 型は将来実装予定.
    """

    type: Literal["image_pixel"]
    x: int = Field(..., ge=0, description="ROI 左上 X 座標 [px]")
    y: int = Field(..., ge=0, description="ROI 左上 Y 座標 [px]")
    width: int = Field(..., ge=1, description="ROI 幅 [px]")
    height: int = Field(..., ge=1, description="ROI 高さ [px]")


class ZonesConfig(_FrozenBase):
    """3 ゾーン分析設定(中心 / 中間 / 周辺)."""

    enabled: bool = Field(default=True, description="ゾーン分析の有効/無効")
    boundaries: tuple[float, float] = Field(
        default=(0.33, 0.66),
        description="境界値 (b1, b2).0 < b1 < b2 < 1 を満たす正規化距離",
    )

    @field_validator("boundaries")
    @classmethod
    def _validate_boundaries(cls, v: tuple[float, float]) -> tuple[float, float]:
        """境界値が 0 < b1 < b2 < 1 を満たすことを検証する."""
        b1, b2 = v
        if not (0.0 < b1 < b2 < 1.0):
            msg = f"boundaries must satisfy 0 < b1 < b2 < 1, got ({b1}, {b2})"
            raise ValueError(msg)
        return v


class MetricsConfig(_FrozenBase):
    """算出する均一性指標フラグ設定."""

    cov: bool = Field(default=True, description="変動係数 (CoV)")
    std: bool = Field(default=True, description="標準偏差")
    max_min: bool = Field(default=True, description="最大値/最小値比")
    radial_profile: bool = Field(default=True, description="放射状プロファイル")


class EvaluationConfig(_FrozenBase):
    """評価設定(ROI・ゾーン・指標)."""

    roi: RoiConfig
    zones: ZonesConfig | None = None
    metrics: MetricsConfig | None = None


# ---------------------------------------------------------------------------
# Top-level SceneConfig
# ---------------------------------------------------------------------------


class SceneConfig(_FrozenBase):
    """シーン定義全体を表すトップレベルモデル.

    SPEC_02_シーン定義仕様 に準拠した構造を持つ.
    model_validator でメッシュ・光源 ID の一意性,および材料参照整合性を検証する.
    """

    name: str = Field(..., description="設計識別名(出力ディレクトリ名に使用)")
    description: str | None = Field(default=None, description="自由記述")
    geometry: GeometryConfig
    materials: dict[str, MaterialConfig]
    lights: list[LightConfig] = Field(..., min_length=1, description="光源リスト(1 個以上)")
    sensor: SensorConfig
    render: RenderConfig = Field(default_factory=RenderConfig)
    evaluation: EvaluationConfig | None = None

    @model_validator(mode="after")
    def _validate_cross_references(self) -> SceneConfig:
        """メッシュ/光源 ID の一意性と材料参照整合性を検証する.

        検証内容:
            1. geometry.meshes の各 id が重複していないこと.
            2. lights の各 id が重複していないこと.
            3. 各 mesh.material が materials の dict キーに存在すること.

        Returns:
            検証済みの SceneConfig 自身.

        Raises:
            ValueError: 重複 ID または未定義材料参照が見つかった場合.
        """
        # 1. mesh.id ユニーク性チェック
        seen_mesh_ids: set[str] = set()
        for mesh in self.geometry.meshes:
            if mesh.id in seen_mesh_ids:
                msg = f"Duplicate mesh id: '{mesh.id}'"
                raise ValueError(msg)
            seen_mesh_ids.add(mesh.id)

        # 2. light.id ユニーク性チェック
        # LightConfig は discriminated union なので各エントリは必ず id を持つ
        seen_light_ids: set[str] = set()
        for light in self.lights:
            if light.id in seen_light_ids:
                msg = f"Duplicate light id: '{light.id}'"
                raise ValueError(msg)
            seen_light_ids.add(light.id)

        # 3. mesh.material 参照整合性チェック
        material_keys = set(self.materials.keys())
        for mesh in self.geometry.meshes:
            if mesh.material not in material_keys:
                msg = f"mesh '{mesh.id}' references unknown material '{mesh.material}'"
                raise ValueError(msg)

        return self


# ---------------------------------------------------------------------------
# YAML Loader
# ---------------------------------------------------------------------------


def load_scene_yaml(path: str | Path) -> SceneConfig:
    """YAML ファイルを読み込み,検証済み SceneConfig を返す.

    Args:
        path: YAML ファイルのパス.

    Returns:
        検証済みの SceneConfig インスタンス.

    Raises:
        FileNotFoundError: ファイルが存在しない場合.
        pydantic.ValidationError: スキーマ違反.
        yaml.YAMLError: YAML 構文エラー.
    """
    resolved = Path(path)
    if not resolved.exists():
        msg = f"Scene YAML file not found: {resolved}"
        raise FileNotFoundError(msg)

    with resolved.open(encoding="utf-8") as f:
        data = yaml.safe_load(f)

    return SceneConfig.model_validate(data)