def _wasserstein_distance(a, b):
"""
Wasserstein distance function
:param a: data points from distribution A
:param b: data points from distribution B
:return:
"""
# Compute the cdf distance
a_sorter = tf.argsort(a)
b_sorter = tf.argsort(b)
# Pooled value from both a and b
all_values = K.concatenate([a, b])
all_values = tf.sort(all_values)
# Compute the difference between the sorted pooled values
deltas = _diff(all_values)
# Get the positions of the values of a and b between the 2 distributions
a_cdf_indices = tf.searchsorted(tf.gather(a, a_sorter),
all_values[:-1],
side='right')
b_cdf_indices = tf.searchsorted(tf.gather(b, b_sorter),
all_values[:-1],
side='right')
# Calculate CDF
a_cdf = tf.cast(a_cdf_indices / tf.size(a), 'float')
b_cdf = tf.cast(b_cdf_indices / tf.size(b), 'float')
# Wasserstein distance
return K.sum(tf.multiply(K.abs(a_cdf - b_cdf), deltas))
def emd(self, x, y, nbins=255):
x = tf.reshape(x, [-1])
y = tf.reshape(y, [-1])
range_h = [tf.reduce_min(tf.concat([x, y], axis=-1)), tf.reduce_max(tf.concat([x, y], axis=-1))]
histo_x = tf.cast(tf.histogram_fixed_width(x, range_h, nbins=nbins, dtype=tf.int32), dtype=tf.float32)
histo_y = tf.cast(tf.histogram_fixed_width(y, range_h, nbins=nbins, dtype=tf.int32), dtype=tf.float32)
all_sorted_xy = tf.sort(tf.concat([histo_x, histo_y], axis=-1))
all_sorted_xy_delta = tf.cast(all_sorted_xy[1:] - all_sorted_xy[:-1], dtype=tf.float32)
histo_x_sorted = tf.sort(histo_x)
histo_y_sorted = tf.sort(histo_y)
histo_x_indices = tf.searchsorted(histo_x_sorted, all_sorted_xy[:-1], side='right')
histo_y_indices = tf.searchsorted(histo_y_sorted, all_sorted_xy[:-1], side='right')
cmdf_x = tf.cast(tf.math.divide(histo_x_indices, nbins), dtype=tf.float32)
cmdf_y = tf.cast(tf.math.divide(histo_y_indices, nbins), dtype=tf.float32)
return tf.math.sqrt(tf.reduce_sum(tf.math.multiply(tf.math.squared_difference(cmdf_x, cmdf_y), all_sorted_xy_delta)))
def tf_wasserstein_dist(label, predicted):
u_values = K.flatten(predicted)
v_values = K.flatten(label)
u_weights = weightsMatrix[0]
v_weights = weightsMatrix[0]
u_sorter = tf.argsort(u_values)
v_sorter = tf.argsort(v_values)
all_values = tf.concat((u_values, v_values), axis=0)
all_values = tf.sort(all_values)
def tf_diff_axis_0(a):
return a[1:] - a[:-1]
deltas = tf_diff_axis_0(all_values)
u_cdf_indices = tf.gather(u_values, u_sorter)
u_cdf_indices = tf.searchsorted(u_cdf_indices, all_values[:-1], 'right')
v_cdf_indices = tf.gather(v_values, v_sorter)
v_cdf_indices = tf.searchsorted(v_cdf_indices, all_values[:-1], 'right')
u_sorted_weights = K.cumsum(tf.gather(u_weights, u_sorter))
u_sorted_weights = tf.concat(([0], u_sorted_weights), axis=0)
u_cdf = tf.gather(u_sorted_weights, u_cdf_indices) / u_sorted_weights[-1]
v_sorted_weights = K.cumsum(tf.gather(v_weights, v_sorter))
v_sorted_weights = tf.concat(([0], v_sorted_weights), axis=0)
v_cdf = tf.gather(v_sorted_weights, v_cdf_indices) / v_sorted_weights[-1]
result = tf.cast(K.abs(u_cdf - v_cdf), tf.float32)
return K.sum(tf.math.multiply(result, deltas))
def update_state(self, y_true, y_pred, sample_weight=None):
"""Accumulates ranks.
Args:
y_true: actual rank values.
y_pred: predicted rank values.
sample_weight (optional): Ignored.
Returns:
Update op.
"""
if y_true.shape and y_true.shape[0]:
i = tf.searchsorted(
self.actual_cuts,
tf.cast(tf.reshape(y_true, -1), self.actual_cuts.dtype),
)
j = tf.searchsorted(
self.preds_cuts, tf.cast(tf.reshape(y_pred, -1), self.preds_cuts.dtype)
)
def body(k, n, m, nrow, ncol):
return (
k + 1,
n + 1,
tf.sparse.add(
m,
tf.SparseTensor(
[[i[k], j[k]]],
tf.cast([1], dtype=self.m.dtype),
self.m.shape,
),
),
tf.sparse.add(
nrow,
tf.SparseTensor(
[[i[k]]],
tf.cast([1], dtype=self.nrow.dtype),
self.nrow.shape,
),
),
tf.sparse.add(
ncol,
tf.SparseTensor(
[[j[k]]],
tf.cast([1], dtype=self.ncol.dtype),
self.ncol.shape,
),
),
)
_, self.n, self.m, self.nrow, self.ncol = tf.while_loop(
lambda k, n, m, nrow, ncol: k < i.shape[0],
body=body,
loop_vars=(0, self.n, self.m, self.nrow, self.ncol),
)
def replace(x, mapping, zero=True):
"""Replace values in tensor `x` using dictionary `mapping`.
Parameters
----------
x: tensor, values to replace.
mapping: dict, dictionary mapping original values to new values. Values in
x equal to a key in the mapping are replaced with the corresponding
value. Keys and values may overlap.
zero: boolean, zero values in `x` not in `mapping.keys()`.
Returns
-------
Modified tensor.
"""
x = tf.cast(x, dtype=tf.int32)
keys = tf.convert_to_tensor(list(mapping.keys()))
vals = tf.convert_to_tensor(list(mapping.values()))
sidx = tf.argsort(keys)
ks = tf.gather(keys, sidx)
vs = tf.gather(vals, sidx)
idx = tf.searchsorted(ks, tf.reshape(x, (-1, )))
idx = tf.reshape(idx, x.shape)
# Zero values that are equal to len(vs).
idx = tf.multiply(idx, tf.cast(tf.not_equal(idx, vs.shape[0]), tf.int32))
mask = tf.equal(tf.gather(ks, idx), x)
if zero:
out = tf.where(mask, tf.gather(vs, idx), 0)
else:
out = tf.where(mask, tf.gather(vs, idx), x)
return out
def map_cat_sparse(fn, Nr, batchsize, index_dtype=tf.int64):
CURROW = 0
WIDTH = None
def go(st, en):
nonlocal CURROW, WIDTH
g = fn(st, en)
# check for consistant widths
if WIDTH is None:
WIDTH = g.shape[1]
else:
assert WIDTH == g.shape[1]
rows_n_cols = tf.cast(tf.where(g), dtype=index_dtype)
data = tf.gather_nd(g, rows_n_cols)
rows = rows_n_cols[:, 0] + st
cols = rows_n_cols[:, 1]
return rows, cols, data
rows, cols, data = map_accum(go, Nr, batchsize, 'ccc')
row_indptr = tf.searchsorted(rows, tf.range(0, Nr + 1, dtype=rows.dtype))
return sparsematrix.CSRMatrix(row_indptr, rows, cols, data, (Nr, WIDTH),
rows.shape[0])
def rle2mask(rle, mask_shape):
'''
Converts a run lenght encoding (RLE) into a mask of shape mask_shape
Args:
rle: (str or bytestring) run lenght encoding. A series of space
separated start-pixel run pairs.
mask_shape: (tuple of 2 ints) the 2D expected shape of the mask
Returns: mask of shape mask_shape
'''
size = tf.math.reduce_prod(mask_shape)
s = tf.strings.split(rle)
s = tf.strings.to_number(s, tf.int32)
starts = s[0::2] - 1
lens = s[1::2]
total_ones = tf.reduce_sum(lens)
ones = tf.ones([total_ones], tf.int32)
r = tf.range(total_ones)
lens_cum = tf.math.cumsum(lens)
s = tf.searchsorted(lens_cum, r, 'right')
idx = r + tf.gather(starts - tf.pad(lens_cum[:-1], [(1, 0)]), s)
mask_flat = tf.scatter_nd(tf.expand_dims(idx, 1), ones, [size])
mask = tf.reshape(mask_flat, (mask_shape[1], mask_shape[0]))
return tf.transpose(mask)
def get_all_N_bit_intervals(self, Z):
"""
:param Z: batch of latent representations of shape B x C
:return: two tensors of shape C x (N+1) x B
"""
N = self.max_bits_per_coord
# backend has to be tensorflow >= 1.15 for searchsorted to work below (numpy's version doesn't support matrices)
Z_repeated = tf.repeat(tf.transpose(Z)[:, None, :], N + 1,
axis=1) # C x (N+1) x B
right_endpoint_idx = tf.searchsorted(
self._search_grids, Z_repeated,
side='left') # search inner-most dim
right_endpoint_idx = tf.clip_by_value(right_endpoint_idx, 0,
2**N - 1) # C x (N+1) x B
left_endpoint_idx = tf.clip_by_value(right_endpoint_idx - 1, 0,
2**N - 1)
right_endpoints = tf.gather(self._search_grids,
right_endpoint_idx,
batch_dims=2)
left_endpoints = tf.gather(self._search_grids,
left_endpoint_idx,
batch_dims=2)
return left_endpoints, right_endpoints
def from_dense_tensor(cls,t,index_dtype=tf.int64):
rows_n_cols = tf.cast(tf.where(t),dtype=index_dtype)
data=tf.gather_nd(t,rows_n_cols)
rows=rows_n_cols[:,0]
cols=rows_n_cols[:,1]
row_indptr = tf.searchsorted(rows,tf.range(0,t.shape[0]+1,dtype=rows.dtype))
return CSRMatrix(row_indptr,cols,data,t.shape,rows.shape[0])
def sample_pdf(self, u, bins, weights):
""" Creates samples along ray.
Args:
u: Distribution of sample probabilities along rays
bins: Number of samples to compute
weights: Unnormalized pdf"""
weights += 1e-5 # prevent nans
pdf = weights / tf.reduce_sum(weights, -1, keepdims=True)
print(f"pdf {pdf}")
cdf = tf.cumsum(pdf, -1)
cdf = tf.concat([tf.zeros_like(cdf[..., :1]), cdf], -1)
inds = tf.searchsorted(cdf, u, side='right')
print(f"inds {inds}")
below = tf.maximum(0, inds - 1)
above = tf.minimum(tf.shape(cdf)[-1] - 1, inds)
inds_g = tf.stack([below, above], -1)
cdf_g = tf.gather(cdf,
inds_g,
axis=-1,
batch_dims=len(tf.shape(inds_g)) - 2)
bins_g = tf.gather(bins,
inds_g,
axis=-1,
batch_dims=len(tf.shape(inds_g)) - 2)
denom = (cdf_g[..., 1] - cdf_g[..., 0])
denom = tf.where(denom < 1e-5, tf.ones_like(denom), denom)
t = (u - cdf_g[..., 0]) / denom
samples = bins_g[..., 0] + t * (bins_g[..., 1] - bins_g[..., 0])
return samples
def _inverse(self, y):
flat_y = tf.reshape(y, shape=[-1])
# Search for the indices of self.map_values that are closest to flat_y.
# Since self.map_values is strictly increasing, the closest is either the
# first one that is strictly greater than flat_y, or the one before it.
upper_candidates = tf.minimum(
tf.size(input=self.map_values) - 1,
tf.searchsorted(self.map_values, values=flat_y, side='right'))
lower_candidates = tf.maximum(0, upper_candidates - 1)
candidates = tf.stack([lower_candidates, upper_candidates], axis=-1)
lower_cand_diff = tf.abs(flat_y - self._forward(lower_candidates))
upper_cand_diff = tf.abs(flat_y - self._forward(upper_candidates))
if self.validate_args:
with tf.control_dependencies([
tf.compat.v1.assert_near(
tf.minimum(lower_cand_diff, upper_cand_diff),
0,
message='inverse value not found')
]):
candidates = tf.identity(candidates)
candidate_selector = tf.stack([
tf.range(tf.size(input=flat_y), dtype=tf.int32),
tf.argmin(
input=[lower_cand_diff, upper_cand_diff], output_type=tf.int32)
],
axis=-1)
return tf.reshape(
tf.gather_nd(candidates, candidate_selector), shape=y.shape)
def sample_pdf(bins, weights, n_samples, det=False):
"""Function for sampling a probability distribution."""
# Get pdf
weights += 1e-5 # prevent nans
pdf = weights / tf.reduce_sum(weights, -1, keepdims=True)
cdf = tf.cumsum(pdf, -1)
cdf = tf.concat([tf.zeros_like(cdf[Ellipsis, :1]), cdf], -1)
# Take uniform samples
u_shape = tf.concat([tf.shape(cdf)[:-1], tf.constant([n_samples])], axis=0)
if det:
u = tf.linspace(0., 1., n_samples)
u = tf.broadcast_to(u, u_shape)
else:
u = tf.random.uniform(u_shape)
# Invert CDF
inds = tf.searchsorted(cdf, u, side='right')
below = tf.maximum(0, inds - 1)
above = tf.minimum(tf.shape(cdf)[-1] - 1, inds)
inds_g = tf.stack([below, above], -1)
cdf_g = tf.gather(cdf, inds_g, axis=-1, batch_dims=len(tf.shape(inds_g)) - 2)
bins_g = tf.gather(
bins, inds_g, axis=-1, batch_dims=len(tf.shape(inds_g)) - 2)
denom = (cdf_g[Ellipsis, 1] - cdf_g[Ellipsis, 0])
denom = tf.where(denom < 1e-5, tf.ones_like(denom), denom)
t = (u - cdf_g[Ellipsis, 0]) / denom
samples = bins_g[Ellipsis, 0] + t * (bins_g[Ellipsis, 1] - bins_g[Ellipsis, 0])
return samples
def sample_pdf(bins, weights, N_samples, det=False):
# Get pdf
weights += 1e-5 # prevent nans
pdf = weights / tf.reduce_sum(input_tensor=weights, axis=-1, keepdims=True)
cdf = tf.cumsum(pdf, -1)
cdf = tf.concat([tf.zeros_like(cdf[..., :1]), cdf], -1)
# Take uniform samples
if det:
u = tf.linspace(0., 1., N_samples)
u = tf.broadcast_to(u, list(cdf.shape[:-1]) + [N_samples])
else:
u = tf.random.uniform(list(cdf.shape[:-1]) + [N_samples])
# Invert CDF
inds = tf.searchsorted(cdf, u, side='right')
below = tf.maximum(0, inds - 1)
above = tf.minimum(cdf.shape[-1] - 1, inds)
inds_g = tf.stack([below, above], -1)
cdf_g = tf.gather(cdf, inds_g, axis=-1, batch_dims=len(inds_g.shape) - 2)
bins_g = tf.gather(bins, inds_g, axis=-1, batch_dims=len(inds_g.shape) - 2)
denom = (cdf_g[..., 1] - cdf_g[..., 0])
denom = tf.compat.v1.where(denom < 1e-5, tf.ones_like(denom), denom)
t = (u - cdf_g[..., 0]) / denom
samples = bins_g[..., 0] + t * (bins_g[..., 1] - bins_g[..., 0])
return samples
def alphas_neighbour_knots(a_q2, padded_q2, actual_values):
"""
Parameters
----------
a_q2: tf.tensor
tensor of values of q2
padded_q2: tf.tensor
values of log(q2) of the grid
actual_values: tf.tensor
values of the grid
Returns
----------
q2_id: tf.tensor of shape [None]
q2 bin for each query point
corn_q2: tf.tensor of shape [4,None]
q2 values of the 4 knots around the query point
A: tf.tensor of shape [4,None]
alphas values of the 4 grid knots around the query point
"""
# print('nk')
q2_id = tf.searchsorted(padded_q2[1:-1], a_q2, out_type=DTYPEINT)
s = tf.size(padded_q2, out_type=DTYPEINT)
q2_id = tf.clip_by_value(q2_id, tf.constant([0], dtype=DTYPEINT), s - 3)
piu = tf.reshape(tf.range(-1, 3, dtype=DTYPEINT), (4, 1))
corn_q2_id = tf.repeat(tf.reshape(q2_id, (1, -1)), 4, axis=0) + piu
corn_q2 = tf.gather(padded_q2, corn_q2_id, name="fnk_2")
A = tf.gather(actual_values, corn_q2_id, name="fnk_3")
return q2_id, corn_q2, A
def offsets_to_segment_ids(offsets):
'''Transforms offsets to segment_ids,
the segment_ids will be used in tf.segment_sum/segment_mean
[3, 0, 1, 2] -> [0, 0, 0, 1, 3, 3].
'''
c = tf.cumsum(offsets)
return tf.searchsorted(c, tf.range(c[-1]), side='right')
def _piecewise_constant_function(x, jump_locations, values,
batch_rank, side='left'):
"""Computes value of the piecewise constant function."""
# Initializer already verified that `jump_locations` and `values` have the
# same shape
batch_shape = jump_locations.shape.as_list()[:-1]
# Check that the batch shape of `x` is the same as of `jump_locations` and
# `values`
batch_shape_x = x.shape.as_list()[:batch_rank]
if batch_shape_x != batch_shape:
raise ValueError('Batch shape of `x` is {1} but should be {0}'.format(
batch_shape, batch_shape_x))
if x.shape.as_list()[:batch_rank]:
no_batch_shape = False
else:
no_batch_shape = True
x = tf.expand_dims(x, 0)
# Expand batch size to one if there is no batch shape
if not batch_shape:
jump_locations = tf.expand_dims(jump_locations, 0)
values = tf.expand_dims(values, 0)
indices = tf.searchsorted(jump_locations, x, side=side)
index_matrix = _prepare_index_matrix(
indices.shape.as_list()[:-1], indices.shape.as_list()[-1], indices.dtype)
indices_nd = tf.concat(
[index_matrix, tf.expand_dims(indices, -1)], -1)
res = tf.gather_nd(values, indices_nd)
if no_batch_shape:
return tf.squeeze(res, 0)
else:
return res
def quantize(self, lab_images):
l = lab_images[:, :, :, :1] * 50.0
ab = (lab_images[:, :, :, 1:] * 255.0) - 128.0
e = np.linspace(1, 50, 50, dtype=np.float32)
linspace = np.tile(e, [ab.shape[0], ab.shape[1], ab.shape[2], 1])
l_inds = tf.searchsorted(linspace, l)
l_inds = tf.squeeze(l_inds)
bs, h, w, c = ab.shape
ab = tf.reshape(ab, [bs * h * w, c])
(dists, inds) = self.nbrs.kneighbors(ab)
# Sigma = 5
sigma = 5
wts = tf.exp(dists**2 / (2 * sigma**2))
#wts = tf.reduce_mean(wts, axis=1)
wts = tf.nn.softmax(wts)
inds = tf.expand_dims(inds, -1)
batch_ind = tf.expand_dims(tf.range(0, bs * h * w, 1, dtype=tf.int64),
-1)
batch_ind = tf.tile(batch_ind, tf.constant([1, 5], tf.int32))
batch_ind = tf.expand_dims(batch_ind, -1)
inds = tf.concat([batch_ind, inds], axis=-1)
#wts = tf.expand_dims(wts, -2)
#hot_mixed_l = tf.scatter_nd(indices=l_inds, updates=1, shape=[128, 32, 32, 50])
hot_mixed_l = tf.one_hot(l_inds, 50)
hot_mixed_ab = tf.scatter_nd(indices=inds,
updates=wts,
shape=[bs * h * w, 313])
hot_mixed_ab = tf.reshape(hot_mixed_ab, [bs, h, w, 313])
return hot_mixed_l, hot_mixed_ab
def rle_decode_tf(
mask_rle # the run-length encoded mask, as a string
,
shape # the image dimension e.g. (640,480)
):
"""
function for converting masks in Run-Length Encoding into a tensor
source: https://stackoverflow.com/questions/58693261/decoding-rle-run-length-encoding-mask-with-tensorflow-datasets
"""
shape = tf.convert_to_tensor(shape, tf.int64)
size = tf.math.reduce_prod(shape)
# Split string
s = tf.strings.split(mask_rle)
s = tf.strings.to_number(s, tf.int64)
# Get starts and lengths
starts = s[::2] - 1
lens = s[1::2]
# Make ones to be scattered
total_ones = tf.reduce_sum(lens)
ones = tf.ones([total_ones], tf.uint8)
# Make scattering indices
r = tf.range(total_ones)
lens_cum = tf.math.cumsum(lens)
s = tf.searchsorted(lens_cum, r, 'right')
idx = r + tf.gather(starts - tf.pad(lens_cum[:-1], [(1, 0)]), s)
# Scatter ones into flattened mask
mask_flat = tf.scatter_nd(tf.expand_dims(idx, 1), ones, [size])
# Reshape into mask
return tf.transpose(tf.reshape(mask_flat, shape))
def _discrete_percentile_function(spacings,
n_particles,
on_log,
weights=None,
log_weights=None):
"""vectorised resampling function, can be used for systematic/stratified/multinomial resampling
"""
if on_log:
cumlogsumexp = tf.math.cumulative_logsumexp(log_weights, axis=1)
log_spacings = tf.math.log(spacings)
indices = tf.searchsorted(cumlogsumexp, log_spacings, side='left')
else:
cum_sum = tf.math.cumsum(weights, axis=1)
indices = tf.searchsorted(cum_sum, spacings, side='left')
return tf.clip_by_value(indices, 0, n_particles - 1)
def _interleave_ecdfs(x1: tf.Tensor, y1: tf.Tensor, x2: tf.Tensor,
y2: tf.Tensor) -> Tuple[tf.Tensor, tf.Tensor, tf.Tensor]:
"""
Interleave two eCDFs by their argument
"""
x = tf.contrib.framework.sort(tf.concat([x1, x2], axis=0), axis=0)
y1 = tf.concat([tf.zeros(shape=(1, tf.shape(y1)[1]), dtype=y1.dtype), y1],
axis=0)
y2 = tf.concat([tf.zeros(shape=(1, tf.shape(y2)[1]), dtype=y2.dtype), y2],
axis=0)
return (x,
_T(
tf.batch_gather(_T(y1),
tf.searchsorted(_T(x1), _T(x), side='right'))),
_T(
tf.batch_gather(_T(y2),
tf.searchsorted(_T(x2), _T(x), side='right'))))
def __call__(self, t):
"Evaluate p(t)"
# does not use horners scheme but it's probably ok
t = tf.convert_to_tensor([t], dtype=tf.float64)
i = tf.searchsorted(self.x, t, side="right") - 1
ci = tf.gather(self.c, i, axis=1) # [D, T]
ti = t - tf.gather(self.x, i, axis=0) # [T]
i = tf.range(ci.shape[0], dtype=tf.float64)[::-1, None]
return tf.reduce_sum(ci * tf.math.pow(ti[None], i), axis=0)[0] # [T]
def compress_latents(self, posterior_means, posterior_logvars, lambs):
float_type = self.float_type
posterior_vars = tf.exp(posterior_logvars)
C = posterior_vars.shape[-1] # number of latent channels
assert C == self.num_channels
batch_means, batch_vars = map(
lambda r: tf.reshape(r, (-1, C)),
[posterior_means, posterior_vars]) # num_samples x C
batch_stds = batch_vars**0.5
Z_hat_dict, raw_num_bits_dict = self.compress_batch_channel_latents(
batch_means, batch_stds, lambs, return_np=False)
out_keys = ('Z_hat', 'raw_num_bits', 'num_bits_cl', 'num_bits')
output = {key: dict() for key in out_keys}
for lamb in lambs:
tmp_res_for_lamb = dict()
Z_hat = Z_hat_dict[lamb]
raw_num_bits = raw_num_bits_dict[lamb]
tmp_res_for_lamb['Z_hat'] = Z_hat
tmp_res_for_lamb['raw_num_bits'] = raw_num_bits
# # calculate code length using raw num bits + codelength overhead approach
# raw_code_length_entropy_model = self.raw_code_length_entropy_models[lamb] # C x (N+1)
# raw_code_length_overhead_num_bits = tf.gather(raw_code_length_entropy_model,
# tf.transpose(raw_num_bits), batch_dims=1) # C x B
# raw_code_length_overhead_num_bits = tf.transpose(raw_code_length_overhead_num_bits) # B x C
# num_bits_cl = tf.cast(raw_num_bits, float_type) + \
# tf.cast(raw_code_length_overhead_num_bits, float_type) # B x C
# calculate code length using direct entropy coding approach
I = tf.searchsorted(self.code_points_by_channel,
tf.transpose(Z_hat)) # C x B; requires TF
# assert tf.reduce_all(tf.equal(tf.gather(self.code_points_by_channel, I, batch_dims=1),
# tf.transpose(Z_hat)))
entropy_model = self.entropy_models[
lamb] # C x quantization_levels
num_bits = tf.gather(entropy_model, I,
batch_dims=1) # gather across innermost axis
num_bits = tf.transpose(num_bits) # B x C
# tmp_res_for_lamb['num_bits_cl'] = num_bits_cl
if self.raw_code_length_entropy_models:
tmp_res_for_lamb['num_bits_cl'] = raw_num_bits
tmp_res_for_lamb['num_bits'] = num_bits
for key in out_keys:
# tmp_res_for_lamb[key] = np.reshape(item, posterior_means.shape)
output[key][lamb] = np.reshape(
tmp_res_for_lamb[key], posterior_means.shape
) # same shape as latents; np.reshape moves to CPU
return output
def interp(x, xp, fp):
i = tf.clip_by_value(tf.searchsorted(xp, x, side='right'), 0,
xp.shape[0] - 1)
df = tf.gather(fp, i) - tf.gather(fp, i - 1)
dx = tf.gather(xp, i) - tf.gather(xp, i - 1)
delta = x - tf.gather(xp, i - 1)
f = tf.where((dx == 0), tf.cast(tf.gather(fp, i), tf.float32),
tf.cast(tf.gather(fp, i - 1), tf.float32) +
(delta / dx) * tf.cast(df, tf.float32))
return f * 0.99999
def sample_batch(_):
values = tf.random.uniform((sample_batch_size, ), 0.0,
weight_sum[-1])
ind = tf.searchsorted(weight_sum, values)
return (tf.gather(self.states, ind,
0), tf.gather(self.actions, ind, 0),
tf.gather(self.next_states, ind,
0), tf.gather(self.rewards, ind,
0), tf.gather(self.masks, ind, 0),
tf.gather(self.weights, ind,
0), tf.gather(self.steps, ind, 0))
def weighted_quantile(x, quantiles, weights=None, side="middle"):
"""Very close to numpy.percentile, but supports weights.
NOTE: quantiles should be in [0, 1]!
:param x: tensor with data
:param quantiles: array-like with many quantiles needed
:param weights: array-like of the same length as `x`
:return: numpy.array with computed quantiles.
"""
if weights is None:
return tfp.stats.percentile(x, 100 * quantiles)
x = znp.array(x)
quantiles = znp.array(quantiles)
quantiles = znp.reshape(quantiles, (-1, ))
weights = znp.array(weights)
sorter = znp.argsort(x)
x = tf.gather(x, sorter)
weights = tf.gather(weights, sorter)
weighted_quantiles = znp.cumsum(weights) - 0.5 * weights
weighted_quantiles /= znp.sum(weights)
if side == "middle":
quantile_index_left = tf.searchsorted(weighted_quantiles,
quantiles,
side="left")
quantile_index_right = tf.searchsorted(weighted_quantiles,
quantiles,
side="right")
calculated_left = tf.gather(x, quantile_index_left)
calculated_right = tf.gather(x, quantile_index_right)
calculated = (calculated_left + calculated_right) / 2
elif side in ("left", "right"):
quantile_index = tf.searchsorted(weighted_quantiles,
quantiles,
side=side)
calculated = tf.gather(x, quantile_index)
return calculated
def inverse_transform_sampling_1d(
bins: TensorLike,
pdf: TensorLike,
num_samples: int,
name="inverse_transform_sampling_1d") -> tf.Tensor:
"""Sampling 1D points from a distribution using the inverse transform.
The target distrubution is defined by its probability density function and
the spatial 1D location of its bins. The new random samples correspond to
the centers of the bins.
Args:
bins: A tensor of shape `[A1, ..., An, M]` containing 1D location of M bins.
For example, a tensor [a, b, c, d] corresponds to
the bin structure |--a--|-b-|--c--|d|.
pdf: A tensor of shape `[A1, ..., An, M]` containing the probability
distribution in M bins.
num_samples: The number N of new samples.
name: A name for this op that defaults to "inverse_transform_sampling_1d".
Returns:
A tensor of shape `[A1, ..., An, N]` indicating the new N random points.
"""
with tf.name_scope(name):
bins = tf.convert_to_tensor(value=bins)
pdf = tf.convert_to_tensor(value=pdf)
shape.check_static(tensor=bins,
tensor_name="bins",
has_rank_greater_than=0)
shape.check_static(tensor=pdf,
tensor_name="pdf",
has_rank_greater_than=0)
shape.compare_batch_dimensions(tensors=(bins, pdf),
tensor_names=("bins", "pdf"),
last_axes=-2,
broadcast_compatible=True)
shape.compare_dimensions(tensors=(bins, pdf),
tensor_names=("bins", "pdf"),
axes=-1)
# Do not consider the last bin, as the cdf contains has +1 dimension.
pdf = _normalize_pdf(pdf[..., :-1])
cdf = _get_cdf(pdf)
batch_shape = tf.shape(pdf)[:-1]
# TODO(krematas): Use dynamic values
batch_dims = tf.get_static_value(tf.rank(pdf) - 1)
target_shape = tf.concat([batch_shape, [num_samples]], axis=-1)
uniform_samples = tf.random.uniform(target_shape)
bin_indices = tf.searchsorted(cdf, uniform_samples, side="right")
bin_indices = tf.maximum(0, bin_indices - 1)
z_values = tf.gather(bins, bin_indices, axis=-1, batch_dims=batch_dims)
return z_values
def tf_argsort_based_on_two_cols(v):
c2_log = tf.log(1 + v[:, 1])
max_c2 = tf.reduce_max(c2_log) + 1
us_c1_u, _ = tf.unique(v[:, 0])
c1_u = tf.sort(us_c1_u)
c1_inds = tf.searchsorted(c1_u, v[:, 0], side='left')
eq_num = tf.cast(c1_inds, tf.float64) * max_c2 + c2_log
return tf.argsort(eq_num)
def backtest_rolling_forecast_aggregates(params, historic_types, historic_times, event_types, event_times,
batch_size, window, n_events, n_simulations, score_against, score_shape=[], pred_window_multiple=1, n_splits=1): # our time step is the same as the window size. params may contain multiple processes once again
interval_event_types, interval_event_times, n_intervals = chop_record_by_time_interval(event_types, event_times, window)
leftover_excitation = tf.zeros_like(params.excitation_coef)
leftover_suppression = tf.zeros_like(params.suppression_coef)
time_of_leftover = historic_times[0]
n_evaluations = (n_intervals - 1 if historic_types is None else n_intervals) - pred_window_multiple + 1 # if we are given (historic_types, historic_times), then we may evaluate every single interval by kickstarting on e.g. the training set. Helps especially when the test set is tiny.
n_processes = params.excitation_coef.shape[0] if len(params.excitation_coef.shape) == 3 else 1 # static version of `tf.rank`
scores = tf.zeros([n_evaluations, n_processes] + list(score_shape), dtype=event_times.dtype)
anchor_time = tf.cast(0, dtype=event_times.dtype) # should really be set with the first interval, otherwise might produce garbage until changed
prev_hard_window = tf.cast(window, event_times.dtype)
for interval in tf.range(n_evaluations): # all completely dynamic. funny that this is a desideratum now that static unrolling has become a bottleneck, contraty to how it is in most programming environments!
if historic_times is None:
prelude_types, prelude_times = interval_event_types[interval, ...], interval_event_times[interval, ...]
else:
if interval == 0:
prelude_types, prelude_times = historic_types, historic_times
else:
prelude_types, prelude_times = interval_event_types[interval-1, ...], interval_event_times[interval-1, ...]
if tf.size(prelude_times) == 0:
hard_window = prev_hard_window
else:
anchor_time = prelude_times[-1] # basically `prelude_times[-1]`, with the added functionality of tracking the last non-empty one
# this shouldn't be needed now that the last interval is thrown away anyways. keep it as a failsafe
hard_window = tf.math.minimum(event_times[-1] - anchor_time, window * pred_window_multiple) # if we are at the end of the record, our last interval must be subjected to a hard cutoff. pred_window_multiple allows us to achieve a crude version of overlapping sliding windows: make the intervals small, then forecast over a collection of them.
sub_windows = tf.linspace(tf.cast(0, event_times.dtype), hard_window, n_splits+1)[1:]
aggregates, max_n_events_counted, leftover_excitation, leftover_suppression, time_of_leftover = \
forecast_window_aggregates(params, prelude_types, prelude_times, batch_size, hard_window, n_events, n_simulations, sub_windows,
warm_start_excitation=leftover_excitation, warm_start_suppression=leftover_suppression, time_of_warm_start=time_of_leftover, return_leftovers=True)
tf.print("Counted at most", max_n_events_counted, "simulated events out of", n_events, "in interval", interval+1, "out of", n_evaluations, "\b.")
real_deal_intervals = interval_event_types[(interval+1):(interval+1+pred_window_multiple), ...] if historic_times is None \
else interval_event_types[interval:(interval+pred_window_multiple), ...]
real_deal_types = tf.concat([real_deal_intervals[i, ...] for i in range(pred_window_multiple)], axis=0) # how is it so hard to unstack and flatten a ragged tensor?
real_deal_time_intervals = interval_event_times[(interval+1):(interval+1+pred_window_multiple), ...] if historic_times is None \
else interval_event_times[interval:(interval+pred_window_multiple), ...]
real_deal_times = tf.concat([real_deal_time_intervals[i, ...] for i in range(pred_window_multiple)], axis=0)
real_instances = tf.math.equal(real_deal_types, tf.range(params.n_dims, dtype=real_deal_types.dtype)[:, None]) # (n_types x window events)
if n_splits == 1:
real_aggregates = tf.math.count_nonzero(real_instances, axis=1, dtype=event_times.dtype)
else:
real_instance_indicators = tf.transpose(tf.cast(real_instances, tf.int32)) # (window events x n_types)
real_assignments = tf.searchsorted(sub_windows, real_deal_times - anchor_time) # shift the frame of reference
split_real_aggregates = tf.math.segment_sum( # see `forecast_window_aggregates` for explanation
tf.concat([real_instance_indicators, tf.broadcast_to([[0]], [1, params.n_dims])], axis=0),
tf.concat([real_assignments, [n_splits]], axis=0) ) # (n_splits+1 x n_types)
real_aggregates = tf.cast(tf.transpose(split_real_aggregates[:-1, :]), event_times.dtype) # touched-up/refined version of the above
interval_scores = score_against(aggregates, real_aggregates)# real_aggregate if n_splits==1 else real_aggregates -- for some reason this is not treated purely like a static evaluation.. must be something with the decorator
# ^ must return a tensor with one entry (row? hyper-row?) for each process
selector = tf.reshape(tf.one_hot(interval, n_evaluations, dtype=scores.dtype), [-1, 1] + [1]*len(score_shape))
scores += tf.ensure_shape(interval_scores, [n_processes] + score_shape) * selector
prev_hard_window = hard_window
return scores
def _cdf(self, x):
x = tf.convert_to_tensor(value=x, name='x')
flat_x = tf.reshape(x, shape=[-1])
upper_bound = tf.searchsorted(self.outcomes, values=flat_x, side='right')
values_at_ub = tf.gather(
self.outcomes,
indices=tf.minimum(upper_bound,
dist_util.prefer_static_shape(self.outcomes)[-1] -
1))
should_use_upper_bound = self._is_equal_or_close(flat_x, values_at_ub)
indices = tf.where(should_use_upper_bound, upper_bound, upper_bound - 1)
return self._categorical.cdf(
tf.reshape(indices, shape=dist_util.prefer_static_shape(x)))
def inverse(x, R_in):
''' Require tensorflow '''
''' Input:
x: a real number
R: a neuron netweek with one hidden layer and RELU activation
Return:
It computes the inverst of R_in by regarding R_in as a piecewise linear continuous function
Example of R_in (built in keras):
R_in = keras.Sequential(
[
keras.Input(shape=(1,)),
layers.Dense(H, activation="relu", name="hidden"),
layers.Dense(1, activation=None, name="output"),
]
)
'''
x = tf.dtypes.cast(x, tf.float32)
w, b = R_in.get_layer('hidden').weights
v, v0 = R_in.get_layer('output').weights
num_layer = w.shape[1]
b_over_w = -b / w # the negative sign is important
b_over_w = tf.reshape(b_over_w, [-1]) # change it to row vector
# Then we need to compute range for the inverse function
r_x = tf.reshape([
tf.matmul((tf.nn.relu(w * b_over_w[i] + b)), v) + v0 for i in range(1)
], [-1])
# pad r_x with 0 and inf
# right side gives us the right index
pos = tf.searchsorted(tf.sort(tf.concat([[0], r_x, [np.inf]], 0)),
x,
side='right') # get the interval index
# tf.print('pos',pos)
pos = tf.reshape(pos, []) # get only numerical value
# change it to row vector
v = tf.reshape(v, [-1])
w = tf.reshape(w, [-1])
index = tf.argsort(r_x)
# tf.print(pos,"\n")
# sort according to the index
v_b = tf.gather(v * b, index)
v_w = tf.gather(v * w, index)
num = x - v0 - tf.reduce_sum(v_b[:(pos - 1)])
deo = tf.reduce_sum(v_w[:(pos - 1)])
return tf.math.divide_no_nan(num, deo)[0]
