reflectorch 1.5.1__py3-none-any.whl

reflectorch/inference/preprocess_exp/interpolation.py

@@ -0,0 +1,19 @@
+import numpy as np
+
+
+def interp_reflectivity(q_interp, q, reflectivity, min_value: float = 1e-10, logspace = False):
+    """Interpolate data on a base 10 logarithmic scale
+
+    Args:
+        q_interp (array-like): reciprocal space points used for the interpolation
+        q (array-like): reciprocal space points of the measured reflectivity curve
+        reflectivity (array-like): reflectivity curve measured at the points ``q``
+        min_value (float, optional): minimum intensity of the reflectivity curve. Defaults to 1e-10.
+
+    Returns:
+        array-like: interpolated reflectivity curve
+    """
+    if not(logspace):
+        return 10 ** np.interp(q_interp, q, np.log10(np.clip(reflectivity, min_value, None)))
+    else:
+        return 10 ** np.interp(np.log10(q_interp), np.log10(q), np.log10(np.clip(reflectivity, min_value, None)))

reflectorch/inference/preprocess_exp/normalize.py

@@ -0,0 +1,21 @@
+try:
+    from typing import Literal
+except ImportError:
+    from typing_extensions import Literal
+
+import numpy as np
+from numpy import ndarray
+
+NORMALIZE_MODE = Literal["first", "max", "incoming_intensity"]
+
+
+def intensity2reflectivity(intensity: ndarray, mode: NORMALIZE_MODE, incoming_intensity=None) -> np.ndarray:
+    if mode == "first":
+        return intensity / intensity[0]
+    if mode == "max":
+        return intensity / intensity.max()
+    if mode == "incoming_intensity":
+        if incoming_intensity is None:
+            raise ValueError("incoming_intensity is None")
+        return intensity / incoming_intensity
+    raise ValueError(f"Unknown mode {mode}")

reflectorch/inference/preprocess_exp/preprocess.py

@@ -0,0 +1,121 @@
+from dataclasses import dataclass
+
+import numpy as np
+
+from reflectorch.inference.preprocess_exp.interpolation import interp_reflectivity
+from reflectorch.inference.preprocess_exp.footprint import apply_footprint_correction, BEAM_SHAPE
+from reflectorch.inference.preprocess_exp.normalize import intensity2reflectivity, NORMALIZE_MODE
+from reflectorch.inference.preprocess_exp.attenuation import apply_attenuation_correction
+from reflectorch.inference.preprocess_exp.cut_with_q_ratio import cut_curve
+from reflectorch.utils import angle_to_q
+
+
+def standard_preprocessing(
+        intensity: np.ndarray,
+        scattering_angle: np.ndarray,
+        attenuation: np.ndarray,
+        q_interp: np.ndarray,
+        wavelength: float,
+        beam_width: float,
+        sample_length: float,
+        min_intensity: float = 1e-10,
+        beam_shape: BEAM_SHAPE = "gauss",
+        normalize_mode: NORMALIZE_MODE = "max",
+        incoming_intensity: float = None,
+        max_q: float = None,  # if provided, max_angle is ignored
+        max_angle: float = None,
+) -> dict:
+    """Preprocesses a raw experimental reflectivity curve by applying attenuation correction, footprint correction, cutting at a maximum q value and interpolation
+
+    Args:
+        intensity (np.ndarray): array of intensities of the reflectivity curve
+        scattering_angle (np.ndarray): array of scattering angles
+        attenuation (np.ndarray): attenuation factors for each measured point
+        q_interp (np.ndarray): reciprocal space points used for the interpolation
+        wavelength (float): the wavelength of the beam
+        beam_width (float): the beam width
+        sample_length (float): the sample length
+        min_intensity (float, optional): intensities lower than this value are removed. Defaults to 1e-10.
+        beam_shape (BEAM_SHAPE, optional): the shape of the beam, either "gauss" or "box". Defaults to "gauss".
+        normalize_mode (NORMALIZE_MODE, optional): normalization mode, either "first", "max" or "incoming_intensity". Defaults to "max".
+        incoming_intensity (float, optional): array of intensities for the "incoming_intensity" normalization. Defaults to None.
+        max_q (float, optional): the maximum q value at which the curve is cut. Defaults to None.
+        max_angle (float, optional): the maximum scattering angle at which the curve is cut; only used if max_q is not provided. Defaults to None.
+
+    Returns:
+        dict: dictionary containing the interpolated reflectivity curve, the curve before interpolation, the q values before interpolation, the q values after interpolation and the q ratio of the cutting
+    """
+    intensity = apply_attenuation_correction(
+        intensity,
+        attenuation,
+        scattering_angle,
+    )
+
+    intensity = apply_footprint_correction(
+        intensity, scattering_angle, beam_width=beam_width, sample_length=sample_length, beam_shape=beam_shape
+    )
+
+    curve = intensity2reflectivity(intensity, normalize_mode, incoming_intensity)
+
+    curve, scattering_angle = remove_low_statistics(curve, scattering_angle, thresh=min_intensity)
+
+    q = angle_to_q(scattering_angle, wavelength)
+
+    q, curve, q_ratio = cut_curve(q, curve, max_q, max_angle, wavelength)
+
+    curve_interp = interp_reflectivity(q_interp, q, curve)
+
+    assert np.all(np.isfinite(curve_interp))
+    assert np.all(np.isfinite(curve))
+    assert np.all(np.isfinite(q))
+    assert np.all(np.isfinite(q_interp))
+    assert np.all(curve > 0.)
+    assert np.all(curve_interp > 0.)
+
+    return {
+        "curve_interp": curve_interp, "curve": curve, "q_values": q, "q_interp": q_interp, "q_ratio": q_ratio,
+    }
+
+
+def remove_low_statistics(curve, scattering_angle, thresh: float = 1e-7):
+    indices = (curve > thresh) & np.isfinite(curve)
+    return curve[indices], scattering_angle[indices]
+
+
+@dataclass
+class StandardPreprocessing:
+    q_interp: np.ndarray = None
+    wavelength: float = 1.
+    beam_width: float = None
+    sample_length: float = None
+    beam_shape: BEAM_SHAPE = "gauss"
+    normalize_mode: NORMALIZE_MODE = "max"
+    incoming_intensity: float = None
+
+    def preprocess(self,
+                   intensity: np.ndarray,
+                   scattering_angle: np.ndarray,
+                   attenuation: np.ndarray,
+                   **kwargs
+                   ) -> dict:
+        attrs = self._get_updated_attrs(**kwargs)
+        return standard_preprocessing(
+            intensity,
+            scattering_angle,
+            attenuation,
+            **attrs
+        )
+
+    __call__ = preprocess
+
+    def set_parameters(self, **kwargs) -> None:
+        for k, v in kwargs.items():
+            if k in self.__annotations__:
+                setattr(self, k, v)
+            else:
+                raise KeyError(f'Unknown parameter {k}.')
+
+    def _get_updated_attrs(self, **kwargs):
+        current_attrs = {k: getattr(self, k) for k in self.__annotations__.keys()}
+        current_attrs.update(kwargs)
+        return current_attrs

