From 9266eb54c6bd17408a0c1eb57454a1c3624190e4 Mon Sep 17 00:00:00 2001 From: Aaron Date: Wed, 15 Mar 2023 00:43:56 +0000 Subject: [PATCH] [SymForce] Add full stats to Python result object This adds all the fields from `cc_sym.OptimizationStats` to `Optimizer.Result` in Python. Adding these as cached properties, since 1) we don't want to expose things from the stats directly, since different types might make more sense 2) conversions of some of these things from C++ can be expensive, e.g. best_linearization, and we'd rather only convert if actually requested To allow conversions to happen lazily we keep the original stats as a private member. Would be good to have a test at some point that the Result has all the stuff the OptimizationStats does so they don't get out of sync again, but not doing that for now. Also passed through other arguments to `Optimizer.optimize`, i.e. `num_iterations` and `populate_best_linearization`. Topic: sf-stats Reviewers: bradley,nathan,ryan-b,chao,hayk GitOrigin-RevId: 751e3aa4095d8d2daed2d0c777902b84818e83b8 --- .../robot_2d_localization/plotting.py | 6 +- .../robot_2d_localization.py | 2 +- .../robot_3d_localization/plotting.py | 6 +- .../robot_3d_localization.py | 2 +- symforce/opt/optimization_stats.h | 2 + symforce/opt/optimizer.py | 91 +++++++++++++++---- symforce/pybind/cc_optimization_stats.cc | 4 +- symforce/pybind/cc_sym.pyi | 6 +- ...rce_examples_robot_2d_localization_test.py | 2 +- ...rce_examples_robot_3d_localization_test.py | 2 +- test/symforce_py_optimizer_test.py | 4 +- 11 files changed, 93 insertions(+), 34 deletions(-) diff --git a/symforce/examples/robot_2d_localization/plotting.py b/symforce/examples/robot_2d_localization/plotting.py index 8a6a0e5cd..c1a06d542 100644 --- a/symforce/examples/robot_2d_localization/plotting.py +++ b/symforce/examples/robot_2d_localization/plotting.py @@ -20,9 +20,7 @@ def plot_solution(optimizer: Optimizer, result: Optimizer.Result, animated: bool matplotlib animation instead of providing an interactive slider. """ # Pull out values from the result - values_per_iter = [ - optimizer.load_iteration_values(stats.values) for stats in result.iteration_stats - ] + values_per_iter = [optimizer.load_iteration_values(stats.values) for stats in result.iterations] # Create the layout fig = plt.figure() @@ -101,7 +99,7 @@ def update_plot(slider_value: np.float64) -> None: num = int(slider_value) # Set iteration text and abort if we rejected this iteration - stats = result.iteration_stats[num] + stats = result.iterations[num] if num > 0 and not stats.update_accepted: text.set_text(f"Iteration: {num} (rejected)\nError: {stats.new_error:.6f}") return diff --git a/symforce/examples/robot_2d_localization/robot_2d_localization.py b/symforce/examples/robot_2d_localization/robot_2d_localization.py index 9937351b0..186c46d65 100644 --- a/symforce/examples/robot_2d_localization/robot_2d_localization.py +++ b/symforce/examples/robot_2d_localization/robot_2d_localization.py @@ -122,7 +122,7 @@ def main() -> None: result = optimizer.optimize(initial_values) # Print some values - print(f"Num iterations: {len(result.iteration_stats) - 1}") + print(f"Num iterations: {len(result.iterations) - 1}") print(f"Final error: {result.error():.6f}") for i, pose in enumerate(result.optimized_values["poses"]): diff --git a/symforce/examples/robot_3d_localization/plotting.py b/symforce/examples/robot_3d_localization/plotting.py index 2d8a05a4d..6bfc67192 100644 --- a/symforce/examples/robot_3d_localization/plotting.py +++ b/symforce/examples/robot_3d_localization/plotting.py @@ -24,9 +24,7 @@ def plot_solution( matplotlib animation instead of providing an interactive slider. """ # Pull out values from the result - values_per_iter = [ - optimizer.load_iteration_values(stats.values) for stats in result.iteration_stats - ] + values_per_iter = [optimizer.load_iteration_values(stats.values) for stats in result.iterations] # Create the layout fig = plt.figure() @@ -109,7 +107,7 @@ def update_plot(slider_value: np.float64) -> None: # Set iteration text and abort if we rejected this iteration if show_iteration_text: - stats = result.iteration_stats[num] + stats = result.iterations[num] if num > 0 and not stats.update_accepted: text.set_text(f"Iteration: {num} (rejected)\nError: {stats.new_error:.1f}") return diff --git a/symforce/examples/robot_3d_localization/robot_3d_localization.py b/symforce/examples/robot_3d_localization/robot_3d_localization.py index bdf58ce12..351c30c91 100644 --- a/symforce/examples/robot_3d_localization/robot_3d_localization.py +++ b/symforce/examples/robot_3d_localization/robot_3d_localization.py @@ -249,7 +249,7 @@ def main() -> None: result = optimizer.optimize(values) # Print some values - print(f"Num iterations: {len(result.iteration_stats) - 1}") + print(f"Num iterations: {len(result.iterations) - 1}") print(f"Final error: {result.error():.6f}") for i, pose in enumerate(result.optimized_values["world_T_body"]): diff --git a/symforce/opt/optimization_stats.h b/symforce/opt/optimization_stats.h index d105d164f..aba27f4a6 100644 --- a/symforce/opt/optimization_stats.h +++ b/symforce/opt/optimization_stats.h @@ -25,6 +25,8 @@ struct OptimizationStats { // good step) bool early_exited{false}; + // The linearization at best_index (at optimized_values), filled out if + // populate_best_linearization=true optional> best_linearization{}; sparse_matrix_structure_t jacobian_sparsity; diff --git a/symforce/opt/optimizer.py b/symforce/opt/optimizer.py index 66cd20dee..05603ce60 100644 --- a/symforce/opt/optimizer.py +++ b/symforce/opt/optimizer.py @@ -6,11 +6,16 @@ from __future__ import annotations import dataclasses +import warnings from dataclasses import dataclass +from functools import cached_property + +import numpy as np from lcmtypes.sym._index_entry_t import index_entry_t from lcmtypes.sym._optimization_iteration_t import optimization_iteration_t from lcmtypes.sym._optimizer_params_t import optimizer_params_t +from lcmtypes.sym._sparse_matrix_structure_t import sparse_matrix_structure_t from lcmtypes.sym._values_t import values_t from symforce import cc_sym @@ -112,29 +117,79 @@ class Result: optimized_values: The best Values achieved during the optimization (Values with the smallest error) - iteration_stats: + iterations: Per-iteration stats, if requested, like the error per iteration. If debug stats are turned on, also the Values and linearization per iteration. + best_index: + The index into iterations for the iteration that produced the smallest error. I.e. + `result.iterations[best_index].values == optimized_values`. This is not guaranteed + to be the last iteration, if the optimizer tried additional steps which did not reduce + the error + early_exited: Did the optimization early exit? This can happen because it converged successfully, of because it was unable to make progress - best_index: - The index into iteration_stats for the iteration that produced the smallest error. I.e. - `result.iteration_stats[best_index].values == optimized_values`. This is not guaranteed - to be the last iteration, if the optimizer tried additional steps which did not reduce - the error + best_linearization: + The linearization at best_index (at optimized_values), filled out if + populate_best_linearization=True + + jacobian_sparsity: + The sparsity pattern of the jacobian, filled out if debug_stats=True + + linear_solver_ordering: + The ordering used for the linear solver, filled out if debug_stats=True + + cholesky_factor_sparsity: + The sparsity pattern of the cholesky factor L, filled out if debug_stats=True """ initial_values: Values optimized_values: Values - iteration_stats: T.Sequence[optimization_iteration_t] - early_exited: bool - best_index: int + + # Private field holding the original stats - we expose fields of this through properties, + # since some of the conversions out of this are expensive + _stats: cc_sym.OptimizationStats + + @cached_property + def iterations(self) -> T.List[optimization_iteration_t]: + return self._stats.iterations + + @cached_property + def best_index(self) -> int: + return self._stats.best_index + + @cached_property + def early_exited(self) -> bool: + return self._stats.early_exited + + @cached_property + def best_linearization(self) -> T.Optional[cc_sym.Linearization]: + return self._stats.best_linearization + + @cached_property + def jacobian_sparsity(self) -> sparse_matrix_structure_t: + return self._stats.jacobian_sparsity + + @cached_property + def linear_solver_ordering(self) -> np.ndarray: + return self._stats.linear_solver_ordering + + @cached_property + def cholesky_factor_sparsity(self) -> sparse_matrix_structure_t: + return self._stats.cholesky_factor_sparsity + + @cached_property + def iteration_stats(self) -> T.Sequence[optimization_iteration_t]: + warnings.warn("iteration_stats is deprecated, use iterations", FutureWarning) + return self.iterations def error(self) -> float: - return self.iteration_stats[self.best_index].new_error + """ + The lowest error achieved by the optimization (the error at optimized_values) + """ + return self.iterations[self.best_index].new_error def __init__( self, @@ -232,13 +287,17 @@ def _cc_values(self, values: Values) -> cc_sym.Values: return cc_values - def optimize(self, initial_guess: Values) -> Optimizer.Result: + def optimize(self, initial_guess: Values, **kwargs: T.Any) -> Optimizer.Result: """ Optimize from the given initial guess, and return the optimized Values and stats Args: initial_guess: A Values containing the initial guess, should contain at least all the - keys required by the `factors` passed to the constructor + keys required by the `factors` passed to the constructor + num_iterations: If < 0 (the default), uses the number of iterations specified by the + params at construction + populate_best_linearization: If true, the linearization at the best values will be + filled out in the stats Returns: The optimization results, with additional stats and debug information. See the @@ -247,7 +306,7 @@ def optimize(self, initial_guess: Values) -> Optimizer.Result: cc_values = self._cc_values(initial_guess) try: - stats = self._cc_optimizer.optimize(cc_values) + stats = self._cc_optimizer.optimize(cc_values, **kwargs) except ZeroDivisionError as ex: raise ZeroDivisionError("ERROR: Division by zero - check your use of epsilon!") from ex @@ -259,11 +318,7 @@ def optimize(self, initial_guess: Values) -> Optimizer.Result: ) return Optimizer.Result( - initial_values=initial_guess, - optimized_values=optimized_values, - iteration_stats=stats.iterations, - best_index=stats.best_index, - early_exited=stats.early_exited, + initial_values=initial_guess, optimized_values=optimized_values, _stats=stats ) def linearize(self, values: Values) -> cc_sym.Linearization: diff --git a/symforce/pybind/cc_optimization_stats.cc b/symforce/pybind/cc_optimization_stats.cc index 941608e59..c820a3646 100644 --- a/symforce/pybind/cc_optimization_stats.cc +++ b/symforce/pybind/cc_optimization_stats.cc @@ -49,7 +49,9 @@ void AddOptimizationStatsWrapper(pybind11::module_ module) { } else { stats.best_linearization = *best_linearization; } - }) + }, + "The linearization at best_index (at optimized_values), filled out if " + "populate_best_linearization=True") .def("get_lcm_type", &sym::OptimizationStatsd::GetLcmType) .def(py::pickle( [](const sym::OptimizationStatsd& stats) { // __getstate__ diff --git a/symforce/pybind/cc_sym.pyi b/symforce/pybind/cc_sym.pyi index 3c01a75f8..7a1b27d04 100644 --- a/symforce/pybind/cc_sym.pyi +++ b/symforce/pybind/cc_sym.pyi @@ -291,11 +291,15 @@ class OptimizationStats: @property def best_linearization(self) -> typing.Optional[Linearization]: """ + The linearization at best_index (at optimized_values), filled out if populate_best_linearization=True + :type: object """ @best_linearization.setter def best_linearization(self, arg1: Linearization) -> None: - pass + """ + The linearization at best_index (at optimized_values), filled out if populate_best_linearization=True + """ @property def cholesky_factor_sparsity(self) -> sparse_matrix_structure_t: """ diff --git a/test/symforce_examples_robot_2d_localization_test.py b/test/symforce_examples_robot_2d_localization_test.py index 60bd60f0a..7e2531f24 100644 --- a/test/symforce_examples_robot_2d_localization_test.py +++ b/test/symforce_examples_robot_2d_localization_test.py @@ -41,7 +41,7 @@ def test_optimize(self) -> None: # Solve and return the result result = optimizer.optimize(initial_values) - self.assertAlmostEqual(result.iteration_stats[0].new_error, 6.396357536315918) + self.assertAlmostEqual(result.iterations[0].new_error, 6.396357536315918) self.assertLess(result.error(), 1e-3) diff --git a/test/symforce_examples_robot_3d_localization_test.py b/test/symforce_examples_robot_3d_localization_test.py index 35a8208ab..f714196ad 100644 --- a/test/symforce_examples_robot_3d_localization_test.py +++ b/test/symforce_examples_robot_3d_localization_test.py @@ -42,7 +42,7 @@ def test_optimize(self) -> None: # Solve and return the result result = optimizer.optimize(values) - self.assertAlmostEqual(result.iteration_stats[0].new_error, 463700.5625) + self.assertAlmostEqual(result.iterations[0].new_error, 463700.5625) self.assertLess(result.error(), 140) diff --git a/test/symforce_py_optimizer_test.py b/test/symforce_py_optimizer_test.py index ded00e3b5..596524dcf 100644 --- a/test/symforce_py_optimizer_test.py +++ b/test/symforce_py_optimizer_test.py @@ -66,11 +66,11 @@ def prior_residual(x: sf.Rot3, epsilon: sf.Scalar, x_prior: sf.Rot3) -> sf.V3: logger.debug(f"Initial values: {result.initial_values}") logger.debug(f"Optimized values: {result.optimized_values}") - logger.debug(f"Num iterations: {len(result.iteration_stats)}") + logger.debug(f"Num iterations: {len(result.iterations)}") logger.debug(f"Final error: {result.error()}") # Check values - self.assertEqual(len(result.iteration_stats), 7) + self.assertEqual(len(result.iterations), 7) self.assertAlmostEqual(result.error(), 0.039, places=3) # Check that we can pull out the variable blocks