reflectorch 1.0.0__py3-none-any.whl

reflectorch/inference/record_time.py

@@ -0,0 +1,43 @@
+from time import perf_counter
+from contextlib import contextmanager
+from functools import wraps
+
+
+class EvaluateTime(list):
+    @contextmanager
+    def __call__(self, name: str, *args, **kwargs):
+        start = perf_counter()
+        yield
+        self.action(perf_counter() - start, name, *args, **kwargs)
+
+    @staticmethod
+    def action(delta_time, name, *args, **kwargs):
+        print(f"Time for {name} = {delta_time:.2f} sec")
+
+    def __repr__(self):
+        return f'EvaluateTime(total={sum(self)}, num_records={len(self)})'
+
+
+def print_time(name: str or callable):
+    if isinstance(name, str):
+        return _print_time_context(name)
+    else:
+        return _print_time_wrap(name)
+
+
+def _print_time_wrap(func, name: str = None):
+    name = name or func.__name__
+
+    @wraps(func)
+    def wrapped_func(*args, **kwargs):
+        with _print_time_context(name):
+            return func(*args, **kwargs)
+
+    return wrapped_func
+
+
+@contextmanager
+def _print_time_context(name: str):
+    start = perf_counter()
+    yield
+    print(f"Time for {name} = {(perf_counter() - start):.2f} sec")

reflectorch/inference/sampler_solution.py

@@ -0,0 +1,56 @@
+import torch
+from torch import Tensor
+
+from reflectorch.data_generation.priors.utils import uniform_sampler
+from reflectorch.data_generation.priors.subprior_sampler import UniformSubPriorParams
+from reflectorch.data_generation.priors.params import Params
+from reflectorch.data_generation.likelihoods import LogLikelihood
+
+
+def simple_sampler_solution(
+        likelihood: LogLikelihood,
+        predicted_params: UniformSubPriorParams,
+        total_min_bounds: Tensor,
+        total_max_bounds: Tensor,
+        num: int = 2 ** 15,
+        coef: float = 0.1,
+) -> UniformSubPriorParams:
+    sampled_params_t = sample_around_params(predicted_params, total_min_bounds, total_max_bounds, num=num, coef=coef)
+    sampled_params = Params.from_tensor(sampled_params_t)
+    return get_best_mse_param(sampled_params, likelihood, predicted_params.min_bounds, predicted_params.max_bounds)
+
+
+def sample_around_params(predicted_params: UniformSubPriorParams,
+                         total_min_bounds: Tensor,
+                         total_max_bounds: Tensor,
+                         num: int = 2 ** 15,
+                         coef: float = 0.1,
+                         ) -> Tensor:
+    params_t = predicted_params.as_tensor(add_bounds=False)
+
+    delta = (predicted_params.max_bounds - predicted_params.min_bounds) * coef
+    min_bounds = torch.clamp(params_t - delta, total_min_bounds, total_max_bounds)
+    max_bounds = torch.clamp(params_t + delta, total_min_bounds, total_max_bounds)
+
+    sampled_params_t = uniform_sampler(min_bounds, max_bounds, num, params_t.shape[-1])
+    sampled_params_t[0] = params_t[0]
+
+    return sampled_params_t
+
+
+def get_best_mse_param(
+        params: Params,
+        likelihood: LogLikelihood,
+        min_bounds: Tensor = None,
+        max_bounds: Tensor = None,
+):
+    sampled_curves = params.reflectivity(likelihood.q)
+    log_probs = likelihood.calc_log_likelihood(sampled_curves)
+    best_idx = torch.argmax(log_probs)
+    best_param = params[best_idx:best_idx + 1]
+
+    if min_bounds is not None:
+        best_param = UniformSubPriorParams.from_tensor(
+            torch.cat([best_param.as_tensor(), torch.atleast_2d(min_bounds), torch.atleast_2d(max_bounds)], -1)
+        )
+    return best_param

reflectorch/inference/scipy_fitter.py

