"""

# --- CIPIC Database Definitions ---

AZIMUTH_ANGLES = [

-90, -65, -55, -45, -40, -35, -30, -25, -20, -15, -10, -5, 0, 5, 10, 15,

20, 25, 30, 35, 40, 45, 55, 65, 90,

]

ELEVATION_ANGLES = -45 + 5.625 * np.arange(0, 50)

POINTS_NP = np.array(list(itertools.product(AZIMUTH_ANGLES, ELEVATION_ANGLES)))

AZ_DICT = dict(zip(AZIMUTH_ANGLES, np.arange(len(AZIMUTH_ANGLES))))

EL_DICT = dict(zip(ELEVATION_ANGLES, np.arange(len(ELEVATION_ANGLES))))

# --- Anthropometry Data Loading ---

try:

anthro_data = loadmat('CIPIC_hrtf_database/anthropometry/anthro.mat')

PARAMETERS = {

'info': ['Age:', 'Sex:', 'Weight:'],

'X': ['Head width:', 'Head height:', 'Head depth:', 'Pinna offset down:',

'Pinna offset back:', 'Neck width:', 'Neck height:', 'Neck depth:',

'Torso top width:', 'Torso top height:', 'Torso top depth:',

'Shoulder width:', 'Head offset forward:', 'Height:',

'Seated height:', 'Head circumference:', 'Shoulder circumference:'],

'D': ['Cavum concha height:', 'Cymba concha height:', 'Cavum concha width:',

'Fossa height:', 'Pinna height:', 'Pinna width:',

'Intertragal incisure width:', 'Cavum concha depth:'],

'theta': ['Pinna rotation angle:', 'Pinna flare angle:']

}

anthro_ids_raw = np.squeeze(anthro_data['id'])

ANTHRO_ID = [str(int(id_)) for id_ in anthro_ids_raw]

ID_TO_IDX = dict(zip([int(id_) for id_ in ANTHRO_ID], range(len(ANTHRO_ID))))

for key, value in anthro_data.items():

if key not in ['__header__', '__version__', '__globals__', 'id', 'sex']:

if key == 'age':

anthro_data[key][np.isnan(anthro_data[key])] = 0

anthro_data[key] = np.squeeze(value.astype('int')).astype('str')

anthro_data[key][anthro_data[key] == '0'] = '-'

else:

anthro_data[key] = np.around(np.squeeze(value), 1).astype('str')

anthro_data[key][anthro_data[key] == 'nan'] = '-'

except Exception as e:

print(f"CRITICAL WARNING: Could not parse 'id' from anthro.mat: {e}")

ID_TO_IDX = {}

def __init__(self, num_speakers, block_size: int, subject='018', device='cuda', dtype=torch.float32,

grid_size: Tuple[int, int] = (256, 256)):

"""

if not Head.ID_TO_IDX:

raise RuntimeError("Head class failed to initialize. Could not load subject IDs.")

idx = Head.ID_TO_IDX[int(subject)]

self.properties_cm = {}

for param, values in Head.PARAMETERS.items():

if param != 'info':

LR_const = 2 if param in ('D', 'theta') else 1

for j in range(LR_const * len(values)):

if Head.anthro_data[param][idx][j] != '-':

key = values[j % len(values)].strip().lower().replace(':', '').replace(' ', '_')

self.properties_cm[key] = float(Head.anthro_data[param][idx][j])

self.device = device

self.dtype = dtype

self.subject = subject

self.block_size = block_size

self.native_sample_rate = 44100 # Hard-coded, we never resample

if not (isinstance(grid_size, tuple) and len(grid_size) == 2):

raise ValueError("grid_size must be a tuple of two integers (e.g., (256, 256))")

self.grid_size = grid_size

# --- Cache Management ---

cache_dir = "hrir_cache"

# FFT cache is specific to subject, block size, and grid.

cache_file = f"subject_{self.subject}_b{self.block_size}_grid{self.grid_size[0]}x{self.grid_size[1]}_NATIVE_FFT.pt"

cache_path = os.path.join(cache_dir, cache_file)

