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