entropy.py

import numpy as np
from scipy.stats import norm
from joblib import Parallel, delayed

def ncdf_mapping_and_class_assignment(x, c):
    """
    Equation (1) & Section III.A: NCDF mapping and strict boundary handling.
    z_i^c = max(1, min(c, round(c * y_j + 0.5)))
    """
    mu, sigma = np.mean(x), np.std(x)
    if sigma < 1e-8:
        return np.ones(len(x), dtype=int)
        
    y = norm.cdf(x, loc=mu, scale=sigma)
    z = np.round(c * y + 0.5)
    z = np.clip(z, 1, c).astype(int) # Continuous-to-discrete vector quantization
    return z

def get_dispersion_patterns_prob(z, m, c):
    """
    Equation (2) & Section III.A: Calculate relative frequency using base-c index mapping.
    Index(pi) = sum_{k=0}^{m-1} (v_k - 1) c^{m-1-k} + 1
    """
    N = len(z)
    if N < m:
        return np.array([])
        
    # Efficient rolling window for sequence extraction
    shape = (N - m + 1, m)
    strides = (z.strides[0], z.strides[0])
    windows = np.lib.stride_tricks.as_strided(z, shape=shape, strides=strides)
    
    # Base-c index calculation
    powers = c ** np.arange(m-1, -1, -1)
    indices = np.sum((windows - 1) * powers, axis=1)
    
    counts = np.bincount(indices, minlength=c**m)
    probs = counts / (N - m + 1)
    return probs

def calculate_rcmde_single_scale(u, tau, m=2, c=3):
    """
    Equations (6), (7), (8): Refined Composite MDE calculation.
    """
    if tau == 1:
        z = ncdf_mapping_and_class_assignment(u, c)
        probs = get_dispersion_patterns_prob(z, m, c)
        if len(probs) == 0: return 0.0
        probs = probs[probs > 0]
        return -np.sum(probs * np.log(probs))
        
    L = len(u)
    avg_probs = np.zeros(c**m)
    valid_shifts = 0
    
    # Eq 6: Coarse-graining with varying starting points k
    for k in range(1, tau + 1):
        start_idx = k - 1
        cg_len = (L - start_idx) // tau
        if cg_len < m:
            continue
            
        # Coarse-graining procedure
        segments = u[start_idx : start_idx + cg_len * tau].reshape(cg_len, tau)
        x_k = np.mean(segments, axis=1)
        
        z_k = ncdf_mapping_and_class_assignment(x_k, c)
        p_k = get_dispersion_patterns_prob(z_k, m, c)
        avg_probs += p_k
        valid_shifts += 1
        
    if valid_shifts == 0:
        return 0.0
        
    # Eq 7 & 8: Average probabilities and Shannon entropy
    avg_probs /= valid_shifts
    avg_probs = avg_probs[avg_probs > 0] # Avoid log(0)
    return -np.sum(avg_probs * np.log(avg_probs))

def extract_2d_entropy_map(trial_data, tau_max=24, m=2, c=3, n_jobs=-1):
    """
    Section III.E: Constructs the 2D Spatial-Scale Feature Map (Channels x Scales).
    """
    channels = trial_data.shape[0]
    
    def process_channel(ch_data):
        return [calculate_rcmde_single_scale(ch_data, tau, m, c) for tau in range(1, tau_max + 1)]

    # Parallel scaling (Section IV.L)
    feature_map = Parallel(n_jobs=n_jobs)(
        delayed(process_channel)(trial_data[ch]) for ch in range(channels)
    )
    return np.array(feature_map)