Improve imports and update core modules

PyRES/__init__.py

@@ -0,0 +1,71 @@
+"""
+PyRES: Python library for Reverberation Enhancement System development and simulation.
+"""
+
+from physical_room import PhRoom, PhRoom_dataset, PhRoom_wgn
+from virtual_room import (
+    VrRoom,
+    unitary_parallel_connections,
+    unitary_mixing_matrix,
+    random_FIRs,
+    phase_cancellation,
+    FDN,
+    unitary_reverberator,
+)
+from res import RES
+from loss_functions import MSE_evs_mod, MSE_evs_idxs, colorless_reverb
+from functional import (
+    energy_coupling,
+    direct_to_reverb_ratio,
+    system_equalization_curve,
+    one_pole_filter,
+    resonance_filter,
+    modal_reverb,
+    reverb_time,
+    simulate_setup,
+)
+from plots import (
+    plot_evs_distribution,
+    plot_evs_compare,
+    plot_irs_compare,
+    plot_spectrograms_compare,
+    plot_room_setup,
+    plot_coupling,
+    plot_DRR,
+    plot_distributions,
+)
+from dataset_api import (
+    get_hl_info,
+    get_ll_info,
+    get_transducer_number,
+    get_transducer_positions,
+    get_rirs,
+    get_rir_metadata,
+)
+from utils import find_direct_path, expand_to_dimension, limit_frequency_points, next_power_of_2
+
+__all__ = [
+    # Physical room
+    'PhRoom', 'PhRoom_dataset', 'PhRoom_wgn',
+    # Virtual room
+    'VrRoom', 'unitary_parallel_connections', 'unitary_mixing_matrix',
+    'random_FIRs', 'phase_cancellation', 'FDN', 'unitary_reverberator',
+    # RES
+    'RES',
+    # Loss functions
+    'MSE_evs_mod', 'MSE_evs_idxs', 'colorless_reverb',
+    # Functional
+    'energy_coupling', 'direct_to_reverb_ratio', 'system_equalization_curve',
+    'one_pole_filter', 'resonance_filter', 'modal_reverb', 'reverb_time',
+    'simulate_setup',
+    # Plots
+    'plot_evs_distribution', 'plot_evs_compare', 'plot_irs_compare',
+    'plot_spectrograms_compare', 'plot_room_setup', 'plot_coupling',
+    'plot_DRR', 'plot_distributions',
+    # Dataset API
+    'get_hl_info', 'get_ll_info', 'get_transducer_number',
+    'get_transducer_positions', 'get_rirs', 'get_rir_metadata',
+    # Utils
+    'find_direct_path', 'expand_to_dimension', 'limit_frequency_points',
+    'next_power_of_2',
+]

PyRES/dataset_api.py

@@ -4,8 +4,7 @@
 import json
 # PyTorch
 import torch, torchaudio
-# PyRES
-from PyRES.functional import energy_coupling
+from functional import energy_coupling
 
 
 # ==================================================================
