添加双耳渲染功能

src/rosella_model.py

@@ -0,0 +1,517 @@
+"""Parser for ``.personalized_headphone`` and raw ``rp`` models."""
+
+from __future__ import annotations
+
+import hashlib
+import json
+import math
+from numbers import Real
+import struct
+from dataclasses import dataclass
+from pathlib import Path
+
+import numpy as np
+
+
+Q15 = np.float32(1.0 / 32768.0)
+
+
+def _f32(value) -> np.float32:
+    return np.float32(value)
+
+
+def _q15(value: int) -> np.float32:
+    return _f32(_f32(value) * Q15)
+
+
+def _q15_exp(value: int, exponent: int) -> np.float32:
+    return _f32(_q15(value) * _f32(np.ldexp(1.0, exponent)))
+
+
+@dataclass(frozen=True)
+class DistanceProfile:
+    bounds: np.ndarray
+    distance_scale_m: np.float32
+    inverse_distance_per_m: np.float32
+    axis_scales_internal: np.ndarray
+    minimum_normalized_radius: np.float32
+
+    @property
+    def floats(self) -> np.ndarray:
+        return np.concatenate((
+            self.bounds,
+            np.asarray([self.distance_scale_m,
+                        self.inverse_distance_per_m], dtype=np.float32),
+            self.axis_scales_internal,
+            np.asarray([self.minimum_normalized_radius], dtype=np.float32),
+        ))
+
+
+@dataclass(frozen=True)
+class RosellaModel:
+    source_path: str
+    coefficients: np.ndarray
+    coefficient_sha256: str
+    coefficient_version: str | None
+    room_model: str | None
+    table_a_dimension: int
+    table_a_option: int
+    table_a_extra: int
+    table_a_header_field: int
+    table_a_header_25: int
+    table_a_control: int
+    table_a_option_ids: np.ndarray
+    table_a_option_values: np.ndarray
+    table_a_scalar: np.float32
+    table_a_filter_16x64_padded: np.ndarray
+    table_a_four_integers: np.ndarray
+    table_a_integer: int
+    table_a_filter_8x64_padded: np.ndarray
+    table_a_vector16: np.ndarray
+    table_a_filter_4x64_padded: np.ndarray
+    table_a_extra_indices: np.ndarray
+    table_a_extra_fields_padded: np.ndarray
+    table_a_extra_vectors: np.ndarray
+    sample_rate: int
+    matrix_exponent: int
+    field_exponent: int
+    matrix_left: np.ndarray
+    matrix_right: np.ndarray
+    vector_left: np.ndarray
+    vector_right: np.ndarray
+    field_left_padded: np.ndarray
+    field_right_padded: np.ndarray
+    field_left_odd_serialized_zero: bool
+    hybrid_flags: np.ndarray
+    hybrid_values: np.ndarray
+    model_scalars: np.ndarray
+    header_float_scalars: np.ndarray
+    header_integer_fields: np.ndarray
+    profiles: tuple[DistanceProfile, ...]
+    profile_tail: np.ndarray
+    post_fields: np.ndarray
+    table_a_main_serialized: np.ndarray
+
+
+def _lane(data: bytes, index: int) -> int:
+    if (index + 1) * 4 > len(data):
+        raise ValueError(f"Rosella rp truncated before int32 lane {index}")
+    return struct.unpack_from("<I", data, index * 4)[0]
+
+
+def inspect_rp(data: bytes) -> dict:
+    """Return the active-lane layout and checksum status for one raw rp image."""
+    if len(data) < 20 or len(data) % 4:
+        raise ValueError("Rosella rp must contain whole little-endian int32 lanes")
+    if _lane(data, 0) != 0x7072:
+        raise ValueError(f"bad Rosella rp magic: 0x{_lane(data, 0):08X}")
+
+    low16 = lambda value: value & 0xFFFF
+    checksum = low16(_lane(data, 1))
+    table_a_present = low16(_lane(data, 2))
+    table_b_present = low16(_lane(data, 3))
+    table_c_present = low16(_lane(data, 4))
+    index = 5
+    if table_a_present:
+        table_a_dimension = low16(_lane(data, index))
+        table_a_option = low16(_lane(data, index + 1))
+        table_a_extra = low16(_lane(data, index + 2))
+        index += 5
+    else:
+        table_a_dimension, table_a_option, table_a_extra = 77, 0, 0
+    if table_b_present:
+        if not table_a_present:
+            raise ValueError("Rosella rp table B cannot be present without table A")
+        table_b_dimension = low16(_lane(data, index))
+        table_b_extra = low16(_lane(data, index + 1))
+        table_b_groups = low16(_lane(data, index + 2))
+        index += 3
+    else:
+        table_b_dimension = table_b_extra = table_b_groups = 0
+    if table_c_present:
+        table_c_dimension = low16(_lane(data, index))
+        index += 1
+    else:
+        table_c_dimension = 0
+
+    payload_words = (
+        (index - 2)
+        + table_b_present * (
+            table_b_dimension + 380 * table_b_groups + table_b_extra + 79)
+        + table_a_present * (
+            171 * table_a_extra + 79 + 2 * (table_a_option + 14 * table_a_dimension))
+        + 11
+        + table_c_present * (314 * table_c_dimension + 1)
+    )
+    total_lanes = 2 + payload_words
+    if len(data) < total_lanes * 4:
+        raise ValueError(
+            f"Rosella rp truncated: need {total_lanes * 4} bytes, have {len(data)}")
+    computed = 0xA569
+    for lane_index in range(2, total_lanes):
+        computed ^= low16(_lane(data, lane_index))
+    computed &= 0xFFFF
+    return {
+        "stored_checksum": checksum,
+        "computed_checksum": computed,
+        "checksum_valid": computed == checksum,
+        "table_a_present": table_a_present,
+        "table_b_present": table_b_present,
+        "table_c_present": table_c_present,
+        "table_a_dimension": table_a_dimension,
+        "table_a_option": table_a_option,
+        "table_a_extra": table_a_extra,
+        "table_b_dimension": table_b_dimension,
+        "table_b_extra": table_b_extra,
+        "table_b_groups": table_b_groups,
+        "table_c_dimension": table_c_dimension,
+        "active_int32_lanes": total_lanes,
+    }
+
+
+def _json_int32(values) -> np.ndarray:
+    if not isinstance(values, list):
+        raise ValueError("rosella_coefficients must be a JSON array")
+    result = np.empty(len(values), dtype=np.int32)
+    for index, value in enumerate(values):
+        if isinstance(value, bool) or not isinstance(value, Real):
+            raise ValueError(f"rosella_coefficients[{index}] is not a number")
+        numeric = float(value)
+        if not math.isfinite(numeric) or numeric != math.trunc(numeric):
+            raise ValueError(
+                f"rosella_coefficients[{index}] is not an exact integer: {value!r}")
+        integer = int(value)
+        if integer < -(1 << 31) or integer > (1 << 31) - 1:
+            raise ValueError(
+                f"rosella_coefficients[{index}] is outside signed int32: {integer}")
+        result[index] = integer
+    return result
+
+
+def _load_coefficients(path: Path) -> tuple[np.ndarray, str | None, str | None]:
+    if not path.is_file():
+        raise FileNotFoundError(path)
+    source = path.read_bytes()
+    stripped = source.lstrip()
+    if stripped.startswith(b"{"):
+        try:
+            document = json.loads(source.decode("utf-8"))
+            virtualizer = document["personalized_hrtf"]["virtualizer_parameters"]
+            coefficients = _json_int32(virtualizer["rosella_coefficients"])
+        except (UnicodeDecodeError, json.JSONDecodeError, KeyError, TypeError) as exc:
+            raise ValueError(f"invalid personalized_headphone JSON: {exc}") from exc
+        version = virtualizer.get("rosella_coefficients_version")
+        room = virtualizer.get("room_model")
+    else:
+        if len(source) % 4:
+            raise ValueError("raw rp payload must contain complete int32 lanes")
+        coefficients = np.frombuffer(source, dtype="<i4").copy()
+        version = room = None
+    return coefficients, version, room
+
+
+def _unpack_field(serialized: np.ndarray, directions: int,
+                  exponent: int) -> np.ndarray:
+    expected = 154 * directions
+    if serialized.size != expected:
+        raise ValueError(f"expected {expected} field lanes, got {serialized.size}")
+    padded = np.zeros(160 * directions, dtype=np.float32)
+    stride8 = 8 * directions
+    stride2 = 2 * directions
+    for source_index, value in enumerate(serialized):
+        group4 = (source_index % stride8) // stride2
+        destination = ((group4 & 3) + 4 * (
+            source_index % stride2 +
+            2 * directions * (source_index // stride8 + (group4 >> 2))))
+        padded[destination] = _q15_exp(int(value), exponent)
+    return padded
+
+
+def _unpack_table_a_grid(serialized: np.ndarray, dimension: int,
+                         serialized_rows: int, padded_rows: int,
+                         lane_group: int) -> np.ndarray:
+    if serialized.size != serialized_rows * dimension:
+        raise ValueError("unexpected table-A grid size")
+    padded = np.zeros(padded_rows * dimension, dtype=np.float32)
+    group_width = lane_group * 4
+    for source_index, value in enumerate(serialized):
+        remainder = source_index % group_width
+        destination = ((remainder // lane_group) + 4 * (
+            remainder % lane_group +
+            group_width // 4 * (source_index // group_width)))
+        padded[destination] = _q15(int(value))
+    return padded
+
+
+def _unpack_table_a_extra(serialized: np.ndarray) -> np.ndarray:
+    if serialized.size != 154:
+        raise ValueError("table-A extra field must contain 154 serialized values")
+    padded = np.zeros(160, dtype=np.float32)
+    for source_index, value in enumerate(serialized):
+        remainder = source_index & 7
+        destination = ((remainder >> 1) + 4 * (
+            (source_index & 1) + 2 * (source_index >> 3)))
+        padded[destination] = _q15(int(value))
+    return padded
+
+
+def _parse_profile(values: np.ndarray, position: int) -> tuple[DistanceProfile, int]:
+    bounds = np.asarray([_q15(int(value)) for value in values[position:position + 6]],
+                        dtype=np.float32)
+    position += 6
+    distance = _q15_exp(int(values[position]), int(values[position + 1]))
+    position += 2
+    remaining = np.asarray(
+        [_q15(int(value)) for value in values[position:position + 5]],
+        dtype=np.float32)
+    position += 5
+    return DistanceProfile(
+        bounds=bounds,
+        distance_scale_m=distance,
+        inverse_distance_per_m=remaining[0],
+        axis_scales_internal=remaining[1:4],
+        minimum_normalized_radius=remaining[4],
+    ), position
+
+
+def load_personalized_headphone(path: str | Path) -> RosellaModel:
+    path = Path(path).resolve()
+    coefficients, version, room = _load_coefficients(path)
+    raw = coefficients.astype("<i4", copy=False).tobytes()
+    header = inspect_rp(raw)
+    if not header["checksum_valid"] or header["active_int32_lanes"] != coefficients.size:
+        raise ValueError("invalid or non-active Rosella rp coefficient sequence")
+    if not header["table_a_present"] or not header["table_b_present"] or header["table_c_present"]:
+        raise NotImplementedError("current local renderer requires table A+B and no table C")
+    if header["table_a_dimension"] != 64 or header["table_a_option"] != 3:
+        raise NotImplementedError("current local renderer requires the observed 64-channel HQMF layout")
+    if header["table_b_dimension"] != 20 or header["table_b_groups"] != 36:
+        raise NotImplementedError("current local renderer requires 20 hybrid groups and 36 direction terms")
+
+    values = coefficients
+    extra = header["table_a_extra"]
+    table_a_main_start = 13
+    position = table_a_main_start
+    table_a_control = int(values[position]) & 0xFFFF
+    field_exponent = int(values[position])
+    position += 1
+    option_count = header["table_a_option"]
+    option_ids = (values[position:position + option_count].astype(np.int64) &
+                  0xFFFF).astype(np.int32)
+    position += option_count
+    option_values = np.asarray(
+        [_q15(int(value)) for value in values[position:position + option_count]],
+        dtype=np.float32)
+    position += option_count
+    table_a_scalar = _q15(int(values[position]))
+    position += 1
+    dimension = header["table_a_dimension"]
+    table_a_filter_16x64 = _unpack_table_a_grid(
+        values[position:position + 16 * dimension], dimension, 16, 20, 16)
+    position += 16 * dimension
+    table_a_four_integers = (values[position:position + 4].astype(np.int64) &
+                             0xFFFF).astype(np.int32)
+    position += 4
+    table_a_integer = int(values[position]) & 0xFFFF
+    position += 1
+    table_a_filter_8x64 = _unpack_table_a_grid(
+        values[position:position + 8 * dimension], dimension, 8, 10, 8)
+    position += 8 * dimension
+    table_a_vector16 = np.asarray(
+        [_q15(int(value)) for value in values[position:position + 16]],
+        dtype=np.float32)
+    position += 16
+    table_a_filter_4x64 = _unpack_table_a_grid(
+        values[position:position + 4 * dimension], dimension, 4, 5, 4)
+    position += 4 * dimension
+    extra_indices = (values[position:position + extra].astype(np.int64) &
+                     0xFFFF).astype(np.int32)
+    position += extra
+    extra_fields = np.empty((extra, 160), dtype=np.float32)
+    for index in range(extra):
+        extra_fields[index] = _unpack_table_a_extra(values[position:position + 154])
+        position += 154
+    extra_vectors = np.empty((extra, 16), dtype=np.float32)
+    for index in range(extra):
+        extra_vectors[index] = np.asarray(
+            [_q15(int(value)) for value in values[position:position + 16]],
+            dtype=np.float32)
+        position += 16
+    table_b_start = position
+    expected_table_b_start = table_a_main_start + 1821 + 171 * extra
+    if table_b_start != expected_table_b_start:
+        raise AssertionError(
+            f"table-A parser ended at {table_b_start}, expected {expected_table_b_start}")
+    table_a_main = values[table_a_main_start:table_b_start].copy()
+
+    sample_rate = 2 * (int(values[position]) & 0xFFFF)
+    position += 1
+    matrix_exponent = int(values[position])
+    position += 1
+    matrix_count = 36 * 36
+    scale_matrix = lambda block: np.asarray(
+        [_q15_exp(int(value), matrix_exponent) for value in block],
+        dtype=np.float32).reshape(36, 36)
+    matrix_left = scale_matrix(values[position:position + matrix_count])
+    position += matrix_count
+    matrix_right = scale_matrix(values[position:position + matrix_count])
+    position += matrix_count
+    vector_left = np.asarray(
+        [_q15_exp(int(value), matrix_exponent)
+         for value in values[position:position + 36]], dtype=np.float32)
+    position += 36
+    vector_right = np.asarray(
+        [_q15_exp(int(value), matrix_exponent)
+         for value in values[position:position + 36]], dtype=np.float32)
+    position += 36
+
+    serialized_count = 154 * 36
+    field_left_serialized = values[position:position + serialized_count]
+    field_left = _unpack_field(field_left_serialized, 36, field_exponent)
+    field_left_odd_zero = not np.any(
+        np.abs(np.asarray([_q15_exp(int(value), field_exponent)
+                           for value in field_left_serialized[1::2]],
+                          dtype=np.float32)) > np.float32(1e-6))
+    position += serialized_count
+    field_right = _unpack_field(
+        values[position:position + serialized_count], 36, field_exponent)
+    position += serialized_count
+
+    hybrid_flags = (values[position:position + 20].astype(np.int64) & 0xFFFF).astype(np.int32)
+    position += 20
+    active_hybrid_values = int(np.count_nonzero(hybrid_flags == 1))
+    if active_hybrid_values != header["table_b_extra"]:
+        raise ValueError(
+            f"hybrid value count {active_hybrid_values} != header {header['table_b_extra']}")
+    hybrid_values = np.asarray(
+        [_q15(int(value)) for value in values[position:position + active_hybrid_values]],
+        dtype=np.float32)
+    position += active_hybrid_values
+    model_scalars = np.asarray(
+        [_q15(int(value)) for value in values[position:position + 5]],
+        dtype=np.float32)
+    position += 5
+
+    expected_table_a_tail = table_b_start + (
+        header["table_b_dimension"] +
+        380 * header["table_b_groups"] +
+        header["table_b_extra"] + 79)
+    if position != expected_table_a_tail:
+        raise AssertionError(f"table-B parser ended at {position}, expected {expected_table_a_tail}")
+
+    header_float_scalars = np.asarray([
+        _q15(int(values[position])),
+        _f32(_q15(int(values[position + 1])) * _f32(16.0)),
+    ], dtype=np.float32)
+    header_integer_fields = np.asarray([
+        int(values[position + 2]),
+        int(values[position + 3]) & 0xFFFF,
+    ], dtype=np.int32)
+    position += 4
+    profiles = []
+    for _ in range(4):
+        profile, position = _parse_profile(values, position)
+        profiles.append(profile)
+    profile_tail = np.asarray(
+        [_q15(int(value)) for value in values[position:position + 8]],
+        dtype=np.float32)
+    position += 8
+    post_fields = values[position:position + 3].astype(np.int32, copy=True)
+    position += 3
+    if position != values.size:
+        raise AssertionError(f"unparsed coefficient lanes: {values.size - position}")
+
+    return RosellaModel(
+        source_path=str(path),
+        coefficients=coefficients,
+        coefficient_sha256=hashlib.sha256(raw).hexdigest(),
+        coefficient_version=version,
+        room_model=room,
+        table_a_dimension=header["table_a_dimension"],
+        table_a_option=header["table_a_option"],
+        table_a_extra=extra,
+        table_a_header_field=int(values[8]) & 0xFFFF,
+        table_a_header_25=int(values[9]) & 0xFFFF,
+        table_a_control=table_a_control,
+        table_a_option_ids=option_ids,
+        table_a_option_values=option_values,
+        table_a_scalar=table_a_scalar,
+        table_a_filter_16x64_padded=table_a_filter_16x64,
+        table_a_four_integers=table_a_four_integers,
+        table_a_integer=table_a_integer,
+        table_a_filter_8x64_padded=table_a_filter_8x64,
+        table_a_vector16=table_a_vector16,
+        table_a_filter_4x64_padded=table_a_filter_4x64,
+        table_a_extra_indices=extra_indices,
+        table_a_extra_fields_padded=extra_fields,
+        table_a_extra_vectors=extra_vectors,
+        sample_rate=sample_rate,
+        matrix_exponent=matrix_exponent,
+        field_exponent=field_exponent,
+        matrix_left=matrix_left,
+        matrix_right=matrix_right,
+        vector_left=vector_left,
+        vector_right=vector_right,
+        field_left_padded=field_left,
+        field_right_padded=field_right,
+        field_left_odd_serialized_zero=field_left_odd_zero,
+        hybrid_flags=hybrid_flags,
+        hybrid_values=hybrid_values,
+        model_scalars=model_scalars,
+        header_float_scalars=header_float_scalars,
+        header_integer_fields=header_integer_fields,
+        profiles=tuple(profiles),
+        profile_tail=profile_tail,
+        post_fields=post_fields,
+        table_a_main_serialized=table_a_main,
+    )
+
+
+def direction_basis(x: float, y: float, z: float,
+                    dtype=np.float64) -> np.ndarray:
+    """Return the observed 36-term Rosella direction basis."""
+    f = dtype
+    x, y, z = f(x), f(y), f(z)
+    out = np.empty(36, dtype=dtype)
+    yz = f(y * z)
+    x2 = f(x * x)
+    y2 = f(y * y)
+    x2m02 = f(x2 - f(0.2))
+    xy = f(x * y)
+    out[0:4] = (f(1.0), x, y, z)
+    out[4] = f(x2 - f(1.0 / 3.0))
+    out[5] = xy
+    out[6] = f(x * z)
+    out[7] = f(y2 - f(1.0 / 3.0))
+    out[8] = yz
+    out[9] = f(f(x2 - f(0.6)) * x)
+    out[10] = f(x2m02 * y)
+    out[11] = f(x2m02 * z)
+    out[12] = f(f(y2 - f(0.2)) * x)
+    out[13] = f(yz * x)
+    out[14] = f(f(y2 - f(0.6)) * y)
+    out[15] = f(f(y2 - f(0.2)) * z)
+    out[16] = f(f(x2 * x2) - f(0.2))
+    out[17] = f(xy * x2)
+    out[18] = f(f(x * z) * x2)
+    out[19] = f(f(y2 * x2) - f(1.0 / 15.0))
+    out[20] = f(yz * x2)
+    out[21] = f(x * y2 * y)
+    out[22] = f(x * y2 * z)
+    out[23] = f(f(y2 * y2) - f(0.2))
+    x4 = f(x2 * x2)
+    x2y2 = f(y2 * x2)
+    y4 = f(y2 * y2)
+    out[24] = f(yz * y2)
+    out[25] = f(f(x4 - f(3.0 / 7.0)) * x)
+    out[26] = f(f(x4 - f(3.0 / 35.0)) * y)
+    out[27] = f(f(x4 - f(3.0 / 35.0)) * z)
+    out[28] = f(f(x2y2 - f(3.0 / 35.0)) * x)
+    out[29] = f(f(x2 * z) * xy)
+    out[30] = f(f(x2y2 - f(3.0 / 35.0)) * y)
+    out[31] = f(f(x2y2 - f(1.0 / 35.0)) * z)
+    out[32] = f(f(y4 - f(3.0 / 35.0)) * x)
+    out[33] = f(f(y2 * z) * xy)
+    out[34] = f(f(y4 - f(3.0 / 7.0)) * y)
+    out[35] = f(f(y4 - f(3.0 / 35.0)) * z)
+    return out