Compute RIR up to a K (earliest or strongest) reflection (or max time)

[Chutlhu]

This function prunes the later source's images so that the computed RIR features contain only the first K strongest (or earliest) reflections or all the reflections up to Tmax. It might be useful for some echo people.

It would be nice to have it implemented in Shoebox (or have a callback in the ISM computation)

def keep_k_reflection_in_rir(pra_room, R=None, Tmax=None, sort_by='strongest', source_index=0, ref_pos=None):
    ## compute the ISM
    pra_room.image_source_model()
    ## remove images here
    s = source_index
    if ref_pos is None:
        ref_pos = np.mean(pra_room.mic_array.R, axis=1)
    # get the ISM fields
    visibility = pra_room.visibility[s]
    images = pra_room.sources[s].images
    orders = pra_room.sources[s].orders
    orders_xyz = pra_room.sources[s].orders_xyz
    walls = pra_room.sources[s].walls
    dampings = pra_room.sources[s].damping
    generators = pra_room.sources[s].generators
    # sort images by distance to the reference microphone (center?)
        
    distances = np.linalg.norm(images - ref_pos[:,None], axis=0)
    # print(distances)
    alpha = dampings.squeeze() / distances 
    # sort by decresasing alpha
    if sort_by == 'strongest':
        idx = np.argsort(alpha)[::-1]
    elif sort_by == 'earliest':
        idx = np.argsort(distances)
    # crop rirs based on the maximum time
    max_dist = Tmax * pra_room.c
    idx = idx[distances[idx] < max_dist]
    # consider only the first R paths
    idx = idx[:R]
    # update the ISM fields
    pra_room.visibility[0] = visibility[:,idx]
    pra_room.sources[0].images = images[:,idx]
    pra_room.sources[0].orders = orders[idx]
    pra_room.sources[0].orders_xyz = orders_xyz[:,idx]
    pra_room.sources[0].walls = walls[idx]
    pra_room.sources[0].damping = dampings[:,idx]
    pra_room.sources[0].generators = generators[idx]
    # compute the RIR with the updated images
    pra_room.compute_rir()
    return pra_room

Usage

room_dim = [5, 6]
s1 = np.random.randn(16000)
src_pos = [3, 2]
mic_pos = pra.circular_2D_array([2,3], 6, 0.0, 0.15)

e_absorption, max_order = pra.inverse_sabine(RT60, room_dim)
eroom = pra.ShoeBox(room_dim, fs=fs, materials=pra.Material(e_absorption), max_order=max_order, sigma2_awgn=0.)
eroom.add_source(src_pos, signal=s1)
eroom.add_microphone_array(pra.MicrophoneArray(mic_pos, room.fs))
eroom.image_source_model() # added here for clarity, but the function compute it also inside
# !!! important: it must be used before simulate(), it computes the ISM internally
eroom = compute_k_strongest_reflection_rir(eroom, Tmax=0.050)
eroom.simulate()
eroom.plot_rir()