# Source code for imaginaire.evaluation.msid

```# Copyright (C) 2021 NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
#
# This work is made available under the Nvidia Source Code License-NC.

# flake8: noqa
from scipy.sparse import lil_matrix, diags, eye
import math

import numpy as np
import torch

EPSILON = 1e-6
NORMALIZATION = 1e6

def _get_msid(x, y, ts=np.logspace(-1, 1, 256), k=5, m=10, niters=100,
msid_mode='max', normalized_laplacian=True, normalize='empty'):
"""
Compute the msid score between two samples, x and y
Arguments:
x: x samples
y: y samples
ts: temperature values
k: number of neighbours for graph construction
m: Lanczos steps in SLQ
niters: number of starting random vectors for SLQ
distributions, else sample standard normal distribution
graph_builder: if 'kgraph', uses faster graph construction (options:
'sparse', 'kgraph')
msid_mode: 'l2' to compute the l2 norm of the distance between `msid1`
and `msid2`; 'max' to find the maximum abosulute difference between
two  descriptors over temperature
normalized_laplacian: if True, use normalized Laplacian
normalize: 'empty' for average heat kernel (corresponds to the empty
graph normalization of NetLSD), 'complete' for the complete, 'er'
for erdos-renyi normalization, 'none' for no normalization
Returns:
msid_score: the scalar value of the distance between discriptors
"""
normed_msidx = msid_descriptor(x, ts, k, m, niters, rademacher,
graph_builder, normalized_laplacian,
normalize)
normed_msidy = msid_descriptor(y, ts, k, m, niters, rademacher,
graph_builder, normalized_laplacian,
normalize)

c = np.exp(-2 * (ts + 1 / ts))

if msid_mode == 'l2':
score = np.linalg.norm(normed_msidx - normed_msidy)
elif msid_mode == 'max':
score = np.amax(c * np.abs(normed_msidx - normed_msidy))
else:
raise Exception('Use either l2 or max mode.')

return {'IMD': score}

[docs]def msid_descriptor(x, ts=np.logspace(-1, 1, 256), k=5, m=10, niters=100,
normalized_laplacian=True, normalize='empty'):
"""
Compute the msid descriptor for a single sample x
Arguments:
x: x samples
ts: temperature values
k: number of neighbours for graph construction
m: Lanczos steps in SLQ
niters: number of starting random vectors for SLQ
distributions, else sample standard normal distribution
graph_builder: if 'kgraph', uses faster graph construction
(options: '
rse', 'kgraph')
normalized_laplacian: if True, use normalized Laplacian
normalize: 'empty' for average heat kernel (corresponds to the empty
graph normalization of NetLSD), 'complete' for the complete, 'er'
for erdos-renyi normalization, 'none' for no normalization
Returns:
normed_msidx: normalized msid descriptor
"""
Lx = _build_graph(x, k, graph_builder, normalized_laplacian)

nx = Lx.shape[0]
msidx = slq_red_var(Lx, m, niters, ts, rademacher)

normed_msidx = _normalize_msid(msidx, normalize, nx, k, ts) * NORMALIZATION

return normed_msidx

def _build_graph(data, k=5, graph_builder='full', normalized=True):
"""
Return Laplacian from data or load preconstructed from path

Arguments:
data: samples
k: number of neighbours for graph construction
graph_builder: if 'kgraph', use faster graph construction
normalized: if True, use nnormalized Laplacian
Returns:
L: Laplacian of the graph constructed with data
"""
if graph_builder == 'sparse':
A = construct_graph_sparse(data.cpu().numpy(), k)
elif graph_builder == 'kgraph':
A = construct_graph_kgraph(data.cpu().numpy(), k)
elif graph_builder == 'full':
A = lil_matrix(construct_graph(data, k).cpu().numpy()).tocsr()
else:
raise Exception('Please specify graph builder: sparse or kgraph.')
A = (A + A.T) / 2
A.data = np.ones(A.data.shape)
L = _laplacian_sparse(A, normalized)
return L

def _normalize_msid(msid, normalization, n, k, ts):
normed_msid = msid.copy()
if normalization == 'empty':
normed_msid /= n
elif normalization == 'complete':
normed_msid /= (1 + (n - 1) * np.exp(-(1 + 1 / (n - 1)) * ts))
elif normalization == 'er':
xs = np.linspace(0, 1, n)
er_spectrum = 4 / np.sqrt(k) * xs + 1 - 2 / np.sqrt(k)
er_msid = np.exp(-np.outer(ts, er_spectrum)).sum(-1)
normed_msid = normed_msid / (er_msid + EPSILON)
elif normalization == 'none' or normalization is None:
pass
else:
raise ValueError('Unknown normalization parameter!')
return normed_msid

def _lanczos_m(A, m, nv, rademacher, SV=None):
"""
Lanczos algorithm computes symmetric m x m tridiagonal matrix T and
matrix V with orthogonal rows constituting the basis of the Krylov
subspace K_m(A, x), where x is an arbitrary starting unit vector. This
implementation parallelizes `nv` starting vectors.

Arguments:
m: number of Lanczos steps
nv: number of random vectors
False - standard normal for random vectors
SV: specified starting vectors

Returns: T: a nv x m x m tensor, T[i, :, :] is the ith symmetric
tridiagonal matrix V: a n x m x nv tensor, V[:, :, i] is the ith matrix
with orthogonal rows
"""
orthtol = 1e-5
if type(SV) != np.ndarray:
SV = np.sign(np.random.randn(A.shape[0], nv))
else:
SV = np.random.randn(A.shape[0],
nv)  # init random vectors in columns: n x nv
V = np.zeros((SV.shape[0], m, nv))
T = np.zeros((nv, m, m))

np.divide(SV, np.linalg.norm(SV, axis=0), out=SV)  # normalize each column
V[:, 0, :] = SV

w = A.dot(SV)
alpha = np.einsum('ij,ij->j', w, SV)
w -= alpha[None, :] * SV
beta = np.einsum('ij,ij->j', w, w)
np.sqrt(beta, beta)

T[:, 0, 0] = alpha
T[:, 0, 1] = beta
T[:, 1, 0] = beta

np.divide(w, beta[None, :], out=w)
V[:, 1, :] = w
t = np.zeros((m, nv))

for i in range(1, m):
SVold = V[:, i - 1, :]
SV = V[:, i, :]

w = A.dot(SV)  # sparse @ dense
w -= beta[None, :] * SVold  # n x nv
np.einsum('ij,ij->j', w, SV, out=alpha)

T[:, i, i] = alpha

if i < m - 1:
w -= alpha[None, :] * SV  # n x nv
# reortho
np.einsum('ijk,ik->jk', V, w, out=t)
w -= np.einsum('ijk,jk->ik', V, t)
np.einsum('ij,ij->j', w, w, out=beta)
np.sqrt(beta, beta)
np.divide(w, beta[None, :], out=w)

T[:, i, i + 1] = beta
T[:, i + 1, i] = beta

# more reotho
innerprod = np.einsum('ijk,ik->jk', V, w)
reortho = False
for _ in range(100):
if not (innerprod > orthtol).sum():
reortho = True
break
np.einsum('ijk,ik->jk', V, w, out=t)
w -= np.einsum('ijk,jk->ik', V, t)
np.divide(w, np.linalg.norm(w, axis=0)[None, :], out=w)
innerprod = np.einsum('ijk,ik->jk', V, w)

V[:, i + 1, :] = w

if (np.abs(beta) > 1e-6).sum() == 0 or not reortho:
break
return T, V

"""
Compute the trace of matrix exponential

Arguments:
A: square matrix in trace(exp(A))
m: number of Lanczos steps
niters: number of quadratures (also, the number of random vectors in the
hutchinson trace estimator)
for random vectors in Hutchinson
Returns: trace: estimate of trace of matrix exponential
"""
T, _ = _lanczos_m(A, m, niters, rademacher)
eigvals, eigvecs = np.linalg.eigh(T)
expeig = np.exp(eigvals)
sqeigv1 = np.power(eigvecs[:, 0, :], 2)
trace = A.shape[-1] * (expeig * sqeigv1).sum() / niters
return trace

def _slq_ts(A, m, niters, ts, rademacher):
"""
Compute the trace of matrix exponential

Arguments:
A: square matrix in trace(exp(-t*A)), where t is temperature
m: number of Lanczos steps
niters: number of quadratures (also, the number of random vectors in the
hutchinson trace estimator)
ts: an array with temperatures
for random vectors in Hutchinson
Returns:
trace: estimate of trace of matrix exponential across temperatures `ts`
"""
T, _ = _lanczos_m(A, m, niters, rademacher)
eigvals, eigvecs = np.linalg.eigh(T)
expeig = np.exp(-np.outer(ts, eigvals)).reshape(ts.shape[0], niters, m)
sqeigv1 = np.power(eigvecs[:, 0, :], 2)
traces = A.shape[-1] * (expeig * sqeigv1).sum(-1).mean(-1)
return traces

def _slq_ts_fs(A, m, niters, ts, rademacher, fs):
"""
Compute the trace of matrix functions

Arguments:
A: square matrix in trace(exp(-t*A)), where t is temperature
m: number of Lanczos steps
niters: number of quadratures (also, the number of random vectors in the
hutchinson trace estimator)
ts: an array with temperatures
for random vectors in Hutchinson
fs: a list of functions
Returns:
traces: estimate of traces for each of the functions in fs
"""
T, _ = _lanczos_m(A, m, niters, rademacher)
eigvals, eigvecs = np.linalg.eigh(T)
traces = np.zeros((len(fs), len(ts)))
for i, f in enumerate(fs):
expeig = f(-np.outer(ts, eigvals)).reshape(ts.shape[0], niters, m)
sqeigv1 = np.power(eigvecs[:, 0, :], 2)
traces[i, :] = A.shape[-1] * (expeig * sqeigv1).sum(-1).mean(-1)
return traces

[docs]def slq_red_var(A, m, niters, ts, rademacher):
"""
Compute the trace of matrix exponential with reduced variance

Arguments:
A: square matrix in trace(exp(-t*A)), where t is temperature
m: number of Lanczos steps
niters: number of quadratures (also, the number of random vectors in the
hutchinson trace estimator)
ts: an array with temperatures
Returns:
traces: estimate of trace for each temperature value in `ts`
"""
fs = [np.exp, lambda x: x]

traces = _slq_ts_fs(A, m, niters, ts, rademacher, fs)
subee = traces[0, :] - traces[1, :] / np.exp(ts)
sub = - ts * A.shape[0] / np.exp(ts)
return subee + sub

[docs]def np_euc_cdist(data):
dd = np.sum(data * data, axis=1)
dist = -2 * np.dot(data, data.T)
dist += dd + dd[:, np.newaxis]
np.fill_diagonal(dist, 0)
np.sqrt(dist, dist)
return dist

[docs]def construct_graph_sparse(data, k):
n = len(data)
spmat = lil_matrix((n, n))
dd = np.sum(data * data, axis=1)

for i in range(n):
dists = dd - 2 * data[i, :].dot(data.T)
inds = np.argpartition(dists, k + 1)[:k + 1]
inds = inds[inds != i]
spmat[i, inds] = 1

return spmat.tocsr()

[docs]def construct_graph_kgraph(data, k):
raise NotImplementedError
#
# n = len(data)
# spmat = lil_matrix((n, n))
# index = pykgraph.KGraph(data, 'euclidean')
# index.build(reverse=0, K=2 * k + 1, L=2 * k + 50)
# result = index.search(data, K=k + 1)[:, 1:]
# spmat[np.repeat(np.arange(n), k, 0), result.ravel()] = 1
# return spmat.tocsr()

[docs]def construct_graph(input_features, k, num_splits=10):
num_samples = input_features.shape[0]
indices = []
for i in range(num_splits):
batch_size = math.ceil(num_samples / num_splits)
start_idx = i * batch_size
end_idx = min((i + 1) * batch_size, num_samples)
dist = torch.cdist(input_features[start_idx:end_idx],
input_features)
indices.append(torch.topk(dist, k + 1, dim=-1, largest=False)[1].cpu())
indices = torch.cat(indices, dim=0)
graph = torch.zeros(num_samples, num_samples, device=indices.device)
graph.scatter_(1, indices, 1)
graph -= torch.eye(num_samples, device=graph.device)
return graph

def _laplacian_sparse(A, normalized=True):
D = A.sum(1).A1
if normalized:
Dsqrt = diags(1 / np.sqrt(D))
L = eye(A.shape[0]) - Dsqrt.dot(A).dot(Dsqrt)
else:
L = diags(D) - A
return L
```