N_AZ, N_EL = self.grid_size

print(f"\rFFT HRIR cache initialized ({N_AZ}x{N_EL}) ")

self.cache_az_step = 360.0 / (N_AZ - 1)

self.cache_el_step = 135.0 / (N_EL - 1)

# --- Head State (Global) ---

self.head_position = torch.tensor([0.0, 0.0, 0.0], device=device, dtype=dtype)

self.head_rotation = torch.eye(3, device=device, dtype=dtype)

self.head_rotation_dirty = True

# Add to __init__ or cache initialization:

self.cache_az_offset = 180.0

self.cache_el_offset = 45.0

self.cache_az_step_inv = 1.0 / self.cache_az_step

self.cache_el_step_inv = 1.0 / self.cache_el_step

self.hypot_buf = torch.empty(num_speakers*5*8, device=device, dtype=dtype)

self.arctan2_buf = torch.empty(num_speakers*5*8, device=device, dtype=dtype)

self.arctan2_buf2 = torch.empty(num_speakers*5*8, device=device, dtype=dtype)

self.az_idx = torch.empty(num_speakers*5*8*8, device=device, dtype=torch.long)

self.el_idx = torch.empty(num_speakers*5*8*8, device=device, dtype=torch.long)

self.azimuth_add_buf = torch.empty(num_speakers*8, dtype=dtype, device=device)

self.elevation_add_buf = torch.empty(num_speakers*8, dtype=dtype, device=device)

if os.path.exists(cache_path):

try:

print(f"Loading cached FFT Head data from {cache_path}...")

cached_data = torch.load(cache_path, map_location=self.device, weights_only=True)

# Validate cache

if (tuple(cached_data.get('grid_size', (0, 0))) != self.grid_size or

cached_data.get('block_size', 0) != self.block_size):

print("Cache parameter mismatch! Re-initializing...")

raise ValueError("Cache parameter mismatch")

self.hrir_length = cached_data['hrir_length']

self.fft_size = cached_data['fft_size']

self.overlap_size = cached_data['overlap_size']

self.hrir_fft_cache = cached_data['hrir_fft_cache']

self.T_inv = cached_data['T_inv']

self.cache_az_step = cached_data['cache_az_step']

self.cache_el_step = cached_data['cache_el_step']

self.hrtf_points = torch.from_numpy(self.POINTS_NP).to(device=self.device, dtype=self.dtype)

self.triangulation = Delaunay(self.POINTS_NP)

print(f"Cache loaded. FFT size: {self.fft_size}, Block: {self.block_size}")

except Exception as e:

print(f"Failed to load cache: {e}. Re-initializing...")

self._full_initialization(subject)

else:

print(f"Cache not found. Initializing FFT HRIR cache...")

self._full_initialization(subject)

@staticmethod

def _muffle(hrir_np: np.ndarray, a: float = 0.6) -> np.ndarray:

