reflectorch 1.5.1__py3-none-any.whl

reflectorch/ml/schedulers.py

@@ -0,0 +1,356 @@
+import math
+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',
+    'CosineAnnealingWithWarmup',
+]
+
+
+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()
+
+    def simulate_and_plot(self, total_steps: int, initial_lr: float, log_scale: bool = False):
+        import torch
+        import matplotlib.pyplot as plt
+
+        dummy_optim = torch.optim.Adam([torch.zeros(1)], lr=initial_lr)
+        scheduler = self.lr_scheduler_cls(dummy_optim, **self.kwargs)
+
+        lrs = []
+        for step in range(total_steps):
+            lrs.append(dummy_optim.param_groups[0]['lr'])
+            scheduler.step()
+
+        plt.figure(figsize=(10, 6))
+        plt.plot(lrs, label='Learning Rate')
+        plt.xlabel('Steps')
+        plt.ylabel('Learning Rate')
+        plt.title('Learning Rate Schedule')
+
+        if log_scale:
+            plt.yscale('log')
+
+        plt.grid(True, which="both", linestyle='--', linewidth=0.5)
+        plt.legend()
+        plt.show()
+
+
+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)
+
+    def simulate_and_plot(self, total_steps: int, log_scale: bool = True):
+        import matplotlib.pyplot as plt
+        lrs = [self.get_lr(batch_num) for batch_num in range(total_steps)]
+
+        plt.figure(figsize=(10, 6))
+        plt.plot(lrs, label='Learning Rate')
+        plt.xlabel('Steps')
+        plt.ylabel('Learning Rate')
+        plt.title('Learning Rate Schedule')
+        if log_scale:
+            plt.yscale('log')
+        plt.grid(True, which='both', linestyle='--', linewidth=0.5)
+        plt.legend()
+        plt.show()
+
+
+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
+        )
+
+class CosineAnnealingWithWarmup(TrainerCallback):
+    """
+    Cosine annealing scheduler with a warm-up stage.
+
+    Args:
+        max_lr (float): The maximum learning rate after the warm-up phase.
+        min_lr (float): The minimum learning rate after the warm-up phase.
+        warmup_iters (int): The number of iterations for the warm-up phase.
+        total_iters (int): The total number of iterations for the scheduler (including warm-up).
+    """
+    def __init__(self, max_lr=None, min_lr=1.0e-6, warmup_iters=100, total_iters=100000):
+        self.max_lr = max_lr
+        self.min_lr = min_lr
+        self.warmup_iters = warmup_iters
+        self.total_iters = total_iters
+
+    def get_lr(self, step):
+        """
+        Compute the learning rate for a given iteration.
+
+        Args:
+            step (int): The current iteration.
+
+        Returns:
+            float: The learning rate for the current iteration.
+        """
+        if step < self.warmup_iters:
+            # Warm-up stage: Linear increase from 0 to max_lr
+            return self.max_lr * step / self.warmup_iters
+        elif step < self.total_iters:
+            # Cosine annealing stage
+            t = (step - self.warmup_iters) / (self.total_iters - self.warmup_iters)
+            return self.min_lr + 0.5 * (self.max_lr - self.min_lr) * (1 + math.cos(math.pi * t))
+        else:
+            # Beyond total iterations: Return min_lr
+            return self.min_lr
+    
+    def start_training(self, trainer: Trainer) -> None:
+        self.max_lr = trainer.lr()
+
+    def end_batch(self, trainer: Trainer, batch_num: int):
+        """
+        Updates the learning rate at the end of each batch.
+
+        Args:
+            trainer (Trainer): The trainer object.
+            batch_num (int): The current batch number.
+        """
+        lr = self.get_lr(batch_num)
+        trainer.set_lr(lr)
+
+    def simulate_and_plot(self, total_steps: int = None, log_scale: bool = False):
+        """
+        Simulates and plots the learning rate evolution.
+
+        Args:
+            total_batches (int, optional): Total number of batches to simulate. If None, uses self.total_iters.
+        """
+        
+        total_steps = total_steps or self.total_iters
+        lrs = [self.get_lr(step) for step in range(total_steps)]
+
+        import matplotlib.pyplot as plt
+        plt.figure(figsize=(10, 6))
+        plt.plot(lrs, label='Learning Rate')
+        plt.xlabel('Steps')
+        plt.ylabel('Learning Rate')
+        plt.title('Learning Rate Scheduler')
+        if log_scale:
+            plt.yscale('log')
+        plt.grid(True, which='both', linestyle='--', linewidth=0.5)
+        plt.legend()
+        plt.show()

reflectorch/ml/trainers.py

