<commit_before>"""Compute driving functions for various systems."""

import numpy as np
from numpy.core.umath_tests import inner1d  # element-wise inner product
from scipy.special import sph_jn, sph_yn, jn, hankel2
from .. import util
from .. import defs


def wfs_2d_line(omega, x0, n0, xs, c=None):
    """Line source by 2-dimensional WFS.

    ::

        D(x0,k) = j/2 k (x0-xs) n0 / |x0-xs| * H1(k |x0-xs|)

    """
    x0 = np.asarray(x0)
    n0 = np.asarray(n0)
    xs = np.squeeze(np.asarray(xs))
    k = util.wavenumber(omega, c)
    ds = x0 - xs
    r = np.linalg.norm(ds, axis=1)
    return -1j/2 * k * inner1d(ds, n0) / r * hankel2(1, k * r)


def _wfs_point(omega, x0, n0, xs, c=None):
    """Point source by two- or three-dimensional WFS.

    ::

                       (x0-xs) n0
        D(x0,k) = j k ------------- e^(-j k |x0-xs|)
                      |x0-xs|^(3/2)

    """
    x0 = np.asarray(x0)
    n0 = np.asarray(n0)
    xs = np.squeeze(np.asarray(xs))
    k = util.wavenumber(omega, c)
    ds = x0 - xs
    r = np.linalg.norm(ds, axis=1)
    return 1j * k * inner1d(ds, n0) / r ** (3 / 2) * np.exp(-1j * k * r)


wfs_2d_point = _wfs_point


def wfs_25d_point(omega, x0, n0, xs, xref=[0, 0, 0], c=None, omalias=None):
    """Point source by 2.5-dimensional WFS.

    ::

                    ____________   (x0-xs) n0
        D(x0,k) = \|j k |xref-x0| ------------- e^(-j k |x0-xs|)
                                  |x0-xs|^(3/2)

    """
    x0 = np.asarray(x0)
    n0 = np.asarray(n0)
    xs = np.squeeze(np.asarray(xs))
    xref = np.squeeze(np.asarray(xref))
    k = util.wavenumber(omega, c)
    ds = x0 - xs
    r = np.linalg.norm(ds, axis=1)

    return wfs_25d_preeq(omega, omalias, c) * \
        np.sqrt(np.linalg.norm(xref - x0)) * inner1d(ds, n0) / \
        r ** (3 / 2) * np.exp(-1j * k * r)


wfs_3d_point = _wfs_point


def _wfs_plane(omega, x0, n0, n=[0, 1, 0], c=None):
    """Plane wave by two- or three-dimensional WFS.

    Eq.(17) from [Spors et al, 2008]::

        D(x0,k) =  j k n n0  e^(-j k n x0)

    """
    x0 = np.asarray(x0)
    n0 = np.asarray(n0)
    n = np.squeeze(np.asarray(n))
    k = util.wavenumber(omega, c)
    return 2j * k * np.inner(n, n0) * np.exp(-1j * k * np.inner(n, x0))


wfs_2d_plane = _wfs_plane


def wfs_25d_plane(omega, x0, n0, n=[0, 1, 0], xref=[0, 0, 0], c=None,
                  omalias=None):
    """Plane wave by 2.5-dimensional WFS.

    ::

                         ____________
        D_2.5D(x0,w) = \|j k |xref-x0| n n0 e^(-j k n x0)

    """
    x0 = np.asarray(x0)
    n0 = np.asarray(n0)
    n = np.squeeze(np.asarray(n))
    xref = np.squeeze(np.asarray(xref))
    k = util.wavenumber(omega, c)
    return wfs_25d_preeq(omega, omalias, c) * \
        np.sqrt(2*np.pi * np.linalg.norm(xref - x0)) * \
        np.inner(n, n0) * np.exp(-1j * k * np.inner(n, x0))


wfs_3d_plane = _wfs_plane


