Skip to content

GitLab

Unverified Commit d9942bae authored by moto's avatar moto Committed by GitHub
Browse files

Remove 2d impl from ray tracing (#3633) 

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

We only add 3D RIR.
parent 47f502a6
Show whitespace changes
Inline Side-by-side
Showing with 56 additions and 195 deletions
test/cpp/rir/wall_collision.cpp
View file @ d9942bae
......@@ -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
View file @ d9942bae
......@@ -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);
}
}
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,15 +127,6 @@ 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
......@@ -166,7 +150,6 @@ class RayTracer {
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 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);
source, num_rays, time_thres, energy_thres, sound_speed, num_bins);
});
}
default:
TORCH_CHECK(false, "Only 2D and 3D are supported.");
}
}
TORCH_LIBRARY_IMPL(torchaudio, CPU, m) {
......
torchaudio/csrc/rir/wall.h
View file @ d9942bae
......@@ -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
View file @ d9942bae
......@@ -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
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment