Remove 2d impl from ray tracing (#3633)

* Fix some issues
* Remove 2D implementation from ray tracing

We only add 3D RIR.

test/cpp/rir/wall_collision.cpp

@@ -28,59 +28,6 @@ CollisionTestParam par(
       hit_distance};
 }
 
-//////////////////////////////////////////////////////////////////////////////
-// 2D test
-//////////////////////////////////////////////////////////////////////////////
-
-class Simple2DRoomCollisionTest
-    : public ::testing::TestWithParam<CollisionTestParam> {};
-
-TEST_P(Simple2DRoomCollisionTest, CollisionTest2D) {
-  //
-  //  ^
-  //  |        3
-  //  |     ______
-  //  |    |      |
-  //  |  0 |      | 1
-  //  |    |______|
-  //  |        2
-  // -+---------------->
-  //
-  auto room = torch::tensor({1, 1});
-
-  auto param = GetParam();
-  auto [hit_point, next_wall_index, hit_distance] =
-      find_collision_wall<float, 2>(room, param.origin, param.direction);
-
-  EXPECT_EQ(param.next_wall_index, next_wall_index);
-  EXPECT_FLOAT_EQ(param.hit_distance, hit_distance);
-  EXPECT_TRUE(torch::allclose(
-      param.hit_point, hit_point, /*rtol*/ 1e-05, /*atol*/ 1e-07));
-}
-
-#define ISQRT2 0.70710678118
-
-INSTANTIATE_TEST_CASE_P(
-    Collision2DTests,
-    Simple2DRoomCollisionTest,
-    ::testing::Values(
-        // From 0
-        par({0.0, 0.5}, {1.0, 0.0}, {1.0, 0.5}, 1, 1.0),
-        par({0.0, 0.5}, {1.0, -1.}, {0.5, 0.0}, 2, ISQRT2),
-        par({0.0, 0.5}, {1.0, 1.0}, {0.5, 1.0}, 3, ISQRT2),
-        // From 1
-        par({1.0, 0.5}, {-1., 0.0}, {0.0, 0.5}, 0, 1.0),
-        par({1.0, 0.5}, {-1., -1.}, {0.5, 0.0}, 2, ISQRT2),
-        par({1.0, 0.5}, {-1., 1.0}, {0.5, 1.0}, 3, ISQRT2),
-        // From 2
-        par({0.5, 0.0}, {-1., 1.0}, {0.0, 0.5}, 0, ISQRT2),
-        par({0.5, 0.0}, {1.0, 1.0}, {1.0, 0.5}, 1, ISQRT2),
-        par({0.5, 0.0}, {0.0, 1.0}, {0.5, 1.0}, 3, 1.0),
-        // From 3
-        par({0.5, 1.0}, {-1., -1.}, {0.0, 0.5}, 0, ISQRT2),
-        par({0.5, 1.0}, {1.0, -1.}, {1.0, 0.5}, 1, ISQRT2),
-        par({0.5, 1.0}, {0.0, -1.}, {0.5, 0.0}, 2, 1.0)));
-
 //////////////////////////////////////////////////////////////////////////////
 // 3D test
 //////////////////////////////////////////////////////////////////////////////
@@ -103,7 +50,7 @@ TEST_P(Simple3DRoomCollisionTest, CollisionTest3D) {
 
   auto param = GetParam();
   auto [hit_point, next_wall_index, hit_distance] =
-      find_collision_wall<float, 3>(room, param.origin, param.direction);
+      find_collision_wall<float>(room, param.origin, param.direction);
 
   EXPECT_EQ(param.next_wall_index, next_wall_index);
   EXPECT_FLOAT_EQ(param.hit_distance, hit_distance);
@@ -111,6 +58,8 @@ TEST_P(Simple3DRoomCollisionTest, CollisionTest3D) {
       param.hit_point, hit_point, /*rtol*/ 1e-05, /*atol*/ 1e-07));
 }
 
