ms_thesis/fd_comparison.py at main · probabilis/ms_thesis · 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
import matplotlib.pyplot as plt
from dataclasses import asdict, replace

from utils.env_utils import PATHS, print_bars, get_args, plotting_style, plotting_schematic, log_data
from utils.pattern_formation import initialize_u0_random

from params.opt_params import labyrinth_data_params, sim_config, get_DataParameters, get_SimulationParamters
from params.opt_params import pgd_sim_params as ngd_sim_params

from optimization.gd_nesterov import gradient_descent_nesterov


if __name__ == "__main__":

    SINGLE_COMPARISON = True
    GAMMA_SWEEP = False

    plotting_style()
    FOLDER_PATH = PATHS.PATH_COMPARISON

    LIVE_PLOT = False
    DATA_LOG = False
    N = 64
    labyrinth_data_params = replace(labyrinth_data_params, N = N, gamma = 0.002, epsilon = 0.01)

    #gridsize, N, th, epsilon, gamma = get_DataParameters(labyrinth_data_params)
    u0 = initialize_u0_random(N, REAL = True)

    print_bars()
    print(labyrinth_data_params)
    print(ngd_sim_params)
    print_bars()

    _types = ["Finite Differences | PBC", "Finite Differences | Von Neumann", "Spectral"]
    PBC_ls = [True, False, True]

    if SINGLE_COMPARISON:
        fig, axs = plt.subplots(len(_types), 2, figsize = (12,12))
        for ii in range(0, len(_types)):
            if ii == 2:
                sim_config = replace(sim_config, LAPLACE_SPECTRAL = True)

            u, e = gradient_descent_nesterov(u0, LIVE_PLOT, DATA_LOG, FOLDER_PATH, **asdict(labyrinth_data_params),**asdict(ngd_sim_params), **asdict(sim_config), PBC = PBC_ls[ii])

            axs[ii, 0].imshow(u)
            axs[ii, 0].set_title(f"LaPlace: {_types[ii]}")

            axs[ii, 1].loglog(e)
            axs[ii, 1].hlines(e[-1],1,len(e), label = f"E[-1] = {e[-1]:.3f}", color = "black", linestyle = ":")
            axs[ii, 1].set_title("Energy evolution")


            axs[ii, 0].set_box_aspect(1)
            axs[ii, 0].axes.get_xaxis().set_ticks([])
            axs[ii, 0].axes.get_yaxis().set_ticks([])

            axs[ii, 1].grid(color = "gray")
            axs[ii, 1].legend()
            print_bars()

        plt.savefig(FOLDER_PATH / "laplace_evaluation_comparison.png", dpi = 300)
        plt.show()


    ngd_sim_params = replace(ngd_sim_params, tau = 0.01)

    if GAMMA_SWEEP:
        gamma_ls = [1/80, 1/100, 1/200, 1/500, 1/1000, 1/2000, 1/4000]

        fig, axs = plt.subplots( len(gamma_ls), len(_types), figsize = (12,6) )

        sim_config = replace(sim_config, LAPLACE_SPECTRAL = False)
        u0 = initialize_u0_random(N, REAL = True)


        for tt, type in enumerate(_types):
            if tt == 2:
                sim_config = replace(sim_config, LAPLACE_SPECTRAL = True)

            for ii, _gamma in enumerate(gamma_ls):
                labyrinth_data_params = replace(labyrinth_data_params, gamma = _gamma)
                u, e = gradient_descent_nesterov(u0, LIVE_PLOT, DATA_LOG, FOLDER_PATH, **asdict(labyrinth_data_params),**asdict(ngd_sim_params), **asdict(sim_config), PBC=PBC_ls[tt])

                axs[ii, tt].imshow(u.cpu().numpy(), extent=(0,1,0,1))

                axs[ii, tt].set_box_aspect(1)
                axs[ii, tt].axes.get_xaxis().set_ticks([])
                axs[ii, tt].axes.get_yaxis().set_ticks([])

                axs[ii, tt].set_title(f"$\\gamma = {_gamma:.5f}$")


        fig.canvas.draw()  # ensures positions are compute

        for kk in range(0, len(_types)):

            bbox = axs[0, kk].get_position()
            x_center = 0.5 * (bbox.x0 + bbox.x1)
            y_top = bbox.y1 + 0.02
            title = f"LaPlace: {_types[kk]}"
            fig.text(x_center, y_top, title, ha="center", va = "bottom", fontsize=12)

        plt.savefig(FOLDER_PATH / "laplace_evaluation_gamma_sweep.png", dpi = 300)
        plt.show()