flamo 0.1.13__py3-none-any.whl → 0.2.0__py3-none-any.whl

flamo/auxiliary/config/config.py

@@ -23,6 +23,8 @@ class HomogeneousFDNConfig(BaseModel):
     nfft: int = 96000
     # device to run the model
     device: str = 'cpu'
+    # data type
+    dtype: torch.dtype = torch.float32
     # delays in samples
     delays: Optional[List[int]] = None
     # delay lengths range in ms
@@ -76,4 +78,4 @@ class HomogeneousFDNConfig(BaseModel):
             ), "CUDA is not available for training"
     
     # forbid extra fields - adding this to help prevent errors in config file creation
-    model_config = ConfigDict(extra="forbid")
+    model_config = ConfigDict(extra="forbid", arbitrary_types_allowed=True)

flamo/auxiliary/eq.py

@@ -5,7 +5,7 @@ from flamo.functional import db2mag, shelving_filter, peak_filter, probe_sos
 from flamo.auxiliary.minimize import minimize_LBFGS
 
 
-def eq_freqs(interval: int = 1, start_freq: float = 31.25, end_freq: float = 16000.0):
+def eq_freqs(interval: int = 1, start_freq: float = 31.25, end_freq: float = 16000.0, device: str = "cpu", dtype: torch.dtype = torch.float32):
     r"""
     Calculate the center frequencies and shelving crossover frequencies for an equalizer.
 
@@ -19,13 +19,13 @@ def eq_freqs(interval: int = 1, start_freq: float = 31.25, end_freq: float = 160
 
     """
     center_freq = torch.tensor(
-        octave_bands(interval=interval, start_freq=start_freq, end_freq=end_freq)
+        octave_bands(interval=interval, start_freq=start_freq, end_freq=end_freq), device=device, dtype=dtype
     )
     shelving_crossover = torch.tensor(
         [
             center_freq[0] / np.power(2, 1 / interval / 2),
             center_freq[-1] * np.power(2, 1 / interval / 2),
-        ]
+        ], device=device, dtype=dtype
     )
 
     return center_freq, shelving_crossover
@@ -61,6 +61,7 @@ def geq(
     gain_db: torch.Tensor,
     fs: int = 48000,
     device: str = "cpu",
+    dtype: torch.dtype = torch.float32,
 ):
     r"""
     Computes the second-order sections coefficients of a graphic equalizer.
@@ -73,6 +74,7 @@ def geq(
         - **gain_db** (torch.Tensor): Tensor containing the gain values in decibels for each frequency band.
         - **fs** (int, optional): Sampling frequency. Default: 48000 Hz.
         - **device** (str, optional): Device to use for constructing tensors. Default: cpu.
+        - **dtype** (torch.dtype, optional): Data type for tensors. Default: torch.float32.
 
     **Returns**:
         - tuple: A tuple containing the numerator and denominator coefficients of the GEQ filter.
@@ -82,25 +84,25 @@ def geq(
     assert (
         len(gain_db) == num_bands
     ), "The number of gains must be equal to the number of frequencies."
-    sos = torch.zeros((6, num_bands), device=device)
+    sos = torch.zeros((6, num_bands), device=device, dtype=dtype)
 
     for band in range(num_bands):
         if band == 0:
-            b = torch.zeros(3, device=device)
+            b = torch.zeros(3, device=device, dtype=dtype)
             b[0] = db2mag(gain_db[band])
-            a = torch.tensor([1, 0, 0], device=device)
+            a = torch.tensor([1, 0, 0], device=device, dtype=dtype)
         elif band == 1:
             b, a = shelving_filter(
-                shelving_freq[0], db2mag(gain_db[band]), "low", fs=fs, device=device
+                shelving_freq[0], db2mag(gain_db[band]), "low", fs=fs, device=device, dtype=dtype
             )
         elif band == num_bands - 1:
             b, a = shelving_filter(
-                shelving_freq[1], db2mag(gain_db[band]), "high", fs=fs, device=device
+                shelving_freq[1], db2mag(gain_db[band]), "high", fs=fs, device=device, dtype=dtype
             )
         else:
             Q = torch.sqrt(R) / (R - 1)
             b, a = peak_filter(
-                center_freq[band - 2], db2mag(gain_db[band]), Q, fs=fs, device=device
+                center_freq[band - 2], db2mag(gain_db[band]), Q, fs=fs, device=device, dtype=dtype
             )
 
         sos_band = torch.hstack((b, a))
@@ -115,6 +117,7 @@ def accurate_geq(
     shelving_crossover: torch.Tensor,
     fs=48000,
     device: str = "cpu",
+    dtype: torch.dtype = torch.float32,
 ):
     r"""
     Design a Graphic Equalizer (GEQ) filter.
