def test_2d_scalar_valued_function(self):
x_ref_min = np.array([0., 1.], dtype=np.float32)
x_ref_max = np.array([1.3, 2.], dtype=np.float32)
ny = [100, 110]
# Build y_ref.
x0s, x1s = tf.meshgrid(
tf.linspace(x_ref_min[0], x_ref_max[0], ny[0]),
tf.linspace(x_ref_min[1], x_ref_max[1], ny[1]),
indexing='ij')
def func(x0, x1):
return tf.sin(x0) * tf.cos(x1)
# Shape ny
y_ref = self.evaluate(func(x0s, x1s))
# Shape [10, 2]
seed = test_util.test_seed_stream()
x = tf.stack([
tf.random.uniform(
shape=(10,), minval=x_ref_min[0], maxval=x_ref_max[0], seed=seed()),
tf.random.uniform(
shape=(10,), minval=x_ref_min[1], maxval=x_ref_max[1], seed=seed()),
],
axis=-1)
x = self.evaluate(x)
expected_y = func(x[:, 0], x[:, 1])
actual_y = tfp.math.batch_interp_regular_nd_grid(
x=x, x_ref_min=x_ref_min, x_ref_max=x_ref_max, y_ref=y_ref, axis=-2)
self.assertAllClose(
self.evaluate(expected_y), self.evaluate(actual_y), atol=0.02)
def first_order_coeff_fn(t, coord_grid):
del t
y, x = tf.meshgrid(*coord_grid, indexing='ij')
return [
y**2 * (1 + 3 * x**2) + 2 * tf.math.sin(y),
x * (1 + 2 * x**2 * y) - 2 * tf.math.cos(x)
]
def _generate_boxes(self):
"""Generates multiscale anchor boxes.
Returns:
a Tensor of shape [N, 4], represneting anchor boxes of all levels
concatenated together.
"""
boxes_all = []
for level in range(self.min_level, self.max_level + 1):
boxes_l = []
for scale in range(self.num_scales):
for aspect_ratio in self.aspect_ratios:
stride = 2 ** level
intermidate_scale = 2 ** (scale / float(self.num_scales))
base_anchor_size = self.anchor_size * stride * intermidate_scale
aspect_x = aspect_ratio ** 0.5
aspect_y = aspect_ratio ** -0.5
half_anchor_size_x = base_anchor_size * aspect_x / 2.0
half_anchor_size_y = base_anchor_size * aspect_y / 2.0
x = tf.range(stride / 2, self.image_size[1], stride)
y = tf.range(stride / 2, self.image_size[0], stride)
xv, yv = tf.meshgrid(x, y)
xv = tf.cast(tf.reshape(xv, [-1]), dtype=tf.float32)
yv = tf.cast(tf.reshape(yv, [-1]), dtype=tf.float32)
# Tensor shape Nx4.
boxes = tf.stack([yv - half_anchor_size_y, xv - half_anchor_size_x,
yv + half_anchor_size_y, xv + half_anchor_size_x],
axis=1)
boxes_l.append(boxes)
# Concat anchors on the same level to tensor shape NxAx4.
boxes_l = tf.stack(boxes_l, axis=1)
boxes_l = tf.reshape(boxes_l, [-1, 4])
boxes_all.append(boxes_l)
return tf.concat(boxes_all, axis=0)
def binary_image_to_points(features, normalize_coords=True, keys=("image", )):
"""Converts a (binary) image into a 2D point cloud.
Args:
features: Dictionary of data features to preprocess.
normalize_coords: Normalize coords to be in [0,1] by preserving the aspect
ratio.
keys: On which keys to apply this function.
Returns:
Features with the image as a point cloud.
"""
for key in keys:
image = features[key] # [HxW] or [HxWx1]
image = tf.reshape(image, [image.shape[0], image.shape[1], 1])
# We map background pixels to the origin, which may be suboptimal
# but saves us some engineering work.
coords = tf.cast(
tf.stack(tf.meshgrid(tf.range(image.shape[0]),
tf.range(image.shape[1]),
indexing="ij"),
axis=-1), tf.float32)
if normalize_coords:
coords /= tf.cast(tf.reduce_max(image.shape[:2]), tf.float32)
mask = tf.tile(image > 0, [1, 1, 2])
features[key] = tf.reshape(tf.cast(mask, tf.float32) * coords, [-1, 2])
return features
def _generate_boxes(self):
"""Generates multiscale anchor boxes.
Returns:
a Tensor of shape [N, 4], represneting anchor boxes of all levels
concatenated together.
"""
boxes_all = []
for level in range(self.min_level, self.max_level + 1):
boxes_l = []
for anchor_size in anchors[self.anchor_masks[level - 2]]:
stride = 2**level
half_anchor_size_x = anchor_size / 2.0
half_anchor_size_y = anchor_size / 2.0
x = tf.range(stride / 2, self.image_size[1], stride)
y = tf.range(stride / 2, self.image_size[0], stride)
xv, yv = tf.meshgrid(x, y)
xv = tf.cast(tf.reshape(xv, [-1]), dtype=tf.float32)
yv = tf.cast(tf.reshape(yv, [-1]), dtype=tf.float32)
# Tensor shape Nx4.
boxes = tf.stack([
yv - half_anchor_size_y, xv - half_anchor_size_x,
yv + half_anchor_size_y, xv + half_anchor_size_x
],
axis=1)
boxes_l.append(boxes)
# Concat anchors on the same level to tensor shape NxAx4.
boxes_l = tf.stack(boxes_l, axis=1)
boxes_l = tf.reshape(boxes_l, [-1, 4])
boxes_all.append(boxes_l)
return tf.concat(boxes_all, axis=0)
def basis(sample_paths, time_index):
"""Computes polynomial basis expansion at the given sample points.
Args:
sample_paths: A `Tensor` of either `flaot32` or `float64` dtype and of
shape `[num_samples, num_times, dim]`.
time_index: An integer scalar `Tensor` that corresponds to the time
coordinate at which the basis function is computed.
Returns:
A `Tensor`s of shape `[(degree + 1)**dim, num_samples]`.
"""
sample_paths = tf.convert_to_tensor(sample_paths, name="sample_paths")
shape = sample_paths.shape.as_list()
num_samples = shape[0]
dim = shape[-1]
slice_samples = tf.slice(sample_paths, [0, time_index, 0],
[num_samples, 1, dim])
slice_samples = tf.squeeze(slice_samples, 1)
dim = sample_paths.shape.as_list()[-1]
samples_centered = slice_samples - tf.math.reduce_mean(slice_samples,
axis=0)
samples_centered = tf.expand_dims(samples_centered, axis=-2)
grid = tf.range(degree + 1, dtype=samples_centered.dtype)
# Creates a grid of 'power' expansions, i.e., a `Tensor` of shape
# [(degree + 1)**dim, dim] with entries [k_1, .., k_dim] where
## 0 <= k_i <= dim.
grid = tf.meshgrid(*(dim * [grid]))
# Shape [(degree + 1)**3, dim]
grid = tf.reshape(tf.stack(grid, -1), [-1, dim])
# `samples_centered` has shape [num_samples, 1, dim],
# `samples_centered**grid` has shape `[num_samples, (degree + 1)**dim, dim]`
# so that the output shape is [num_samples, (degree + 1)**dim]
basis_expansion = tf.reduce_prod(samples_centered**grid, -1)
return tf.transpose(basis_expansion)
def plot_prior_latent(self, intervals, figsize=None, **kwargs):
if len(intervals) != self.ndim_latent or self.ndim_latent not in (1,
2):
raise ValueError(
"This method is only defined for 1D or 2D models.")
y = [
tf.linspace(float(i[0]), float(i[1]), int(i[2])) for i in intervals
]
y = tf.meshgrid(*y, indexing="ij")
y = tf.stack(y, axis=-1)
prob = tfpd.Independent(self.prior(**kwargs), 1).prob(y)
if self.ndim_latent == 1:
plt.figure(figsize=figsize or (16, 6))
plt.plot(y, prob)
plt.xlabel("latent space")
plt.ylabel("prior")
plt.grid(True)
plt.xticks(np.arange(*np.ceil(intervals[0][:2])))
else:
plt.figure(figsize=figsize or (16, 14))
plt.imshow(prob,
vmin=0,
vmax=np.max(prob),
origin="lower",
extent=(y[0, 0, 1], y[0, -1, 1], y[0, 0, 0], y[-1, 0,
0]))
plt.axis("image")
plt.grid(False)
plt.xlim(y[0, 0, 1], y[0, -1, 1])
plt.ylim(y[0, 0, 0], y[-1, 0, 0])
plt.xlabel("latent dimension 1")
plt.ylabel("latent dimension 2")
def _cartesian_product(*supports):
"""Construct "cartesian product" of tensors.
Args:
*supports: a sequence of tensors `s1, ..., sn`.
Returns:
This function computes a tensor analogous to the cartesian
product of sets.
If `t = _cartesian_product(s1, ..., sn)` then
`t[i1, ..., in] = s1[i1] s2[i2] ... sn[in]`
where the elements on the right hand side are concatenated
together.
In particular, if `s1, ..., sn` are the supports of `n`
distributions, the cartesian product represents the support of the
product distribution.
For example if `a = [0, 1]`, b = [10, 20]` and
`c = _cartesian_product(a, b)` then
`c = [[[0, 10], [0, 20]], [[1, 10], [1, 20]]]`.
In this case note (for example) that
`a[0] = 0`, `b[1] = 20` and so `c[0, 1] = [0, 20]`.
"""
return tf.stack(tf.meshgrid(*supports, indexing='ij'), axis=-1)
def _coord_grid_to_mesh_grid(coord_grid):
if len(coord_grid) == 1:
return tf.expand_dims(coord_grid[0], 0)
x_meshgrid = tf.stack(values=tf.meshgrid(*coord_grid, indexing='ij'),
axis=-1)
perm = [len(coord_grid)] + list(range(len(coord_grid)))
return tf.transpose(x_meshgrid, perm=perm)
def create_hmm(self, num_steps):
"""Same as the original CREPE viterbi decdoding, but in TF."""
# Initial distribution is uniform.
initial_distribution = tfp.distributions.Categorical(
probs=tf.ones([360]) / 360)
# Transition probabilities inducing continuous pitch.
bins = tf.range(360, dtype=tf.float32)
xx, yy = tf.meshgrid(bins, bins)
min_transition = 1e-5 # For training stabiity.
transition = tf.maximum(12 - abs(xx - yy), min_transition)
transition = transition / tf.reduce_sum(transition, axis=1)[:, None]
transition = tf.cast(transition, tf.float32)
transition_distribution = tfp.distributions.Categorical(
probs=transition)
# Emission probability = fixed probability for self, evenly distribute the
# others.
self_emission = 0.1
emission = (tf.eye(360) * self_emission + tf.ones(shape=(360, 360)) *
((1 - self_emission) / 360.))
emission = tf.cast(emission, tf.float32)[None, ...]
observation_distribution = tfp.distributions.Multinomial(
total_count=1, probs=emission)
return tfp.distributions.HiddenMarkovModel(
initial_distribution=initial_distribution,
transition_distribution=transition_distribution,
observation_distribution=observation_distribution,
num_steps=num_steps,
)
def _two_d_integration(grid, value_grid):
"""Perform 2-D integration numerically."""
log_spot_grid, variance_grid = tf.meshgrid(*grid)
delta_v = variance_grid[1:, :] - variance_grid[:-1, :]
delta_s = log_spot_grid[:, 1:] - log_spot_grid[:, :-1]
integral = tf.math.reduce_sum(value_grid[0, :-1, :] * delta_v, axis=0)
integral = tf.math.reduce_sum(integral[:-1] * delta_s[0, :])
return integral
def call(self, inputs):
value, index = inputs
if self.cache.shape == inputs[0].shape:
self.cache = value
return value
shape = self.cache.shape.as_list()
num_index_axes = index.shape[0]
num_batch_axes = self.num_batch_axes
num_feature_axes = len(shape) - num_index_axes - num_batch_axes
features_shape = shape[num_batch_axes + num_index_axes:]
batch_shape = shape[:num_batch_axes]
value_index_shape = tf.shape(value)[num_batch_axes:-num_feature_axes]
if tf.reduce_max(value_index_shape) > 1:
# This is a block update starting at index.
value_ranges = []
for i, s in enumerate(tf.unstack(value_index_shape)):
curr_range = tf.range(index[i], index[i] + s)
value_ranges.append(curr_range)
batch_ranges = [tf.range(s) for s in batch_shape]
mesh = tf.meshgrid(*(batch_ranges + value_ranges), indexing='ij')
indices = tf.stack(mesh, axis=-1)
indices = tf.reshape(indices,
[-1, num_index_axes + num_batch_axes])
else:
# This is a single update at index position.
batch_ranges = [tf.range(s) for s in batch_shape]
mesh = tf.meshgrid(*batch_ranges, indexing='ij')
batch_indices = tf.stack(mesh, axis=-1)
batch_indices = tf.reshape(batch_indices, [-1, num_batch_axes])
# Add leading axes to nd-index and tile to get batched indices.
shape_indices = tf.reshape(index, [1] * num_batch_axes + [-1])
shape_indices = tf.tile(shape_indices, batch_shape + [1])
shape_indices = tf.reshape(shape_indices, [-1, num_index_axes])
indices = tf.concat([batch_indices, shape_indices], axis=-1)
# We need to squeeze nd-axes from value before updating.
value = tf.reshape(value, [-1] + features_shape)
self.cache = tf.tensor_scatter_nd_update(self.cache, indices, value)
return self.cache
def second_order_coeff_fn(t, location_grid):
del t
z, y, x = tf.meshgrid(*location_grid, indexing='ij')
u_zz = 1
u_xx = 1
u_yy = 1
u_xy = tf.math.sin(x) * tf.math.cos(y) / 2
u_yz = tf.math.cos(y) * tf.math.cos(z) / 2
return [[u_zz, u_yz, None], [u_yz, u_yy, u_xy], [None, u_xy, u_xx]]
def test_3d_vector_valued_function_and_fill_value(self):
x_ref_min = np.array([1.0, 0.0, -1.2], dtype=np.float32)
x_ref_max = np.array([2.3, 3.0, 1.0], dtype=np.float32)
ny = [200, 210, 211]
# Build y_ref.
x0s, x1s, x2s = tf.meshgrid(tf.linspace(x_ref_min[0], x_ref_max[0],
ny[0]),
tf.linspace(x_ref_min[1], x_ref_max[1],
ny[1]),
tf.linspace(x_ref_min[2], x_ref_max[2],
ny[2]),
indexing='ij')
def func(x0, x1, x2):
# Shape [..., 2] output.
return tf.stack([tf.sin(x0 * x1 * x2),
tf.cos(x0 * x1 * x2)],
axis=-1)
# Shape ny + [2]
y_ref = self.evaluate(func(x0s, x1s, x2s))
seed = test_util.test_seed_stream()
# Shape [10, 3]
x = tf.stack([
tf.random.uniform(shape=(10, ),
minval=x_ref_min[0],
maxval=x_ref_max[0],
seed=seed()),
tf.random.uniform(shape=(10, ),
minval=x_ref_min[1],
maxval=x_ref_max[1],
seed=seed()),
tf.random.uniform(shape=(10, ),
minval=x_ref_min[2],
maxval=x_ref_max[2],
seed=seed()),
],
axis=-1)
x = np.array(self.evaluate(x))
x[0, 0] = -3 # Outside the grid, so `fill_value` will be imputed.
expected_y = np.array(self.evaluate(func(x[:, 0], x[:, 1], x[:, 2])))
fill_value = -42
expected_y[0, :] = fill_value
actual_y = tfp.math.batch_interp_regular_nd_grid(x=x,
x_ref_min=x_ref_min,
x_ref_max=x_ref_max,
y_ref=y_ref,
axis=-4,
fill_value=fill_value)
self.assertAllClose(expected_y, self.evaluate(actual_y), atol=0.02)
def call(self, x, training=None):
x = self._conv2(x)
x = self._conv3(x)
x = self._conv4(x)
posx, posy = tf.meshgrid(tf.linspace(-1., 1., num=16),
tf.linspace(-1., 1., num=16))
stacked_pos = tf.expand_dims(tf.stack([posx, posy], axis=-1), axis=0)
stacked_pos = tf.tile(stacked_pos, [tf.shape(x)[0], 1, 1, 1])
x = tf.concat([x, stacked_pos], axis=-1)
x = self._conv5(x)
return self._reshape(x)
def test_multidim_with_x(self):
y = tf.ones((4, 5), dtype=tf.float64)
v0 = tf.cast(tf.linspace(0., 1., 4), tf.float64)
v1 = tf.cast(tf.linspace(0., 4.2, 5), tf.float64)
x0, x1 = tf.meshgrid(v0, v1, indexing='ij')
integral_0 = self.evaluate(tfp_math.trapz(y, x0, axis=0))
integral_1 = self.evaluate(tfp_math.trapz(y, x1, axis=1))
self.assertTupleEqual(integral_0.shape, (5, ))
self.assertTupleEqual(integral_1.shape, (4, ))
self.assertAllClose(integral_0, np.ones((5, )) * 1.0)
self.assertAllClose(integral_1, np.ones((4, )) * 4.2)
def _initial_value():
"""Computes initial value as a delta function delta(log_spot(t), var(0))."""
log_spot, variance = tf.meshgrid(*grid)
init_value = tf.where(
tf.math.logical_and(
tf.math.abs(log_spot - scaled_initial_point[0]) <
delta_x + pde_grid_tol,
tf.math.abs(variance - scaled_initial_point[1]) <
delta_y + pde_grid_tol), 1.0 / (delta_x * delta_y * 4), 0.0)
# initial_value.shape = (1, num_grid_x, num_grid_y)
return tf.expand_dims(init_value, axis=0)
def test_compare_monte_carlo_to_backward_pde(self):
dtype = tf.float64
kappa = 0.3
theta = 0.05
epsilon = 0.02
rho = 0.1
maturity_time = 1.0
initial_log_spot = 3.0
initial_vol = 0.05
strike = 15
discounting = 0.5
heston = HestonModel(kappa=kappa,
theta=theta,
epsilon=epsilon,
rho=rho,
dtype=dtype)
initial_state = np.array([initial_log_spot, initial_vol])
samples = heston.sample_paths(
times=[maturity_time / 2, maturity_time],
initial_state=initial_state,
time_step=0.01,
num_samples=1000,
random_type=tff.math.random.RandomType.PSEUDO_ANTITHETIC,
seed=42)
self.assertEqual(samples.shape, [1000, 2, 2])
log_spots = samples[:, -1, 0]
monte_carlo_price = (
tf.constant(np.exp(-discounting * maturity_time), dtype=dtype) *
tf.math.reduce_mean(tf.nn.relu(tf.math.exp(log_spots) - strike)))
s_min, s_max = 2, 4
v_min, v_max = 0.03, 0.07
grid_size_s, grid_size_v = 101, 101
time_step = 0.01
grid = grids.uniform_grid(minimums=[s_min, v_min],
maximums=[s_max, v_max],
sizes=[grid_size_s, grid_size_v],
dtype=dtype)
s_mesh, _ = tf.meshgrid(grid[0], grid[1], indexing="ij")
final_value_grid = tf.nn.relu(tf.math.exp(s_mesh) - strike)
value_grid = heston.fd_solver_backward(
start_time=1.0,
end_time=0.0,
coord_grid=grid,
values_grid=final_value_grid,
time_step=time_step,
discounting=lambda *args: discounting)[0]
pde_price = value_grid[int(grid_size_s / 2), int(grid_size_v / 2)]
self.assertAllClose(monte_carlo_price, pde_price, atol=0.1, rtol=0.1)
def testReferenceEquation_WithTransformationYieldingMixedTerm(self):
"""Tests an equation with mixed terms against exact solution.
Take the reference equation `v_{t} = v_{xx} + v_{yy}` and substitute
`v(x, y, t) = u(x, 2y - x, t)`. This yields
`u_{t} = u_{xx} + 5u_{zz} - 2u_{xz}`, where `z = 2y - x`.
Having `u(x, z, t) = v(x, (x+z)/2, t)` where `v(x, y, t)` is the known
solution of the reference equation, we derive the boundary conditions
and the expected solution for `u(x, y, t)`.
"""
grid = grids.uniform_grid(minimums=[0, 0],
maximums=[1, 1],
sizes=[201, 251],
dtype=tf.float32)
final_t = 0.1
time_step = 0.002
def second_order_coeff_fn(t, coord_grid):
del t, coord_grid
return [[-5, 1], [None, -1]]
@dirichlet
def boundary_lower_z(t, coord_grid):
x = coord_grid[1]
return _reference_pde_solution(x, t) * _reference_pde_solution(
x / 2, t)
@dirichlet
def boundary_upper_z(t, coord_grid):
x = coord_grid[1]
return _reference_pde_solution(x, t) * _reference_pde_solution(
(x + 1) / 2, t)
z_mesh, x_mesh = tf.meshgrid(grid[0], grid[1], indexing='ij')
initial = (_reference_pde_initial_cond(x_mesh) *
_reference_pde_initial_cond((x_mesh + z_mesh) / 2))
expected = (_reference_pde_solution(x_mesh, final_t) *
_reference_pde_solution((x_mesh + z_mesh) / 2, final_t))
actual = fd_solvers.solve_forward(
start_time=0,
end_time=final_t,
coord_grid=grid,
values_grid=initial,
time_step=time_step,
second_order_coeff_fn=second_order_coeff_fn,
boundary_conditions=[(boundary_lower_z, boundary_upper_z),
(_zero_boundary, _zero_boundary)])[0]
self.assertAllClose(expected, actual, atol=1e-3, rtol=1e-3)
def testLogitNormalMeanAndVariance(self):
locs, scales = tf.meshgrid(tf.linspace(-10.0, 10.0, 10),
tf.exp(tf.linspace(-3.0, 3.0, 10)))
dist = tfd.LogitNormal(loc=locs,
scale=scales,
validate_args=True,
gauss_hermite_scale_limit=1.,
num_probit_terms_approx=6)
means = dist.mean_approx()
trap_means = logit_normal_mean_trapezoid(locs, scales)
self.assertAllClose(trap_means, means, rtol=1e-4)
variances = dist.variance_approx()
trap_variances = logit_normal_variance_trapezoid(locs, scales)
self.assertAllClose(trap_variances, variances, rtol=1e-4)
def test_2d_vector_valued_function_with_batch_dims(self):
# No batch dims, will broadcast.
x_ref_min = np.array([0., 0.], dtype=np.float32)
# No batch dims, will broadcast.
x_ref_max = np.array([1., 1.], dtype=np.float32)
ny = [200, 210]
# Build y_ref.
# First step is to build up two batches of x0 and x1.
x0s, x1s = tf.meshgrid(tf.linspace(x_ref_min[0], x_ref_max[0], ny[0]),
tf.linspace(x_ref_min[1], x_ref_max[1], ny[1]),
indexing='ij')
x0s = tf.stack([x0s, x0s], axis=0)
x1s = tf.stack([x1s, x1s], axis=0)
def func(batch_of_x0, batch_of_x1):
"""Function that does something different for batch 0 and batch 1."""
# batch_0_result.shape = [..., 2].
x0, x1 = batch_of_x0[0, ...], batch_of_x1[0, ...]
batch_0_result = tf.stack(
[tf.sin(x0 * x1), tf.cos(x0 * x1)], axis=-1)
x0, x1 = batch_of_x0[1, ...], batch_of_x1[1, ...]
batch_1_result = tf.stack(
[tf.sin(2 * x0), tf.cos(2 * x1)], axis=-1)
return tf.stack([batch_0_result, batch_1_result], axis=0)
# Shape [2] + ny + [2]
y_ref = self.evaluate(func(x0s, x1s))
# Shape [2, 10, 2]. The batch shape is [2], the [10] is the number of
# interpolants per batch.
x = tf.random.uniform(shape=[2, 10, 2], seed=test_util.test_seed())
x = self.evaluate(x)
expected_y = func(x[..., 0], x[..., 1])
actual_y = tfp.math.batch_interp_regular_nd_grid(x=x,
x_ref_min=x_ref_min,
x_ref_max=x_ref_max,
y_ref=y_ref,
axis=-3)
self.assertAllClose(self.evaluate(expected_y),
self.evaluate(actual_y),
atol=0.02)
def _conditional_expected_variance_from_pde_solution(grid, value_grid, dtype):
"""Computes E[variance|log_spot=k]."""
# value_grid.shape = [1, num_x, num_y]
log_spot_grid, variance_grid = tf.meshgrid(*grid)
delta_s = variance_grid[1:, :] - variance_grid[:-1, :]
# Calculate I(0)
integral_0 = tf.math.reduce_sum(value_grid[0, :-1, :] * delta_s, axis=0)
# Calculate I(1)
integral_1 = tf.math.reduce_sum(variance_grid[:-1, :] *
value_grid[0, :-1, :] * delta_s,
axis=0)
variance_given_logspot = tf.math.divide_no_nan(integral_1, integral_0)
return functools.partial(linear.interpolate,
x_data=log_spot_grid[0, :],
y_data=variance_given_logspot,
dtype=dtype)
def meshgrid(*xi, **kwargs):
"""This currently requires copy=True and sparse=False."""
sparse = kwargs.get('sparse', False)
if sparse:
raise ValueError('tf.numpy doesnt support returning sparse arrays yet')
copy = kwargs.get('copy', True)
if not copy:
raise ValueError('tf.numpy only supports copy=True')
indexing = kwargs.get('indexing', 'xy')
xi = [array_ops.asarray(arg).data for arg in xi]
kwargs = {'indexing': indexing}
outputs = tf.meshgrid(*xi, **kwargs)
outputs = [utils.tensor_to_ndarray(output) for output in outputs]
return outputs
def __getitem__(self, slices):
if not isinstance(slices, tuple):
slices = [slices]
else:
slices = list(slices)
if Ellipsis in slices:
idx = slices.index(Ellipsis)
slices[idx:idx+1] = [slice(None)] * (self.rank - len(slices) + 1)
# Remove trailing full slices for performance.
while (slices
and isinstance(slices[-1], slice)
and slices[-1] == slice(None)):
slices.pop()
grid = tf.meshgrid(*(rng[sl] for rng, sl in zip(self.ranges, slices)),
indexing='ij')
stack = tf.stack(grid, axis=-1)
to_squeeze = [i for i, sl in enumerate(slices) if not isinstance(sl, slice)]
if to_squeeze:
stack = tf.squeeze(stack, axis=to_squeeze)
return stack
def basis(sample_paths):
"""Computes polynomial basis expansion at the given sample points.
Args:
sample_paths: A `Tensor`s of either `flot32` or `float64` dtype and of
shape `[num_samples, dim]` where `dim` has to be statically known.
Returns:
A `Tensor`s of shape `[degree * dim, num_samples]`.
"""
samples = tf.convert_to_tensor(sample_paths)
dim = samples.shape.as_list()[-1]
grid = tf.range(0, degree + 1, dtype=samples.dtype)
samples_centered = samples - tf.math.reduce_mean(samples, axis=0)
samples_centered = tf.expand_dims(samples_centered, -2)
grid = tf.meshgrid(*(dim * [grid]))
grid = tf.reshape(tf.stack(grid, -1), [-1, dim])
# Shape [num_samples, degree * dim]
basis_expansion = tf.reduce_prod(samples_centered**grid, -1)
return tf.transpose(basis_expansion)
def test_2d_vector_valued_function(self):
x_ref_min = [1., 0.]
x_ref_max = [2.3, 1.]
ny = [200, 210]
# Build y_ref.
x0s, x1s = tf.meshgrid(tf.linspace(x_ref_min[0], x_ref_max[0], ny[0]),
tf.linspace(x_ref_min[1], x_ref_max[1], ny[1]),
indexing='ij')
def func(x0, x1):
# Shape [..., 2] output.
return tf.stack([tf.sin(x0 * x1), tf.cos(x0 * x1)], axis=-1)
# Shape ny + [2]
y_ref = self.evaluate(func(x0s, x1s))
# Shape [10, 2]
x = tf.stack([
tf.random.uniform(
shape=(10, ), minval=x_ref_min[0], maxval=x_ref_max[0],
seed=0),
tf.random.uniform(
shape=(10, ), minval=x_ref_min[1], maxval=x_ref_max[1],
seed=1),
],
axis=-1)
x = self.evaluate(x)
expected_y = func(x[:, 0], x[:, 1])
actual_y = tfp.math.batch_interp_regular_nd_grid(x=x,
x_ref_min=x_ref_min,
x_ref_max=x_ref_max,
y_ref=y_ref,
axis=-3)
self.assertAllClose(self.evaluate(expected_y),
self.evaluate(actual_y),
atol=0.02)
def basis(sample_paths: types.RealTensor,
time_index: types.IntTensor) -> types.RealTensor:
"""Computes polynomial basis expansion at the given sample points.
Args:
sample_paths: A `Tensor` of either `flaot32` or `float64` dtype and of
either shape `[num_samples, num_times, dim]` or
`[batch_size, num_samples, num_times, dim]`.
time_index: An integer scalar `Tensor` that corresponds to the time
coordinate at which the basis function is computed.
Returns:
A `Tensor`s of shape `[batch_size, (degree + 1)**dim, num_samples]`.
"""
sample_paths = tf.convert_to_tensor(sample_paths, name="sample_paths")
if sample_paths.shape.rank == 3:
sample_paths = tf.expand_dims(sample_paths, axis=0)
shape = tf.shape(sample_paths)
num_samples = shape[1]
batch_size = shape[0]
dim = sample_paths.shape[-1] # Dimension should statically known
# Shape [batch_size, num_samples, 1, dim]
slice_samples = tf.slice(sample_paths, [0, 0, time_index, 0],
[batch_size, num_samples, 1, dim])
# Shape [batch_size, num_samples, 1, dim]
samples_centered = slice_samples - tf.math.reduce_mean(
slice_samples, axis=1, keepdims=True)
grid = tf.range(degree + 1, dtype=samples_centered.dtype)
# Creates a grid of 'power' expansions, i.e., a `Tensor` of shape
# [(degree + 1)**dim, dim] with entries [k_1, .., k_dim] where
## 0 <= k_i <= dim.
grid = tf.meshgrid(*(dim * [grid]))
# Shape [(degree + 1)**3, dim]
grid = tf.reshape(tf.stack(grid, -1), [-1, dim])
# `samples_centered` has shape [batch_size, num_samples, 1, dim],
# `samples_centered**grid` has shape
# `[batch_size, num_samples, (degree + 1)**dim, dim]`
# so that the output shape is `[batch_size, num_samples, (degree + 1)**dim]`
basis_expansion = tf.reduce_prod(samples_centered**grid, axis=-1)
return tf.transpose(basis_expansion, [0, 2, 1])
def testLogProb(self):
# Test that numerically integrating over some portion of the domain yields a
# normalization constant of close to 1.
# pyformat: disable
scale = tf.linalg.LinearOperatorFullMatrix(
self._input([[1., -0.5], [-0.5, 1.]]))
# pyformat: enable
dist = tfd.MultivariateStudentTLinearOperator(loc=self._input([1.,
1.]),
df=self._input(5.),
scale=scale)
spacings = tf.cast(tf.linspace(-20., 20., 100), self.dtype)
x, y = tf.meshgrid(spacings, spacings)
points = tf.concat([x[..., tf.newaxis], y[..., tf.newaxis]], -1)
log_probs = dist.log_prob(points)
normalization = tf.exp(
tf.reduce_logsumexp(log_probs)) * (spacings[1] - spacings[0])**2
self.assertAllClose(1., self.evaluate(normalization), atol=1e-3)
mode_log_prob = dist.log_prob(dist.mode())
self.assertTrue(np.all(self.evaluate(mode_log_prob >= log_probs)))
def testLogitNormalVarianceGH(self):
locs, scales = tf.meshgrid(tf.linspace(-10.0, 10.0, 10),
tf.exp(tf.linspace(-3.0, 0.0, 10)))
ghs = ln_lib.logit_normal_variance_gh(locs, scales, deg=50)
traps = logit_normal_variance_trapezoid(locs, scales)
self.assertAllClose(traps, ghs, rtol=1e-4)
def call(self, net, training):
keep_prob = self.keep_prob
dropblock_size = self.dropblock_size
data_format = self.data_format
if not training or keep_prob is None:
return net
tf.logging.info(
'Applying DropBlock: dropblock_size {}, net.shape {}'.format(
dropblock_size, net.shape))
if data_format == 'channels_last':
_, width, height, _ = net.get_shape().as_list()
else:
_, _, width, height = net.get_shape().as_list()
if width != height:
raise ValueError(
'Input tensor with width!=height is not supported.')
dropblock_size = min(dropblock_size, width)
# seed_drop_rate is the gamma parameter of DropBlcok.
seed_drop_rate = (1.0 - keep_prob) * width**2 / dropblock_size**2 / (
width - dropblock_size + 1)**2
# Forces the block to be inside the feature map.
w_i, h_i = tf.meshgrid(tf.range(width), tf.range(width))
valid_block_center = tf.logical_and(
tf.logical_and(w_i >= int(dropblock_size // 2),
w_i < width - (dropblock_size - 1) // 2),
tf.logical_and(h_i >= int(dropblock_size // 2),
h_i < width - (dropblock_size - 1) // 2))
valid_block_center = tf.expand_dims(valid_block_center, 0)
valid_block_center = tf.expand_dims(
valid_block_center, -1 if data_format == 'channels_last' else 0)
randnoise = tf.random_uniform(net.shape, dtype=tf.float32)
block_pattern = (
1 - tf.cast(valid_block_center, dtype=tf.float32) + tf.cast(
(1 - seed_drop_rate), dtype=tf.float32) + randnoise) >= 1
block_pattern = tf.cast(block_pattern, dtype=tf.float32)
if dropblock_size == width:
block_pattern = tf.reduce_min(
block_pattern,
axis=[1, 2] if data_format == 'channels_last' else [2, 3],
keepdims=True)
else:
if data_format == 'channels_last':
ksize = [1, dropblock_size, dropblock_size, 1]
else:
ksize = [1, 1, dropblock_size, dropblock_size]
block_pattern = -tf.nn.max_pool(-block_pattern,
ksize=ksize,
strides=[1, 1, 1, 1],
padding='SAME',
data_format='NHWC' if data_format
== 'channels_last' else 'NCHW')
percent_ones = (tf.cast(tf.reduce_sum((block_pattern)), tf.float32) /
tf.cast(tf.size(block_pattern), tf.float32))
net = net / tf.cast(percent_ones, net.dtype) * tf.cast(
block_pattern, net.dtype)
return net
