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

reflectorch/inference/query_matcher.py

@@ -1,82 +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)
-
+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

@@ -1,43 +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")
+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

@@ -1,56 +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
+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