Refactor the parametric energy decay curve function

pyrato/parametric.py

@@ -5,82 +5,95 @@
 """
 import numpy as np
 from typing import Union
+import pyfar as pf
 
 
-def energy_decay_curve_analytic(
-        surfaces, alphas, volume, times, source=None,
-        receiver=None, method='eyring', c=343.4, frequency=None,
-        air_absorption=True):
-    """Calculate the energy decay curve analytically by using Eyring's or
-    Sabine's equation.
+def energy_decay_curve(
+        times : np.ndarray,
+        reverberation_time : float | np.ndarray,
+        energy : float | np.ndarray = 1,
+    ) -> pf.TimeData:
+    r"""Calculate the energy decay curve for the reverberation time and energy.
 
-    Calculation according to [#]_.
+    The energy decay curve is calculated as
+
+    .. math::
+        E(t) = E_0 e^{-\frac{6 \ln(10)}{T_{60}} t}
+
+    where :math:`E_0` is the initial energy, :math:`T_{60}` the reverberation
+    time, and :math:`t` the time [#]_.
 
     Parameters
     ----------
-    surfaces : ndarray, double
-        Surface areas of all surfaces in the room
-    alphas : ndarray, double
-        Absorption coefficients corresponding to each surface
-    volume : double
-        Room volume
-    times : ndarray, double
-        Time vector for which the decay curve is calculated
-    source : Coordinates
-        Coordinate object with the source coordinates
-    receiver : Coordinates
-        Coordinate object with the receiver coordinates
-    method : 'eyring', 'sabine'
-        Use either Eyring's or Sabine's equation
-    c : double
-        Speed of sound
-    frequency : double, optional
-        Center frequency of the respective octave band. This is only used for
-        the air absorption calculation.
-    air_absorption : bool, optional
-        If True, the air absorption is included in the calculation.
-        Default is True.
+    times : numpy.ndarray[float]
+        The times at which the energy decay curve is evaluated.
+    reverberation_time : float | numpy.ndarray[float]
+        The reverberation time in seconds. If an array is passed, an energy
+        decay curve is calculated for each reverberation time.
+    energy : float | numpy.ndarray[float], optional
+        The initial energy of the sound field, by default 1. If
+        `reverberation_time` is an array, the shape of `energy` is required
+        to match the shape or be broadcastable to the shape of
+        `reverberation_time`.
 
     Returns
     -------
-    energy_decay_curve : ndarray, double
-        The energy decay curve
+    energy_decay_curve : pyfar.TimeData
+        The energy decay curve with a ``cshape`` equal to the shape of
+        the passed ``reverberation_time``.
+
+    Example
+    -------
+    Calculate and plot an energy decay curve with a reverberation time of
+    2 seconds.
+
+    .. plot::
+
+        >>> import numpy as np
+        >>> import pyrato
+        >>> import pyfar as pf
+        >>>
+        >>> times = np.linspace(0, 3, 50)
+        >>> T_60 = [2, 1]
+        >>> edc = pyrato.parametric.energy_decay_curve(times, T_60)
+        >>> pf.plot.time(edc, log_prefix=10, dB=True)
+
 
     References
     ----------
     .. [#] H. Kuttruff, Room acoustics, 4th Ed. Taylor & Francis, 2009.
 
     """
+    reverberation_time = np.asarray(reverberation_time)
+    energy = np.asarray(energy)
+    times = np.asarray(times)
+
+    if np.any(reverberation_time <= 0):
+        raise ValueError("Reverberation time must be greater than zero.")
+
+    if np.any(energy < 0):
+        raise ValueError("Energy must be greater than or equal to zero.")
+
+    if reverberation_time.shape != energy.shape:
+        try:
+            energy = np.broadcast_to(energy, reverberation_time.shape)
+        except ValueError as error:
+            raise ValueError(
+                "Reverberation time and energy must be broadcastable to the "
+                "same shape.",
+            ) from error
+
+    matching_shape = reverberation_time.shape
+    reverberation_time = reverberation_time.flatten()
+    energy = energy.flatten()
+
+    reverberation_time = np.atleast_2d(reverberation_time)
+    energy = np.atleast_2d(energy)
+
+    damping_term = (3*np.log(10) / reverberation_time).T
+    edc = energy.T * np.exp(-2*damping_term*times)
 
