procgne/Plonk

1import torch2from torch import nn3import numpy as np4 5 6class NormGPS(nn.Module):7    def __init__(self, input_key="gps", output_key="x_0", normalize=True):8        super().__init__()9        self.input_key = input_key10        self.output_key = output_key11        self.normalize = normalize12        if self.normalize:13            self.register_buffer(14                "gps_normalize", 1 / torch.Tensor([np.pi * 0.5, np.pi]).unsqueeze(0)15            )16 17    def forward(self, batch):18        """Normalize latitude longtitude radians to -1, 1."""  # not used currently19        x = batch[self.input_key]20        if self.normalize:21            x = x * self.gps_normalize22        batch[self.output_key] = x23        return batch24 25class GPStoCartesian(nn.Module):26    def __init__(self, input_key="gps", output_key="x_0"):27        super().__init__()28        self.input_key = input_key29        self.output_key = output_key30 31    def forward(self, batch):32        """Project latitude longtitude radians to 3D coordinates."""33        x = batch[self.input_key]34        lat, lon = x[:, 0], x[:, 1]35        x = torch.stack([lat.cos() * lon.cos(), lat.cos() * lon.sin(), lat.sin()], dim=-1)36        batch[self.output_key] = x37        return batch38 39class PrecomputedPreconditioning:40    def __init__(41        self,42        input_key="emb",43        output_key="emb",44    ):45        self.input_key = input_key46        self.output_key = output_key47 48    def __call__(self, batch, device=None):49        batch[self.output_key] = batch[self.input_key]50        return batch51