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	
124
125def jensenshannon_gpu(p, q, axis=0, keepdims=False, base=None):
126	"""
127	Compute the Jensen-Shannon distance between two probability arrays on GPU.
128	
129	This is a port of scipy.spatial.distance.jensenshannon using a 
130	numerically robust rel_entr kernel.
131
132	https://github.com/scipy/scipy/blob/v1.17.0/scipy/spatial/distance.py#L1205-L1290
133	"""
134	js = jensenshannondivergence_gpu( p=p, q=q, axis=axis, keepdims=keepdims, base=base )
135
136	return cp.sqrt(js / 2.0)
137