@@ -0,0 +1,171 @@
+import warnings
+
+import numpy as np
+from scipy.optimize import minimize, curve_fit
+
+from reflectorch.data_generation.reflectivity import abeles_np
+
+__all__ = [
+    "standard_refl_fit",
+    "fit_refl_curve",
+    "restore_masked_params",
+    "get_fit_with_growth",
+]
+
+
+def standard_restore_params(fitted_params) -> dict:
+    num_layers = (fitted_params.size - 2) // 3
+
+    return dict(
+        thickness=fitted_params[:num_layers],
+        roughness=fitted_params[num_layers:2 * num_layers + 1],
+        sld=fitted_params[2 * num_layers + 1:],
+    )
+
+
+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,
+        bounds: np.ndarray = None,
+        refl_generator=abeles_np,
+        restore_params_func=standard_restore_params,
+        scale_curve_func=np.log10,
+        **kwargs
+):
+    if bounds is not None:
+        kwargs['bounds'] = bounds
+        init_params = np.clip(init_params, *bounds)
+
+    res = curve_fit(
+        get_scaled_curve_func(
+            refl_generator=refl_generator,
+            restore_params_func=restore_params_func,
+            scale_curve_func=scale_curve_func,
+        ),
+        q, scale_curve_func(curve),
+        p0=init_params, **kwargs
+    )
+
+    curve = refl_generator(q, **restore_params_func(res[0]))
+    return res[0], curve
+
+
+def get_fit_with_growth(
+        q: np.ndarray, curve: np.ndarray,
+        init_params: np.ndarray,
+        bounds: np.ndarray = None,
+        init_d_change: float = 0.,
+        max_d_change: float = 30.,
+        scale_curve_func=np.log10,
+        **kwargs
+):
+    init_params = np.array(list(init_params) + [init_d_change])
+    if bounds is not None:
+        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,
+        restore_params_func=get_restore_params_with_growth_func(q_size=q.size, d_idx=0),
+        scale_curve_func=scale_curve_func, **kwargs
+    )
+    params[0] += params[-1] / 2
+    return params, curve
+
+
+def fit_refl_curve(q: np.ndarray, curve: np.ndarray,
+                   init_params: np.ndarray,
+                   bounds: np.ndarray = None,
+                   refl_generator=abeles_np,
+                   restore_params_func=standard_restore_params,
+                   scale_curve_func=np.log10,
+                   **kwargs
+                   ) -> np.ndarray:
+    fitting_func = get_fitting_func(
+        q=q, curve=curve,
+        refl_generator=refl_generator,
+        restore_params_func=restore_params_func,
+        scale_curve_func=scale_curve_func,
+    )
+
+    res = minimize(fitting_func, init_params, bounds=bounds, **kwargs)
+
+    if not res.success:
+        warnings.warn(f"Minimization did not converge.")
+    return res.x
+
+
+def get_scaled_curve_func(
+        refl_generator=abeles_np,
+        restore_params_func=standard_restore_params,
+        scale_curve_func=np.log10,
+):
+    def scaled_curve_func(q, *fitted_params):
+        fitted_params = restore_params_func(np.asarray(fitted_params))
+        fitted_curve = refl_generator(q, **fitted_params)
+        scaled_curve = scale_curve_func(fitted_curve)
+        return scaled_curve
+
+    return scaled_curve_func
+
+
+def get_fitting_func(
+        q: np.ndarray,
+        curve: np.ndarray,
+        refl_generator=abeles_np,
+        restore_params_func=standard_restore_params,
+        scale_curve_func=np.log10,
+        loss_func=mse_loss,
+):
+    scaled_curve = scale_curve_func(curve)
+
+    def fitting_func(fitted_params):
+        fitted_params = restore_params_func(fitted_params)
+        fitted_curve = refl_generator(q, **fitted_params)
+        loss = loss_func(scale_curve_func(fitted_curve), scaled_curve)
+        return loss
+
+    return fitting_func
+
+
+def restore_masked_params(fixed_params, fixed_mask):
+    def restore_params(fitted_params) -> dict:
+        params = np.empty_like(fixed_mask).astype(fitted_params.dtype)
+        params[fixed_mask] = fixed_params
+        params[~fixed_mask] = fitted_params
+        return standard_restore_params(params)
+
+    return restore_params
+
+
+def base_params2growth(base_params: dict, d_shift: np.ndarray, d_idx: int = 0) -> dict:
+    d_init = base_params['thickness'][None]
+    q_size = d_shift.size
+    d = d_init.repeat(q_size, 0)
+    d[:, d_idx] = d[:, d_idx] + d_shift
+
+    roughness = np.broadcast_to(base_params['roughness'][None], (q_size, base_params['roughness'].size))
+    sld = np.broadcast_to(base_params['sld'][None], (q_size, base_params['sld'].size))
+
+    return {
+        'thickness': d,
+        'roughness': roughness,
+        'sld': sld,
+    }
+
+
+def get_restore_params_with_growth_func(q_size: int, d_idx: int = 0):
+    def restore_params_with_growth(fitted_params) -> dict:
+        fitted_params, delta_d = fitted_params[:-1], fitted_params[-1]
+        base_params = standard_restore_params(fitted_params)
+        d_shift = np.linspace(0, delta_d, q_size)
+        return base_params2growth(base_params, d_shift, d_idx)
+
+    return restore_params_with_growth
+
+
+def growth_reflectivity(q: np.ndarray, **kwargs):
+    return abeles_np(q[..., None], **kwargs).flatten()