@@ -125,6 +128,7 @@ def accurate_geq(
             - shelving_crossover (torch.Tensor): Crossover frequencies for shelving filters.
             - fs (int, optional): Sampling frequency. Default: 48000 Hz.
             - device (str, optional): Device to use for constructing tensors. Default: 'cpu'.
+            - dtype (torch.dtype, optional): Data type for tensors. Default: torch.float32.
 
         **Returns**:
             - tuple: A tuple containing the numerator and denominator coefficients of the GEQ filter.
@@ -141,38 +145,38 @@ def accurate_geq(
 
     nfft = 2**16
     num_freq = len(center_freq) + len(shelving_crossover)
-    R = torch.tensor(2.7)
+    R = torch.tensor(2.7, dtype=dtype)
     # Control frequencies are spaced logarithmically
     num_control = 100
     control_freq = torch.round(
-        torch.logspace(np.log10(1), np.log10(fs / 2.1), num_control + 1)
+        torch.logspace(np.log10(1), np.log10(fs / 2.1), num_control + 1, dtype=dtype)
     )
     # interpolate the target gain values at control frequencies
-    target_freq = torch.cat((torch.tensor([1]), center_freq, torch.tensor([fs / 2.1])))
+    target_freq = torch.cat((torch.tensor([1], dtype=dtype), center_freq, torch.tensor([fs / 2.1], dtype=dtype)))
     # targetInterp = torch.tensor(np.interp(control_freq, target_freq, target_gain.squeeze()))
     interp = RegularGridInterpolator([target_freq], target_gain)
     targetInterp = interp([control_freq])
 
     # Design prototype of the biquad sections
     prototype_gain = 10  # dB
-    prototype_gain_array = torch.full((num_freq + 1, 1), prototype_gain)
+    prototype_gain_array = torch.full((num_freq + 1, 1), prototype_gain, dtype=dtype)
     prototype_b, prototype_a = geq(
-        center_freq, shelving_crossover, R, prototype_gain_array, fs
+        center_freq, shelving_crossover, R, prototype_gain_array, fs, dtype=dtype
     )
     prototype_sos = torch.vstack((prototype_b, prototype_a))
-    G, _, _ = probe_sos(prototype_sos, control_freq, nfft, fs)
+    G, _, _ = probe_sos(prototype_sos, control_freq, nfft, fs, dtype=dtype)
     G = G / prototype_gain  # dB vs control frequencies
 
     # Define the optimization bounds
     upperBound = torch.tensor(
-        [torch.inf] + [2 * prototype_gain] * num_freq, device=device
+        [torch.inf] + [2 * prototype_gain] * num_freq, device=device, dtype=dtype
     )
-    lowerBound = torch.tensor([-val for val in upperBound], device=device)
+    lowerBound = torch.tensor([-val for val in upperBound], device=device, dtype=dtype)
 
     # Optimization
     opt_gains = minimize_LBFGS(G, targetInterp, lowerBound, upperBound, num_freq)
 
     # Generate the SOS coefficients
-    b, a = geq(center_freq, shelving_crossover, R, opt_gains, fs, device=device)
+    b, a = geq(center_freq, shelving_crossover, R, opt_gains, fs, device=device, dtype=dtype)
 
     return b, a

flamo/auxiliary/reverb.py

@@ -104,11 +104,12 @@ class HomogeneousFDN:
     def set_model(self, input_layer=None, output_layer=None):
         # set the input and output layers of the FDN model
         if input_layer is None:
-            input_layer = dsp.FFT(self.config_dict.nfft)
+            input_layer = dsp.FFT(self.config_dict.nfft, dtype=self.config_dict.dtype)
         if output_layer is None:
             output_layer = dsp.iFFTAntiAlias(
                 nfft=self.config_dict.nfft,
                 alias_decay_db=self.config_dict.alias_decay_db,
+                dtype=self.config_dict.dtype,
             )
 
         self.model = self.get_shell(input_layer, output_layer)
@@ -125,6 +126,7 @@ class HomogeneousFDN:
             requires_grad=self.config_dict.input_gain_grad,
             alias_decay_db=self.config_dict.alias_decay_db,
             device=self.config_dict.device,
+            dtype=self.config_dict.dtype,
         )
         output_gain = dsp.Gain(
             size=(1, self.N),
@@ -132,6 +134,7 @@ class HomogeneousFDN:
             requires_grad=self.config_dict.output_gain_grad,
             alias_decay_db=self.config_dict.alias_decay_db,
             device=self.config_dict.device,
+            dtype=self.config_dict.dtype,
         )
 
         # RECURSION
@@ -144,6 +147,7 @@ class HomogeneousFDN:
             requires_grad=self.config_dict.delays_grad,
             alias_decay_db=self.config_dict.alias_decay_db,
             device=self.config_dict.device,
+            dtype=self.config_dict.dtype,
         )
         # assign the required delay line lengths
         delays.assign_value(delays.sample2s(delay_lines))