"""

# Numerator (b) coefficients: applied to input x[n]

b = [1.0 - a]

# Denominator (a) coefficients: applied to output y[n]

# The '1.0' corresponds to y[n], the '-a' corresponds to y[n-1]

# (scipy expects coefficients on the left side of the difference equation)

a_coeffs = [1.0, -a]

return lfilter(b, a_coeffs, hrir_np)

def _full_initialization(self, subject):

"""Full initialization with NATIVE 44.1kHz HRIRs."""

hrir = loadmat(f'CIPIC_hrtf_database/standard_hrir_database/subject_{subject}/hrir_final.mat')

hrir_l = np.array(hrir['hrir_l'])

hrir_r = np.array(hrir['hrir_r'])

hrir_combined = np.concatenate([hrir_l[..., None], hrir_r[..., None]], -1)

# Transpose to [Az, El, Channel, Samples]

hrir_combined = hrir_combined.transpose(0, 1, 3, 2)

self.hrir_data_native = torch.from_numpy(hrir_combined).to(device=self.device, dtype=self.dtype)

self.hrir_length = self.hrir_data_native.shape[-1] # Should be 200

self.hrtf_points = torch.from_numpy(self.POINTS_NP).to(device=self.device, dtype=self.dtype)

# --- FFT Size Calculation ---

self.fft_size = next_fast_len(self.hrir_length + self.block_size - 1)

self.overlap_size = self.hrir_length - 1

print(f"Native HRIR length: {self.hrir_length} samples.")

print(f"FFT size: {self.fft_size}, Block: {self.block_size}")

# Initialize FFT cache

self.N_AZ, self.N_EL = self.grid_size

self.hrir_fft_cache = torch.zeros((self.N_AZ, self.N_EL, 2, self.fft_size),

device=self.device, dtype=torch.complex64)

self._initialize_triangulation()

self._initialize_fft_cache_360()

# Save cache atomically

self._save_cache()

def _save_cache(self):

"""Save FFT cache atomically to prevent corruption."""

cache_dir = "hrir_cache"

cache_file = f"subject_{self.subject}_b{self.block_size}_grid{self.grid_size[0]}x{self.grid_size[1]}_NATIVE_FFT.pt"

cache_path = os.path.join(cache_dir, cache_file)

temp_cache_path = cache_path + ".tmp"

print(f"Saving initialized FFT cache to {cache_path}...")

try:

os.makedirs(cache_dir, exist_ok=True)

data_to_save = {

'hrir_length': self.hrir_length,

'fft_size': self.fft_size,

'overlap_size': self.overlap_size,

'hrir_fft_cache': self.hrir_fft_cache.cpu(),

'T_inv': self.T_inv.cpu(),

'cache_az_step': self.cache_az_step,

'cache_el_step': self.cache_el_step,

'grid_size': self.grid_size,

'block_size': self.block_size,

}

torch.save(data_to_save, temp_cache_path, _use_new_zipfile_serialization=True)

os.replace(temp_cache_path, cache_path)

print(f"FFT cache saved successfully.")

except Exception as e:

print(f"WARNING: Failed to save RFFT cache: {e}")

if os.path.exists(temp_cache_path):

try:

os.remove(temp_cache_path)

except Exception as e:

print(f"WARNING: Failed to remove failed RFFT cache: {e}")

pass

def _initialize_fft_cache_360(self):

"""Initialize 360-degree FFT cache."""

self.cache_az_range = np.linspace(-180, 180, self.N_AZ)

self.cache_el_range = np.linspace(-45, 90, self.N_EL)

print(f"Caching 360° FFT HRIRs ({self.N_AZ}x{self.N_EL})...", end="", flush=True)

# Pre-allocate padding buffers

padded_left = torch.zeros(self.fft_size, device=self.device, dtype=self.dtype)

padded_right = torch.zeros(self.fft_size, device=self.device, dtype=self.dtype)

for i, azimuth in enumerate(self.cache_az_range):

for j, elevation in enumerate(self.cache_el_range):

interp_az = azimuth

is_rear = False

if azimuth > 90.0:

interp_az = 180.0 - azimuth

is_rear = True

elif azimuth < -90.0:

interp_az = -180.0 - azimuth

is_rear = True

hrir_left_np, hrir_right_np = self._interpolate_hrir_triangulation(interp_az, elevation)

if is_rear:

hrir_left_np = self._muffle(hrir_left_np)

hrir_right_np = self._muffle(hrir_right_np)

padded_left.zero_()

padded_right.zero_()

padded_left[:self.hrir_length] = torch.from_numpy(hrir_left_np).to(self.device)

padded_right[:self.hrir_length] = torch.from_numpy(hrir_right_np).to(self.device)

# --- Perform FFT ---

self.hrir_fft_cache[i, j, 0] = torch.fft.fft(padded_left, n=self.fft_size)

self.hrir_fft_cache[i, j, 1] = torch.fft.fft(padded_right, n=self.fft_size)

if i % (max(1, self.N_AZ // 20)) == 0 or i == self.N_AZ - 1:

progress = int((i + 1) / self.N_AZ * 100)

print(f"\rCaching 360° FFT HRIRs... {progress}%", end="", flush=True)

def _initialize_triangulation(self):

"""Initialize Delaunay triangulation."""

self.triangulation = Delaunay(self.POINTS_NP)

T_inv_np = self._calculate_T_inv_np()

self.T_inv = torch.from_numpy(T_inv_np).to(device=self.device, dtype=self.dtype)

def _calculate_T_inv_np(self):

"""Calculate inverse transformation matrices for triangulation."""

A = self.POINTS_NP[self.triangulation.simplices][:, 0, :]

B = self.POINTS_NP[self.triangulation.simplices][:, 1, :]

C = self.POINTS_NP[self.triangulation.simplices][:, 2, :]

T = np.empty((2 * A.shape[0], A.shape[1]))

T[::2, :] = A - C

T[1::2, :] = B - C

T = T.reshape(-1, 2, 2)

return np.linalg.inv(T)

def _interpolate_hrir_triangulation(self, azimuth: float, elevation: float) -> Tuple[np.ndarray, np.ndarray]:

"""Interpolate HRIR using triangulation (used during caching)."""

azimuth = np.clip(azimuth, -90, 90)

elevation = np.clip(elevation, -45, 231.875)

position = np.array([azimuth, elevation])

triangle_idx = self.triangulation.find_simplex(position)

if triangle_idx == -1:

# Fallback: nearest neighbor

distances = np.sum((self.POINTS_NP - position) ** 2, axis=1)

nearest_idx = np.argmin(distances)

az_idx = self.AZ_DICT[self.AZIMUTH_ANGLES[nearest_idx % len(self.AZIMUTH_ANGLES)]]

el_idx = self.EL_DICT[self.ELEVATION_ANGLES[nearest_idx // len(self.AZIMUTH_ANGLES)]]

return (self.hrir_data_native[az_idx, el_idx, 0].cpu().numpy(),

self.hrir_data_native[az_idx, el_idx, 1].cpu().numpy())

vertices = self.POINTS_NP[self.triangulation.simplices[triangle_idx]]

T_inv_np = self.T_inv[triangle_idx].cpu().numpy()

g = np.dot(position - vertices[2], T_inv_np)

g_1, g_2 = g[0], g[1]

g_3 = 1 - g_1 - g_2

if g_1 >= 0 and g_2 >= 0 and g_3 >= 0:

# Barycentric interpolation

v_indices = [

(self.AZ_DICT[v[0]], self.EL_DICT[v[1]]) for v in vertices

]

hrir_l = (g_1 * self.hrir_data_native[v_indices[0][0], v_indices[0][1], 0] +

g_2 * self.hrir_data_native[v_indices[1][0], v_indices[1][1], 0] +

g_3 * self.hrir_data_native[v_indices[2][0], v_indices[2][1], 0])

hrir_r = (g_1 * self.hrir_data_native[v_indices[0][0], v_indices[0][1], 1] +

g_2 * self.hrir_data_native[v_indices[1][0], v_indices[1][1], 1] +

g_3 * self.hrir_data_native[v_indices[2][0], v_indices[2][1], 1])

return hrir_l.cpu().numpy(), hrir_r.cpu().numpy()

else:

# Fallback: nearest vertex in triangle

distances = [np.linalg.norm(position - v) for v in vertices]

nearest_vertex = vertices[np.argmin(distances)]

az_idx, el_idx = self.AZ_DICT[nearest_vertex[0]], self.EL_DICT[nearest_vertex[1]]

return (self.hrir_data_native[az_idx, el_idx, 0].cpu().numpy(),

self.hrir_data_native[az_idx, el_idx, 1].cpu().numpy())

def get_cached_fft_hrir(self, azimuth, elevation) -> torch.Tensor:

"""