reflectorch/inference/torch_fitter.py

@@ -0,0 +1,87 @@
+from tqdm import trange
+
+import torch
+from torch import nn, Tensor
+
+from reflectorch.data_generation import LogLikelihood, reflectivity, PriorSampler
+
+
+class ReflGradientFit(object):
+    """Directly optimizes the thin film parameters using a Pytorch optimizer
+
+    Args:
+        q (Tensor): the q positions
+        exp_curve (Tensor): the experimental reflectivity curve
+        prior_sampler (PriorSampler): the prior sampler
+        params (Tensor): the initial thin film parameters
+        fit_indices (Tensor): the indices of the thin film parameters which are to be fitted
+        sigmas (Tensor, optional): error bars of the reflectivity curve, if not provided they are derived from ``rel_err`` and ``abs_err``. Defaults to None.
+        optim_cls (Type[torch.optim.Optimizer], optional): the Pytorch optimizer class. Defaults to None.
+        lr (float, optional): the learning rate. Defaults to 1e-2.
+        rel_err (float, optional): the relative error in the reflectivity curve. Defaults to 0.1.
+        abs_err (float, optional): the absolute error in the reflectivity curve. Defaults to 1e-7.
+    """
+    def __init__(self,
+                 q: Tensor,
+                 exp_curve: Tensor,
+                 prior_sampler: PriorSampler,
+                 params: Tensor,
+                 fit_indices: Tensor,
+                 sigmas: Tensor = None,
+                 optim_cls=None,
+                 lr: float = 1e-2,
+                 rel_err: float = 0.1,
+                 abs_err: float = 1e-7,
+                 ):
+        self.q = q
+
+        if sigmas is None:
+            sigmas = exp_curve * rel_err + abs_err
+
+        self.likelihood = LogLikelihood(q, exp_curve, prior_sampler, sigmas)
+
+        self.num_layers = params.shape[-1] // 3
+        self.fit_indices = fit_indices
+        self.init_params = params.clone()
+        self.params_to_fit = nn.Parameter(self.init_params[fit_indices].clone())
+
+        optim_cls = optim_cls or torch.optim.Adam
+        self.optim = optim_cls([self.params_to_fit], lr)
+
+        self.losses = []
+
+    @property
+    def params(self):
+        params = self.init_params.clone()
+        params[self.fit_indices] = self.params_to_fit
+        return params
+
+    def calc_log_likelihood(self):
+        return self.likelihood.calc_log_likelihood(self.refl())
+
+    def calc_log_prob_loss(self):
+        return - self.calc_log_likelihood().mean()
+
+    def refl(self):
+        d, sigma, rho = torch.split(self.params, [self.num_layers, self.num_layers + 1, self.num_layers + 1], -1)
+        return reflectivity(self.q, d, sigma, rho)
+
+    def run(self, num_iterations: int = 500, disable_tqdm: bool = False):
+        """Runs the optimization process
+
+        Args:
+            num_iterations (int, optional): number of iterations the optimization is run for. Defaults to 500.
+            disable_tqdm (bool, optional): whether to disable the prograss bar. Defaults to False.
+        """
+        pbar = trange(num_iterations, disable=disable_tqdm)
+
+        for _ in pbar:
+            self.optim.zero_grad()
+            loss = self.calc_log_prob_loss()
+            loss.backward()
+            self.optim.step()
+            self.losses.append(loss.item())
+            pbar.set_description(f'Loss = {loss.item():.2e}')
+
+    def clear(self):
+        self.losses.clear()

reflectorch/ml/__init__.py

