[WIP] Source/microphone directivity for ray tracing

.gitmodules

@@ -1,3 +1,6 @@
 [submodule "pyroomacoustics/libroom_src/ext/eigen"]
 	path = pyroomacoustics/libroom_src/ext/eigen
 	url = https://github.com/eigenteam/eigen-git-mirror.git
+[submodule "pyroomacoustics/libroom_src/ext/nanoflann"]
+	path = pyroomacoustics/libroom_src/ext/nanoflann
+	url = https://github.com/jlblancoc/nanoflann.git

MANIFEST.in

@@ -11,9 +11,11 @@ include pyroomacoustics/libroom_src/*.h
 include pyroomacoustics/libroom_src/*.hpp
 include pyroomacoustics/libroom_src/*.cpp
 
-# The compiled extension code rely on Eigen, which is included
+# The compiled extension code rely on Eigen and nanoflann,
+# which are included
 include pyroomacoustics/libroom_src/ext/eigen/COPYING.*
 graft pyroomacoustics/libroom_src/ext/eigen/Eigen
+graft pyroomacoustics/libroom_src/ext/nanoflann/include
 
 include pyproject.toml
 include requirements.txt

examples/plot_directivity_2D.py

@@ -0,0 +1,33 @@
+import matplotlib.pyplot as plt
+import numpy as np
+
+from pyroomacoustics import dB
+from pyroomacoustics.directivities import (
+    CardioidFamily,
+    DirectionVector,
+    DirectivityPattern,
+)
+from pyroomacoustics.doa import spher2cart
+
+ORIENTATION = DirectionVector(azimuth=0, colatitude=90, degrees=True)
+LOWER_GAIN = -20
+
+# plot each directivity
+angles = np.linspace(start=0, stop=360, num=361, endpoint=True)
+angles = np.radians(angles)
+
+# plot each pattern
+fig = plt.figure()
+ax = plt.subplot(111, projection="polar")
+for pattern in DirectivityPattern:
+
+    dir_obj = CardioidFamily(orientation=ORIENTATION, pattern_enum=pattern)
+    resp = dir_obj.get_response(angles, magnitude=True, degrees=False)
+    resp_db = dB(np.array(resp))
+    ax.plot(angles, resp_db, label=pattern.name)
+
+plt.legend(bbox_to_anchor=(1, 1))
+plt.ylim([LOWER_GAIN, 0])
+ax.yaxis.set_ticks(np.arange(start=LOWER_GAIN, stop=5, step=5))
+plt.tight_layout()
+plt.show()

examples/plot_directivity_3D.py

@@ -0,0 +1,21 @@
+import matplotlib.pyplot as plt
+import numpy as np
+
+from pyroomacoustics.directivities import (
+    CardioidFamily,
+    DirectionVector,
+    DirectivityPattern,
+)
+
+PATTERN = DirectivityPattern.HYPERCARDIOID
+ORIENTATION = DirectionVector(azimuth=0, colatitude=45, degrees=True)
+
+# create cardioid object
+dir_obj = CardioidFamily(orientation=ORIENTATION, pattern_enum=PATTERN)
+
+# plot
+azimuth = np.linspace(start=0, stop=360, num=361, endpoint=True)
+colatitude = np.linspace(start=0, stop=180, num=180, endpoint=True)
+# colatitude = None   # for 2D plot
+dir_obj.plot_response(azimuth=azimuth, colatitude=colatitude, degrees=True)
+plt.show()

examples/plot_directivity_one_shot.py

@@ -0,0 +1,60 @@
+import matplotlib.pyplot as plt
+import numpy as np
+
+from pyroomacoustics import all_combinations, dB
+from pyroomacoustics.directivities import (
+    CardioidFamily,
+    DirectionVector,
+    DirectivityPattern,
+    cardioid_func,
+)
+from pyroomacoustics.doa import spher2cart
+
+ORIENTATION = DirectionVector(azimuth=0, colatitude=90, degrees=True)
+azimuth = np.radians(np.linspace(start=0, stop=360, num=361, endpoint=True))
+colatitude = np.radians(np.linspace(start=0, stop=180, num=180, endpoint=True))
+LOWER_GAIN = -40
+
+""" 2D """
+# get cartesian coordinates
+cart = spher2cart(azimuth=azimuth)
+direction = spher2cart(azimuth=225, degrees=True)
+
+# compute response
+resp = cardioid_func(x=cart, direction=direction, coef=0.5, magnitude=True)
+resp_db = dB(np.array(resp))
+
+# plot
+plt.figure()
+plt.polar(azimuth, resp_db)
+plt.ylim([LOWER_GAIN, 0])
+ax = plt.gca()
+ax.yaxis.set_ticks(np.arange(start=LOWER_GAIN, stop=5, step=10))
+plt.tight_layout()
+
+""" 3D """
+# get cartesian coordinates
+spher_coord = all_combinations(azimuth, colatitude)
+cart = spher2cart(azimuth=spher_coord[:, 0], colatitude=spher_coord[:, 1])
+direction = spher2cart(azimuth=0, colatitude=45, degrees=True)
+
+# compute response
+resp = cardioid_func(x=cart, direction=direction, coef=0.25, magnitude=True)
+
+# plot (surface plot)
+fig = plt.figure()
+RESP_2D = resp.reshape(len(azimuth), len(colatitude))
+AZI, COL = np.meshgrid(azimuth, colatitude)
+X = RESP_2D.T * np.sin(COL) * np.cos(AZI)
+Y = RESP_2D.T * np.sin(COL) * np.sin(AZI)
+Z = RESP_2D.T * np.cos(COL)
+ax = fig.add_subplot(1, 1, 1, projection="3d")
+ax.plot_surface(X, Y, Z)
+ax.set_xlabel("x")
+ax.set_ylabel("y")
+ax.set_zlabel("z")
+ax.set_xlim([-1, 1])
+ax.set_ylim([-1, 1])
+ax.set_zlim([-1, 1])
+
+plt.show()

pyroomacoustics/__init__.py

@@ -122,7 +122,7 @@
 
 from . import adaptive, bss, datasets, denoise, doa, experimental
 from . import libroom as libroom
-from . import phase, transform
+from . import phase, random, transform
 from .acoustics import *
 from .beamforming import *
 from .directivities import *

pyroomacoustics/directivities/__init__.py

@@ -84,6 +84,7 @@
 from .analytic import (
     Cardioid,
     CardioidFamily,
+    CardioidFamilySampler,
     FigureEight,
     HyperCardioid,
     Omnidirectional,

pyroomacoustics/directivities/analytic.py

@@ -59,6 +59,7 @@
 """
 import numpy as np
 
+from .. import random
 from ..doa import spher2cart
 from ..utilities import all_combinations, requires_matplotlib
 from .base import Directivity
@@ -122,6 +123,12 @@ def __init__(self, orientation, p, gain=1.0):
         self._gain = gain
         self._pattern_name = f"cardioid family, p={self._p}"
 
+        # this is the object that will allow to sample rays according to the
+        # distribution corresponding to the source directivity
+        self._ray_sampler = CardioidFamilySampler(
+            loc=self._orientation.unit_vector, p=self._p
+        )
+
     @property
     def is_impulse_response(self):
         # this is not an impulse response, do not make docstring to avoid clutter in the
@@ -200,6 +207,9 @@ def get_response(
             else:
                 return resp
 
+    def sample_rays(self, n_rays):
+        return self._ray_sampler(n_rays).T
+
     @requires_matplotlib
     def plot_response(
         self, azimuth, colatitude=None, degrees=True, ax=None, offset=None
@@ -383,6 +393,39 @@ def __init__(self, gain=1.0):
         self._pattern_name = "omni"
 
 
+class CardioidFamilySampler(random.sampler.DirectionalSampler):
+    """
+    This object draws samples from a cardioid shaped distribution on the sphere
+
+    Parameters
+    ----------
+    loc: array_like
+        The unit vector pointing in the main direction of the cardioid
+    p: float
+        Parameter of the cardioid pattern. A value of 0 corresponds to a
+        figure-eight pattern, 0.5 to a cardioid pattern, and 1 to an omni
+        pattern
+        The parameter must be between 0 and 1
+    """
+
+    def __init__(self, loc=None, p=0.5):
+        super().__init__(loc=loc)
+        self._coeff = p
+
+    def _pattern(self, x):
+        response = cardioid_func(
+            x.T,
+            direction=self._loc,
+            p=self._coeff,
+            gain=1.0,
+            normalize=False,
+            magnitude=True,
+        )
+        # The number of rays needs to be proportional to the
+        # response energy.
+        return response**2
+
+
 def cardioid_func(x, direction, p, gain=1.0, normalize=True, magnitude=False):
     """
     One-shot function for computing cardioid response.

pyroomacoustics/directivities/base.py

@@ -92,3 +92,21 @@ def get_response(
             Response at provided angles.
         """
         raise NotImplementedError
+
+    @abc.abstractmethod
+    def sample_rays(self, n_rays):
+        """
+        This method samples unit vectors from the sphere according to
+        the distribution of the source
+
+        Parameters
+        ----------
+        n_rays: int
+            The number of rays to sample
+
+        Returns
+        -------
+        ray_directions: numpy.ndarray, shape (n_dim, n_rays)
+            An array containing the unit vectors in its columns
+        """
+        raise NotImplementedError

pyroomacoustics/directivities/measured.py

@@ -75,6 +75,7 @@
 from scipy.interpolate import griddata
 from scipy.spatial import cKDTree
 
+from .. import random
 from ..datasets import SOFADatabase
 from ..doa import Grid, GridSphere, cart2spher, fibonacci_spherical_sampling, spher2cart
 from ..utilities import requires_matplotlib
@@ -84,6 +85,34 @@
 from .sofa import open_sofa_file
 
 
+class MeasuredDirectivitySampler(random.sampler.DirectionalSampler):
+    """
+    This object draws samples from the distribution defined by the energy
+    of the measured directional response object.
+
+    Parameters
+    ----------
+    loc: array_like
+        The unit vector pointing in the main direction of the cardioid
+    p: float
+        Parameter of the cardioid pattern. A value of 0 corresponds to a
+        figure-eight pattern, 0.5 to a cardioid pattern, and 1 to an omni
+        pattern
+        The parameter must be between 0 and 1
+    """
+
+    def __init__(self, kdtree, energy):
+        super().__init__()
+        self._kdtree = kdtree
+        # Normalize to maximum energy 1 because that is also the
+        # maximum value of the proposal unnormalized uniform distribution.
+        self._energy = energy / energy.max()
+
+    def _pattern(self, x):
+        _, index = self._kdtree.query(x)
+        return self._energy[index]
+
+
 class MeasuredDirectivity(Directivity):
     """
     A class to store directivity patterns obtained by measurements.
@@ -144,6 +173,10 @@ def set_orientation(self, orientation):
         # create the kd-tree
         self._kdtree = cKDTree(self._grid.cartesian.T)
 
+        # create the ray sampler
+        ir_energy = np.square(self._irs).mean(axis=-1)
+        self._ray_sampler = MeasuredDirectivitySampler(self._kdtree, ir_energy)
+
     def get_response(
         self, azimuth, colatitude=None, magnitude=False, frequency=None, degrees=True
     ):
@@ -172,6 +205,9 @@ def get_response(
         _, index = self._kdtree.query(cart.T)
         return self._irs[index, :]
 
+    def sample_rays(self, n_rays):
+        return self._ray_sampler(n_rays).T
+
     @requires_matplotlib
     def plot(self, freq_bin=0, n_grid=100, ax=None, depth=False, offset=None):
         """
@@ -238,7 +274,7 @@ def plot(self, freq_bin=0, n_grid=100, ax=None, depth=False, offset=None):
             Y *= V
             Z *= V
 
-        surf = ax.plot_surface(
+        _ = ax.plot_surface(
             X, Y, Z, facecolors=cmap.to_rgba(V), linewidth=0, antialiased=False
         )
 

pyroomacoustics/doa/__init__.py

@@ -112,6 +112,7 @@
 from .doa import *
 from .frida import *
 from .grid import *
+from .histogram import SphericalHistogram
 from .music import *
 from .normmusic import *
 from .srp import *

pyroomacoustics/doa/grid.py

@@ -209,7 +209,7 @@ def __init__(
             # If no list was provided, samples points on the sphere
             # as uniformly as possible
 
-            self.x, self.y, self.z = fibonacci_spherical_sampling(n_points)
+            self.x[:], self.y[:], self.z[:] = fibonacci_spherical_sampling(n_points)
 
             # Create convenient arrays
             # to access both in cartesian and spherical coordinates
@@ -389,10 +389,7 @@ def plot(
     def plot_old(self, plot_points=False, mark_peaks=0):
         """Plot the points on the sphere with their values"""
 
-        from scipy import rand
-
         try:
-            import matplotlib.colors as colors
             import matplotlib.pyplot as plt
 
             # from mpl_toolkits.mplot3d import Axes3D