amachine.am_fast.distance API documentation

# amachine.am_fast.distance

View Source

      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

def jensenshannondivergence_cpu(p, q, axis=0, keepdims=False, base=None): View Source

    42def jensenshannondivergence_cpu(p, q, axis=0, keepdims=False, base=None):
    43	"""
    44	Compute the Jensen-Shannon divergence between two probability arrays on GPU.
    45
    46	This is a port of scipy.spatial.distance.jensenshannon using a
    47	numerically robust rel_entr kernel, except not taking the square root
    48
    49	https://github.com/scipy/scipy/blob/v1.17.0/scipy/spatial/distance.py#L1205-L1290
    50	"""
    51
    52	p = np.asarray(p)
    53	q = np.asarray(q)
    54
    55	# Normalize inputs along the specified axis
    56	p = p / np.sum(p, axis=axis, keepdims=True)
    57	q = q / np.sum(q, axis=axis, keepdims=True)
    58
    59	# Calculate the midpoint distribution
    60	m = (p + q) / 2.0
    61
    62	# Calculate relative entropy using the robust CUDA kernel
    63	# Note: rel_entr results are in nats
    64	left = rel_entr(p, m)
    65	right = rel_entr(q, m)
    66
    67	# Sum across the specified axis
    68	left_sum = np.sum(left, axis=axis, keepdims=keepdims)
    69	right_sum = np.sum(right, axis=axis, keepdims=keepdims)
    70
    71	# JSDiv = (D_kl(P||M) + D_kl(Q||M)) / 2
    72	js = (left_sum + right_sum) / 2.0
    73
    74	# Convert units if base is provided (e.g., base=2 for bits)
    75	if base is not None:
    76		js /= np.log(base)
    77
    78	return js / 2.0

Compute the Jensen-Shannon divergence between two probability arrays on GPU.

This is a port of scipy.spatial.distance.jensenshannon using a numerically robust rel_entr kernel, except not taking the square root

https://github.com/scipy/scipy/blob/v1.17.0/scipy/spatial/distance.py#L1205-L1290

def jensenshannondivergence_gpu(p, q, axis=0, keepdims=False, base=None): View Source

     81def jensenshannondivergence_gpu(p, q, axis=0, keepdims=False, base=None):
     82	"""
     83	Compute the Jensen-Shannon divergence between two probability arrays on GPU.
     84
     85	This is a port of scipy.spatial.distance.jensenshannon using a
     86	numerically robust rel_entr kernel, except not taking the square root
     87
     88	https://github.com/scipy/scipy/blob/v1.17.0/scipy/spatial/distance.py#L1205-L1290
     89	"""
     90
     91	p = cp.asarray(p)
     92	q = cp.asarray(q)
     93
     94	if p.dtype not in (cp.float32, cp.float64) or p.dtype != q.dtype :
     95		raise TypeError(
     96			f"Both arrays must be the same dtype and either float32 or float64, "
     97			f"got p.dtype={p.dtype}, q.dtype={q.dtype}"
     98		)
     99
    100	# Normalize inputs along the specified axis
    101	p = p / cp.sum(p, axis=axis, keepdims=True)
    102	q = q / cp.sum(q, axis=axis, keepdims=True)
    103
    104	# Calculate the midpoint distribution
    105	m = (p + q) / 2.0
    106
    107	# Calculate relative entropy using the robust CUDA kernel
    108	# Note: rel_entr results are in nats
    109	left  = _rel_entr_kernel(p, m)
    110	right = _rel_entr_kernel(q, m)
    111
    112	# Sum across the specified axis
    113	left_sum  = cp.sum(left, axis=axis, keepdims=keepdims)
    114	right_sum = cp.sum(right, axis=axis, keepdims=keepdims)
    115
    116	# JSDiv = (D_kl(P||M) + D_kl(Q||M)) / 2
    117	js = (left_sum + right_sum) / 2.0
    118
    119	# Convert units if base is provided (e.g., base=2 for bits)
    120	if base is not None:
    121		js /= cp.log(base)
    122
    123	return js / 2.0

Compute the Jensen-Shannon divergence between two probability arrays on GPU.

This is a port of scipy.spatial.distance.jensenshannon using a numerically robust rel_entr kernel, except not taking the square root

https://github.com/scipy/scipy/blob/v1.17.0/scipy/spatial/distance.py#L1205-L1290

def jensenshannon_gpu(p, q, axis=0, keepdims=False, base=None): View Source

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

Compute the Jensen-Shannon distance between two probability arrays on GPU.

This is a port of scipy.spatial.distance.jensenshannon using a numerically robust rel_entr kernel.

https://github.com/scipy/scipy/blob/v1.17.0/scipy/spatial/distance.py#L1205-L1290