reflectorch 1.2.1__py3-none-any.whl → 1.4.0__py3-none-any.whl

reflectorch/inference/scipy_fitter.py

@@ -2,11 +2,14 @@ import warnings
 
 import numpy as np
 from scipy.optimize import minimize, curve_fit
+import torch
 
+from reflectorch.data_generation.priors.base import PriorSampler
 from reflectorch.data_generation.reflectivity import abeles_np
 
 __all__ = [
     "standard_refl_fit",
+    "refl_fit",
     "fit_refl_curve",
     "restore_masked_params",
     "get_fit_with_growth",
@@ -26,7 +29,6 @@ def standard_restore_params(fitted_params) -> dict:
 def mse_loss(curve1, curve2):
     return np.sum((curve1 - curve2) ** 2)
 
-
 def standard_refl_fit(
         q: np.ndarray, curve: np.ndarray,
         init_params: np.ndarray,
@@ -41,7 +43,7 @@ def standard_refl_fit(
         init_params = np.clip(init_params, *bounds)
 
     res = curve_fit(
-        get_scaled_curve_func(
+        standard_get_scaled_curve_func(
             refl_generator=refl_generator,
             restore_params_func=restore_params_func,
             scale_curve_func=scale_curve_func,
@@ -53,9 +55,73 @@ def standard_refl_fit(
     curve = refl_generator(q, **restore_params_func(res[0]))
     return res[0], curve
 
+def refl_fit(
+        q: np.ndarray, 
+        curve: np.ndarray,
+        init_params: np.ndarray,
+        prior_sampler: PriorSampler,
+        bounds: np.ndarray = None,
+        error_bars: np.ndarray = None,
+        scale_curve_func=np.log10,
+        method: str = 'trf', #'lm', 'trf'
+        polishing_max_nfev: int = None,
+        reflectivity_kwargs: dict = None,
+        **kwargs
+):
+    if bounds is not None:
+        # introduce a small perturbation for fixed bounds
+        epsilon = 1e-6
+        adjusted_bounds = bounds.copy()
+
+        for i in range(bounds.shape[1]): 
+            if bounds[0, i] == bounds[1, i]:
+                adjusted_bounds[0, i] -= epsilon
+                adjusted_bounds[1, i] += epsilon
+
+        init_params = np.clip(init_params, *adjusted_bounds)
+        if method != 'lm':
+            kwargs['bounds'] = adjusted_bounds
+
+    reflectivity_kwargs = reflectivity_kwargs or {}
+    for key, value in reflectivity_kwargs.items():
+        if isinstance(value, float):
+            reflectivity_kwargs[key] = torch.tensor([[value]], dtype=torch.float64)
+        elif isinstance(value, np.ndarray):
+            reflectivity_kwargs[key] = torch.tensor(value, dtype=torch.float32).unsqueeze(0)     
+
+    curve = np.clip(curve, a_min=1e-12, a_max=None)
+
+    if error_bars is not None and scale_curve_func == np.log10:
+        error_bars = np.clip(error_bars, a_min=1e-20, a_max=None)
+        scaled_error_bars = error_bars / (curve * np.log(10))
+    else:
+        scaled_error_bars = None  
+
+    res = curve_fit(
+        f=get_scaled_curve_func(
+            scale_curve_func=scale_curve_func,
+            prior_sampler=prior_sampler,
+            reflectivity_kwargs=reflectivity_kwargs,
+        ),
+        xdata=q, 
+        ydata=scale_curve_func(curve).reshape(-1),
+        p0=init_params,
+        sigma=scaled_error_bars,
+        absolute_sigma=True,
+        method=method,
+        max_nfev=polishing_max_nfev,
+        **kwargs
+    )
+
+    curve = prior_sampler.param_model.reflectivity(torch.tensor(q, dtype=torch.float64), 
+                                                   torch.tensor(res[0], dtype=torch.float64).unsqueeze(0), 
+                                                   **reflectivity_kwargs).squeeze().numpy()
+    return res[0], curve
+
 
 def get_fit_with_growth(
-        q: np.ndarray, curve: np.ndarray,
+        q: np.ndarray, 
+        curve: np.ndarray,
         init_params: np.ndarray,
         bounds: np.ndarray = None,
         init_d_change: float = 0.,
@@ -68,10 +134,16 @@ def get_fit_with_growth(
         bounds = np.concatenate([bounds, np.array([0, max_d_change])[..., None]], -1)
 
     params, curve = standard_refl_fit(
-        q, curve, init_params, bounds, refl_generator=growth_reflectivity,
+        q, 
+        curve, 
+        init_params, 
+        bounds, 
+        refl_generator=growth_reflectivity,
         restore_params_func=get_restore_params_with_growth_func(q_size=q.size, d_idx=0),
-        scale_curve_func=scale_curve_func, **kwargs
+        scale_curve_func=scale_curve_func, 
+        **kwargs
     )
+
     params[0] += params[-1] / 2
     return params, curve
 
@@ -97,8 +169,7 @@ def fit_refl_curve(q: np.ndarray, curve: np.ndarray,
         warnings.warn(f"Minimization did not converge.")
     return res.x
 
-
-def get_scaled_curve_func(
+def standard_get_scaled_curve_func(
         refl_generator=abeles_np,
         restore_params_func=standard_restore_params,
         scale_curve_func=np.log10,
@@ -111,6 +182,26 @@ def get_scaled_curve_func(
 
     return scaled_curve_func
 
+def get_scaled_curve_func(
+        scale_curve_func=np.log10,
+        prior_sampler: PriorSampler = None,
+        reflectivity_kwargs: dict = None,
+):  
+    reflectivity_kwargs = reflectivity_kwargs or {}
+    
+    def scaled_curve_func(q, *fitted_params):
+        q_tensor = torch.from_numpy(q).to(torch.float64)
+        fitted_params_tensor = torch.tensor(fitted_params, dtype=torch.float64).unsqueeze(0)
+        
+        fitted_curve_tensor = prior_sampler.param_model.reflectivity(q_tensor, fitted_params_tensor, **reflectivity_kwargs)
+        fitted_curve = fitted_curve_tensor.squeeze().numpy()
+        
+        scaled_curve = scale_curve_func(fitted_curve)
+
+        return scaled_curve.reshape(-1)
+
+    return scaled_curve_func
+
 
 def get_fitting_func(
         q: np.ndarray,

reflectorch/ml/__init__.py

@@ -15,6 +15,7 @@ __all__ = [
     'Logger',
     'Loggers',
     'PrintLogger',
+    'TensorBoardLogger',
     'ScheduleBatchSize',
     'ScheduleLR',
     'StepLR',
@@ -22,6 +23,7 @@ __all__ = [
     'LogCyclicLR',
     'ReduceLROnPlateau',
     'OneCycleLR',
+    'CosineAnnealingWithWarmup',
     'ReflectivityDataLoader',
     'MultilayerDataLoader',
     'RealTimeSimTrainer',

reflectorch/ml/basic_trainer.py

@@ -1,7 +1,8 @@
 from typing import Optional, Tuple, Iterable, Any, Union, Type
 from collections import defaultdict
 
-from tqdm.notebook import trange
+from tqdm import tqdm as standard_tqdm
+from tqdm.notebook import tqdm as notebook_tqdm
 import numpy as np
 
 import torch
@@ -31,7 +32,6 @@ class Trainer(object):
         logger (Union[Logger, Tuple[Logger, ...], Loggers], optional): logger. Defaults to None.
         optim_cls (Type[torch.optim.Optimizer], optional): Pytorch optimizer. Defaults to torch.optim.Adam.
         optim_kwargs (dict, optional): optimizer arguments. Defaults to None.
-        train_with_q_input (bool, optional): if ``True`` the q values are also used as input. Defaults to False.
     """
 
     TOTAL_LOSS_KEY: str = 'total_loss'
@@ -42,7 +42,6 @@ class Trainer(object):
                  lr: float,
                  batch_size: int,
                  clip_grad_norm_max: Optional[int] = None,
-                 train_with_q_input: bool = False,
                  logger: Union[Logger, Tuple[Logger, ...], Loggers] = None,
                  optim_cls: Type[torch.optim.Optimizer] = torch.optim.Adam,
                  optim_kwargs: dict = None,
@@ -53,7 +52,6 @@ class Trainer(object):
         self.loader = loader
         self.batch_size = batch_size
         self.clip_grad_norm_max = clip_grad_norm_max
-        self.train_with_q_input = train_with_q_input
 
         self.optim = self.configure_optimizer(optim_cls, lr=lr, **(optim_kwargs or {}))
         self.lrs = []
@@ -78,7 +76,8 @@ class Trainer(object):
                     num_batches: int,
                     callbacks: Union[Tuple['TrainerCallback', ...], 'TrainerCallback'] = (),
                     disable_tqdm: bool = False,
-                    update_tqdm_freq: int = 10,
+                    use_notebook_tqdm: bool = False,
+                    update_tqdm_freq: int = 1,
                     grad_accumulation_steps: int = 1,
                     ):
         """starts the training process
@@ -87,6 +86,7 @@ class Trainer(object):
             num_batches (int): total number of training iterations
             callbacks (Union[Tuple['TrainerCallback'], 'TrainerCallback']): the trainer callbacks. Defaults to ().
             disable_tqdm (bool, optional): if ``True``, the progress bar is disabled. Defaults to False.
+            use_notebook_tqdm (bool, optional): should be set to ``True`` when used in a Jupyter Notebook. Defaults to False.
             update_tqdm_freq (int, optional): frequency for updating the progress bar. Defaults to 10.
             grad_accumulation_steps (int, optional): number of gradient accumulation steps. Defaults to 1.
         """
@@ -96,7 +96,8 @@ class Trainer(object):
 
         callbacks = _StackedTrainerCallbacks(list(callbacks) + [self.loader])
 
-        pbar = trange(num_batches, disable=disable_tqdm)
+        tqdm_class = notebook_tqdm if use_notebook_tqdm else standard_tqdm
+        pbar = tqdm_class(range(num_batches), disable=disable_tqdm)
 
         callbacks.start_training(self)
 
@@ -121,6 +122,7 @@ class Trainer(object):
 
             if self.clip_grad_norm_max is not None:
                 torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=self.clip_grad_norm_max)
+
             self.optim.step()
 
             avr_loss_dict = {k: np.mean(v) for k, v in avr_loss_dict.items()}
@@ -141,6 +143,16 @@ class Trainer(object):
             last_loss = np.mean(self.losses[self.TOTAL_LOSS_KEY][-10:])
             pbar.set_description(f'Loss = {last_loss:.2e}')
 
+            postfix = {}
+            for key in self.losses.keys():
+                if key != self.TOTAL_LOSS_KEY:
+                    last_value = self.losses[key][-1]
+                    postfix[key] = f'{last_value:.4f}'
+
+            postfix['lr'] = f'{self.lr():.2e}'
+
+            pbar.set_postfix(postfix)
+
     def get_batch_by_idx(self, batch_num: int) -> Any:
         raise NotImplementedError
 

reflectorch/ml/callbacks.py

@@ -74,7 +74,8 @@ class LogLosses(TrainerCallback):
             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
+        for loss_name, loss_values in trainer.losses.items():
+            try:
+                trainer.log(f'train/{loss_name}', loss_values[-1])
+            except IndexError:
+                continue

reflectorch/ml/loggers.py

@@ -1,7 +1,10 @@
+from torch.utils.tensorboard import SummaryWriter
+
 __all__ = [
     'Logger',
     'Loggers',
     'PrintLogger',
+    'TensorBoardLogger',
 ]
 
 
@@ -29,3 +32,25 @@ class PrintLogger(Logger):
     """Logger which prints to the console"""
     def log(self, name: str, data):
         print(name, ': ', data)
+
+class TensorBoardLogger(Logger):
+    def __init__(self, log_dir: str):
+        """
+        Args:
+            log_dir (str): Directory where TensorBoard logs will be written
+        """
+        super().__init__()
+        self.writer = SummaryWriter(log_dir=log_dir)
+        self.step = 1
+        
+    def log(self, name: str, data):
+        """Log scalar data to TensorBoard
+        
+        Args:
+            name (str): Name/tag for the data
+            data: Scalar value to log
+        """
+        if hasattr(data, 'item'):
+            data = data.item()
+        self.writer.add_scalar(name, data, self.step)
+        self.step += 1

reflectorch/ml/schedulers.py

@@ -1,3 +1,4 @@
+import math
 from torch.optim import lr_scheduler
 
 import numpy as np
@@ -12,6 +13,7 @@ __all__ = [
     'LogCyclicLR',
     'ReduceLROnPlateau',
     'OneCycleLR',
+    'CosineAnnealingWithWarmup',
 ]
 
 
@@ -69,6 +71,31 @@ class ScheduleLR(TrainerCallback):
         """
         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.
@@ -176,6 +203,21 @@ class LogCyclicLR(TrainerCallback):
         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
@@ -238,3 +280,77 @@ class OneCycleLR(ScheduleLR):
             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

@@ -2,6 +2,8 @@ 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
@@ -14,6 +16,18 @@ __all__ = [
 ]
 
 
+@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
@@ -34,46 +48,140 @@ class RealTimeSimTrainer(Trainer):
 
 
 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)
+    """Point estimator trainer for the inverse problem."""
+
+    def init(self):
+        if getattr(self, 'use_l1_loss', False):
+            self.criterion = nn.L1Loss(reduction='none')
         else:
-            scaled_q_values = None
+            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 scaled_params, scaled_bounds, scaled_curves, scaled_q_values
+        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
 
-    def get_loss_dict(self, batch_data):
-        """computes the loss dictionary"""
+            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
 
-        scaled_params, scaled_bounds, scaled_curves, scaled_q_values = batch_data
+            loss = (width_factors * base_loss).mean()
 
-        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 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
+    
+        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"""

reflectorch/models/__init__.py

@@ -5,11 +5,12 @@ __all__ = [
     "ConvEncoder",
     "ConvDecoder",
     "ConvAutoencoder",
-    "ConvVAE",
     "FnoEncoder",
+    "IntegralConvEmbedding",
     "SpectralConv1d",
     "ConvResidualNet1D",
     "ResidualMLP",
+    "NetworkWithPriors",
     "NetworkWithPriorsConvEmb",
     "NetworkWithPriorsFnoEmb",
 ]

reflectorch/models/encoders/__init__.py

@@ -2,9 +2,9 @@ from reflectorch.models.encoders.conv_encoder import (
     ConvEncoder,
     ConvDecoder,
     ConvAutoencoder,
-    ConvVAE,
 )
 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
 
 
@@ -12,8 +12,8 @@ __all__ = [
     "ConvEncoder",
     "ConvDecoder",
     "ConvAutoencoder",
-    "ConvVAE",
     "ConvResidualNet1D",
     "FnoEncoder",
     "SpectralConv1d",
+    "IntegralConvEmbedding",
 ]