reflectorch/inference/query_matcher.py

@@ -0,0 +1,82 @@
+import os
+import tempfile
+import yaml
+from huggingface_hub import hf_hub_download, list_repo_files
+
+class HuggingfaceQueryMatcher:
+    """Downloads the available configurations files to a temporary directory and provides functionality for filtering those configuration files matching user specified queries.
+    
+    Args:
+        repo_id (str): The Hugging Face repository ID.
+        config_dir (str): Directory within the repo where YAML files are stored.
+    """
+    def __init__(self, repo_id='valentinsingularity/reflectivity', config_dir='configs'):
+        self.repo_id = repo_id
+        self.config_dir = config_dir
+        self.cache = {
+            'parsed_configs': None,
+            'temp_dir': None
+        }
+        self._renew_cache()
+
+    def _renew_cache(self):
+        temp_dir = tempfile.mkdtemp()
+        print(f"Temporary directory created at: {temp_dir}")
+
+        repo_files = list_repo_files(self.repo_id, repo_type='model')
+        config_files = [file for file in repo_files if file.startswith(self.config_dir) and file.endswith('.yaml')]
+
+        downloaded_files = []
+        for file in config_files:
+            file_path = hf_hub_download(repo_id=self.repo_id, filename=file, local_dir=temp_dir, repo_type='model')
+            downloaded_files.append(file_path)
+        
+        parsed_configs = {}
+        for file_path in downloaded_files:
+            with open(file_path, 'r') as file:
+                config_data = yaml.safe_load(file)
+                file_name = os.path.basename(file_path)
+                parsed_configs[file_name] = config_data
+        
+        self.cache['parsed_configs'] = parsed_configs
+        self.cache['temp_dir'] = temp_dir
+
+    def get_matching_configs(self, query):
+        """retrieves configuration files that match the user specified query.
+        
+        Args:
+            query (dict): Dictionary of key-value pairs to filter configurations, e.g. ``query = {'dset.prior_sampler.kwargs.max_num_layers': 3, 'dset.prior_sampler.kwargs.param_ranges.slds': [0., 100.]}``.
+                          For keys containing the ``param_ranges`` subkey a configuration is selected if the value of the query (i.e. desired parameter range) 
+                          is a subrange of the parameter range in the configuration, in all other cases the values must match exactly.
+        
+        Returns:
+            list: List of file names that match the query.
+        """
+        
+        filtered_configs = []
+        
+        for file_name, config_data in self.cache['parsed_configs'].items():
+            if self.matches_query(config_data, query):
+                filtered_configs.append(file_name)
+        
+        return filtered_configs
+    
+    def matches_query(self, config_data, query):
+        for q_key, q_value in query.items():
+            keys = q_key.split('.')
+            value = self.deep_get(config_data, keys)
+            if 'param_ranges' in keys:
+                if q_value[0] < value[0] or q_value[1] > value[1]:
+                    return False
+            else: 
+                if value != q_value:
+                    return False
+                
+        return True
+
+    def deep_get(self, d, keys):
+        for key in keys:
+            if isinstance(d, dict):
+                d = d.get(key, None)
+
+        return d

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