From 2fcd3c9f55460979cbe662ba35a14d57b20aa075 Mon Sep 17 00:00:00 2001 From: Frederike Duembgen Date: Fri, 22 Aug 2025 19:31:45 +0200 Subject: [PATCH 01/10] Change output of get_A0 to be always lists, improve doc --- popcor/base_lifters/_base_class.py | 2 +- popcor/examples/rotation_lifter.py | 85 ++++++++++++++---------------- tests/test_rotation_lifter.py | 7 +-- 3 files changed, 46 insertions(+), 48 deletions(-) diff --git a/popcor/base_lifters/_base_class.py b/popcor/base_lifters/_base_class.py index 6342d0e..d1b4572 100644 --- a/popcor/base_lifters/_base_class.py +++ b/popcor/base_lifters/_base_class.py @@ -682,7 +682,7 @@ def get_A0(self, var_subset=None): var_dict = self.var_dict A0 = PolyMatrix() A0[self.HOM, self.HOM] = 1.0 - return A0.get_matrix(var_dict) + return [A0.get_matrix(var_dict)], [1.0] def get_A_b_list(self, A_list, var_subset=None): """get equality constraint tuples (Ai, bi) s.t. x.t @ Ai @ bi, 0 diff --git a/popcor/examples/rotation_lifter.py b/popcor/examples/rotation_lifter.py index ff7f5d2..711e60b 100644 --- a/popcor/examples/rotation_lifter.py +++ b/popcor/examples/rotation_lifter.py @@ -1,14 +1,5 @@ """ A class for solving rotation averaging problems. - -Some notes and conventions: -- The rotation matrices are called R_i. -- We consider two different forms of "lifting": - - "no" corresponds to the rank-1 version, where we define x = [1, c_1, ..., c_N] with each c_i = R_i.flatten("C"), x in (d**2 x N + 1) - - "bm" corresponds to the rank-d version, where we define X.T = [R_0.T; R_1.T; ...; R_N.T], X in (d x (N + 1)*d), and R_0 is the world frame. -- According to above conventions, HOM is either 1 or R_0, the world frame. -- Relative measurements are defined such that R_j = R_i @ R_ij - """ from typing import List @@ -28,6 +19,11 @@ def get_orthogonal_constraints(key, hom, d, level): + """Return A, b lists enforcing orthogonality constraints for variable `key`. + + Produces constraints encoding R'R = I (diagonal == 1 and off-diagonal == 0) + for either the rank-1 ("no") or Burer-Monteiro ("bm") formulation. + """ A_list = [] b_list = [] for i in range(d): @@ -66,7 +62,7 @@ def get_orthogonal_constraints(key, hom, d, level): class RotationLifter(StateLifter): - """Rotation averaging problem. + """Rotation averaging problem lifter. We solve the following optimization problem: @@ -81,14 +77,13 @@ class RotationLifter(StateLifter): .. math:: \\theta = \\begin{bmatrix} R_1 & R_2 & \\ldots & R_N \\end{bmatrix} - We can alternatively replace the absolute-measurement terms by .. math:: || R_i - R_w \\tilde{R}_{i} ||_F^2 where :math:`R_w` is an arbitrary world frame that we can also optimize over, transforming the solutions - by :math:`R_w^{-1}R_i` after to move the world frame to the origin. Using this formulation, all + by :math:`R_w^{-1}R_i` afterwards to move the world frame to the origin. Using this formulation, all measurements are binary factors, which may simplify implementation. We consider two different formulations of the problem: @@ -102,6 +97,8 @@ class RotationLifter(StateLifter): .. math:: X = \\begin{bmatrix} R_1^\\top \\\\ \\vdots \\\\ R_N^\\top \\end{bmatrix} + + According to the above conventions, HOM is either 1 or R_0, the world frame. """ LEVELS = ["no", "bm"] @@ -110,10 +107,8 @@ class RotationLifter(StateLifter): # whether or not to include the determinant constraints in the known constraints. ADD_DETERMINANT = False - NOISE = 1e-3 - # Add any parameters here that describe the problem (e.g. number of landmarks etc.) def __init__( self, level="no", @@ -124,6 +119,7 @@ def __init__( n_rel=1, sparsity="chain", ): + """Initialize a RotationLifter instance with problem dimensions and options.""" assert n_rel in [ 0, 1, @@ -141,6 +137,7 @@ def __init__( @property def var_dict(self): + """Return dictionary mapping variable names to their dimensions in the lifted representation.""" if self.level == "no": # self.HOM is the scalar homogenization variable var_dict = {self.HOM: 1} @@ -152,7 +149,7 @@ def var_dict(self): return var_dict def sample_theta(self): - """Generate a random new feasible point.""" + """Generate a random feasible set of rotations theta shaped (d, n_rot * d).""" C = np.empty((self.d, self.n_rot * self.d)) for i in range(self.n_rot): if self.d == 2: @@ -165,7 +162,7 @@ def sample_theta(self): return C def get_x(self, theta=None, parameters=None, var_subset=None) -> np.ndarray: - """Get the lifted vector x given theta and parameters.""" + """Return the lifted variable x (vector or stacked matrices) for given theta and var_subset.""" if theta is None: theta = self.theta if parameters is None: @@ -180,6 +177,7 @@ def get_x(self, theta=None, parameters=None, var_subset=None) -> np.ndarray: x_data.append(1.0) elif "c" in key: i = int(key.split("_")[1]) + # column-major flattening to match vec(R) x_data += list(theta[:, i * self.d : (i + 1) * self.d].flatten("F")) else: raise ValueError(f"untreated key {key}") @@ -200,10 +198,11 @@ def get_x(self, theta=None, parameters=None, var_subset=None) -> np.ndarray: raise ValueError(f"Unknown level {self.level} for RotationLifter") def get_theta(self, x: np.ndarray) -> np.ndarray: + """Recover theta (d x n_rot*d) from lifted variable x for current level.""" if self.level == "no": if np.ndim(x) == 2: assert x.shape[1] == 1 - + x = x.flatten() C_flat = x[1 : 1 + self.n_rot * self.d**2] return C_flat.reshape((self.d, self.n_rot * self.d), order="F") elif self.level == "bm": @@ -211,22 +210,21 @@ def get_theta(self, x: np.ndarray) -> np.ndarray: # We want to return theta in the form [R*_1, R*_2, ..., R*_N] # where each R*_i := R_0.T @ R_i - # d x d + # d x d (original world frame) R0 = x[: self.d, : self.d].T # nd x d: [R_1.T; R_2.T; ...; R_N.T] - Ri = np.array(x[self.d : (self.n_rot + 1) * self.d, : self.d]) + Ri = np.array(x[self.d : (self.n_rot + 1) * self.d, : self.d]).T # return d x nd: [R*_1, R*_2, ..., R*_N] - Ri_world = R0.T @ Ri.T + Ri_world = R0.T @ Ri return Ri_world else: raise ValueError(f"Unknown level {self.level} for RotationLifter") def add_relative_measurement(self, i, j, noise) -> np.ndarray: - """ - || R_j - R_i @ R_ij ||_F^2 - + """Create a noisy relative measurement R_ij = R_i.T @ R_j with additive rotation noise. + error terms: || R_j - R_i @ R_ij ||_F^2 measurements: R_ij = R_i.T @ R_j """ R_i = self.theta[:, i * self.d : (i + 1) * self.d] @@ -246,11 +244,9 @@ def add_relative_measurement(self, i, j, noise) -> np.ndarray: return R_gt def add_absolute_measurement(self, i, noise, n_meas=1) -> List[np.ndarray]: - """ - || R_i - R_w @ R_wi || - - measurements: - R_wi = R_w.T @ R_i + """Create one or more noisy absolute measurements of rotation R_i (relative to world). + error terms: || R_i - R_w @ R_wi ||_F^2 + measurements: R_wi = R_w.T @ R_i """ R_gt = self.theta[:, i * self.d : (i + 1) * self.d] y = [] @@ -271,6 +267,7 @@ def add_absolute_measurement(self, i, noise, n_meas=1) -> List[np.ndarray]: return y def simulate_y(self, noise: float | None = None) -> dict: + """Simulate measurement dictionary y given current theta and noise level.""" if noise is None: noise = self.NOISE @@ -279,7 +276,8 @@ def simulate_y(self, noise: float | None = None) -> dict: for i in range(self.n_rot): y[i] = self.add_absolute_measurement(i, noise, self.n_abs) else: - y[0] = self.add_absolute_measurement(0, 0.0, 1) # add prior + # add prior if no absolute measurements exist + y[0] = self.add_absolute_measurement(0, 0.0, 1) if self.n_rel > 0: if self.sparsity == "chain": @@ -291,6 +289,7 @@ def simulate_y(self, noise: float | None = None) -> dict: return y def get_Q(self, noise: float | None = None, output_poly: bool = False): + """Return the cost matrix Q (poly or ndarray) constructed from simulated measurements.""" if noise is None: noise = self.NOISE if self.y_ is None: @@ -299,7 +298,7 @@ def get_Q(self, noise: float | None = None, output_poly: bool = False): return self.get_Q_from_y(self.y_, output_poly=output_poly) def get_Q_from_y(self, y, output_poly=False): - """param y: list of noisy rotation matrices.""" + """Construct the quadratic cost (PolyMatrix or ndarray) from measurement dictionary y.""" Q = PolyMatrix() for key, R in y.items(): @@ -308,10 +307,8 @@ def get_Q_from_y(self, y, output_poly=False): assert isinstance(R, list) for Rk in R: if self.level == "no": - # treat unary factors - # Rk is measured - # f(R) = sum_k || R - Rk ||_F^2 - # = sum_k 2tr(I) - 2tr(R'Rk) + # unary factors: f(R) = sum_k || R - Rk ||_F^2 = sum_k 2tr(I) - 2tr(R'Rk) + # Rk is measured rotation Q_test = PolyMatrix() Q_test[self.HOM, f"c_{key}"] -= Rk.flatten("F")[None, :] # Not adding below to be consistent with "bm" case, where we cannot @@ -326,8 +323,7 @@ def get_Q_from_y(self, y, output_poly=False): Q += Q_test elif self.level == "bm": - # in this case, we use self.HOM as a world frame. - # f(R) = sum_i || R - HOM @ Rk || + # use HOM as a world frame: f(R) = sum_i || R - HOM @ Rk || # compare with below: # R corresponds to Rj, HOM corresponds to Ri Q_test = PolyMatrix() @@ -395,6 +391,7 @@ def get_Q_from_y(self, y, output_poly=False): return Q.get_matrix(self.var_dict) def test_and_add(self, A_list, Ai, output_poly, b_list=[], bi=0.0): + """Test that constraint Ai holds at current theta then append it to A_list and b_list.""" x = self.get_x() Ai_sparse = Ai.get_matrix(self.var_dict) err = np.trace(np.atleast_2d(x.T @ Ai_sparse @ x)) - bi @@ -406,6 +403,7 @@ def test_and_add(self, A_list, Ai, output_poly, b_list=[], bi=0.0): b_list.append(bi) def get_A0(self, var_subset=None): + """Return the homogenization constraint A0 for chosen level (either h^2=1, or H'H=I).""" if var_subset is None: var_subset = self.var_dict if self.level == "no": @@ -418,11 +416,16 @@ def get_A0(self, var_subset=None): return [Ai.get_matrix(var_subset) for Ai in A_orth], b_orth def get_A_known(self, var_dict=None, output_poly=False, add_redundant=False): + """Return known linear constraints (A and b). If level 'no' returns A_list; if 'bm' returns (A_list, b_list).""" A_list = [] b_list = [] if var_dict is None: var_dict = self.var_dict + if self.level == "bm" and add_redundant: + print("No known redundant constraints for level bm") + add_redundant = False + for k in range(self.n_rot): if f"c_{k}" in var_dict: A_orth, b_orth = get_orthogonal_constraints( @@ -467,10 +470,6 @@ def get_A_known(self, var_dict=None, output_poly=False, add_redundant=False): ) if add_redundant and f"c_{k}" in var_dict: - if self.level == "bm": - print("No known redundant constraints for level bm") - continue - # enforce diagonal == 1 for RR' = I for i in range(self.d): Ei = np.zeros((self.d, self.d)) @@ -504,12 +503,10 @@ def get_A_known(self, var_dict=None, output_poly=False, add_redundant=False): Ai = PolyMatrix(symmetric=True) Ai[f"c_{k}", f"c_{k}"] = constraint self.test_and_add(A_list, Ai, output_poly, b_list, 0.0) - if self.level == "bm": - return A_list, b_list - else: - return A_list + return A_list, b_list def plot(self, estimates={}): + """Plot ground-truth frames and optional estimated frames; returns (fig, ax).""" import itertools import matplotlib.pyplot as plt diff --git a/tests/test_rotation_lifter.py b/tests/test_rotation_lifter.py index 06e3039..42668c7 100644 --- a/tests/test_rotation_lifter.py +++ b/tests/test_rotation_lifter.py @@ -151,8 +151,9 @@ def test_solve_sdp(): Q = lifter.get_Q_from_y(y=y, output_poly=False) assert isinstance(Q, sp.csc_matrix) - A_known = lifter.get_A_known() - constraints = lifter.get_A_b_list(A_known) + A_known, b_known = lifter.get_A_known() + A_0, b_0 = lifter.get_A0() + constraints = list(zip(A_0 + A_known, b_0 + b_known)) if noise == 0 and PLOT: plot_matrices(A_known, Q) @@ -163,7 +164,7 @@ def test_solve_sdp(): print(f"EVR at noise {noise:.2f}: {info_rank['EVR']:.2e}") if noise <= 1: assert info_rank["EVR"] > 1e8 - theta_sdp = lifter.get_theta(x.flatten()) + theta_sdp = lifter.get_theta(x) if noise == 0.0: np.testing.assert_allclose(theta_sdp, lifter.theta, atol=1e-5) From 93e076598d5bef1d4eb4ba119532b5cc8950f335 Mon Sep 17 00:00:00 2001 From: Frederike Duembgen Date: Fri, 22 Aug 2025 19:50:22 +0200 Subject: [PATCH 02/10] Apply same constraints convention everywhere --- popcor/auto_template.py | 50 ++++++++------ popcor/auto_tight.py | 40 +++++++++--- popcor/examples/poly4_lifter.py | 4 +- popcor/examples/stereo1d_lifter.py | 4 +- popcor/examples/stereo3d_lifter.py | 101 +++++++++++++++++------------ popcor/utils/constraint.py | 3 + tests/test_autotight.py | 50 +++++++------- tests/test_quickstart.py | 40 ++++++------ 8 files changed, 169 insertions(+), 123 deletions(-) diff --git a/popcor/auto_template.py b/popcor/auto_template.py index 7293076..6839f2b 100644 --- a/popcor/auto_template.py +++ b/popcor/auto_template.py @@ -267,25 +267,30 @@ def is_tight(self, verbose=False, data_dict={}, tightness=None): elif tightness == "cost": return cost_tight - def get_A_list(self, var_dict=None): + def get_A_b_list(self, var_dict=None): + A_list, b_list = self.lifter.get_A0() + A_b_list = list(zip(A_list, b_list)) if var_dict is None: - A_known = [] if self.use_known: - A_known += [constraint.A_sparse_ for constraint in self.templates_known] - return A_known + [constraint.A_sparse_ for constraint in self.constraints] + A_b_list += [ + (constraint.A_sparse_, constraint.rhs) + for constraint in self.templates_known + ] + return A_b_list + [ + (constraint.A_sparse_, constraint.rhs) + for constraint in self.constraints + ] else: - A_known = [] + A_b_list = [] if self.use_known: - A_known += [constraint.A_poly_ for constraint in self.templates_known] - A_list_poly = A_known + [ - constraint.A_poly_ for constraint in self.constraints + A_b_list += [ + (constraint.A_poly_, constraint.rhs) + for constraint in self.templates_known + ] + A_list_poly = A_b_list + [ + (constraint.A_poly_, constraint.rhs) for constraint in self.constraints ] - return [A.get_matrix(var_dict) for A in A_list_poly] - - def get_A_b_list(self): - A_list = self.get_A_list() - A_b_list_all = self.lifter.get_A_b_list(A_list) - return A_b_list_all + return [(A.get_matrix(var_dict), b) for A, b in A_list_poly] def generate_minimal_subset( self, @@ -453,7 +458,7 @@ def find_local_solution(self, n_inits=None, verbose=False, plot=False): self.solver_vars.update( { f"local {solution_idx} {error_name}": err - for error_name, err in error_dict.items() + for error_name, err in error_dict.items() # type: ignore } ) @@ -517,7 +522,7 @@ def find_global_solution(self, data_dict={}): if x is not None: theta = self.lifter.get_theta(x) - cost = self.lifter.get_cost(theta, self.solver_vars["y"]) + cost = self.lifter.get_cost(theta, self.solver_vars["y"]) # type: ignore data_dict["global theta"] = theta data_dict["global cost"] = cost return True @@ -825,7 +830,7 @@ def sort_fun_sparsity(series): ) ) df = pd.DataFrame( - series, dtype="Sparse[float]", index=templates_poly.variable_dict_i + series, dtype="Sparse[float]", index=templates_poly.variable_dict_i # type: ignore ) df.dropna(axis=1, how="all", inplace=True) @@ -1033,7 +1038,7 @@ def save_matrices_sparsity(self, A_matrices=None, fname_root="", title=""): vmin = min(-np.max(Q), np.min(Q)) vmax = max(np.max(Q), -np.min(Q)) - norm = matplotlib.colors.SymLogNorm(10**-5, vmin=vmin, vmax=vmax) + norm = matplotlib.colors.SymLogNorm(10**-5, vmin=vmin, vmax=vmax) # type: ignore im1 = axs[1].matshow(Q, norm=norm) for ax in axs: @@ -1227,11 +1232,14 @@ def run( def apply(self, lifter: StateLifter, use_known: bool = False) -> list: """Apply the learned templates to a new lifter.""" constraints = lifter.apply_templates(self.templates) - + A_0, b_0 = lifter.get_A0() if use_known: # if we set use_known=True in running AutoTemplate, then we learned only # constraints that were not already known, so we need to add them to the # overall set of constraints. - A_known = lifter.get_A_known() + A_known, b_known = lifter.get_A_known() assert isinstance(A_known, list) - return A_known + [c.A_sparse_ for c in constraints] # type: ignore + A_b_list = list(zip(A_0 + A_known, b_0 + b_known)) + else: + A_b_list = list(zip(A_0, b_0)) + return A_b_list + [(c.A_sparse_, c.rhs) for c in constraints] # type: ignore diff --git a/popcor/auto_tight.py b/popcor/auto_tight.py index 7c65ee4..4697279 100644 --- a/popcor/auto_tight.py +++ b/popcor/auto_tight.py @@ -1,3 +1,5 @@ +import warnings + import matplotlib.pylab as plt import numpy as np import scipy.sparse as sp @@ -109,7 +111,7 @@ def get_A_learned( var_dict=None, method=METHOD, verbose=False, - ) -> list: + ) -> tuple[list, np.ndarray | list]: """Generate list of learned constraints by sampling the lifter. :param lifter: StateLifter object @@ -131,10 +133,19 @@ def get_A_learned( if verbose: print(f"get basis ({basis.shape})): {time.time() - t1:4.4f}") t1 = time.time() - A_learned = AutoTight.generate_matrices(lifter, basis, var_dict=var_dict) - if verbose: - print(f"get matrices ({len(A_learned)}): {time.time() - t1:4.4f}") - return A_learned + if lifter.level == "bm": + A_learned = AutoTight.generate_matrices( + lifter, basis[:, 1:], var_dict=var_dict + ) + b_learned = -basis[:, 0] + if verbose: + print(f"get matrices ({len(A_learned)}): {time.time() - t1:4.4f}") + return A_learned, b_learned + else: + A_learned = AutoTight.generate_matrices(lifter, basis, var_dict=var_dict) + if verbose: + print(f"get matrices ({len(A_learned)}): {time.time() - t1:4.4f}") + return A_learned, [0] * len(A_learned) @staticmethod def get_A_learned_simple( @@ -223,6 +234,8 @@ def generate_Y_sparse(lifter, var_subset, param_subset, factor=FACTOR, ax=None): def generate_Y(lifter, factor=FACTOR, ax=None, var_subset=None, param_subset=None): # need at least dim_Y different random setups dim_Y = lifter.get_dim_Y(var_subset, param_subset) + if lifter.level == "bm": + dim_Y += lifter.get_dim_P(param_subset) n_seeds = int(dim_Y * factor) Y = np.empty((n_seeds, dim_Y)) for seed in range(n_seeds): @@ -237,12 +250,20 @@ def generate_Y(lifter, factor=FACTOR, ax=None, var_subset=None, param_subset=Non ax.scatter(*theta[:, :2].T) x = lifter.get_x(theta=theta, parameters=parameters, var_subset=var_subset) - X = np.outer(x, x) + if np.ndim(x) == 1: + x = x[:, None] + X = x @ x.T # generates [1*x, a1*x, ..., aK*x] p = lifter.get_p(parameters=parameters, param_subset=param_subset) assert p[0] == 1 - Y[seed, :] = np.kron(p, get_vec(X)) + + row = get_vec(X) + if lifter.level == "bm": + # in Burer-Monteiro, we need to add the leading 1. + row = np.hstack([np.array([1.0]), np.array(row)]) + + Y[seed, :] = np.kron(p, row) return Y @staticmethod @@ -266,8 +287,9 @@ def get_basis( # if there is a known list of constraints, add them to the Y so that resulting nullspace is orthogonal to them if basis_known is not None: if len(A_known): - print( - "Warning: ignoring given A_known because basis_all is also given." + warnings.warn( + "Ignoring given A_known because basis_all is also given.", + UserWarning, ) Y = np.vstack([Y, basis_known.T]) elif len(A_known): diff --git a/popcor/examples/poly4_lifter.py b/popcor/examples/poly4_lifter.py index 1fd394c..754c978 100644 --- a/popcor/examples/poly4_lifter.py +++ b/popcor/examples/poly4_lifter.py @@ -57,9 +57,9 @@ def get_A_known(self, output_poly=False, add_redundant=False, var_dict=None): A_1[self.HOM, "z0"] = -1 A_1["t", "t"] = 2 if output_poly: - return [A_1] + return [A_1], [0] else: - return [A_1.get_matrix(self.var_dict)] + return [A_1.get_matrix(self.var_dict)], [0] def generate_random_setup(self): self.theta_ = np.array([-1]) diff --git a/popcor/examples/stereo1d_lifter.py b/popcor/examples/stereo1d_lifter.py index 366c48d..a9a5a2a 100644 --- a/popcor/examples/stereo1d_lifter.py +++ b/popcor/examples/stereo1d_lifter.py @@ -156,7 +156,7 @@ def get_A_known(self, var_dict=None, output_poly=False, add_redundant=False): if add_redundant: A_known += self.get_A_known_redundant(var_dict, output_poly) - return A_known + return A_known, [0] * len(A_known) def get_A_known_redundant(self, var_dict=None, output_poly=False): import itertools @@ -181,7 +181,7 @@ def get_A_known_redundant(self, var_dict=None, output_poly=False): A_known.append(A) else: A_known.append(A.get_matrix(variables=self.var_dict)) - return A_known + return A_known, [0] * len(A_known) def get_cost(self, theta, y): if np.ndim(theta) == 0 or (np.ndim(theta) == 1 and len(theta) == 1): diff --git a/popcor/examples/stereo3d_lifter.py b/popcor/examples/stereo3d_lifter.py index 63d571c..514a1d2 100644 --- a/popcor/examples/stereo3d_lifter.py +++ b/popcor/examples/stereo3d_lifter.py @@ -1,3 +1,9 @@ +"""Stereo 3D lifter example. + +Contains Stereo3DLifter for stereo-camera localization in 3D. The main +information about the example is in the class docstring. +""" + import pickle import numpy as np @@ -8,9 +14,15 @@ def change_dimensions(a, y): - p_w = np.concatenate([a, np.ones((a.shape[0], 1))], axis=1) - y_mat = np.c_[[*y]] # N x 2 - return p_w[:, :, None], y_mat[:, :, None] + """Convert landmarks and measurements to solver-friendly shapes. + + Returns p_w with shape (N, 4, 1) (homogeneous world points) and y with + shape (N, 2, 1) (measurements) for subsequent vectorized processing. + """ + a = np.asarray(a) + p_w = np.hstack((a, np.ones((a.shape[0], 1), dtype=a.dtype)))[:, :, None] + y_mat = np.asarray(y).reshape((a.shape[0], -1)) + return p_w, y_mat[:, :, None] GTOL = 1e-6 @@ -27,7 +39,7 @@ class Stereo3DLifter(StereoLifter): where - :math:`p_j` are known landmarks (in homogeneous coordinates), - - :math:`u_j` are pixel measurements (4 elements: two pixel coordinates in left image and two in right image), + - :math:`u_j` are pixel measurements (four elements: two pixel coordinates in the left image and two in the right image), - :math:`q_j = T(\\theta) p_j` are the (homogeneous) coordinates of landmark j in the (unknown) camera frame, parameterized by :math:`T(\\theta)`, and - :math:`M` is the stereo camera calibration matrix. Here, it is given by @@ -42,28 +54,21 @@ class Stereo3DLifter(StereoLifter): where :math:`f_u, f_v` are horizontal and vertical focal lengths, :math:`c_u,c_v` are image center points in pixels and :math:`b` is the camera baseline. - This example is treated in more details in `this paper `_. + This example is treated in more detail in `this paper `_. """ def __init__(self, n_landmarks, level="no", param_level="no", variable_list=None): - self.W = np.stack([np.eye(4)] * n_landmarks) + """Create a Stereo3DLifter and initialize default per-landmark weights.""" + # Pre-allocate W as repeated identity matrices; repeat is slightly faster than building a list. + self.W = np.repeat(np.eye(4)[None, :, :], n_landmarks, axis=0) - super().__init__( - n_landmarks=n_landmarks, - level=level, - param_level=param_level, - d=3, - variable_list=variable_list, - ) - - @property - def M_matrix(self): + # Precompute the M_matrix once for efficiency. f_u = 484.5 f_v = 484.5 c_u = 322 c_v = 247 b = 0.24 - return np.array( + self._M_matrix = np.array( [ [f_u, 0, c_u, f_u * b / 2], [0, f_v, c_v, 0], @@ -72,16 +77,24 @@ def M_matrix(self): ] ) + super().__init__( + n_landmarks=n_landmarks, + level=level, + param_level=param_level, + d=3, + variable_list=variable_list, + ) + + @property + def M_matrix(self): + """Stereo camera calibration matrix used by the cost and solver.""" + return self._M_matrix + @staticmethod def from_file(fname): + """Load a Stereo3DLifter instance state from a file created by to_file.""" with open(fname, "rb") as f: - y_ = pickle.load(f) - landmarks = pickle.load(f) - theta = pickle.load(f) - - level = pickle.load(f) - param_level = pickle.load(f) - variable_list = pickle.load(f) + y_, landmarks, theta, level, param_level, variable_list = pickle.load(f) lifter = Stereo3DLifter( n_landmarks=landmarks.shape[0], level=level, @@ -95,22 +108,27 @@ def from_file(fname): return lifter def to_file(self, fname): + """Serialize the lifter's minimal state to a file for later restoration.""" with open(fname, "wb") as f: - pickle.dump(self.y_, f) - pickle.dump(self.landmarks, f) - - pickle.dump(self.theta, f) - pickle.dump(self.level, f) - pickle.dump(self.param_level, f) - pickle.dump(self.variable_list, f) + pickle.dump( + ( + self.y_, + self.landmarks, + self.theta, + self.level, + self.param_level, + self.variable_list, + ), + f, + ) def get_cost(self, theta, y, W=None): - """ - :param t: can be either - - x, y, z, yaw, pitch roll: vector of unknowns, or - - [c1, c2, c3, x, y, z], the theta vector (flattened C and x, y, z) - """ + """Evaluate the reprojection cost for a given parameter vector theta. + theta can be either: + - a pose vector (x, y, z, yaw, pitch, roll), or + - the full theta vector containing flattened C and x,y,z depending on parameterization. + """ if W is None: W = self.W a = self.landmarks @@ -126,10 +144,10 @@ def get_cost(self, theta, y, W=None): return cost def local_solver(self, t0, y, W=None, verbose=False, **kwargs): - """ - :param t_init: same options asfor t in cost. - """ + """Run the local solver starting from initial pose t0 and measurements y. + Returns an estimated pose (theta), solver info dict, and the final cost. + """ if W is None: W = self.W @@ -149,10 +167,9 @@ def local_solver(self, t0, y, W=None, verbose=False, **kwargs): ) if verbose: - print("Stereo3D local solver:", info["msg"]) + print("Stereo3D local solver:", info.get("msg", "")) if StereoLifter.NORMALIZE: - cost /= self.n_landmarks * self.d x_hat = get_theta_from_T(T_hat) @@ -165,7 +182,7 @@ def local_solver(self, t0, y, W=None, verbose=False, **kwargs): info["cost"] = cost - if info["success"]: + if info.get("success", False): return x_hat, info, cost else: return None, info, cost diff --git a/popcor/utils/constraint.py b/popcor/utils/constraint.py index 4365433..6d2570e 100644 --- a/popcor/utils/constraint.py +++ b/popcor/utils/constraint.py @@ -143,6 +143,7 @@ def __init__( mat_param_dict=None, known=False, template_idx=0, + rhs=0, ): self.index = index self.mat_var_dict = mat_var_dict @@ -162,6 +163,8 @@ def __init__( # list of applied constraints derived from this constraint. self.applied_list = [] + # what this cosntraint equals to: Ax = rhs + self.rhs = rhs @staticmethod # @profile diff --git a/tests/test_autotight.py b/tests/test_autotight.py index 70253da..6401229 100644 --- a/tests/test_autotight.py +++ b/tests/test_autotight.py @@ -1,3 +1,5 @@ +import warnings + import numpy as np from popcor import AutoTight @@ -15,16 +17,16 @@ def test_constraints(): # testing qrp, our go-to method, for all example lifters for lifter in all_lifters(SEED): - A_learned = AutoTight.get_A_learned(lifter=lifter, method="qrp") - constraints_test_with_tol(lifter, A_learned, tol=1e-4) + A_learned, b_learned = AutoTight.get_A_learned(lifter=lifter, method="qrp") + constraints_test_with_tol(lifter, A_learned, b_learned, tol=1e-4) def test_constraints_params(): # for a couple of representative lifters, checking that all parameter levels are working for param_level in ["no", "p", "ppT"]: for lifter in example_lifters(seed=SEED, param_level=param_level): - A_learned = AutoTight.get_A_learned(lifter=lifter, verbose=False) - constraints_test_with_tol(lifter, A_learned, tol=1e-4) + A_learned, b_learned = AutoTight.get_A_learned(lifter=lifter, verbose=False) + constraints_test_with_tol(lifter, A_learned, b_learned, tol=1e-4) def test_constraints_methods(): @@ -32,10 +34,12 @@ def test_constraints_methods(): for lifter in example_lifters(seed=SEED, param_level="p"): num_learned = None for method in ["qrp", "svd", "qr"]: - A_learned = AutoTight.get_A_learned( - lifter=lifter, verbose=False, method=method - ) - constraints_test_with_tol(lifter, A_learned, tol=1e-4) + with warnings.catch_warnings(): + warnings.filterwarnings("ignore", category=UserWarning) + A_learned, b_learned = AutoTight.get_A_learned( + lifter=lifter, verbose=False, method=method + ) + constraints_test_with_tol(lifter, A_learned, b_learned, tol=1e-4) # make sure each method finds the same number of matrices if num_learned is None: @@ -46,35 +50,27 @@ def test_constraints_methods(): def test_constraints_stereo(): np.random.seed(0) - n_landmarks = 1 # z_0 and z_1 - + warnings.filterwarnings("ignore", category=UserWarning) for param_level in ["p", "ppT"]: - print(f"1D, {param_level}") - lifter = Stereo1DLifter(n_landmarks=n_landmarks, param_level=param_level) - A_learned = AutoTight.get_A_learned(lifter) + print(f"learning Stereo1D, {param_level}") + lifter = Stereo1DLifter(n_landmarks=1, param_level=param_level) + A_learned, b_learned = AutoTight.get_A_learned(lifter) lifter.test_constraints(A_learned, errors="raise", n_seeds=1) - print(f"2D, {param_level}") - lifter = Stereo2DLifter( - n_landmarks=n_landmarks, param_level=param_level, level="urT" - ) - A_learned = AutoTight.get_A_learned(lifter) + print(f"learning Stereo2D, {param_level}") + lifter = Stereo2DLifter(n_landmarks=1, param_level=param_level, level="urT") + A_learned, b_learned = AutoTight.get_A_learned(lifter) lifter.test_constraints(A_learned, errors="raise", n_seeds=1) - print(f"3D, {param_level}") - lifter = Stereo3DLifter( - n_landmarks=n_landmarks, param_level=param_level, level="urT" - ) - A_learned = AutoTight.get_A_learned(lifter) + print(f"learning Stereo3D, {param_level}") + lifter = Stereo3DLifter(n_landmarks=1, param_level=param_level, level="urT") + A_learned, b_learned = AutoTight.get_A_learned(lifter) lifter.test_constraints(A_learned, errors="raise", n_seeds=1) if __name__ == "__main__": - # import pytest - # pytest.main([__file__, "-s"]) - # print("all tests passed") - test_constraints_stereo() test_constraints() test_constraints_params() test_constraints_methods() + print("all tests passed") diff --git a/tests/test_quickstart.py b/tests/test_quickstart.py index e25e06c..74aa3b9 100644 --- a/tests/test_quickstart.py +++ b/tests/test_quickstart.py @@ -36,12 +36,13 @@ def test_solve_sdp(): Q = lifter.get_Q() # the equality constraints - A_known = lifter.get_A_known() + A_known, b_known = lifter.get_A_known() # the homogenization constraint - A_0 = lifter.get_A0() + A_0, b_0 = lifter.get_A0() + constraints = list(zip(A_0 + A_known, b_0 + b_known)) - X, info = solve_sdp(Q, [(A_0, 1.0)] + [(A_i, 0.0) for A_i in A_known]) + X, info = solve_sdp(Q, constraints) assert X is not None # if X is rank one, the global optimum can be found in element X_10 of the matrix. @@ -68,15 +69,15 @@ def test_autotight(): # solve the SDP -- it is not cost tight! Q = lifter.get_Q(noise=1e-1) - A_known = lifter.get_A_known() - A_0 = lifter.get_A0() + A_known, b_known = lifter.get_A_known() + A_0, b_0 = lifter.get_A0() # solve locally starting at ground truth x, info_local, cost = lifter.local_solver(lifter.theta, y=lifter.y_) assert x is not None assert info_local["success"] - constraints = lifter.get_A_b_list(A_list=A_known) + constraints = list(zip(A_0 + A_known, b_0 + b_known)) X, info_sdp = solve_sdp(Q, constraints) gap = auto_tight.get_duality_gap( @@ -85,8 +86,9 @@ def test_autotight(): assert gap > 0.1 # learn matrices and solve again - A_learned = auto_tight.get_A_learned(lifter) - constraints = lifter.get_A_b_list(A_list=A_learned) + A_learned, b_learned = auto_tight.get_A_learned(lifter) + A_0, b_0 = lifter.get_A0() + constraints = list(zip(A_0 + A_learned, b_0 + b_learned)) X, info_sdp_learned = solve_sdp(Q, constraints) gap = auto_tight.get_duality_gap( info_local["cost"], info_sdp_learned["cost"], relative=True @@ -124,12 +126,10 @@ def test_autotemplate(): # apply the templates to a different lifter new_lifter = Stereo1DLifter(n_landmarks=5, param_level="p") - A_learned = auto_template.apply(new_lifter, use_known=True) - + constraints = auto_template.apply(new_lifter, use_known=True) Q = new_lifter.get_Q(noise=1e-1) # adds homogenization constraint and all b_i terms - constraints = new_lifter.get_A_b_list(A_list=A_learned) X, info_sdp = solve_sdp(Q, constraints) # evaluate duality gap @@ -149,9 +149,10 @@ def test_autotemplate_literal(): new_lifter = Stereo1DLifter(n_landmarks=5, param_level="p") Q = new_lifter.get_Q(noise=1e-1) - A_known = new_lifter.get_A_known() - A_red = new_lifter.get_A_known_redundant() - constraints = new_lifter.get_A_b_list(A_list=A_known + A_red) + A_0, b_0 = new_lifter.get_A0() + A_known, b_known = new_lifter.get_A_known() + A_red, b_red = new_lifter.get_A_known_redundant() + constraints = list(zip(A_0 + A_known + A_red, b_0 + b_known + b_red)) X, info_sdp = solve_sdp(Q, constraints) @@ -163,16 +164,15 @@ def test_autotemplate_literal(): if __name__ == "__main__": - # import warnings - # make sure warnings raise errors, for debugging - # warnings.simplefilter("error") + import warnings - # import pytest - # pytest.main([__file__, "-s"]) - # print("all tests passed") + # make sure warnings raise errors, for debugging + warnings.simplefilter("error") + warnings.simplefilter("ignore", UserWarning) test_setup_problem() test_solve_sdp() test_autotight() test_autotemplate() test_autotemplate_literal() + print("all tests passed") From 18f2c2066506aef75b1de957eee28a773f2559e7 Mon Sep 17 00:00:00 2001 From: Frederike Duembgen Date: Fri, 22 Aug 2025 19:54:00 +0200 Subject: [PATCH 03/10] Fix sampling and use UserWarnings --- popcor/base_lifters/range_only_lifter.py | 18 +++++++++--------- 1 file changed, 9 insertions(+), 9 deletions(-) diff --git a/popcor/base_lifters/range_only_lifter.py b/popcor/base_lifters/range_only_lifter.py index b9768db..258606f 100644 --- a/popcor/base_lifters/range_only_lifter.py +++ b/popcor/base_lifters/range_only_lifter.py @@ -1,3 +1,4 @@ +import warnings from abc import ABC, abstractmethod import autograd.numpy as anp @@ -178,16 +179,15 @@ def sample_theta(self): sample = ( np.random.rand(self.landmarks.shape[1]) ) * self.SCALE # between 0 and SCALE - if np.all( - np.linalg.norm(sample[None, :] - self.landmarks, axis=1) - > self.MIN_DIST - ): - samples[i, :] = sample - elif j == MAX_TRIALS - 1: - print( - f"Warning: did not find valid sample in {MAX_TRIALS} trials. Using last sample." + distances = np.linalg.norm(sample[None, :] - self.landmarks, axis=1) + if np.all(distances > self.MIN_DIST): + break + if j == MAX_TRIALS - 1: + warnings.warn( + f"Did not find valid sample in {MAX_TRIALS} trials. Using last sample with distances {distances.round(4)}.", + UserWarning, ) - samples[i, :] = sample + samples[i, :] = sample return samples.flatten("C") def get_residuals(self, t, y, squared=True, ad=False): From 418db45cfeab47506879e3a932ea0ab818fd673c Mon Sep 17 00:00:00 2001 From: Frederike Duembgen Date: Fri, 22 Aug 2025 20:02:21 +0200 Subject: [PATCH 04/10] More UserWarnings and fix tests to include variable bs --- popcor/base_lifters/_base_class.py | 42 +++++++++++------------------ popcor/base_lifters/state_lifter.py | 16 +++++++---- 2 files changed, 26 insertions(+), 32 deletions(-) diff --git a/popcor/base_lifters/_base_class.py b/popcor/base_lifters/_base_class.py index d1b4572..82bcdec 100644 --- a/popcor/base_lifters/_base_class.py +++ b/popcor/base_lifters/_base_class.py @@ -639,7 +639,9 @@ def get_vec_around_gt(self, delta: float = 0): """ return self.theta + np.random.normal(size=self.theta.shape, scale=delta) - def test_constraints(self, A_list, errors: str = "raise", n_seeds: int = 3): + def test_constraints( + self, A_list, b_list=None, errors: str = "raise", n_seeds: int = 3 + ): """ :param A_list: can be either list of sparse matrices, or poly matrices :param errors: "raise" or "print" detected violations. @@ -650,17 +652,24 @@ def test_constraints(self, A_list, errors: str = "raise", n_seeds: int = 3): for j, A in enumerate(A_list): if isinstance(A, PolyMatrix): A = A.get_matrix(self.get_var_dict(unroll_keys=True)) + if b_list is None: + b = 0.0 + else: + b = b_list[j] for i in range(n_seeds): if i == 0: - x = self.get_x().flatten() + x = self.get_x() else: np.random.seed(i) t = self.sample_theta() p = self.sample_parameters() - x = self.get_x(theta=t, parameters=p).flatten() + x = self.get_x(theta=t, parameters=p) - constraint_violation = abs(float(x.T @ A @ x)) + if np.ndim(x) == 2 and x.shape[1] > 1: + constraint_violation = abs(np.trace(x.T @ A @ x) - b) + else: + constraint_violation = abs(float(x.T @ A @ x - b)) max_violation = max(max_violation, constraint_violation) if constraint_violation > self.EPS_ERROR: msg = f"big violation at {j}: {constraint_violation:.1e}" @@ -675,35 +684,14 @@ def test_constraints(self, A_list, errors: str = "raise", n_seeds: int = 3): raise ValueError(errors) return max_violation, j_bad - def get_A0(self, var_subset=None): + def get_A0(self, var_subset=None) -> tuple[list, np.ndarray | list]: if var_subset is not None: var_dict = {k: self.var_dict[k] for k in var_subset} else: var_dict = self.var_dict A0 = PolyMatrix() A0[self.HOM, self.HOM] = 1.0 - return [A0.get_matrix(var_dict)], [1.0] - - def get_A_b_list(self, A_list, var_subset=None): - """get equality constraint tuples (Ai, bi) s.t. x.t @ Ai @ bi, 0 - - :param A_list: Normally, this is just the list of equality constaints that equal zero. We will add the homogenization. - For certain cases, such as the RotationLifter with level="bm", this is already a tuple of (Ai, bi). - - """ - if var_subset is None: - var_subset = self.var_dict - - if isinstance(A_list, list): - # TODO(FD): do more with var_subset - assert self.HOM in var_subset - return [(self.get_A0(var_subset), 1.0)] + [(A, 0.0) for A in A_list] - else: - assert isinstance(A_list, tuple) - A0_list, b0_list = self.get_A0(var_subset) - return [ - (Ai, bi) for Ai, bi in zip(A0_list + A_list[0], b0_list + A_list[1]) - ] + return [A0.get_matrix(var_dict)], np.array([1.0]) def sample_parameters_landmarks(self, landmarks: np.ndarray): """Used by RobustPoseLifter, RangeOnlyLocLifter: the default way of adding landmarks to parameters.""" diff --git a/popcor/base_lifters/state_lifter.py b/popcor/base_lifters/state_lifter.py index 0000f34..e7ca7cc 100644 --- a/popcor/base_lifters/state_lifter.py +++ b/popcor/base_lifters/state_lifter.py @@ -1,3 +1,4 @@ +import warnings from abc import abstractmethod import numpy as np @@ -45,7 +46,10 @@ def __init__( self.variable_list = self.VARIABLE_LIST if (param_level != "no") and (n_parameters == 1): - print("Warning: make sure to give the correct n_parameters for the level.") + warnings.warn( + "Warning: make sure to give the correct n_parameters for the level.", + category=UserWarning, + ) super().__init__(d, param_level, n_parameters) @@ -150,9 +154,9 @@ def get_hess(self, theta, y=None) -> np.ndarray: def get_cost(self, theta, y: np.ndarray | None = None) -> float: """Compute the cost of the given state theta. This uses the simple form x.T @ Q @ x. Consider overwriting this for more efficient computations.""" - print( - "Warning: using default get_cost, which may be less efficient than a custom one." - ) + # print( + # "Warning: using default get_cost, which may be less efficient than a custom one." + # ) if y is not None: Q = self.get_Q_from_y(y) else: @@ -198,7 +202,9 @@ def local_solver( f"Warning: ignoreing method argument {method} in local_solver, using default (IPOPT)." ) - Constraints = self.get_A_b_list(A_list=self.get_A_known()) + A_0, b_0 = self.get_A0() + A_known, b_known = self.get_A_known() + Constraints = list(zip(A_0 + A_known, b_0 + b_known)) x0 = self.get_x(theta=t0) if np.ndim(x0) == 1: x0 = x0.reshape((-1, 1)) From 4a3c108ef0af7e252dce39f37b9f8be6b8782d73 Mon Sep 17 00:00:00 2001 From: Frederike Duembgen Date: Fri, 22 Aug 2025 20:03:42 +0200 Subject: [PATCH 05/10] Add new RotationLifter to general tests, make b general --- popcor/utils/test_tools.py | 28 ++++++++++++++++++---------- 1 file changed, 18 insertions(+), 10 deletions(-) diff --git a/popcor/utils/test_tools.py b/popcor/utils/test_tools.py index 0c80508..0929514 100644 --- a/popcor/utils/test_tools.py +++ b/popcor/utils/test_tools.py @@ -51,8 +51,10 @@ (Stereo1DLifter, dict(n_landmarks=n_landmarks, param_level="p")), # skip (Stereo2DLifter, dict(n_landmarks=n_landmarks)), (Stereo3DLifter, dict(n_landmarks=n_landmarks)), - (RotationLifter, dict(d=2)), # ok - (RotationLifter, dict(d=3)), # ok + (RotationLifter, dict(d=2, level="no")), + (RotationLifter, dict(d=3, level="no")), + (RotationLifter, dict(d=2, level="bm")), + (RotationLifter, dict(d=3, level="bm")), ] ExampleLifters = [ @@ -66,21 +68,27 @@ ] -def constraints_test_with_tol(lifter, A_list, tol): - x = lifter.get_x().astype(float).reshape((-1, 1)) - for Ai in A_list: - err = abs((x.T @ Ai @ x)[0, 0]) +def constraints_test_with_tol(lifter, A_list, b_list, tol): + x = lifter.get_x().astype(float) + if np.ndim(x) == 1: + x = x[:, None] + for i, Ai in enumerate(A_list): + if b_list is None: + bi = 0.0 + else: + bi = b_list[i] + err = abs(np.trace(x.T @ Ai @ x) - bi) assert err < tol, err ai = get_vec(Ai.toarray()) - xvec = get_vec(np.outer(x, x)) + xvec = get_vec(x @ x.T) assert isinstance(ai, np.ndarray) - np.testing.assert_allclose(ai @ xvec, 0.0, atol=tol) + np.testing.assert_allclose(ai @ xvec, bi, atol=tol) ai = get_vec(Ai) - xvec = get_vec(np.outer(x, x)) + xvec = get_vec(x @ x.T) assert isinstance(ai, np.ndarray) - np.testing.assert_allclose(ai @ xvec, 0.0, atol=tol) + np.testing.assert_allclose(ai @ xvec, bi, atol=tol) # Below, we always reset seeds to make sure tests are reproducible. From 2e4cc3241f7e46eea3ced598b0ea2fff560c3dcd Mon Sep 17 00:00:00 2001 From: Frederike Duembgen Date: Fri, 22 Aug 2025 20:05:36 +0200 Subject: [PATCH 06/10] Allow for csr and scs matrices --- popcor/utils/common.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/popcor/utils/common.py b/popcor/utils/common.py index fd4f0fc..8318339 100644 --- a/popcor/utils/common.py +++ b/popcor/utils/common.py @@ -133,7 +133,7 @@ def get_vec(mat, correct=True, sparse=False) -> np.ndarray | sp.csr_matrix | Non mat = deepcopy(mat) if correct: - if isinstance(mat, sp.csc_matrix): + if isinstance(mat, (sp.csc_matrix, sp.csr_matrix)): ii, jj = mat.nonzero() mat[ii, jj] *= np.sqrt(2.0) diag = ii == jj From 7523f0d62f2f55331ac78238e02539560f2d8c56 Mon Sep 17 00:00:00 2001 From: Frederike Duembgen Date: Fri, 22 Aug 2025 20:25:58 +0200 Subject: [PATCH 07/10] Fix rotation averaging test --- popcor/auto_template.py | 5 +++-- popcor/base_lifters/_base_class.py | 4 ++-- popcor/examples/rotation_lifter.py | 9 +++++---- 3 files changed, 10 insertions(+), 8 deletions(-) diff --git a/popcor/auto_template.py b/popcor/auto_template.py index 6839f2b..a0e70fb 100644 --- a/popcor/auto_template.py +++ b/popcor/auto_template.py @@ -268,8 +268,8 @@ def is_tight(self, verbose=False, data_dict={}, tightness=None): return cost_tight def get_A_b_list(self, var_dict=None): - A_list, b_list = self.lifter.get_A0() - A_b_list = list(zip(A_list, b_list)) + A_0, b_0 = self.lifter.get_A0() + A_b_list = list(zip(A_0, b_0)) if var_dict is None: if self.use_known: A_b_list += [ @@ -1232,6 +1232,7 @@ def run( def apply(self, lifter: StateLifter, use_known: bool = False) -> list: """Apply the learned templates to a new lifter.""" constraints = lifter.apply_templates(self.templates) + A_0, b_0 = lifter.get_A0() if use_known: # if we set use_known=True in running AutoTemplate, then we learned only diff --git a/popcor/base_lifters/_base_class.py b/popcor/base_lifters/_base_class.py index 82bcdec..d7bef4e 100644 --- a/popcor/base_lifters/_base_class.py +++ b/popcor/base_lifters/_base_class.py @@ -684,14 +684,14 @@ def test_constraints( raise ValueError(errors) return max_violation, j_bad - def get_A0(self, var_subset=None) -> tuple[list, np.ndarray | list]: + def get_A0(self, var_subset=None) -> tuple[list, list]: if var_subset is not None: var_dict = {k: self.var_dict[k] for k in var_subset} else: var_dict = self.var_dict A0 = PolyMatrix() A0[self.HOM, self.HOM] = 1.0 - return [A0.get_matrix(var_dict)], np.array([1.0]) + return [A0.get_matrix(var_dict)], [1.0] def sample_parameters_landmarks(self, landmarks: np.ndarray): """Used by RobustPoseLifter, RangeOnlyLocLifter: the default way of adding landmarks to parameters.""" diff --git a/popcor/examples/rotation_lifter.py b/popcor/examples/rotation_lifter.py index 711e60b..15393ee 100644 --- a/popcor/examples/rotation_lifter.py +++ b/popcor/examples/rotation_lifter.py @@ -402,7 +402,7 @@ def test_and_add(self, A_list, Ai, output_poly, b_list=[], bi=0.0): A_list.append(Ai_sparse) b_list.append(bi) - def get_A0(self, var_subset=None): + def get_A0(self, var_subset=None) -> tuple[list, list]: """Return the homogenization constraint A0 for chosen level (either h^2=1, or H'H=I).""" if var_subset is None: var_subset = self.var_dict @@ -413,7 +413,7 @@ def get_A0(self, var_subset=None): A_orth, b_orth = get_orthogonal_constraints( self.HOM, None, self.d, self.level ) - return [Ai.get_matrix(var_subset) for Ai in A_orth], b_orth + return [Ai.get_matrix(var_subset) for Ai in A_orth], list(b_orth) def get_A_known(self, var_dict=None, output_poly=False, add_redundant=False): """Return known linear constraints (A and b). If level 'no' returns A_list; if 'bm' returns (A_list, b_list).""" @@ -590,8 +590,9 @@ def __repr__(self): plt.show(block=False) Q = lifter.get_Q_from_y(y=y, output_poly=False) - A_known = lifter.get_A_known(output_poly=False, add_redundant=True) - constraints = lifter.get_A_b_list(A_known) + A_known, b_known = lifter.get_A_known(output_poly=False, add_redundant=True) + A_0, b_0 = lifter.get_A0() + constraints = list(zip(A_known + A_0, b_known + b_0)) fig, axs = plt.subplots(1, len(constraints) + 1) fig.set_size_inches(3 * (len(constraints) + 1), 3) From 136f7b38a79a47baa9762ce0ed0e93122d04a42b Mon Sep 17 00:00:00 2001 From: Frederike Duembgen Date: Sat, 23 Aug 2025 10:36:34 +0200 Subject: [PATCH 08/10] Fix get_A_known for all, make sure constraint rank check passes for bm --- popcor/auto_template.py | 5 ++-- popcor/base_lifters/_base_class.py | 14 ++++++++--- popcor/base_lifters/robust_pose_lifter.py | 2 +- popcor/base_lifters/stereo_lifter.py | 8 +++++-- popcor/examples/poly6_lifter.py | 6 +++-- popcor/examples/ro_nsq_lifter.py | 2 +- popcor/examples/ro_sq_lifter.py | 2 +- popcor/utils/constraint.py | 2 ++ tests/test_solvers.py | 5 +++- tests/test_state_lifter.py | 29 +++++++++++++++++------ 10 files changed, 55 insertions(+), 20 deletions(-) diff --git a/popcor/auto_template.py b/popcor/auto_template.py index a0e70fb..8c83055 100644 --- a/popcor/auto_template.py +++ b/popcor/auto_template.py @@ -737,8 +737,8 @@ def get_known_templates(self, unroll=False): # TODO(FD) we should not always recompute from scratch, but it's not very expensive so it's okay for now. target_dict = self.lifter.get_var_dict(unroll_keys=unroll) - for i, Ai in enumerate( - self.lifter.get_A_known(var_dict=target_dict, output_poly=True) + for i, (Ai, bi) in enumerate( + zip(*self.lifter.get_A_known(var_dict=target_dict, output_poly=True)) ): template = Constraint.init_from_A_poly( lifter=self.lifter, @@ -748,6 +748,7 @@ def get_known_templates(self, unroll=False): template_idx=self.constraint_index, mat_var_dict=self.lifter.var_dict, compute_polyrow_b=True, + rhs=bi, ) self.constraint_index += 1 templates_known.append(template) diff --git a/popcor/base_lifters/_base_class.py b/popcor/base_lifters/_base_class.py index d7bef4e..776dde9 100644 --- a/popcor/base_lifters/_base_class.py +++ b/popcor/base_lifters/_base_class.py @@ -508,9 +508,11 @@ def compute_Ai(self, templates, var_dict, param_dict): A_list.append(A_poly) return A_list - def get_constraint_rank(self, A_list_poly, output_sorted=False): + def get_constraint_rank(self, A_list_poly, b_known=None, output_sorted=False): """Find the number of independent constraints when they are of the form A_i @ x = 0. + For bm lifting, X is nxd, thus we check the linear independence of vec(A_i @ X). + :param A_list_poly: list of constraints matrices :return: rank (int) and sorted matrices (list, if output_sorted=True) where the first rank matrices correspond to @@ -521,12 +523,18 @@ def get_constraint_rank(self, A_list_poly, output_sorted=False): current_rank = 0 independent_indices = [] dependent_indices = [] - basis_incremental = np.zeros((len(x), 0)) + if self.level == "bm": # type: ignore + basis_incremental = np.zeros((x.size + 1, 0)) + else: + basis_incremental = np.zeros((x.size, 0)) for i, Ai in enumerate(A_list_poly): if isinstance(Ai, PolyMatrix): new_candidate = (Ai.get_matrix(self.var_dict) @ x).reshape((-1, 1)) else: - new_candidate = (Ai @ x).reshape((-1, 1)) + new_candidate = (Ai @ x).flatten()[:, None] + if self.level == "bm": # type: ignore + assert b_known is not None + new_candidate = np.vstack([new_candidate, b_known[i]]) basis_candidate = np.hstack([basis_incremental, new_candidate]) new_rank = np.linalg.matrix_rank(basis_candidate) if new_rank == current_rank + 1: diff --git a/popcor/base_lifters/robust_pose_lifter.py b/popcor/base_lifters/robust_pose_lifter.py index 6ae0c5c..02ff1cf 100644 --- a/popcor/base_lifters/robust_pose_lifter.py +++ b/popcor/base_lifters/robust_pose_lifter.py @@ -510,7 +510,7 @@ def get_A_known(self, var_dict=None, output_poly=False, add_redundant=False): constraint[j] = -1.0 Ai[f"c", f"w_{i}"] = constraint[:, None] self.test_and_add(A_list, Ai, output_poly=output_poly) - return A_list + return A_list, [0.0] * len(A_list) def get_B_known(self): """Get inequality constraints of the form x.T @ B @ x <= 0. diff --git a/popcor/base_lifters/stereo_lifter.py b/popcor/base_lifters/stereo_lifter.py index 3bfb9ff..f288951 100644 --- a/popcor/base_lifters/stereo_lifter.py +++ b/popcor/base_lifters/stereo_lifter.py @@ -1,3 +1,4 @@ +import warnings from abc import ABC, abstractmethod import numpy as np @@ -227,8 +228,11 @@ def get_A_known(self, var_dict=None, output_poly=False, add_redundant=False): a3) cz @ pj * u_zj + tz*u_zj - h*h = 0 ------1z------- --2z--- """ - print("not using known stereo templates because they depend on the landmarks.") - return [] + warnings.warn( + "Not using known stereo templates because they depend on the landmarks.", + UserWarning, + ) + return [], [] # x contains: [c1, c2, c3, t] # z contains: [u_xj, u_yj, u_zj, H.O.T.] diff --git a/popcor/examples/poly6_lifter.py b/popcor/examples/poly6_lifter.py index 1537cbc..8243b0f 100644 --- a/popcor/examples/poly6_lifter.py +++ b/popcor/examples/poly6_lifter.py @@ -68,9 +68,11 @@ def get_A_known(self, output_poly=False, add_redundant=True, var_dict=None): A_list.append(B_0) if output_poly: - return A_list + return A_list, [0.0] * len(A_list) else: - return [A_i.get_matrix(self.var_dict) for A_i in A_list] + return [A_i.get_matrix(self.var_dict) for A_i in A_list], [0.0] * len( + A_list + ) def get_D(self, that): D = np.array( diff --git a/popcor/examples/ro_nsq_lifter.py b/popcor/examples/ro_nsq_lifter.py index 3818220..e0eb54c 100644 --- a/popcor/examples/ro_nsq_lifter.py +++ b/popcor/examples/ro_nsq_lifter.py @@ -128,7 +128,7 @@ def get_A_known(self, var_dict=None, output_poly=False, add_redundant=False): raise NotImplementedError( "get_A_known not implemented yet for simple level" ) - return A_list + return A_list, [0.0] * len(A_list) def get_residuals(self, t, y, ad=False): return super().get_residuals(t, y, ad=ad, squared=False) diff --git a/popcor/examples/ro_sq_lifter.py b/popcor/examples/ro_sq_lifter.py index e5e4697..1f352c1 100644 --- a/popcor/examples/ro_sq_lifter.py +++ b/popcor/examples/ro_sq_lifter.py @@ -139,7 +139,7 @@ def get_A_known(self, var_dict=None, output_poly=False, add_redundant=False): A_list.append(A) else: A_list.append(A.get_matrix(self.var_dict)) - return A_list + return A_list, [0.0] * len(A_list) def get_residuals(self, t, y, ad=False): return super().get_residuals(t, y, ad=ad, squared=True) diff --git a/popcor/utils/constraint.py b/popcor/utils/constraint.py index 6d2570e..a54d13a 100644 --- a/popcor/utils/constraint.py +++ b/popcor/utils/constraint.py @@ -223,6 +223,7 @@ def init_from_A_poly( index: int = 0, template_idx: int = 0, compute_polyrow_b=False, + rhs=0, ): Ai_sparse_small = A_poly.get_matrix(variables=mat_var_dict) ai = get_vec(Ai_sparse_small, correct=True) @@ -244,6 +245,7 @@ def init_from_A_poly( index=index, mat_var_dict=mat_var_dict, template_idx=template_idx, + rhs=rhs, ) @staticmethod diff --git a/tests/test_solvers.py b/tests/test_solvers.py index 19ee91c..8b35eea 100644 --- a/tests/test_solvers.py +++ b/tests/test_solvers.py @@ -134,7 +134,10 @@ def test_cost(noise=0.0): cost = lifter.get_cost(theta, y) x = lifter.get_x(theta) - costQ = x.T @ Q @ x + if np.ndim(x) == 2 and x.shape[1] > 1: + costQ = np.trace(x.T @ Q @ x) + else: + costQ = x.T @ Q @ x assert abs(cost - costQ) < 1e-6, (cost, costQ) if noise == 0 and not isinstance(lifter, PolyLifter) and not lifter.robust: diff --git a/tests/test_state_lifter.py b/tests/test_state_lifter.py index b005d33..f271da3 100644 --- a/tests/test_state_lifter.py +++ b/tests/test_state_lifter.py @@ -44,33 +44,48 @@ def test_ravel(): def test_known_constraints(): + raise_error = False for lifter in all_lifters(): - A_known = lifter.get_A_known() - constraints_test_with_tol(lifter, A_known, tol=1e-10) + try: + A_known, b_known = lifter.get_A_known() + except ValueError: + raise_error = True + print(f"Error in {lifter}") + continue + + constraints_test_with_tol(lifter, A_known, b_known, tol=1e-10) B_known = lifter.get_B_known() x = lifter.get_x(theta=lifter.theta) for Bi in B_known: assert x.T @ Bi @ x <= 0 + if raise_error: + raise ValueError("Some lifters could not provide known constraints.") def test_constraint_rank(): + raise_error = False for lifter in all_lifters(): assert isinstance(lifter, StateLifter) try: - A_known = lifter.get_A_known(add_redundant=False) + A_known, b_known = lifter.get_A_known(add_redundant=False) pass except TypeError: - A_known = lifter.get_A_known() - rank = lifter.get_constraint_rank(A_known) + print(f"lifter {lifter} does not support add_redundant!") + A_known, b_known = lifter.get_A_known() + raise_error = True + + rank = lifter.get_constraint_rank(A_known, b_known=b_known) assert rank == len(A_known) + if raise_error: + raise ValueError("Some lifters could not provide known constraints.") def test_vec_mat(): """Make sure that we can go back and forth from vec to mat.""" for lifter in all_lifters(): try: - A_known = lifter.get_A_known() + A_known, b_known = lifter.get_A_known() except AttributeError: print(f"could not get A_known of {lifter}") A_known = [] @@ -96,7 +111,7 @@ def test_vec_mat(): a_poly = lifter.convert_a_to_polyrow(a) a_test = lifter.convert_polyrow_to_a(a_poly) - np.testing.assert_allclose(a, a_test) + np.testing.assert_allclose(np.asarray(a), np.asarray(a_test)) def test_levels(): From eccdcdba059f71935fc7c5bf33bad73c9e195fc7 Mon Sep 17 00:00:00 2001 From: Frederike Duembgen Date: Sat, 23 Aug 2025 12:10:29 +0200 Subject: [PATCH 09/10] Add type hints and docstrings everywhere using new copilot instructions --- .github/copilot-instructions.md | 20 ++ docs/source/conf.py | 3 + popcor/base_lifters/robust_pose_lifter.py | 28 +-- popcor/examples/mono_lifter.py | 132 +++++------- popcor/examples/poly4_lifter.py | 61 +++--- popcor/examples/poly6_lifter.py | 39 +++- popcor/examples/ro_nsq_lifter.py | 109 +++++++--- popcor/examples/ro_sq_lifter.py | 137 +++++++----- popcor/examples/rotation_lifter.py | 247 +++++++--------------- popcor/examples/stereo1d_lifter.py | 98 ++++++--- popcor/examples/stereo2d_lifter.py | 49 +++-- popcor/examples/stereo3d_lifter.py | 74 +++---- popcor/examples/wahba_lifter.py | 105 +++++---- 13 files changed, 575 insertions(+), 527 deletions(-) create mode 100644 .github/copilot-instructions.md diff --git a/.github/copilot-instructions.md b/.github/copilot-instructions.md new file mode 100644 index 0000000..feac66a --- /dev/null +++ b/.github/copilot-instructions.md @@ -0,0 +1,20 @@ +# Copilot Instructions + +YOU MUST UBSOLUTELY FOLLOW THESE INSTRUCTIONS AT ALL TIMES. + +## General guidelines to be used throughout: +- Use spaces, not tabs: 4 spaces per indentation level +- Variables are snake_case, Classes are CamelCase. + +## When I ask you to overhaul a module, follow these instructions: +- Add type hints to all functions and methods. If a variable is initialized to None, make sure to use | None in the type hint. +- ABSOLUTELY do not document standard functions such as __init__, __str__, __repr__, unless there is something special to document. +- ABSOLUTELY do not document property funcitons (those that use @property decorator), unless there is something special to document. +- Never remove inline comments, those starting with #. They are important for understanding the code. You may correct or complete them. +- Add at least one-line docstrings to each function, unless the function starts with an underscore, and unless it is clear from the functon name what it does. +- ABSOLUTELY if the function or class already has a docstring, do not rewrite it, only correct or complete it. +- Fix typos in documentation. Do not make any unnecessary changes to language, only if something is grammatically incorrect. +- Make code more efficient without making it less readable. +- If the module contains mostly one class, just provide one very brief module-level docstring. All the main information should be in the class top-level docstring. +- If this module contains many helper functions, then add a top-level docstring if it does not exist yet. If there is one, make sure that it features all important aspects of the code. Do not add information such as “Requires:, Author:, Date:”, only the technical content of the file. +- Do not introduce dummy variables, like d=self.d inside a function. Always use self.d for better readability. \ No newline at end of file diff --git a/docs/source/conf.py b/docs/source/conf.py index e06c07c..01b37b8 100644 --- a/docs/source/conf.py +++ b/docs/source/conf.py @@ -120,6 +120,9 @@ # Do not show "View Page Source" html_show_sourcelink = False +# Make sure that inheriting classes do not show the docstring of the parent class +autodoc_inherit_docstrings = False + def setup(app): app.add_css_file("custom_css") diff --git a/popcor/base_lifters/robust_pose_lifter.py b/popcor/base_lifters/robust_pose_lifter.py index 02ff1cf..83ba415 100644 --- a/popcor/base_lifters/robust_pose_lifter.py +++ b/popcor/base_lifters/robust_pose_lifter.py @@ -31,6 +31,20 @@ class RobustPoseLifter(StateLifter, ABC): + """RobustPoseLifter is a general class for point-to-point, point-to-line, and point-to-plane registration problems, + with starndard or robust loss functions. + + The goal is to regress an unknown pose based on extrinsic measurements. + + See :class:`~popcor.examples.WahbaLifter` for point-to-point registration, + and :class:`~popcor.examples.MonoLifter`) for point-to-line registration. + + Implemented lifting functions are: + + - xwT: :math:`x \\otimes w` + - xxT: :math:`x \\otimes x` + """ + LEVELS = ["no", "xwT", "xxT"] PARAM_LEVELS = ["no", "p", "ppT"] LEVEL_NAMES = {"no": "no", "xwT": "x kron w", "xxT": "x kron x"} @@ -65,18 +79,6 @@ def __init__( robust=False, beta=BETA, ): - """RobustPoseLifter is a general class for point-to-point, point-to-line, and point-to-plane registration problems, - with starndard or robust loss functions. - - The goal is to regress an unknown pose based on extrinsic measurements. - - See class:`~popcor.examples.WahbaLifter` for point-to-point registration and :class:`~popcor.examples.MonoLifter`) for point-to-line registration. - - Implemented lifting functions are: - - - xwT: :math:`x \\otimes w` - - xxT: :math:`x \\otimes x` - """ self.beta = beta self.n_landmarks = n_landmarks @@ -84,8 +86,6 @@ def __init__( self.level = level if variable_list == "all": variable_list = self.get_all_variables() - # elif variable_list is None: - # self.variable_list = self.VARIABLE_LIST if not robust: assert level == "no" diff --git a/popcor/examples/mono_lifter.py b/popcor/examples/mono_lifter.py index e35b672..a134671 100644 --- a/popcor/examples/mono_lifter.py +++ b/popcor/examples/mono_lifter.py @@ -1,6 +1,7 @@ +"""MonoLifter class for robust pose registration with monocular measurements.""" + from copy import deepcopy -# import autograd.numpy as np import numpy as np from poly_matrix.poly_matrix import PolyMatrix @@ -9,7 +10,6 @@ from popcor.utils.plotting_tools import plot_frame FOV = np.pi / 2 # camera field of view - N_TRYS = 10 # TODO(FD) for some reason this is not required as opposed to what is stated in Heng's paper @@ -20,27 +20,29 @@ class MonoLifter(RobustPoseLifter): + """Robust pose lifter for monocular point-to-line registration. + + Doc under construction. - NOISE = 1e-3 # inlier noise - NOISE_OUT = 0.1 # outlier noise + In the meantime, this example is treated in more detail in `this paper `_, + under the name "PLR". + """ + + NOISE: float = 1e-3 # inlier noise + NOISE_OUT: float = 0.1 # outlier noise def __init__(self, *args, **kwargs): - """This example is treated in more details in `this paper `_, - under the name "PLR" (point-to-line registration). - """ - # only introduced this for the documentation -- otherwise the RobustPoseLifter __init__ is shown. - return super().__init__(*args, **kwargs) + super().__init__(*args, **kwargs) @property - def TIGHTNESS(self): + def TIGHTNESS(self) -> str: return "cost" if self.robust else "rank" - def h_list(self, t, *args, **kwargs): + def h_list(self, t: np.ndarray, *args, **kwargs) -> list: """ - We want to inforce that + Returns constraints h_j(t) <= 0: - norm(t) <= 10 (default) - tan(a/2)*t3 >= sqrt(t1**2 + t2**2) - as constraints h_j(t)<=0 """ default = super().h_list(t, *args, **kwargs) try: @@ -56,31 +58,21 @@ def h_list(self, t, *args, **kwargs): -t[-1], ] - def get_random_position(self, *args, **kwargs): + def get_random_position(self, *args, **kwargs) -> np.ndarray: pc_cw = np.random.rand(self.d) * 0.1 - # make sure all landmarks are in field of view: - # min_dist = max(np.linalg.norm(self.landmarks[:, :self.d-1], axis=1)) pc_cw[self.d - 1] = np.random.uniform(1, self.MAX_DIST) return pc_cw - def get_B_known(self, *args, **kwargs): - """Get inequality constraints of the form x.T @ B @ x <= 0""" - - # TODO(FD) for some reason this is not required as opposed to what is stated in Heng's paper - # and it currently breaks tightness (might be a bug in my implementation though) + def get_B_known(self, *args, **kwargs) -> list: + """Get inequality constraints of the form x.T @ B @ x <= 0.""" if not USE_INEQ: return [] - default = super().get_B_known(*args, **kwargs) - # B2 and B3 enforce that tan(FOV/2)*t3 >= sqrt(t1**2 + t2**2) - # 0 <= - tan**2(FOV/2)*t3**2 + t1**2 + t2**2 B3 = PolyMatrix(symmetric=True) constraint = np.zeros((self.d, self.d)) constraint[range(self.d - 1), range(self.d - 1)] = 1.0 constraint[self.d - 1, self.d - 1] = -np.tan(FOV / 2) ** 2 B3["t", "t"] = constraint - - # t3 >= 0 constraint = np.zeros(self.d) constraint[self.d - 1] = -1 B2 = PolyMatrix(symmetric=True) @@ -90,41 +82,43 @@ def get_B_known(self, *args, **kwargs): B3.get_matrix(self.var_dict), ] - def term_in_norm(self, R, t, pi, ui, *args, **kwargs): + def term_in_norm( + self, + R: np.ndarray, + t: np.ndarray, + pi: np.ndarray, + ui: np.ndarray, + *args, + **kwargs, + ) -> np.ndarray: return R @ pi + t def residual_sq( self, - R, - t, - pi, - ui, + R: np.ndarray, + t: np.ndarray, + pi: np.ndarray, + ui: np.ndarray, *args, **kwargs, - ): + ) -> float: W = np.eye(self.d) - np.outer(ui, ui) term = self.term_in_norm(R, t, pi, ui, *args, **kwargs) - res = term.T @ W @ term + res = float(term.T @ W @ term) if NORMALIZE: res /= (self.n_landmarks * self.d) ** 2 return res - def plot_setup(self, *args, **kwargs): + def plot_setup(self, *args, **kwargs) -> tuple: if self.d != 2: print("Plotting currently only supported for d=2") - return + return None, None import matplotlib.pylab as plt assert self.landmarks is not None - fig, ax = plt.subplots() - - # R, t = get_C_r_from_theta(self.theta, self.d) - # ax.scatter(*t, color="k", label="pose") - ax.axis("equal") t_wc_w, C_cw = plot_frame(ax, self.theta, label="pose", color="gray", d=2) - if self.y_ is not None: for i in range(self.y_.shape[0]): ax.scatter( @@ -133,11 +127,8 @@ def plot_setup(self, *args, **kwargs): color=f"C{i}", label="landmarks", ) - - # this vector is in camera coordinates ui_c = self.y_[i] assert abs(np.linalg.norm(ui_c) - 1.0) < 1e-10 - ax.plot( [t_wc_w[0], self.landmarks[i][0]], [t_wc_w[1], self.landmarks[i][1]], @@ -153,30 +144,22 @@ def plot_setup(self, *args, **kwargs): ) return fig, ax - def simulate_y(self, noise: float | None = None): + def simulate_y(self, noise: float | None = None) -> np.ndarray: if noise is None: noise = self.NOISE - y = np.zeros((self.n_landmarks, self.d)) theta = self.theta[: self.d + self.d**2] outlier_index = self.get_outlier_index() - R, t = get_C_r_from_theta(theta, self.d) for i in range(self.n_landmarks): pi = self.landmarks[i] - # ui = deepcopy(pi) #R @ pi + t ui = R @ pi + t ui /= ui[self.d - 1] - - # random unit vector inside the FOV cone - # tan(a/2)*t3 >= sqrt(t1**2 + t2**2) or t3 >= 1 if np.tan(FOV / 2) * ui[self.d - 1] < np.sqrt( np.sum(ui[: self.d - 1] ** 2) ): print("warning: inlier not in FOV!!") - if i in outlier_index: - # randomly sample a vector success = False for _ in range(N_TRYS): ui_test = deepcopy(ui) @@ -192,13 +175,12 @@ def simulate_y(self, noise: float | None = None): if not success: raise ValueError("did not find valid outlier ui") else: - ui[: self.d - 1] += np.random.normal( # type: ignore + ui[: self.d - 1] += np.random.normal( scale=noise, loc=0, size=self.d - 1 ) assert ui[self.d - 1] == 1.0 ui /= np.linalg.norm(ui) y[i] = ui - self.y_ = y return y @@ -209,73 +191,59 @@ def get_Q( use_cliques: list = [], *args, **kwargs, - ): + ) -> np.ndarray | PolyMatrix: assert self.landmarks is not None, "landmarks must be set before calling get_Q" if noise is None: noise = self.NOISE - if self.y_ is None: self.y_ = self.simulate_y(noise=noise) - Q = self.get_Q_from_y( self.y_, output_poly=output_poly, use_cliques=use_cliques, *args, **kwargs ) return Q def get_Q_from_y( - self, y, output_poly: bool = False, use_cliques: list = [], *args, **kwargs - ): + self, + y: np.ndarray, + output_poly: bool = False, + use_cliques: list = [], + *args, + **kwargs, + ) -> np.ndarray | PolyMatrix: """ - every cost term can be written as - (1 + wi)/b**2 [l x'] Qi [l; x] / norm + 1 - wi - = [l x'] Qi/b**2 [l; x] /norm + wi * [l x']Qi/b**2[l;x] / norm + 1 - wi - - cost term: - (Rpi + t) (I - uiui') (Rpi + t) + Returns the cost matrix Q for the current y. """ assert ( self.landmarks is not None ), "landmarks must be set before calling get_Q_from_y" - Q = PolyMatrix(symmetric=True) if NORMALIZE: norm = (self.n_landmarks * self.d) ** 2 - - if len(use_cliques): - js = use_cliques - else: - js = list(range(self.n_landmarks)) - + js = use_cliques if len(use_cliques) else list(range(self.n_landmarks)) for i in js: pi = self.landmarks[i] ui = y[i] - Pi = np.c_[np.eye(self.d), np.kron(pi, np.eye(self.d))] # I, pi x I + Pi = np.c_[np.eye(self.d), np.kron(pi, np.eye(self.d))] Wi = np.eye(self.d) - np.outer(ui, ui) - Qi = Pi.T @ Wi @ Pi # "t,t, t,c, c,c: Wi, Wi @ kron, kron.T @ Wi @ kron + Qi = Pi.T @ Wi @ Pi if NORMALIZE: Qi /= norm - if self.robust: Qi /= self.beta**2 - # last two terms, should not be affected by norm Q[self.HOM, self.HOM] += 1 Q[self.HOM, f"w_{i}"] += -0.5 if self.level == "xwT": - # Q[f"z_{i}", "x"] += 0.5 * Qi Q[f"z_{i}", "t"] += 0.5 * Qi[:, : self.d] Q[f"z_{i}", "c"] += 0.5 * Qi[:, self.d :] - # Q["x", "x"] += Qi Q["t", "t"] += Qi[: self.d, : self.d] Q["t", "c"] += Qi[: self.d, self.d :] Q["c", "c"] += Qi[self.d :, self.d :] elif self.level == "xxT": Q["z_0", f"w_{i}"] += 0.5 * Qi.flatten()[:, None] - # Q["x", "x"] += Qi Q["t", "t"] += Qi[: self.d, : self.d] Q["t", "c"] += Qi[: self.d, self.d :] Q["c", "c"] += Qi[self.d :, self.d :] else: - # Q["x", "x"] += Qi Q["t", "t"] += Qi[: self.d, : self.d] Q["t", "c"] += Qi[: self.d, self.d :] Q["c", "c"] += Qi[self.d :, self.d :] @@ -284,6 +252,6 @@ def get_Q_from_y( Q_sparse = 0.5 * Q.get_matrix(variables=self.var_dict) return Q_sparse - def __repr__(self, *args, **kwargs): + def __repr__(self, *args, **kwargs) -> str: appendix = "_robust" if self.robust else "" return f"mono_{self.d}d_{self.level}_{self.param_level}{appendix}" diff --git a/popcor/examples/poly4_lifter.py b/popcor/examples/poly4_lifter.py index 754c978..ae4edeb 100644 --- a/popcor/examples/poly4_lifter.py +++ b/popcor/examples/poly4_lifter.py @@ -1,10 +1,12 @@ +"""Poly4Lifter class for fourth-degree polynomial lifter examples.""" + import numpy as np from popcor.base_lifters import PolyLifter class Poly4Lifter(PolyLifter): - """Fourth-degree polynomial examples. + """Fourth-degree polynomial lifter. Two types are provided: @@ -13,40 +15,38 @@ class Poly4Lifter(PolyLifter): """ @property - def VARIABLE_LIST(self): + def VARIABLE_LIST(self) -> list[list[str]]: return [[self.HOM, "t", "z0"]] - def __init__(self, poly_type="A"): + def __init__(self, poly_type: str = "A") -> None: # actual minimum assert poly_type in ["A", "B"] - self.poly_type = poly_type + self.poly_type: str = poly_type super().__init__(degree=4) - def get_Q(self, output_poly=False, noise=None): + def get_Q( + self, output_poly: bool = False, noise: float | None = None + ) -> np.ndarray: + """Returns the Q matrix for the selected polynomial type.""" if output_poly: raise ValueError("output_poly not implemented for Poly4Lifter.") if self.poly_type == "A": - # fmt: off - # noqa - Q = np.r_[ - np.c_[2, 1, 0], - np.c_[1, -1 / 2, -1 / 3], - np.c_[0, -1 / 3, 1 / 4] + # Q matrix for type A + return np.r_[ + np.c_[2, 1, 0], np.c_[1, -1 / 2, -1 / 3], np.c_[0, -1 / 3, 1 / 4] ] - # fmt: on elif self.poly_type == "B": - # below is constructed such that f'(t) = (t-1)*(t-2)*(t-3) - # fmt: off - # noqa - Q = np.r_[ - np.c_[3, -3, 0], - np.c_[-3, 11 / 2, -1], - np.c_[0, -1, 1 / 4] - ] - # fmt: on - return Q - - def get_A_known(self, output_poly=False, add_redundant=False, var_dict=None): + # Q matrix for type B, constructed such that f'(t) = (t-1)*(t-2)*(t-3) + return np.r_[np.c_[3, -3, 0], np.c_[-3, 11 / 2, -1], np.c_[0, -1, 1 / 4]] + else: + raise ValueError(f"Unknown poly_type: {self.poly_type}") + + def get_A_known( + self, + output_poly: bool = False, + add_redundant: bool = False, + var_dict: dict | None = None, + ) -> tuple[list[np.ndarray] | list, list[int]]: from poly_matrix import PolyMatrix if add_redundant: @@ -61,11 +61,10 @@ def get_A_known(self, output_poly=False, add_redundant=False, var_dict=None): else: return [A_1.get_matrix(self.var_dict)], [0] - def generate_random_setup(self): - self.theta_ = np.array([-1]) + def generate_random_setup(self) -> None: + self.theta_: np.ndarray = np.array([-1]) - def get_D(self, that): - """Not currently used.""" + def get_D(self, that: float) -> np.ndarray: D = np.array( [ [1.0, 0.0, 0.0], @@ -82,10 +81,10 @@ def get_D(self, that): import matplotlib.pylab as plt # Get the directory two levels up from this file - base_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..")) + base_dir: str = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..")) - thetas = np.linspace(-2, 3, 100) - poly_lifter = Poly4Lifter(poly_type="A") + thetas: np.ndarray = np.linspace(-2, 3, 100) + poly_lifter: Poly4Lifter = Poly4Lifter(poly_type="A") fig, ax = poly_lifter.plot(thetas) fig.savefig(os.path.join(base_dir, "docs", "figures", "poly4_lifter_A.png")) diff --git a/popcor/examples/poly6_lifter.py b/popcor/examples/poly6_lifter.py index 8243b0f..265b47c 100644 --- a/popcor/examples/poly6_lifter.py +++ b/popcor/examples/poly6_lifter.py @@ -1,4 +1,8 @@ +"""Poly6Lifter class for sixth-degree polynomial examples.""" + import numpy as np +import scipy.sparse as sp +from poly_matrix import PolyMatrix from popcor.base_lifters import PolyLifter @@ -8,22 +12,25 @@ class Poly6Lifter(PolyLifter): Two types are provided: - - poly_type="A": one global minimum, two local minima, 2 local maxima + - poly_type="A": one global minimum, two local minima, two local maxima - poly_type="B": one global minimum, one local minimum, one local maximum """ @property - def VARIABLE_LIST(self): + def VARIABLE_LIST(self) -> list[list[str]]: return [[self.HOM, "t", "z0", "z1"]] - def __init__(self, poly_type="A"): + def __init__(self, poly_type: str = "A") -> None: assert poly_type in ["A", "B"] - self.poly_type = poly_type + self.poly_type: str = poly_type super().__init__(degree=6) - def get_Q(self, output_poly=False, noise=None): + def get_Q( + self, output_poly: bool = False, noise: float | None = None + ) -> np.ndarray: + """Returns the Q matrix for the selected polynomial type.""" if output_poly: - raise ValueError("output_poly not implemented for Poly4Lifter.") + raise ValueError("output_poly not implemented for Poly6Lifter.") if self.poly_type == "A": return 0.1 * np.array( [ @@ -42,11 +49,19 @@ def get_Q(self, output_poly=False, noise=None): [0, 0.2685, -0.0667, 0.0389], ] ) - - def get_A_known(self, output_poly=False, add_redundant=True, var_dict=None): + else: + raise ValueError(f"Unknown poly_type: {self.poly_type}") + + def get_A_known( + self, + output_poly: bool = False, + add_redundant: bool = True, + var_dict: dict | None = None, + ) -> tuple[list, list[float]]: + """Returns the list of known A matrices and their corresponding values.""" from poly_matrix import PolyMatrix - A_list = [] + A_list: list = [] # z_0 = t^2 A_1 = PolyMatrix(symmetric=True) @@ -74,7 +89,8 @@ def get_A_known(self, output_poly=False, add_redundant=True, var_dict=None): A_list ) - def get_D(self, that): + def get_D(self, that: float) -> np.ndarray: + """Returns the D matrix for the given value.""" D = np.array( [ [1.0, 0.0, 0.0, 0.0], @@ -85,7 +101,8 @@ def get_D(self, that): ) return D - def generate_random_setup(self): + def generate_random_setup(self) -> None: + """Initializes a random setup for theta_.""" self.theta_ = np.array([-1]) diff --git a/popcor/examples/ro_nsq_lifter.py b/popcor/examples/ro_nsq_lifter.py index e0eb54c..9daa0d7 100644 --- a/popcor/examples/ro_nsq_lifter.py +++ b/popcor/examples/ro_nsq_lifter.py @@ -1,3 +1,5 @@ +"""RangeOnlyNonSqLifter class for localization in 2D or 3D using non-squared lifting.""" + import numpy as np import scipy.sparse as sp from poly_matrix import PolyMatrix @@ -33,18 +35,20 @@ class RangeOnlyNsqLifter(RangeOnlyLifter): - :math:`\\theta` is now the flattened vector of positions :math:`p_n` and also normal vectors :math:`z_{nk}`. """ - TIGHTNESS = "rank" - LEVELS = ["normals", "simple"] - LEVEL_NAMES = { - "normals": "$\\boldymbol{y}_n$", + TIGHTNESS: str = "rank" + LEVELS: list[str] = ["normals", "simple"] + LEVEL_NAMES: dict[str, str] = { + "normals": "$\\boldsymbol{y}_n$", "simple": "$z_n$", } - MONOMIAL_DEGREE = 1 - - SCALE = 1.0 + MONOMIAL_DEGREE: int = 1 + SCALE: float = 1.0 @staticmethod - def create_good(n_positions, n_landmarks, d=2): + def create_good( + n_positions: int, n_landmarks: int, d: int = 2 + ) -> "RangeOnlyNsqLifter": + """Create a lifter with good initial positions.""" landmarks, theta = RangeOnlyLifter.create_good(n_positions, n_landmarks, d) lifter = RangeOnlyNsqLifter(n_positions, n_landmarks, d) lifter.overwrite_theta(theta) @@ -52,7 +56,10 @@ def create_good(n_positions, n_landmarks, d=2): return lifter @staticmethod - def create_bad(n_positions, n_landmarks, d=2): + def create_bad( + n_positions: int, n_landmarks: int, d: int = 2 + ) -> "RangeOnlyNsqLifter": + """Create a lifter with bad initial positions.""" landmarks, theta = RangeOnlyLifter.create_bad(n_positions, n_landmarks, d) lifter = RangeOnlyNsqLifter(n_positions, n_landmarks, d) lifter.overwrite_theta(theta) @@ -60,7 +67,10 @@ def create_bad(n_positions, n_landmarks, d=2): return lifter @staticmethod - def create_bad_fixed(n_positions, n_landmarks, d=2): + def create_bad_fixed( + n_positions: int, n_landmarks: int, d: int = 2 + ) -> "RangeOnlyNsqLifter": + """Create a lifter with fixed bad initial positions.""" landmarks, theta = RangeOnlyLifter.create_bad_fixed(n_positions, n_landmarks, d) lifter = RangeOnlyNsqLifter(n_positions, n_landmarks, d) lifter.overwrite_theta(theta) @@ -69,13 +79,13 @@ def create_bad_fixed(n_positions, n_landmarks, d=2): def __init__( self, - n_positions, - n_landmarks, - d, - W=None, - level="normals", - variable_list=None, - param_level="no", + n_positions: int, + n_landmarks: int, + d: int, + W: np.ndarray | None = None, + level: str = "normals", + variable_list: list | None = None, + param_level: str = "no", ): if level == "simple": raise NotImplementedError("simple is not implemented yet.") @@ -90,13 +100,14 @@ def __init__( ) @property - def VARIABLE_LIST(self): + def VARIABLE_LIST(self) -> list[list[str]]: return [ [self.HOM, "x_0"], [self.HOM, "x_0"] + [f"z_0_{i}" for i in range(self.n_landmarks)], ] - def get_all_variables(self): + def get_all_variables(self) -> list[list[str]]: + """Return all variables used in the lifting.""" vars = [self.HOM] vars += [f"x_{i}" for i in range(self.n_positions)] vars += [ @@ -106,7 +117,13 @@ def get_all_variables(self): ] return [vars] - def get_A_known(self, var_dict=None, output_poly=False, add_redundant=False): + def get_A_known( + self, + var_dict: dict | None = None, + output_poly: bool = False, + add_redundant: bool = False, + ) -> list[np.ndarray | PolyMatrix]: + """Return known quadratic constraints for the lifted variables.""" if var_dict is None: var_dict = self.var_dict @@ -124,20 +141,31 @@ def get_A_known(self, var_dict=None, output_poly=False, add_redundant=False): A_list.append(A.get_matrix(self.var_dict)) elif self.level == "simple": - # enfore that y_{nk}^2 = ||p_n - a_k||^2 = ||p_n||^2 - 2a_k^T p_n + ||a_k||^2 + # enforce that y_{nk}^2 = ||p_n - a_k||^2 = ||p_n||^2 - 2a_k^T p_n + ||a_k||^2 raise NotImplementedError( "get_A_known not implemented yet for simple level" ) - return A_list, [0.0] * len(A_list) + return A_list - def get_residuals(self, t, y, ad=False): + def get_residuals( + self, t: np.ndarray, y: np.ndarray, ad: bool = False + ) -> np.ndarray: + """Compute residuals for the cost function.""" return super().get_residuals(t, y, ad=ad, squared=False) - def get_cost(self, theta, y, sub_idx=None, ad=False): + def get_cost( + self, + theta: np.ndarray, + y: np.ndarray, + sub_idx: int | None = None, + ad: bool = False, + ) -> float: + """Compute the cost from theta and y.""" residuals = self.get_residuals(theta, y, ad=ad) return self.get_cost_from_res(residuals, sub_idx, ad=ad) - def get_normals(self, theta=None): + def get_normals(self, theta: np.ndarray | None = None) -> np.ndarray: + """Compute normal vectors for each position-landmark pair.""" if theta is None: theta = self.theta @@ -148,7 +176,13 @@ def get_normals(self, theta=None): self.landmarks[None, :, :] - theta, axis=2 )[:, :, None] - def get_x(self, theta=None, parameters=None, var_subset=None): + def get_x( + self, + theta: np.ndarray | None = None, + parameters: np.ndarray | dict | None = None, + var_subset: list[str] | dict | None = None, + ) -> np.ndarray: + """Return the lifted variable vector x for the given theta.""" if var_subset is None: var_subset = self.var_dict if theta is None: @@ -175,7 +209,8 @@ def get_x(self, theta=None, parameters=None, var_subset=None): assert len(x_data) == self.get_dim_x(var_subset) return np.array(x_data) - def get_J_lifting(self, t): + def get_J_lifting(self, t: np.ndarray) -> sp.csr_array: + """Return the Jacobian of the lifting (currently not implemented).""" pos = t.reshape((-1, self.d)) ii = [] jj = [] @@ -192,7 +227,10 @@ def get_J_lifting(self, t): ) return J_lifting - def get_Q_from_y(self, y, output_poly: bool = False): + def get_Q_from_y( + self, y: np.ndarray, output_poly: bool = False + ) -> np.ndarray | PolyMatrix | sp.csr_matrix | sp.csc_matrix: + """Return the quadratic cost matrix Q for the given y.""" Q = PolyMatrix() for k in range(self.n_landmarks): for i in range(self.n_positions): @@ -214,11 +252,14 @@ def get_Q_from_y(self, y, output_poly: bool = False): else: return Q.get_matrix(self.var_dict) - def simulate_y(self, noise: float | None = None, squared: bool = True): + def simulate_y( + self, noise: float | None = None, squared: bool = True + ) -> np.ndarray: + """Simulate measurements y (distances) with optional noise.""" return super().simulate_y(noise=noise, squared=False) @property - def var_dict(self): + def var_dict(self) -> dict[str, int]: var_dict = {self.HOM: 1} var_dict.update({f"x_{n}": self.d for n in range(self.n_positions)}) var_dict.update( @@ -239,10 +280,8 @@ def size_z(self) -> int: else: raise ValueError(f"Unknown level {self.level}") - def get_valid_samples( - self, - n_samples, - ): + def get_valid_samples(self, n_samples: int) -> np.ndarray: + """Return valid samples for the lifted variables.""" samples = super().get_valid_samples(n_samples) # TODO(FD): maybe this should be moved to theta. @@ -250,7 +289,7 @@ def get_valid_samples( normals /= np.linalg.norm(normals, axis=2)[:, :, None] return np.hstack([samples, normals.reshape(normals.shape[0], -1)]) - def __repr__(self): + def __repr__(self) -> str: return f"rangeonlyloc{self.d}d_{self.level}" diff --git a/popcor/examples/ro_sq_lifter.py b/popcor/examples/ro_sq_lifter.py index 1f352c1..4cb49ab 100644 --- a/popcor/examples/ro_sq_lifter.py +++ b/popcor/examples/ro_sq_lifter.py @@ -1,39 +1,41 @@ +"""RangeOnlySqLifter class for localization in 2D or 3D using squared lifting.""" + import itertools +from typing import Any, Dict, List, Tuple import numpy as np import scipy.sparse as sp +from poly_matrix import PolyMatrix from poly_matrix.least_squares_problem import LeastSquaresProblem from popcor.base_lifters import RangeOnlyLifter from popcor.utils.common import diag_indices, upper_triangular NOISE = 1e-2 # std deviation of distance noise - NORMALIZE = True class RangeOnlySqLifter(RangeOnlyLifter): """Range-only localization in 2D or 3D. - We minimize the cost function + Minimizes the cost function .. math:: f(\\theta) = \\sum_{n=0}^{N-1} \\sum_{k=0}^{K-1} w_{nk} (d_{nk}^2 - ||p_n - a_k||^2)^2 where - - :math:`w_{nk}` is the weight for the nth point and kth landmark (currently assumed binary to mark missing edges). + - :math:`w_{nk}` is the weight for the nth point and kth landmark (binary for missing edges). - :math:`\\theta` is the flattened vector of positions :math:`p_n`. - :math:`d_{nk}` is the distance measurement from point n to landmark k. - :math:`a_k` is the kth landmark. - Note that in the current implementation, there is no regularization term so the problem could be split into individual points. - - We experiment with two different substitutions to turn the cost function into aquadratic form: + No regularization term is present, so the problem could be split into individual points. - - level "no" uses substitution :math:`z_i=||p_i||^2=x_i^2 + y_i^2` (or equivalent 3D version). - - level "quad" uses substitution :math:`y_i=[x_i^2, x_iy_i, y_i^2]` (or equivalent 3D version). + Two substitutions for quadratic form: + - level "no": :math:`z_i=||p_i||^2=x_i^2 + y_i^2` (or 3D version). + - level "quad": :math:`y_i=[x_i^2, x_iy_i, y_i^2]` (or 3D version). - This example is treated in more details in `this paper `_. + See https://arxiv.org/abs/2308.05783 for details. """ TIGHTNESS = "rank" @@ -45,7 +47,9 @@ class RangeOnlySqLifter(RangeOnlyLifter): MONOMIAL_DEGREE = 2 @staticmethod - def create_good(n_positions, n_landmarks, d=2): + def create_good( + n_positions: int, n_landmarks: int, d: int = 2 + ) -> "RangeOnlySqLifter": landmarks, theta = RangeOnlyLifter.create_good(n_positions, n_landmarks, d) lifter = RangeOnlySqLifter(n_positions, n_landmarks, d) lifter.overwrite_theta(theta) @@ -53,7 +57,9 @@ def create_good(n_positions, n_landmarks, d=2): return lifter @staticmethod - def create_bad(n_positions, n_landmarks, d=2): + def create_bad( + n_positions: int, n_landmarks: int, d: int = 2 + ) -> "RangeOnlySqLifter": landmarks, theta = RangeOnlyLifter.create_bad(n_positions, n_landmarks, d) lifter = RangeOnlySqLifter(n_positions, n_landmarks, d) lifter.overwrite_theta(theta) @@ -61,7 +67,9 @@ def create_bad(n_positions, n_landmarks, d=2): return lifter @staticmethod - def create_bad_fixed(n_positions, n_landmarks, d=2): + def create_bad_fixed( + n_positions: int, n_landmarks: int, d: int = 2 + ) -> "RangeOnlySqLifter": landmarks, theta = RangeOnlyLifter.create_bad_fixed(n_positions, n_landmarks, d) lifter = RangeOnlySqLifter(n_positions, n_landmarks, d) lifter.overwrite_theta(theta) @@ -70,14 +78,14 @@ def create_bad_fixed(n_positions, n_landmarks, d=2): def __init__( self, - n_positions, - n_landmarks, - d, - W=None, - level="no", - variable_list=None, - param_level="no", - ): + n_positions: int, + n_landmarks: int, + d: int, + W: np.ndarray | None = None, + level: str = "no", + variable_list: list | None = None, + param_level: str = "no", + ) -> None: super().__init__( n_positions=n_positions, n_landmarks=n_landmarks, @@ -89,26 +97,31 @@ def __init__( ) @property - def VARIABLE_LIST(self): + def VARIABLE_LIST(self) -> list: return [ [self.HOM, "x_0"], [self.HOM, "x_0", "z_0"], ] - def get_all_variables(self): - vars = [self.HOM] + def get_all_variables(self) -> list: + vars: List[str] = [self.HOM] vars += [f"x_{i}" for i in range(self.n_positions)] vars += [f"z_{i}" for i in range(self.n_positions)] return [vars] - def get_A_known(self, var_dict=None, output_poly=False, add_redundant=False): + def get_A_known( + self, + var_dict: dict | None = None, + output_poly: bool = False, + add_redundant: bool = False, + ) -> Tuple[list, list]: from poly_matrix.poly_matrix import PolyMatrix if var_dict is None: var_dict = self.var_dict positions = self.get_variable_indices(var_dict) - A_list = [] + A_list: List[Any] = [] for n in positions: if self.level == "no": A = PolyMatrix(symmetric=True) @@ -118,7 +131,6 @@ def get_A_known(self, var_dict=None, output_poly=False, add_redundant=False): A_list.append(A) else: A_list.append(A.get_matrix(self.var_dict)) - elif self.level == "quad": count = 0 for i in range(self.d): @@ -141,14 +153,27 @@ def get_A_known(self, var_dict=None, output_poly=False, add_redundant=False): A_list.append(A.get_matrix(self.var_dict)) return A_list, [0.0] * len(A_list) - def get_residuals(self, t, y, ad=False): + def get_residuals( + self, t: np.ndarray, y: np.ndarray, ad: bool = False + ) -> np.ndarray: return super().get_residuals(t, y, ad=ad, squared=True) - def get_cost(self, theta, y, sub_idx=None, ad=False): + def get_cost( + self, + theta: np.ndarray, + y: np.ndarray, + sub_idx: list | None = None, + ad: bool = False, + ) -> float: residuals = self.get_residuals(theta, y, ad=ad) return self.get_cost_from_res(residuals, sub_idx, ad=ad) - def get_x(self, theta=None, parameters=None, var_subset=None): + def get_x( + self, + theta: np.ndarray | None = None, + parameters: np.ndarray | dict | None = None, + var_subset: dict | None = None, + ) -> np.ndarray: if var_subset is None: var_subset = self.var_dict if theta is None: @@ -158,7 +183,7 @@ def get_x(self, theta=None, parameters=None, var_subset=None): positions = theta.reshape(self.n_positions, -1) - x_data = [] + x_data: List[float] = [] for key in var_subset: if key == self.HOM: x_data.append(1.0) @@ -168,17 +193,17 @@ def get_x(self, theta=None, parameters=None, var_subset=None): elif "z" in key: n = int(key.split("_")[-1]) if self.level == "no": - x_data.append(np.linalg.norm(positions[n]) ** 2) + x_data.append(float(np.linalg.norm(positions[n]) ** 2)) elif self.level == "quad": x_data += list(upper_triangular(positions[n])) assert len(x_data) == self.get_dim_x(var_subset) return np.array(x_data) - def get_J_lifting(self, t): + def get_J_lifting(self, t: np.ndarray) -> sp.csr_array: pos = t.reshape((-1, self.d)) - ii = [] - jj = [] - data = [] + ii: List[int] = [] + jj: List[int] = [] + data: List[float] = [] idx = 0 for n in range(self.n_positions): @@ -187,7 +212,6 @@ def get_J_lifting(self, t): jj += list(range(n * self.d, (n + 1) * self.d)) data += list(2 * pos[n]) elif self.level == "quad": - # it seemed easier to do this manually that programtically if self.d == 3: x, y, z = pos[n] jj += [n * self.d + j for j in [0, 0, 1, 0, 2, 1, 1, 2, 2]] @@ -205,9 +229,9 @@ def get_J_lifting(self, t): ) return J_lifting - def get_hess_lifting(self, t): - """return list of the hessians of the M lifting functions.""" - hessians = [] + def get_hess_lifting(self, t: np.ndarray) -> List[sp.csr_array]: + """Return list of the Hessians of the M lifting functions.""" + hessians: List[sp.csr_array] = [] for n in range(self.n_positions): idx = range(n * self.d, (n + 1) * self.d) if self.level == "no": @@ -227,7 +251,7 @@ def get_hess_lifting(self, t): return hessians @property - def fixed_hessian_list(self): + def fixed_hessian_list(self) -> List[np.ndarray]: if self.d == 2: return [ np.array([[2, 0], [0, 0]]), @@ -246,8 +270,12 @@ def fixed_hessian_list(self): else: raise ValueError(f"Unsupported dimension {self.d} for fixed hessians.") - def get_Q_from_y(self, y, output_poly: bool = False): + def get_Q_from_y( + self, y: np.ndarray, output_poly: bool = False + ) -> np.ndarray | sp.csr_matrix | sp.csc_matrix | PolyMatrix: """ + Returns the quadratic cost matrix Q from distance measurements. + :param y: the distance measurements, shape (n_positions, n_landmarks). IMPORTANT: these are not squared! """ self.ls_problem = LeastSquaresProblem() @@ -283,14 +311,15 @@ def get_Q_from_y(self, y, output_poly: bool = False): return Q / np.sum(self.W > 0) return Q - def simulate_y(self, noise: float | None = None, squared: bool = True): - # purposefully not using squared=True here: - # the noise should always be added to the non-squared distances. + def simulate_y( + self, noise: float | None = None, squared: bool = True + ) -> np.ndarray: + # The noise should always be added to the non-squared distances. return super().simulate_y(noise=noise, squared=False) @property - def var_dict(self): - var_dict = {self.HOM: 1} + def var_dict(self) -> Dict[str, int]: + var_dict: Dict[str, int] = {self.HOM: 1} var_dict.update({f"x_{n}": self.d for n in range(self.n_positions)}) var_dict.update({f"z_{n}": self.size_z for n in range(self.n_positions)}) return var_dict @@ -304,12 +333,17 @@ def size_z(self) -> int: else: raise ValueError(f"Unknown level {self.level}") - def __repr__(self): + def __repr__(self) -> str: return f"rangeonlyloc{self.d}d_{self.level}" # ============ below are currently not used anymore, but it is an elegant way to compute the ============= # gradient and hessian when there are no constraints - def get_grad(self, t, y, sub_idx=None): + def get_grad( + self, + t: np.ndarray, + y: np.ndarray, + sub_idx: list | None = None, + ) -> np.ndarray: J = self.get_J(t, y) x = self.get_x(t) Q = self.get_Q_from_y(y) @@ -319,16 +353,17 @@ def get_grad(self, t, y, sub_idx=None): sub_idx_x = self.get_sub_idx_x(sub_idx) return 2 * J.T[:, sub_idx_x] @ Q[sub_idx_x, :][:, sub_idx_x] @ x[sub_idx_x] # type: ignore - def get_J(self, t, y): + def get_J(self, t: np.ndarray, y: np.ndarray) -> sp.csr_array: J = sp.csr_array( (np.ones(self.N), (range(1, self.N + 1), range(self.N))), shape=(self.N + 1, self.N), ) J_lift = self.get_J_lifting(t) J = sp.vstack([J, J_lift]) + assert isinstance(J, sp.csr_array) return J - def get_hess(self, t, y): + def get_hess(self, t: np.ndarray, y: np.ndarray) -> np.ndarray: x = self.get_x(t) Q = self.get_Q_from_y(y) J = self.get_J(t, y) @@ -343,7 +378,7 @@ def get_hess(self, t, y): hess += 2 * factor * h return hess - def get_D(self, that): + def get_D(self, that: np.ndarray) -> sp.csc_array: D = np.eye(1 + self.n_positions * self.d + self.size_z) x = self.get_x(theta=that) J = self.get_J_lifting(t=that) diff --git a/popcor/examples/rotation_lifter.py b/popcor/examples/rotation_lifter.py index 15393ee..6a7cf97 100644 --- a/popcor/examples/rotation_lifter.py +++ b/popcor/examples/rotation_lifter.py @@ -1,31 +1,32 @@ -""" -A class for solving rotation averaging problems. -""" +"""RotationAveraging class for rotation averaging and synchronization problems.""" -from typing import List +from typing import Any, Dict, List, Tuple import numpy as np +import scipy.sparse as sp from poly_matrix.poly_matrix import PolyMatrix from scipy.spatial.transform import Rotation from popcor.base_lifters import StateLifter -METHOD = "CG" -SOLVER_KWARGS = dict( +METHOD: str = "CG" +SOLVER_KWARGS: dict = dict( min_gradient_norm=1e-7, max_iterations=10000, min_step_size=1e-8, verbosity=1 ) -DEBUG = True +DEBUG: bool = True -def get_orthogonal_constraints(key, hom, d, level): +def get_orthogonal_constraints( + key: Any, hom: Any, d: int, level: str +) -> Tuple[List[PolyMatrix], List[float]]: """Return A, b lists enforcing orthogonality constraints for variable `key`. Produces constraints encoding R'R = I (diagonal == 1 and off-diagonal == 0) for either the rank-1 ("no") or Burer-Monteiro ("bm") formulation. """ - A_list = [] - b_list = [] + A_list: List[PolyMatrix] = [] + b_list: List[float] = [] for i in range(d): # enforce diagonal == 1 for R'R = I Ei = np.zeros((d, d)) @@ -64,71 +65,35 @@ def get_orthogonal_constraints(key, hom, d, level): class RotationLifter(StateLifter): """Rotation averaging problem lifter. - We solve the following optimization problem: - - .. math:: - f(\\theta) = \\min_{R_0, R_1, \\ldots, R_N \\in \\mathrm{SO}^d} - \\sum_{i,j \\in \\mathcal{E}} || R_i - R_j \\tilde{R}_{ij} ||_F^2 - + \\sum_{i=\\in\\mathcal{A}} || R_i - \\tilde{R}_i ||_F^2 - - where :math:`\\tilde{R}_{ij}` are the relative measurements, :math:`\\tilde{R}_{i}` are the absolute measurements, - and the unknowns are - - .. math:: - \\theta = \\begin{bmatrix} R_1 & R_2 & \\ldots & R_N \\end{bmatrix} - - We can alternatively replace the absolute-measurement terms by - - .. math:: - || R_i - R_w \\tilde{R}_{i} ||_F^2 - - where :math:`R_w` is an arbitrary world frame that we can also optimize over, transforming the solutions - by :math:`R_w^{-1}R_i` afterwards to move the world frame to the origin. Using this formulation, all - measurements are binary factors, which may simplify implementation. - - We consider two different formulations of the problem: - - - level "no" corresponds to the rank-1 version: - - .. math:: - x = \\begin{bmatrix} 1, \\mathrm{vec}(R_1), \\ldots, \\mathrm{vec}(R_N) \\end{bmatrix}^T - - - level "bm" corresponds to the rank-d version (bm=Burer-Monteiro). - - .. math:: - X = \\begin{bmatrix} R_1^\\top \\\\ \\vdots \\\\ R_N^\\top \\end{bmatrix} - - According to the above conventions, HOM is either 1 or R_0, the world frame. + See class-level docstring for details on problem formulation and lifted representations. """ - LEVELS = ["no", "bm"] - HOM = "h" - VARIABLE_LIST = [["h", "c_0"], ["h", "c_0", "c_1"]] + LEVELS: List[str] = ["no", "bm"] + HOM: str = "h" + VARIABLE_LIST: List[List[str]] = [["h", "c_0"], ["h", "c_0", "c_1"]] - # whether or not to include the determinant constraints in the known constraints. - ADD_DETERMINANT = False - NOISE = 1e-3 + ADD_DETERMINANT: bool = False + NOISE: float = 1e-3 def __init__( self, - level="no", - param_level="no", - d=2, - n_abs=2, - n_rot=1, - n_rel=1, - sparsity="chain", - ): - """Initialize a RotationLifter instance with problem dimensions and options.""" + level: str = "no", + param_level: str = "no", + d: int = 2, + n_abs: int = 2, + n_rot: int = 1, + n_rel: int = 1, + sparsity: str = "chain", + ) -> None: assert n_rel in [ 0, 1, ], "do not support more than 1 relative measurement per pair currently." - self.n_rot = n_rot - self.n_abs = n_abs - self.n_rel = n_rel - self.level = level - self.sparsity = sparsity + self.n_rot: int = n_rot + self.n_abs: int = n_abs + self.n_rel: int = n_rel + self.level: str = level + self.sparsity: str = sparsity super().__init__( level=level, param_level=param_level, @@ -136,19 +101,17 @@ def __init__( ) @property - def var_dict(self): + def var_dict(self) -> Dict[str, int]: """Return dictionary mapping variable names to their dimensions in the lifted representation.""" if self.level == "no": - # self.HOM is the scalar homogenization variable var_dict = {self.HOM: 1} - var_dict.update({f"c_{i}": self.d**2 for i in range(self.n_rot)}) + var_dict.update({f"c_{i}": int(self.d**2) for i in range(self.n_rot)}) else: - # self.HOM is the world frame, which is d x d - var_dict = {self.HOM: self.d} - var_dict.update({f"c_{i}": self.d for i in range(self.n_rot)}) + var_dict = {self.HOM: int(self.d)} + var_dict.update({f"c_{i}": int(self.d) for i in range(self.n_rot)}) return var_dict - def sample_theta(self): + def sample_theta(self) -> np.ndarray: """Generate a random feasible set of rotations theta shaped (d, n_rot * d).""" C = np.empty((self.d, self.n_rot * self.d)) for i in range(self.n_rot): @@ -161,7 +124,12 @@ def sample_theta(self): C[:, i * self.d : (i + 1) * self.d] = Rotation.random().as_matrix() return C - def get_x(self, theta=None, parameters=None, var_subset=None) -> np.ndarray: + def get_x( + self, + theta: np.ndarray | None = None, + parameters: Any = None, + var_subset: Any = None, + ) -> np.ndarray: """Return the lifted variable x (vector or stacked matrices) for given theta and var_subset.""" if theta is None: theta = self.theta @@ -170,14 +138,13 @@ def get_x(self, theta=None, parameters=None, var_subset=None) -> np.ndarray: if var_subset is None: var_subset = self.var_dict.keys() - x_data = [] + x_data: list = [] if self.level == "no": for key in var_subset: if key == self.HOM: x_data.append(1.0) elif "c" in key: i = int(key.split("_")[1]) - # column-major flattening to match vec(R) x_data += list(theta[:, i * self.d : (i + 1) * self.d].flatten("F")) else: raise ValueError(f"untreated key {key}") @@ -206,32 +173,19 @@ def get_theta(self, x: np.ndarray) -> np.ndarray: C_flat = x[1 : 1 + self.n_rot * self.d**2] return C_flat.reshape((self.d, self.n_rot * self.d), order="F") elif self.level == "bm": - # Remember that x is composed of R_0.T, R_1.T, ..., R_N.T - # We want to return theta in the form [R*_1, R*_2, ..., R*_N] - # where each R*_i := R_0.T @ R_i - - # d x d (original world frame) R0 = x[: self.d, : self.d].T - - # nd x d: [R_1.T; R_2.T; ...; R_N.T] Ri = np.array(x[self.d : (self.n_rot + 1) * self.d, : self.d]).T - - # return d x nd: [R*_1, R*_2, ..., R*_N] Ri_world = R0.T @ Ri return Ri_world else: raise ValueError(f"Unknown level {self.level} for RotationLifter") - def add_relative_measurement(self, i, j, noise) -> np.ndarray: - """Create a noisy relative measurement R_ij = R_i.T @ R_j with additive rotation noise. - error terms: || R_j - R_i @ R_ij ||_F^2 - measurements: R_ij = R_i.T @ R_j - """ + def add_relative_measurement(self, i: int, j: int, noise: float) -> np.ndarray: + """Create a noisy relative measurement R_ij = R_i.T @ R_j with additive rotation noise.""" R_i = self.theta[:, i * self.d : (i + 1) * self.d] R_j = self.theta[:, j * self.d : (j + 1) * self.d] R_gt = R_i.T @ R_j - # Generate a random small rotation as noise and apply it if noise > 0: noise_rotvec = np.random.normal(scale=noise, size=(self.d,)) Rnoise = ( @@ -243,17 +197,14 @@ def add_relative_measurement(self, i, j, noise) -> np.ndarray: else: return R_gt - def add_absolute_measurement(self, i, noise, n_meas=1) -> List[np.ndarray]: - """Create one or more noisy absolute measurements of rotation R_i (relative to world). - error terms: || R_i - R_w @ R_wi ||_F^2 - measurements: R_wi = R_w.T @ R_i - """ + def add_absolute_measurement( + self, i: int, noise: float, n_meas: int = 1 + ) -> List[np.ndarray]: + """Create one or more noisy absolute measurements of rotation R_i (relative to world).""" R_gt = self.theta[:, i * self.d : (i + 1) * self.d] - y = [] + y: List[np.ndarray] = [] for _ in range(n_meas): - # noise model: R_i = R.T @ Rnoise if noise > 0: - # Generate a random small rotation as noise and apply it noise_rotvec = np.random.normal(scale=noise, size=(self.d,)) Rnoise = ( Rotation.from_rotvec(noise_rotvec).as_matrix() @@ -266,17 +217,16 @@ def add_absolute_measurement(self, i, noise, n_meas=1) -> List[np.ndarray]: y.append(Ri) return y - def simulate_y(self, noise: float | None = None) -> dict: + def simulate_y(self, noise: float | None = None) -> Dict[Any, Any]: """Simulate measurement dictionary y given current theta and noise level.""" if noise is None: noise = self.NOISE - y = {} + y: Dict[Any, Any] = {} if self.n_abs > 0: for i in range(self.n_rot): y[i] = self.add_absolute_measurement(i, noise, self.n_abs) else: - # add prior if no absolute measurements exist y[0] = self.add_absolute_measurement(0, 0.0, 1) if self.n_rel > 0: @@ -288,7 +238,9 @@ def simulate_y(self, noise: float | None = None) -> dict: raise ValueError(f"Unknown sparsity {self.sparsity}") return y - def get_Q(self, noise: float | None = None, output_poly: bool = False): + def get_Q( + self, noise: float | None = None, output_poly: bool = False + ) -> PolyMatrix | np.ndarray | sp.csr_matrix | sp.csc_matrix: """Return the cost matrix Q (poly or ndarray) constructed from simulated measurements.""" if noise is None: noise = self.NOISE @@ -297,23 +249,19 @@ def get_Q(self, noise: float | None = None, output_poly: bool = False): return self.get_Q_from_y(self.y_, output_poly=output_poly) - def get_Q_from_y(self, y, output_poly=False): + def get_Q_from_y( + self, y: Dict[Any, Any], output_poly: bool = False + ) -> PolyMatrix | np.ndarray | sp.csr_matrix | sp.csc_matrix: """Construct the quadratic cost (PolyMatrix or ndarray) from measurement dictionary y.""" Q = PolyMatrix() for key, R in y.items(): if isinstance(key, int): - # loop over all absolute measurements of this rotation. assert isinstance(R, list) for Rk in R: if self.level == "no": - # unary factors: f(R) = sum_k || R - Rk ||_F^2 = sum_k 2tr(I) - 2tr(R'Rk) - # Rk is measured rotation Q_test = PolyMatrix() Q_test[self.HOM, f"c_{key}"] -= Rk.flatten("F")[None, :] - # Not adding below to be consistent with "bm" case, where we cannot - # add a constant term to the cost. - # Q_test[self.HOM, self.HOM] += 2 * self.d if DEBUG: x = self.get_x() cost_x = x.T @ Q_test.get_matrix(self.var_dict) @ x @@ -323,9 +271,6 @@ def get_Q_from_y(self, y, output_poly=False): Q += Q_test elif self.level == "bm": - # use HOM as a world frame: f(R) = sum_i || R - HOM @ Rk || - # compare with below: - # R corresponds to Rj, HOM corresponds to Ri Q_test = PolyMatrix() Q_test[self.HOM, f"c_{key}"] -= Rk if DEBUG: @@ -337,20 +282,10 @@ def get_Q_from_y(self, y, output_poly=False): cost_R = np.linalg.norm(Ri - Rk) ** 2 - 2 * self.d assert abs(cost_x - cost_R) < 1e-10 Q += Q_test - # treat binary factors - # f(R) = sum_ij || Rj - Ri @ Rij ||_F^2 - # = sum_ij tr((Rj - Ri @ Rij)'(Rj - Ri @ Rij)) - # = sum_ij tr(Rj'Rj) - 2 tr(Rj'Ri Rij) + tr(Rij'Ri'RiRij) - # = sum_ij 2tr(I) - 2tr(Rij Rj' Ri) - # = sum_ij 2tr(I) - 2c_i' (Rij kron I) c_j elif isinstance(key, tuple): i, j = key if self.level == "no": Q_test = PolyMatrix() - - # Not adding below to be consistent with "bm" case, where we cannot - # add a constant term to the cost. - # Q_test[self.HOM, self.HOM] += 2 * self.d Q_test[f"c_{i}", f"c_{j}"] = -np.kron(R, np.eye(self.d)) if DEBUG: x = self.get_x() @@ -390,8 +325,14 @@ def get_Q_from_y(self, y, output_poly=False): else: return Q.get_matrix(self.var_dict) - def test_and_add(self, A_list, Ai, output_poly, b_list=[], bi=0.0): - """Test that constraint Ai holds at current theta then append it to A_list and b_list.""" + def test_and_add( + self, + A_list: list, + Ai: PolyMatrix, + output_poly: bool, + b_list: list = [], + bi: float = 0.0, + ) -> None: x = self.get_x() Ai_sparse = Ai.get_matrix(self.var_dict) err = np.trace(np.atleast_2d(x.T @ Ai_sparse @ x)) - bi @@ -402,23 +343,27 @@ def test_and_add(self, A_list, Ai, output_poly, b_list=[], bi=0.0): A_list.append(Ai_sparse) b_list.append(bi) - def get_A0(self, var_subset=None) -> tuple[list, list]: + def get_A0(self, var_subset: Any = None) -> Tuple[list, list]: """Return the homogenization constraint A0 for chosen level (either h^2=1, or H'H=I).""" if var_subset is None: var_subset = self.var_dict if self.level == "no": return super().get_A0(var_subset=var_subset) else: - # using self.HOM, None just to make it clear that the second argument is not used. A_orth, b_orth = get_orthogonal_constraints( self.HOM, None, self.d, self.level ) return [Ai.get_matrix(var_subset) for Ai in A_orth], list(b_orth) - def get_A_known(self, var_dict=None, output_poly=False, add_redundant=False): + def get_A_known( + self, + var_dict: Dict[str, int] | None = None, + output_poly: bool = False, + add_redundant: bool = False, + ) -> Tuple[list, list]: """Return known linear constraints (A and b). If level 'no' returns A_list; if 'bm' returns (A_list, b_list).""" - A_list = [] - b_list = [] + A_list: list = [] + b_list: list = [] if var_dict is None: var_dict = self.var_dict @@ -434,22 +379,7 @@ def get_A_known(self, var_dict=None, output_poly=False, add_redundant=False): for Ai, bi in zip(A_orth, b_orth): self.test_and_add(A_list, Ai, output_poly, b_list, bi) - # enforce that determinant is one. if self.d == 2 and self.ADD_DETERMINANT: - # level "no": - # C = [a b; c d]; ad - bc - 1 = 0 - # a b c d - # a 1 - # b -1 - # c -1 - # d 1 - # level "bm" - # C = [a b - # c d] - # C @ C.T - # [a b] [a c] [a^2 + b^2 a*c + b*d] - # [c d] [b d] = [a*c + b*d c^2 + d^2] - # cannot be implemented... if self.level == "bm": raise NotImplementedError( "Cannot add determinant constraint for level bm" @@ -462,15 +392,11 @@ def get_A_known(self, var_dict=None, output_poly=False, add_redundant=False): Ai[f"c_{k}", f"c_{k}"] = constraint self.test_and_add(A_list, Ai, output_poly, b_list, 0.0) elif self.d == 3 and self.ADD_DETERMINANT: - # c11 c12 c13 c21 * c32 - c31 * c22 = c13 - # C = [c21, c22, c23]; c1 x c2 = c3: c31 * c12 - c11 * c12 = c23 - # c31 c32 c33 c11 * c22 - c21 * c12 = c33 print( "Warning: consider implementing the determinant constraint for RobustPoseLifter, d=3" ) if add_redundant and f"c_{k}" in var_dict: - # enforce diagonal == 1 for RR' = I for i in range(self.d): Ei = np.zeros((self.d, self.d)) Ei[i, i] = 1.0 @@ -481,10 +407,6 @@ def get_A_known(self, var_dict=None, output_poly=False, add_redundant=False): self.test_and_add(A_list, Ai, output_poly, b_list, 0.0) if self.d == 2: - # enforce structure: - # [cos(x) -sin(x)] - # [sin(x) cos(x)] - # => [c s -s c] Ai = PolyMatrix(symmetric=True) Ai[self.HOM, f"c_{k}"] = np.array([1.0, 0, 0, -1.0])[None, :] self.test_and_add(A_list, Ai, output_poly, b_list, 0.0) @@ -493,7 +415,6 @@ def get_A_known(self, var_dict=None, output_poly=False, add_redundant=False): Ai[self.HOM, f"c_{k}"] = np.array([0, 1.0, 1.0, 0.0])[None, :] self.test_and_add(A_list, Ai, output_poly, b_list, 0.0) - # enforce off-diagonal == 0 for RR' = I for i in range(self.d): for j in range(i + 1, self.d): Ei = np.zeros((self.d, self.d)) @@ -505,7 +426,7 @@ def get_A_known(self, var_dict=None, output_poly=False, add_redundant=False): self.test_and_add(A_list, Ai, output_poly, b_list, 0.0) return A_list, b_list - def plot(self, estimates={}): + def plot(self, estimates: Dict[str, np.ndarray] = {}) -> Tuple[Any, Any]: """Plot ground-truth frames and optional estimated frames; returns (fig, ax).""" import itertools @@ -523,7 +444,7 @@ def plot(self, estimates={}): ls="-", scale=0.5, marker="", - r_wc_w=np.hstack([i * 2.0] + [0.0] * (self.d - 1)), # type.ignore + r_wc_w=np.hstack([i * 2.0] + [0.0] * (self.d - 1)), # type: ignore ) label = None @@ -537,7 +458,7 @@ def plot(self, estimates={}): ls=next(linestyles), scale=1.0, marker="", - r_wc_w=np.hstack([i * 2.0] + [0.0] * (self.d - 1)), # type.ignore + r_wc_w=np.hstack([i * 2.0] + [0.0] * (self.d - 1)), # type: ignore ) label = None @@ -545,7 +466,7 @@ def plot(self, estimates={}): ax.legend() return fig, ax - def __repr__(self): + def __repr__(self) -> str: return f"rotation_lifter{self.d}d_{self.level}" @@ -557,22 +478,12 @@ def __repr__(self): from popcor.utils.plotting_tools import plot_matrix - level = "no" - # level = "bm" - + level: str = "no" np.random.seed(0) - # lifter = RotationLifter( - # d=2, n_abs=0, n_rel=1, n_rot=4, sparsity="chain", level=level - # ) lifter = RotationLifter( d=2, n_abs=1, n_rel=0, n_rot=4, sparsity="chain", level=level ) - # add relative measurements between all subsequent rotations. - # lifter.add_relative_measurement(0, 1, noise=1e-3) - # add one absolute measurement to fix Gauge freedom - # lifter.add_absolute_measurement(0, noise=1e-5) - y = lifter.simulate_y(noise=1e-10) x = lifter.get_x() diff --git a/popcor/examples/stereo1d_lifter.py b/popcor/examples/stereo1d_lifter.py index a9a5a2a..9a399e1 100644 --- a/popcor/examples/stereo1d_lifter.py +++ b/popcor/examples/stereo1d_lifter.py @@ -1,4 +1,14 @@ +"""Stereo1DLifter class for 1D stereo localization example. + +This module provides the Stereo1DLifter class, a pedagogical example for 1D stereo localization. +It demonstrates cost minimization for landmark-based localization, as described in +https://arxiv.org/abs/2308.05783 and used in the Quick Start Guide. +""" + +from typing import Iterable + import numpy as np +import scipy.sparse as sp from poly_matrix.least_squares_problem import LeastSquaresProblem from poly_matrix.poly_matrix import PolyMatrix @@ -22,29 +32,31 @@ class Stereo1DLifter(StateLifter): NOISE = 0.1 - def __init__(self, n_landmarks, param_level="no"): - self.n_landmarks = n_landmarks - self.d = 1 - self.W = 1.0 + def __init__(self, n_landmarks: int, param_level: str = "no") -> None: + self.n_landmarks: int = n_landmarks + self.d: int = 1 + self.W: float = 1.0 # will be initialized later - self.landmarks_ = None + self.landmarks_: np.ndarray | None = None super().__init__(param_level=param_level, d=self.d, n_parameters=n_landmarks) @property - def landmarks(self): + def landmarks(self) -> np.ndarray: if self.landmarks_ is None: self.landmarks_ = np.random.rand(self.n_landmarks, self.d) return self.landmarks_ - def sample_parameters(self, theta=None): + def sample_parameters(self, theta: np.ndarray | None = None) -> dict: + """Sample parameters for the landmarks.""" if self.parameters_ is None: return self.sample_parameters_landmarks(self.landmarks) landmarks = np.random.rand(self.n_landmarks, self.d) return self.sample_parameters_landmarks(landmarks) - def sample_theta(self): + def sample_theta(self) -> np.ndarray | None: + """Sample a valid theta not too close to any landmark.""" x_try = np.random.rand(1) counter = 0 while np.min(np.abs(x_try - self.landmarks)) <= 1e-2: @@ -52,10 +64,15 @@ def sample_theta(self): counter += 1 if counter >= 1000: print("Warning: couldn't find valid setup") - return + return None return x_try - def get_x(self, theta=None, parameters=None, var_subset=None): + def get_x( + self, + theta: np.ndarray | None = None, + parameters: dict | None = None, + var_subset: Iterable | None = None, + ) -> np.ndarray: """ :param var_subset: list of variables to include in x vector. Set to None for all. """ @@ -91,22 +108,26 @@ def get_x(self, theta=None, parameters=None, var_subset=None): return np.hstack(x_data) @property - def var_dict(self): + def var_dict(self) -> dict: vars = [self.HOM, "x"] + [f"z_{j}" for j in range(self.n_landmarks)] return {v: 1 for v in vars} @property - def param_dict(self): + def param_dict(self) -> dict: return self.param_dict_landmarks - def simulate_y(self, noise: float | None = None): + def simulate_y(self, noise: float | None = None) -> np.ndarray: + """Simulate noisy measurements y.""" if noise is None: noise = self.NOISE return 1 / (self.theta - self.landmarks.flatten()) + np.random.normal( scale=noise, loc=0, size=self.n_landmarks ) - def get_Q(self, noise: float | None = None, output_poly: bool = False): + def get_Q( + self, noise: float | None = None, output_poly: bool = False + ) -> np.ndarray | PolyMatrix | sp.csr_matrix | sp.csc_matrix: + """Return the quadratic cost matrix Q for the current measurements.""" if self.landmarks is None: raise ValueError("self.landmarks must be initialized before calling get_Q.") if noise is None: @@ -118,7 +139,10 @@ def get_Q(self, noise: float | None = None, output_poly: bool = False): return self.get_Q_from_y(y, output_poly=output_poly) - def get_Q_from_y(self, y, output_poly: bool = False): + def get_Q_from_y( + self, y: np.ndarray, output_poly: bool = False + ) -> np.ndarray | PolyMatrix | sp.csr_matrix | sp.csc_matrix: + """Return the quadratic cost matrix Q from given measurements y.""" ls_problem = LeastSquaresProblem() for j in range(len(y)): # (z_j - y[j])**2 @@ -128,7 +152,13 @@ def get_Q_from_y(self, y, output_poly: bool = False): else: return ls_problem.get_Q().get_matrix(self.var_dict) - def get_A_known(self, var_dict=None, output_poly=False, add_redundant=False): + def get_A_known( + self, + var_dict: dict | None = None, + output_poly: bool = False, + add_redundant: bool = False, + ) -> tuple[list, list]: + """Return known constraint matrices for the system.""" if var_dict is None: var_dict = self.var_dict @@ -139,7 +169,7 @@ def get_A_known(self, var_dict=None, output_poly=False, add_redundant=False): # enforce that z_j = 1/(x - a_j) <=> 1 - z_j*x + a_j*z_j = 0 if not ("x" in var_dict and self.HOM in var_dict): - return [] + return [], [] landmark_indices = [ int(key.split("_")[-1]) for key in var_dict if key.startswith("z_") @@ -155,10 +185,15 @@ def get_A_known(self, var_dict=None, output_poly=False, add_redundant=False): A_known.append(A.get_matrix(variables=self.var_dict)) if add_redundant: - A_known += self.get_A_known_redundant(var_dict, output_poly) + redundant_A, redundant_b = self.get_A_known_redundant(var_dict, output_poly) + A_known += redundant_A + return A_known, [0] * len(A_known) return A_known, [0] * len(A_known) - def get_A_known_redundant(self, var_dict=None, output_poly=False): + def get_A_known_redundant( + self, var_dict: dict | None = None, output_poly: bool = False + ) -> tuple[list, list]: + """Return redundant known constraint matrices for the system.""" import itertools if var_dict is None: @@ -183,17 +218,26 @@ def get_A_known_redundant(self, var_dict=None, output_poly=False): A_known.append(A.get_matrix(variables=self.var_dict)) return A_known, [0] * len(A_known) - def get_cost(self, theta, y): + def get_cost(self, theta: np.ndarray, y: np.ndarray) -> float: + """Compute the cost function value for given theta and y.""" if np.ndim(theta) == 0 or (np.ndim(theta) == 1 and len(theta) == 1): - return np.sum((y - (1 / (theta - self.landmarks.flatten()))) ** 2) + return float(np.sum((y - (1 / (theta - self.landmarks.flatten()))) ** 2)) else: u = theta[1:] - return np.sum((y - u) ** 2) + return float(np.sum((y - u) ** 2)) def local_solver( - self, t0, y, num_iters=100, eps=1e-5, W=None, verbose=False, **kwargs - ): - info = {} + self, + t0: np.ndarray, + y: np.ndarray, + num_iters: int = 100, + eps: float = 1e-5, + W: float | None = None, + verbose: bool = False, + **kwargs, + ) -> tuple[np.ndarray | None, dict, float | None]: + """Run a local solver to minimize the cost function.""" + info: dict = {} a = self.landmarks.flatten() x_op = t0 for i in range(num_iters): @@ -212,9 +256,9 @@ def local_solver( else: msg = f"converged in du after {i} it" cost = self.get_cost(x_op, y) - info = {"msg": msg, "cost": self.get_cost(x_op, y), "success": True} + info = {"msg": msg, "cost": cost, "success": True} return x_op, info, cost return None, {"msg": "didn't converge", "cost": None, "success": False}, None - def __repr__(self): + def __repr__(self) -> str: return f"stereo1d_{self.param_level}" diff --git a/popcor/examples/stereo2d_lifter.py b/popcor/examples/stereo2d_lifter.py index 4b70fb8..2acb734 100644 --- a/popcor/examples/stereo2d_lifter.py +++ b/popcor/examples/stereo2d_lifter.py @@ -1,4 +1,5 @@ -# import autograd.numpy as np +"""Stereo2DLifter class for 2D stereo localization example.""" + import numpy as np from popcor.base_lifters import StereoLifter @@ -6,13 +7,16 @@ from popcor.utils.stereo2d_problem import _cost, local_solver -def change_dimensions(a, y, x): +def change_dimensions( + a: np.ndarray, y: np.ndarray, x: np.ndarray +) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + """Reshape and concatenate arrays for stereo localization.""" p_w = np.concatenate([a, np.ones((a.shape[0], 1))], axis=1) y_mat = np.c_[[*y]] # N x 2 return p_w[:, :, None], y_mat[:, :, None], x[:, None] -GTOL = 1e-6 +GTOL: float = 1e-6 class Stereo2DLifter(StereoLifter): @@ -39,12 +43,17 @@ class Stereo2DLifter(StereoLifter): where :math:`f_u, f_v` are horizontal and vertical focal lengths, :math:`c_u,c_v` are image center points in pixels and :math:`b` is the camera baseline. - This example is treated in more details in `this paper `_. + This example is treated in more detail in `this paper `_. """ - def __init__(self, n_landmarks, level="no", param_level="no", variable_list=None): - self.W = np.stack([np.eye(2)] * n_landmarks) - + def __init__( + self, + n_landmarks: int, + level: str = "no", + param_level: str = "no", + variable_list: list | None = None, + ): + self.W: np.ndarray = np.stack([np.eye(2)] * n_landmarks) super().__init__( n_landmarks=n_landmarks, level=level, @@ -54,14 +63,16 @@ def __init__(self, n_landmarks, level="no", param_level="no", variable_list=None ) @property - def M_matrix(self): - f_u = 484.5 - c_u = 322 - b = 0.24 + def M_matrix(self) -> np.ndarray: + f_u: float = 484.5 + c_u: float = 322 + b: float = 0.24 return np.array([[f_u, c_u, f_u * b / 2], [f_u, c_u, -f_u * b / 2]]) - def get_cost(self, theta, y, W=None): - + def get_cost( + self, theta: np.ndarray, y: np.ndarray, W: np.ndarray | None = None + ) -> float: + """Compute the cost for given parameters.""" if W is None: W = self.W a = self.landmarks @@ -74,8 +85,15 @@ def get_cost(self, theta, y, W=None): else: return cost - def local_solver(self, t0, y, W=None, verbose=False, **kwargs): - + def local_solver( + self, + t0: np.ndarray, + y: np.ndarray, + W: np.ndarray | None = None, + verbose: bool = False, + **kwargs + ) -> tuple[np.ndarray | None, dict, float]: + """Run local solver for stereo localization.""" if W is None: W = self.W a = self.landmarks @@ -87,7 +105,6 @@ def local_solver(self, t0, y, W=None, verbose=False, **kwargs): ) if StereoLifter.NORMALIZE: cost /= self.n_landmarks * self.d - # cost /= self.n_landmarks * self.d theta_hat = convert_phi_to_theta(phi_hat) info = {"success": success, "msg": "converged", "cost": cost} if success: diff --git a/popcor/examples/stereo3d_lifter.py b/popcor/examples/stereo3d_lifter.py index 514a1d2..1063e54 100644 --- a/popcor/examples/stereo3d_lifter.py +++ b/popcor/examples/stereo3d_lifter.py @@ -1,8 +1,4 @@ -"""Stereo 3D lifter example. - -Contains Stereo3DLifter for stereo-camera localization in 3D. The main -information about the example is in the class docstring. -""" +"""Stereo3DLifter class for 2D stereo localization example.""" import pickle @@ -13,19 +9,15 @@ from popcor.utils.stereo3d_problem import _cost, local_solver -def change_dimensions(a, y): - """Convert landmarks and measurements to solver-friendly shapes. - - Returns p_w with shape (N, 4, 1) (homogeneous world points) and y with - shape (N, 2, 1) (measurements) for subsequent vectorized processing. - """ +def change_dimensions(a: np.ndarray, y: np.ndarray) -> tuple[np.ndarray, np.ndarray]: + """Convert landmarks and measurements to solver-friendly shapes.""" a = np.asarray(a) p_w = np.hstack((a, np.ones((a.shape[0], 1), dtype=a.dtype)))[:, :, None] y_mat = np.asarray(y).reshape((a.shape[0], -1)) return p_w, y_mat[:, :, None] -GTOL = 1e-6 +GTOL: float = 1e-6 class Stereo3DLifter(StereoLifter): @@ -57,18 +49,23 @@ class Stereo3DLifter(StereoLifter): This example is treated in more detail in `this paper `_. """ - def __init__(self, n_landmarks, level="no", param_level="no", variable_list=None): - """Create a Stereo3DLifter and initialize default per-landmark weights.""" + def __init__( + self, + n_landmarks: int, + level: str = "no", + param_level: str = "no", + variable_list: list | None = None, + ): # Pre-allocate W as repeated identity matrices; repeat is slightly faster than building a list. - self.W = np.repeat(np.eye(4)[None, :, :], n_landmarks, axis=0) + self.W: np.ndarray = np.repeat(np.eye(4)[None, :, :], n_landmarks, axis=0) # Precompute the M_matrix once for efficiency. - f_u = 484.5 - f_v = 484.5 - c_u = 322 - c_v = 247 - b = 0.24 - self._M_matrix = np.array( + f_u: float = 484.5 + f_v: float = 484.5 + c_u: float = 322 + c_v: float = 247 + b: float = 0.24 + self._M_matrix: np.ndarray = np.array( [ [f_u, 0, c_u, f_u * b / 2], [0, f_v, c_v, 0], @@ -86,12 +83,12 @@ def __init__(self, n_landmarks, level="no", param_level="no", variable_list=None ) @property - def M_matrix(self): + def M_matrix(self) -> np.ndarray: """Stereo camera calibration matrix used by the cost and solver.""" return self._M_matrix @staticmethod - def from_file(fname): + def from_file(fname: str) -> "Stereo3DLifter": """Load a Stereo3DLifter instance state from a file created by to_file.""" with open(fname, "rb") as f: y_, landmarks, theta, level, param_level, variable_list = pickle.load(f) @@ -107,7 +104,7 @@ def from_file(fname): lifter.theta = theta return lifter - def to_file(self, fname): + def to_file(self, fname: str) -> None: """Serialize the lifter's minimal state to a file for later restoration.""" with open(fname, "wb") as f: pickle.dump( @@ -122,19 +119,18 @@ def to_file(self, fname): f, ) - def get_cost(self, theta, y, W=None): - """Evaluate the reprojection cost for a given parameter vector theta. - - theta can be either: - - a pose vector (x, y, z, yaw, pitch, roll), or - - the full theta vector containing flattened C and x,y,z depending on parameterization. - """ + def get_cost( + self, + theta: np.ndarray, + y: np.ndarray, + W: np.ndarray | None = None, + ) -> float: + """Evaluate the reprojection cost for a given parameter vector theta.""" if W is None: W = self.W a = self.landmarks p_w, y = change_dimensions(a, y) - T = get_T(theta=theta, d=3) cost = _cost(p_w=p_w, y=y, T=T, M=self.M_matrix, W=W) @@ -143,11 +139,15 @@ def get_cost(self, theta, y, W=None): else: return cost - def local_solver(self, t0, y, W=None, verbose=False, **kwargs): - """Run the local solver starting from initial pose t0 and measurements y. - - Returns an estimated pose (theta), solver info dict, and the final cost. - """ + def local_solver( + self, + t0: np.ndarray, + y: np.ndarray, + W: np.ndarray | None = None, + verbose: bool = False, + **kwargs, + ) -> tuple[np.ndarray | None, dict, float]: + """Run the local solver starting from initial pose t0 and measurements y.""" if W is None: W = self.W diff --git a/popcor/examples/wahba_lifter.py b/popcor/examples/wahba_lifter.py index 4d55155..661a7a4 100644 --- a/popcor/examples/wahba_lifter.py +++ b/popcor/examples/wahba_lifter.py @@ -1,78 +1,85 @@ -# import autograd.numpy as np +"""WahbaLifter class for robust pose registration with point-to-point measurements.""" + import numpy as np +import scipy.sparse as sp +from poly_matrix import PolyMatrix from popcor.base_lifters import RobustPoseLifter from popcor.utils.geometry import get_C_r_from_theta from popcor.utils.plotting_tools import plot_frame -N_TRYS = 10 +N_TRYS: int = 10 # TODO(FD) for some reason this is not required as opposed to what is stated in Heng's paper # and it currently breaks tightness (might be a bug in my implementation though) -USE_INEQ = False +USE_INEQ: bool = False -NORMALIZE = False +NORMALIZE: bool = False class WahbaLifter(RobustPoseLifter): + """Robust pose lifter for point-to-point registration. - NOISE = 1e-2 # inlier noise - NOISE_OUT = 1.0 # outlier noise + Doc under construction. - def __init__(self, *args, **kwargs): - """This example is treated in more details in `this paper `_, - under the name "PPR" (point-to-point registration). - """ - # only introduced this for the documentation -- otherwise the RobustPoseLifter __init__ is shown. + In the meantime, this example is treated in more detail in `this paper `_, + under the name "PPR". + + """ + + NOISE: float = 1e-2 # inlier noise + NOISE_OUT: float = 1.0 # outlier noise + + def __init__(self, *args, **kwargs) -> None: return super().__init__(*args, **kwargs) - def h_list(self, t): + def h_list(self, t: np.ndarray) -> list: """ - We want to inforce that - - norm(t) <= 10 (default) - as constraints h_j(t)<=0 + Returns constraints h_j(t) <= 0, e.g., norm(t) <= 10 (default). """ default = super().h_list(t) return default - def get_random_position(self): + def get_random_position(self) -> np.ndarray: + """Generates a random position within the allowed distance.""" return np.random.uniform( -0.5 * self.MAX_DIST ** (1 / self.d), 0.5 * self.MAX_DIST ** (1 / self.d), size=self.d, ) - def get_B_known(self): - """Get inequality constraints of the form x.T @ B @ x <= 0""" + def get_B_known(self) -> list: + """Get inequality constraints of the form x.T @ B @ x <= 0.""" if not USE_INEQ: return [] - default = super().get_B_known() return default - def term_in_norm(self, R, t, pi, ui): + def term_in_norm( + self, R: np.ndarray, t: np.ndarray, pi: np.ndarray, ui: np.ndarray + ) -> np.ndarray: + """Computes the term inside the norm for residual calculation.""" return R @ pi + t - ui - def residual_sq(self, R, t, pi, ui): - # TODO: can easily extend below to matrix-weighted + def residual_sq( + self, R: np.ndarray, t: np.ndarray, pi: np.ndarray, ui: np.ndarray + ) -> float: + """Computes the squared residual for a landmark measurement.""" W = np.eye(self.d) term = self.term_in_norm(R, t, pi, ui) - res_sq = term.T @ W @ term + res_sq = float(term.T @ W @ term) if NORMALIZE: return res_sq / (self.n_landmarks * self.d) ** 2 return res_sq - def plot_setup(self): + def plot_setup(self) -> tuple | None: + """Plots the pose and landmarks setup for d=2.""" if self.d != 2: print("Plotting currently only supported for d=2") - return + return None import matplotlib.pylab as plt fig, ax = plt.subplots() - - # R, t = get_C_r_from_theta(self.theta, self.d) - # ax.scatter(*t, color="k", label="pose") - ax.axis("equal") t_wc_w, C_cw = plot_frame(ax, self.theta, label="pose", color="gray", d=2) @@ -80,11 +87,7 @@ def plot_setup(self): w = self.theta[-self.n_landmarks :] for i in range(self.y_.shape[0]): ax.scatter(*self.landmarks[i], color=f"C{i}", label="landmarks") - - # this vector is in camera coordinates t_cpi_c = self.y_[i] - # t_cpi_w: vector from camera to pi in world coordinates - ax.plot( [t_wc_w[0], self.landmarks[i][0]], [t_wc_w[1], self.landmarks[i][1]], @@ -101,7 +104,8 @@ def plot_setup(self): ax.axis("equal") return fig, ax - def simulate_y(self, noise: float | None = None): + def simulate_y(self, noise: float | None = None) -> np.ndarray: + """Simulates landmark measurements with noise and outliers.""" if noise is None: noise = self.NOISE @@ -132,7 +136,7 @@ def simulate_y(self, noise: float | None = None): break if not valid_measurement and self.robust: self.plot_setup() - raise ValueError("did not find a valid measurement.") + raise ValueError("Did not find a valid measurement.") y[i] = y_i self.y_ = y return y @@ -142,7 +146,8 @@ def get_Q( noise: float | None = None, output_poly: bool = False, use_cliques: list = [], - ): + ) -> np.ndarray | PolyMatrix | sp.csr_matrix | sp.csc_matrix: + """Returns the quadratic cost matrix Q for the current measurements.""" if noise is None: noise = self.NOISE @@ -151,21 +156,17 @@ def get_Q( Q = self.get_Q_from_y(self.y_, output_poly=output_poly, use_cliques=use_cliques) return Q - def get_Q_from_y(self, y, output_poly: bool = False, use_cliques: list = []): + def get_Q_from_y( + self, + y: np.ndarray, + output_poly: bool = False, + use_cliques: list = [], + ) -> np.ndarray | PolyMatrix | sp.csr_matrix | sp.csc_matrix: """ - every cost term can be written as + Returns the quadratic cost matrix Q from measurements y. + Every cost term can be written as: (1 + wi)/b^2 r^2(x, zi) + (1 - wi) - - residual term: - (Rpi + t - ui).T Wi (Rpi + t - ui) = - [t', vec(R)'] @ [I (pi x I)]' @ Wi @ [I (pi x I)] @ [t ; vec(R)] - ------x'----- -----Pi'----- - - 2 [t', vec(R)'] @ [I (pi x I)]' Wi @ ui - -----x'------ ---------Pi_xl-------- - + ui.T @ Wi @ ui - -----Pi_ll------ """ - if len(use_cliques): js = use_cliques else: @@ -195,12 +196,10 @@ def get_Q_from_y(self, y, output_poly: bool = False, use_cliques: list = []): Qi /= self.beta**2 Pi_ll /= self.beta**2 Pi_xl /= self.beta**2 - # Q["x", "x"] += Qi Q["t", "t"] += Qi[: self.d, : self.d] Q["t", "c"] += Qi[: self.d, self.d :] Q["c", "c"] += Qi[self.d :, self.d :] - # Q["x", self.HOM] += Pi_xl Q["t", self.HOM] += Pi_xl[: self.d, :] Q["c", self.HOM] += Pi_xl[self.d :, :] Q[self.HOM, self.HOM] += ( @@ -210,7 +209,6 @@ def get_Q_from_y(self, y, output_poly: bool = False, use_cliques: list = []): self.HOM, f"w_{i}" ] += -0.5 # from (1 - wi), 0.5 cause on off-diagonal if self.level == "xwT": - # Q[f"z_{i}", "x"] += 0.5 * Qi Q[f"z_{i}", "t"] += 0.5 * Qi[:, : self.d] Q[f"z_{i}", "c"] += 0.5 * Qi[:, self.d :] @@ -220,18 +218,15 @@ def get_Q_from_y(self, y, output_poly: bool = False, use_cliques: list = []): elif self.level == "xxT": Q["z_0", f"w_{i}"] += 0.5 * Qi.flatten()[:, None] - # Q["x", f"w_{i}"] += Pi_xl Q["t", f"w_{i}"] += Pi_xl[: self.d, :] Q["c", f"w_{i}"] += Pi_xl[self.d :, :] Q[self.HOM, f"w_{i}"] += 0.5 * Pi_ll else: - # Q["x", "x"] += Qi Q["t", "t"] += Qi[: self.d, : self.d] Q["t", "c"] += Qi[: self.d, self.d :] Q["c", "c"] += Qi[self.d :, self.d :] - # Q["x", self.HOM] += Pi_xl Q["t", self.HOM] += Pi_xl[: self.d, :] Q["c", self.HOM] += Pi_xl[self.d :, :] Q[self.HOM, self.HOM] += Pi_ll # on diagonal @@ -240,6 +235,6 @@ def get_Q_from_y(self, y, output_poly: bool = False, use_cliques: list = []): Q_sparse = 0.5 * Q.get_matrix(variables=self.var_dict) return Q_sparse - def __repr__(self): + def __repr__(self) -> str: appendix = "_robust" if self.robust else "" return f"wahba_{self.d}d_{self.level}_{self.param_level}{appendix}" From 05794eeb8a26e7169e9eb8d30fa549476397e4a2 Mon Sep 17 00:00:00 2001 From: Frederike Duembgen Date: Sat, 23 Aug 2025 13:30:25 +0200 Subject: [PATCH 10/10] Fix and extend type hints --- popcor/auto_template.py | 4 +- popcor/auto_tight.py | 206 ++++++++----------- popcor/base_lifters/_base_class.py | 312 +++++++++++++++-------------- popcor/examples/example_lifter.py | 35 +++- popcor/examples/mono_lifter.py | 25 +-- popcor/examples/ro_nsq_lifter.py | 2 +- popcor/examples/ro_sq_lifter.py | 6 +- popcor/examples/wahba_lifter.py | 8 +- popcor/utils/common.py | 89 ++++---- popcor/utils/constraint.py | 124 +++++++----- popcor/utils/plotting_tools.py | 2 +- tests/test_parameters.py | 6 +- 12 files changed, 418 insertions(+), 401 deletions(-) diff --git a/popcor/auto_template.py b/popcor/auto_template.py index 8c83055..66cf4d6 100644 --- a/popcor/auto_template.py +++ b/popcor/auto_template.py @@ -454,7 +454,7 @@ def find_local_solution(self, n_inits=None, verbose=False, plot=False): for key, qcqp_that_local in info.items(): if key.startswith("local solution"): solution_idx = key.strip("local solution ") - error_dict = self.lifter.get_error(qcqp_that_local) + error_dict = self.lifter.get_error(np.asarray(qcqp_that_local)) self.solver_vars.update( { f"local {solution_idx} {error_name}": err @@ -589,11 +589,13 @@ def learn_templates(self, plot=False, data_dict=None): if self.use_incremental: for c in self.templates: ai = get_vec(c.A_poly_.get_matrix(mat_var_dict)) + assert isinstance(ai, np.ndarray) bi = self.lifter.augment_using_zero_padding(ai, param_dict) a_vectors.append(bi) if self.use_known: for c in self.templates_known_sub: ai = get_vec(c.A_poly_.get_matrix(mat_var_dict)) + assert isinstance(ai, np.ndarray) bi = self.lifter.augment_using_zero_padding(ai, param_dict) a_vectors.append(bi) Y = np.vstack([Y] + a_vectors) diff --git a/popcor/auto_tight.py b/popcor/auto_tight.py index 4697279..4f1dec5 100644 --- a/popcor/auto_tight.py +++ b/popcor/auto_tight.py @@ -1,4 +1,7 @@ +"""Module for automatic constraint generation using nullspace methods.""" + import warnings +from typing import Iterable import matplotlib.pylab as plt import numpy as np @@ -14,36 +17,22 @@ class AutoTight(object): """Class for automatic constraint generation.""" - # consider singular value zero below this - EPS_SVD = 1e-5 - - # basis pursuit method, can be - # - qr: qr decomposition - # - qrp: qr decomposition with permutations (sparser), recommended - # - svd: svd - METHOD = "qrp" - - # normalize learned Ai or not - NORMALIZE = False - - # how much to oversample (>= 1) - FACTOR = 1.2 - - # number of times we remove bad samples from data matrix - N_CLEANING_STEPS = 1 # was 3 - - # maximum number of iterations of local solver - LOCAL_MAXITER = 100 - - # find and remove linearly dependent constraints - REDUCE_DEPENDENT = False + EPS_SVD: float = 1e-5 # consider singular value zero below this + METHOD: str = "qrp" # basis pursuit method: 'qr', 'qrp', 'svd' + NORMALIZE: bool = False # normalize learned Ai or not + FACTOR: float = 1.2 # how much to oversample (>= 1) + N_CLEANING_STEPS: int = 1 # number of times we remove bad samples from data matrix + LOCAL_MAXITER: int = 100 # maximum number of iterations of local solver + REDUCE_DEPENDENT: bool = False # find and remove linearly dependent constraints - def __init__(self): + def __init__(self) -> None: pass @staticmethod - def clean_Y(basis_new, Y, s, plot=False): - errors = np.abs(basis_new @ Y.T) # Nb x n x n x Ns = Nb x Ns + def clean_Y( + basis_new: np.ndarray, Y: np.ndarray, s: np.ndarray, plot: bool = False + ) -> list: + errors = np.abs(basis_new @ Y.T) # Nb x Ns if np.all(errors < 1e-10): return [] bad_bins = np.unique(np.argmax(errors, axis=1)) @@ -56,22 +45,26 @@ def clean_Y(basis_new, Y, s, plot=False): return bad_bins @staticmethod - def test_S_cutoff(S, corank, eps_svd=None): + def test_S_cutoff(S: np.ndarray, corank: int, eps_svd: float | None = None) -> None: if eps_svd is None: eps_svd = AutoTight.EPS_SVD if corank > 1: try: - assert abs(S[-corank]) / eps_svd < 1e-1 # 1e-1 1e-10 - assert abs(S[-corank - 1]) / eps_svd > 10 # 1e-11 1e-10 + assert abs(S[-corank]) / eps_svd < 1e-1 + assert abs(S[-corank - 1]) / eps_svd > 10 except AssertionError: print(f"there might be a problem with the chosen threshold {eps_svd}:") print(S[-corank], eps_svd, S[-corank - 1]) @staticmethod def get_basis_sparse( - lifter, var_list, param_list, A_known=[], plot=False, eps_svd=None - ): - + lifter, + var_list: list, + param_list: list, + A_known: list = [], + plot: bool = False, + eps_svd: float | None = None, + ) -> list: Y = AutoTight.generate_Y_sparse( lifter, var_subset=var_list, param_subset=param_list, factor=1.0 ) @@ -91,7 +84,6 @@ def get_basis_sparse( ) if plot: plot_matrix = np.vstack([t.b_[None, :] for t in constraints]) - cmap, norm, colorbar_yticks = initialize_discrete_cbar(plot_matrix) X_dim = lifter.get_dim_X(var_list) fig, ax = plt.subplots() @@ -107,21 +99,12 @@ def get_basis_sparse( @staticmethod def get_A_learned( lifter, - A_known=[], - var_dict=None, - method=METHOD, - verbose=False, - ) -> tuple[list, np.ndarray | list]: - """Generate list of learned constraints by sampling the lifter. - - :param lifter: StateLifter object - :param A_known: list of known constraints, if given, will generate basis that is orthogonal to these given constraints. - :param var_dict: variable dictionary, if None, will use all variables - :param method: method to use for basis generation, can be 'qr', 'qrp', or 'svd'. 'qrp' is recommended. - :param verbose: if True, will print timing information - - :return: list of learned constraints. - """ + A_known: list = [], + var_dict: dict | list | None = None, + method: str = METHOD, + verbose: bool = False, + ) -> tuple[list, list]: + """Generate list of learned constraints by sampling the lifter.""" import time t1 = time.time() @@ -140,7 +123,7 @@ def get_A_learned( b_learned = -basis[:, 0] if verbose: print(f"get matrices ({len(A_learned)}): {time.time() - t1:4.4f}") - return A_learned, b_learned + return A_learned, list(b_learned) else: A_learned = AutoTight.generate_matrices(lifter, basis, var_dict=var_dict) if verbose: @@ -150,10 +133,10 @@ def get_A_learned( @staticmethod def get_A_learned_simple( lifter, - A_known=[], - var_dict=None, - method=METHOD, - verbose=False, + A_known: list = [], + var_dict: Iterable | None = None, + method: str = METHOD, + verbose: bool = False, ) -> list: """Simplified version of get_A_learned that does not consider parameters.""" import time @@ -185,14 +168,15 @@ def get_A_learned_simple( return A_learned @staticmethod - def generate_Y_simple(lifter, var_subset, factor): + def generate_Y_simple( + lifter, var_subset: Iterable | None = None, factor: float = 1.0 + ) -> np.ndarray: # need at least dim_Y different random setups dim_Y = lifter.get_dim_X(var_subset) n_seeds = int(dim_Y * factor) Y = np.empty((n_seeds, dim_Y)) for seed in range(n_seeds): np.random.seed(seed) - theta = lifter.sample_theta() x = lifter.get_x(theta=theta, parameters=None, var_subset=var_subset) X = np.outer(x, x) @@ -200,39 +184,40 @@ def generate_Y_simple(lifter, var_subset, factor): return Y @staticmethod - def generate_Y_sparse(lifter, var_subset, param_subset, factor=FACTOR, ax=None): + def generate_Y_sparse( + lifter, var_subset: list, param_subset: list, factor: float = FACTOR, ax=None + ) -> np.ndarray: from popcor.base_lifters import StateLifter assert isinstance(lifter, StateLifter) assert lifter.HOM in param_subset - # need at least dim_Y different random setups dim_Y = lifter.get_dim_Y(var_subset, param_subset) n_seeds = int(dim_Y * factor) Y = np.empty((n_seeds, dim_Y)) for seed in range(n_seeds): np.random.seed(seed) - theta = lifter.sample_theta() parameters = lifter.sample_parameters(theta) - if seed < 10 and ax is not None: if np.ndim(lifter.theta) == 1: ax.scatter(np.arange(len(theta)), theta) else: ax.scatter(*theta[:, :2].T) - x = lifter.get_x(theta=theta, parameters=parameters, var_subset=var_subset) X = np.outer(x, x) - - # generates [1*x, a1*x, ..., aK*x] p = lifter.get_p(parameters=parameters, param_subset=param_subset) Y[seed, :] = np.kron(p, get_vec(X)) return Y @staticmethod - def generate_Y(lifter, factor=FACTOR, ax=None, var_subset=None, param_subset=None): - # need at least dim_Y different random setups + def generate_Y( + lifter, + factor: float = FACTOR, + ax=None, + var_subset: Iterable | None = None, + param_subset: Iterable | None = None, + ) -> np.ndarray: dim_Y = lifter.get_dim_Y(var_subset, param_subset) if lifter.level == "bm": dim_Y += lifter.get_dim_P(param_subset) @@ -240,7 +225,6 @@ def generate_Y(lifter, factor=FACTOR, ax=None, var_subset=None, param_subset=Non Y = np.empty((n_seeds, dim_Y)) for seed in range(n_seeds): np.random.seed(seed) - theta = lifter.sample_theta() parameters = lifter.sample_parameters(theta) if seed < 10 and ax is not None: @@ -248,43 +232,30 @@ def generate_Y(lifter, factor=FACTOR, ax=None, var_subset=None, param_subset=Non ax.scatter(np.arange(len(theta)), theta) else: ax.scatter(*theta[:, :2].T) - x = lifter.get_x(theta=theta, parameters=parameters, var_subset=var_subset) if np.ndim(x) == 1: x = x[:, None] X = x @ x.T - - # generates [1*x, a1*x, ..., aK*x] p = lifter.get_p(parameters=parameters, param_subset=param_subset) assert p[0] == 1 - row = get_vec(X) if lifter.level == "bm": - # in Burer-Monteiro, we need to add the leading 1. row = np.hstack([np.array([1.0]), np.array(row)]) - Y[seed, :] = np.kron(p, row) return Y @staticmethod def get_basis( lifter, - Y, + Y: np.ndarray, A_known: list = [], basis_known: np.ndarray | None = None, - method=METHOD, - eps_svd=None, - ): - """Generate basis from lifted state matrix Y. - - :param A_known: if given, will generate basis that is orthogonal to these given constraints. - - :return: basis, S - """ + method: str = METHOD, + eps_svd: float | None = None, + ) -> tuple[np.ndarray, np.ndarray]: + """Generate basis from lifted state matrix Y.""" if eps_svd is None: eps_svd = AutoTight.EPS_SVD - - # if there is a known list of constraints, add them to the Y so that resulting nullspace is orthogonal to them if basis_known is not None: if len(A_known): warnings.warn( @@ -297,21 +268,19 @@ def get_basis( [lifter.augment_using_zero_padding(get_vec(a)) for a in A_known] ) Y = np.vstack([Y, A]) - basis, info = get_nullspace(Y, method=method, tolerance=eps_svd) - basis[np.abs(basis) < lifter.EPS_SPARSE] = 0.0 return basis, info["values"] @staticmethod def generate_matrices_simple( lifter, - basis, - normalize=NORMALIZE, - sparse=True, - trunc_tol=1e-10, - var_dict=None, - ): + basis: np.ndarray, + normalize: bool = NORMALIZE, + sparse: bool = True, + trunc_tol: float = 1e-10, + var_dict: Iterable | None = None, + ) -> list: """ Generate constraint matrices from the rows of the nullspace basis matrix. """ @@ -319,92 +288,85 @@ def generate_matrices_simple( n_basis = len(basis) except Exception: n_basis = basis.shape[0] - if isinstance(var_dict, list): var_dict = lifter.get_var_dict(var_dict) - from popcor.base_lifters import StateLifter assert isinstance(lifter, StateLifter) - A_list = [] for i in range(n_basis): ai = basis[i] Ai = lifter.get_mat(ai, sparse=sparse, correct=True, var_dict=None) - # Normalize the matrix if normalize and not sparse: - # Ai /= np.max(np.abs(Ai)) assert isinstance(Ai, np.ndarray) - Ai /= np.linalg.norm(Ai, p=2) # type: ignore + Ai /= np.linalg.norm(Ai, ord=2) elif normalize and sparse: Ai /= splinalg.norm(Ai, ord="fro") - # Sparsify and truncate if sparse: - Ai.eliminate_zeros() # type: ignore + assert isinstance(Ai, (sp.csr_matrix, sp.csc_matrix)) + Ai.eliminate_zeros() else: - Ai[np.abs(Ai) < trunc_tol] = 0.0 # type: ignore - # add to list + assert isinstance(Ai, np.ndarray) + Ai[np.abs(Ai) < trunc_tol] = 0.0 A_list.append(Ai) return A_list @staticmethod def generate_matrices( lifter, - basis, - normalize=NORMALIZE, - sparse=True, - trunc_tol=1e-10, - var_dict=None, - ): + basis: np.ndarray, + normalize: bool = NORMALIZE, + sparse: bool = True, + trunc_tol: float = 1e-10, + var_dict: dict | list | None = None, + ) -> list: """ Generate constraint matrices from the rows of the nullspace basis matrix. """ from popcor.base_lifters import StateLifter assert isinstance(lifter, StateLifter) - try: n_basis = len(basis) except Exception: n_basis = basis.shape[0] - if isinstance(var_dict, list): var_dict = lifter.get_var_dict(var_dict) - A_list = [] basis_reduced = [] for i in range(n_basis): ai = lifter.get_reduced_a(bi=basis[i], var_subset=var_dict, sparse=True) basis_reduced.append(ai) basis_reduced = sp.vstack(basis_reduced) - + assert isinstance(basis_reduced, sp.csr_matrix) if AutoTight.REDUCE_DEPENDENT: bad_idx = find_dependent_columns(basis_reduced.T, tolerance=1e-6) else: bad_idx = [] - - for i in range(basis_reduced.shape[0]): # type: ignore + for i in range(basis_reduced.shape[0]): if i in bad_idx: continue - ai = basis_reduced[[i], :].toarray().flatten() # type: ignore + ai = basis_reduced[[i], :].toarray().flatten() Ai = lifter.get_mat(ai, sparse=sparse, correct=True, var_dict=None) - # Normalize the matrix + if normalize and not sparse: - # Ai /= np.max(np.abs(Ai)) - Ai /= np.linalg.norm(Ai, p=2) # type: ignore + assert isinstance(Ai, np.ndarray) + Ai /= np.linalg.norm(Ai, ord=2) elif normalize and sparse: Ai /= splinalg.norm(Ai, ord="fro") - # Sparsify and truncate if sparse: - Ai.eliminate_zeros() # type: ignore + assert isinstance(Ai, (sp.csr_matrix, sp.csc_matrix)) + Ai.eliminate_zeros() else: - Ai[np.abs(Ai) < trunc_tol] = 0.0 # type: ignore - # add to list + assert isinstance(Ai, np.ndarray) + Ai[np.abs(Ai) < trunc_tol] = 0.0 A_list.append(Ai) return A_list @staticmethod - def get_duality_gap(cost_local, cost_sdp, relative=True): + def get_duality_gap( + cost_local: float, cost_sdp: float, relative: bool = True + ) -> float: if relative: return (cost_local - cost_sdp) / abs(cost_sdp) else: diff --git a/popcor/base_lifters/_base_class.py b/popcor/base_lifters/_base_class.py index 776dde9..e340f2b 100644 --- a/popcor/base_lifters/_base_class.py +++ b/popcor/base_lifters/_base_class.py @@ -19,22 +19,25 @@ class BaseClass(object): # Homogenization variable name - HOM = "h" + HOM: str = "h" # set elements below this threshold to zero. - EPS_SPARSE = 1e-9 + EPS_SPARSE: float = 1e-9 # properties of template scaling - ALL_PAIRS = True + ALL_PAIRS: bool = True # Below only have effect if ALL_PAIRS is False. # Then, they determine the clique size hierarchy. - CLIQUE_SIZE = 5 - STEP_SIZE = 1 + CLIQUE_SIZE: int = 5 + STEP_SIZE: int = 1 # tolerance for feasibility error of learned constraints - EPS_ERROR = 1e-8 + EPS_ERROR: float = 1e-8 - PARAM_LEVELS = ["no", "p", "ppT"] + PARAM_LEVELS: list[str] = ["no", "p", "ppT"] + + theta_: np.ndarray | None = None + parameters_: dict | None = None @property @abstractmethod @@ -51,16 +54,16 @@ def sample_theta(self) -> np.ndarray: pass @abstractmethod - def sample_parameters(self) -> np.ndarray: + def sample_parameters(self) -> dict: pass @staticmethod - def get_variable_indices(var_subset, variable="z"): + def get_variable_indices(var_subset: Iterable, variable: str = "z") -> np.ndarray: return np.unique( [int(v.split("_")[-1]) for v in var_subset if v.startswith(f"{variable}_")] ) - def __init__(self, d, param_level, n_parameters): + def __init__(self, d: int, param_level: str, n_parameters: int) -> None: assert param_level in self.PARAM_LEVELS self.param_level = param_level self.d = d @@ -71,11 +74,11 @@ def __init__(self, d, param_level, n_parameters): self.n_parameters = n_parameters else: raise ValueError(f"Unknown param_level: {param_level}") - # n_parameters * (self.d + (self.d * (self.d + 1)) // 2) self.generate_random_setup() - # Functionalities related to var_dict - def get_var_dict(self, var_subset=None, unroll_keys=False): + def get_var_dict( + self, var_subset: Iterable | None = None, unroll_keys: bool = False + ) -> dict: if var_subset is not None: var_dict = {k: v for k, v in self.var_dict.items() if k in var_subset} if unroll_keys: @@ -86,20 +89,19 @@ def get_var_dict(self, var_subset=None, unroll_keys=False): return unroll(self.var_dict) return self.var_dict - def get_param_dict(self, param_subset=None): + def get_param_dict(self, param_subset: Iterable | None = None) -> dict: if param_subset is not None: return {k: v for k, v in self.param_dict.items() if k in param_subset} return self.param_dict - # Functionalities related to random setups @property - def theta(self): + def theta(self) -> np.ndarray: if self.theta_ is None: self.theta_ = self.sample_theta() return self.theta_ @theta.setter - def theta(self, t): + def theta(self, t: np.ndarray) -> None: assert ( self.theta_ is None ), "The property self.theta is only meant to be set once!" @@ -110,17 +112,19 @@ def parameters(self) -> dict: if self.parameters_ is None: self.parameters_ = self.sample_parameters() assert isinstance(self.parameters_, dict) - return self.parameters_ # type: ignore + return self.parameters_ @parameters.setter - def parameters(self, p): + def parameters(self, p: dict) -> None: assert ( self.parameters_ is None ), "The property self.parameters is only meant to be set once!" assert isinstance(p, dict) self.parameters_ = p - def extract_A_known(self, A_known, var_subset, output_type="csc"): + def extract_A_known( + self, A_known: list, var_subset: dict | list, output_type: str = "csc" + ) -> list | np.ndarray: """ Extract from the list of constraint matrices only the ones that touch only a subset of var_subset. @@ -134,23 +138,21 @@ def extract_A_known(self, A_known, var_subset, output_type="csc"): assert len(A_poly.get_variables()) > 0 - # if all of the non-zero elements of A_poly are in var_subset, - # we can use this matrix. if np.all([v in var_subset for v in A_poly.get_variables()]): Ai = A_poly.get_matrix( self.get_var_dict(var_subset), output_type=output_type ) if output_type == "dense": - ai = self.augment_using_zero_padding( - get_vec(Ai, correct=True), var_subset - ) + a = get_vec(Ai, correct=True) + assert isinstance(a, np.ndarray) + ai = self.augment_using_zero_padding(a, var_subset) sub_A_known = np.r_[sub_A_known, ai[None, :]] else: assert isinstance(sub_A_known, list) sub_A_known.append(Ai) return sub_A_known - def get_param_idx_dict(self, var_subset=None): + def get_param_idx_dict(self, var_subset: Iterable) -> dict: """ Give the current subset of variables, extract the parameter dictionary to use. Example: @@ -160,24 +162,21 @@ def get_param_idx_dict(self, var_subset=None): - if param_level == 'p': {'l': 0, 'p_0:0': 1, ..., 'p_0:d-1': d} - if param_level == 'ppT': {'l': 0, 'p_0:0.p_0:0': 1, ..., 'p_0:d-1:.p_0:d-1': 1} """ - # TODO(FD) change to use param_subset instead. param_subset = [self.HOM] + [ f"p_{i}" for i in self.get_variable_indices(var_subset) ] return unroll(self.get_param_dict(param_subset)) - def get_mat(self, vec, sparse=False, var_dict=None, correct=True): - """Convert (N+1)N/2 vectorized matrix to NxN Symmetric matrix in a way that preserves inner products. - - In particular, this means that we divide the off-diagonal elements by sqrt(2). - - :param vec (ndarray): vector of upper-diagonal elements - :return: symmetric matrix filled with vec. - """ - # len(vec) = k = n(n+1)/2 -> dim_x = n = + def get_mat( + self, + vec: np.ndarray | sp.csr_matrix | sp.csc_matrix, + sparse: bool = False, + var_dict: Iterable | None = None, + correct: bool = True, + ) -> np.ndarray | sp.csr_matrix | sp.csc_matrix: + """Convert (N+1)N/2 vectorized matrix to NxN Symmetric matrix in a way that preserves inner products.""" if var_dict is None: pass - elif not isinstance(var_dict, dict): var_dict = {k: v for k, v in self.var_dict.items() if k in var_dict} @@ -187,8 +186,6 @@ def get_mat(self, vec, sparse=False, var_dict=None, correct=True): if var_dict is None: return Ai - # if var_dict is not None, then Ai corresponds to the subblock - # defined by var_dict, of the full constraint matrix. Ai_poly, __ = PolyMatrix.init_from_sparse(Ai, var_dict, unfold=True) from poly_matrix.poly_matrix import augment @@ -197,23 +194,25 @@ def get_mat(self, vec, sparse=False, var_dict=None, correct=True): return Ai_poly.get_matrix(all_var_dict) def var_list_row( - self, var_subset=None, param_subset=None, force_parameters_off=False - ): + self, + var_subset: list | dict | None = None, + param_subset: list | dict | None = None, + force_parameters_off: bool = False, + ) -> list[str]: if var_subset is None: var_subset = list(self.var_dict.keys()) elif isinstance(var_subset, dict): var_subset = list(var_subset.keys()) + assert var_subset is not None if param_subset is None: param_subset = self.param_dict - label_list = [] + label_list: list[str] = [] if force_parameters_off: param_dict = {self.HOM: 0} else: - param_dict = unroll( - self.get_param_dict(param_subset) - ) # self.get_param_idx_dict(var_subset) + param_dict = unroll(self.get_param_dict(param_subset)) for idx, key in enumerate(param_dict.keys()): for i in range(len(var_subset)): @@ -232,12 +231,12 @@ def var_list_row( for dj in djs: keyj = f"{zj}:{dj}" if sizej > 1 else f"{zj}" label_list.append(f"{key}-{keyi}.{keyj}") - # for zi, zj in vectorized_var_list: - # label_list += self.get_labels(key, zi, zj) assert len(label_list) == (idx + 1) * self.get_dim_X(var_subset) return label_list - def var_dict_row(self, var_subset=None, force_parameters_off=False): + def var_dict_row( + self, var_subset: list | dict | None = None, force_parameters_off: bool = False + ) -> dict: return { label: 1 for label in self.var_list_row( @@ -245,14 +244,14 @@ def var_dict_row(self, var_subset=None, force_parameters_off=False): ) } - def get_basis_from_poly_rows(self, basis_poly_list, var_subset=None): + def get_basis_from_poly_rows( + self, basis_poly_list: list, var_subset: list | dict | None = None + ) -> np.ndarray: var_dict = self.get_var_dict(var_subset=var_subset) all_dict = {label: 1 for label in self.var_list_row(var_subset)} basis_reduced = np.empty((0, len(all_dict))) for i, bi_poly in enumerate(basis_poly_list): - # test that this constraint holds - bi = bi_poly.get_matrix(([self.HOM], all_dict)) if bi.shape[1] == self.get_dim_X(var_subset) * self.get_dim_P(): @@ -266,7 +265,6 @@ def get_basis_from_poly_rows(self, basis_poly_list, var_subset=None): print(f"found b{i} has error: {err}") continue - # test that this constraint is lin. independent of previous ones. basis_reduced_test = np.vstack([basis_reduced, bi.toarray()]) rank = np.linalg.matrix_rank(basis_reduced_test) if rank == basis_reduced_test.shape[0]: @@ -276,13 +274,10 @@ def get_basis_from_poly_rows(self, basis_poly_list, var_subset=None): print(f"left with {basis_reduced.shape} total constraints") return basis_reduced - def convert_polyrow_to_Apoly(self, poly_row, correct=True): - """Convert poly-row to reduced a. - - poly-row has elements with keys "pk:l.xi:m.xj:n", - meaning this element corresponds to the l-th element of parameter i, - and the m-n-th element of xj times xk. - """ + def convert_polyrow_to_Apoly( + self, poly_row: PolyMatrix, correct: bool = True + ) -> PolyMatrix: + """Convert poly-row to reduced a.""" parameters = self.get_p() param_dict = dict(zip(unroll(self.param_dict), parameters)) @@ -295,7 +290,6 @@ def convert_polyrow_to_Apoly(self, poly_row, correct=True): m = keyi_m.split(":")[-1] n = keyj_n.split(":")[-1] - # divide off-diagonal elements by sqrt(2) newval = poly_row[self.HOM, key] * param_val if correct and not ((keyi_m == keyj_n) and (m == n)): newval /= np.sqrt(2) @@ -303,34 +297,36 @@ def convert_polyrow_to_Apoly(self, poly_row, correct=True): poly_mat[keyi_m, keyj_n] += newval return poly_mat - def convert_polyrow_to_Asparse(self, poly_row, var_subset=None): + def convert_polyrow_to_Asparse( + self, poly_row: PolyMatrix, var_subset: Iterable | None = None + ) -> sp.csr_matrix | sp.csc_matrix: poly_mat = self.convert_polyrow_to_Apoly(poly_row, correct=False) var_dict = self.get_var_dict(var_subset) - mat_var_list = [] + mat_var_list: list[str] = [] for var, size in var_dict.items(): if size == 1: mat_var_list.append(var) else: mat_var_list += [f"{var}:{i}" for i in range(size)] mat_sparse = poly_mat.get_matrix({m: 1 for m in mat_var_list}) + assert isinstance(mat_sparse, (sp.csr_matrix, sp.csc_matrix)) return mat_sparse - def convert_polyrow_to_a(self, poly_row, var_subset=None, sparse=False): - """Convert poly-row to reduced a. - - poly-row has elements with keys "pk:l.xi:m.xj:n", - meaning this element corresponds to the l-th element of parameter i, - and the m-n-th element of xj times xk. - """ + def convert_polyrow_to_a( + self, + poly_row: PolyMatrix, + var_subset: Iterable | None = None, + sparse: bool = False, + ) -> np.ndarray | sp.csr_matrix: + """Convert poly-row to reduced a.""" mat_sparse = self.convert_polyrow_to_Asparse(poly_row, var_subset) return get_vec(mat_sparse, correct=False, sparse=sparse) - # TODO(FD) consider removing this cause only used in tests. def convert_a_to_polyrow( self, - a, - var_subset=None, + a: np.ndarray, + var_subset: Iterable | None = None, ) -> PolyMatrix: """Convert a array to poly-row.""" if var_subset is None: @@ -354,7 +350,6 @@ def convert_a_to_polyrow( if keyi != keyj: vals = val.flatten() else: - # TODO: use get_vec instead? vals = val[np.triu_indices(val.shape[0])] assert len(labels) == len(vals) for label, v in zip(labels, vals): @@ -363,7 +358,11 @@ def convert_a_to_polyrow( return poly_row def convert_b_to_polyrow( - self, b, var_subset, param_subset=None, tol=1e-10 + self, + b: np.ndarray, + var_subset: list | dict, + param_subset: list | dict | None = None, + tol: float = 1e-10, ) -> PolyMatrix: """Convert (augmented) b array to poly-row.""" if isinstance(b, PolyMatrix): @@ -375,7 +374,6 @@ def convert_b_to_polyrow( poly_row = PolyMatrix(symmetric=False) mask = np.abs(b) > tol - # get the variable names such as p_0:0-x:0.x:4 whch corresponds to p_0[0]*x[0]*x[4] var_list_row = self.var_list_row(var_subset, param_subset) assert len(b) == len(var_list_row) @@ -383,29 +381,38 @@ def convert_b_to_polyrow( poly_row[self.HOM, var_list_row[idx]] = b[idx] return poly_row - def get_dim_x(self, var_subset=None): + def get_dim_x(self, var_subset: Iterable | None = None) -> int: var_dict = self.get_var_dict(var_subset) return sum([val for val in var_dict.values()]) - def get_dim_p(self, param_subset=None): + def get_dim_p(self, param_subset: Iterable | None = None) -> int: param_dict = self.get_param_dict(param_subset) return sum([val for val in param_dict.values()]) - def get_dim_Y(self, var_subset=None, param_subset=None): + def get_dim_Y( + self, + var_subset: Iterable | None = None, + param_subset: list | dict | None = None, + ) -> int: dim_X = self.get_dim_X(var_subset=var_subset) dim_P = self.get_dim_P(param_subset=param_subset) return int(dim_X * dim_P) - def get_dim_X(self, var_subset=None): + def get_dim_X(self, var_subset: Iterable | None = None) -> int: dim_x = self.get_dim_x(var_subset) return int(dim_x * (dim_x + 1) / 2) - def get_dim_P(self, param_subset=None): + def get_dim_P(self, param_subset: list | dict | None = None) -> int: return len(self.get_p(param_subset=param_subset)) def get_reduced_a( - self, bi, param_here=None, var_subset=None, param_subset=None, sparse=False - ): + self, + bi: np.ndarray | PolyMatrix | sp.csc_matrix | sp.csr_matrix, + param_here: np.ndarray | None = None, + var_subset: list | dict | None = None, + param_subset: list | dict | None = None, + sparse: bool = False, + ) -> np.ndarray | sp.csr_array: """ Extract first block of bi by summing over other blocks times the parameters. """ @@ -415,10 +422,10 @@ def get_reduced_a( if isinstance(bi, np.ndarray): len_b = len(bi) elif isinstance(bi, PolyMatrix): - bi = bi.get_matrix(([self.HOM], self.var_dict_row(var_subset))) - len_b = bi.shape[1] # type: ignore + var_dict_row = self.var_dict_row(var_subset) + bi = bi.get_matrix(([self.HOM], var_dict_row)) + len_b = bi.shape[1] else: - # bi can be a scipy sparse matrix, len_b = bi.shape[1] n_params = self.get_dim_P(param_subset=param_subset) @@ -443,30 +450,34 @@ def get_reduced_a( else: return ai - def augment_using_zero_padding(self, ai, param_subset=None): + def augment_using_zero_padding( + self, ai: np.ndarray, param_subset: list | dict | None = None + ) -> np.ndarray: n_parameters = self.get_dim_P(param_subset) return np.hstack([ai, np.zeros((n_parameters - 1) * len(ai))]) - def augment_using_parameters(self, x, param_subset=None): + def augment_using_parameters( + self, x: np.ndarray, param_subset: dict | None = None + ) -> np.ndarray: p = self.get_p(param_subset) return np.kron(p, x) - def compute_Ai(self, templates, var_dict, param_dict): + def compute_Ai( + self, + templates: list, + var_dict: dict, + param_dict: dict, + ) -> list: """ Take all elements from the list of templates and apply them to the given pair of var_list and param_list. """ - - A_list = [] + A_list: list = [] for k, template in enumerate(templates): assert template.b_ is not None assert isinstance(template, Constraint) - # First, we find the current parameters, so that we can factor - # them into b and compute a from it. p_here = self.get_p(param_subset=param_dict) - # We need to partition the vector b into its subblocks - # so that we can compute a from it. X_dim = self.get_dim_X(template.mat_var_dict) assert self.get_dim_X(var_dict) == X_dim p_dim = self.get_dim_p(template.mat_param_dict) @@ -492,13 +503,8 @@ def compute_Ai(self, templates, var_dict, param_dict): assert isinstance(a_test, np.ndarray) np.testing.assert_allclose(a, a_test) - # Get a symmetric matrix where the upper and lower parts have been filled with a, - # and applying the correction to the diagonal. - # Note that we do not set var_dict because otherwise A would already - # be the zero-padded large matrix. A = self.get_mat(a, sparse=True, correct=True) - # Get the corresponding PolyMatrix. A_poly, __ = PolyMatrix.init_from_sparse( A, var_dict=self.get_var_dict(var_dict), @@ -508,22 +514,19 @@ def compute_Ai(self, templates, var_dict, param_dict): A_list.append(A_poly) return A_list - def get_constraint_rank(self, A_list_poly, b_known=None, output_sorted=False): - """Find the number of independent constraints when they are of the form A_i @ x = 0. - - For bm lifting, X is nxd, thus we check the linear independence of vec(A_i @ X). - - :param A_list_poly: list of constraints matrices - - :return: rank (int) and sorted matrices (list, if output_sorted=True) where the first rank matrices correspond to - linearly independent matrices. - """ + def get_constraint_rank( + self, + A_list_poly: list, + b_known: list | None = None, + output_sorted: bool = False, + ) -> int | tuple[int, list]: + """Find the number of independent constraints when they are of the form A_i @ x = 0.""" x = self.get_x() current_rank = 0 - independent_indices = [] - dependent_indices = [] - if self.level == "bm": # type: ignore + independent_indices: list[int] = [] + dependent_indices: list[int] = [] + if getattr(self, "level", None) == "bm": basis_incremental = np.zeros((x.size + 1, 0)) else: basis_incremental = np.zeros((x.size, 0)) @@ -532,7 +535,7 @@ def get_constraint_rank(self, A_list_poly, b_known=None, output_sorted=False): new_candidate = (Ai.get_matrix(self.var_dict) @ x).reshape((-1, 1)) else: new_candidate = (Ai @ x).flatten()[:, None] - if self.level == "bm": # type: ignore + if getattr(self, "level", None) == "bm": assert b_known is not None new_candidate = np.vstack([new_candidate, b_known[i]]) basis_candidate = np.hstack([basis_incremental, new_candidate]) @@ -550,13 +553,17 @@ def get_constraint_rank(self, A_list_poly, b_known=None, output_sorted=False): ] return current_rank, A_list_sorted - def apply_template(self, bi_poly, n_parameters=None, verbose=False): + def apply_template( + self, + bi_poly: PolyMatrix, + n_parameters: int | None = None, + verbose: bool = False, + ) -> list[PolyMatrix]: if n_parameters is None: n_parameters = len(self.parameters) - new_poly_rows = [] - # find the number of variables that this constraint touches. - unique_idx = set() + new_poly_rows: list[PolyMatrix] = [] + unique_idx: set[int] = set() for key in bi_poly.variable_dict_j: param, var_keys = key.split("-") vars = var_keys.split(".") @@ -573,11 +580,9 @@ def apply_template(self, bi_poly, n_parameters=None, verbose=False): raise ValueError("unexpected triple dependencies!") variable_indices = self.get_variable_indices(self.var_dict) - # if z_0 is in this constraint, repeat the constraint for each landmark. for idx in itertools.combinations(variable_indices, len(unique_idx)): new_poly_row = PolyMatrix(symmetric=False) for key in bi_poly.variable_dict_j: - # need intermediate variables cause otherwise z_0 -> z_1 -> z_2 etc. can happen. key_ij = key for from_, to_ in zip(unique_idx, idx): key_ij = key_ij.replace(f"x_{from_}", f"xi_{to_}") @@ -603,22 +608,25 @@ def apply_template(self, bi_poly, n_parameters=None, verbose=False): raise IndexError( "something went wrong in augment_basis_list" ) - except ValueError as e: + except ValueError: pass new_poly_row[self.HOM, key_ij] = bi_poly[self.HOM, key] new_poly_rows.append(new_poly_row) return new_poly_rows def apply_templates( - self, templates, starting_index=0, var_dict=None, all_pairs=None - ): - + self, + templates: list, + starting_index: int = 0, + var_dict: dict | None = None, + all_pairs: bool | None = None, + ) -> list: if all_pairs is None: all_pairs = self.ALL_PAIRS if var_dict is None: var_dict = self.var_dict - new_constraints = [] + new_constraints: list = [] index = starting_index for template in templates: constraints = self.apply_template(template.polyrow_b_) @@ -641,21 +649,23 @@ def apply_templates( remove_dependent_constraints(new_constraints) return new_constraints - def get_vec_around_gt(self, delta: float = 0): - """Sample around ground truth. - :param delta: sample from gt + std(delta) (set to 0 to start from gt.) - """ + def get_vec_around_gt(self, delta: float = 0) -> np.ndarray: + """Sample around ground truth.""" return self.theta + np.random.normal(size=self.theta.shape, scale=delta) def test_constraints( - self, A_list, b_list=None, errors: str = "raise", n_seeds: int = 3 - ): + self, + A_list: list, + b_list: list | None = None, + errors: str = "raise", + n_seeds: int = 3, + ) -> tuple[float, set[int]]: """ :param A_list: can be either list of sparse matrices, or poly matrices :param errors: "raise" or "print" detected violations. """ max_violation = -np.inf - j_bad = set() + j_bad: set[int] = set() for j, A in enumerate(A_list): if isinstance(A, PolyMatrix): @@ -692,7 +702,7 @@ def test_constraints( raise ValueError(errors) return max_violation, j_bad - def get_A0(self, var_subset=None) -> tuple[list, list]: + def get_A0(self, var_subset: Iterable | None = None) -> tuple[list, list]: if var_subset is not None: var_dict = {k: self.var_dict[k] for k in var_subset} else: @@ -701,7 +711,7 @@ def get_A0(self, var_subset=None) -> tuple[list, list]: A0[self.HOM, self.HOM] = 1.0 return [A0.get_matrix(var_dict)], [1.0] - def sample_parameters_landmarks(self, landmarks: np.ndarray): + def sample_parameters_landmarks(self, landmarks: np.ndarray) -> dict: """Used by RobustPoseLifter, RangeOnlyLocLifter: the default way of adding landmarks to parameters.""" parameters = {self.HOM: 1.0} n_landmarks = landmarks.shape[0] @@ -713,19 +723,19 @@ def sample_parameters_landmarks(self, landmarks: np.ndarray): if self.param_level == "p": parameters[f"p_{i}"] = landmarks[i] elif self.param_level == "ppT": + quadratic_terms = get_vec( + np.outer(landmarks[i], landmarks[i]), + correct=False, + sparse=False, + ) parameters[f"p_{i}"] = np.hstack( # type: ignore - [ - landmarks[i], - get_vec( # type: ignore - np.outer(landmarks[i], landmarks[i]), - correct=False, - sparse=False, - ), - ] + [landmarks[i], np.asarray(quadratic_terms)] ) return parameters - def get_error(self, theta_hat, error_type="mse", *args, **kwargs) -> float: + def get_error( + self, theta_hat: np.ndarray, error_type: str = "mse", *args, **kwargs + ) -> float: if error_type in ("mse", "MSE"): return float(np.mean((theta_hat - self.theta) ** 2)) elif error_type in ("rmse", "RMSE"): @@ -733,7 +743,7 @@ def get_error(self, theta_hat, error_type="mse", *args, **kwargs) -> float: else: raise ValueError(f"unknown error type {error_type}") - def generate_random_setup(self): + def generate_random_setup(self) -> None: if self.parameters is None: self.parameters = self.sample_parameters() if self.theta is None: @@ -743,15 +753,11 @@ def get_x( self, theta: np.ndarray | None = None, var_subset: Iterable | None = None, + parameters: dict | None = None, **kwargs, ) -> np.ndarray: """ Get the lifted state vector x. - - :param theta: if given, use this theta instead of the ground truth one. - :param var_subset: list of parameter keys to use. If None, use all. - - :returns: lifted vector x """ if theta is None: theta = self.theta @@ -760,7 +766,7 @@ def get_x( assert theta is not None - x_data = [] + x_data: list[float] = [] for key in var_subset: if key == self.HOM: x_data.append(1.0) @@ -772,8 +778,10 @@ def get_x( return np.hstack(x_data) def get_p( - self, parameters: dict | None = None, param_subset: dict | list | None = None - ): + self, + parameters: dict | None = None, + param_subset: dict | list | None = None, + ) -> np.ndarray: if parameters is None: parameters = self.parameters if param_subset is None: @@ -781,7 +789,7 @@ def get_p( assert isinstance(parameters, dict) - p_data = [] + p_data: list[float] = [] for key in param_subset: if key == self.HOM: p_data.append(1.0) diff --git a/popcor/examples/example_lifter.py b/popcor/examples/example_lifter.py index 5daaea9..e7337d8 100644 --- a/popcor/examples/example_lifter.py +++ b/popcor/examples/example_lifter.py @@ -1,3 +1,8 @@ +"""Example module for implementing a custom StateLifter. + +This module provides an ExampleLifter class that demonstrates the minimal implementation required to use AutoTight. +""" + from collections.abc import Iterable import numpy as np @@ -16,49 +21,57 @@ class ExampleLifter(StateLifter): You can take a look at the :ref:`Examples` for inspiration. """ - # choose your homogenization variable here - HOM = "h" - - # chose the "lifting" levels when going up in the sparse Lasserre's hierarchy. - LEVELS = ["no"] + HOM: str = "h" + LEVELS: list[str] = ["no"] - def __init__(self, param_level="no"): - # you can choose if you want to use parameters. Otherwise remove param_level or set to "no" + def __init__(self, param_level: str = "no") -> None: super().__init__(param_level=param_level) @property - def var_dict(self): + def var_dict(self) -> dict[str, int]: var_dict = {self.HOM: 1} return var_dict @property - def param_dict(self): + def param_dict(self) -> dict[str, int]: param_dict = {self.HOM: 1} return param_dict def get_x( self, theta: np.ndarray | None = None, - var_subset: Iterable | None = None, + var_subset: dict | list | None = None, + parameters: np.ndarray | dict | None = None, **kwargs ) -> np.ndarray: + """Returns the lifted variable vector x for the given subset.""" if theta is None: theta = self.theta if var_subset is None: var_subset = self.var_dict - x_data = [] + x_data: list[float] = [] for key in var_subset: if key == self.HOM: x_data.append(1.0) + + # possibly use the parameter, if x depends on it. + if isinstance(parameters, dict): + pi = parameters[key] + else: + pi = parameters + # fill in the rest of x according to var_subset. # elif "x" in key: # elif "z" in key: + assert len(x_data) == self.get_dim_x(var_subset) return np.array(x_data) def sample_parameters(self, theta: np.ndarray | None = None) -> dict: + """Samples parameters for the lifter.""" return {} def sample_theta(self) -> np.ndarray: + """Samples theta for the lifter.""" return np.array([]) diff --git a/popcor/examples/mono_lifter.py b/popcor/examples/mono_lifter.py index a134671..8be6d02 100644 --- a/popcor/examples/mono_lifter.py +++ b/popcor/examples/mono_lifter.py @@ -2,6 +2,7 @@ from copy import deepcopy +import autograd.numpy as anp import numpy as np from poly_matrix.poly_matrix import PolyMatrix @@ -45,18 +46,10 @@ def h_list(self, t: np.ndarray, *args, **kwargs) -> list: - tan(a/2)*t3 >= sqrt(t1**2 + t2**2) """ default = super().h_list(t, *args, **kwargs) - try: - import autograd.numpy as anp - - return default + [ - anp.sum(t[:-1] ** 2) - anp.tan(FOV / 2) ** 2 * t[-1] ** 2, # type: ignore - -t[-1], - ] - except ModuleNotFoundError: - return default + [ - np.sum(t[:-1] ** 2) - np.tan(FOV / 2) ** 2 * t[-1] ** 2, - -t[-1], - ] + return default + [ + anp.sum(t[:-1] ** 2) - anp.tan(FOV / 2) ** 2 * t[-1] ** 2, # type: ignore + -t[-1], + ] def get_random_position(self, *args, **kwargs) -> np.ndarray: pc_cw = np.random.rand(self.d) * 0.1 @@ -101,10 +94,14 @@ def residual_sq( ui: np.ndarray, *args, **kwargs, - ) -> float: + ) -> float | anp.numpy_boxes.ArrayBox: W = np.eye(self.d) - np.outer(ui, ui) term = self.term_in_norm(R, t, pi, ui, *args, **kwargs) - res = float(term.T @ W @ term) + if isinstance(term, np.ndarray): + res = float(term.T @ W @ term) + else: + res = term.T @ W @ term + if NORMALIZE: res /= (self.n_landmarks * self.d) ** 2 return res diff --git a/popcor/examples/ro_nsq_lifter.py b/popcor/examples/ro_nsq_lifter.py index 9daa0d7..4dbcb2b 100644 --- a/popcor/examples/ro_nsq_lifter.py +++ b/popcor/examples/ro_nsq_lifter.py @@ -180,7 +180,7 @@ def get_x( self, theta: np.ndarray | None = None, parameters: np.ndarray | dict | None = None, - var_subset: list[str] | dict | None = None, + var_subset: list[str] | dict[str, int] | None = None, ) -> np.ndarray: """Return the lifted variable vector x for the given theta.""" if var_subset is None: diff --git a/popcor/examples/ro_sq_lifter.py b/popcor/examples/ro_sq_lifter.py index 4cb49ab..e101c1b 100644 --- a/popcor/examples/ro_sq_lifter.py +++ b/popcor/examples/ro_sq_lifter.py @@ -353,14 +353,14 @@ def get_grad( sub_idx_x = self.get_sub_idx_x(sub_idx) return 2 * J.T[:, sub_idx_x] @ Q[sub_idx_x, :][:, sub_idx_x] @ x[sub_idx_x] # type: ignore - def get_J(self, t: np.ndarray, y: np.ndarray) -> sp.csr_array: - J = sp.csr_array( + def get_J(self, t: np.ndarray, y: np.ndarray) -> sp.csr_matrix: + J = sp.csr_matrix( (np.ones(self.N), (range(1, self.N + 1), range(self.N))), shape=(self.N + 1, self.N), ) J_lift = self.get_J_lifting(t) J = sp.vstack([J, J_lift]) - assert isinstance(J, sp.csr_array) + assert isinstance(J, sp.csr_matrix) return J def get_hess(self, t: np.ndarray, y: np.ndarray) -> np.ndarray: diff --git a/popcor/examples/wahba_lifter.py b/popcor/examples/wahba_lifter.py index 661a7a4..4d6fd2a 100644 --- a/popcor/examples/wahba_lifter.py +++ b/popcor/examples/wahba_lifter.py @@ -1,5 +1,6 @@ """WahbaLifter class for robust pose registration with point-to-point measurements.""" +import autograd.numpy as anp import numpy as np import scipy.sparse as sp from poly_matrix import PolyMatrix @@ -63,11 +64,14 @@ def term_in_norm( def residual_sq( self, R: np.ndarray, t: np.ndarray, pi: np.ndarray, ui: np.ndarray - ) -> float: + ) -> float | anp.numpy_boxes.ArrayBox: """Computes the squared residual for a landmark measurement.""" W = np.eye(self.d) term = self.term_in_norm(R, t, pi, ui) - res_sq = float(term.T @ W @ term) + if isinstance(term, np.ndarray): + res_sq = float(term.T @ W @ term) + else: + res_sq = term.T @ W @ term if NORMALIZE: return res_sq / (self.n_landmarks * self.d) ** 2 return res_sq diff --git a/popcor/utils/common.py b/popcor/utils/common.py index 8318339..28aef64 100644 --- a/popcor/utils/common.py +++ b/popcor/utils/common.py @@ -1,25 +1,29 @@ +"""Common utilities for symmetric matrix vectorization and sparsity handling.""" + import itertools +from typing import Any, Dict, List, Sequence, Tuple import numpy as np import scipy.sparse as sp -def upper_triangular(p): - """Given vector, get the half kronecker product.""" +def upper_triangular(p: np.ndarray) -> np.ndarray: + """Return the vectorized upper-triangular part of the outer product of p.""" return np.outer(p, p)[np.triu_indices(len(p))] -def diag_indices(n): - """Given the half kronecker product, return diagonal elements""" +def diag_indices(n: int) -> np.ndarray: + """Return indices of diagonal elements in the vectorized upper-triangular matrix.""" z = np.empty((n, n)) z[np.triu_indices(n)] = range(int(n * (n + 1) / 2)) return np.diag(z).astype(int) -def get_aggregate_sparsity(matrix_list_sparse): - agg_ii = [] - agg_jj = [] - for i, A_sparse in enumerate(matrix_list_sparse): +def get_aggregate_sparsity(matrix_list_sparse: Sequence[sp.spmatrix]) -> sp.csr_matrix: + """Aggregate sparsity pattern from a list of sparse matrices.""" + agg_ii: List[int] = [] + agg_jj: List[int] = [] + for A_sparse in matrix_list_sparse: assert isinstance(A_sparse, sp.spmatrix) ii, jj = A_sparse.nonzero() # type: ignore agg_ii += list(ii) @@ -27,12 +31,12 @@ def get_aggregate_sparsity(matrix_list_sparse): return sp.csr_matrix(([1.0] * len(agg_ii), (agg_ii, agg_jj)), A_sparse.shape) -def unravel_multi_index_triu(flat_indices, shape): - """Equivalent of np.multi_index_triu, but using only the upper-triangular part of matrix.""" - i_upper = [] - j_upper = [] - - # for 4 x 4, this would give [4, 7, 9, 11] +def unravel_multi_index_triu( + flat_indices: Sequence[int], shape: Tuple[int, int] +) -> Tuple[np.ndarray, np.ndarray]: + """Convert flat indices to (i, j) indices for upper-triangular part of a matrix.""" + i_upper: List[int] = [] + j_upper: List[int] = [] cutoffs = np.cumsum(list(range(1, shape[0] + 1))[::-1]) for idx in flat_indices: i = np.where(idx < cutoffs)[0][0] @@ -45,31 +49,33 @@ def unravel_multi_index_triu(flat_indices, shape): return np.array(i_upper), np.array(j_upper) -def ravel_multi_index_triu(index_tuple, shape): - """Equivalent of np.multi_index_triu, but using only the upper-triangular part of matrix.""" +def ravel_multi_index_triu( + index_tuple: Tuple[np.ndarray, np.ndarray], shape: Tuple[int, int] +) -> List[int]: + """Convert (i, j) indices to flat indices for upper-triangular part of a matrix.""" ii, jj = index_tuple - triu_mask = jj >= ii i_upper = ii[triu_mask] j_upper = jj[triu_mask] - flat_indices = [] + flat_indices: List[int] = [] for i, j in zip(i_upper, j_upper): - # for i == 0: idx = j - # for i == 1: idx = shape[0] + j - # for i == 2: idx = shape[0] + shape[0]-1 + j idx = np.sum(range(shape[0] - i, shape[0])) + j flat_indices.append(idx) return flat_indices -def create_symmetric(vec, eps_sparse, correct=False, sparse=False): - """Create a symmetric matrix from the vectorized elements of the upper half""" +def create_symmetric( + vec: np.ndarray | sp.csr_matrix | sp.csc_matrix, + eps_sparse: float, + correct: bool = False, + sparse: bool = False, +) -> np.ndarray | sp.csr_array: + """Create a symmetric matrix from the vectorized upper-triangular elements.""" - def get_dim_x(len_vec): + def get_dim_x(len_vec: int) -> int: return int(0.5 * (-1 + np.sqrt(1 + 8 * len_vec))) - try: - # vec is dense + if isinstance(vec, np.ndarray): len_vec = len(vec) dim_x = get_dim_x(len_vec) triu = np.triu_indices(n=dim_x) @@ -77,16 +83,14 @@ def get_dim_x(len_vec): triu_i_nnz = triu[0][mask] triu_j_nnz = triu[1][mask] vec_nnz = vec[mask] - except Exception: - # vec is sparse + else: len_vec = vec.shape[1] dim_x = get_dim_x(len_vec) vec.data[np.abs(vec.data) < eps_sparse] = 0 vec.eliminate_zeros() - ii, jj = vec.nonzero() # vec is 1 x jj + ii, jj = vec.nonzero() triu_i_nnz, triu_j_nnz = unravel_multi_index_triu(jj, (dim_x, dim_x)) vec_nnz = np.array(vec[ii, jj]).flatten() - # assert dim_x == self.get_dim_x(var_dict) if sparse: offdiag = triu_i_nnz != triu_j_nnz @@ -95,7 +99,6 @@ def get_dim_x(len_vec): triu_j = triu_j_nnz[offdiag] diag_i = triu_i_nnz[diag] if correct: - # divide off-diagonal elements by sqrt(2) vec_nnz_off = vec_nnz[offdiag] / np.sqrt(2) else: vec_nnz_off = vec_nnz[offdiag] @@ -110,12 +113,9 @@ def get_dim_x(len_vec): ) else: Ai = np.zeros((dim_x, dim_x)) - if correct: - # divide all elements by sqrt(2) Ai[triu_i_nnz, triu_j_nnz] = vec_nnz / np.sqrt(2) Ai[triu_j_nnz, triu_i_nnz] = vec_nnz / np.sqrt(2) - # undo operation for diagonal Ai[range(dim_x), range(dim_x)] *= np.sqrt(2) else: Ai[triu_i_nnz, triu_j_nnz] = vec_nnz @@ -123,12 +123,12 @@ def get_dim_x(len_vec): return Ai -def get_vec(mat, correct=True, sparse=False) -> np.ndarray | sp.csr_matrix | None: - """Convert NxN Symmetric matrix to (N+1)N/2 vectorized version that preserves inner product. - - :param mat: (spmatrix or ndarray) symmetric matrix - :return: ndarray - """ +def get_vec( + mat: np.ndarray | sp.csc_matrix | sp.csr_matrix, + correct: bool = True, + sparse: bool = False, +) -> np.ndarray | sp.csr_matrix: + """Convert NxN symmetric matrix to vectorized upper-triangular form preserving inner product.""" from copy import deepcopy mat = deepcopy(mat) @@ -146,9 +146,8 @@ def get_vec(mat, correct=True, sparse=False) -> np.ndarray | sp.csr_matrix | Non ii, jj = mat.nonzero() if len(ii) == 0: # got an empty matrix -- this can happen depending on the parameter values. - return None + raise ValueError triu_mask = jj >= ii - flat_indices = ravel_multi_index_triu([ii[triu_mask], jj[triu_mask]], mat.shape) # type: ignore data = np.array(mat[ii[triu_mask], jj[triu_mask]]).flatten() # type: ignore vec_size = int(mat.shape[0] * (mat.shape[0] + 1) / 2) # type: ignore @@ -159,12 +158,12 @@ def get_vec(mat, correct=True, sparse=False) -> np.ndarray | sp.csr_matrix | Non return np.array(mat[np.triu_indices(n=mat.shape[0])]).flatten() # type: ignore -def get_labels(p, zi, zj, var_dict): - labels = [] +def get_labels(p: str, zi: str, zj: str, var_dict: Dict[str, int]) -> List[str]: + """Generate labels for matrix/vector elements based on variable sizes.""" + labels: List[str] = [] size_i = var_dict[zi] size_j = var_dict[zj] if zi == zj: - # only upper diagonal for i == j key_pairs = itertools.combinations_with_replacement(range(size_i), 2) else: key_pairs = itertools.product(range(size_i), range(size_j)) diff --git a/popcor/utils/constraint.py b/popcor/utils/constraint.py index a54d13a..1bb15b8 100644 --- a/popcor/utils/constraint.py +++ b/popcor/utils/constraint.py @@ -1,3 +1,9 @@ +""" +Module for handling constraints and their representations, including conversion between different formats and plotting. +""" + +from typing import Any, Iterable + import numpy as np import scipy.sparse as sp from poly_matrix.poly_matrix import PolyMatrix @@ -6,10 +12,13 @@ from popcor.utils.plotting_tools import plot_basis -def remove_dependent_constraints(constraints, verbose=False): +def remove_dependent_constraints(constraints: list, verbose: bool = False) -> None: + """ + Removes linearly dependent constraints from the list. + """ from cert_tools.linalg_tools import find_dependent_columns - # find which constraints are lin. dep. + # Find which constraints are linearly dependent. A_vec = sp.vstack( [constraint.a_full_ for constraint in constraints], format="coo" ).T @@ -17,12 +26,17 @@ def remove_dependent_constraints(constraints, verbose=False): bad_idx = find_dependent_columns(A_vec, verbose=verbose) if len(bad_idx): np.testing.assert_allclose(bad_idx, sorted(bad_idx)) - # important: by changing the order we + # Remove constraints in reverse order to avoid index shift. for idx in sorted(bad_idx)[::-1]: del constraints[idx] -def generate_poly_matrix(constraints, factor_out_parameters=False, lifter=None): +def generate_poly_matrix( + constraints: list, factor_out_parameters: bool = False, lifter: Any = None +) -> PolyMatrix: + """ + Generates a PolyMatrix from a list of constraints. + """ plot_rows = [] plot_row_labels = [] j = -1 @@ -30,7 +44,7 @@ def generate_poly_matrix(constraints, factor_out_parameters=False, lifter=None): for constraint in constraints: mat_vars = constraint.mat_var_dict i = constraint.index - if factor_out_parameters: # use a and not b. + if factor_out_parameters: if constraint.polyrow_a_ is not None: plot_rows.append(constraint.polyrow_a_) else: @@ -65,7 +79,6 @@ def generate_poly_matrix(constraints, factor_out_parameters=False, lifter=None): if mat_vars != old_mat_vars: j += 1 plot_row_labels.append(f"{j}:b{i}") - # plot_row_labels.append(f"{j}{mat_vars}:b{i}") old_mat_vars = mat_vars else: plot_row_labels.append(f"{j}:b{i}") @@ -77,14 +90,20 @@ def generate_poly_matrix(constraints, factor_out_parameters=False, lifter=None): def plot_poly_matrix( - poly_matrix, variables_j=None, variables_i=None, simplify=True, hom="h" -): + poly_matrix: PolyMatrix, + variables_j: Iterable | None = None, + variables_i: Iterable | None = None, + simplify: bool = True, + hom: str = "h", +) -> tuple: + """ + Plots the PolyMatrix using the provided variables. + """ if variables_i is None: variables_i = poly_matrix.variable_dict_i if variables_j is None: variables_j = poly_matrix.variable_dict_j - # plot the templates stored in poly_matrix. fig, ax = plot_basis( poly_matrix, variables_j=variables_j, @@ -100,14 +119,14 @@ def plot_poly_matrix( new_xticks = [] for lbl in ax.get_xticklabels(): lbl = lbl.get_text() - if "_" in lbl: # avoid double subscript + if "_" in lbl: new_lbl = f"${lbl.replace(f'{hom}.', '').replace(':', '^')}$" else: new_lbl = f"${lbl.replace(f'{hom}.', '').replace(':', '_')}$" new_xticks.append(new_lbl) ax.set_xticklabels(new_xticks, fontsize=7) - # plot a red vertical line at each new block of parameters. + # Plot a red vertical line at each new block of parameters. params = [v.split("-")[0] for v in variables_j] old_param = params[0] for i, p in enumerate(params): @@ -125,25 +144,25 @@ def plot_poly_matrix( class Constraint(object): """ - This class serves the main purpose of not recomputing representations of constraints more than once. + Stores representations of constraints and avoids redundant computations. """ def __init__( self, - index=0, - polyrow_a=None, - polyrow_b=None, - A_poly=None, - A_sparse=None, - b=None, - a=None, - a_full=None, - b_full=None, - mat_var_dict=None, - mat_param_dict=None, - known=False, - template_idx=0, - rhs=0, + index: int = 0, + polyrow_a: PolyMatrix | None = None, + polyrow_b: PolyMatrix | None = None, + A_poly: PolyMatrix | None = None, + A_sparse: sp.csr_matrix | sp.csc_matrix | None = None, + b: np.ndarray | sp.csr_matrix | sp.csc_matrix | None = None, + a: np.ndarray | sp.csr_matrix | sp.csc_matrix | None = None, + a_full: np.ndarray | sp.csr_matrix | sp.csc_matrix | None = None, + b_full: np.ndarray | sp.csr_matrix | sp.csc_matrix | None = None, + mat_var_dict: dict | list | None = None, + mat_param_dict: dict | list | None = None, + known: bool = False, + template_idx: int = 0, + rhs: int | float = 0, ): self.index = index self.mat_var_dict = mat_var_dict @@ -161,31 +180,35 @@ def __init__( self.known = known self.template_idx = template_idx - # list of applied constraints derived from this constraint. - self.applied_list = [] - # what this cosntraint equals to: Ax = rhs + # List of applied constraints derived from this constraint. + self.applied_list: list = [] + # What this constraint equals to: Ax = rhs self.rhs = rhs @staticmethod - # @profile def init_from_b( index: int, b: np.ndarray, - mat_var_dict: dict, + mat_var_dict: dict | list, lifter=None, - mat_param_dict: dict | None = None, + mat_param_dict: dict | list | None = None, convert_to_polyrow: bool = True, known: bool = True, template_idx: int = 0, - ): - a = None - A_sparse = None - a_full = None + ) -> "Constraint | None": + """ + Initializes a Constraint from vector b. + """ + a: np.ndarray | sp.csr_matrix | sp.csc_matrix | None = None + A_sparse: sp.csr_matrix | sp.csc_matrix | None = None + a_full: np.ndarray | sp.csr_matrix | sp.csc_matrix | None = None if lifter is not None: a = lifter.get_reduced_a( b, var_subset=mat_var_dict, param_subset=mat_param_dict, sparse=True ) - A_sparse = lifter.get_mat(a, var_dict=mat_var_dict, sparse=True) + out = lifter.get_mat(a, var_dict=mat_var_dict, sparse=True) + assert isinstance(out, (sp.csr_matrix, sp.csc_matrix)) + A_sparse = out a_full = get_vec(A_sparse, sparse=True) if a_full is None: return None @@ -216,15 +239,18 @@ def init_from_b( @staticmethod def init_from_A_poly( - lifter, + lifter: Any, A_poly: PolyMatrix, mat_var_dict: dict, known: bool = False, index: int = 0, template_idx: int = 0, - compute_polyrow_b=False, - rhs=0, - ): + compute_polyrow_b: bool = False, + rhs: int | float = 0, + ) -> "Constraint": + """ + Initializes a Constraint from a PolyMatrix A_poly. + """ Ai_sparse_small = A_poly.get_matrix(variables=mat_var_dict) ai = get_vec(Ai_sparse_small, correct=True) bi = lifter.augment_using_zero_padding(ai) @@ -234,6 +260,7 @@ def init_from_A_poly( polyrow_b = None polyrow_a = lifter.convert_a_to_polyrow(ai, mat_var_dict) Ai_sparse = A_poly.get_matrix(variables=lifter.var_dict) + assert isinstance(Ai_sparse, (sp.csr_matrix, sp.csc_matrix)) return Constraint( a=ai, polyrow_a=polyrow_a, @@ -251,12 +278,15 @@ def init_from_A_poly( @staticmethod def init_from_polyrow_b( polyrow_b: PolyMatrix, - lifter, + lifter: Any, index: int = 0, known: bool = False, template_idx: int = 0, mat_var_dict: dict | None = None, - ): + ) -> "Constraint": + """ + Initializes a Constraint from a PolyMatrix polyrow_b. + """ if mat_var_dict is None: mat_var_dict = lifter.var_dict A_poly = lifter.convert_polyrow_to_Apoly(polyrow_b) @@ -274,16 +304,16 @@ def init_from_polyrow_b( mat_var_dict=mat_var_dict, ) - def scale_to_new_lifter(self, lifter): + def scale_to_new_lifter(self, lifter: Any) -> "Constraint": + """ + Scales the constraint to a new lifter's variable dictionary. + """ if self.known: assert self.A_poly_ is not None - # known matrices are stored in origin variables, not unrolled form self.A_sparse_ = self.A_poly_.get_matrix(lifter.var_dict) self.a_full_ = get_vec(self.A_sparse_, sparse=True) - else: assert self.A_poly_ is not None - # known matrices are stored in origin variables, not unrolled form target_dict_unroll = lifter.get_var_dict(unroll_keys=True) self.A_sparse_ = self.A_poly_.get_matrix(target_dict_unroll) self.a_full_ = get_vec(self.A_sparse_, sparse=True) diff --git a/popcor/utils/plotting_tools.py b/popcor/utils/plotting_tools.py index 070c5d0..94f1633 100644 --- a/popcor/utils/plotting_tools.py +++ b/popcor/utils/plotting_tools.py @@ -196,7 +196,7 @@ def initialize_discrete_cbar(values): def plot_basis( basis_poly: PolyMatrix, - variables_j: dict, + variables_j: Iterable, variables_i: Iterable | None = None, fname_root: str = "", discrete: bool = True, diff --git a/tests/test_parameters.py b/tests/test_parameters.py index c9899b3..e55f6cf 100644 --- a/tests/test_parameters.py +++ b/tests/test_parameters.py @@ -12,7 +12,7 @@ def test_canonical_operations(): n_landmarks = 1 # z_0 and z_1 lifter = Stereo2DLifter(n_landmarks=n_landmarks, param_level="p", level="no") - var_subset = "x" + var_subset = ["x"] # fmt: off Ai_sub = np.array( [ @@ -31,6 +31,7 @@ def test_canonical_operations(): ai = get_vec(Ai_sub) # zero-pad to emulate an augmented basis vector + assert isinstance(ai, np.ndarray) bi = lifter.augment_using_zero_padding(ai) ai_test = lifter.get_reduced_a(bi, var_subset=var_subset) assert isinstance(ai_test, np.ndarray) @@ -38,7 +39,7 @@ def test_canonical_operations(): np.testing.assert_allclose(ai, ai_test) # will return a 9 x 9 matrix with zero padding. - Ai_test = lifter.get_mat(ai_test, var_dict={var_subset: 6}) + Ai_test = lifter.get_mat(ai_test, var_dict={"x": 6}) assert isinstance(Ai, sp.csc_matrix) assert isinstance(Ai_test, sp.csc_matrix) np.testing.assert_allclose(Ai.toarray(), Ai_test.toarray()) @@ -60,6 +61,7 @@ def test_b_to_a(): x_sub = get_vec(np.outer(x, x)) # test that bi_sub @ x_aug holds (including parameters in x_aug) + assert isinstance(x_sub, np.ndarray) x_sub_aug = lifter.augment_using_parameters(x_sub) assert len(x_sub_aug) == len(bi_sub)