def _inverse_event_shape_tensor(self, output_shape):
perm = self._make_perm(tf.size(input=output_shape),
tf.argsort(self.perm))
return tf.gather(output_shape, perm)
def _resample_using_log_points(log_probs, sample_shape, log_points, name=None):
"""Resample from `log_probs` using supplied points in interval `[0, 1]`."""
# We divide up the unit interval [0, 1] according to the provided
# probability distributions using `cumulative_logsumexp`.
# At the end of each division we place a 'marker'.
# We use points on the unit interval supplied by caller.
# We sort the combination of points and markers. The number
# of points between the markers defining a division gives the number
# of samples we require in that division.
# For example, suppose `probs` is `[0.2, 0.3, 0.5]`.
# We divide up `[0, 1]` using 3 markers:
#
# | | |
# 0. 0.2 0.5 1.0 <- markers
#
# Suppose we are given four points: [0.1, 0.25, 0.9, 0.75]
# After sorting the combination we get:
#
# 0.1 0.25 0.75 0.9 <- points
# * | * | * *|
# 0. 0.2 0.5 1.0 <- markers
#
# We have one sample in the first category, one in the second and
# two in the last.
#
# All of these computations are carried out in batched form.
with tf.name_scope(name or 'resample_using_log_points') as name:
points_shape = ps.shape(log_points)
batch_shape, [num_markers] = ps.split(ps.shape(log_probs),
num_or_size_splits=[-1, 1])
# `working_shape` specifies the total number of events
# we will be generating.
working_shape = ps.concat([sample_shape, batch_shape], axis=0)
# `markers_shape` is the shape of the markers we temporarily insert.
markers_shape = ps.concat([working_shape, [num_markers]], axis=0)
markers = ps.concat([
tf.ones(markers_shape, dtype=tf.int32),
tf.zeros(points_shape, dtype=tf.int32)
],
axis=-1)
log_marker_positions = tf.broadcast_to(
log_cumsum_exp(log_probs, axis=-1), markers_shape)
log_markers_and_points = ps.concat([log_marker_positions, log_points],
axis=-1)
# Stable sort is used to ensure that no points get sorted between
# markers that have zero distance between them. This ensures that
# there will never be a sample drawn whose probability is intended
# to be zero even when a point falls on the edge of the
# corresponding zero-width bucket.
indices = tf.argsort(log_markers_and_points, axis=-1, stable=True)
sorted_markers = tf.gather_nd(
markers,
indices[..., tf.newaxis],
batch_dims=(ps.rank_from_shape(sample_shape) +
ps.rank_from_shape(batch_shape)))
markers_and_samples = ps.cast(tf.cumsum(sorted_markers, axis=-1),
dtype=tf.int32)
markers_and_samples = tf.math.minimum(markers_and_samples,
num_markers - np.int32(1))
# Collect up samples, omitting markers.
samples_mask = tf.equal(sorted_markers, 0)
# The following block of code is equivalent to
# `samples = markers_and_samples[samples_mask]` however boolean mask
# indices are not supported by XLA.
# Instead we use `argsort` to pick out the top `num_samples`
# elements of `markers_and_samples` when sorted using `samples_mask`
# as key.
num_samples = points_shape[-1]
sample_locations = tf.argsort(ps.cast(samples_mask, dtype=tf.int32),
direction='DESCENDING',
stable=True)
samples = tf.gather_nd(markers_and_samples,
sample_locations[..., :num_samples, tf.newaxis],
batch_dims=(ps.rank_from_shape(sample_shape) +
ps.rank_from_shape(batch_shape)))
return tf.reshape(samples, points_shape)
def argsort_fn(x, axis=-1): return tf.cast(tf.argsort(x, axis=axis), dtype=x.dtype)
def _inverse(self, y):
return self._transpose(y, tf.argsort(self.perm))
def _inverse(self, y): return self._transpose(y, tf.argsort(self._get_perm()))
def _invert_permutation(perm): # TODO(b/130217510): Remove this function.
return tf.cast(tf.argsort(perm, axis=-1), perm.dtype)
def minimize(objective_function,
initial_simplex=None,
initial_vertex=None,
step_sizes=None,
objective_at_initial_simplex=None,
objective_at_initial_vertex=None,
batch_evaluate_objective=False,
func_tolerance=1e-8,
position_tolerance=1e-8,
parallel_iterations=1,
max_iterations=None,
reflection=None,
expansion=None,
contraction=None,
shrinkage=None,
name=None):
"""Minimum of the objective function using the Nelder Mead simplex algorithm.
Performs an unconstrained minimization of a (possibly non-smooth) function
using the Nelder Mead simplex method. Nelder Mead method does not support
univariate functions. Hence the dimensions of the domain must be 2 or greater.
For details of the algorithm, see
[Press, Teukolsky, Vetterling and Flannery(2007)][1].
Points in the domain of the objective function may be represented as a
`Tensor` of general shape but with rank at least 1. The algorithm proceeds
by modifying a full rank simplex in the domain. The initial simplex may
either be specified by the user or can be constructed using a single vertex
supplied by the user. In the latter case, if `v0` is the supplied vertex,
the simplex is the convex hull of the set:
```None
S = {v0} + {v0 + step_i * e_i}
```
Here `e_i` is a vector which is `1` along the `i`-th axis and zero elsewhere
and `step_i` is a characteristic length scale along the `i`-th axis. If the
step size is not supplied by the user, a unit step size is used in every axis.
Alternately, a single step size may be specified which is used for every
axis. The most flexible option is to supply a bespoke step size for every
axis.
### Usage:
The following example demonstrates the usage of the Nelder Mead minimzation
on a two dimensional problem with the minimum located at a non-differentiable
point.
```python
# The objective function
def sqrt_quadratic(x):
return tf.sqrt(tf.reduce_sum(x ** 2, axis=-1))
start = tf.constant([6.0, -21.0]) # Starting point for the search.
optim_results = tfp.optimizer.nelder_mead_minimize(
sqrt_quadratic, initial_vertex=start, func_tolerance=1e-8,
batch_evaluate_objective=True)
# Check that the search converged
assert(optim_results.converged)
# Check that the argmin is close to the actual value.
np.testing.assert_allclose(optim_results.position, np.array([0.0, 0.0]),
atol=1e-7)
# Print out the total number of function evaluations it took.
print("Function evaluations: %d" % optim_results.num_objective_evaluations)
```
### References:
[1]: William Press, Saul Teukolsky, William Vetterling and Brian Flannery.
Numerical Recipes in C++, third edition. pp. 502-507. (2007).
http://numerical.recipes/cpppages/chap0sel.pdf
[2]: Jeffrey Lagarias, James Reeds, Margaret Wright and Paul Wright.
Convergence properties of the Nelder-Mead simplex method in low dimensions,
Siam J. Optim., Vol 9, No. 1, pp. 112-147. (1998).
http://www.math.kent.edu/~reichel/courses/Opt/reading.material.2/nelder.mead.pdf
[3]: Fuchang Gao and Lixing Han. Implementing the Nelder-Mead simplex
algorithm with adaptive parameters. Computational Optimization and
Applications, Vol 51, Issue 1, pp 259-277. (2012).
https://pdfs.semanticscholar.org/15b4/c4aa7437df4d032c6ee6ce98d6030dd627be.pdf
Args:
objective_function: A Python callable that accepts a point as a
real `Tensor` and returns a `Tensor` of real dtype containing
the value of the function at that point. The function
to be minimized. If `batch_evaluate_objective` is `True`, the callable
may be evaluated on a `Tensor` of shape `[n+1] + s ` where `n` is
the dimension of the problem and `s` is the shape of a single point
in the domain (so `n` is the size of a `Tensor` representing a
single point).
In this case, the expected return value is a `Tensor` of shape `[n+1]`.
Note that this method does not support univariate functions so the problem
dimension `n` must be strictly greater than 1.
initial_simplex: (Optional) `Tensor` of real dtype. The initial simplex to
start the search. If supplied, should be a `Tensor` of shape `[n+1] + s`
where `n` is the dimension of the problem and `s` is the shape of a
single point in the domain. Each row (i.e. the `Tensor` with a given
value of the first index) is interpreted as a vertex of a simplex and
hence the rows must be affinely independent. If not supplied, an axes
aligned simplex is constructed using the `initial_vertex` and
`step_sizes`. Only one and at least one of `initial_simplex` and
`initial_vertex` must be supplied.
initial_vertex: (Optional) `Tensor` of real dtype and any shape that can
be consumed by the `objective_function`. A single point in the domain that
will be used to construct an axes aligned initial simplex.
step_sizes: (Optional) `Tensor` of real dtype and shape broadcasting
compatible with `initial_vertex`. Supplies the simplex scale along each
axes. Only used if `initial_simplex` is not supplied. See description
above for details on how step sizes and initial vertex are used to
construct the initial simplex.
objective_at_initial_simplex: (Optional) Rank `1` `Tensor` of real dtype
of a rank `1` `Tensor`. The value of the objective function at the
initial simplex. May be supplied only if `initial_simplex` is
supplied. If not supplied, it will be computed.
objective_at_initial_vertex: (Optional) Scalar `Tensor` of real dtype. The
value of the objective function at the initial vertex. May be supplied
only if the `initial_vertex` is also supplied.
batch_evaluate_objective: (Optional) Python `bool`. If True, the objective
function will be evaluated on all the vertices of the simplex packed
into a single tensor. If False, the objective will be mapped across each
vertex separately. Evaluating the objective function in a batch allows
use of vectorization and should be preferred if the objective function
allows it.
func_tolerance: (Optional) Scalar `Tensor` of real dtype. The algorithm
stops if the absolute difference between the largest and the smallest
function value on the vertices of the simplex is below this number.
position_tolerance: (Optional) Scalar `Tensor` of real dtype. The
algorithm stops if the largest absolute difference between the
coordinates of the vertices is below this threshold.
parallel_iterations: (Optional) Positive integer. The number of iterations
allowed to run in parallel.
max_iterations: (Optional) Scalar positive `Tensor` of dtype `int32`.
The maximum number of iterations allowed. If `None` then no limit is
applied.
reflection: (Optional) Positive Scalar `Tensor` of same dtype as
`initial_vertex`. This parameter controls the scaling of the reflected
vertex. See, [Press et al(2007)][1] for details. If not specified,
uses the dimension dependent prescription of [Gao and Han(2012)][3].
expansion: (Optional) Positive Scalar `Tensor` of same dtype as
`initial_vertex`. Should be greater than `1` and `reflection`. This
parameter controls the expanded scaling of a reflected vertex.
See, [Press et al(2007)][1] for details. If not specified, uses the
dimension dependent prescription of [Gao and Han(2012)][3].
contraction: (Optional) Positive scalar `Tensor` of same dtype as
`initial_vertex`. Must be between `0` and `1`. This parameter controls
the contraction of the reflected vertex when the objective function at
the reflected point fails to show sufficient decrease.
See, [Press et al(2007)][1] for more details. If not specified, uses
the dimension dependent prescription of [Gao and Han(2012][3].
shrinkage: (Optional) Positive scalar `Tensor` of same dtype as
`initial_vertex`. Must be between `0` and `1`. This parameter is the scale
by which the simplex is shrunk around the best point when the other
steps fail to produce improvements.
See, [Press et al(2007)][1] for more details. If not specified, uses
the dimension dependent prescription of [Gao and Han(2012][3].
name: (Optional) Python str. The name prefixed to the ops created by this
function. If not supplied, the default name 'minimize' is used.
Returns:
optimizer_results: A namedtuple containing the following items:
converged: Scalar boolean tensor indicating whether the minimum was
found within tolerance.
num_objective_evaluations: The total number of objective
evaluations performed.
position: A `Tensor` containing the last argument value found
during the search. If the search converged, then
this value is the argmin of the objective function.
objective_value: A tensor containing the value of the objective
function at the `position`. If the search
converged, then this is the (local) minimum of
the objective function.
final_simplex: The last simplex constructed before stopping.
final_objective_values: The objective function evaluated at the
vertices of the final simplex.
initial_simplex: The starting simplex.
initial_objective_values: The objective function evaluated at the
vertices of the initial simplex.
num_iterations: The number of iterations of the main algorithm body.
Raises:
ValueError: If any of the following conditions hold
1. If none or more than one of `initial_simplex` and `initial_vertex` are
supplied.
2. If `initial_simplex` and `step_sizes` are both specified.
"""
with tf.name_scope(name or 'minimize'):
(
dim,
_,
simplex,
objective_at_simplex,
num_evaluations
) = _prepare_args(objective_function,
initial_simplex,
initial_vertex,
step_sizes,
objective_at_initial_simplex,
objective_at_initial_vertex,
batch_evaluate_objective)
domain_dtype = simplex.dtype
(
reflection,
expansion,
contraction,
shrinkage
) = _resolve_parameters(dim,
reflection,
expansion,
contraction,
shrinkage,
domain_dtype)
closure_kwargs = dict(
objective_function=objective_function,
dim=dim,
func_tolerance=func_tolerance,
position_tolerance=position_tolerance,
batch_evaluate_objective=batch_evaluate_objective,
reflection=reflection,
expansion=expansion,
contraction=contraction,
shrinkage=shrinkage)
def _loop_body(_, iterations, simplex, objective_at_simplex,
num_evaluations):
(
converged,
next_simplex,
next_objective,
evaluations
) = nelder_mead_one_step(simplex, objective_at_simplex, **closure_kwargs)
return (converged, iterations + 1, next_simplex, next_objective,
num_evaluations + evaluations)
initial_args = (False, 0, simplex, objective_at_simplex,
num_evaluations)
# Loop until either we have converged or if the max iterations are supplied
# then until we have converged or exhausted the available iteration budget.
def _is_converged(converged, num_iterations, *ignored_args): # pylint:disable=unused-argument
# It is important to ensure that not_converged is a tensor. If
# converged is not a tensor but a Python bool, then the overloaded
# op '~' acts as bitwise complement so ~True = -2 and ~False = -1.
# In that case, the loop will never terminate.
not_converged = tf.logical_not(converged)
return (not_converged if max_iterations is None
else (not_converged & (num_iterations < max_iterations)))
(converged, num_iterations, final_simplex, final_objective_values,
final_evaluations) = tf.while_loop(
cond=_is_converged,
body=_loop_body,
loop_vars=initial_args,
parallel_iterations=parallel_iterations)
order = tf.argsort(
final_objective_values, direction='ASCENDING', stable=True)
best_index = order[0]
# The explicit cast to Tensor below is done to avoid returning a mixture
# of Python types and Tensors which cause problems with session.run.
# In the eager mode, converged may remain a Python bool. Trying to evaluate
# the whole tuple in one evaluate call will raise an exception because
# of the presence of non-tensors. This is very annoying so we explicitly
# cast those arguments to Tensors.
return NelderMeadOptimizerResults(
converged=tf.convert_to_tensor(converged),
num_objective_evaluations=final_evaluations,
position=final_simplex[best_index],
objective_value=final_objective_values[best_index],
final_simplex=final_simplex,
final_objective_values=final_objective_values,
num_iterations=tf.convert_to_tensor(num_iterations),
initial_simplex=simplex,
initial_objective_values=objective_at_simplex)
def generate_one(dim, seed):
return tf.argsort(samplers.uniform([num_results, dim], seed=seed),
axis=-1)
def ranks(inputs, axis=-1):
"""Returns the ranks of the input values among the given axis."""
return 1 + tf.cast(tf.argsort(tf.argsort(inputs, axis=axis), axis=axis),
dtype=inputs.dtype)
def nelder_mead_one_step(current_simplex,
current_objective_values,
objective_function=None,
dim=None,
func_tolerance=None,
position_tolerance=None,
batch_evaluate_objective=False,
reflection=None,
expansion=None,
contraction=None,
shrinkage=None,
name=None):
"""A single iteration of the Nelder Mead algorithm."""
with tf.name_scope(name or 'nelder_mead_one_step'):
domain_dtype = dtype_util.base_dtype(current_simplex.dtype)
order = tf.argsort(
current_objective_values, direction='ASCENDING', stable=True)
(
best_index,
worst_index,
second_worst_index
) = order[0], order[-1], order[-2]
worst_vertex = current_simplex[worst_index]
(
best_objective_value,
worst_objective_value,
second_worst_objective_value
) = (
current_objective_values[best_index],
current_objective_values[worst_index],
current_objective_values[second_worst_index]
)
# Compute the centroid of the face opposite the worst vertex.
face_centroid = tf.reduce_sum(
current_simplex, axis=0) - worst_vertex
face_centroid /= tf.cast(dim, domain_dtype)
# Reflect the worst vertex through the opposite face.
reflected = face_centroid + reflection * (face_centroid - worst_vertex)
objective_at_reflected = objective_function(reflected)
num_evaluations = 1
has_converged = _check_convergence(current_simplex,
current_simplex[best_index],
best_objective_value,
worst_objective_value,
func_tolerance,
position_tolerance)
def _converged_fn():
return (True, current_simplex, current_objective_values, np.int32(0))
case0 = has_converged, _converged_fn
accept_reflected = (
(objective_at_reflected < second_worst_objective_value) &
(objective_at_reflected >= best_objective_value))
accept_reflected_fn = _accept_reflected_fn(current_simplex,
current_objective_values,
worst_index,
reflected,
objective_at_reflected)
case1 = accept_reflected, accept_reflected_fn
do_expansion = objective_at_reflected < best_objective_value
expansion_fn = _expansion_fn(objective_function,
current_simplex,
current_objective_values,
worst_index,
reflected,
objective_at_reflected,
face_centroid,
expansion)
case2 = do_expansion, expansion_fn
do_outside_contraction = (
(objective_at_reflected < worst_objective_value) &
(objective_at_reflected >= second_worst_objective_value)
)
outside_contraction_fn = _outside_contraction_fn(
objective_function,
current_simplex,
current_objective_values,
face_centroid,
best_index,
worst_index,
reflected,
objective_at_reflected,
contraction,
shrinkage,
batch_evaluate_objective)
case3 = do_outside_contraction, outside_contraction_fn
default_fn = _inside_contraction_fn(objective_function,
current_simplex,
current_objective_values,
face_centroid,
best_index,
worst_index,
worst_objective_value,
contraction,
shrinkage,
batch_evaluate_objective)
(
converged,
next_simplex,
next_objective_at_simplex,
case_evals) = prefer_static.case([case0, case1, case2, case3],
default=default_fn, exclusive=False)
tensorshape_util.set_shape(next_simplex, current_simplex.shape)
tensorshape_util.set_shape(next_objective_at_simplex,
current_objective_values.shape)
return (
converged,
next_simplex,
next_objective_at_simplex,
num_evaluations + case_evals
)
def _sample_from_edpp(eigenvectors, vector_onehot, seed):
"""Samples a batch of subsets from a DPP given pre-selected elementary DPPs.
Recall that an elementary DPP is a DPP with eigenvalues all exactly 0 or 1.
This function implements the second step of standard sampling algorithm for
DPPs, by sampling subsets based on the E-DPPs obtained by selecting
`vector_onehot` against the DPP's original eigenvectors.
Args:
eigenvectors: A Tensor of `float32` of shape `[..., num_points, num_vecs]`
representing the eigenvectors of a DPP's L-ensemble matrix, eigenvectors
in columns. Generally, `num_vecs == num_points`; we name separately to
distinguish axes.
vector_onehot: A Tensor of shape `[..., n_vecs]` whose innermost
dimension corresponds to 1-hot subset encodings. The subsets represent the
subset of eigenvectors of the original DPP that define an elementary DPP.
seed: The random seed.
Returns:
samples: A many-hot `bool` Tensor of shape `[..., n_points]`
representing a batch of 1-hot subset encodings.
"""
with tf.name_scope('sample_from_edpp'):
seed = samplers.sanitize_seed(seed)
# Sort the 1's to the front, and sort corresponding eigenvectors, then mask.
vector_onehot = tf.cast(vector_onehot, eigenvectors.dtype)
vector_indices = tf.argsort(vector_onehot,
axis=-1,
direction='DESCENDING')
vector_onehot = tf.gather(vector_onehot,
vector_indices,
axis=-1,
batch_dims=len(vector_indices.shape) - 1)
eigenvectors = tf.gather(eigenvectors,
vector_indices,
axis=-1,
batch_dims=len(vector_indices.shape) - 1)
eigenvectors = eigenvectors * vector_onehot[..., tf.newaxis, :]
sample_size = tf.reduce_sum(tf.cast(vector_onehot, tf.int32), axis=-1)
max_sample_size = tf.reduce_max(sample_size)
d = ps.shape(eigenvectors)[-2]
n = ps.shape(eigenvectors)[-1]
# Slice eigvecs to do less work in eager/non-XLA modes.
if FAST_PATH_ENABLED and not JAX_MODE and (
tf.executing_eagerly()
or not control_flow_util.GraphOrParentsInXlaContext(
tf1.get_default_graph())):
# We can save some work in non-XLA contexts by reducing the size of the
# eigenvectors.
eigenvectors = eigenvectors[..., :max_sample_size]
n = max_sample_size
def cond(i, *_):
return i < max_sample_size
def body(i, vecs, cur_sample, seed):
sample_seed, next_seed = samplers.split_seed(seed)
# squared norm at each coord across active subspace
is_active = (i < sample_size)
coord_prob = tf.reduce_sum(tf.square(vecs), axis=-1)
coord_logits = tf.where(is_active[..., tf.newaxis],
tf.math.log(coord_prob), 0.)
idx = categorical.Categorical(logits=coord_logits).sample(
seed=sample_seed)
new_vecs = tf.where(
(tf.range(n) < sample_size[..., tf.newaxis, tf.newaxis] - i -
1) & ~cur_sample[..., tf.newaxis],
_orthogonal_complement_e_i(vecs,
i=tf.where(is_active, idx, 0),
gram_schmidt_iters=max_sample_size -
i), 0.)
# Since range(n) may have unknown shape in the stmt above, we clarify.
tensorshape_util.set_shape(new_vecs, vecs.shape)
vecs = tf.where(is_active[..., tf.newaxis, tf.newaxis], new_vecs,
vecs)
cur_sample = (cur_sample |
(tf.equal(tf.range(d), idx[..., tf.newaxis])
& is_active[..., tf.newaxis]))
return i + 1, vecs, cur_sample, next_seed
_, _, sample, _ = tf.while_loop(
cond, body,
(tf.zeros([], tf.int32, name='i'), eigenvectors,
tf.zeros(ps.shape(eigenvectors)[:-1], dtype=tf.bool), seed))
return tf.cast(sample, tf.int32)
def random_choice_noreplace(m, n, axis=-1):
# Generate m random permuations of range (0, n)
# NumPy version: np.random.rand(m,n).argsort(axis=axis)
return tf.cast(tf.argsort(tf.random.uniform((m, n)), axis=axis), tf.int64)
def _prepare_grid(times,
time_step,
dtype,
*params,
num_time_steps=None,
times_grid=None):
"""Prepares grid of times for path generation.
Args:
times: Rank 1 `Tensor` of increasing positive real values. The times at
which the path points are to be evaluated.
time_step: Rank 0 real `Tensor`. Maximal distance between points in
resulting grid.
dtype: `tf.Dtype` of the input and output `Tensor`s.
*params: Parameters of the Heston model. Either scalar `Tensor`s of the
same `dtype` or instances of `PiecewiseConstantFunc`.
num_time_steps: Number of points on the grid. If suppied, a uniform grid
is constructed for `[time_step, times[-1] - time_step]` consisting of
max(0, num_time_steps - len(times)) points that is then concatenated with
times. This parameter guarantees the number of points on the time grid
is `max(len(times), num_time_steps)` and that `times` are included to the
grid.
Default value: `None`, which means that a uniform grid is created.
containing all points from 'times` and the uniform grid of points between
`[0, times[-1]]` with grid size equal to `time_step`.
times_grid: An optional rank 1 `Tensor` representing time discretization
grid. If `times` are not on the grid, then the nearest points from the
grid are used.
Default value: `None`, which means that times grid is computed using
`time_step` and `num_time_steps`.
Returns:
Tuple `(all_times, mask)`.
`all_times` is a 1-D real `Tensor` containing all points from 'times`, the
uniform grid of points between `[0, times[-1]]` with grid size equal to
`time_step`, and jump locations of piecewise constant parameters The
`Tensor` is sorted in ascending order and may contain duplicates.
`mask` is a boolean 1-D `Tensor` of the same shape as 'all_times', showing
which elements of 'all_times' correspond to THE values from `times`.
Guarantees that times[0]=0 and mask[0]=False.
"""
additional_times = []
for param in params:
if isinstance(param, piecewise.PiecewiseConstantFunc):
additional_times.append(param.jump_locations())
if times_grid is None:
if time_step is not None:
grid = tf.range(0.0, times[-1], time_step, dtype=dtype)
all_times = tf.concat([grid, times] + additional_times, axis=0)
elif num_time_steps is not None:
grid = tf.linspace(tf.convert_to_tensor(0.0, dtype=dtype),
times[-1] - time_step, num_time_steps)
all_times = tf.concat([grid, times] + additional_times, axis=0)
additional_times_mask = [
tf.zeros_like(times, dtype=tf.bool) for times in additional_times
]
mask = tf.concat([
tf.zeros_like(grid, dtype=tf.bool),
tf.ones_like(times, dtype=tf.bool)
] + additional_times_mask,
axis=0)
perm = tf.argsort(all_times, stable=True)
all_times = tf.gather(all_times, perm)
mask = tf.gather(mask, perm)
else:
all_times, mask, _ = utils.prepare_grid(times=times,
time_step=time_step,
times_grid=times_grid,
dtype=dtype)
return all_times, mask
def resample_deterministic_minimum_error(
log_probs,
event_size,
sample_shape,
seed=None,
name='resample_deterministic_minimum_error'):
"""Deterministic minimum error resampler for sequential Monte Carlo.
The return value of this function is similar to sampling with
```python
expanded_sample_shape = tf.concat([sample_shape, [event_size]]), axis=-1)
tfd.Categorical(logits=log_probs).sample(expanded_sample_shape)`
```
but with values chosen deterministically so that the empirical distribution
is as close as possible to the specified distribution.
(Note that the empirical distribution can only exactly equal the requested
distribution if multiplying every probability by `event_size` gives
an integer. So in general this is a biased "sampler".)
It is intended to provide a good representative sample, suitable for use
with some Sequential Monte Carlo algorithms.
This function is based on Algorithm #3 in [Maskell et al. (2006)][1].
Args:
log_probs: a tensor-valued batch of discrete log probability distributions.
event_size: the dimension of the vector considered a single draw.
sample_shape: the `sample_shape` determining the number of draws. Because
this resampler is deterministic it simply replicates the draw you
would get for `sample_shape=[1]`.
seed: This argument is unused but is present so that this function shares
its interface with the other resampling functions.
Default value: None
name: Python `str` name for ops created by this method.
Default value: `None` (i.e., `'resample_deterministic_minimum_error'`).
Returns:
resampled_indices: a tensor of samples.
#### References
[1]: S. Maskell, B. Alun-Jones and M. Macleod. A Single Instruction Multiple
Data Particle Filter.
In 2006 IEEE Nonlinear Statistical Signal Processing Workshop.
http://people.ds.cam.ac.uk/fanf2/hermes/doc/antiforgery/stats.pdf
"""
del seed
with tf.name_scope(name or 'resample_deterministic_minimum_error'):
sample_shape = tf.convert_to_tensor(sample_shape, dtype_hint=tf.int32)
log_probs = dist_util.move_dimension(log_probs,
source_idx=0,
dest_idx=-1)
probs = tf.math.exp(log_probs)
prob_shape = ps.shape(probs)
pdf_size = prob_shape[-1]
# If we could draw fractional numbers of samples we would
# choose `ideal_numbers` for the number of each element.
ideal_numbers = event_size * probs
# We approximate the ideal numbers by truncating to integers
# and then repair the counts starting with the one with the
# largest fractional error and working our way down.
first_approximation = tf.floor(ideal_numbers)
missing_fractions = ideal_numbers - first_approximation
first_approximation = ps.cast(first_approximation, dtype=tf.int32)
fraction_order = tf.argsort(missing_fractions, axis=-1)
# We sort the integer parts and fractional parts together.
batch_dims = ps.rank_from_shape(prob_shape) - 1
first_approximation = tf.gather_nd(first_approximation,
fraction_order[..., tf.newaxis],
batch_dims=batch_dims)
missing_fractions = tf.gather_nd(missing_fractions,
fraction_order[..., tf.newaxis],
batch_dims=batch_dims)
sample_defect = event_size - tf.reduce_sum(
first_approximation, axis=-1, keepdims=True)
unpermuted = tf.broadcast_to(tf.range(pdf_size), prob_shape)
increments = tf.cast(unpermuted >= pdf_size - sample_defect,
dtype=first_approximation.dtype)
counts = first_approximation + increments
samples = _samples_from_counts(fraction_order, counts, event_size)
result_shape = tf.concat([sample_shape, prob_shape[:-1], [event_size]],
axis=0)
# Replicate sample up to batch size.
# TODO(dpiponi): rather than replicating, spread the "error" over
# multiple samples with a minimum-discrepancy sequence.
resampled = tf.broadcast_to(samples, result_shape)
return dist_util.move_dimension(resampled, source_idx=-1, dest_idx=0)
def _argsort(a, axis, stable):
if axis is None:
a = tf.reshape(a, [-1])
axis = 0
return tf.argsort(a, axis, stable=stable)
def _mode(self):
return tf.cast(
tf.argsort(self.scores, axis=-1, direction='DESCENDING'),
self.dtype)
def resample_independent(log_probs, event_size, sample_shape,
seed=None, name=None):
"""Categorical resampler for sequential Monte Carlo.
This function is based on Algorithm #1 in the paper
[Maskell et al. (2006)][1].
Args:
log_probs: A tensor-valued batch of discrete log probability distributions.
event_size: the dimension of the vector considered a single draw.
sample_shape: the `sample_shape` determining the number of draws.
seed: Python '`int` used to seed calls to `tf.random.*`.
Default value: None (i.e. no seed).
name: Python `str` name for ops created by this method.
Default value: `None` (i.e., `'resample_independent'`).
Returns:
resampled_indices: The result is similar to sampling with
```python
expanded_sample_shape = tf.concat([[event_size], sample_shape]), axis=-1)
tfd.Categorical(logits=log_probs).sample(expanded_sample_shape)`
```
but with values sorted along the first axis. It can be considered to be
sampling events made up of a length-`event_size` vector of draws from
the `Categorical` distribution. For large input values this function should
give better performance than using `Categorical`.
The sortedness is an unintended side effect of the algorithm that is
harmless in the context of simple SMC algorithms.
#### References
[1]: S. Maskell, B. Alun-Jones and M. Macleod. A Single Instruction Multiple
Data Particle Filter.
In 2006 IEEE Nonlinear Statistical Signal Processing Workshop.
http://people.ds.cam.ac.uk/fanf2/hermes/doc/antiforgery/stats.pdf
"""
with tf.name_scope(name or 'resample_independent') as name:
log_probs = tf.convert_to_tensor(log_probs, dtype_hint=tf.float32)
log_probs = dist_util.move_dimension(log_probs, source_idx=0, dest_idx=-1)
batch_shape = ps.shape(log_probs)[:-1]
num_markers = ps.shape(log_probs)[-1]
# `working_shape` specifies the total number of events
# we will be generating.
working_shape = ps.concat([sample_shape, batch_shape], axis=0)
# `points_shape` is the shape of the final result.
points_shape = ps.concat([working_shape, [event_size]], axis=0)
# `markers_shape` is the shape of the markers we temporarily insert.
markers_shape = ps.concat([working_shape, [num_markers]], axis=0)
# Generate one real point for each particle.
log_points = -exponential.Exponential(
rate=tf.constant(1.0, dtype=log_probs.dtype)).sample(
points_shape, seed=seed)
# We divide up the unit interval [0, 1] according to the provided
# probability distributions using `cumsum`.
# At the end of each division we place a 'marker'.
# We generate random points on the unit interval.
# We sort the combination of points and markers. The number
# of points between the markers defining a division gives the number
# of samples we require in that division.
# For example, suppose `probs` is `[0.2, 0.3, 0.5]`.
# We divide up `[0, 1]` using 3 markers:
#
# | | |
# 0. 0.2 0.5 1.0 <- markers
#
# Suppose we generate four points: [0.1, 0.25, 0.9, 0.75]
# After sorting the combination we get:
#
# 0.1 0.25 0.75 0.9 <- points
# * | * | * *|
# 0. 0.2 0.5 1.0 <- markers
#
# We have one sample in the first category, one in the second and
# two in the last.
#
# All of these computations are carried out in batched form.
markers = ps.concat(
[tf.zeros(points_shape, dtype=tf.int32),
tf.ones(markers_shape, dtype=tf.int32)],
axis=-1)
log_marker_positions = tf.broadcast_to(
tf.math.cumulative_logsumexp(log_probs, axis=-1),
markers_shape)
log_points_and_markers = ps.concat(
[log_points, log_marker_positions], axis=-1)
indices = tf.argsort(log_points_and_markers, axis=-1, stable=False)
sorted_markers = tf.gather_nd(
markers,
indices[..., tf.newaxis],
batch_dims=(
ps.rank_from_shape(sample_shape) +
ps.rank_from_shape(batch_shape)))
markers_and_samples = ps.cast(
tf.cumsum(sorted_markers, axis=-1), dtype=tf.int32)
markers_and_samples = tf.minimum(markers_and_samples, num_markers - 1)
# Collect up samples, omitting markers.
resampled = tf.reshape(markers_and_samples[tf.equal(sorted_markers, 0)],
points_shape)
resampled = dist_util.move_dimension(resampled, source_idx=-1, dest_idx=0)
return resampled
def _points_sort(points, key, stable=False):
indices = tf.argsort(key(points), stable=stable)
return tf.gather(points, indices)
def get_dice_pose_results(bounding_boxes, classes, scores, y_rotation_angles, camera_matrix : np.ndarray, distortion_coefficients : np.ndarray, score_threshold : float = 0.5):
"""Estimates pose results for all die, given estimates for bounding box, die (top face) classes, scores and threshold, rotation angles around vertical axes, and camera information."""
scores_in_threshold = tf.math.greater(scores, score_threshold)
classes_in_score = tf.boolean_mask(classes, scores_in_threshold)
boxes_in_scores = tf.boolean_mask(bounding_boxes, scores_in_threshold)
y_angles_in_scores = tf.boolean_mask(y_rotation_angles, scores_in_threshold)
classes_are_dots = tf.equal(classes_in_score, 0)
classes_are_dice = tf.logical_not(classes_are_dots)
dice_bounding_boxes = tf.boolean_mask(boxes_in_scores, classes_are_dice)
dice_y_angles = tf.boolean_mask(y_angles_in_scores, classes_are_dice)
dice_classes = tf.boolean_mask(classes_in_score, classes_are_dice)
dot_bounding_boxes = tf.boolean_mask(boxes_in_scores, classes_are_dots)
dot_centers = _get_dot_centers(dot_bounding_boxes)
dot_sizes = _get_dot_sizes(dot_bounding_boxes)
#NB Largest box[2] is the box lower bound
dice_bb_lower_y = dice_bounding_boxes[:,2]
dice_indices = tf.argsort(dice_bb_lower_y, axis = -1, direction='DESCENDING')
def get_area(bb):
return tf.math.maximum(bb[:, 3] - bb[:, 1], 0) * tf.math.maximum(bb[:, 2] - bb[:, 0], 0)
dice_indices_np = dice_indices.numpy()
bounding_box_pose_results = [_get_die_image_bounding_box_pose(dice_bounding_boxes[index, :], camera_matrix, distortion_coefficients) for index in dice_indices_np]
approximate_dice_up_vector_pyrender = _get_approximate_dice_up_vector(bounding_box_pose_results, in_pyrender_coords=True)
pose_results = []
for index, bounding_box_pose_result in zip(dice_indices_np, bounding_box_pose_results):
die_box = dice_bounding_boxes[index, :]
die_y_angle = dice_y_angles[index]
die_class = dice_classes[index]
die_box_size = (-die_box[0:2] + die_box[2:4])
dot_centers_fraction_of_die_box = (dot_centers - die_box[0:2]) / die_box_size
dot_centers_rounded_rectangle_distance = tf.norm(tf.math.maximum(tf.math.abs(dot_centers_fraction_of_die_box - 0.5) - 0.5 + rounded_rectangle_radius,0.0), axis = -1) - rounded_rectangle_radius
dots_are_in_rounded_rectangle = dot_centers_rounded_rectangle_distance < 0
dot_bb_intersection_left = tf.math.maximum(dot_bounding_boxes[:, 1], die_box[1])
dot_bb_intersection_right = tf.math.minimum(dot_bounding_boxes[:, 3], die_box[3])
dot_bb_intersection_top = tf.math.maximum(dot_bounding_boxes[:, 0], die_box[0])
dot_bb_intersection_bottom = tf.math.minimum(dot_bounding_boxes[:, 2], die_box[2])
dot_bb_intersection = tf.stack([dot_bb_intersection_top, dot_bb_intersection_left, dot_bb_intersection_bottom, dot_bb_intersection_right], axis = 1)
dot_bb_intersection_area = get_area(dot_bb_intersection)
dot_bb_area = get_area(dot_bounding_boxes)
dot_bb_intersection_over_area = dot_bb_intersection_area / dot_bb_area
dots_have_sufficient_bb_intersection_over_area = tf.greater(dot_bb_intersection_over_area, 0.9)
dots_are_in_box = tf.logical_and(dots_have_sufficient_bb_intersection_over_area, dots_are_in_rounded_rectangle)
dot_centers_in_box = tf.boolean_mask(dot_centers, dots_are_in_box)
dot_centers_cv = _convert_tensorflow_points_to_opencv(dot_centers_in_box)
die_pose_result = _get_die_pose(die_box, die_class, die_y_angle, dot_centers_cv, bounding_box_pose_result, approximate_dice_up_vector_pyrender, camera_matrix, distortion_coefficients)
die_pose_result.calculate_comparison(dot_centers_cv, camera_matrix, distortion_coefficients)
die_pose_result.calculate_inliers(_convert_tensorflow_points_to_opencv(dot_sizes))
pose_results.append(die_pose_result)
indices_in_box = tf.where(dots_are_in_box)
inlier_indices_in_box = tf.gather(indices_in_box, die_pose_result.comparison_inlier_indices)
dot_centers = _delete_tf(dot_centers, inlier_indices_in_box)
dot_sizes = _delete_tf(dot_sizes, inlier_indices_in_box)
dot_bounding_boxes = _delete_tf(dot_bounding_boxes, inlier_indices_in_box)
return pose_results
