time4apint/old/graph_based_simplification.py at gaia · Parallel-in-Time/time4apint · 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
from PintGraph import PintGraph


class GraphBasedApproach(PintGraph):
    def __init__(self, cost_fine: float, cost_coarse: float, cost_copy: float = 0, cost_correction: float = 0,
                 simplify_graph=True, *args: object, **kwargs: object) -> None:
        """
        Constructor

        :param cost_fine: Cost of the fine propagator
        :param cost_coarse: Cost of the coarse propagator
        :param cost_copy: Cost of a copy operation
        :param cost_correction: Cost of a correction
        :param args:
        :param kwargs:
        """
        super().__init__(*args, **kwargs)
        self.cost_coarse = cost_coarse
        self.cost_fine = cost_fine
        self.cost_correction = cost_correction
        self.cost_copy = cost_copy
        self.sim_graph = simplify_graph

    def block_node_simple(self, k, n):
        if k >= self.iterations[n]:
            pred_k = self.iterations[n]
            pred_k_plus_1 = self.iterations[n]
        else:
            pred_k = k
            pred_k_plus_1 = k + 1
        self.add_node(op=f"F",
                      predecessors=[f'u_{pred_k}_{n}'],
                      set_values=[f'v_1_{k}_{n + 1}_1'],
                      cost=self.cost_fine,
                      point=n + 1,
                      pos=(n + 1 - .3, k + 1), )
        self.add_node(op=f"G",
                      predecessors=[f'u_{pred_k}_{n}'],
                      set_values=[f'v_1_{k}_{n + 1}_2'],
                      cost=self.cost_coarse,
                      point=n + 1,
                      pos=(n + 1 + .3, k + 1))
        self.add_node(op=f"G",
                      predecessors=[f'u_{pred_k_plus_1}_{n}'],
                      set_values=[f'v_2_{k + 1}_{n + 1}'],
                      cost=0,
                      point=n + 1,
                      pos=(n + 1 - .3, k + 1 + .5))
        self.add_node(op="+",
                      predecessors=[f'v_1_{k}_{n + 1}_1', f'v_1_{k}_{n + 1}_2', f'v_2_{k + 1}_{n + 1}'],
                      set_values=[f'u_{k + 1}_{n + 1}'],
                      cost=0,
                      point=n + 1,
                      pos=(n + 1 + .3, k + 1 + .5))

    def build_graph(self):
        for k in range(max(self.iterations)):
            for n in range(self.nt - 1):
                if k <= self.iterations[n]:
                    print(k, n)
                    self.block_node_simple(k=k, n=n)

    def boundary_condition(self):
        self.add_node(op="init_cond",
                      predecessors=['u_0'],
                      set_values=['u_0_0'],
                      cost=0,
                      point=0,
                      pos=(0, 0))
        for n in range(self.nt - 1):
            self.add_node(op='G',
                          predecessors=[f'u_0_{n}'],
                          set_values=[f'u_0_{n + 1}'],
                          cost=self.cost_coarse,
                          point=n + 1,
                          pos=(n + 1, 0))

    def compute(self):
        self.boundary_condition()
        self.build_graph()
        if self.sim_graph:
            self.simplify_graph()


parareal_model = GraphBasedApproach(cost_fine=4, cost_coarse=1, nt=4, iters=[0, 1, 2], conv_crit=1, simplify_graph=False)
parareal_model.compute()
parareal_model.plot_dag()

parareal_model = GraphBasedApproach(cost_fine=4, cost_coarse=1, nt=4, iters=[0, 1, 2], conv_crit=1, simplify_graph=True)
parareal_model.compute()
parareal_model.plot_dag()