-
-
Notifications
You must be signed in to change notification settings - Fork 17
/
Copy pathmultimodal.py
311 lines (263 loc) · 13.7 KB
/
multimodal.py
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
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
"""The default composite model architecture for PVNet"""
from collections import OrderedDict
from typing import Optional
import torch
from ocf_datapipes.batch import BatchKey, NWPBatchKey
from torch import nn
import pvnet
from pvnet.models.base_model import BaseModel
from pvnet.models.multimodal.basic_blocks import ImageEmbedding
from pvnet.models.multimodal.encoders.basic_blocks import AbstractNWPSatelliteEncoder
from pvnet.models.multimodal.linear_networks.basic_blocks import AbstractLinearNetwork
from pvnet.models.multimodal.site_encoders.basic_blocks import AbstractPVSitesEncoder
from pvnet.optimizers import AbstractOptimizer
class Model(BaseModel):
"""Neural network which combines information from different sources
Architecture is roughly as follows:
- Satellite data, if included, is put through an encoder which transforms it from 4D, with time,
channel, height, and width dimensions to become a 1D feature vector.
- NWP, if included, is put through a similar encoder.
- PV site-level data, if included, is put through an encoder which transforms it from 2D, with
time and system-ID dimensions, to become a 1D feature vector.
- The satellite features*, NWP features*, PV site-level features*, GSP ID embedding*, and sun
paramters* are concatenated into a 1D feature vector and passed through another neural
network to combine them and produce a forecast.
* if included
"""
name = "conv3d_sat_nwp"
def __init__(
self,
output_network: AbstractLinearNetwork,
output_quantiles: Optional[list[float]] = None,
nwp_encoders_dict: Optional[dict[AbstractNWPSatelliteEncoder]] = None,
sat_encoder: Optional[AbstractNWPSatelliteEncoder] = None,
pv_encoder: Optional[AbstractPVSitesEncoder] = None,
wind_encoder: Optional[AbstractPVSitesEncoder] = None, # TODO Change to SensorEncoder
add_image_embedding_channel: bool = False,
include_gsp_yield_history: bool = True,
include_sun: bool = True,
include_gsp: bool = True,
embedding_dim: Optional[int] = 16,
forecast_minutes: int = 30,
history_minutes: int = 60,
sat_history_minutes: Optional[int] = None,
min_sat_delay_minutes: Optional[int] = 30,
nwp_forecast_minutes: Optional[int] = None,
nwp_history_minutes: Optional[int] = None,
pv_history_minutes: Optional[int] = None,
wind_history_minutes: Optional[int] = None,
optimizer: AbstractOptimizer = pvnet.optimizers.Adam(),
target_key: str = "gsp",
interval_minutes: int = 30,
):
"""Neural network which combines information from different sources.
Notes:
In the args, where it says a module `m` is partially instantiated, it means that a
normal pytorch module will be returned by running `mod = m(**kwargs)`. In this library,
this partial instantiation is generally achieved using partial instantiation via hydra.
However, the arg is still valid as long as `m(**kwargs)` returns a valid pytorch module
- for example if `m` is a regular function.
Args:
output_network: A partially instatiated pytorch Module class used to combine the 1D
features to produce the forecast.
output_quantiles: A list of float (0.0, 1.0) quantiles to predict values for. If set to
None the output is a single value.
nwp_encoders_dict: A dictionary of partially instatiated pytorch Module class used to
encode the NWP data from 4D into an 1D feature vector from different sources.
sat_encoder: A partially instatiated pytorch Module class used to encode the satellite
data from 4D into an 1D feature vector.
pv_encoder: A partially instatiated pytorch Module class used to encode the site-level
PV data from 2D into an 1D feature vector.
add_image_embedding_channel: Add a channel to the NWP and satellite data with the
embedding of the GSP ID.
include_gsp_yield_history: Include GSP yield data.
include_sun: Include sun azimuth and altitude data.
embedding_dim: Number of embedding dimensions to use for GSP ID. Not included if set to
`None`.
forecast_minutes: The amount of minutes that should be forecasted.
history_minutes: The default amount of historical minutes that are used.
sat_history_minutes: Length of recent observations used for satellite inputs. Defaults
to `history_minutes` if not provided.
min_sat_delay_minutes: Minimum delay with respect to t0 of the latest available
satellite image.
nwp_forecast_minutes: Period of future NWP forecast data used as input. Defaults to
`forecast_minutes` if not provided.
nwp_history_minutes: Period of historical NWP forecast used as input. Defaults to
`history_minutes` if not provided.
pv_history_minutes: Length of recent site-level PV data data used as input. Defaults to
`history_minutes` if not provided.
optimizer: Optimizer factory function used for network.
target_key: The key of the target variable in the batch.
interval_minutes: The interval between each sample of the target data
"""
self.include_gsp_yield_history = include_gsp_yield_history
self.include_sat = sat_encoder is not None
self.include_nwp = nwp_encoders_dict is not None and len(nwp_encoders_dict) != 0
self.include_pv = pv_encoder is not None
self.include_sun = include_sun
self.include_gsp = include_gsp
self.include_wind = wind_encoder is not None
self.embedding_dim = embedding_dim
self.add_image_embedding_channel = add_image_embedding_channel
self.target_key_name = target_key
self.interval_minutes = interval_minutes
super().__init__(
history_minutes=history_minutes,
forecast_minutes=forecast_minutes,
optimizer=optimizer,
output_quantiles=output_quantiles,
target_key=target_key,
interval_minutes=interval_minutes,
)
# Number of features expected by the output_network
# Add to this as network pices are constructed
fusion_input_features = 0
if self.include_sat:
# Param checks
assert sat_history_minutes is not None
assert nwp_forecast_minutes is not None
self.sat_sequence_len = (sat_history_minutes - min_sat_delay_minutes) // 5 + 1
self.sat_encoder = sat_encoder(
sequence_length=self.sat_sequence_len,
in_channels=sat_encoder.keywords["in_channels"] + add_image_embedding_channel,
)
if add_image_embedding_channel:
self.sat_embed = ImageEmbedding(
318, self.sat_sequence_len, self.sat_encoder.image_size_pixels
)
# Update num features
fusion_input_features += self.sat_encoder.out_features
if self.include_nwp:
# Param checks
assert nwp_forecast_minutes is not None
assert nwp_history_minutes is not None
# For each NWP encoder the forecast and history minutes must be set
assert set(nwp_encoders_dict.keys()) == set(nwp_forecast_minutes.keys())
assert set(nwp_encoders_dict.keys()) == set(nwp_history_minutes.keys())
self.nwp_encoders_dict = torch.nn.ModuleDict()
if add_image_embedding_channel:
self.nwp_embed_dict = torch.nn.ModuleDict()
for nwp_source in nwp_encoders_dict.keys():
nwp_sequence_len = (
nwp_history_minutes[nwp_source] // 60
+ nwp_forecast_minutes[nwp_source] // 60
+ 1
)
self.nwp_encoders_dict[nwp_source] = nwp_encoders_dict[nwp_source](
sequence_length=nwp_sequence_len,
in_channels=(
nwp_encoders_dict[nwp_source].keywords["in_channels"]
+ add_image_embedding_channel
),
)
if add_image_embedding_channel:
self.nwp_embed_dict[nwp_source] = ImageEmbedding(
318, nwp_sequence_len, self.nwp_encoders_dict[nwp_source].image_size_pixels
)
# Update num features
fusion_input_features += self.nwp_encoders_dict[nwp_source].out_features
if self.include_pv:
assert pv_history_minutes is not None
self.pv_encoder = pv_encoder(
sequence_length=pv_history_minutes // 15,
)
# Update num features
fusion_input_features += self.pv_encoder.out_features
if self.include_wind:
if wind_history_minutes is None:
wind_history_minutes = history_minutes
self.wind_encoder = wind_encoder(sequence_length=self.history_len_30)
# Update num features
fusion_input_features += self.wind_encoder.out_features
if self.embedding_dim:
self.embed = nn.Embedding(num_embeddings=318, embedding_dim=embedding_dim)
# Update num features
fusion_input_features += embedding_dim
if self.include_sun:
# the minus 12 is bit of hard coded smudge for pvnet
self.sun_fc1 = nn.Linear(
in_features=2 * (self.forecast_len_30 + self.history_len_30 + 1),
out_features=16,
)
# Update num features
fusion_input_features += 16
if include_gsp_yield_history:
# Update num features
fusion_input_features += self.history_len_30
self.output_network = output_network(
in_features=fusion_input_features,
out_features=self.num_output_features,
)
self.save_hyperparameters()
def forward(self, x):
"""Run model forward"""
modes = OrderedDict()
# ******************* Satellite imagery *************************
if self.include_sat:
# Shape: batch_size, seq_length, channel, height, width
sat_data = x[BatchKey.satellite_actual][:, : self.sat_sequence_len]
sat_data = torch.swapaxes(sat_data, 1, 2).float() # switch time and channels
if self.add_image_embedding_channel:
id = x[BatchKey.wind_id][:, 0].int()
sat_data = self.sat_embed(sat_data, id)
modes["sat"] = self.sat_encoder(sat_data)
# *********************** NWP Data ************************************
if self.include_nwp:
# Loop through potentially many NMPs
for nwp_source in self.nwp_encoders_dict:
# shape: batch_size, seq_len, n_chans, height, width
nwp_data = x[BatchKey.nwp][nwp_source][NWPBatchKey.nwp].float()
nwp_data = torch.swapaxes(nwp_data, 1, 2) # switch time and channels
if self.add_image_embedding_channel:
id = x[BatchKey.gsp_id][:, 0].int()
nwp_data = self.nwp_embed_dict[nwp_source](nwp_data, id)
modes[f"nwp/{nwp_source}"] = self.nwp_encoders_dict[nwp_source](nwp_data)
# *********************** PV Data *************************************
# Add site-level PV yield
if self.include_pv:
if self.target_key_name != "pv":
modes["pv"] = self.pv_encoder(x)
else:
# Target is PV, so only take the history
# Copy batch
x_tmp = x.copy()
x_tmp[BatchKey.pv] = x_tmp[BatchKey.pv][:, : self.history_len_30]
modes["pv"] = self.pv_encoder(x_tmp)
# *********************** GSP Data ************************************
# add gsp yield history
if self.include_gsp_yield_history:
gsp_history = x[BatchKey.gsp][:, : self.history_len_30].float()
gsp_history = gsp_history.reshape(gsp_history.shape[0], -1)
modes["gsp"] = gsp_history
# ********************** Embedding of GSP ID ********************
if self.embedding_dim:
if self.target_key_name == "wind":
id = x[BatchKey.wind_id][:, 0].int()
elif self.target_key_name == "pv":
id = x[BatchKey.pv_id][:, 0].int()
else:
id = x[BatchKey.gsp_id][:, 0].int()
id_embedding = self.embed(id)
modes["id"] = id_embedding
# *********************** Sensor Data ************************************
# add sensor yield history
if self.include_wind:
if self.target_key_name != "wind":
modes["wind"] = self.wind_encoder(x)
else:
# Have to be its own Batch format
x_tmp = x.copy()
x_tmp[BatchKey.wind] = x_tmp[BatchKey.wind][:, : self.history_len_30]
# This needs to be a Batch as input
modes["wind"] = self.wind_encoder(x_tmp)
if self.include_sun:
sun = torch.cat(
(x[BatchKey.gsp_solar_azimuth], x[BatchKey.gsp_solar_elevation]), dim=1
).float()
sun = self.sun_fc1(sun)
modes["sun"] = sun
out = self.output_network(modes)
if self.use_quantile_regression:
# Shape: batch_size, seq_length * num_quantiles
out = out.reshape(out.shape[0], self.forecast_len_30, len(self.output_quantiles))
return out