jonasmaltebecker/vae_drilling

1"""2Alle transforms sind grundsätzlich auf batches bezogen!3Vae transforms sind invertierbar4"""5import pickle6from dataclasses import dataclass7from functools import partial, reduce, wraps8 9import numpy as np10import torch11 12# Allgemeine Funktionen -------------------------------------------------------------13# Transformations in Pytorch sind am einfachsten.14 15 16def load(p):17    with open(p, "rb") as stream:18        return pickle.load(stream)19 20 21def save(obj, p):22    with open(p, "wb") as stream:23        pickle.dump(obj, stream)24 25 26def sequential_function(*functions):27    return lambda x: reduce(lambda res, func: func(res), functions, x)28 29 30def np_sample(func):31    rtn = sequential_function(32        lambda x: torch.from_numpy(x).float(),33        lambda x: torch.unsqueeze(x, 0),34        func,35        lambda x: x[0].numpy(),36    )37    return rtn38 39 40# Inverseabvle41class SequentialInversable(torch.nn.Sequential):42    def __init__(self, *functions):43        super().__init__(*functions)44 45        self.inv_funcs = [f.inv for f in functions]46        self.inv_funcs.reverse()47 48    # def forward(self, x):49    #     return sequential_function(*self.functions)(x)50 51    def inv(self, x):52        return sequential_function(*self.inv_funcs)(x)53 54 55class LatentSelector(torch.nn.Module):56    """Verarbeitet Tensoren und numpy arrays"""57 58    def __init__(self, ldim: int, selectdim: int):59        super().__init__()60        self.ldim = ldim61        self.selectdim = selectdim62 63    def forward(self, x: torch.Tensor):64        return x[:, : self.selectdim]65 66    def inv(self, x: torch.Tensor):67        rtn = torch.cat(68            [x, torch.zeros((x.shape[0], self.ldim - x.shape[1]), device=x.device)],69            dim=1,70        )71        return rtn72 73 74class MinMaxScaler(torch.nn.Module):75    #! Bei mehreren Signalen vorsicht mit dem Broadcasting.76    def __init__(77        self,78        _min: torch.Tensor,79        _max: torch.Tensor,80        min_norm: float = 0.0,81        max_norm: float = 1.0,82    ):83        super().__init__()84        self._min = _min85        self._max = _max86        self.min_norm = min_norm87        self.max_norm = max_norm88 89    def forward(self, ts):90        """None, no_signals"""91        std = (ts - self._min) / (self._max - self._min)92        rtn = std * (self.max_norm - self.min_norm) + self.min_norm93        return rtn94 95    def inv(self, ts):96        std = (ts - self.min_norm) / (self.max_norm - self.min_norm)97        rtn = std * (self._max - self._min) + self._min98        return rtn99 100    @classmethod101    def from_array(cls, arr: torch.Tensor):102        _min = torch.min(arr, axis=0).values103        _max = torch.max(arr, axis=0).values104 105        return cls(_min, _max)106 107 108class LatentSorter(torch.nn.Module):109    def __init__(self, kl_dict: dict):110        super().__init__()111        self.kl_dict = kl_dict112 113    def forward(self, latent):114        """115        unsorted -> sorted116        latent: (None, latent_dim)117        """118        return latent[:, list(self.kl_dict.keys())]119 120    def inv(self, latent):121        keys = np.array(list(self.kl_dict.keys()))122        return latent[:, torch.from_numpy(keys.argsort())]123 124    @property125    def names(self):126        rtn = ["{} KL{:.2f}".format(k, v) for k, v in self.kl_dict.items()]127        return rtn128 129 130def apply_along_axis(function, x, axis: int = 0):131    return torch.stack([function(x_i) for x_i in torch.unbind(x, dim=axis)], dim=axis)132 133 134# Eingangsshapes bleiben wie sie sind!135class SumField(torch.nn.Module):136    """137    time series: [idx, time_step, signal]138    image: [idx, signal, time_step, time_step]139    """140 141    def forward(self, ts: torch.Tensor):142        """ts2img"""143 144        samples = ts.shape[0]145        time = ts.shape[1]146        channels = ts.shape[2]147 148        ts = torch.swapaxes(ts, 1, 2)  # Zeitachse ans Ende149        ts = torch.reshape(150            ts, (samples * channels, time)151        )  # Zusammenfassen von Channel + idx152        #! TODO: Schleife besser lösen153        rtn = apply_along_axis(self._mtf_forward, ts, 0)154        rtn = torch.reshape(rtn, (samples, channels, time, time))155 156        return rtn157 158    def inv(self, img: torch.Tensor):159        """img2ts"""160        rtn = torch.diagonal(img, dim1=2, dim2=3)161        rtn = torch.swapaxes(rtn, 1, 2)  # Channel und Zeitachse tauschen162 163        return rtn164 165    @staticmethod166    def _mtf_forward(ts):167        """For one dimensional time series ts"""168        return torch.add(*torch.meshgrid(ts, ts, indexing="ij")) / 2169