modules.py

"""
Custom modules for ProTeVa tone recognition model

Authors
 * St Germes BENGONO OBIANG 2024
"""

import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import torchyin
from scipy.interpolate import interp1d
from speechbrain.lobes.models.VanillaNN import VanillaNN
from torch.nn import LeakyReLU, ReLU
from speechbrain.nnet.containers import ModuleList


class F0Extractor(torch.nn.Module):
    """This module extracts F0 of sound and returns it as embedding vector

    Arguments
    ---------
    device : str
        Device to run computations on ('cpu' or 'cuda')
    sample_rate : int
        The signal sample rate (default: 16000)
    frame_stride : float
        Length of the sliding window used for F0 extraction (default: 0.018)
    pitch_min : float
        The minimum value of pitch (default: 50)
    pitch_max : float
        The maximum value of pitch (default: 500)

    Example
    -------
    >>> compute_f0 = F0Extractor(sample_rate=16000)
    >>> input_feats = torch.rand([1, 23000])
    >>> outputs = compute_f0(input_feats, target_size=220)
    >>> outputs.shape
    torch.Size([1, 220, 1])

    Authors
    -------
    * St Germes BENGONO OBIANG 2024
    """

    def __init__(
        self,
        device="cpu",
        sample_rate=16000,
        frame_stride=0.018,
        pitch_min=50,
        pitch_max=500,
    ):
        super().__init__()
        self.device = device
        self.sample_rate = sample_rate
        self.pitch_min = pitch_min
        self.pitch_max = pitch_max
        self.frame_stride = frame_stride

    def interpolate_spline(self, H, N):
        """Interpolate pitch values to target size using cubic spline interpolation"""
        # Generate indices for the original and new tensors
        idx_original = np.arange(len(H))
        idx_new = np.linspace(0, len(H) - 1, N)

        # Create the interpolation function
        interpolator = interp1d(idx_original, H, kind='cubic')

        # Perform interpolation
        H_interpolated = interpolator(idx_new)

        # Create a mask for values below minimum pitch
        mask = H_interpolated < self.pitch_min
        H_interpolated[mask] = 0

        return torch.as_tensor(H_interpolated.tolist())

    def forward(self, wavs, target_size):
        """Extract F0 from waveforms and interpolate to target size"""
        results = []
        for wav in wavs:
            pitch = torchyin.estimate(
                wav,
                self.sample_rate,
                pitch_min=self.pitch_min,
                pitch_max=self.pitch_max,
                frame_stride=self.frame_stride
            )

            # Interpolate the pitch
            pitch = self.interpolate_spline(pitch.cpu().numpy(), target_size)

            # Reshape the pitch output
            pitch = pitch.view(pitch.shape[0], 1)
            results.append(pitch.tolist())

        return torch.as_tensor(results).to(self.device)


class PitchDecoderLayer(torch.nn.Module):
    """Layer for decoding latent vector to pitch

    This decoder reconstructs F0 contours from encoded representations
    using stacked VanillaNN layers.

    Arguments
    ---------
    input_shape : list
        Shape of input tensor [None, None, feature_dim]
    dnn_blocks : list
        Number of blocks for each DNN layer
    dnn_neurons : list
        Number of neurons for each DNN layer

    Authors
    -------
    * St Germes BENGONO OBIANG 2024
    """

    def __init__(
        self,
        input_shape=[None, None, 256],
        dnn_blocks=[2, 2],
        dnn_neurons=[256, 128],
    ):
        super().__init__()
        if len(dnn_blocks) != len(dnn_neurons):
            raise ValueError(
                f"dnn_blocks and dnn_neurons should have the same size but we received {len(dnn_blocks)} and {len(dnn_neurons)}"
            )

        layers = []
        for index in range(len(dnn_neurons)):
            if index == 0:
                layers.append(
                    VanillaNN(
                        activation=LeakyReLU,
                        dnn_blocks=dnn_blocks[index],
                        dnn_neurons=dnn_neurons[index],
                        input_shape=input_shape
                    )
                )
            else:
                # The input shape is equal to the output of the previous layer
                layers.append(
                    VanillaNN(
                        activation=LeakyReLU,
                        dnn_blocks=dnn_blocks[index],
                        dnn_neurons=dnn_neurons[index],
                        input_shape=[None, None, dnn_neurons[index - 1]]
                    )
                )

        # Add the last required layer. The input shape is equal to the last DNN block output
        layers.append(
            VanillaNN(
                activation=ReLU,
                dnn_blocks=1,
                dnn_neurons=1,
                input_shape=[None, None, dnn_neurons[len(dnn_neurons) - 1]]
            )
        )

        self.decoder = ModuleList(*layers)

    def forward(self, latent_vector):
        """Decode latent vector to F0 prediction"""
        return self.decoder(latent_vector)


# ============ HELPER FUNCTIONS FOR PROTOTYPE LAYER ============

def distance_to_prototype(latent_vector, prototypes):
    """
    Compute the L2 squared distance between each timestamp in the latent_vector and each prototype.

    Args:
        latent_vector (torch.Tensor): Tensor of shape [batch, timesteps, features].
        prototypes (torch.Tensor): Tensor of shape [n_prototypes, features].

    Returns:
        torch.Tensor: Tensor of shape [batch, timesteps, n_prototypes] with L2 squared distances.
    """
    # Expand the dimensions of prototypes to match the shape for broadcasting
    prototypes = prototypes.unsqueeze(0).unsqueeze(0)  # Shape: [1, 1, n_prototypes, features]

    # Expand latent_vector to match the shape for broadcasting
    latent_vector = latent_vector.unsqueeze(2)  # Shape: [batch, timesteps, 1, features]

    # Compute the L2 squared distance
    distance = torch.sum((latent_vector - prototypes) ** 2, dim=-1)  # Shape: [batch, timesteps, n_prototypes]

    return distance


