Finalizes the bugfix of the angle_function function

CHANGELOG.rst

@@ -35,6 +35,8 @@ Bugfix
   @hrosseel
 - Fixes a bug when setting the air absorption coefficients to custom values (#191),
   adds a test, and more details in the doc
+- Fixes a bug in the utilities.angle_function in the calculation of the colatitude (#329)
+  by @fabiodimarco
 
 
 `0.7.3`_ - 2022-12-05

pyroomacoustics/tests/test_angle_function.py

@@ -10,17 +10,28 @@
 a1 = np.array([[0, 0, 1], [0, 0, 1], [0, 1, 1]])
 a2 = np.array([[0], [0], [1]])
 a3 = np.array([0, 0, 1]).T
+a4 = np.array([[0, 0, 1], [0, 1, 0], [1, 0, 0]])
 
 b1 = np.array([[0], [0], [0]])
 b2 = np.array([1, -1, 1]).T
+b3 = np.array([1, 0, 0]).T
 
 
-a1_b1 = np.array([[0, 0, pi / 4], [0, 0, pi / 4]])
-a1_b2 = np.array([[3 * pi / 4, 3 * pi / 4, pi / 2], [3 * pi / 4, pi / 2, pi / 2]])
+a1_b1 = np.array([[0, 0, pi / 4], [0, 0, np.arctan(np.sqrt(2.0))]])
+a1_b2 = np.array(
+    [
+        [3 * pi / 4, 3 * pi / 4, pi / 2],
+        [pi / 2 + np.arctan(1.0 / np.sqrt(2)), pi / 2, pi / 2],
+    ]
+)
 a2_b1 = np.array([[0], [0]])
 a2_b2 = np.array([[3 * pi / 4], [pi / 2]])
 a3_b1 = np.array([[0], [0]])
 a3_b2 = np.array([[3 * pi / 4], [pi / 2]])
+a4_b1 = np.array([[0, pi / 2, 0], [0, pi / 2, pi / 2]])
+a4_b3 = np.array([[pi, 3 * pi / 4, 0], [pi / 4, pi / 2, 0]])
+a2_b3 = np.array([[pi], [pi / 4]])
+a3_b3 = np.array([[pi], [pi / 4]])
 
 
 # for 2-D coordinates
@@ -40,92 +51,38 @@
 c3_d2 = np.array([[0], [pi / 2]])
 
 
-class TestAngleFunction(TestCase):
-    def test_set_3d(self):
-        self.assertTrue(angle_function(a1, b1).all() == a1_b1.all())
-        self.assertTrue(angle_function(a1, b2).all() == a1_b2.all())
-
-    def test_point_3d_1(self):
-        self.assertTrue(angle_function(a2, b1).all() == a2_b1.all())
-        self.assertTrue(angle_function(a2, b2).all() == a2_b2.all())
-
-    def test_point_3d_2(self):
-        self.assertTrue(angle_function(a3, b1).all() == a3_b1.all())
-        self.assertTrue(angle_function(a3, b2).all() == a3_b2.all())
-
-    def test_set_2d(self):
-        self.assertTrue(angle_function(c1, d1).all() == c1_d1.all())
-        self.assertTrue(angle_function(c1, d2).all() == c1_d2.all())
-
-    def test_point_2d_1(self):
-        self.assertTrue(angle_function(c2, d1).all() == c2_d1.all())
-        self.assertTrue(angle_function(c2, d2).all() == c2_d2.all())
-
-    def test_point_2d_2(self):
-        self.assertTrue(angle_function(c3, d1).all() == c3_d1.all())
-        self.assertTrue(angle_function(c3, d2).all() == c3_d2.all())
-
-
-def find_error(type_coordinates):
-    if type_coordinates == "3-D":
-        print("-" * 40)
-        print("type_coordinates = 3-D")
-        print("-" * 40)
-
-        a_range = [a1, a2, a3]
-        b_range = [b1, b2]
-        a_b_range = [a1_b1, a1_b2, a2_b1, a2_b2, a3_b1, a3_b2]
-
-        a_b_index = 0
-        for a in a_range:
-            for b in b_range:
-                error_azimuth = (angle_function(a, b) - a_b_range[a_b_index])[0]
-                error_colatitude = (angle_function(a, b) - a_b_range[a_b_index])[1]
-
-                print("for points :")
-                print(a)
-                print(b)
-                print(
-                    "error in azimuth calculation: {}".format(np.average(error_azimuth))
-                )
-                print(
-                    "error in colatitude calculation: {}".format(
-                        np.average(error_colatitude)
-                    )
-                )
-                print()
-                a_b_index += 1
-
-    elif type_coordinates == "2-D":
-        print("-" * 40)
-        print("type_coordinates = 2-D")
-        print("-" * 40)
-
-        c_range = [c1, c2, c3]
-        d_range = [d1, d2]
-        c_d_range = [c1_d1, c1_d2, c2_d1, c2_d2, c3_d1, c3_d2]
-
-        c_d_index = 0
-        for c in c_range:
-            for d in d_range:
-                error_azimuth = (angle_function(c, d) - c_d_range[c_d_index])[0]
-                error_colatitude = (angle_function(c, d) - c_d_range[c_d_index])[1]
-
-                print("for points :")
-                print(c)
-                print(d)
-                print(
-                    "error in azimuth calculation: {}".format(np.average(error_azimuth))
-                )
-                print(
-                    "error in colatitude calculation: {}".format(
-                        np.average(error_colatitude)
-                    )
-                )
-                print()
-                c_d_index += 1
-
-
-if __name__ == "__main__":
-    find_error("3-D")
-    find_error("2-D")
+def test_set_3d():
+    assert np.allclose(angle_function(a1, b1), a1_b1)
+    assert np.allclose(angle_function(a1, b2), a1_b2)
+
+
+def test_point_3d_1():
+    assert np.allclose(angle_function(a2, b1), a2_b1)
+    assert np.allclose(angle_function(a2, b2), a2_b2)
+    assert np.allclose(angle_function(a2, b3), a2_b3)
+
+
+def test_point_3d_2():
+    assert np.allclose(angle_function(a3, b1), a3_b1)
+    assert np.allclose(angle_function(a3, b2), a3_b2)
+    assert np.allclose(angle_function(a3, b3), a3_b3)
+
+
+def test_point_3d_3():
+    assert np.allclose(angle_function(a4, b1), a4_b1)
+    assert np.allclose(angle_function(a4, b3), a4_b3)
+
+
+def test_set_2d():
+    assert np.allclose(angle_function(c1, d1), c1_d1)
+    assert np.allclose(angle_function(c1, d2), c1_d2)
+
+
+def test_point_2d_1():
+    assert np.allclose(angle_function(c2, d1), c2_d1)
+    assert np.allclose(angle_function(c2, d2), c2_d2)
+
+
+def test_point_2d_2():
+    assert np.allclose(angle_function(c3, d1), c3_d1)
+    assert np.allclose(angle_function(c3, d2), c3_d2)

pyroomacoustics/utilities.py

@@ -29,6 +29,7 @@
 from scipy import signal
 from scipy.io import wavfile
 
+from .doa import cart2spher
 from .parameters import constants, eps
 from .sync import correlate
 
@@ -774,7 +775,8 @@ def all_combinations(lst1, lst2):
 
 def angle_function(s1, v2):
     """
-    Compute azimuth and colatitude angles in radians for a given set of points `s1` and a singular point `v2`.
+    Compute azimuth and colatitude angles in radians for a given set of points `s1`
+    with respect to a reference point `v2`.
 
     Parameters
     -----------
@@ -786,8 +788,9 @@ def angle_function(s1, v2):
     Returns
     -----------
     numpy array
-        2×N numpy array with azimuth and colatitude angles in radians.
-
+        2×N numpy array with azimuth and colatitude angles in radians in the
+        first and second row, respectively.
+        If the input vectors are 2-D, the colatitude is always fixed to pi/2.
     """
 
     if len(s1.shape) == 1:
@@ -797,26 +800,19 @@ def angle_function(s1, v2):
 
     assert s1.shape[0] == v2.shape[0]
 
-    x_vals = s1[0]
-    y_vals = s1[1]
-    x2 = v2[0]
-    y2 = v2[1]
-
-    # colatitude calculation for 3-D coordinates
-    if s1.shape[0] == 3 and v2.shape[0] == 3:
-        z2 = v2[2]
-        z_vals = s1[2]
-
-        colatitude = np.arctan2(
-            ((x_vals - x2) ** 2 + (y_vals - y2) ** 2) ** 1 / 2, (z_vals - z2)
-        )
+    ndim = s1.shape[0]
+    if ndim == 2:
+        s1 = np.concatenate((s1, np.zeros((1, s1.shape[1]))), axis=0)
+        v2 = np.concatenate((v2, np.zeros((1, v2.shape[1]))), axis=0)
 
-    # colatitude calculation for 2-D coordinates
-    elif s1.shape[0] == 2 and v2.shape[0] == 2:
-        num_points = s1.shape[1]
-        colatitude = np.ones(num_points) * np.pi / 2
+    # this is slightly wasteful for 2d points, but is safer as we are relying
+    # on tried and tested code
+    az, co, r = cart2spher(s1 - v2)
 
-    # azimuth calculation (same for 2-D and 3-D)
-    azimuth = np.arctan2((y_vals - y2), (x_vals - x2))
+    if ndim == 2:
+        # this is only necessary to handle correctly the case s1 - v2 = 0
+        # in this case cart2spher returns zero, but we would like to have
+        # colatitude = pi/2 for consistency
+        co[:] = np.pi / 2
 
-    return np.vstack((azimuth, colatitude))
+    return np.vstack((az, co))