@@ -0,0 +1,201 @@
+import numpy as np
+import torch
+import torch.nn.functional as F
+from torch import nn
+from dataclasses import dataclass
+from typing import Optional
+
+from reflectorch.data_generation import BATCH_DATA_TYPE
+from reflectorch.ml.basic_trainer import Trainer
+from reflectorch.ml.dataloaders import ReflectivityDataLoader
+
+__all__ = [
+    'RealTimeSimTrainer',
+    'DenoisingAETrainer',
+    'PointEstimatorTrainer',
+]
+
+
+@dataclass
+class BasicBatchData:
+    scaled_curves: torch.Tensor
+    scaled_bounds: torch.Tensor
+    scaled_params: torch.Tensor = None
+    scaled_sigmas: Optional[torch.Tensor] = None
+    scaled_q_values: Optional[torch.Tensor] = None
+    scaled_denoised_curves: Optional[torch.Tensor] = None
+    key_padding_mask: Optional[torch.Tensor] = None
+    scaled_conditioning_params: Optional[torch.Tensor] = None
+    unscaled_q_values: Optional[torch.Tensor] = None
+
+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):
+    """Point estimator trainer for the inverse problem."""
+
+    def init(self):
+        if getattr(self, 'use_l1_loss', False):
+            self.criterion = nn.L1Loss(reduction='none')
+        else:
+            self.criterion = nn.MSELoss(reduction='none')
+        self.use_curve_reconstruction_loss = getattr(self, 'use_curve_reconstruction_loss', False)
+        self.rescale_loss_interval_width = getattr(self, 'rescale_loss_interval_width', False)
+        if self.use_curve_reconstruction_loss:
+            self.loader.calc_denoised_curves = True
+
+        self.train_with_q_input = getattr(self, 'train_with_q_input', False)
+        self.train_with_sigmas = getattr(self, 'train_with_sigmas', False)
+        self.condition_on_q_resolutions = getattr(self, 'condition_on_q_resolutions', False)
+    
+    def _get_batch(self, batch_data: BATCH_DATA_TYPE) -> BasicBatchData:
+        def get_scaled_or_none(key, scaler=None):
+            value = batch_data.get(key)
+            if value is None:
+                return None
+            scale_func = scaler or (lambda x: x)
+            return scale_func(value).to(torch.float32)
+
+        scaled_params = batch_data['scaled_params'].to(torch.float32)
+        scaled_curves = batch_data['scaled_noisy_curves'].to(torch.float32)
+        scaled_denoised_curves = get_scaled_or_none('curves', self.loader.curves_scaler.scale)
+        scaled_q_values = get_scaled_or_none('q_values', self.loader.q_generator.scale_q) if self.train_with_q_input else None
+        key_padding_mask = batch_data.get('key_padding_mask', None)
+
+        scaled_q_resolutions = get_scaled_or_none('q_resolutions', self.loader.smearing.scale_resolutions) if self.condition_on_q_resolutions else None
+        conditioning_params = []
+        if scaled_q_resolutions is not None:
+            conditioning_params.append(scaled_q_resolutions)
+        scaled_conditioning_params = torch.cat(conditioning_params, dim=-1) if len(conditioning_params) > 0 else 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 BasicBatchData(
+            scaled_params=scaled_params,
+            scaled_bounds=scaled_bounds,
+            scaled_curves=scaled_curves,
+            scaled_q_values=scaled_q_values,
+            scaled_denoised_curves=scaled_denoised_curves,
+            scaled_conditioning_params=scaled_conditioning_params,
+            unscaled_q_values=batch_data['q_values'],
+            key_padding_mask=key_padding_mask,
+        )
+
+    def get_loss_dict(self, batch_data: BasicBatchData):
+        """Returns the regression loss"""
+        scaled_params=batch_data.scaled_params
+        scaled_curves=batch_data.scaled_curves
+        scaled_bounds=batch_data.scaled_bounds
+        scaled_q_values=batch_data.scaled_q_values
+        key_padding_mask=batch_data.key_padding_mask
+        scaled_conditioning_params=batch_data.scaled_conditioning_params
+        unscaled_q_values=batch_data.unscaled_q_values
+
+        predicted_params = self.model(
+            curves = scaled_curves,
+            bounds = scaled_bounds,
+            q_values = scaled_q_values,
+            conditioning_params = scaled_conditioning_params,
+            key_padding_mask = key_padding_mask,
+            unscaled_q_values = unscaled_q_values,
+        )
+
+        if not self.rescale_loss_interval_width:
+            loss = self.criterion(predicted_params, scaled_params).mean()
+        else:
+            n_params = scaled_params.shape[-1]
+            b_min = scaled_bounds[..., :n_params]
+            b_max = scaled_bounds[..., n_params:]
+            interval_width = b_max - b_min
+
+            base_loss  = self.criterion(predicted_params, scaled_params)
+            if isinstance(self.criterion, torch.nn.MSELoss):
+                width_factors = (interval_width / 2) ** 2
+            elif isinstance(self.criterion, torch.nn.L1Loss):
+                width_factors = interval_width / 2
+
+            loss = (width_factors * base_loss).mean()
+
+        return {'loss': loss}
+
+
+# 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.loader.calc_denoised_curves = True
+    
+        if getattr(self, 'use_l1_loss', False):
+            self.criterion = nn.L1Loss()
+        else:
+            self.criterion = nn.MSELoss()
+
+    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}
+    

reflectorch/ml/utils.py

@@ -0,0 +1,2 @@
+def is_divisor(num: int, div: int):
+    return num and not num % div

reflectorch/models/__init__.py

@@ -0,0 +1,16 @@
+from reflectorch.models.encoders import *
+from reflectorch.models.networks import *
+
+__all__ = [
+    "ConvEncoder",
+    "ConvDecoder",
+    "ConvAutoencoder",
+    "FnoEncoder",
+    "IntegralConvEmbedding",
+    "SpectralConv1d",
+    "ConvResidualNet1D",
+    "ResidualMLP",
+    "NetworkWithPriors",
+    "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,19 @@
+from reflectorch.models.encoders.conv_encoder import (
+    ConvEncoder,
+    ConvDecoder,
+    ConvAutoencoder,
+)
+from reflectorch.models.encoders.fno import FnoEncoder, SpectralConv1d
+from reflectorch.models.encoders.integral_kernel_embedding import IntegralConvEmbedding
+from reflectorch.models.encoders.conv_res_net import ConvResidualNet1D
+
+
+__all__ = [
+    "ConvEncoder",
+    "ConvDecoder",
+    "ConvAutoencoder",
+    "ConvResidualNet1D",
+    "FnoEncoder",
+    "SpectralConv1d",
+    "IntegralConvEmbedding",
+]