reflectorch 1.5.1__py3-none-any.whl

reflectorch/inference/scipy_fitter.py

@@ -0,0 +1,364 @@
+
+import warnings
+import joblib
+from joblib import Parallel, delayed
+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",
+    "batch_standard_refl_fit",
+]
+
+
+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(
+        standard_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 batch_refl_fit(
+        q: np.ndarray, 
+        curves: np.ndarray,
+        init_params: np.ndarray, # (n_curves, n_params)
+        prior_sampler: PriorSampler,
+        bounds: np.ndarray = None,
+        error_bars: np.ndarray = None,
+        scale_curve_func=np.log10,
+        method: str = 'trf', #'lm', 'trf'
+        polishing_max_steps: int = None,
+        reflectivity_kwargs: dict = None,
+        n_jobs: int = -1,
+        verbose: int = 5,
+        **kwargs
+):
+    """
+    Fit (polished fit) multiple reflectivity curves in parallel using joblib.
+
+    Parameters
+    ----------
+    q : np.ndarray
+        1D array of momentum transfer values (same for all curves).
+    curves : np.ndarray
+        2D array of reflectivity curves with shape (n_curves, n_q).
+    init_params : np.ndarray
+        2D array of initial parameter guesses (n_curves, n_params).
+    prior_sampler : PriorSampler
+        The prior sampler.
+    bounds : np.ndarray, optional
+        Bounds for the parameters, shape (2, n_params). Shared by all the curves. Default: None.
+    error_bars : np.ndarray, optional
+        Error bars for the curves, shape (n_curves, n_q). Default: None.
+    scale_curve_func : callable, optional
+        Function to scale the curves. Default: `np.log10`.
+    method : str, optional
+        The method to use for the fitting. Default: 'trf'.
+    polishing_max_steps : int, optional
+        The maximum number of function evaluations for the polishing step. Default: None.
+    reflectivity_kwargs : dict, optional
+        Keyword arguments for the reflectivity function. Default: None.
+    n_jobs : int, optional
+        The number of jobs to run in parallel. Default: -1 (all CPUs).
+    verbose : int, optional
+        The verbosity level for joblib. Default: 5.
+    **kwargs : dict
+        Extra keyword arguments passed to `scipy.optimize.curve_fit`.
+
+    Returns
+    -------
+    params_array : np.ndarray
+        Array of fitted parameter values for each curve, shape (n_curves, n_params).
+    error_bars_array : np.ndarray
+        Array of error bars for the fitted parameter values, shape (n_curves, n_params).
+    curves_array : np.ndarray
+        Array of fitted reflectivity curves, shape (n_curves, n_q).
+    """
+    if bounds is not None:
+        if bounds.ndim == 2:
+            bounds = bounds[None].repeat(curves.shape[0], 0)
+        elif bounds.ndim == 3:
+            assert bounds.shape[0] == curves.shape[0], f"Bounds must have the same number of curves as the number of curves, got {bounds.shape[0]} and {curves.shape[0]}"
+        else:
+            raise ValueError(f"Bounds must be a 2D or 3D array, got {bounds.ndim}D array")
+    else:
+        bounds = [None] * curves.shape[0]
+    
+    results = Parallel(n_jobs=n_jobs, verbose=verbose)(
+        delayed(refl_fit)(
+            q=q, curve=curve, init_params=init_params,
+            bounds=bound,
+            prior_sampler=prior_sampler,
+            error_bars=error_bars,
+            method=method,
+            scale_curve_func=scale_curve_func,
+            polishing_max_steps=polishing_max_steps,
+            reflectivity_kwargs=reflectivity_kwargs,
+            **kwargs
+        )
+        for curve, init_params, bound in zip(curves, init_params, bounds)
+    )
+
+    params_array, error_bars, curves_array = zip(*results)
+    return np.array(params_array), np.array(error_bars), np.array(curves_array)
+
+
+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_steps: int = None,
+        reflectivity_kwargs: dict = None,
+        **kwargs
+):
+    if bounds is not None:
+        if bounds.ndim != 2:
+            raise ValueError(f"Bounds must be a 2D array, got {bounds.ndim}D array")
+        # 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  
+
+    if polishing_max_steps is not None:
+        if method == 'lm':
+            kwargs['maxfev'] = polishing_max_steps
+        else:
+            kwargs['max_nfev'] = polishing_max_steps
+
+    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,
+        **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()
+    # cov matrix --> variance of the parameter estimate
+    if res[1] is not None and np.ndim(res[1]) == 2 and np.all(np.isfinite(res[1])):
+        pol_param_errs = np.sqrt(np.diag(res[1]))
+    else: 
+        pol_param_errs = np.full_like(res[1], np.nan)
+    return res[0], pol_param_errs, 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 standard_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_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,
+        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,32 @@
+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',
+    'TensorBoardLogger',
+    'ScheduleBatchSize',
+    'ScheduleLR',
+    'StepLR',
+    'CyclicLR',
+    'LogCyclicLR',
+    'ReduceLROnPlateau',
+    'OneCycleLR',
+    'CosineAnnealingWithWarmup',
+    'ReflectivityDataLoader',
+    'MultilayerDataLoader',
+    'RealTimeSimTrainer',
+    'DenoisingAETrainer',
+    'PointEstimatorTrainer',
+]