# Track if input was scalar-like (scalar or single-element 1D)

is_scalar_like = azimuth.dim() == 0 or (azimuth.dim() == 1 and azimuth.shape[0] == 1)

size = azimuth.flatten().shape[0]

# Convert scalar to 1D for uniform processing

if azimuth.dim() == 0:

azimuth = azimuth[None]

elevation = elevation[None]

# Map to cache indices - no flattening needed!

az_idx = make_shape(azimuth, self.az_idx).copy_(torch.add(azimuth, self.cache_az_offset, out=make_shape(azimuth, self.azimuth_add_buf)).mul_(self.cache_az_step_inv))

el_idx = make_shape(azimuth, self.el_idx).copy_(torch.add(elevation, self.cache_el_offset, out=make_shape(elevation, self.elevation_add_buf)).mul_(self.cache_el_step_inv))

# Advanced indexing preserves the shape of index tensors

fft_stereo = self.hrir_fft_cache[az_idx, el_idx] # [..., 2, fft_size]

# Return scalar result if input was scalar-like

return fft_stereo[0] if is_scalar_like else fft_stereo

# --- Coordinate Transforms & Setters ---

def world_to_head_coordinates(self, world_pos: torch.Tensor) -> torch.Tensor:

"""Transform world to head-relative coordinates. Supports batched [..., 3]."""

