reflectorch 1.0.0__py3-none-any.whl

reflectorch/ml/callbacks.py

@@ -0,0 +1,86 @@
+# -*- coding: utf-8 -*-
+#
+#
+# This source code is licensed under the GPL license found in the
+# LICENSE file in the root directory of this source tree.
+
+import torch
+
+import numpy as np
+
+from reflectorch.ml.basic_trainer import (
+    TrainerCallback,
+    Trainer,
+)
+from reflectorch.ml.utils import is_divisor
+
+__all__ = [
+    'SaveBestModel',
+    'LogLosses',
+]
+
+
+class SaveBestModel(TrainerCallback):
+    """Callback for periodically saving the best model weights
+
+    Args:
+        path (str): path for saving the model weights
+        freq (int, optional): frequency in iterations at which the current average loss is evaluated. Defaults to 50.
+        average (int, optional): number of recent iterations over which the average loss is computed. Defaults to 10.
+    """
+
+    def __init__(self, path: str, freq: int = 50, average: int = 10):
+        self.path = path
+        self.average = average
+        self._best_loss = np.inf
+        self.freq = freq
+
+    def end_batch(self, trainer: Trainer, batch_num: int) -> None:
+        """checks if the current average loss has improved from the previous save, if true the model is saved
+
+        Args:
+            trainer (Trainer): the trainer object
+            batch_num (int): the current iteration / batch
+        """
+        if is_divisor(batch_num, self.freq):
+
+            loss = np.mean(trainer.losses['total_loss'][-self.average:])
+
+            if loss < self._best_loss:
+                self._best_loss = loss
+                self.save(trainer, batch_num)
+
+    def save(self, trainer: Trainer, batch_num: int):
+        """saves a dictionary containing the network weights, the learning rates, the losses and the current \
+            best loss with its corresponding iteration to the disk
+
+        Args:
+            trainer (Trainer): the trainer object
+            batch_num (int): the current iteration / batch
+        """
+        prev_save = trainer.callback_params.pop('saved_iteration', 0)
+        trainer.callback_params['saved_iteration'] = batch_num
+        save_dict = {
+            'model': trainer.model.state_dict(),
+            'lrs': trainer.lrs,
+            'losses': trainer.losses,
+            'prev_save': prev_save,
+            'batch_num': batch_num,
+            'best_loss': self._best_loss
+        }
+        torch.save(save_dict, self.path)
+
+
+class LogLosses(TrainerCallback):
+    """Callback for logging the training losses"""
+    def end_batch(self, trainer: Trainer, batch_num: int) -> None:
+        """log loss at the current iteration
+
+        Args:
+            trainer (Trainer): the trainer object
+            batch_num (int): the index of the current iteration / batch
+        """
+        try:
+            trainer.log('train/total_loss', trainer.losses[trainer.TOTAL_LOSS_KEY][-1])
+        except IndexError:
+            pass

reflectorch/ml/dataloaders.py

@@ -0,0 +1,27 @@
+from torch import Tensor
+
+
+from reflectorch.data_generation import BasicDataset
+from reflectorch.data_generation.reflectivity import kinematical_approximation
+from reflectorch.data_generation.priors import BasicParams
+from reflectorch.ml.basic_trainer import DataLoader
+
+
+__all__ = [
+    "ReflectivityDataLoader",
+    "MultilayerDataLoader",
+]
+
+
+class ReflectivityDataLoader(BasicDataset, DataLoader):
+    """Dataloader for reflectivity data, combining functionality from the ``BasicDataset`` (basic dataset class for reflectivity) and the ``DataLoader`` (which inherits from ``TrainerCallback``) classes"""
+    pass
+
+
+class MultilayerDataLoader(ReflectivityDataLoader):
+    """Dataloader for reflectivity curves simulated using the kinematical approximation"""
+    def _sample_from_prior(self, batch_size: int):
+        return self.prior_sampler.optimized_sample(batch_size)
+
+    def _calc_curves(self, q_values: Tensor, params: BasicParams):
+        return kinematical_approximation(q_values, params.thicknesses, params.roughnesses, params.slds)

