COMAP_2025/StepDistribution.py at main · ErickWhiteDev/COMAP_2025 · 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 numpy as np
from scipy.stats import multivariate_normal
from scipy.integrate import dblquad
import matplotlib.pyplot as plt

class StepDistribution():
    def __init__(self, stair, pos='r', dir='up'):
        self.half_width = stair.xmax
        self.half_length = stair.ymax

        self.X = stair.X
        self.Y = stair.Y

        self.xres = stair.xres
        self.yres = stair.yres

        if pos == 'c':
            self.meanl = np.array([-0.09, 0])
            self.meanr = np.array([0.09, 0])
        elif pos == 'l':
            self.meanl = np.array([-self.half_width / 2 - 0.09, 0])
            self.meanr = np.array([-self.half_width / 2 + 0.09, 0])
        elif pos == 'r':
            self.meanl = np.array([self.half_width / 2 - 0.09, 0])
            self.meanr = np.array([self.half_width / 2 + 0.09, 0])

        if dir == 'down':
            deg = np.pi * 8 / 180
            rot = np.array([[np.cos(deg), -np.sin(deg)],[np.sin(deg),np.cos(deg)]])
            self.covr = rot @ np.array([[0.01, 0], [0, .035]]) @ rot.T
            deg = 2*np.pi - deg
            rot2= np.array([[np.cos(deg), -np.sin(deg)],[np.sin(deg),np.cos(deg)]])
            self.covl = rot2 @ np.array([[0.01, 0], [0, .035]]) @ rot2.T
        else:
            deg = np.pi - np.pi * 8/180
            rot = np.array([[np.cos(deg), -np.sin(deg)],[np.sin(deg),np.cos(deg)]])
            self.covr = rot @ np.array([[0.01, 0], [0, .035]]) @ rot.T
            deg = 2*np.pi - deg
            rot2 = np.array([[np.cos(deg), -np.sin(deg)],[np.sin(deg),np.cos(deg)]])
            self.covl = rot2 @ np.array([[0.01, 0], [0, .035]]) @ rot2.T

        self.rvl = multivariate_normal(self.meanl, self.covl)
        self.rvr = multivariate_normal(self.meanr, self.covr)

        f = lambda x, y: 0 if ((np.abs(x) >= self.half_width) | (np.abs(y) >= self.half_length)) else self.rvr.pdf([x, y]) + self.rvl.pdf([x, y])

        q, _ = dblquad(f, -self.half_length, self.half_length, -self.half_width, self.half_width)
        self.norm = 1 / q

        self.pdf = lambda x: self.norm * f(x[0], x[1])

    def _pdf(self, x, y):
        rvl = multivariate_normal(self.meanl, self.covl)
        rvr = multivariate_normal(self.meanr, self.covr)

        if type(x) != np.ndarray:
            x = np.array([x])
        if type(y) != np.ndarray:
            y = np.array([y])

        dim = x.shape

        x = x.flatten()
        y = y.flatten()

        coords = np.vstack((x, y)).T

        res_flattened = ~((np.abs(x) >= self.half_width) | (np.abs(y) >= self.half_length)) * (rvr.pdf(coords) + rvl.pdf(coords))

        return np.reshape(res_flattened, dim) * self.norm


    def find_pdf_max_bound(self):
        maxl = self.rvl.pdf(self.meanl)
        maxr = self.rvr.pdf(self.meanr)

        return (maxl + maxr) * self.norm

    def rejection_sampling(self, size=1, use_res=True):
        pdf_max = self.find_pdf_max_bound()

        q_dist = multivariate_normal((self.meanl + self.meanr) / 2, self.covl + self.covr)

        p = lambda x: self._pdf(x[:, 0], x[:, 1])
        q = lambda x: q_dist.pdf(x)

        x = self.X.flatten()
        y = self.Y.flatten()

        coords = np.vstack((x, y)).T

        k = max(p(coords) / q(coords))

        samples = np.empty(shape=(size, 2))

        x = k * q_dist.rvs(size=size)
        z = np.random.uniform(0, pdf_max, size=size)

        mask = z < self._pdf(x[:, 0], x[:, 1])

        samples = x[mask]

        if use_res:
            return np.array([self.xres * (samples[:, 0] + self.half_width) / (2 * self.half_width), self.yres * (samples[:, 1] + self.half_length) / (2 * self.half_length)]).astype(int)
        else:
            return samples

    def plot_dist(self):
        Z1 = self.rvr.pdf(np.dstack((self.X, self.Y)))
        Z2 = self.rvl.pdf(np.dstack((self.X, self.Y)))

        fig, ax = plt.subplots(3,1)
        ax[0].contourf(self.X, self.Y, Z1, cmap='viridis')
        ax[1].contourf(self.X, self.Y, Z2, cmap='viridis')
        ax[2].contourf(self.X, self.Y, self._pdf(self.X,self.Y), cmap='viridis')
        plt.show()

        plt.contourf(self.X, self.Y, self._pdf(self.X,self.Y), cmap='viridis')
        plt.colorbar()
        plt.xlabel('X')
        plt.ylabel('Y')
        plt.show()