# This is the exact transformation used by well tested transform_global_point_to_screen_space()

relative_pos = world_pos - self.head_position

return relative_pos @ self.head_rotation

def head_coords_to_cipic_angles(self, head_coords: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:

"""Convert head coords to CIPIC angles. Supports batched [..., 3]."""

x = head_coords[..., 0]

y = head_coords[..., 1]

z = head_coords[..., 2]

azimuth_deg = torch.rad2deg_(torch.arctan2(x, z, out=make_shape(z, self.arctan2_buf)))

elevation_deg = torch.rad2deg_(torch.arctan2(

y, torch.hypot(x, z, out=make_shape(x, self.hypot_buf)), out=make_shape(x, self.arctan2_buf2))

)

return azimuth_deg, elevation_deg

def set_head_rotation(self, rotation_matrix: torch.Tensor):

"""Set head rotation matrix."""

if rotation_matrix.shape != (3, 3):

raise ValueError("Rotation matrix must be 3x3")

# This unfortunately stalls the GPU to a level where it's better to just always make head_rotation_dirty

# Interleave all position and HRIR assignments to prevent stalls

# if not torch.allclose(self.head_rotation, rotation_matrix):

self.head_rotation = rotation_matrix

def __init__(self, device, dtype, num_speakers=1, speed_of_sound=343.0):

self.memory_lock = threading.RLock()

self.device = device

self.dtype = dtype

self.num_speakers = num_speakers

self.speed_of_sound = speed_of_sound

self.late_spatial_regions = 4 # Reduced from 8

self.late_reverb_mix = 0.9

# --- BATCHIFIED ---

self.speaker_positions = torch.zeros((self.num_speakers, 3), device=device, dtype=dtype)

self.fine_speaker_positions = torch.zeros((self.num_speakers, 3), device=device, dtype=dtype)

self.course_speaker_positions = torch.zeros((self.num_speakers, 3), device=device, dtype=dtype)

self.fine_speaker_positions[:, 1] = 0

self.course_speaker_positions[:, 1] = 0

# Set a default position

self.speaker_positions[:, 2] = 15.0

# --- GLOBAL ---

self.room_dimensions = torch.tensor([40.0, 20.0, 60.0], device=device, dtype=dtype)

self.half_dims = self.room_dimensions / 2.0

self.absorption_db = torch.tensor(-4.8, device=device, dtype=dtype)

self.params_dirty = True

self.small_val_tensor = torch.tensor(1e9, device=device, dtype=dtype)

self.rot60 = self.compute_rt60()

self.speaker_indices = torch.arange(self.num_speakers, device=self.device, dtype=torch.long)

self._position_change_intermediate_buffer = [torch.empty(

[self.num_speakers, 3],

dtype=self.dtype,

device=self.device

) for _ in range(4)]

self._position_sub_intermediate_buffer = [torch.empty(

[self.num_speakers, 3],

dtype=self.dtype,

device=self.device

) for _ in range(4)]

self._position_minimum_intermediate_buffer = [torch.empty(

(self.num_speakers, 3),

dtype=self.dtype,

device=self.device

) for _ in range(4)]

self._position_greater_intermediate_buffer = [torch.empty(

[self.num_speakers, 3],

dtype=torch.bool,

device=self.device

) for _ in range(4)]

self._position_sum_intermediate_buffer = [torch.empty(

self.num_speakers,

dtype=self.dtype,

device=self.device

) for _ in range(4)]

self._position_thresh_intermediate_buffer = [torch.empty(

self.num_speakers,

dtype=self.dtype,

device=self.device

) for _ in range(4)]

self._position_greater_output_intermediate_buffer = [torch.empty(

self.num_speakers,

dtype=torch.bool,

device=self.device

) for _ in range(4)]

# --- GLOBAL (Image source indices are constant) ---

self.er_indices_np = self._get_early_reflection_indices()

self.er_indices = torch.from_numpy(self.er_indices_np).to(device=self.device, dtype=self.dtype)

self.num_er_sources = len(self.er_indices) + 1 # +1 for direct path

self.dist_to_walls = torch.zeros([self.num_speakers, 6], dtype=self.dtype, device=device)

# --- GLOBAL (Optimized reflection index calculations) ---

W, H, D = self.room_dimensions

self.scaled_er_indices_w = self.er_indices[:, 0] * W

self.scaled_er_indices_h = self.er_indices[:, 1] * H

self.scaled_er_indices_d = self.er_indices[:, 2] * D

self.centered_er_indices_w = (self.er_indices[:, 0] % 2 * 2 - 1)

self.centered_er_indices_h = (self.er_indices[:, 1] % 2 * 2 - 1)

self.centered_er_indices_d = (self.er_indices[:, 2] % 2 - 1)

self.mirrors = torch.tensor([

[-1, 1, 1], [1, -1, 1], [1, 1, -1],

[-1, 1, 1], [1, -1, 1], [1, 1, -1]

], device=device, dtype=dtype)

self.offsets = torch.tensor([

[-W, 0, 0], [0, -H, 0], [0, 0, -D],

[W, 0, 0], [0, H, 0], [0, 0, D]

], device=device, dtype=dtype)

# --- GLOBAL (Cluster definitions are global) ---

self.enhanced_clusters = torch.tensor(

LocalizationBlueprint._define_optimized_clusters(), device=self.device, dtype=self.dtype

)

self.num_enhanced_clusters = len(self.enhanced_clusters)

# --- GLOBAL Late Reverb Setup ---

self.comb_delays_s = torch.tensor([0.0297, 0.0371, 0.0411, 0.0437], device=self.device, dtype=self.dtype)

self.allpass_delays_s = torch.tensor([0.0051, 0.0037], device=self.device, dtype=self.dtype)

self.late_weights = torch.tensor([0.3, 0.25, 0.25, 0.2], device=self.device, dtype=self.dtype) # Added

self.max_speakers = self.num_speakers + 50 + self.num_speakers // 2

self.buf_delta = torch.zeros((self.max_speakers, 3), device=device, dtype=dtype)

self.buf_dist_direct = torch.zeros(self.max_speakers, device=device, dtype=dtype)

self.buf_g_direct = torch.zeros(self.max_speakers, device=device, dtype=dtype)

self.buf_amp_direct = torch.zeros(self.max_speakers, device=device, dtype=dtype) # Output Accumulator

# 2. Proximity Effect Buffers

self.buf_walls_neg = torch.zeros((self.max_speakers, 3), device=device, dtype=dtype)

self.buf_walls_pos = torch.zeros((self.max_speakers, 3), device=device, dtype=dtype)

self.buf_min_wall = torch.zeros(self.max_speakers, device=device, dtype=dtype)

self.buf_indices = torch.zeros(self.max_speakers, device=device, dtype=torch.long) # for torch.min

self.buf_prox_gain = torch.zeros(self.max_speakers, device=device, dtype=dtype)

# 3. Reflection & Image Source Buffers

self.buf_img_src = torch.zeros((self.max_speakers, 6, 3), device=device, dtype=dtype)

self.buf_d_refs = torch.zeros((self.max_speakers, 6), device=device, dtype=dtype)

self.buf_amp_refs = torch.zeros((self.max_speakers, 6), device=device, dtype=dtype)

self.buf_rms_refs = torch.zeros(self.max_speakers, device=device, dtype=dtype)

# 4. Cluster & Scratch Buffers

self.buf_dist_sq = torch.zeros(self.max_speakers, device=device, dtype=dtype)

self.buf_reverb_scratch = torch.zeros(4, device=device, dtype=dtype) # Tiny buffer for reverb

self.k_decay = 0.67

self.asymptote = 0.985

self.dynamic_range = 1.8 - 0.985

self.proximity_dirty = torch.ones(self.num_speakers, dtype=torch.bool, device=self.device)

self.reflections_dirty = torch.ones(self.num_speakers, dtype=torch.bool, device=self.device)

self.hrirs_dirty = torch.ones(self.num_speakers, dtype=torch.bool, device=self.device)

self.absorption_gains_dirty = True

self.any_reflections_dirty = True

@staticmethod

def _define_optimized_clusters() -> List[Tuple[float, float]]:

"""Define optimized angular clusters (Global)."""

return [

(0, 0), # Front center

(-45, 0), # Front left

(45, 0), # Front right

(-90, 0), # Left

(90, 0), # Right

(0, 45), # Above

(0, -30), # Below

(-90, 30), # Left up (represents back/up region)

]

@staticmethod

def _get_early_reflection_indices():

"""Get image source indices for early reflections - (NumPy Init, Global)."""

return np.array([

[0, 1, 0], # top

[0, -1, 0], # bottom

[-1, 0, 0], # left

[1, 0, 0], # right

], dtype=np.float32)

def compute_rt60(self):

"""Compute RT60 (PyTorch) (Global)."""

# Remove this, that if check is a GPU Stall

# if self.absorption_db >= -0.0001:

# return self.small_val_tensor

volume = torch.prod(self.room_dimensions)

surface_area = 2 * torch.sum(self.room_dimensions * torch.roll(self.room_dimensions, 1))

alpha = 1.0 - 10 ** (self.absorption_db / 10.0)

rt60 = 0.161 * volume / (surface_area * alpha)

b0_cupy: cupy.ndarray, b1_cupy: cupy.ndarray, a1_cupy: cupy.ndarray,

x_cp: cupy.ndarray,

z_cupy: cupy.ndarray,

y_out_cupy: cupy.ndarray,

stream: cupy.cuda.ExternalStream

):

