Support excitation used in Wavehax

Makefile

@@ -69,7 +69,7 @@ test: tool
 	python -m pytest $$module $(OPT)
 
 test-example: tool
-	[ -n "$(MODULE)" ] && module=modules/$(MODULE).py || module=modules; \
+	[ -n "$(MODULE)" ] && module=modules/$(MODULE).py || module=; \
 	. .venv/bin/activate && export NUMBA_CUDA_LOW_OCCUPANCY_WARNINGS=0 && \
 	python -m pytest --doctest-modules --no-cov --ignore=$(PROJECT)/third_party $(PROJECT)/$$module
 

diffsptk/functional.py

@@ -635,7 +635,7 @@ def excite(
         The frame period in samples, :math:`P`.
 
     voiced_region : ['pulse', 'sinusoidal', 'sawtooth', 'inverted-sawtooth', \
-                     'triangle', 'square']
+                     'triangle', 'square', 'harmonic-pulse']
         The type of voiced region.
 
     unvoiced_region : ['zeros', 'gauss', 'm-sequence', 'uniform']

diffsptk/modules/excite.py

@@ -19,13 +19,114 @@
 import torch
 import torch.nn.functional as F
 
-from ..signals import mseq
+from ..signals import mseq_like
 from ..typing import Precomputed
 from ..utils.private import TAU, UNVOICED_SYMBOL, filter_values
 from .base import BaseFunctionalModule
 from .linear_intpl import LinearInterpolation
 
 
