lamp-testbench/crates/tm30/tools/extract_data.py

"""Extract the spectral tables used by the TM-30-18 path of colour-science.

Writes plain CSV files into ``crates/tm30/data`` containing the *effective*
arrays colour-science 0.4.7 uses when ``colour_fidelity_index_ANSIIESTM3018``
is called on a 380-780 nm / 1 nm spectral distribution (which is what
``tm30.py`` always builds). "Effective" matters: colour re-aligns its 1 nm
CMFs and 5 nm D-series basis functions with a Sprague interpolator, which is
not bit-exact even on the original nodes, so the arrays are dumped exactly as
the pipeline sees them (same reshape calls, same order).

Usage (from the repository root, with a Python that has colour-science 0.4.7):

    python -W ignore crates/tm30/tools/extract_data.py

Values are written with ``repr`` (shortest round-trip), so Rust's
``str::parse::<f64>`` recovers the identical doubles.
"""

from __future__ import annotations

import os
import sys
import warnings

import numpy as np

with warnings.catch_warnings():
    warnings.simplefilter("ignore")
    import colour
    from colour.appearance.ciecam02 import CAT_INVERSE_CAT02, MATRIX_XYZ_TO_HPE
    from colour.adaptation import CAT_CAT02
    from colour.colorimetry import (
        MSDS_CMFS,
        SDS_BASIS_FUNCTIONS_CIE_ILLUMINANT_D_SERIES,
        SDS_ILLUMINANTS,
        SPECTRAL_SHAPE_DEFAULT,
        SpectralShape,
        reshape_msds,
        reshape_sd,
    )
    from colour.quality.cfi2017 import load_TCS_CIE2017

HERE = os.path.dirname(os.path.abspath(__file__))
DATA = os.path.normpath(os.path.join(HERE, "..", "data"))

SHAPE = SpectralShape(380, 780, 1)


def fmt(x: float) -> str:
    return repr(float(x))


def write_csv(name: str, header: list[str], rows) -> None:
    path = os.path.join(DATA, name)
    with open(path, "w", newline="\n") as f:
        f.write(",".join(header) + "\n")
        for row in rows:
            f.write(",".join(fmt(v) for v in row) + "\n")
    print(f"wrote {path}")


def main() -> int:
    assert colour.__version__ == "0.4.7", colour.__version__
    os.makedirs(DATA, exist_ok=True)

    with warnings.catch_warnings():
        warnings.simplefilter("ignore")

        # --- CIE 1931 2 degree observer -------------------------------------
        # handle_spectral_arguments() reshapes the default observer to
        # SPECTRAL_SHAPE_DEFAULT (360-780 nm, 1 nm); this array feeds the
        # Ohno (2013) Planckian table (CCT_to_uv_Planck1900).
        cmfs_2_default = reshape_msds(
            MSDS_CMFS["CIE 1931 2 Degree Standard Observer"],
            SPECTRAL_SHAPE_DEFAULT,
            copy=False,
        )
        assert cmfs_2_default.shape == SpectralShape(360, 780, 1)
        # sd_to_XYZ_integration() then re-aligns *that* array to the test
        # shape (380-780 nm) for the test SPD CCT and the reference mixing
        # normalisation.
        cmfs_2_test = reshape_msds(cmfs_2_default, SHAPE, copy=False)

        # --- CIE 1964 10 degree observer ------------------------------------
        cmfs_10 = reshape_msds(
            MSDS_CMFS["CIE 1964 10 Degree Standard Observer"], SHAPE, copy=False
        )

        # --- Illuminant E (implicit weighting in the integration) -----------
        for shape in (SPECTRAL_SHAPE_DEFAULT, SHAPE):
            E = reshape_sd(SDS_ILLUMINANTS["E"], shape, copy=False)
            assert np.all(E.values == 100.0), "illuminant E is not exactly 100"
        E2 = reshape_sd(
            reshape_sd(SDS_ILLUMINANTS["E"], SPECTRAL_SHAPE_DEFAULT, copy=False),
            SHAPE,
            copy=False,
        )
        assert np.all(E2.values == 100.0)

        # --- CIE D-series basis functions -----------------------------------
        basis = [
            reshape_sd(SDS_BASIS_FUNCTIONS_CIE_ILLUMINANT_D_SERIES[n], SHAPE, copy=False)
            for n in ("S0", "S1", "S2")
        ]

        # --- CIE 2017 test colour samples -----------------------------------
        tcs = load_TCS_CIE2017(SHAPE)
        assert tcs.shape == SHAPE and tcs.values.shape == (401, 99)

    wl_default = cmfs_2_default.wavelengths
    write_csv(
        "cmf_cie1931_2deg_360_780.csv",
        ["wavelength_nm", "x_bar", "y_bar", "z_bar"],
        (
            [w, *v]
            for w, v in zip(wl_default, cmfs_2_default.values, strict=True)
        ),
    )
    write_csv(
        "cmf_cie1931_2deg_380_780.csv",
        ["wavelength_nm", "x_bar", "y_bar", "z_bar"],
        ([w, *v] for w, v in zip(SHAPE.wavelengths, cmfs_2_test.values, strict=True)),
    )
    write_csv(
        "cmf_cie1964_10deg_380_780.csv",
        ["wavelength_nm", "x_bar", "y_bar", "z_bar"],
        ([w, *v] for w, v in zip(SHAPE.wavelengths, cmfs_10.values, strict=True)),
    )
    write_csv(
        "d_series_basis_380_780.csv",
        ["wavelength_nm", "S0", "S1", "S2"],
        (
            [w, a, b, c]
            for w, a, b, c in zip(
                SHAPE.wavelengths,
                basis[0].values,
                basis[1].values,
                basis[2].values,
                strict=True,
            )
        ),
    )
    write_csv(
        "tcs_cie2017_380_780.csv",
        ["wavelength_nm", *[f"TCS{i:02d}" for i in range(99)]],
        ([w, *v] for w, v in zip(SHAPE.wavelengths, tcs.values, strict=True)),
    )

    # CIECAM02 matrices exactly as colour uses them: CAT02 (literal) and the
    # numpy product MATRIX_XYZ_TO_HPE @ inv(CAT02) used by RGB_to_rgb().
    hpe_cat02_inv = np.matmul(MATRIX_XYZ_TO_HPE, CAT_INVERSE_CAT02)
    path = os.path.join(DATA, "ciecam02_matrices.csv")
    with open(path, "w", newline="\n") as f:
        f.write("name,m00,m01,m02,m10,m11,m12,m20,m21,m22\n")
        for name, m in (("CAT02", CAT_CAT02), ("HPE_CAT02_INV", hpe_cat02_inv)):
            f.write(name + "," + ",".join(fmt(v) for v in np.ravel(m)) + "\n")
    print(f"wrote {path}")
    return 0


if __name__ == "__main__":
    sys.exit(main())