@@ -156,6 +160,7 @@ class HomogeneousFDN:
             requires_grad=self.config_dict.mixing_matrix_grad,
             alias_decay_db=self.config_dict.alias_decay_db,
             device=self.config_dict.device,
+            dtype=self.config_dict.dtype,
         )
 
         # homogeneous attenuation
@@ -165,6 +170,7 @@ class HomogeneousFDN:
             requires_grad=self.config_dict.attenuation_grad,
             alias_decay_db=self.config_dict.alias_decay_db,
             device=self.config_dict.device,
+            dtype=self.config_dict.dtype,
         )
         attenuation.map = map_gamma(delay_lines)
         attenuation.assign_value(
@@ -328,7 +334,8 @@ class parallelFDNAccurateGEQ(dsp.parallelAccurateGEQ):
         alias_decay_db: float = 0.0,
         start_freq: float = 31.25,
         end_freq: float = 16000.0,
-        device=None
+        device=None, 
+        dtype=torch.float32
     ):
         assert (delays is not None), "Delays must be provided"
         self.delays = delays
@@ -342,7 +349,8 @@ class parallelFDNAccurateGEQ(dsp.parallelAccurateGEQ):
             alias_decay_db=alias_decay_db,
             start_freq=start_freq,
             end_freq=end_freq,
-            device=device
+            device=device, 
+            dtype=dtype
         )
 
 
@@ -394,7 +402,8 @@ class parallelGFDNAccurateGEQ(parallelFDNAccurateGEQ):
         alias_decay_db: float = 0.0,
         start_freq: float = 31.25,
         end_freq: float = 16000.0,
-        device=None
+        device=None,
+        dtype=torch.float32
     ):
         assert (delays is not None), "Delays must be provided"
         self.delays = delays
@@ -482,6 +491,7 @@ class parallelFDNGEQ(dsp.parallelGEQ):
         requires_grad: bool = False,
         alias_decay_db: float = 0.0,
         device: Optional[str] = None,