+def generate_pulse(
+    pitch: torch.Tensor,
+    phase: torch.Tensor,
+    bipolar: bool,
+) -> torch.Tensor:
+    def get_pulse_pos(p):
+        return torch.ge(torch.diff(torch.ceil(p)), 1)
+
+    e = torch.zeros_like(pitch)
+    pulse_pos = get_pulse_pos(phase)
+    e[pulse_pos] = torch.sqrt(pitch[pulse_pos])
+
+    if bipolar:
+        pulse_pos_double = get_pulse_pos(0.5 * phase)
+        e[pulse_pos & ~pulse_pos_double] *= -1
+
+    return e
+
+
+def generate_harmonic_pulse(
+    pitch: torch.Tensor,
+    phase: torch.Tensor,
+    bipolar: bool,
+) -> torch.Tensor:
+    # The number of harmonics is floor(0.5 * fs / f0) = floor(0.5 * 1 / T)
+    n_harmonics = torch.floor(0.5 * pitch)
+
+    # The summation of sinusoids can be computed efficiently using the closed-form
+    # expression of the Dirichlet kernel.
+    theta = TAU * phase[..., :-1]
+    half_theta = 0.5 * theta
+    if bipolar:
+        numer = torch.cos(half_theta) - torch.cos((n_harmonics + 0.5) * theta)
+    else:
+        numer = -torch.sin(half_theta) + torch.sin((n_harmonics + 0.5) * theta)
+    denom = 2 * torch.sin(half_theta)
+
+    # Handle singularities at phase is zero (i.e., harmonics).
+    eps = 1e-6
+    is_on_peak = denom.abs() < eps
+    safe_denom = torch.where(is_on_peak, torch.ones_like(denom), denom)
+    e = numer / safe_denom
+    if bipolar:
+        e[is_on_peak] = 0
+    else:
+        e[is_on_peak] = n_harmonics[is_on_peak]
+
+    # Normalize as the same energy as the pulse excitation.
+    norm_factor = torch.sqrt(2 / n_harmonics.clip(min=1))
+    return norm_factor * e
+
+
+def generate_sinusoidal(phase: torch.Tensor, bipolar: bool) -> torch.Tensor:
+    if bipolar:
+        e = torch.sin(TAU * phase)
+    else:
+        e = 0.5 * (1 - torch.cos(TAU * phase))
+    return e
+
+
+def generate_sawtooth(phase: torch.Tensor, bipolar: bool) -> torch.Tensor:
+    e = torch.fmod(phase, 1)
+    if bipolar:
+        e = 2 * e - 1
+    return e
+
+
+def generate_inverted_sawtooth(phase: torch.Tensor, bipolar: bool) -> torch.Tensor:
+    e = 1 - torch.fmod(phase, 1)
+    if bipolar:
+        e = 2 * e - 1
+    return e
+
+
+def generate_triangle(phase: torch.Tensor, bipolar: bool) -> torch.Tensor:
+    if bipolar:
+        e = 2 * torch.abs(2 * torch.fmod(phase + 0.75, 1) - 1) - 1
+    else:
+        e = torch.abs(2 * torch.fmod(phase + 0.5, 1) - 1)
+    return e
+
+
+def generate_square(phase: torch.Tensor, bipolar: bool) -> torch.Tensor:
+    e = torch.le(torch.fmod(phase, 1), 0.5).to(phase.dtype)
+    if bipolar:
+        e = 2 * e - 1
+    return e
+
+
+def generate_gauss(source: torch.Tensor) -> torch.Tensor:
+    return torch.randn_like(source)
+
+
+def generate_mseq(source: torch.Tensor) -> torch.Tensor:
+    return mseq_like(source)
+
+
+def generate_uniform(source: torch.Tensor) -> torch.Tensor:
+    return math.sqrt(12) * torch.rand_like(source)
+
+
 class ExcitationGeneration(BaseFunctionalModule):
     """See `this page <https://sp-nitech.github.io/sptk/latest/main/excite.html>`_
     for details.
@@ -36,7 +137,7 @@ class ExcitationGeneration(BaseFunctionalModule):
         The frame period in samples, :math:`P`.
 
     voiced_region : ['pulse', 'sinusoidal', 'sawtooth', 'inverted-sawtooth', \
-                     'triangle', 'square']
+                     'triangle', 'square', 'harmonic-pulse']
         The type of voiced region.
 
     unvoiced_region : ['zeros', 'gauss', 'm-sequence', 'uniform']
@@ -141,13 +242,7 @@ def _forward(
             p = p.transpose(-2, -1)
         p *= mask
 
-        # Compute phase.
-        voiced_pos = torch.gt(p, 0)
-        q = torch.zeros_like(p)
-        q[voiced_pos] = torch.reciprocal(p[voiced_pos])
-        s = torch.cumsum(q.double(), dim=-1)
-        bias, _ = torch.cummax(s * ~mask, dim=-1)
-        phase = (s - bias).to(p.dtype)
+        # Compute initial phase.
         if not isinstance(init_phase, str):
             shift = init_phase / TAU
         elif init_phase == "zeros":
@@ -156,74 +251,56 @@ def _forward(
             shift = torch.rand_like(p[..., :1])
         else:
             raise ValueError(f"init_phase {init_phase} is not supported.")
-        if isinstance(shift, torch.Tensor) or shift != 0.0:
-            phase += shift
 
-        # Generate excitation signal using phase.
+        # Compute phase.
+        voiced_pos = torch.gt(p, 0)
+        q = torch.zeros_like(p)
+        q[voiced_pos] = torch.reciprocal(p[voiced_pos])
+        s = torch.cumsum(q.double(), dim=-1)
+        bias, _ = torch.cummax(s * ~mask, dim=-1)
+        phase = (s - bias).to(p.dtype)
+
         if polarity == "auto":
-            unipolar = voiced_region == "pulse"
+            bipolar = voiced_region != "pulse"
         elif polarity in ("unipolar", "bipolar"):
-            unipolar = polarity == "unipolar"
+            bipolar = polarity == "bipolar"
         else:
             raise ValueError(f"polarity {polarity} is not supported.")
-        e = torch.zeros_like(p)
-        if voiced_region == "pulse":
-
-            def get_pulse_pos(p):
-                r = torch.ceil(p)
-                return torch.ge(torch.diff(r), 1)
-
-            if isinstance(shift, float):
-                padded_phase = F.pad(phase, (1, 0), value=shift)
-            else:
-                padded_phase = torch.cat([shift, phase], dim=-1)
-
-            if unipolar:
-                pulse_pos = get_pulse_pos(padded_phase)
-                e[pulse_pos] = torch.sqrt(p[pulse_pos])
-            else:
-                pulse_pos1 = get_pulse_pos(padded_phase)
-                pulse_pos2 = get_pulse_pos(0.5 * padded_phase)
-                e[pulse_pos1] = torch.sqrt(p[pulse_pos1])
-                e[pulse_pos1 & ~pulse_pos2] *= -1
-        elif voiced_region == "sinusoidal":
-            if unipolar:
-                e[mask] = 0.5 * (1 - torch.cos(TAU * phase[mask]))
-            else:
-                e[mask] = torch.sin(TAU * phase[mask])
-        elif voiced_region == "sawtooth":
-            if unipolar:
-                e[mask] = torch.fmod(phase[mask], 1)
-            else:
-                e[mask] = 2 * torch.fmod(phase[mask], 1) - 1
-        elif voiced_region == "inverted-sawtooth":
-            if unipolar:
-                e[mask] = 1 - torch.fmod(phase[mask], 1)
-            else:
-                e[mask] = 1 - 2 * torch.fmod(phase[mask], 1)
-        elif voiced_region == "triangle":
-            if unipolar:
-                e[mask] = torch.abs(2 * torch.fmod(phase[mask] + 0.5, 1) - 1)
-            else:
-                e[mask] = 2 * torch.abs(2 * torch.fmod(phase[mask] + 0.75, 1) - 1) - 1
-        elif voiced_region == "square":
-            if unipolar:
-                e[mask] = torch.le(torch.fmod(phase[mask], 1), 0.5).to(e.dtype)
-            else:
-                e[mask] = 2 * torch.le(torch.fmod(phase[mask], 1), 0.5).to(e.dtype) - 1
+
+        if "pulse" in voiced_region:
+            generators = {
+                "pulse": generate_pulse,
+                "harmonic-pulse": generate_harmonic_pulse,
+            }
+            if voiced_region not in generators:
+                raise ValueError(f"voiced_region {voiced_region} is not supported.")
+            phase = F.pad(phase, (1, 0))
+            phase += shift
+            e = generators[voiced_region](p, phase, bipolar)
         else:
-            raise ValueError(f"voiced_region {voiced_region} is not supported.")
+            generators = {
+                "sinusoidal": generate_sinusoidal,
+                "sawtooth": generate_sawtooth,
+                "inverted-sawtooth": generate_inverted_sawtooth,
+                "triangle": generate_triangle,
+                "square": generate_square,
+            }
+            if voiced_region not in generators:
+                raise ValueError(f"voiced_region {voiced_region} is not supported.")
+            phase += shift
+            e = torch.zeros_like(p)
+            e[mask] = generators[voiced_region](phase[mask], bipolar)
 
         if unvoiced_region == "zeros":
             pass
-        elif unvoiced_region == "gauss":
-            e[~mask] = torch.randn(torch.sum(~mask), device=e.device, dtype=e.dtype)
-        elif unvoiced_region == "m-sequence":
-            e[~mask] = mseq(torch.sum(~mask) - 1, device=e.device, dtype=e.dtype)
-        elif unvoiced_region == "uniform":
-            e[~mask] = math.sqrt(12) * (
-                torch.rand(torch.sum(~mask), device=e.device, dtype=e.dtype) - 0.5
-            )
         else:
-            raise ValueError(f"unvoiced_region {unvoiced_region} is not supported.")
+            generators = {
+                "gauss": generate_gauss,
+                "m-sequence": generate_mseq,
+                "uniform": generate_uniform,
+            }
+            if unvoiced_region not in generators:
+                raise ValueError(f"unvoiced_region {unvoiced_region} is not supported.")
+            e[~mask] = generators[unvoiced_region](e[~mask])
+
         return e

diffsptk/signals.py

@@ -301,6 +301,36 @@ def mseq(*order: int, **kwargs) -> torch.Tensor:
     return out.reshape(order)
 
 
+def mseq_like(tensor: torch.Tensor, **kwargs) -> torch.Tensor:
+    """Generate M-sequence with the same shape as the input tensor.
+
+    Parameters
+    ----------
+    tensor : Tensor
+        The input tensor. The shape of the output is the same as this tensor.
+
+    **kwargs : additional keyword arguments
+        See `torch.ones <https://pytorch.org/docs/stable/generated/torch.ones.html>`_.
+
+    Returns
+    -------
+    out : Tensor
+        The random value sequence.
+
+    Examples
+    --------
+    >>> import diffsptk
+    >>> x = torch.empty(4)
+    >>> y = diffsptk.mseq_like(x)
+    >>> y
+    tensor([-1.,  1., -1.,  1.])
+
+    """
+    shape = list(tensor.shape)
+    shape[-1] -= 1
+    return mseq(*shape, device=tensor.device, dtype=tensor.dtype, **kwargs)
+
+
 def nrand(
     *order: int, mean: float = 0, stdv: float = 1, var: float | None = None, **kwargs
 ) -> torch.Tensor:

tests/test_excite.py

@@ -92,7 +92,15 @@ def compute_error(infile):
 
 @pytest.mark.parametrize(
     "voiced_region",
-    ["pulse", "sinusoidal", "sawtooth", "inverted-sawtooth", "triangle", "square"],
+    [
+        "pulse",
+        "sinusoidal",
+        "sawtooth",
+        "inverted-sawtooth",
+        "triangle",
+        "square",
+        "harmonic-pulse",
+    ],
 )
 @pytest.mark.parametrize("polarity", ["unipolar", "bipolar"])
 @pytest.mark.parametrize("init_phase", ["zeros", "random", np.round(np.pi / 2, 2)])
@@ -108,7 +116,7 @@ def test_waveform(voiced_region, polarity, init_phase, P=80, verbose=False):
         U.call(f"x2x +sd tools/SPTK/asset/data.short | pitch -s 16 -p {P} -o 0 -a 2")
     )
     e = excite(pitch)
-    if voiced_region == "pulse":
+    if "pulse" in voiced_region:
         e = e / e.abs().max()
     if verbose:
         sf.write(f"excite_{voiced_region}_{polarity}_{init_phase}.wav", e, 16000)