matthewplotlib/examples/lissajous.py at main · matomatical/matthewplotlib · 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
import tyro
import numpy as np
import sklearn.decomposition # pip install scikit-learn

import matthewplotlib as mp


DIMENSION = 10_000
NUM_STEPS = 1_000
SEED = 42


def main(save: str | None = None):
    """Brownian motion PCA visualisation."""
    # sample some high-dimensional brownian motion
    np.random.seed(SEED)
    traj = np.random.normal(size=(NUM_STEPS, DIMENSION)).cumsum(axis=0)

    # 3-dimensional PCA reduction
    proj = sklearn.decomposition.PCA(n_components=3).fit_transform(traj)
    time = np.linspace(.1, 1., NUM_STEPS)

    # construct visualisation
    plot = mp.vstack(
        mp.text("BROWNIAN MOTION"),
        mp.text("first two dimensions"),
        mp.scatter(
            (traj[:, (0,1)], (.5,.5,.5)),
            height=20,
            width=75,
        ),
        mp.text("principal components 1, 2, 3 over time"),
        mp.scatter(
            (time, proj[:,0], "red"),
            (time, proj[:,1], "green"),
            (time, proj[:,2], "blue"),
            width=75,
        ),
        mp.text("principal components 1, 2, 3 paired, increasing lightness over time"),
        mp.hstack(
            mp.text("PC1 v PC2")
            / mp.scatter((proj[:, (0,1)], mp.yellows(time)), width=25),
            mp.text("PC1 v PC3")
            / mp.scatter((proj[:, (0,2)], mp.magentas(time)), width=25),
            mp.text("PC2 v PC3")
            / mp.scatter((proj[:, (1,2)], mp.cyans(time)), width=25),
        ),
    )

    print(plot)
    if save:
        plot.saveimg(save)

if __name__ == "__main__":
    tyro.cli(main)