"""

B, N_samples = x_cp.shape

# --- Convert DYNAMIC input (signal) ---

first_order_iir_batched_N_filters_kernel(

(B,), (1,), # Grid = (num_speakers,), Block = (1,)

(x_cp, y_out_cupy, z_cupy, z_cupy,

b0_cupy, b1_cupy, a1_cupy, N_samples),

stream=stream

)

# No return needed; y_out_cupy and z_cupy are modified in-place (and share

x_cp: cupy.ndarray, # [num_speakers, N_samples] (Dynamic)

y_cupy: cupy.ndarray, # [num_speakers, N_filters, N_samples] (Persistent)

packed_buffer_cupy: cupy.ndarray, # [num_speakers, sum(N_delays)] (Persistent)

buffer_offsets_cupy: cupy.ndarray, # [N_filters] (Persistent)

buffer_idx_cupy: cupy.ndarray, # [num_speakers, N_filters] (Persistent)

feedback_cupy: cupy.ndarray, # [num_speakers, N_filters] (Persistent, smoothed)

N_delays_cupy: cupy.ndarray, # [N_filters] (Persistent)

total_N_delays: int, # sum(N_delays)

stream: cupy.cuda.ExternalStream

) -> None:

"""

num_speakers, N_samples = x_cp.shape