+#define ISQRT2 0.70710678118
+
 INSTANTIATE_TEST_CASE_P(
     Collision3DTests,
     Simple3DRoomCollisionTest,

torchaudio/csrc/rir/ray_tracing.cpp

@@ -45,22 +45,15 @@ const int ISM_ORDER = 10;
 #define MAX(x) (VAL((x).max()))
 #define IN_RANGE(x, y) ((-EPS < (x)) && ((x) < (y) + EPS))
 
-template <typename scalar_t, unsigned int D>
-const std::array<Wall<scalar_t>, D * 2> make_walls(
+template <typename scalar_t>
+const std::array<Wall<scalar_t>, 6> make_walls(
     const torch::Tensor& room,
     const torch::Tensor& absorption,
     const torch::Tensor& scattering) {
-  if constexpr (D == 2) {
-    auto w = room.index({0}).item<scalar_t>();
-    auto l = room.index({1}).item<scalar_t>();
-    return make_room<scalar_t>(w, l, absorption, scattering);
-  }
-  if constexpr (D == 3) {
-    auto w = room.index({0}).item<scalar_t>();
-    auto l = room.index({1}).item<scalar_t>();
-    auto h = room.index({2}).item<scalar_t>();
-    return make_room<scalar_t>(w, l, h, absorption, scattering);
-  }
+  auto w = room.index({0}).item<scalar_t>();
+  auto l = room.index({1}).item<scalar_t>();
+  auto h = room.index({2}).item<scalar_t>();
+  return make_room<scalar_t>(w, l, h, absorption, scattering);
 }
 
 inline double get_energy_coeff(
@@ -73,7 +66,7 @@ inline double get_energy_coeff(
 
 /// RayTracer class helper for ray tracing.
 /// For attribute description, Python wrapper.
-template <typename scalar_t, unsigned int D>
+template <typename scalar_t>
 class RayTracer {
   // Provided parameters
   const torch::Tensor& room;
@@ -84,7 +77,7 @@ class RayTracer {
   const int num_bands;
   const double mic_radius_sq;
   const bool do_scattering; // Whether scattering is needed (scattering != 0)
-  const std::array<Wall<scalar_t>, D * 2> walls; // The walls of the room
+  const std::array<Wall<scalar_t>, 6> walls; // The walls of the room
 
   // Runtime value caches
   // Updated at the beginning of the simulation
@@ -106,7 +99,7 @@ class RayTracer {
         num_bands(absorption.size(0)),
         mic_radius_sq(mic_radius * mic_radius),
         do_scattering(MAX(scattering) > 0.),
-        walls(make_walls<scalar_t, D>(room, absorption, scattering)) {}
+        walls(make_walls<scalar_t>(room, absorption, scattering)) {}
 
   // The main (and only) public entry point of this class. The histograms Tensor
   // reference is passed along and modified in the subsequent private method
@@ -134,38 +127,28 @@ class RayTracer {
     // TODO: the for loop can be parallelized over num_rays by creating
     // `num_threads` histograms and then sum-reducing them into a single
     // histogram.
-    static_assert(D == 2 || D == 3, "Only 2D and 3D are supported.");
-    if constexpr (D == 2) {
-      scalar_t delta = 2. * M_PI / num_rays;
-      for (int i = 0; i < num_rays; ++i) {
-        scalar_t phi = i * delta;
-        auto dir = torch::tensor({cos(phi), sin(phi)}, room.scalar_type());
-        simul_ray(energies, origin, dir, histograms);
-      }
-    } else {
-      scalar_t delta = 2. / num_rays;
-      scalar_t increment = M_PI * (3. - std::sqrt(5.)); // phi increment
+    scalar_t delta = 2. / num_rays;
+    scalar_t increment = M_PI * (3. - std::sqrt(5.)); // phi increment
 
-      for (auto i = 0; i < num_rays; ++i) {
-        auto z = (i * delta - 1) + delta / 2.;
-        auto rho = std::sqrt(1. - z * z);
+    for (auto i = 0; i < num_rays; ++i) {
+      auto z = (i * delta - 1) + delta / 2.;
+      auto rho = std::sqrt(1. - z * z);
 
-        scalar_t phi = i * increment;
+      scalar_t phi = i * increment;
 
-        auto x = cos(phi) * rho;
-        auto y = sin(phi) * rho;
+      auto x = cos(phi) * rho;
+      auto y = sin(phi) * rho;
 
-        auto azimuth = atan2(y, x);
-        auto colatitude = atan2(std::sqrt(x * x + y * y), z);
+      auto azimuth = atan2(y, x);
+      auto colatitude = atan2(std::sqrt(x * x + y * y), z);
 
-        auto dir = torch::tensor(
-            {sin(colatitude) * cos(azimuth),
-             sin(colatitude) * sin(azimuth),
-             cos(colatitude)},
-            room.scalar_type());
+      auto dir = torch::tensor(
+          {sin(colatitude) * cos(azimuth),
+           sin(colatitude) * sin(azimuth),
+           cos(colatitude)},
+          room.scalar_type());
 
-        simul_ray(energies, origin, dir, histograms);
-      }
+      simul_ray(energies, origin, dir, histograms);
     }
     return histograms.transpose(1, 2); // (num_mics, num_bands, num_bins)
   }
@@ -200,7 +183,7 @@ class RayTracer {
     while (true) {
       // Find the next hit point
       auto [hit_point, next_wall_index, hit_distance] =
-          find_collision_wall<scalar_t, D>(room, origin, dir);
+          find_collision_wall<scalar_t>(room, origin, dir);
 
       auto& wall = walls[next_wall_index];
 
@@ -326,38 +309,11 @@ torch::Tensor ray_tracing(
     double hist_bin_size) {
   // TODO: Raise this to Python layer
   auto num_bins = (int)ceil(time_thres / hist_bin_size);
-  switch (room.size(0)) {
-    case 2: {
-      return AT_DISPATCH_FLOATING_TYPES(
-          room.scalar_type(), "ray_tracing_2d", [&] {
-            RayTracer<scalar_t, 2> rt(
-                room, mic_array, absorption, scattering, mic_radius);
-            return rt.compute_histograms(
-                source,
-                num_rays,
-                time_thres,
-                energy_thres,
-                sound_speed,
-                num_bins);
-          });
-    }
-    case 3: {
-      return AT_DISPATCH_FLOATING_TYPES(
-          room.scalar_type(), "ray_tracing_3d", [&] {
-            RayTracer<scalar_t, 3> rt(
-                room, mic_array, absorption, scattering, mic_radius);
-            return rt.compute_histograms(
-                source,
-                num_rays,
-                time_thres,
-                energy_thres,
-                sound_speed,
-                num_bins);
-          });
-    }
-    default:
-      TORCH_CHECK(false, "Only 2D and 3D are supported.");
-  }
+  return AT_DISPATCH_FLOATING_TYPES(room.scalar_type(), "ray_tracing_3d", [&] {
+    RayTracer<scalar_t> rt(room, absorption, scattering, mic_array, mic_radius);
+    return rt.compute_histograms(
+        source, num_rays, time_thres, energy_thres, sound_speed, num_bins);
+  });
 }
 
 TORCH_LIBRARY_IMPL(torchaudio, CPU, m) {

torchaudio/csrc/rir/wall.h

@@ -74,41 +74,6 @@ scalar_t cosine(const Wall<scalar_t>& wall, const torch::Tensor& dir) {
 ///   `find_collision_wall` will search in the order x, y, z and
 ///   wall pairs must be distibguishable on these axis.
 
-/// 2D room
-template <typename T>
-const std::array<Wall<T>, 4> make_room(
-    const T w,
-    const T l,
-    const torch::Tensor& abs,
-    const torch::Tensor& scat) {
-  //
-  //                          (0, 1)
-  //            0:West          ^
-  //            (0, l)          |  3:North
-  // (-1, 0)  <--  + ---------- +  (w, l)
-  //               |            |
-  //               |            |
-  //       (0, 0)  + -----------+ --> (1, 0)
-  //       2:South |            (w, 0)
-  //               v            1:East
-  //              (0, -1)
-  //
-  //  y
-  //  ^
-  //  |
-  //  +-- > x
-  //
-  using namespace torch::indexing;
-#define SLICE(x, i) x.index({Slice(), i})
-  return {
-      Wall<T>({0, l}, {-1, 0}, SLICE(abs, 0), SLICE(scat, 0)), // West
-      Wall<T>({w, 0}, {1, 0}, SLICE(abs, 1), SLICE(scat, 1)), // East
-      Wall<T>({0, 0}, {0, -1}, SLICE(abs, 2), SLICE(scat, 2)), // South
-      Wall<T>({w, l}, {0, 1}, SLICE(abs, 3), SLICE(scat, 3)) // North
-  };
-#undef SLICE
-}
-
 /// 3D room
 template <typename T>
 const std::array<Wall<T>, 6> make_room(
@@ -124,8 +89,8 @@ const std::array<Wall<T>, 6> make_room(
       Wall<T>({w, 0, 0}, {1, 0, 0}, SLICE(abs, 1), SLICE(scat, 1)), // East
       Wall<T>({0, 0, 0}, {0, -1, 0}, SLICE(abs, 2), SLICE(scat, 2)), // South
       Wall<T>({w, l, 0}, {0, 1, 0}, SLICE(abs, 3), SLICE(scat, 3)), // North
-      Wall<T>({w, 0, 0}, {0, 0, -1}, SLICE(abs, 4), SLICE(scat, 3)), // Floor
-      Wall<T>({w, 0, h}, {0, 0, 1}, SLICE(abs, 5), SLICE(scat, 3)) // Ceiling
+      Wall<T>({w, 0, 0}, {0, 0, -1}, SLICE(abs, 4), SLICE(scat, 4)), // Floor
+      Wall<T>({w, 0, h}, {0, 0, 1}, SLICE(abs, 5), SLICE(scat, 5)) // Ceiling
   };
 #undef SLICE
 }
@@ -137,7 +102,7 @@ const std::array<Wall<T>, 6> make_room(
 /// so that it does hit one of the walls.
 /// See also:
 /// https://github.com/LCAV/pyroomacoustics/blob/df8af24c88a87b5d51c6123087cd3cd2d361286a/pyroomacoustics/libroom_src/room.cpp#L609-L716
-template <typename scalar_t, unsigned int Dim>
+template <typename scalar_t>
 std::tuple<torch::Tensor, int, scalar_t> find_collision_wall(
     const torch::Tensor& room,
     const torch::Tensor& origin,
@@ -147,22 +112,16 @@ std::tuple<torch::Tensor, int, scalar_t> find_collision_wall(
 #define INSIDE(x, y) (BOOL(-EPS < (x)) && BOOL((x) < (y + EPS)))
 
   TORCH_INTERNAL_ASSERT_DEBUG_ONLY(
-      Dim == room.size(0),
-      "Expected room to be ",
-      Dim,
-      " dimension, but received ",
+      3 == room.size(0),
+      "Expected room to be 3 dimension, but received ",
       room.sizes());
   TORCH_INTERNAL_ASSERT_DEBUG_ONLY(
-      Dim == origin.size(0),
-      "Expected origin to be ",
-      Dim,
-      " dimension, but received ",
+      3 == origin.size(0),
+      "Expected origin to be 3 dimension, but received ",
       origin.sizes());
   TORCH_INTERNAL_ASSERT_DEBUG_ONLY(
-      Dim == direction.size(0),
-      "Expected direction to be ",
-      Dim,
-      " dimension, but received ",
+      3 == direction.size(0),
+      "Expected direction to be 3 dimension, but received ",
       direction.sizes());
   TORCH_INTERNAL_ASSERT_DEBUG_ONLY(
       BOOL(room > 0), "Room size should be greater than zero. Found: ", room);
@@ -174,7 +133,7 @@ std::tuple<torch::Tensor, int, scalar_t> find_collision_wall(
       room);
 
   // i is the coordinate in the collision is searched.
-  for (unsigned int i = 0; i < Dim; ++i) {
+  for (unsigned int i = 0; i < 3; ++i) {
     auto dir0 = SCALAR(direction[i]);
     auto abs_dir0 = std::abs(dir0);
 

torchaudio/prototype/functional/_rir.py

@@ -110,11 +110,10 @@ def _frac_delay(delay: torch.Tensor, delay_i: torch.Tensor, delay_filter_length:
     return torch.special.sinc(n - delay) * _hann(n - delay, 2 * pad)
 
 
-def _adjust_coeff(dim: int, coeffs: Union[float, torch.Tensor], name: str) -> torch.Tensor:
+def _adjust_coeff(coeffs: Union[float, torch.Tensor], name: str) -> torch.Tensor:
     """Validates and converts absorption or scattering parameters to a tensor with appropriate shape
 
     Args:
-        dim (int): The dimension of the simulation. 2 or 3.
         coeff (float or torch.Tensor): The absorption coefficients of wall materials.
 
             If the dtype is ``float``, the absorption coefficient is identical for all walls and
@@ -129,10 +128,10 @@ def _adjust_coeff(dim: int, coeffs: Union[float, torch.Tensor], name: str) -> to
 
     Returns:
         (torch.Tensor): The expanded coefficient.
-            The shape is `(1, 2*dim)` for single octave band case, and
-            `(7, 2*dim)` for multi octave band case.
+            The shape is `(1, 6)` for single octave band case, and
+            `(7, 6)` for multi octave band case.
     """
-    num_walls = 2 * dim
+    num_walls = 6
     if isinstance(coeffs, float):
         return torch.full((1, num_walls), coeffs)
     if isinstance(coeffs, Tensor):
@@ -154,7 +153,6 @@ def _adjust_coeff(dim: int, coeffs: Union[float, torch.Tensor], name: str) -> to
 
 
 def _validate_inputs(
-    dim: int,
     room: torch.Tensor,
     source: torch.Tensor,
     mic_array: torch.Tensor,
@@ -162,17 +160,16 @@ def _validate_inputs(
     """Validate dimensions of input arguments, and normalize different kinds of absorption into the same dimension.
 
     Args:
-        dim (int): The dimension of the simulation. 2 or 3.
         room (torch.Tensor): The size of the room. width, length (and height)
         source (torch.Tensor): Sound source coordinates. Tensor with dimensions `(dim,)`.
         mic_array (torch.Tensor): Microphone coordinates. Tensor with dimensions `(channel, dim)`.
     """
-    if not (room.ndim == 1 and room.numel() == dim):
-        raise ValueError(f"`room` must be a 1D Tensor with {dim} elements. Found {room.shape}.")
-    if not (source.ndim == 1 and source.numel() == dim):
-        raise ValueError(f"`source` must be 1D Tensor with {dim} elements. Found {source.shape}.")
-    if not (mic_array.ndim == 2 and mic_array.shape[1] == dim):
-        raise ValueError(f"mic_array must be a 2D Tensor with shape (num_channels, {dim}). Found {mic_array.shape}.")
+    if not (room.ndim == 1 and room.numel() == 3):
+        raise ValueError(f"`room` must be a 1D Tensor with 3 elements. Found {room.shape}.")
+    if not (source.ndim == 1 and source.numel() == 3):
+        raise ValueError(f"`source` must be 1D Tensor with 3 elements. Found {source.shape}.")
+    if not (mic_array.ndim == 2 and mic_array.shape[1] == 3):
+        raise ValueError(f"mic_array must be a 2D Tensor with shape (num_channels, 3). Found {mic_array.shape}.")
 
 
 def simulate_rir_ism(
@@ -231,8 +228,8 @@ def simulate_rir_ism(
         of octave bands are fixed to ``[125.0, 250.0, 500.0, 1000.0, 2000.0, 4000.0, 8000.0]``.
         Users need to tune the values of ``absorption`` to the corresponding frequencies.
     """
-    _validate_inputs(3, room, source, mic_array)
-    absorption = _adjust_coeff(3, absorption, "absorption")
+    _validate_inputs(room, source, mic_array)
+    absorption = _adjust_coeff(absorption, "absorption")
     img_location, att = _compute_image_sources(room, source, max_order, absorption)
 
     # compute distances between image sources and microphones