metadit/datapipe.py at main · JessePrince/metadit · 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
"""
Dataset pipeline for model training
"""

import torch.utils

from torch.utils.data.dataset import Dataset
from utils import load_json
import math
import torch.nn.functional as F
from scipy import io


class BaseDataset(Dataset):
    def __init__(
        self,
        data_path: str
    ):
        super().__init__()
        self.data = io.loadmat(data_path)

    def create_3d_grid(self, geometry: dict):
        pass

    def __len__(self):
        return len(self.data["parameter"])

    def __getitem__(self, index):
        pass


class FreeFormDataset(BaseDataset):
    def __init__(self, path):
        super().__init__(path)

    def create_3d_grid(self, geometry):
        # channel 0 for r_index
        # channel 1 for thickness
        # channel 2 for lattice_size
        grid = torch.zeros(3, 32, 32, dtype=torch.float32)
        # get 1/4 because of symmetry
        pattern = torch.tensor(geometry["pattern"][:32, :32], dtype=torch.float32)
        # normalization
        index = geometry["params"][2] / 5.0
        lattice_size = geometry["params"][0] / 3.0

        grid[0][pattern==1] = index
        grid[1][pattern==1] = geometry["params"][1]
        grid[2] = lattice_size

        return grid

    def __getitem__(self, index):
        # geometry
        geometry = {
            "pattern": self.data["pattern"][:,:,index], # np.array(64, 64)
            "params": self.data["parameter"][index] # [Lattice size, Thicknesses, Refractive index]
        }

        """
        metadata: lattice size 2.5 µm to 3 µm
        Thickness: 0.5um to 1um
        refractive index: 3.5 to 5
        """

        grid = self.create_3d_grid(geometry)
        # 301 freq points
        imag = torch.tensor(self.data["imag"][index], dtype=torch.float32)
        real = torch.tensor(self.data["real"][index], dtype=torch.float32)
        condition = torch.stack([real, imag], dim=0)

        data = {
            "inputs": grid,  # 3, 64, 64
            "condition": condition,  # 2, 301
            "labels": grid
        }

        return data


class SurrogateFreeFormDataset(BaseDataset):
    def __init__(self, path):
        super().__init__(path)

    def create_3d_grid(self, geometry):
        # channel 0 for r_index
        # channel 1 for thickness
        # channel 2 for lattice_size
        grid = torch.zeros(3, 64, 64, dtype=torch.float32)
        # get 1/4 because of symmetry
        pattern = torch.tensor(geometry["pattern"], dtype=torch.float32)
        # normalization
        index = geometry["params"][2] / 5.0
        lattice_size = geometry["params"][0] / 3.0

        grid[0][pattern==1] = index
        grid[1][pattern==1] = geometry["params"][1]
        grid[2] = lattice_size

        return grid

    def __getitem__(self, index):
        # geometry
        geometry = {
            "pattern": self.data["pattern"][:,:,index], # np.array(64, 64)
            "params": self.data["parameter"][index] # [Lattice size, Thicknesses, Refractive index]
        }

        """
        metadata: lattice size 2.5 µm to 3 µm
        Thickness: 0.5um to 1um
        refractive index: 3.5 to 5
        """

        grid = self.create_3d_grid(geometry)
        # 301 freq points
        imag = torch.tensor(self.data["imag"][index], dtype=torch.float32)
        real = torch.tensor(self.data["real"][index], dtype=torch.float32)
        condition = torch.stack([real, imag], dim=0)

        data = {
            "inputs": grid,  # 3, 64, 64
            "labels": condition
        }

        return data