def wfs_25d_preeq(omega, omalias, c):
    """Preqeualization for 2.5D WFS."""
    if omalias is None:
        return np.sqrt(1j * util.wavenumber(omega, c))
    else:
        if omega <= omalias:
            return np.sqrt(1j * util.wavenumber(omega, c))
        else:
            return np.sqrt(1j * util.wavenumber(omalias, c))


def delay_3d_plane(omega, x0, n0, n=[0, 1, 0], c=None):
    """Plane wave by simple delay of secondary sources."""
    x0 = np.asarray(x0)
    n = np.squeeze(np.asarray(n))
    k = util.wavenumber(omega, c)
    return np.exp(-1j * k * np.inner(n, x0))


def source_selection_plane(n0, n):
    """Secondary source selection for a plane wave.

    Eq.(13) from [Spors et al, 2008]

    """
    n0 = np.asarray(n0)
    n = np.squeeze(np.asarray(n))
    return np.inner(n, n0) >= defs.selection_tolerance


def source_selection_point(n0, x0, xs):
    """Secondary source selection for a point source.

    Eq.(15) from [Spors et al, 2008]

    """
    n0 = np.asarray(n0)
    x0 = np.asarray(x0)
    xs = np.squeeze(np.asarray(xs))
    ds = x0 - xs
    return inner1d(ds, n0) >= defs.selection_tolerance


def source_selection_all(N):
    """Select all secondary sources."""
    return np.ones(N) >= 0


def nfchoa_2d_plane(omega, x0, r0, n=[0, 1, 0], c=None):
    """Plane wave by two-dimensional NFC-HOA.

    ::
                               __
                       2i     \        i^-m
        D(phi0,w) = - -----   /__   ----------  e^(i m (phi0-phi_pw))
                      pi r0 m=-N..N  (2)
                                    Hm  (w/c r0)
    """
    x0 = np.asarray(x0)
    k = util.wavenumber(omega, c)<commit_msg>Secondary source selection for line source<commit_after>"""Compute driving functions for various systems."""

import numpy as np
from numpy.core.umath_tests import inner1d  # element-wise inner product
from scipy.special import sph_jn, sph_yn, jn, hankel2
from .. import util
from .. import defs


def wfs_2d_line(omega, x0, n0, xs, c=None):
    """Line source by 2-dimensional WFS.

    ::

        D(x0,k) = j/2 k (x0-xs) n0 / |x0-xs| * H1(k |x0-xs|)

    """
    x0 = np.asarray(x0)
    n0 = np.asarray(n0)
    xs = np.squeeze(np.asarray(xs))
    k = util.wavenumber(omega, c)
    ds = x0 - xs
    r = np.linalg.norm(ds, axis=1)
    return -1j/2 * k * inner1d(ds, n0) / r * hankel2(1, k * r)


def _wfs_point(omega, x0, n0, xs, c=None):
    """Point source by two- or three-dimensional WFS.

    ::

                       (x0-xs) n0
        D(x0,k) = j k ------------- e^(-j k |x0-xs|)
                      |x0-xs|^(3/2)

    """
    x0 = np.asarray(x0)
    n0 = np.asarray(n0)
    xs = np.squeeze(np.asarray(xs))
    k = util.wavenumber(omega, c)
    ds = x0 - xs
    r = np.linalg.norm(ds, axis=1)
    return 1j * k * inner1d(ds, n0) / r ** (3 / 2) * np.exp(-1j * k * r)


wfs_2d_point = _wfs_point


def wfs_25d_point(omega, x0, n0, xs, xref=[0, 0, 0], c=None, omalias=None):
    """Point source by 2.5-dimensional WFS.

    ::

                    ____________   (x0-xs) n0
        D(x0,k) = \|j k |xref-x0| ------------- e^(-j k |x0-xs|)
                                  |x0-xs|^(3/2)

    """
    x0 = np.asarray(x0)
    n0 = np.asarray(n0)
    xs = np.squeeze(np.asarray(xs))
    xref = np.squeeze(np.asarray(xref))
    k = util.wavenumber(omega, c)
    ds = x0 - xs
    r = np.linalg.norm(ds, axis=1)

    return wfs_25d_preeq(omega, omalias, c) * \
        np.sqrt(np.linalg.norm(xref - x0)) * inner1d(ds, n0) / \
        r ** (3 / 2) * np.exp(-1j * k * r)