-    alphas = np.asarray(alphas)
-    surfaces = np.asarray(surfaces)
-    surface_room = np.sum(surfaces)
-    alpha_mean = np.sum(surfaces*alphas) / surface_room
-
-    if air_absorption:
-        m = air_attenuation_coefficient(frequency)
-    else:
-        m = 0
-
-    if all([source, receiver]):
-        dist_source_receiver = np.linalg.norm(
-            source.cartesian - receiver.cartesian)
-        delay_direct = dist_source_receiver / c
-    else:
-        delay_direct = 0
-
-    if method == 'eyring':
-        energy_decay_curve = np.exp(
-            -c*(times - delay_direct) *
-            ((-surface_room * np.log(1 - alpha_mean) / 4 / volume) + m))
-    elif method == 'sabine':
-        energy_decay_curve = np.exp(
-            -c*(times - delay_direct) *
-            ((surface_room * alpha_mean / 4 / volume) + m))
-    else:
-        raise ValueError("The method has to be either 'eyring' or 'sabine'.")
-
-    return energy_decay_curve
+    return pf.TimeData(np.reshape(edc, (*matching_shape, times.size)), times)
 
 
 def air_attenuation_coefficient(

tests/test_edc.py

@@ -46,40 +46,49 @@ def test_schroeder_integration_multi_cdim(is_energy):
     npt.assert_almost_equal(energy_decay_curve.time[1, 0], truth)
 
 
-def test_edc_eyring():
-    alphas = [0.9, 0.1]
-    surfaces = [2, 5*2]
-    volume = 2*2*2
+@pytest.mark.parametrize(
+        "edc_function",
+        [ra.edc.energy_decay_curve_chu,
+         ra.edc.energy_decay_curve_lundeby,
+         ra.edc.energy_decay_curve_chu_lundeby],
+)
+def test_multidim_edc(edc_function):
+    """
+    Test if edcs from multichannel signal are equal to corresponding single
+    channel edcs.
+    """
+    rir = pf.signals.files.room_impulse_response()
+    rir_oct = pf.dsp.filter.fractional_octave_bands(rir, 1)
+    shape = rir_oct.time.shape
+    edc = edc_function(rir_oct, channel_independent=True)
+
+    assert shape == edc.time.shape
+
+    edc = edc.flatten()
+    rir_oct = rir_oct.flatten()
+
+    for i in range(edc.cshape[0]):
+        baseline = edc_function(rir_oct[i])
+        npt.assert_array_equal(edc[i].time, baseline.time)
+
+
+def test_edc_function():
+    """Test for the parametric energy decay curve based on previous
+    implementation which also included RT prediction using Sabine's equation.
+    """
     times = np.linspace(0, 0.25, 50)
-    edc = ra.parametric.energy_decay_curve_analytic(
-        surfaces, alphas, volume, times, method='eyring', air_absorption=False)
 
-    truth = array([
-        1.00000000e+00, 8.39817186e-01, 7.05292906e-01, 5.92317103e-01,
-            4.97438083e-01, 4.17757051e-01, 3.50839551e-01, 2.94641084e-01,
-            2.47444646e-01, 2.07808266e-01, 1.74520953e-01, 1.46565696e-01,
-            1.23088390e-01, 1.03371746e-01, 8.68133684e-02, 7.29073588e-02,
-            6.12288529e-02, 5.14210429e-02, 4.31842755e-02, 3.62668968e-02,
-            3.04575632e-02, 2.55787850e-02, 2.14815032e-02, 1.80405356e-02,
-            1.51507518e-02, 1.27238618e-02, 1.06857178e-02, 8.97404943e-03,
-            7.53656094e-03, 6.32933340e-03, 5.31548296e-03, 4.46403394e-03,
-            3.74897242e-03, 3.14845147e-03, 2.64412365e-03, 2.22058049e-03,
-            1.86488165e-03, 1.56615966e-03, 1.31528780e-03, 1.10460130e-03,
-            9.27663155e-04, 7.79067460e-04, 6.54274242e-04, 5.49470753e-04,
-            4.61454981e-04, 3.87537824e-04, 3.25460924e-04, 2.73327678e-04,
-            2.29545281e-04, 1.92776072e-04,
-            ])
-
-    npt.assert_almost_equal(edc, truth)
-
-
-def test_edc_sabine():
-    alphas = [0.9, 0.1]
-    surfaces = [2, 5*2]
+    c = 343.4
+
     volume = 2*2*2
-    times = np.linspace(0, 0.25, 50)
-    edc = ra.parametric.energy_decay_curve_analytic(
-        surfaces, alphas, volume, times, method='sabine', air_absorption=False)
+    alphas = np.asarray([0.9, 0.1])
+    surfaces = np.asarray([2, 5*2])
+    surface_room = np.sum(surfaces)
+    alpha_mean = np.sum(surfaces*alphas) / surface_room
+
+    rt = 24*np.log(10)/c * volume / (surface_room*alpha_mean)
+    energy = 1
+    edc = ra.parametric.energy_decay_curve(times, rt, energy)
 
     truth = array([
         1.00000000e+00, 8.57869258e-01, 7.35939663e-01, 6.31340013e-01,
@@ -96,30 +105,38 @@ def test_edc_sabine():
         1.17632849e-03, 1.00913605e-03, 8.65706791e-04, 7.42663242e-04,
         6.37107964e-04, 5.46555336e-04,
         ])
-    npt.assert_almost_equal(edc, truth)
+    npt.assert_almost_equal(edc.time, np.atleast_2d(truth))
 
 
-@pytest.mark.parametrize(
-        "edc_function",
-        [ra.edc.energy_decay_curve_chu,
-         ra.edc.energy_decay_curve_lundeby,
-         ra.edc.energy_decay_curve_chu_lundeby],
-)
-def test_multidim_edc(edc_function):
-    """
-    Test if edcs from multichannel signal are equal to corresponding single
-    channel edcs.
+def test_parametric_edc():
+    """Test returned edc by the parametric edc function and shape broadcasting.
     """
-    rir = pf.signals.files.room_impulse_response()
-    rir_oct = pf.dsp.filter.fractional_octave_bands(rir, 1)
-    shape = rir_oct.time.shape
-    edc = edc_function(rir_oct, channel_independent=True)
+    times = np.linspace(0, 0.25, 50)
+    T_60 = np.array([2, 1])
 
-    assert shape == edc.time.shape
+    edc = ra.parametric.energy_decay_curve(times, T_60)
+    assert edc.cshape == (2,)
 
-    edc = edc.flatten()
-    rir_oct = rir_oct.flatten()
+    e_0 = np.ones_like(T_60)
+    edc = ra.parametric.energy_decay_curve(times, T_60, energy=e_0)
 
-    for i in range(edc.cshape[0]):
-        baseline = edc_function(rir_oct[i])
-        npt.assert_array_equal(edc[i].time, baseline.time)
+    damping_term = 6*np.log(10) / T_60
+    truth = np.atleast_2d(e_0).T * np.exp(-np.atleast_2d(damping_term).T*times)
+
+    npt.assert_almost_equal(edc.time, truth)
+    assert edc.cshape == (2,)
+
+    edc = ra.parametric.energy_decay_curve(
+        times, np.ones((3, 2)), energy=np.ones((3, 2)))
+
+    assert edc.cshape == (3, 2)
+
+
+def test_parametric_edc_wrong_shapes():
+    """Test error handling for wrong shapes."""
+    times = np.linspace(0, 0.25, 50)
+    T_60 = np.array([2, 1])
+    energy = np.array([1, 1, 1])
+
+    with pytest.raises(ValueError, match="same shape."):
+        ra.parametric.energy_decay_curve(times, T_60, energy=energy)