+        dtype=torch.float32
     ):
         assert (delays is not None), "Delays must be provided"
         self.delays = delays
@@ -497,7 +507,8 @@ class parallelFDNGEQ(dsp.parallelGEQ):
             map=map,
             requires_grad=requires_grad,
             alias_decay_db=alias_decay_db,
-            device=device
+            device=device, 
+            dtype=dtype
         )
 
     def get_poly_coeff(self, param):
@@ -516,14 +527,13 @@ class parallelFDNGEQ(dsp.parallelGEQ):
                     fs=self.fs,
                     device=self.device
                 )
-        b_aa = torch.einsum('p, pon -> pon', self.alias_envelope_dcy.to(torch.double), b.to(torch.double))
-        a_aa = torch.einsum('p, pon -> pon', self.alias_envelope_dcy.to(torch.double), a.to(torch.double))
+        b_aa = torch.einsum('p, pon -> pon', self.alias_envelope_dcy, b)
+        a_aa = torch.einsum('p, pon -> pon', self.alias_envelope_dcy, a)
         B = torch.fft.rfft(b_aa, self.nfft, dim=0)
         A = torch.fft.rfft(a_aa, self.nfft, dim=0)
         H_temp = torch.prod(B, dim=1) / (torch.prod(A, dim=1))
         H = torch.where(torch.abs(torch.prod(A, dim=1)) != 0, H_temp, torch.finfo(H_temp.dtype).eps*torch.ones_like(H_temp))
-        H_type = torch.complex128 if param.dtype == torch.float64 else torch.complex64
-        return H.to(H_type), B, A
+        return H, B, A
 
     def check_param_shape(self):
         assert (
@@ -560,6 +570,7 @@ class parallelFDNPEQ(Filter):
         requires_grad: bool = False,
         alias_decay_db: float = 0.0,
         device: Optional[str] = None,
+        dtype=torch.float32
     ):
         self.delays = delays
         self.is_twostage = is_twostage
@@ -576,12 +587,13 @@ class parallelFDNPEQ(Filter):
         self.center_freq_bias = f_min * (f_max / f_min) ** ((k - 1) / (self.n_bands - 1))
         self.alias_envelope_dcy = gamma ** torch.arange(0, 3, 1, device=device)
         super().__init__(
-            size=(self.n_bands+1 if self.is_twostage else self.n_band, 3, 1 if self.is_proportional else len(delays)),
+            size=(self.n_bands+1 if self.is_twostage else self.n_bands, 3, 1 if self.is_proportional else len(delays)),
             nfft=nfft,
             map=map,
             requires_grad=requires_grad,
             alias_decay_db=alias_decay_db,
             device=device,
+            dtype=dtype
         )
 
     def get_poly_coeff(self, param):
@@ -644,16 +656,13 @@ class parallelFDNPEQ(Filter):
                 type='highshelf'
             )
 
-        b_aa = torch.einsum("p, opn -> opn", self.alias_envelope_dcy.to(torch.double), b.to(torch.double))
-        a_aa = torch.einsum("p, opn -> opn", self.alias_envelope_dcy.to(torch.double), a.to(torch.double))
+        b_aa = torch.einsum("p, opn -> opn", self.alias_envelope_dcy, b)
+        a_aa = torch.einsum("p, opn -> opn", self.alias_envelope_dcy, a)
         B = torch.fft.rfft(b_aa, self.nfft, dim=1)
         A = torch.fft.rfft(a_aa, self.nfft, dim=1)
         H_temp = (torch.prod(B, dim=0) / (torch.prod(A, dim=0)))
-        # H_temp = (torch.prod(B, dim=0) / (torch.prod(A, dim=0)))
-
         H = torch.where(torch.abs(torch.prod(A, dim=0)) != 0, H_temp, torch.finfo(H_temp.dtype).eps*torch.ones_like(H_temp))