@@ -0,0 +1,37 @@
+# -*- 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.ml.basic_trainer import *
+from reflectorch.ml.callbacks import *
+from reflectorch.ml.trainers import *
+from reflectorch.ml.loggers import *
+from reflectorch.ml.schedulers import *
+from reflectorch.ml.dataloaders import *
+
+__all__ = [
+    'Trainer',
+    'TrainerCallback',
+    'DataLoader',
+    'PeriodicTrainerCallback',
+    'SaveBestModel',
+    'LogLosses',
+    'Logger',
+    'Loggers',
+    'PrintLogger',
+    'ScheduleBatchSize',
+    'ScheduleLR',
+    'StepLR',
+    'CyclicLR',
+    'LogCyclicLR',
+    'ReduceLROnPlateau',
+    'OneCycleLR',
+    'ReflectivityDataLoader',
+    'MultilayerDataLoader',
+    'RealTimeSimTrainer',
+    'DenoisingAETrainer',
+    'VAETrainer',
+    'PointEstimatorTrainer',
+]

reflectorch/ml/basic_trainer.py

@@ -0,0 +1,286 @@
+# -*- 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 typing import Optional, Tuple, Iterable, Any, Union, Type
+from collections import defaultdict
+
+from tqdm.notebook import trange
+import numpy as np
+
+import torch
+from torch.nn import Module
+
+from reflectorch.ml.loggers import Logger, Loggers
+
+from .utils import is_divisor
+
+__all__ = [
+    'Trainer',
+    'TrainerCallback',
+    'DataLoader',
+    'PeriodicTrainerCallback',
+]
+
+
+class Trainer(object):
+    """Trainer class
+    
+    Args:
+        model (nn.Module): neural network
+        loader (DataLoader): data loader
+        lr (float): learning rate
+        batch_size (int): batch size
+        clip_grad_norm (int, optional): maximum norm for gradient clipping if it is not ``None``. Defaults to None.
+        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'
+
+    def __init__(self,
+                 model: Module,
+                 loader: 'DataLoader',
+                 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,
+                 **kwargs
+                 ):
+        
+        self.model = model
+        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 = []
+        self.losses = defaultdict(list)
+
+        self.logger = _init_logger(logger)
+        self.callback_params = {}
+        
+        for k, v in kwargs.items():
+            setattr(self, k, v)
+
+        self.init()
+
+    def init(self):
+        pass
+
+    def log(self, name: str, data):
+        """log data"""
+        self.logger.log(name, data)
+
+    def train(self,
+                    num_batches: int,
+                    callbacks: Union[Tuple['TrainerCallback', ...], 'TrainerCallback'] = (),
+                    disable_tqdm: bool = False,
+                    update_tqdm_freq: int = 10,
+                    grad_accumulation_steps: int = 1,
+                    ):
+        """starts the training process
+
+        Args:
+            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.
+            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.
+        """
+
+        if isinstance(callbacks, TrainerCallback):
+            callbacks = (callbacks,)
+
+        callbacks = _StackedTrainerCallbacks(list(callbacks) + [self.loader])
+
+        pbar = trange(num_batches, disable=disable_tqdm)
+
+        callbacks.start_training(self)
+
+        for batch_num in pbar:
+            self.model.train()
+
+            self.optim.zero_grad()
+            total_loss, avr_loss_dict = 0, defaultdict(list)
+
+            for _ in range(grad_accumulation_steps):
+
+                batch_data = self.get_batch_by_idx(batch_num)
+                loss_dict = self.get_loss_dict(batch_data)
+                loss = loss_dict['loss'] / grad_accumulation_steps
+                total_loss += loss.item()
+                _update_loss_dict(avr_loss_dict, loss_dict)
+
+                if not torch.isfinite(loss).item():
+                    raise ValueError('Loss is not finite!')
+
+                loss.backward()
+
+            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()}
+            self._update_losses(avr_loss_dict, total_loss)
+
+            if not disable_tqdm:
+                self._update_tqdm(pbar, batch_num, update_tqdm_freq)
+
+            break_epoch = callbacks.end_batch(self, batch_num)
+
+            if break_epoch:
+                break
+
+        callbacks.end_training(self)
+
+    def _update_tqdm(self, pbar, batch_num: int, update_tqdm_freq: int):
+        if is_divisor(batch_num, update_tqdm_freq):
+            last_loss = np.mean(self.losses[self.TOTAL_LOSS_KEY][-10:])
+            pbar.set_description(f'Loss = {last_loss:.2e}')
+
+    def get_batch_by_idx(self, batch_num: int) -> Any:
+        raise NotImplementedError
+
+    def get_loss_dict(self, batch_data) -> dict:
+        raise NotImplementedError
+
+    def _update_losses(self, loss_dict: dict, loss: float) -> None:
+        _update_loss_dict(self.losses, loss_dict)
+        self.losses[self.TOTAL_LOSS_KEY].append(loss)
+        self.lrs.append(self.lr())
+
+    def configure_optimizer(self, optim_cls, lr: float, **kwargs) -> torch.optim.Optimizer:
+        """configure the optimizer based on the optimizer class, the learning rate and the optimizer keyword arguments
+
+        Args:
+            optim_cls: the class of the optimizer
+            lr (float): the learning rate
+
+        Returns:
+            torch.optim.Optimizer:
+        """
+        optim = optim_cls(self.model.parameters(), lr,  **kwargs)
+        return optim
+
+    def lr(self, param_group: int = 0) -> float:
+        """get the learning rate"""
+        return self.optim.param_groups[param_group]['lr']
+
+    def set_lr(self, lr: float, param_group: int = 0) -> None:
+        """set the learning rate"""
+        self.optim.param_groups[param_group]['lr'] = lr
+
+
+class TrainerCallback(object):
+    """Base class for trainer callbacks
+    """
+    def start_training(self, trainer: Trainer) -> None:
+        """add functionality the start of training
+
+        Args:
+            trainer (Trainer): the trainer object
+        """
+        pass
+
+    def end_training(self, trainer: Trainer) -> None:
+        """add functionality at the end of training
+
+        Args:
+            trainer (Trainer): the trainer object
+        """
+        pass
+
+    def end_batch(self, trainer: Trainer, batch_num: int) -> Union[bool, None]:
+        """add functionality at the end of the iteration / batch
+
+        Args:
+            trainer (Trainer): the trainer object
+            batch_num (int): the index of the current iteration / batch
+
+        Returns:
+            Union[bool, None]:
+        """
+        pass
+
+    def __repr__(self):
+        return f'{self.__class__.__name__}()'
+
+
+class DataLoader(TrainerCallback):
+    pass
+
+
+class PeriodicTrainerCallback(TrainerCallback):
+    """Base class for trainer callbacks which perform an action periodically after a number of iterations
+        
+    Args:
+        step (int, optional): Number of iterations after which the action is repeated. Defaults to 1.
+        last_epoch (int, optional): the last training  iteration for which the action is performed. Defaults to -1.
+    """
+    def __init__(self, step: int = 1, last_epoch: int = -1):
+        self.step = step
+        self.last_epoch = last_epoch
+
+    def end_batch(self, trainer: Trainer, batch_num: int) -> Union[bool, None]:
+        """add functionality at the end of the iteration / batch
+
+        Args:
+            trainer (Trainer): the trainer object
+            batch_num (int): the index of the current iteration / batch
+
+        Returns:
+            Union[bool, None]:
+        """
+        if (
+                is_divisor(batch_num, self.step) and
+                (self.last_epoch == -1 or batch_num < self.last_epoch)
+        ):
+            return self._end_batch(trainer, batch_num)
+
+    def _end_batch(self, trainer: Trainer, batch_num: int) -> Union[bool, None]:
+        pass
+
+
+class _StackedTrainerCallbacks(TrainerCallback):
+    def __init__(self, callbacks: Iterable[TrainerCallback]):
+        self.callbacks = tuple(callbacks)
+
+    def start_training(self, trainer: Trainer) -> None:
+        for c in self.callbacks:
+            c.start_training(trainer)
+
+    def end_training(self, trainer: Trainer) -> None:
+        for c in self.callbacks:
+            c.end_training(trainer)
+
+    def end_batch(self, trainer: Trainer, batch_num: int) -> Union[bool, None]:
+        break_epoch = False
+        for c in self.callbacks:
+            break_epoch += bool(c.end_batch(trainer, batch_num))
+        return break_epoch
+
+    def __repr__(self):
+        callbacks = ", ".join(repr(c) for c in self.callbacks)
+        return f'StackedTrainerCallbacks({callbacks})'
+
+
+def _init_logger(logger: Union[Logger, Tuple[Logger, ...], Loggers] = None):
+    if not logger:
+        return Logger()
+    if isinstance(logger, Logger):
+        return logger
+    return Loggers(*logger)
+
+
+def _update_loss_dict(loss_dict: dict, new_values: dict):
+    for k, v in new_values.items():
+        loss_dict[k].append(v.item())