reflectorch/ml/loggers.py

@@ -0,0 +1,38 @@
+# -*- coding: utf-8 -*-
+#
+#
+# This source code is licensed under the GPL license found in the
+# LICENSE file in the root directory of this source tree.
+
+
+__all__ = [
+    'Logger',
+    'Loggers',
+    'PrintLogger',
+]
+
+
+class Logger(object):
+    "Base class defining a common interface for logging"
+    def log(self, name: str, data):
+        pass
+
+    def __setitem__(self, key, value):
+        """Enable dictionary-style setting to log data."""
+        self.log(key, value)
+
+
+class Loggers(Logger):
+    """Class for using multiple loggers"""
+    def __init__(self, *loggers):
+        self._loggers = tuple(loggers)
+
+    def log(self, name: str, data):
+        for logger in self._loggers:
+            logger.log(name, data)
+
+
+class PrintLogger(Logger):
+    """Logger which prints to the console"""
+    def log(self, name: str, data):
+        print(name, ': ', data)

reflectorch/ml/schedulers.py

@@ -0,0 +1,246 @@
+# -*- coding: utf-8 -*-
+#
+#
+# This source code is licensed under the GPL license found in the
+# LICENSE file in the root directory of this source tree.
+
+from torch.optim import lr_scheduler
+
+import numpy as np
+
+from reflectorch.ml.basic_trainer import Trainer, TrainerCallback, PeriodicTrainerCallback
+
+__all__ = [
+    'ScheduleBatchSize',
+    'ScheduleLR',
+    'StepLR',
+    'CyclicLR',
+    'LogCyclicLR',
+    'ReduceLROnPlateau',
+    'OneCycleLR',
+]
+
+
+class ScheduleBatchSize(PeriodicTrainerCallback):
+    """Batch size scheduler
+
+    Args:
+        step (int): number of iterations after which the batch size is modified.
+        gamma (int, optional): quantity which is added to or multiplied with the current batch size. Defaults to 2.
+        last_epoch (int, optional): the last training iteration for which the batch size is modified. Defaults to -1.
+        mode (str, optional): ``'add'`` for addition or ``'multiply'`` for multiplication. Defaults to 'add'.
+    """
+    def __init__(self, step: int, gamma: int = 2, last_epoch: int = -1, mode: str = 'add'):
+        super().__init__(step, last_epoch)
+
+        assert mode in ('add', 'multiply')
+
+        self.gamma = gamma
+        self.mode = mode
+
+    def _end_batch(self, trainer: Trainer, batch_num: int) -> None:
+        if self.mode == 'add':
+            trainer.batch_size += self.gamma
+        elif self.mode == 'multiply':
+            trainer.batch_size *= self.gamma
+
+
+class ScheduleLR(TrainerCallback):
+    """Base class for learning rate schedulers
+
+    Args:
+        lr_scheduler_cls: class of the learning rate scheduler
+    """
+    
+    def __init__(self, lr_scheduler_cls, **kwargs):
+        self.lr_scheduler_cls = lr_scheduler_cls
+        self.kwargs = kwargs
+        self.lr_scheduler = None
+
+    def start_training(self, trainer: Trainer) -> None:
+        """initializes a learning rate scheduler based on its class and keyword arguments at the start of training
+
+        Args:
+            trainer (Trainer): the trainer object
+        """
+        self.lr_scheduler = self.lr_scheduler_cls(trainer.optim, **self.kwargs)
+        trainer.callback_params['lrs'] = []
+
+    def end_batch(self, trainer: Trainer, batch_num: int) -> None:
+        """modifies the learning rate at the end of each iteration
+
+        Args:
+            trainer (Trainer): the trainer object
+            batch_num (int): index of the current iteration
+        """
+        self.lr_scheduler.step()
+
+
+class StepLR(ScheduleLR):
+    """Learning rate scheduler which decays the learning rate of each parameter group by gamma every ``step_size`` epochs.
+    
+    Args:
+        step_size (int): Period of learning rate decay
+        gamma (float): Multiplicative factor of learning rate decay
+        last_epoch (int, optional): The index of last iteration. Defaults to -1.
+    """
+    def __init__(self, step_size: int, gamma: float, last_epoch: int = -1, **kwargs):
+
+        
+        super().__init__(lr_scheduler.StepLR, step_size=step_size, gamma=gamma, last_epoch=last_epoch, **kwargs)
+
+    def start_training(self, trainer: Trainer) -> None:
+        trainer.optim.param_groups[0]['initial_lr'] = trainer.lr()
+        super().start_training(trainer)
+
+
+class CyclicLR(ScheduleLR):
+    """Cyclic learning rate scheduler
+
+    Args:
+        base_lr (float): Initial learning rate which is the lower boundary in the cycle
+        max_lr (float): Upper learning rate boundary in the cycle
+        step_size_up (int, optional): Number of training iterations in the increasing half of a cycle. Defaults to 2000.
+        cycle_momentum (bool, optional):  If True, momentum is cycled inversely to learning rate between ``base_momentum`` and ``max_momentum``. Defaults to False.
+        gamma (float, optional): Constant in ``‘exp_range’`` mode scaling function: gamma^(cycle iterations). Defaults to 1.
+        mode (str, optional):  One of: ``'triangular'`` (a basic triangular cycle without amplitude scaling),
+              ``'triangular2'`` (a basic triangular cycle that scales initial amplitude by half each cycle), ``'exp_range'`` 
+              (a cycle that scales initial amplitude by gamma^iterations at each cycle iteration). Defaults to 'triangular'.
+    """
+    def __init__(self, base_lr, max_lr, step_size_up: int = 2000,
+                 cycle_momentum: bool = False, gamma: float = 1., mode: str = 'triangular',
+                 **kwargs):
+        super().__init__(
+            lr_scheduler.CyclicLR,
+            base_lr=base_lr,
+            max_lr=max_lr,
+            step_size_up=step_size_up,
+            cycle_momentum=cycle_momentum,
+            gamma=gamma,
+            mode=mode,
+            **kwargs
+        )
+
+
+class LogCyclicLR(TrainerCallback):
+    """Cyclic learning rate scheduler on a logarithmic scale
+    
+    Args:
+        base_lr (float): Lower learning rate boundary in the cycle
+        max_lr (float): Upper learning rate boundary in the cycle
+        period (int, optional): Number of training iterations in the cycle. Defaults to 2000.
+        gamma (float, optional): Constant for scaling the amplitude as ``gamma`` ^ ``iterations``. Defaults to 1.
+        start_period (int, optional): Number of starting iterations with the default learning rate.
+        log (bool, optional): If ``True``, the cycle is in the logarithmic domain.
+        param_groups (tupe, optional): Parameter groups of the optimizer.
+    """
+    def __init__(self,
+                 base_lr,
+                 max_lr,
+                 period: int = 2000,
+                 gamma: float = None,
+                 log: bool = True,
+                 param_groups: tuple = (0,),
+                 start_period: int = 25,
+                 ):
+        self.base_lr = base_lr
+        self.max_lr = max_lr
+        self.period = period
+        self.gamma = gamma
+        self.param_groups = param_groups
+        self.log = log
+        self.start_period = start_period
+        self._axis = None
+        self._period = None
+
+    def get_lr(self, batch_num: int):
+        return self._get_lr(batch_num)
+
+    def _get_lr(self, batch_num):
+        num_period, t = batch_num // self.period, batch_num % self.period
+
+        if self._period != num_period:
+            self._period = num_period
+            if self.gamma and (num_period >= self.start_period):
+                amp = (self.max_lr - self.base_lr) * (self.gamma ** (num_period - self.start_period))
+                max_lr = self.base_lr + amp
+            else:
+                max_lr = self.max_lr
+
+            if self.log:
+                self._axis = np.logspace(np.log10(self.base_lr), np.log10(max_lr), self.period // 2)
+            else:
+                self._axis = np.linspace(self.base_lr, max_lr, self.period // 2)
+        if t < self.period // 2:
+            lr = self._axis[t]
+        else:
+            lr = self._axis[self.period - t - 1]
+        return lr
+
+    def end_batch(self, trainer: Trainer, batch_num: int):
+        lr = self.get_lr(batch_num)
+        for param_group in self.param_groups:
+            trainer.set_lr(lr, param_group)
+
+
+class ReduceLROnPlateau(TrainerCallback):
+    """Learning rate scheduler which reduces the learning rate when the loss stops decreasing
+    
+    Args:
+            gamma (float, optional): Multiplicative factor of learning rate decay. Defaults to 0.5.
+            patience (int, optional): The number of allowed iterations with no improvement after which the learning rate will be reduced. Defaults to 500.
+            average (int, optional): Size of the window over which the average loss is computed. Defaults to 50.
+            loss_key (str, optional): Defaults to 'total_loss'.
+            param_groups (tuple, optional): Defaults to (0,).
+    
+    """
+    def __init__(
+            self,
+            gamma: float = 0.5,
+            patience: int = 500,
+            average: int = 50,
+            loss_key: str = 'total_loss',
+            param_groups: tuple = (0,),
+    ):
+        """
+        """
+        self.patience = patience
+        self.average = average
+        self.gamma = gamma
+        self.loss_key = loss_key
+        self.param_groups = param_groups
+
+    def end_batch(self, trainer: Trainer, batch_num: int) -> None:
+        loss = trainer.losses[self.loss_key]
+
+        if len(loss) < self.patience:
+            return
+
+        if np.mean(loss[-self.patience:-(self.patience - self.average)]) <= np.mean(loss[-self.average:]):
+            for param_group in self.param_groups:
+                trainer.set_lr(trainer.lr(param_group) * self.gamma, param_group)
+                
+                
+class OneCycleLR(ScheduleLR):
+    """One-cycle learning rate scheduler (https://arxiv.org/abs/1708.07120)
+
+    Args:
+        max_lr (float): Upper learning rate boundary in the cycle
+        total_steps (int): The total number of steps in the cycle
+        pct_start (float, optional): The percentage of the cycle (in number of steps) spent increasing the learning rate. Defaults to 0.3.
+        div_factor (float, optional): Determines the initial learning rate via initial_lr = ``max_lr`` / ``div_factor``. Defaults to 25..
+        final_div_factor (float, optional): Determines the minimum learning rate via min_lr = ``initial_lr`` / ``final_div_factor``. Defaults to 1e4.
+        three_phase (bool, optional): If ``True``, use a third phase of the schedule to annihilate the learning rate according to ``final_div_factor`` instead of modifying the second phase. Defaults to True.
+    """
+    def __init__(self, max_lr: float, total_steps: int, pct_start: float = 0.3, div_factor: float = 25., 
+                 final_div_factor: float = 1e4, three_phase: bool = True, **kwargs):
+        super().__init__(
+            lr_scheduler.OneCycleLR,
+            max_lr=max_lr,
+            total_steps=total_steps,
+            pct_start=pct_start,
+            div_factor=div_factor ,
+            final_div_factor=final_div_factor,
+            three_phase=three_phase,
+            **kwargs
+        )

reflectorch/ml/trainers.py

@@ -0,0 +1,126 @@
+# -*- coding: utf-8 -*-
+#
+#
+# This source code is licensed under the GPL license found in the
+# LICENSE file in the root directory of this source tree.
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+from reflectorch.data_generation import BATCH_DATA_TYPE
+from reflectorch.ml.basic_trainer import Trainer
+from reflectorch.ml.dataloaders import ReflectivityDataLoader
+
+__all__ = [
+    'RealTimeSimTrainer',
+    'DenoisingAETrainer',
+    'VAETrainer',
+    'PointEstimatorTrainer',
+]
+
+
+class RealTimeSimTrainer(Trainer):
+    """Trainer with functionality to customize the sampled batch of data"""
+    loader: ReflectivityDataLoader
+
+    def get_batch_by_idx(self, batch_num: int):
+        """Gets a batch of data with the default batch size"""
+        batch_data = self.loader.get_batch(self.batch_size)
+        return self._get_batch(batch_data)
+
+    def get_batch_by_size(self, batch_size: int):
+        """Gets a batch of data with a custom batch size"""
+        batch_data = self.loader.get_batch(batch_size)
+        return self._get_batch(batch_data)
+
+    def _get_batch(self, batch_data: BATCH_DATA_TYPE):
+        """Modify the batch of data sampled from the data loader"""
+        raise NotImplementedError
+
+
+class PointEstimatorTrainer(RealTimeSimTrainer):
+    """Trainer for the regression inverse problem with incorporation of prior bounds"""
+    add_sigmas_to_context: bool = False
+        
+    def _get_batch(self, batch_data: BATCH_DATA_TYPE):
+        scaled_params = batch_data['scaled_params'].to(torch.float32)
+        scaled_curves = batch_data['scaled_noisy_curves'].to(torch.float32)
+        if self.train_with_q_input:
+            q_values = batch_data['q_values'].to(torch.float32)
+            scaled_q_values = self.loader.q_generator.scale_q(q_values)
+        else:
+            scaled_q_values = None
+
+        num_params = scaled_params.shape[-1] // 3
+        assert num_params * 3 == scaled_params.shape[-1]
+        scaled_params, scaled_bounds = torch.split(scaled_params, [num_params, 2 * num_params], dim=-1)
+
+        return scaled_params, scaled_bounds, scaled_curves, scaled_q_values
+
+    def get_loss_dict(self, batch_data):
+        """computes the loss dictionary"""
+
+        scaled_params, scaled_bounds, scaled_curves, scaled_q_values = batch_data
+
+        if self.train_with_q_input:
+            predicted_params = self.model(scaled_curves, scaled_bounds, scaled_q_values)
+        else:
+            predicted_params = self.model(scaled_curves, scaled_bounds)
+            
+        loss = self.mse(predicted_params, scaled_params)
+        return {'loss': loss}
+
+    def init(self):
+        self.mse = nn.MSELoss()
+
+
+class DenoisingAETrainer(RealTimeSimTrainer):
+    """Trainer which can be used for training a denoising autoencoder model. Overrides _get_batch and get_loss_dict methods """
+    def init(self):
+        self.criterion = nn.MSELoss()
+        self.loader.calc_denoised_curves = True
+
+    def _get_batch(self, batch_data: BATCH_DATA_TYPE):
+        """returns scaled curves with and without noise"""
+        scaled_noisy_curves, curves = batch_data['scaled_noisy_curves'], batch_data['curves']
+        scaled_curves = self.loader.curves_scaler.scale(curves)
+
+        scaled_noisy_curves, scaled_curves = scaled_noisy_curves.to(torch.float32), scaled_curves.to(torch.float32)
+
+        return scaled_noisy_curves, scaled_curves
+
+    def get_loss_dict(self, batch_data):
+        """returns the reconstruction loss of the autoencoder"""
+        scaled_noisy_curves, scaled_curves = batch_data
+        restored_curves = self.model(scaled_noisy_curves)
+        loss = self.criterion(scaled_curves, restored_curves)
+        return {'loss': loss}
+    
+
+class VAETrainer(DenoisingAETrainer):
+    """Trainer which can be used for training a denoising autoencoder model. Overrides _get_batch and get_loss_dict methods """
+    def init(self):
+        self.loader.calc_denoised_curves = True
+        self.freebits = 0.05
+    
+    def calc_kl(self, z_mu, z_logvar):
+        return 0.5*(z_mu**2 + torch.exp(z_logvar) - 1 - z_logvar)
+    
+    def gaussian_log_prob(self, z, mu, logvar):
+        return -0.5*(np.log(2*np.pi) + logvar + (z-mu)**2/torch.exp(logvar))
+
+    def get_loss_dict(self, batch_data):
+        """returns the reconstruction loss of the autoencoder"""
+        scaled_noisy_curves, scaled_curves = batch_data
+        _, (z_mu, z_logvar, restored_curves_mu, restored_curves_logvar) = self.model(scaled_noisy_curves)
+
+        l_rec = -torch.mean(self.gaussian_log_prob(scaled_curves, restored_curves_mu, restored_curves_logvar), dim=-1)  
+        l_kl = torch.mean(F.relu(self.calc_kl(z_mu, z_logvar) - self.freebits*np.log(2)) + self.freebits*np.log(2), dim=-1)
+        loss = torch.mean(l_rec + l_kl)/np.log(2)
+
+        l_rec = torch.mean(l_rec)
+        l_kl = torch.mean(l_kl)
+
+        return {'loss': loss}

reflectorch/ml/utils.py

@@ -0,0 +1,9 @@
+# -*- coding: utf-8 -*-
+#
+#
+# This source code is licensed under the GPL license found in the
+# LICENSE file in the root directory of this source tree.
+
+
+def is_divisor(num: int, div: int):
+    return num and not num % div

reflectorch/models/__init__.py

@@ -0,0 +1,22 @@
+# -*- coding: utf-8 -*-
+#
+#
+# This source code is licensed under the GPL license found in the
+# LICENSE file in the root directory of this source tree.
+
+from reflectorch.models.encoders import *
+from reflectorch.models.networks import *
+
+__all__ = [
+    "ConvEncoder",
+    "ConvDecoder",
+    "ConvAutoencoder",
+    "ConvVAE",
+    "TransformerEncoder",
+    "FnoEncoder",
+    "SpectralConv1d",
+    "ConvResidualNet1D",
+    "ResidualMLP",
+    "NetworkWithPriorsConvEmb",
+    "NetworkWithPriorsFnoEmb",
+]

reflectorch/models/activations.py

@@ -0,0 +1,50 @@
+import torch
+from torch import nn
+from torch.nn.functional import relu
+
+class Rowdy(nn.Module):
+    """adaptive activation function"""
+    def __init__(self, K=9):
+        super().__init__()
+        self.K = K
+        self.alpha = nn.Parameter(torch.cat((torch.ones(1), torch.zeros(K-1))))  
+        self.alpha.requiresGrad = True
+        self.omega = nn.Parameter(torch.ones(K))  
+        self.omega.requiresGrad = True
+
+    def forward(self, x):
+        rowdy = self.alpha[0]*relu(self.omega[0]*x)
+        for k in range(1, self.K):
+            rowdy += self.alpha[k]*torch.sin(self.omega[k]*k*x)
+        return rowdy
+    
+
+ACTIVATIONS = {
+    'relu': nn.ReLU,
+    'lrelu': nn.LeakyReLU,
+    'gelu': nn.GELU,
+    'selu': nn.SELU,
+    'elu': nn.ELU,
+    'sigmoid': nn.Sigmoid,    
+    'tanh': nn.Tanh,
+    'silu': nn.SiLU,
+    'mish': nn.Mish,
+    'rowdy': Rowdy,
+}
+
+
+def activation_by_name(name):
+    """returns an activation function module corresponding to its name
+
+    Args:
+        name (str): string denoting the activation function ('relu', 'lrelu', 'gelu', 'selu', 'elu', 'sigmoid', 'silu', 'mish', 'rowdy')
+
+    Returns:
+        nn.Module: Pytorch activation function module
+    """
+    if not isinstance(name, str):
+        return name
+    try:
+        return ACTIVATIONS[name.lower()]
+    except KeyError:
+        raise KeyError(f'Unknown activation function {name}')

reflectorch/models/encoders/__init__.py

@@ -0,0 +1,27 @@
+# -*- coding: utf-8 -*-
+#
+#
+# This source code is licensed under the GPL license found in the
+# LICENSE file in the root directory of this source tree.
+
+from reflectorch.models.encoders.conv_encoder import (
+    ConvEncoder,
+    ConvDecoder,
+    ConvAutoencoder,
+    ConvVAE,
+)
+from reflectorch.models.encoders.fno import FnoEncoder, SpectralConv1d
+from reflectorch.models.encoders.transformers import TransformerEncoder
+from reflectorch.models.encoders.conv_res_net import ConvResidualNet1D
+
+
+__all__ = [
+    "TransformerEncoder",
+    "ConvEncoder",
+    "ConvDecoder",
+    "ConvAutoencoder",
+    "ConvVAE",
+    "ConvResidualNet1D",
+    "FnoEncoder",
+    "SpectralConv1d",
+]