Source code for mesh2hrtf.Output2HRTF.inspect_sofa_files

import os
import glob
import warnings
import numpy as np
import matplotlib.pyplot as plt
from netCDF4 import Dataset
import pyfar as pf


[docs] def inspect_sofa_files(path, pattern=None, plot=None, plane="horizontal", atol=0.1, savedir=None, dB_time=False, dB_freq=True, freq_scale='log'): """ Inspect SOFA files through plots. Generate and save plots from the sofa-files that are contained in the Mesh2HRTF export subfolder `Output2HRTF` and contain the HRIRs and HRTFs at the evaluation grid. Note: If the sofa-files do not exist in your Project folder you have to run ``NumCalc`` and :py:func:`~output2hrtf`. See the online documentation for more information. Parameters ---------- path : str Path to a folder containing Mesh2HRTF projects. SOFA files are searched in `path/Output2HRTF` if it exist and directly in `path` otherwise. The name may contain an asterisk to process data in multiple folders. E.g., if `path` is ``"some/path/HRIRs_*"`` files in all folder starting with "HRIRs" would be scanned for SOFA files. pattern : str Plot only files that contain `pattern` in their filename. The default ``None`` plots all SOFA files. plot : str, optional ``"2D"`` generate line plots of four sources on the horizontal plane (front, back, left, right). The closest sources to the ideal positions are used for plotting. ``"3D"`` generate color coded plots of all sources on the horizontal, median, or frontal plane. See also parameter `plane` and `atol` below. The default ``None`` generate both plots. plane : str, optional Select the plane for which is shown in the 3D plot. Can be ``"horizontal"`` (default), ``"median"``, or ``"frontal"`` atol : float, optional Sources on the `plane` are searched within a range +/- `atol` degree. The default is ``0.1``. savedir : str Directory for saving the merged SOFA files. The default ``None`` saves the files to the directory given by `path`. dB_time : bool, optional Plot the logarithmic time data. The default is ``'False'``. dB_freq : bool, optional Plot the logarithmic magnitude data. The default is ``'True'``. freq_scale : str, optional ``'log'`` to plot on a logarithmic frequency axis and ``'linear'`` to plot on a linear frequency axis. The default is ``'log'``. """ # check input if plot is None: plot = ["2D", "3D"] elif plot not in ["2D", "3D"]: raise ValueError( f"plot is {plot} but must be 2D, 3D, or all") if plane not in ["horizontal", "median", "frontal"]: raise ValueError( f"plane is {plane} but must be horizontal, median, or frontal") # check which data to merge if pattern is None: pattern = "*.sofa" elif not pattern.endswith("sofa"): pattern = f"{pattern}*.sofa" # get all directories containing SOFA files folders = glob.glob(path) # loop directories for folder in folders: # check if Output2HRTF folder exists if os.path.isdir(os.path.join(folder, "Output2HRTF")): folder = os.path.join(folder, "Output2HRTF") # find matching SOFA files files = glob.glob(os.path.join(folder, pattern)) if not files: raise ValueError((f"Did not find any SOFA files in {folder} " f"that are matching {pattern}")) # loop and inspect all SOFA files for file in files: # inspect data save_to = folder if savedir is None else None _inspect_sofa_files(file, save_to, atol, plot, plane, dB_time, dB_freq, freq_scale)
def _inspect_sofa_files(file, savedir, atol, plot, plane, dB_time, dB_freq, freq_scale): with Dataset(file, "r", format="NETCDF4") as sofa_file: data_type = getattr(sofa_file, "DataType") if data_type == "FIR": mode = "hrir" elif data_type == "TF": mode = "hrtf" else: raise ValueError( f"The DataType of {file} is {data_type} but must 'FIR' or 'TF'") # load sofa file and source positions signal, sources, _ = pf.io.read_sofa(file) tail = os.path.basename(file) if "2D" in plot: with pf.plot.context(): # generate plot layout and axes if mode == "hrir": _, ax = plt.subplots(2, 4, figsize=(16, 6), sharey="row") ax_time = ax[0] ax_freq = ax[1] else: _, ax_freq = plt.subplots(1, 4, figsize=(16, 3), sharey="row") # loop sources max_db = -300 for nn, (az, name) in enumerate(zip( [0, 180, 90, 270], ["front", "back", "left", "right"])): # find current source find = pf.Coordinates.from_spherical_elevation( az / 180 * np.pi, 0, np.mean(sources.radius)) idx, _ = sources.find_nearest(find) # exact position for plotting source = sources.spherical_elevation[idx] name += (f" (az. {np.round(source[0] / np.pi * 180)}, " f"el. {np.round(source[1] / np.pi * 180)} deg.)") # plot if mode == "hrir": pf.plot.time_freq( signal[idx], ax=[ax_time[nn], ax_freq[nn]], dB_time=dB_time, dB_freq=dB_freq, freq_scale=freq_scale) ax_time[nn].set_title(name) else: pf.plot.freq(signal[idx], ax=ax_freq[nn], dB=dB_freq, freq_scale=freq_scale) ax_freq[nn].set_title(name) max_db = np.max([ax_freq[nn].get_ylim()[1], max_db]) # format axis and legend ax_freq[nn].set_ylim(max_db - 60, max_db) if signal.cshape[-1] == 2: ax_freq[nn].legend(["left ear", "right ear"], loc=3) elif signal.cshape[-1] > 2: ax_freq[nn].legend( [f"ch. {cc+1}" for cc in range(signal.cshape[-1])], loc=3) # save plt.savefig(os.path.join(savedir, tail[:-5] + "_2D.pdf"), bbox_inches="tight") if "3D" in plot: num_chanel = signal.cshape[-1] with pf.plot.context(): # generate plot layout and axes if mode == "hrir": _, ax = plt.subplots(2, num_chanel, sharex=True, sharey="row", figsize=(4*num_chanel, 6), ) ax_time = np.atleast_1d(ax[0]) ax_freq = np.atleast_1d(ax[1]) else: _, ax_freq = plt.subplots( 1, num_chanel, figsize=(4*num_chanel, 3), sharex=True, sharey="row") ax_freq = np.atleast_1d(ax_freq) # find sources on the desired plane if plane == "horizontal": idx = np.where(np.abs(sources.elevation) <= atol / 180 * np.pi) angles = np.round(sources.azimuth[idx] / np.pi * 180, 10) angle = "azimuth" elif plane == "median": idx = np.where(np.abs(sources.lateral) <= atol / 180 * np.pi) angles = np.round(sources.polar[idx] / np.pi * 180, 10) angle = "polar angle" else: idx = np.where( np.abs(sources.upper - np.pi / 2) <= atol / 180 * np.pi) angles = np.round(sources.frontal[idx] / np.pi * 180, 10) angle = "frontal angle" if not len(idx): warnings.warn(( f"Did not find and sources on the {plane} plane for " f"within +/-{atol} deg. for {file}")) return # sort angles source positions idx = idx[0][np.argsort(angles)] angles = np.sort(angles) # plot titles names = ["left ear", "right ear"] if signal.cshape[-1] == 2 \ else [f"ch. {cc+1}" for cc in range(signal.cshape[-1])] # loop sources for cc, name in enumerate(names): # plot time data if mode == "hrir": _, qm, _ = pf.plot.time_2d( signal[idx, cc], indices=angles, dB=dB_time, ax=ax_time[cc], cmap="coolwarm") ax_time[cc].set_title(name) # set limits of time plot (symmetric for nice colors, clip # to improve visibility) c_lim = qm.get_clim() c_lim = np.round(.6 * np.max(np.abs(c_lim)), 1) qm.set_clim(-c_lim, c_lim) # plot frequency data _, qm, _ = pf.plot.freq_2d( signal[idx, cc], indices=angles, ax=ax_freq[cc], dB=dB_freq, freq_scale=freq_scale, cmap="Reds") if mode == "hrir": ax_freq[cc].set_xlabel(f"{angle} in degree") ax_time[cc].set_xlabel("") else: ax_freq[cc].set_title(name) ax_freq[cc].set_xlabel(f"{angle} in degree") c_lim = qm.get_clim() c_lim = np.round(np.max(c_lim)) qm.set_clim(c_lim - 60, c_lim) # save plt.savefig( os.path.join(savedir, f"{tail[:-5]}_3D_{plane}_plane.jpeg"), dpi=300, bbox_inches="tight")