electrostatic_simulation/Simulation.py at main · sjmoon00/electrostatic_simulation · 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
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
from __future__ import annotations
from typing import Tuple
from typing import TYPE_CHECKING

if TYPE_CHECKING:
    from queue import Queue
    from Charge import ChargeDist

import numpy as np
from PIL import Image

from Calc import Calc

import plots.potential_color as potential_color
import plots.potential_contour as potential_contour


class Simulation:
    def __init__(self, conf: dict):
        self.charges: Tuple[ChargeDist] = conf['charges']
        self.phy_rect: Tuple[float, float, float, float] = conf['phy_rect']
        self.full_phy_rect: Tuple[float, float, float, float] = (0, 0, 0, 0)
        self.mpp: float = conf['mpp']
        self.down_sampling: int = conf['down_sampling']
        self.plots: dict = conf['plots']
        self.device = conf['device']

        self.__init_data()
        self.calc: Calc = Calc(charges=self.charges,
                               phy_rect=self.phy_rect,
                               data=self.data,
                               ref_point=conf['ref_point'],
                               device=self.device)

    def __init_data(self) -> None:
        """
        Adjust physical rect to quantized scale by mpp and down sampling value
        Init data array
        Init image array

        :return: None
        """

        # Adjust physical rect and get data shape
        sizes = Simulation.get_adjusted_size(self.phy_rect, self.down_sampling, self.mpp)
        self.phy_rect = sizes['adj_phy_rect']
        self.full_phy_rect = sizes['full_adj_phy_rect']
        n_data_row, n_data_col = sizes['data_shape']

        # Init data array
        self.data = np.zeros(sizes['data_shape'], dtype=np.float32)

        # Init image array
        n_img_row = n_data_row * self.down_sampling
        n_img_col = n_data_col * self.down_sampling
        self.img = np.zeros((n_img_row, n_img_col, 3), dtype=np.uint8)
        self.img.fill(255)

    @staticmethod
    def get_adjusted_size(phy_rect: Tuple[float, float, float, float], down_sampling: int, mpp: float) -> dict:
        """
        Get
        physical rect which adjusted to quantized scale by mpp and down sampling value
        and shape of data

        n_data_row, n_data_col : shape of data
        adj_phy_rect : boundary of sampling points in physical domain
        full_adj_phy_rect : boundary of full plot in physical domain
        => if down_sampling is 1, phy_rect and full_phy_rect will be same

        :return: dictionary of data shape and adjusted physical rects
        """

        # meter per data (1 data value : data of mpd x mpd size square area)
        mpd = down_sampling * mpp

        # Center point of physical rect (won't be changed)
        phy_cntr = (
            (phy_rect[0] + phy_rect[2]) / 2,
            (phy_rect[1] + phy_rect[3]) / 2
        )

        phy_width = abs((phy_rect[2] - phy_rect[0]) / 2)
        phy_height = abs((phy_rect[3] - phy_rect[1]) / 2)

        # Size of data array
        half_n_data_col = int(np.ceil((phy_width / 2) / mpd))
        half_n_data_row = int(np.ceil((phy_height / 2) / mpd))
        n_data_col = half_n_data_col * 2
        n_data_row = half_n_data_row * 2

        # Adjusted physical rect
        half_adj_phy_width = (half_n_data_col - 0.5) * mpd
        half_adj_phy_height = (half_n_data_row - 0.5) * mpd

        adj_phy_rect = (
            phy_cntr[0] - half_adj_phy_width,
            phy_cntr[1] + half_adj_phy_height,
            phy_cntr[0] + half_adj_phy_width,
            phy_cntr[1] - half_adj_phy_height
        )

        # Adjusted full physical rect
        phy_margin = (down_sampling / 2 - 0.5) * mpp
        full_adj_phy_rect = (
            adj_phy_rect[0] - phy_margin,
            adj_phy_rect[1] + phy_margin,
            adj_phy_rect[2] + phy_margin,
            adj_phy_rect[3] - phy_margin
        )

        return {
            'data_shape': (n_data_row, n_data_col),
            'image_size': (n_data_row * down_sampling, n_data_col * down_sampling),
            'adj_phy_rect': adj_phy_rect,
            'full_adj_phy_rect': full_adj_phy_rect
        }

    def run(self, progress_q: Queue, out_path: str = 'result.png', verbose: bool = True) -> None:
        """
        Run the simulation and save result image at out_path

        :param progress_q: queue for sending progress info to gui thread
                           e.g. {'task': 'calc', 'progress': 25.1, 'el_tm': 1.03, 'est_tm': 11.7}
                                task: current task (e.g. 'calc', 'potential_color', 'potential_contour')
                                progress: percentage of progress (out of 100)
                                el_tm: elapsed time,
                                est_tm: estimated time to done
        :param out_path: path of an output file
        :param verbose: print running information
        :return: None
        """

        # Fill data array
        self.calc.do(progress_q, verbose=verbose)

        # Pile up plots
        plot_module = {
            'potential_color': potential_color,
            'potential_contour': potential_contour
        }
        for plot, conf in self.plots.items():
            if self.device == 'cpu':
                plot_module[plot].cpu(self.img, self.data, conf)
            elif self.device == 'gpu':
                plot_module[plot].gpu(self.img, self.data, conf)

        # Save image
        res = Image.fromarray(self.img)
        res.save(out_path)