amachine.am_fast.distance

  1import cupy as cp
  2import numpy as np 
  3from scipy.special import rel_entr
  4
  5# Relative Entropy Kernel
  6# Ported from SciPy 1.17.1
  7# https://github.com/scipy/scipy/blob/maintenance/1.17.x/scipy/special/_convex_analysis.pxd
  8# Fixes overflow issue not addressed in cupyx implementation
  9# https://github.com/scipy/scipy/issues/20710
 10_rel_entr_kernel = cp.ElementwiseKernel(
 11    'T x, T y',
 12    'T z',
 13    '''
 14    // Select the appropriate minimum normal value for the type at compile time.
 15    // FLT_MIN  ≈ 1.1755e-38  (float32 smallest normal)
 16    // DBL_MIN  ≈ 2.2251e-308 (float64 smallest normal)
 17    const T type_min = (sizeof(T) == 4) ? (T)FLT_MIN : (T)DBL_MIN;
 18
 19    if (isnan(x) || isnan(y)) {
 20        z = (T)CUDART_NAN;
 21    } else if (x <= 0 || y <= 0) {
 22        if (x == 0 && y >= 0) {
 23            z = 0;
 24        } else {
 25            z = (T)CUDART_INF;
 26        }
 27    } else {
 28        T ratio = x / y;
 29        if (ratio > 0.5 && ratio < 2.0) {
 30            z = x * log1p((x - y) / y);
 31        } else if (ratio > type_min && ratio < (T)CUDART_INF) {
 32            z = x * log(ratio);
 33        } else {
 34            z = x * (log(x) - log(y));
 35        }
 36    }
 37    ''',
 38    'rel_entr_robust'
 39)
 40
 41def jensenshannondivergence_cpu(p, q, axis=0, keepdims=False, base=None):
 42    """
 43    Compute the Jensen-Shannon divergence between two probability arrays on GPU.
 44    
 45    This is a port of scipy.spatial.distance.jensenshannon using a 
 46    numerically robust rel_entr kernel, except not taking the square root
 47
 48    https://github.com/scipy/scipy/blob/v1.17.0/scipy/spatial/distance.py#L1205-L1290
 49    """
 50
 51    p = np.asarray(p)
 52    q = np.asarray(q)
 53    
 54    # Normalize inputs along the specified axis
 55    p = p / np.sum(p, axis=axis, keepdims=True)
 56    q = q / np.sum(q, axis=axis, keepdims=True)
 57    
 58    # Calculate the midpoint distribution
 59    m = (p + q) / 2.0
 60    
 61    # Calculate relative entropy using the robust CUDA kernel
 62    # Note: rel_entr results are in nats
 63    left = rel_entr(p, m)
 64    right = rel_entr(q, m)
 65    
 66    # Sum across the specified axis
 67    left_sum = np.sum(left, axis=axis, keepdims=keepdims)
 68    right_sum = np.sum(right, axis=axis, keepdims=keepdims)
 69    
 70    # JSDiv = (D_kl(P||M) + D_kl(Q||M)) / 2
 71    js = (left_sum + right_sum) / 2.0
 72    
 73    # Convert units if base is provided (e.g., base=2 for bits)
 74    if base is not None:
 75        js /= np.log(base)
 76
 77    return js / 2.0
 78
 79
 80def jensenshannondivergence_gpu(p, q, axis=0, keepdims=False, base=None):
 81    """
 82    Compute the Jensen-Shannon divergence between two probability arrays on GPU.
 83    
 84    This is a port of scipy.spatial.distance.jensenshannon using a 
 85    numerically robust rel_entr kernel, except not taking the square root
 86
 87    https://github.com/scipy/scipy/blob/v1.17.0/scipy/spatial/distance.py#L1205-L1290
 88    """
 89    
 90    p = cp.asarray(p)
 91    q = cp.asarray(q)
 92
 93    if p.dtype not in (cp.float32, cp.float64) or p.dtype != q.dtype :
 94        raise TypeError(
 95            f"Both arrays must be the same dtype and either float32 or float64, "
 96            f"got p.dtype={p.dtype}, q.dtype={q.dtype}"
 97        )
 98    
 99    # Normalize inputs along the specified axis
100    p = p / cp.sum(p, axis=axis, keepdims=True)
101    q = q / cp.sum(q, axis=axis, keepdims=True)
102    
103    # Calculate the midpoint distribution
104    m = (p + q) / 2.0
105    
106    # Calculate relative entropy using the robust CUDA kernel
107    # Note: rel_entr results are in nats
108    left  = _rel_entr_kernel(p, m)
109    right = _rel_entr_kernel(q, m)
110    
111    # Sum across the specified axis
112    left_sum  = cp.sum(left, axis=axis, keepdims=keepdims)
113    right_sum = cp.sum(right, axis=axis, keepdims=keepdims)
114    
115    # JSDiv = (D_kl(P||M) + D_kl(Q||M)) / 2
116    js = (left_sum + right_sum) / 2.0
117    
118    # Convert units if base is provided (e.g., base=2 for bits)
119    if base is not None:
120        js /= cp.log(base)
121
122    return js / 2.0
123    
124def jensenshannon_gpu(p, q, axis=0, keepdims=False, base=None):
125    """
126    Compute the Jensen-Shannon distance between two probability arrays on GPU.
127    
128    This is a port of scipy.spatial.distance.jensenshannon using a 
129    numerically robust rel_entr kernel.
130
131    https://github.com/scipy/scipy/blob/v1.17.0/scipy/spatial/distance.py#L1205-L1290
132    """
133    js = jensenshannondivergence_gpu( p=p, q=q, axis=axis, keepdims=keepdims, base=base )
134
135    return cp.sqrt(js / 2.0)
136    
137
138def jensenshannon_pdist_gpu(X: cp.ndarray) -> cp.ndarray:
139
140    n = len(X)
141
142    idx_i, idx_j = cp.triu_indices(n, k=1)  # upper triangle indices
143
144    P = X[idx_i]  # (n_pairs, d)
145    Q = X[idx_j]  # (n_pairs, d)
146
147    js_div  = jensenshannondivergence_gpu(P, Q, axis=1)
148    js_dist = cp.sqrt(js_div)
149
150    return js_dist 
151
152def squareform_gpu(condensed: cp.ndarray, n: int) -> cp.ndarray:
153    """
154    GPU equivalent of scipy squareform.
155    """
156    out = cp.zeros((n, n), dtype=condensed.dtype)
157    idx_i, idx_j = cp.triu_indices(n, k=1)
158    out[idx_i, idx_j] = condensed
159    out[idx_j, idx_i] = condensed
160    return out