Add wall implementation for RIR ray tracing (#3612)

Extracted from #3604

Add Wall helper class and C++ unit test

.github/scripts/unittest-linux/install.sh

@@ -62,6 +62,7 @@ printf "Installing PyTorch with %s\n" "${cudatoolkit}"
 conda install --quiet -y ninja cmake
 
 printf "* Installing torchaudio\n"
+export BUILD_CPP_TEST=1
 python setup.py install
 
 # 3. Install Test tools

.github/scripts/unittest-linux/run_test.sh

@@ -22,6 +22,13 @@ if [[ "${CUDA_TESTS_ONLY}" = "1" ]]; then
   args+=('-k' 'cuda or gpu')
 fi
 
-cd test
-pytest "${args[@]}" torchaudio_unittest
-coverage html
+(
+    cd build/temp*/test/cpp
+    ctest
+)
+
+(
+    cd test
+    pytest "${args[@]}" torchaudio_unittest
+    coverage html
+)

.github/scripts/unittest-windows/install.sh

@@ -47,6 +47,7 @@ printf "* Installing fsspec\n"
 pip install --trusted-host pypi.org --trusted-host files.pythonhosted.org fsspec
 
 printf "* Installing torchaudio\n"
+export BUILD_CPP_TEST=1
 "$root_dir/packaging/vc_env_helper.bat" python setup.py install
 
 # 3. Install Test tools

.github/scripts/unittest-windows/run_test.sh

@@ -11,6 +11,13 @@ source "$this_dir/set_cuda_envs.sh"
 python -m torch.utils.collect_env
 env | grep TORCHAUDIO || true
 
-cd test
-pytest --cov=torchaudio --junitxml=${RUNNER_TEST_RESULTS_DIR}/junit.xml -v --durations 20 torchaudio_unittest
-coverage html
+(
+    cd build/temp*/test/cpp
+    ctest
+)
+
+(
+    cd test
+    pytest --cov=torchaudio --junitxml=${RUNNER_TEST_RESULTS_DIR}/junit.xml -v --durations 20 torchaudio_unittest
+    coverage html
+)

CMakeLists.txt

@@ -185,3 +185,6 @@ if (BUILD_CUDA_CTC_DECODER)
   endif()
   add_subdirectory(torchaudio/csrc/cuctc)
 endif()
+if (BUILD_CPP_TEST)
+  add_subdirectory(test/cpp)
+endif()

test/cpp/CMakeLists.txt

+include(FetchContent)
+FetchContent_Declare(
+  googletest
+  URL https://github.com/google/googletest/archive/refs/tags/v1.14.0.zip
+)
+
+# For Windows: Prevent overriding the parent project's compiler/linker settings
+set(gtest_force_shared_crt ON CACHE BOOL "" FORCE)
+FetchContent_MakeAvailable(googletest)
+
+enable_testing()
+
+add_executable(
+  wall_collision
+  rir/wall_collision.cpp
+)
+target_link_libraries(
+  wall_collision
+  torch
+  GTest::gtest_main
+)
+target_include_directories(
+  wall_collision
+  PRIVATE
+  "${PROJECT_SOURCE_DIR}"
+)
+add_test(NAME wall_collision_test COMMAND wall_collision)

test/cpp/rir/wall_collision.cpp

+#include <gtest/gtest.h>
+#include <torchaudio/csrc/rir/wall.h>
+
+using namespace torchaudio::rir;
+
+struct CollisionTestParam {
+  // Input
+  torch::Tensor origin;
+  torch::Tensor direction;
+  // Expected
+  torch::Tensor hit_point;
+  int next_wall_index;
+  float hit_distance;
+};
+
+CollisionTestParam par(
+    torch::ArrayRef<float> origin,
+    torch::ArrayRef<float> direction,
+    torch::ArrayRef<float> hit_point,
+    int next_wall_index,
+    float hit_distance) {
+  auto dir = torch::tensor(direction);
+  return {
+      torch::tensor(origin),
+      dir / dir.norm(),
+      torch::tensor(hit_point),
+      next_wall_index,
+      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
+//////////////////////////////////////////////////////////////////////////////
+
+class Simple3DRoomCollisionTest
+    : public ::testing::TestWithParam<CollisionTestParam> {};
+
+TEST_P(Simple3DRoomCollisionTest, CollisionTest3D) {
+  //  y                       z
+  //  ^                       ^
+  //  |        3              |      y
+  //  |     ______            |     /
+  //  |    |      |           |    /
+  //  |  0 |      | 1         |  ______
+  //  |    |______|           | /     /  4: floor, 5: ceiling
+  //  |        2              |/     /
+  // -+----------------> x   -+--------------> x
+  //
+  auto room = torch::tensor({1, 1, 1});
+
+  auto param = GetParam();
+  auto [hit_point, next_wall_index, hit_distance] =
+      find_collision_wall<float, 3>(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));
+}
+
+INSTANTIATE_TEST_CASE_P(
+    Collision3DTests,
+    Simple3DRoomCollisionTest,
+    ::testing::Values(
+        // From 0
+        par({0, .5, .5}, {1.0, 0.0, 0.0}, {1., .5, .5}, 1, 1.0),
+        par({0, .5, .5}, {1.0, -1., 0.0}, {.5, .0, .5}, 2, ISQRT2),
+        par({0, .5, .5}, {1.0, 1.0, 0.0}, {.5, 1., .5}, 3, ISQRT2),
+        par({0, .5, .5}, {1.0, 0.0, -1.}, {.5, .5, .0}, 4, ISQRT2),
+        par({0, .5, .5}, {1.0, 0.0, 1.0}, {.5, .5, 1.}, 5, ISQRT2),
+        // From 1
+        par({1, .5, .5}, {-1., 0.0, 0.0}, {.0, .5, .5}, 0, 1.0),
+        par({1, .5, .5}, {-1., -1., 0.0}, {.5, .0, .5}, 2, ISQRT2),
+        par({1, .5, .5}, {-1., 1.0, 0.0}, {.5, 1., .5}, 3, ISQRT2),
+        par({1, .5, .5}, {-1., 0.0, -1.}, {.5, .5, .0}, 4, ISQRT2),
+        par({1, .5, .5}, {-1., 0.0, 1.0}, {.5, .5, 1.}, 5, ISQRT2),
+        // From 2
+        par({.5, 0, .5}, {-1., 1.0, 0.0}, {.0, .5, .5}, 0, ISQRT2),
+        par({.5, 0, .5}, {1.0, 1.0, 0.0}, {1., .5, .5}, 1, ISQRT2),
+        par({.5, 0, .5}, {0.0, 1.0, 0.0}, {.5, 1., .5}, 3, 1.0),
+        par({.5, 0, .5}, {0.0, 1.0, -1.}, {.5, .5, .0}, 4, ISQRT2),
+        par({.5, 0, .5}, {0.0, 1.0, 1.0}, {.5, .5, 1.}, 5, ISQRT2),
+        // From 3
+        par({.5, 1, .5}, {-1., -1., 0.0}, {.0, .5, .5}, 0, ISQRT2),
+        par({.5, 1, .5}, {1.0, -1., 0.0}, {1., .5, .5}, 1, ISQRT2),
+        par({.5, 1, .5}, {0.0, -1., 0.0}, {.5, .0, .5}, 2, 1.0),
+        par({.5, 1, .5}, {0.0, -1., -1.}, {.5, .5, .0}, 4, ISQRT2),
+        par({.5, 1, .5}, {0.0, -1., 1.0}, {.5, .5, 1.}, 5, ISQRT2),
+        // From 4
+        par({.5, .5, 0}, {-1., 0.0, 1.0}, {.0, .5, .5}, 0, ISQRT2),
+        par({.5, .5, 0}, {1.0, 0.0, 1.0}, {1., .5, .5}, 1, ISQRT2),
+        par({.5, .5, 0}, {0.0, -1., 1.0}, {.5, .0, .5}, 2, ISQRT2),
+        par({.5, .5, 0}, {0.0, 1.0, 1.0}, {.5, 1., .5}, 3, ISQRT2),
+        par({.5, .5, 0}, {0.0, 0.0, 1.0}, {.5, .5, 1.}, 5, 1.0),
+        // From 5
+        par({.5, .5, 1}, {-1., 0.0, -1.}, {.0, .5, .5}, 0, ISQRT2),
+        par({.5, .5, 1}, {1.0, 0.0, -1.}, {1., .5, .5}, 1, ISQRT2),
+        par({.5, .5, 1}, {0.0, -1., -1.}, {.5, .0, .5}, 2, ISQRT2),
+        par({.5, .5, 1}, {0.0, 1.0, -1.}, {.5, 1., .5}, 3, ISQRT2),
+        par({.5, .5, 1}, {0.0, 0.0, -1.}, {.5, .5, .0}, 4, 1.0)));

tools/setup_helpers/extension.py

@@ -33,6 +33,7 @@ def _get_build(var, default=False):
     return False
 
 
+_BUILD_CPP_TEST = _get_build("BUILD_CPP_TEST", False)
 _BUILD_SOX = False if platform.system() == "Windows" else _get_build("BUILD_SOX", True)
 _BUILD_RIR = _get_build("BUILD_RIR", True)
 _BUILD_RNNT = _get_build("BUILD_RNNT", True)
@@ -127,6 +128,7 @@ class CMakeBuild(build_ext):
             f"-DCMAKE_INSTALL_PREFIX={extdir}",
             "-DCMAKE_VERBOSE_MAKEFILE=ON",
             f"-DPython_INCLUDE_DIR={distutils.sysconfig.get_python_inc()}",
+            f"-DBUILD_CPP_TEST={'ON' if _BUILD_CPP_TEST else 'OFF'}",
             f"-DBUILD_SOX:BOOL={'ON' if _BUILD_SOX else 'OFF'}",
             f"-DBUILD_RIR:BOOL={'ON' if _BUILD_RIR else 'OFF'}",
             f"-DBUILD_RNNT:BOOL={'ON' if _BUILD_RNNT else 'OFF'}",

torchaudio/csrc/rir/wall.h

+#include <torch/types.h>
+
+#define EPS ((scalar_t)(1e-5))
+#define SCALAR(x) ((x).template item<scalar_t>())
+
+namespace torchaudio {
+namespace rir {
+
+////////////////////////////////////////////////////////////////////////////////
+// Basic Wall implementation
+////////////////////////////////////////////////////////////////////////////////
+
+/// Wall helper class. A wall records its own absorption, reflection and
+/// scattering coefficient, and exposes a few methods for geometrical operations
+/// (e.g. reflection of a ray)
+template <typename scalar_t>
+struct Wall {
+  const torch::Tensor origin;
+  const torch::Tensor normal;
+  const torch::Tensor scattering;
+
+  const torch::Tensor reflection;
+
+  Wall(
+      const torch::ArrayRef<scalar_t>& origin,
+      const torch::ArrayRef<scalar_t>& normal,
+      const torch::Tensor& absorption,
+      const torch::Tensor& scattering)
+      : origin(torch::tensor(origin)),
+        normal(torch::tensor(normal)),
+        scattering(scattering),
+        reflection(1. - absorption) {}
+};
+
+/// Returns the side (-1, 1 or 0) on which a point lies w.r.t. the wall.
+template <typename scalar_t>
+int side(const Wall<scalar_t>& wall, const torch::Tensor& pos) {
+  auto dot = SCALAR((pos - wall.origin).dot(wall.normal));
+
+  if (dot > EPS) {
+    return 1;
+  } else if (dot < -EPS) {
+    return -1;
+  } else {
+    return 0;
+  }
+}
+
+/// Reflects a ray (dir) on the wall. Preserves norm of vector.
+template <typename scalar_t>
+torch::Tensor reflect(const Wall<scalar_t>& wall, const torch::Tensor& dir) {
+  return dir - wall.normal * 2 * dir.dot(wall.normal);
+}
+
+/// Returns the cosine angle of a ray (dir) with the normal of the wall
+template <typename scalar_t>
+scalar_t cosine(const Wall<scalar_t>& wall, const torch::Tensor& dir) {
+  return SCALAR(dir.dot(wall.normal) / dir.norm());
+}
+
+////////////////////////////////////////////////////////////////////////////////
+// Room (multiple walls) and interactions
+////////////////////////////////////////////////////////////////////////////////
+
+/// Creates a shoebox room consists of multiple walls.
+/// Normals are vectors facing *outwards* the room, and origins are arbitrary
+/// corners of each wall.
+///
+/// Note:
+/// The wall has to be ordered in the following way:
+/// - parallel walls are next (W/E, S/N, and F/C)
+/// - The one closer to the origin must come first. (W -> E, S -> N, F -> C)
+/// - The order of wall pair must be W/E, S/N, then F/C because
+///   `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::vector<Wall<T>> 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::vector<Wall<T>> make_room(
+    const T w,
+    const T l,
+    const T h,
+    const torch::Tensor& abs,
+    const torch::Tensor& scat) {
+  using namespace torch::indexing;
+#define SLICE(x, i) x.index({Slice(), i})
+  return {
+      Wall<T>({0, l, 0}, {-1, 0, 0}, SLICE(abs, 0), SLICE(scat, 0)), // West
+      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
+  };
+#undef SLICE
+}
+
+/// Find a wall that the given ray hits.
+/// The room is assumed to be shoebox room and the walls are constructed
+/// in the order used in `make_room`.
+/// The room is shoebox-shape and the ray travels infinite distance
+/// 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>
+std::tuple<torch::Tensor, int, scalar_t> find_collision_wall(
+    const torch::Tensor& room,
+    const torch::Tensor& origin,
+    const torch::Tensor& direction // Unit-vector
+) {
+#define BOOL(x) torch::all(x).template item<bool>()
+#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 ",
+      room.sizes());
+  TORCH_INTERNAL_ASSERT_DEBUG_ONLY(
+      Dim == origin.size(0),
+      "Expected origin to be ",
+      Dim,
+      " dimension, but received ",
+      origin.sizes());
+  TORCH_INTERNAL_ASSERT_DEBUG_ONLY(
+      Dim == direction.size(0),
+      "Expected direction to be ",
+      Dim,
+      " dimension, but received ",
+      direction.sizes());
+  TORCH_INTERNAL_ASSERT_DEBUG_ONLY(
+      BOOL(room > 0), "Room size should be greater than zero. Found: ", room);
+  TORCH_INTERNAL_ASSERT_DEBUG_ONLY(
+      INSIDE(origin, room),
+      "The origin of ray must be inside the room. Origin: ",
+      origin,
+      ", room: ",
+      room);
+
+  // i is the coordinate in the collision is searched.
+  for (unsigned int i = 0; i < Dim; ++i) {
+    auto dir0 = SCALAR(direction[i]);
+    auto abs_dir0 = std::abs(dir0);
+
+    // If the ray is almost parallel to a plane, then we delegate the
+    // computation to the other planes.
+    if (abs_dir0 < EPS) {
+      continue;
+    }
+
+    // Check the distance to the facing wall along the coordinate.
+    scalar_t distance = (dir0 < 0.)
+        ? SCALAR(origin[i]) // Going towards origin
+        : SCALAR(room[i] - origin[i]); // Going away from origin
+    auto ratio = distance / abs_dir0;
+    int i_increment = dir0 > 0.;
+
+    // Compute the intersection of ray and the wall
+    auto intersection = origin + ratio * direction;
+
+    // The intersection can be within the room or outside.
+    // If it's inside, the collision point is found.
+    //      ^
+    //      |           |   Not Good
+    //   ---+-----------+---x----
+    //      |           |  /
+    //      |           | /
+    //      |           |/
+    //      |           x  Found
+    //      |          /|
+    //      |         / |
+    //      |        o  |
+    //      |           |
+    //   ---+-----------+-------->
+    //     O|           |
+    //
+
+    if (INSIDE(intersection, room)) {
+      int i_wall = 2 * i + i_increment;
+      auto dist = SCALAR((intersection - origin).norm());
+      return std::make_tuple(intersection, i_wall, dist);
+    }
+  }
+  // This should not happen
+  TORCH_INTERNAL_ASSERT(
+      false,
+      "Failed to find the intersection. room: ",
+      room,
+      " origin: ",
+      origin,
+      " direction: ",
+      direction);
+#undef INSIDE
+#undef BOOL
+}
+} // namespace rir
+} // namespace torchaudio
+
+#undef EPS
+#undef SCALAR