-        H_type = torch.complex128 if param.dtype == torch.float64 else torch.complex64
-        return H.to(H_type), B, A
+        return H, B, A
 
     def compute_biquad_coeff(self, f, R, G, type='peaking'):
         # f : freq, R : resonance, G : gain in dB
@@ -805,6 +814,7 @@ class parallelFirstOrderShelving(dsp.parallelFilter):
         delays: torch.Tensor =  None,
         alias_decay_db: float = 0.0,
         device: str = None,
+        dtype: torch.dtype = torch.float32
     ):
         size = (2,)      # rt at DC and crossover frequency
         assert (delays is not None), "Delays must be provided"
@@ -816,7 +826,8 @@ class parallelFirstOrderShelving(dsp.parallelFilter):
             nfft=nfft,
             map=map,
             alias_decay_db=alias_decay_db,
-            device=device
+            device=device,
+            dtype=dtype
         )
         gamma = 10 ** (
             -torch.abs(torch.tensor(alias_decay_db, device=device)) / (nfft) / 20

flamo/auxiliary/scattering.py

@@ -4,10 +4,6 @@ import numpy as np
 from typing import Optional
 from flamo.utils import to_complex
 
-torch.random.manual_seed(0)
-np.random.seed(0)
-
-
 class ScatteringMapping(nn.Module):
     r"""
     Class mapping an orthogonal matrix to a paraunitary matrix using sparse scattering.
@@ -47,23 +43,26 @@ class ScatteringMapping(nn.Module):
         m_L: Optional[torch.tensor] = None,
         m_R: Optional[torch.tensor] = None,
         device: str = "cpu",
+        dtype: torch.dtype = torch.float32
     ):
         super(ScatteringMapping, self).__init__()
 
         self.n_stages = n_stages
         self.sparsity = sparsity
         self.gain_per_sample = gain_per_sample
+        self.device = device
+        self.dtype = dtype
         if m_L is None:
-            self.m_L = torch.zeros(N, device=device)
+            self.m_L = torch.zeros(N, device=device, dtype=self.dtype)
         else:
             self.m_L = m_L
         if m_R is None:
-            self.m_R = torch.zeros(N, device=device)
+            self.m_R = torch.zeros(N, device=device, dtype=self.dtype)
         else:
             self.m_R = m_R
-        self.sparsity_vect = torch.ones((n_stages), device=device)
+        self.sparsity_vect = torch.ones((n_stages), device=device, dtype=self.dtype)
         self.sparsity_vect[0] = sparsity
-        self.shifts = get_random_shifts(N, self.sparsity_vect, pulse_size)
+        self.shifts = get_random_shifts(N, self.sparsity_vect, pulse_size, dtype=self.dtype)
 
     def forward(self, U):
         r"""
@@ -81,7 +80,7 @@ class ScatteringMapping(nn.Module):
 
             G = (
                 torch.diag(self.gain_per_sample ** self.shifts[k - 1, :])
-                .to(torch.float32)
+                .to(self.dtype)
                 .to(U.device)
             )
             R = torch.matmul(U[:, :, k], G)
@@ -103,6 +102,7 @@ def cascaded_paraunit_matrix(
     pulse_size: int = 1,
     m_L: Optional[torch.tensor] = None,
     m_R: Optional[torch.tensor] = None,
+    dtype: torch.dtype = torch.float32,
 ):
     r"""
     Creates paraunitary matrix from input orthogonal matrix.
@@ -122,7 +122,7 @@ def cascaded_paraunit_matrix(
     """
 
     K = n_stages + 1
-    sparsity_vect = torch.ones((n_stages), device=U.device)
+    sparsity_vect = torch.ones((n_stages), device=U.device, dtype=dtype)
     sparsity_vect[0] = sparsity
     # check that the input matrix is of correct shape
     assert U.shape[0] == K, "The input matrix must have n_stages+1 stages"
@@ -133,14 +133,14 @@ def cascaded_paraunit_matrix(
     N = V.shape[0]
 
     if m_L is None:
-        m_L = torch.zeros(N, device=U.device)
+        m_L = torch.zeros(N, device=U.device, dtype=dtype)
     if m_R is None:
-        m_R = torch.zeros(N, device=U.device)
+        m_R = torch.zeros(N, device=U.device, dtype=dtype)
 
-    shift_L = get_random_shifts(N, sparsity_vect, pulse_size)
+    shift_L = get_random_shifts(N, sparsity_vect, pulse_size, dtype=dtype)
     for k in range(1, K):
 
-        G = torch.diag(gain_per_sample ** shift_L[k - 1, :]).to(torch.float32)
+        G = torch.diag(gain_per_sample ** shift_L[k - 1, :]).to(dtype)
         R = torch.matmul(U[:, :, k], G)
 
         V = shift_matrix(V, shift_L[k - 1, :], direction="left")
@@ -168,7 +168,7 @@ def poly_matrix_conv(A: torch.tensor, B: torch.tensor):
     if szA[1] != szB[0]:
         raise ValueError("Invalid matrix dimension.")
 
-    C = torch.zeros((szA[0], szB[1], szA[2] + szB[2] - 1), device=A.device)
+    C = torch.zeros((szA[0], szB[1], szA[2] + szB[2] - 1), device=A.device, dtype=A.dtype)
 
     A = A.permute(2, 0, 1)
     B = B.permute(2, 0, 1)
@@ -202,7 +202,7 @@ def shift_matrix(X: torch.tensor, shift: torch.tensor, direction: str = "left"):
         required_space = order + shift.reshape(-1, 1)
         additional_space = int((required_space.max() - X.shape[-1]) + 1)
         X = torch.cat(
-            (X, torch.zeros((N, N, additional_space), device=shift.device)), dim=-1
+            (X, torch.zeros((N, N, additional_space), device=shift.device, dtype=X.dtype)), dim=-1
         )
         for i in range(N):
             X[i, :, :] = torch.roll(X[i, :, :], int(shift[i].item()), dims=-1)
@@ -210,7 +210,7 @@ def shift_matrix(X: torch.tensor, shift: torch.tensor, direction: str = "left"):
         required_space = order + shift.reshape(1, -1)
         additional_space = int((required_space.max() - X.shape[-1]) + 1)
         X = torch.cat(
-            (X, torch.zeros((N, N, additional_space), device=shift.device)), dim=-1
+            (X, torch.zeros((N, N, additional_space), device=shift.device, dtype=X.dtype)), dim=-1
         )
         for i in range(N):
             X[:, i, :] = torch.roll(X[:, i, :], int(shift[i].item()), dims=-1)
@@ -228,14 +228,14 @@ def shift_mat_distribute(X: torch.tensor, sparsity: int, pulse_size: int):
     return (rand_shift * pulse_size).int()
 
 
-def get_random_shifts(N, sparsity_vect, pulse_size):
-    rand_shift = torch.zeros(sparsity_vect.shape[0], N, device=sparsity_vect.device)
+def get_random_shifts(N, sparsity_vect, pulse_size, dtype=torch.float32):
+    rand_shift = torch.zeros(sparsity_vect.shape[0], N, device=sparsity_vect.device, dtype=dtype)
     for k in range(sparsity_vect.shape[0]):
         temp = torch.floor(
             sparsity_vect[k]
             * (
-                torch.arange(0, N, device=sparsity_vect.device)
-                + torch.rand((N), device=sparsity_vect.device) * 0.99
+                torch.arange(0, N, device=sparsity_vect.device, dtype=dtype)
+                + torch.rand((N), device=sparsity_vect.device, dtype=dtype) * 0.99
             )
         )
         rand_shift[k, :] = (temp * pulse_size).int()