N_filters = N_delays_cupy.shape[0]

# --- Convert DYNAMIC input (signal) ---

# Grid: (num_speakers, N_filters), Block: (1,)

multi_comb_filter_N_speakers_kernel(

(num_speakers, N_filters), (1,),

(x_cp, y_cupy,

packed_buffer_cupy, buffer_offsets_cupy, buffer_idx_cupy,

feedback_cupy, N_delays_cupy,

N_samples, N_filters, total_N_delays),

stream=stream

)

x_cupy: cupy.ndarray, # [num_speakers, N_filters, N_samples] (Persistent)

y_cupy: cupy.ndarray, # [num_speakers, N_filters, N_samples] (Persistent)

buffer_cupy: cupy.ndarray, # [num_speakers, N_filters, N_delays] (Persistent)

buffer_idx_cupy: cupy.ndarray, # [num_speakers, N_filters] (Persistent)

gain_cupy: cupy.ndarray, # [num_speakers, N_filters] (Persistent)

N_delays: int,

stream: cupy.cuda.ExternalStream

) -> None:

"""

num_speakers, N_filters, N_samples = x_cupy.shape

# Grid: (num_speakers, N_filters), Block: (1,)

allpass_filter_N_speakers_kernel(

(num_speakers, N_filters), (1,),

(x_cupy, y_cupy,

buffer_cupy, buffer_idx_cupy,

gain_cupy, N_delays,

N_samples, N_filters),

stream=stream

)

def __init__(

self, num_speakers=1, sample_rate=44100, block_size=256, subject='018', device='cuda', dtype=torch.float32,

stream: torch.cuda.Stream = None, present='fast', speed_of_sound=343.0

):

super().__init__(device=device, dtype=dtype, num_speakers=num_speakers, speed_of_sound=speed_of_sound)

self.sample_rate = sample_rate

self.present = present

self.block_size = block_size

self.user_block_size = block_size

self.device = device

self.dtype = dtype

self.frame_counter = 0

self.crossfade_len = block_size

self.crossfade_linspace = torch.linspace(0, math.pi, self.crossfade_len, device=self.device, dtype=self.dtype)

self.reset_mask = torch.zeros(self.num_speakers, dtype=torch.bool, device=self.device)

# --- 3. Early Reflections (Energy Sum) ---

# --- GLOBAL (1 Head) ---

self.head = Head(

num_speakers=num_speakers, block_size=block_size, subject=subject, device=self.device, dtype=self.dtype

)

# We do this for efficient, head delta computation for sparse, speaker parameter updates

self.last_head_positions = torch.zeros((self.num_speakers, 3), device=device, dtype=dtype)

# Use self.num_speakers from base class

print(f"Initializing SpaceSimulationEngine for {self.num_speakers} speakers.")

# --- Physical Parameters (Global) ---

self.mic_radius = 0.05

self.speaker_radius = 0.08

self.min_distance = torch.tensor(self.mic_radius + self.speaker_radius, device=self.device, dtype=self.dtype)

self.num_precise_reflections = 4

# --- BATCHIFIED Early Reflections Setup ---

self.er_delay_line_size = 8192

is_power_of_two = (self.er_delay_line_size & (self.er_delay_line_size - 1) == 0) and self.er_delay_line_size > 0

assert is_power_of_two, "er_delay_line_size MUST be a power of 2 for fast bitwise masking."

self.er_size_mask = self.er_delay_line_size - 1

self.er_delay_line = torch.zeros((self.num_speakers, self.er_delay_line_size), device=self.device,

dtype=self.dtype)

self.er_write_head = 0 # Global write head, as all speakers are processed in sync

# --- BATCHIFIED Per-reflection parameters ---

self.reflection_delays = torch.zeros((self.num_speakers, self.num_er_sources), device=self.device,

dtype=torch.long)

self.reflection_gains = torch.zeros((self.num_speakers, self.num_er_sources), device=self.device,

dtype=self.dtype)

self.reflection_positions = torch.zeros((self.num_speakers, self.num_er_sources, 3), device=self.device,

dtype=self.dtype)

# --- NEW: State for parameter interpolation ---

self.prev_reflection_delays = torch.zeros_like(self.reflection_delays)

self.prev_reflection_gains = torch.zeros_like(self.reflection_gains)