utils/feature_extractors/AF_features_embedding.py

import math
from utils.feature_extractors.dsp_features import compute_log_rms_gated_max, compute_crest_factor, compute_stereo_width, compute_stereo_imbalance, compute_log_spread
import torch

class AF_fourier_embedding:
    def __init__(self,
                input_dim=8,
                output_dim=64,
                sigma=0.2,
                log_rms_shift=-26.5, #calculated as the mean from the dataset
                log_rms_scale=7.0, #calculated as the std from the dataset
                crest_shift=16.7, #calculated as the mean from the dataset
                crest_scale=6.3,
                log_spread_shift=-20.0, #calculated as the mean from the dataset
                log_spread_scale=20.0, #calculated as the std from the dataset
                stereo_width_shift=0.28,
                stereo_width_scale=0.39,
                stereo_imbalance_shift=0.0,
                stereo_imbalance_scale=0.35,
                device="cpu"
                ):
        """
        Deterministic Fourier feature transformer using fixed cosine-based projection
        """

        self.device = device
        # Ensure output_dim is even and >= 2 * input_dim
        self.output_dim = max(input_dim * 2, output_dim)
        if self.output_dim % 2 != 0:
            self.output_dim += 1
        
        self.input_dim = input_dim
        self.sigma = sigma
        
        # Create deterministic projection matrix
        self.projection = self._create_deterministic_projection(input_dim, self.output_dim // 2, sigma)
        self.projection = self.projection.to(self.device)
        
        # Normalization factor
        self.scale_factor = math.sqrt(2.0 / self.output_dim)

        self.log_rms_shift = log_rms_shift
        self.log_rms_scale = log_rms_scale
        self.crest_shift = crest_shift
        self.crest_scale = crest_scale
        self.log_spread_shift = log_spread_shift
        self.log_spread_scale = log_spread_scale
        self.stereo_width_shift = stereo_width_shift
        self.stereo_width_scale = stereo_width_scale
        self.stereo_imbalance_shift = stereo_imbalance_shift
        self.stereo_imbalance_scale = stereo_imbalance_scale
    
    def _create_deterministic_projection(self, input_dim, proj_dim, sigma):
        """
        Create a deterministic projection matrix using a cosine basis
        """
        # Cosine-based matrix (like DCT type-II)
        projection = torch.zeros(input_dim, proj_dim)
        for i in range(input_dim):
            for j in range(proj_dim):
                projection[i, j] = math.cos(math.pi * (i + 0.5) * (j + 1) / proj_dim)
        
        return projection * sigma
    
    def encode(self, x):

        log_rms=compute_log_rms_gated_max(x)
        crest_factor= compute_crest_factor(x)
        log_spread= compute_log_spread(x)
        stereo_width= compute_stereo_width(x)
        stereo_imbalance= compute_stereo_imbalance(x)


        log_rms_std, crest_factor_std, log_spread_std, stereo_width_std, stereo_imbalance_std = self.standardize_features(
            log_rms, crest_factor, log_spread, stereo_width, stereo_imbalance
        )

        embedding= self.transform(
            log_rms_std, crest_factor_std, log_spread_std, stereo_width_std, stereo_imbalance_std
        )


        return embedding, (log_rms, crest_factor, log_spread, stereo_width, stereo_imbalance)
    
    def decode(self, fourier_features):
        """
        Invert Fourier features back to original feature space
        (approximate due to phase-only reconstruction)
        """
        reconstructed = self.inverse_transform(fourier_features)
        
        # Reshape back to original feature dimensions
        log_rms= reconstructed[:,0:2]
        crest_factor = reconstructed[:,2:4]
        log_spread= reconstructed[:,4:6]
        stereo_width = reconstructed[:,6:7]
        stereo_imbalance = reconstructed[:,7:8]

        log_rms, crest_factor, log_spread,stereo_width, stereo_imbalance = self.destandardize_features(
            log_rms, crest_factor, log_spread, stereo_width, stereo_imbalance
        )
        
        return log_rms, crest_factor, log_spread, stereo_width, stereo_imbalance

    def standardize_features(self, log_rms, crest_factor, log_spread, stereo_width, stereo_imbalance):
        """
        Standardize features using pre-computed mean and std
        """
        log_rms = (log_rms - self.log_rms_shift) / self.log_rms_scale
        crest_factor = (crest_factor - self.crest_shift) / self.crest_scale
        log_spread = (log_spread - self.log_spread_shift) / self.log_spread_scale
        stereo_width = (stereo_width - self.stereo_width_shift) / self.stereo_width_scale
        stereo_imbalance = (stereo_imbalance - self.stereo_imbalance_shift) / self.stereo_imbalance_scale
        
        return log_rms, crest_factor, log_spread, stereo_width, stereo_imbalance
    
    def destandardize_features(self, log_rms, crest_factor, log_spread, stereo_width, stereo_imbalance):
        """
        Reverse standardization to get back to original feature space
        """
        log_rms = log_rms * self.log_rms_scale + self.log_rms_shift
        crest_factor = crest_factor * self.crest_scale + self.crest_shift
        log_spread = log_spread * self.log_spread_scale + self.log_spread_shift
        stereo_width = stereo_width * self.stereo_width_scale + self.stereo_width_shift
        stereo_imbalance = stereo_imbalance * self.stereo_imbalance_scale + self.stereo_imbalance_shift
        
        return log_rms, crest_factor, log_spread, stereo_width, stereo_imbalance
        
    def transform(self, log_rms, crest_factor,log_spread, stereo_width, stereo_imbalance):
        """
        Transform features using the stored projection matrix
        """

        flat_features=torch.cat([log_rms, crest_factor, log_spread, stereo_width.unsqueeze(-1), stereo_imbalance.unsqueeze(-1)], dim=-1)
        
        # Project and transform
        projected = flat_features @ self.projection
        cos_features = torch.cos(projected)
        sin_features = torch.sin(projected)
        
        # Concatenate and normalize
        return torch.cat([cos_features, sin_features], dim=-1) * self.scale_factor
    
    def inverse_transform(self, fourier_features):
        """
        Invert Fourier features back to original feature space
        (approximate due to phase-only reconstruction)
        """
        # Split into cosine and sine components
        feature_dim = fourier_features.shape[-1] // 2
        cos_features = fourier_features[:, :feature_dim]
        sin_features = fourier_features[:, feature_dim:]
        
        # Denormalize
        cos_features = cos_features / self.scale_factor
        sin_features = sin_features / self.scale_factor
        
        # Compute phase angles
        phases = torch.atan2(sin_features, cos_features)
        
        # Use pseudo-inverse for approximate inversion
        projection_pinv = torch.pinverse(self.projection)
        reconstructed = phases @ projection_pinv
        
        return reconstructed