biosphere/model/dataset.py at main · CosmosRedshift7/biosphere · 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
import pickle
import torch

from math import floor
from torch.utils.data import Dataset
from torch.nn.functional import conv1d
from icecream import ic


class BiosphereDataset(Dataset):
    def __init__(
        self,
        data_path: str,
        phase: str,
        dim_input: int = 96,
        dim_output: int = 96,
        shuffle_type: str = "random",
    ):
        self.data_path = data_path
        self.phase = phase
        self.dim_input = dim_input
        self.dim_output = dim_output
        self.shuffle_type = shuffle_type

        assert not dim_output > dim_input

        self.shift = self.dim_input - self.dim_output

        with open(f"{self.data_path}/sensors.p", "rb") as f:
            sensors = pickle.load(f)
        with open(f"{self.data_path}/control.p", "rb") as f:
            control = pickle.load(f)

        self.sensors = self.transform(sensors["data"])
        self.control = self.transform(control["data"])

        # self.sensors[72:, :] = torch.roll(
        #     self.sensors[72:, :], -self.dim_output, dims=-1
        # )

        assert not torch.isnan(self.sensors).any()
        assert not torch.isnan(self.control).any()

        self.name_sensors = sensors["name"]
        self.name_control = control["name"]

        self.f_date = sensors["f_date"]
        self.l_date = sensors["l_date"]

        num_inputs = self.sensors.shape[-1] // self.dim_input - 3
        self.len_train = floor(0.8 * num_inputs)
        self.len_val = num_inputs - self.len_train

        torch.manual_seed(0)
        if self.shuffle_type == "random":
            self.get_indices = torch.randperm(num_inputs)
            print("random")
        else:
            self.get_indices = torch.arange(num_inputs)
            # self.get_indices = torch.roll(self.get_indices, self.len_train, dims=-1)
            print("subsequent")

    @property
    def num_sensors(self):
        return self.sensors.shape[0]

    @property
    def num_controls(self):
        return self.control.shape[0]

    def transform(self, x):
        bs = x.shape[0]
        w = torch.tensor([-1, 1]).view(1, 1, -1)

        x = self.normalize(x)

        m = x.mean(-1, keepdim=True)
        s = x.std(-1, keepdim=True)

        x = torch.concat((m, x), dim=-1)

        mask = (x > m + 5 * s) + (x < m - 5 * s)
        v = conv1d(mask.to(torch.long).view(bs, 1, -1), w, padding="same").view(bs, -1)

        r = ((v == -1).nonzero() - (v == 1).nonzero())[:, 1]
        x[mask] = x[v == 1].repeat_interleave(r)
        x = x[:, 1:]
        # x = self.inv_normalize(x)
        return x

    def add_noise(self, x):
        return x + 5 * torch.rand_like(x)

    def normalize(self, x):
        self.x_min = x.min(-1, keepdim=True)[0]
        self.x_max = x.max(-1, keepdim=True)[0]
        self.x_max[self.x_max == 0] = 1

        x -= self.x_min
        x /= (self.x_max - self.x_min)
        return x

    def inv_normalize(self, x):
        x *= (self.x_max - self.x_min).to(x.device)
        x += self.x_min.to(x.device)
        return x

    def __len__(self):
        if self.phase == "train":
            return self.len_train
        else:
            return self.len_val

    def __getitem__(self, index):
        if self.phase == "train":
            index = self.get_indices[index]
        else:
            index = self.get_indices[index + self.len_train]
        # if not self.phase == "train":
        #     index = index + self.len_train

        # a = index * self.dim_input + torch.randint(self.shift, (1,)).item()
        a = index * self.dim_input
        b = a + self.dim_input
        c = b + self.dim_output

        sen = self.add_noise(self.sensors[..., a:b].clone())
        con = self.control[..., b:c]

        return sen, con