wfs_3d_point = _wfs_point


def _wfs_plane(omega, x0, n0, n=[0, 1, 0], c=None):
    """Plane wave by two- or three-dimensional WFS.

    Eq.(17) from [Spors et al, 2008]::

        D(x0,k) =  j k n n0  e^(-j k n x0)

    """
    x0 = np.asarray(x0)
    n0 = np.asarray(n0)
    n = np.squeeze(np.asarray(n))
    k = util.wavenumber(omega, c)
    return 2j * k * np.inner(n, n0) * np.exp(-1j * k * np.inner(n, x0))


wfs_2d_plane = _wfs_plane


def wfs_25d_plane(omega, x0, n0, n=[0, 1, 0], xref=[0, 0, 0], c=None,
                  omalias=None):
    """Plane wave by 2.5-dimensional WFS.

    ::

                         ____________
        D_2.5D(x0,w) = \|j k |xref-x0| n n0 e^(-j k n x0)

    """
    x0 = np.asarray(x0)
    n0 = np.asarray(n0)
    n = np.squeeze(np.asarray(n))
    xref = np.squeeze(np.asarray(xref))
    k = util.wavenumber(omega, c)
    return wfs_25d_preeq(omega, omalias, c) * \
        np.sqrt(2*np.pi * np.linalg.norm(xref - x0)) * \
        np.inner(n, n0) * np.exp(-1j * k * np.inner(n, x0))


wfs_3d_plane = _wfs_plane


def wfs_25d_preeq(omega, omalias, c):
    """Preqeualization for 2.5D WFS."""
    if omalias is None:
        return np.sqrt(1j * util.wavenumber(omega, c))
    else:
        if omega <= omalias:
            return np.sqrt(1j * util.wavenumber(omega, c))
        else:
            return np.sqrt(1j * util.wavenumber(omalias, c))


def delay_3d_plane(omega, x0, n0, n=[0, 1, 0], c=None):
    """Plane wave by simple delay of secondary sources."""
    x0 = np.asarray(x0)
    n = np.squeeze(np.asarray(n))
    k = util.wavenumber(omega, c)
    return np.exp(-1j * k * np.inner(n, x0))


def source_selection_plane(n0, n):
    """Secondary source selection for a plane wave.

    Eq.(13) from [Spors et al, 2008]

    """
    n0 = np.asarray(n0)
    n = np.squeeze(np.asarray(n))
    return np.inner(n, n0) >= defs.selection_tolerance


def source_selection_point(n0, x0, xs):
    """Secondary source selection for a point source.

    Eq.(15) from [Spors et al, 2008]

    """
    n0 = np.asarray(n0)
    x0 = np.asarray(x0)
    xs = np.squeeze(np.asarray(xs))
    ds = x0 - xs
    return inner1d(ds, n0) >= defs.selection_tolerance

def source_selection_line(n0, x0, xs):
    """Secondary source selection for a line source.

    compare Eq.(15) from [Spors et al, 2008]

    """
    return source_selection_point(n0, x0, xs)


def source_selection_focused(ns, xdef source_selection_all(N):
    """Select all secondary sources."""
    return np.ones(N) >= 0


def nfchoa_2d_plane(omega, x0, r0, n=[0, 1, 0], c=None):
    """Plane wave by two-dimensional NFC-HOA.

    ::
                               __
                       2i     \        i^-m
        D(phi0,w) = - -----   /__   ----------  e^(i m (phi0-phi_pw))
                      pi r0 m=-N..N  (2)
                                    Hm  (w/c r0)
    """
    x0 = np.asarray(x0)
    k = util.wavenumber(omega, c)
{'fraction_non_alphanumeric': 0.16779026217228465, 'fraction_numerical': 0.0415195291599786, 'mean_word_length': 2.0012845215157355, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 9, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 6, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1179035', 'n_tokens_mistral': 4134, 'n_tokens_neox': 3906, 'n_words': 898}