@@ -470,4 +469,20 @@ def check_requested_indices(
     assert all(i >= 0 for i in idx), f"{type} indices start from 0. You cannot request a negative index."
     assert len(idx) == len(set(idx)), f"Requested indices of {type} must be unique."
     assert max(idx) <= number, f"For the requested room, the maximum index of {type} is {number-1}. You cannot request the index {max(idx)}." 
-    return None
+    return None
+
+
+__all__ = [
+    'get_hl_info',
+    'get_ll_info',
+    'get_transducer_number',
+    'get_number_of',
+    'get_transducer_positions',
+    'get_positions_of',
+    'get_rir_metadata',
+    'get_rir_foldername_of',
+    'get_rirs',
+    'get_rirs_of',
+    'normalize_rirs',
+    'check_requested_indices',
+]

PyRES/functional.py

@@ -8,8 +8,8 @@
 import torch
 # FLAMO
 from flamo.functional import db2mag
-# PyRES
-from PyRES.utils import expand_to_dimension, find_direct_path
+
+from utils import expand_to_dimension, find_direct_path
 
 
 # ==================================================================
@@ -474,3 +474,19 @@ def system_equalization_curve(
             target = scaling_factor * torch.ones(mean_evs.shape[0],)
 
         return target
+
+__all__ = [
+    # Physical room functions
+    'simulate_setup',
+    'positions_on_surface',
+    'reverb_time',
+    'energy_coupling',
+    'direct_to_reverb_ratio',
+    # Virtual room functions
+    'one_pole_filter',
+    'resonance_filter',
+    'modal_reverb',
+    # Optimization functions
+    'system_equalization_curve',
+]
+

PyRES/loss_functions.py

@@ -150,3 +150,9 @@ def forward(self, y_pred, y_target, model):
         target = torch.ones_like(prediction)
 
         return self.mse_loss(prediction, target)
+
+__all__ = [
+    'MSE_evs_mod',
+    'MSE_evs_idxs',
+    'colorless_reverb',
+]

PyRES/physical_room.py

@@ -8,17 +8,16 @@
 from flamo import dsp
 from flamo.functional import mag2db, WGN_reverb
 # PyRES
-from PyRES.dataset_api import (
+from dataset_api import (
     get_hl_info,
     get_ll_info,
     get_rirs,
     normalize_rirs,
     get_transducer_number,
     get_transducer_positions
 )
-from PyRES.functional import energy_coupling, direct_to_reverb_ratio
-from PyRES.functional import simulate_setup
-from PyRES.plots import (
+from functional import energy_coupling, direct_to_reverb_ratio, simulate_setup
+from plots import (
     plot_room_setup,
     plot_coupling,
     plot_DRR,
@@ -652,4 +651,10 @@ def regenerate_h_LM(self) -> None:
         self.energy_coupling['LM'] = energy_coupling(new_rirs, fs=self.fs)
         self.direct_to_reverb_ratio['LM'] = direct_to_reverb_ratio(new_rirs, fs=self.fs)
 
-        return None
+        return None
+
+__all__ = [
+    'PhRoom',
+    'PhRoom_dataset',
+    'PhRoom_wgn',
+]

PyRES/plots.py

@@ -413,3 +413,14 @@ def plot_spectrograms_compare(ir_1: torch.Tensor, ir_2: torch.Tensor, fs: int, n
     cbar.ax.set_yticks(ticks, ['-100','-80','-60','-40','-20','0'])
 
     plt.show(block=True)
+
+__all__ = [
+    'plot_room_setup',
+    'plot_coupling',
+    'plot_DRR',
+    'plot_distributions',
+    'plot_evs_distribution',
+    'plot_evs_compare',
+    'plot_irs_compare',
+    'plot_spectrograms_compare',
+]

PyRES/res.py

@@ -10,9 +10,9 @@
 from flamo import dsp, system
 from flamo.functional import db2mag, mag2db, get_magnitude, get_eigenvalues
 # PyRES
-from PyRES.physical_room import PhRoom
-from PyRES.virtual_room import VrRoom
-from PyRES.utils import expand_to_dimension
+from physical_room import PhRoom
+from virtual_room import VrRoom
+from utils import expand_to_dimension
 
 
 # ==================================================================
@@ -431,3 +431,7 @@ def save_state_to(self, directory: str) -> None:
         directory = directory.rstrip('/')
         state = self.get_v_ML_state()
         torch.save(state, os.path.join(directory, time.strftime("%Y-%m-%d_%H.%M.%S.pt")))
+
+__all__ = [
+    'RES',
+]

PyRES/utils.py

@@ -111,4 +111,10 @@ def limit_frequency_points(array: torch.Tensor, fs: int, nfft: int, f_interval:
         subset = torch.arange(0, array.shape[-1])
 
     return torch.take_along_dim(array, subset, 0)
-
+
+__all__ = [
+    'next_power_of_2',
+    'expand_to_dimension',
+    'find_direct_path',
+    'limit_frequency_points',
+]

PyRES/virtual_room.py

@@ -7,7 +7,7 @@
 from flamo.functional import db2mag, skew_matrix
 from flamo.auxiliary.reverb import rt2slope
 # PyRES
-from PyRES.functional import modal_reverb, one_pole_filter
+from functional import modal_reverb, one_pole_filter
 
 
 # ==================================================================
@@ -835,4 +835,16 @@ def initialize_class(self):
         self.check_param_shape()
         self.get_io()
         self.get_freq_response()
-        self.get_freq_convolve()
+        self.get_freq_convolve()
+
+__all__ = [
+    'VrRoom',
+    'unitary_parallel_connections',
+    'unitary_mixing_matrix',
+    'random_FIRs',
+    'phase_cancellation',
+    'FDN',
+    'unitary_reverberator',
+    'FDN_one_pole_absorption',
+    'phase_cancelling_modal_reverb',
+]