simple-factorgraph-example/plot_utils.py at main · tussedrotten/simple-factorgraph-example · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
import warnings
from dataclasses import dataclass
from typing import Any

import gtsam
import numpy as np
import visgeom as vg
from shapely import geometry as sg, ops as so
import matplotlib.pyplot as plt
import matplotlib


def plot_result(poses, landmarks):
    matplotlib.use('qtagg')

    plt.figure()
    ax = plt.axes()

    num_draws = 1000

    # Plot pose marginals.
    num_poses = len(poses)
    pose_cmap = plt.get_cmap('autumn')
    for dist, color in zip(poses, pose_cmap([number / num_poses for number in range(num_poses)])):
        mean_translation = dist.mean.translation()

        # Plot covariance from manifold.
        plot_se2_covariance_on_manifold_gtsam(ax, dist, fill_alpha=0.5, fill_color=color)

        # Plot random points from distribution.
        random_points = np.random.multivariate_normal(np.zeros(3), dist.covariance, num_draws).T
        random_translations = np.zeros([2, num_draws])
        for i in range(0, num_draws):
            random_translations[:, i] = dist.mean.compose(gtsam.Pose2.Expmap(random_points[:, i])).translation()

        ax.plot(random_translations[0, :], random_translations[1, :],
                marker='.', markeredgewidth=0, color=color, linestyle='', alpha=0.2)

        # Plot mean.
        ax.plot(mean_translation[0], mean_translation[1], marker='+', color='k', linestyle='')

    # Plot landmark marginals.
    num_landmarks = len(landmarks)
    landmark_cmap = plt.get_cmap('summer')
    for dist, color in zip(landmarks, landmark_cmap([number / num_landmarks for number in range(num_landmarks)])):
        plot_ellipse(ax, dist, fill_alpha=0.5, fill_color=color)

        random_points = np.random.multivariate_normal(dist.mean, dist.covariance, num_draws).T
        ax.plot(random_points[0, :], random_points[1, :],
                marker=".", markeredgewidth=0, color=color, linestyle='', alpha=0.2)

        ax.plot(dist.mean[0], dist.mean[1], marker='+', color='k', linestyle='')

    ax.set_aspect('equal', adjustable='box')

    plt.show()


def plot_se2_covariance_on_manifold_gtsam(ax, dist,
                                          n=50, chi2_val=11.345, right_perturbation=True,
                                          fill_alpha=0., fill_color='lightsteelblue'):
    u, s, _ = np.linalg.svd(dist.covariance)
    scale = np.sqrt(chi2_val * s)

    x, y, z = vg.utils.generate_ellipsoid(n, pose=(u, np.zeros([3, 1])), scale=scale)

    tangent_points = np.vstack((x.flatten(), y.flatten(), z.flatten()))

    num_samples = tangent_points.shape[1]
    transl_points = np.zeros([num_samples, 2])

    if right_perturbation:
        for i in range(num_samples):
            transl_points[i, :] = dist.mean.compose(gtsam.Pose2.Expmap(tangent_points[:, i])).translation()
    else:
        for i in range(num_samples):
            transl_points[i, :] = gtsam.Pose2.Expmap(tangent_points[:, i]).compose(dist.mean).translation()

    p_grid = np.reshape(transl_points, [(n + 1), (n + 1), 2])
    polygons = extract_polygon_slices(p_grid)
    union = so.unary_union(polygons)
    if not union.geom_type == 'Polygon':
        warnings.warn("Could not find a closed boundary", RuntimeWarning)
        return

    ax.fill(*union.exterior.xy, alpha=fill_alpha, facecolor=fill_color)
    ax.plot(*union.exterior.xy, color=fill_color)


def extract_polygon_slices(grid_2d):
    p_a = grid_2d[:-1, :-1]
    p_b = grid_2d[:-1, 1:]
    p_c = grid_2d[1:, 1:]
    p_d = grid_2d[1:, :-1]

    quads = np.concatenate((p_a, p_b, p_c, p_d), axis=2)

    m, n, _ = grid_2d.shape
    quads = quads.reshape(((m-1) * (n-1), 4, 2))

    return [sg.Polygon(t).buffer(0.0001, cap_style=2, join_style=2) for t in quads]


def plot_ellipse(ax, dist, n=50, chi2_val=9.21, fill_alpha=0., fill_color='lightsteelblue'):
    u, s, _ = np.linalg.svd(dist.covariance)
    scale = np.sqrt(chi2_val * s)

    theta = np.linspace(0, 2*np.pi, n+1)
    x = np.cos(theta)
    y = np.sin(theta)

    R = u
    t = np.reshape(dist.mean, [2, 1])
    circle_points = (R * scale) @ np.vstack((x.flatten(), y.flatten())) + t

    ax.fill(circle_points[0, :], circle_points[1, :], alpha=fill_alpha, facecolor=fill_color)
    ax.plot(circle_points[0, :], circle_points[1, :], color=fill_color)


@dataclass
class MultivariateNormalParameters:
    mean: Any
    covariance: np.ndarray