def cosine_similarity_to_prototype(latent_vector, prototypes):
    """
    Compute the cosine similarity between each timestamp in the latent_vector and each prototype.

    Args:
        latent_vector (torch.Tensor): Tensor of shape [batch, timesteps, features].
        prototypes (torch.Tensor): Tensor of shape [n_prototypes, features].

    Returns:
        torch.Tensor: Tensor of shape [batch, timesteps, n_prototypes] with cosine similarities.
    """
    # Normalize the latent vector and prototypes
    latent_vector_norm = F.normalize(latent_vector, p=2, dim=-1)  # Shape: [batch, timesteps, features]
    prototypes_norm = F.normalize(prototypes, p=2, dim=-1)  # Shape: [n_prototypes, features]

    # Expand dimensions to match for broadcasting
    prototypes_norm = prototypes_norm.unsqueeze(0).unsqueeze(0)  # Shape: [1, 1, n_prototypes, features]
    latent_vector_norm = latent_vector_norm.unsqueeze(2)  # Shape: [batch, timesteps, 1, features]

    # Compute the cosine similarity
    similarity = torch.sum(latent_vector_norm * prototypes_norm, dim=-1)  # Shape: [batch, timesteps, n_prototypes]

    return similarity


def distances_to_feature(input_tensor, prototypes):
    """
    Compute the L2 squared distance between each prototype and each timestamp in the input_tensor.

    Args:
        input_tensor (torch.Tensor): Tensor of shape [batch_size, num_timestep, feature_dim].
        prototypes (torch.Tensor): Tensor of shape [num_prototypes, feature_dim].

    Returns:
        torch.Tensor: Tensor of shape [num_prototypes, batch_size, num_timestep] with L2 squared distances.
    """
    # Expand the dimensions of prototypes to match the shape for broadcasting
    prototypes = prototypes.unsqueeze(1).unsqueeze(2)  # Shape: [num_prototypes, 1, 1, feature_dim]

    # Expand input_tensor to match the shape for broadcasting
    input_tensor = input_tensor.unsqueeze(0)  # Shape: [1, batch_size, num_timestep, feature_dim]

    # Compute the L2 squared distance
    distance = torch.sum((input_tensor - prototypes) ** 2, dim=-1)  # Shape: [num_prototypes, batch_size, num_timestep]

    return distance


def compute_prototype_distances(prototypes):
    """
    Compute the L2 squared distance between each pair of prototypes.

    Args:
        prototypes (torch.Tensor): Tensor of shape [n_prototypes, features].

    Returns:
        torch.Tensor: Tensor of shape [n_prototypes, n_prototypes] with L2 squared distances between prototypes.
    """
    # Calculate the squared norms of the prototypes
    squared_norms = torch.sum(prototypes ** 2, dim=1, keepdim=True)  # Shape: [n_prototypes, 1]

    # Calculate the pairwise distance using the formula: (a-b)^2 = a^2 + b^2 - 2ab
    distances = squared_norms + squared_norms.T - 2 * torch.mm(prototypes, prototypes.T)  # Shape: [n_prototypes, n_prototypes]
    distances = distances.fill_diagonal_(1e+6)

    return distances


class PrototypeLayer(torch.nn.Module):
    """
    Prototype Layer for tone representation learning

    Learns M prototypes that represent canonical tone patterns.
    Computes similarity between input features and prototypes.
    Includes regularization losses R_1, R_2, and R_3.

    Arguments
    ---------
    n_prototypes : int
        Number of learnable prototypes (default: 9)
    latent_dims : int
        Dimension of latent space (default: 256)

    Authors
    -------
    * St Germes BENGONO OBIANG 2024
    """

    def __init__(
        self,
        n_prototypes=9,
        latent_dims=256,
    ):
        super().__init__()
        self.n_prototypes = n_prototypes
        self.latent_dims = latent_dims

        # Initialize prototypes with Kaiming uniform initialization
        self.prototypes = torch.nn.Parameter(
            torch.nn.init.kaiming_uniform_(
                torch.empty([n_prototypes, latent_dims]),
                nonlinearity='relu'
            ),
            requires_grad=True
        )

        # Regularization losses
        self.R_1 = 0  # Feature distances regulation
        self.R_2 = 0  # Prototypes distances regulation
        self.R_3 = 0  # Prototypes to prototypes distances

    def setProto(self, proto):
        """Set prototype values (for initialization or transfer learning)"""
        self.prototypes = torch.nn.Parameter(proto, requires_grad=True)

    def forward(self, latent_vector):
        """
        Compute similarity between input and prototypes

        Args:
            latent_vector (torch.Tensor): Input features [batch, time, latent_dims]

        Returns:
            torch.Tensor: Prototype similarities [batch, time, n_prototypes]
        """
        # Compute distances and similarities
        dist2proto = distance_to_prototype(latent_vector, self.prototypes)
        similarity2Proto = cosine_similarity_to_prototype(latent_vector, self.prototypes)
        dist2Feature = distances_to_feature(latent_vector, self.prototypes)
        protoDistance = compute_prototype_distances(self.prototypes)

        if self.training:
            # R_1: Each prototype is near to at least one data in latent space
            self.R_1 = torch.mean(torch.min(dist2Feature, dim=2).values)

            # R_2: Each data in latent space is near to at least one prototype
            self.R_2 = torch.mean(torch.min(dist2proto, dim=2).values)

            # R_3: Prototype is as far as possible to other prototypes
            self.R_3 = 1 / (torch.mean(torch.min(protoDistance, dim=1).values) + 1e-8)

        return similarity2Proto