File size: 3,153 Bytes
bd3493c | 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 | import torch
import torch.nn as nn
class convDeepSet(nn.Module):
"""
ConvDeepSet used to translate ungridded observations to a gridded representation and
vice versa.
"""
def __init__(
self,
init_ls,
mode,
device,
density_channel=True,
step=0.25,
grid=False,
):
super().__init__()
self.init_ls = torch.nn.Parameter(torch.tensor([init_ls]))
self.grid = grid
self.step = step
self.density_channel = density_channel
self.mode = mode
self.init_ls.requires_grad = True
self.device = device
def compute_weights(self, x1, x2):
dists2 = self.pw_dists2(x1.unsqueeze(-1), x2.unsqueeze(-1))
d = torch.exp((-0.5 * dists2) / (self.init_ls.to(x1.device)) ** 2)
return d
def pw_dists2(self, a, b):
norms_a = torch.sum(a**2, axis=-1)[..., :, None]
norms_b = torch.sum(b**2, axis=-1)[..., None, :]
return norms_a + norms_b - 2 * torch.matmul(a, b.permute(0, 2, 1))
def forward(self, x_in, wt, x_out):
# Add a density channel
density_channel = torch.ones_like(wt[:, 0:1, ...])
density_channel[torch.isnan(wt[:, 0:1, ...])] = 0
wt = torch.cat([density_channel, wt], dim=1)
wt[torch.isnan(wt)] = 0
if self.mode == "OffToOn":
# Case 1: converting off-the-grid data to a gridded representation
in_lon_mask = ~torch.isnan(x_in[0])
in_lat_mask = ~torch.isnan(x_in[1])
x_in[0][~in_lon_mask] = 0
x_in[1][~in_lat_mask] = 0
ws = [self.compute_weights(xzi, xi) for xzi, xi in zip(x_in, x_out)]
ws[0] = ws[0] * in_lon_mask.unsqueeze(-1).int()
ws[1] = ws[1] * in_lat_mask.unsqueeze(-1).int()
ee = torch.einsum("...cw,...wx,...wy->...cxy", wt, ws[0], ws[1])
elif self.mode == "OnToOn":
# Case 2: converting between gridded representations
ws = [self.compute_weights(xzi, xi) for xzi, xi in zip(x_in, x_out)]
ee = torch.einsum("...cwh,...wx,...hy->...cxy", wt, ws[0], ws[1])
elif self.mode == "OnToOff":
# Case 3: converting a gridded representation to off-the-grid predictions
out_lon_mask = ~torch.isnan(x_out[0])
out_lat_mask = ~torch.isnan(x_out[1])
x_out[0][~out_lon_mask] = 0
x_out[1][~out_lat_mask] = 0
ws = [self.compute_weights(xzi, xi) for xzi, xi in zip(x_in, x_out)]
ws[0] = ws[0] * out_lon_mask.unsqueeze(-2).int()
ws[1] = ws[1] * out_lat_mask.unsqueeze(-2).int()
ee = torch.einsum("...cwh,...wx,...hx->...cx", wt, ws[0], ws[1])
if self.density_channel:
ee = torch.cat(
[
ee[:, 0:1, ...],
ee[:, 1:, ...] / torch.clamp(ee[:, 0:1, ...], min=1e-6, max=1e5),
],
dim=1,
)
return ee
else:
ee = ee[:, 1:, ...] / torch.clamp(ee[:, 0:1, ...], min=1e-6, max=1e5)
return ee
|