LHSIterator.py

Go to the documentation of this file.

import numpy as np
import pickle
import pandas as pd
from typing import Callable, Dict, List, Any, Union, Iterable
 
class LHSIterator:
    """
    Iterator for Latin Hypercube Sampling (LHS) parameter search.
 
    Parameters
    ----------
     - **suggest_fn** : Callable[[Trial], Dict[str, Any]]
        A function that suggests parameters for each trial. It takes a Trial object as input and returns a dictionary representing the suggested parameters.
     - **n_samples** : int
        The number of samples to generate.
 
    Methods
    -------
     - **save** : Saves the LHSIterator object to a file using pickle serialization.
     - **load** : Loads an LHSIterator object from a file using pickle deserialization.
 
    Trial Methods
    -------------
    ```python
    suggest_categorical(name: str, choices: List)
 
    suggest_int(name: str, low: int, high: int, log: bool = False, step: int = None)
 
    suggest_float(name: str, low: float, high: float, log: bool = False, step: float = None)
    ```
 
 
    Example
    -------
    ```python
    def suggest_params(trial: LHSIterator.Trial) -> Dict[str, Any]:
        return {
            'input':{
                "Cell": {
                    "Common": {
                        'hypothesis': trial.suggest_categorical("hypothesis", [False, True]),
                        'cycleDuration': trial.suggest_float("duration", 10.5, 20.3),
                        'startAge': 0
                        ...
                    }
                }
            }
        }
 
    lhs_iterator = LHSIterator(suggest_params, 10)
 
    # Save the iterator
    lhs_iterator.save("lhs_iterator.pkl")
 
    # Load the iterator
    loaded_iterator = LHSIterator.load("lhs_iterator.pkl")
 
    for i, sample in enumerate(loaded_iterator):
        print(f"Sample {i+1}: {sample}")
    ```
    """
    class Trial:
        def __init__(self, param_space: Dict[str, Any]):
            self.param_space = param_space
            self.params = {}
            self.modeSuggest = False
 
        def suggest_categorical(self, name: str, choices: List[Union[str, int, float]]) -> Any:
            if self.modeSuggest: return self.params[name]
            else:
                if name in self.param_space: raise ValueError(f"Parameter name '{name}' already exists in the parameter space.")
                self.param_space[name] = {"type": "categorical", "choices": choices}
 
        def suggest_int(self, name: str, low: int, high: int, log: bool = False, step: int = None) -> Any:
            if self.modeSuggest: return self.params[name]
            else:
                if name in self.param_space: raise ValueError(f"Parameter name '{name}' already exists in the parameter space.")
                self.param_space[name] = {"type": "int", "low": low, "high": high}
                if log: self.param_space[name]['log'] = log
                if step is not None: self.param_space[name]['step'] = step
 
        def suggest_float(self, name: str, low: float, high: float, log: bool = False, step: float = None) -> Any:
            if self.modeSuggest: return self.params[name]
            else:
                if name in self.param_space: raise ValueError(f"Parameter name '{name}' already exists in the parameter space.")
                self.param_space[name] = {"type": "float", "low": low, "high": high}
                if log: self.param_space[name]['log'] = log
                if step is not None: self.param_space[name]['step'] = step
 
    def __init__(self, suggest_fn: Callable[[Any], Dict[str, Any]], n_samples: int):
        self.param_space: Dict[str, Any] = {}
        self.n_samples: int = n_samples
 
        # Execute the suggestion function with a trial to define the param space
        self.trial = self.Trial(self.param_space)
        self.suggest_fn = suggest_fn
        self.suggest_fn(self.trial)
        self.trial.modeSuggest=True
 
        # Generate LHS samples
        self.lhs_samples = self._lhsclassic(len(self.trial.param_space), self.n_samples)
 
        # Prepare the iterator
        self.index = 0
 
    def __iter__(self) -> 'LHSIterator':
        self.index = 0
        return self
 
    def __len__(self) -> int:
        return self.n_samples
 
    def __next__(self) -> Dict[str, Any]:
        if self.index >= self.n_samples:
            raise StopIteration
        sample = self.lhs_samples[self.index]
        self.trial.params = self._map_sample_to_params(sample)
        self.index += 1
        return self.suggest_fn(self.trial)
 
    def _map_sample_to_params(self, sample: np.ndarray) -> Dict[str, Any]:
        params = {}
        for i, (param_name, param_info) in enumerate(self.trial.param_space.items()):
            if param_info['type'] == 'categorical':
                params[param_name] = param_info['choices'][int(sample[i] * len(param_info['choices']))]
            elif param_info['type'] == 'int':
                if param_info.get('log', False):
                    low = np.log(param_info['low'])
                    high = np.log(param_info['high'])
                    params[param_name] = int(np.exp(sample[i] * (high - low) + low))
                elif 'step' in param_info:
                    params[param_name] = param_info['low'] + int(sample[i] * ((param_info['high'] - param_info['low']) // param_info['step'])) * param_info['step']
                else:
                    params[param_name] = param_info['low'] + int(sample[i] * (param_info['high'] - param_info['low'] + 1))
            elif param_info['type'] == 'float':
                if param_info.get('log', False):
                    low = np.log(param_info['low'])
                    high = np.log(param_info['high'])
                    params[param_name] = np.exp(sample[i] * (high - low) + low)
                elif 'step' in param_info:
                    params[param_name] = param_info['low'] + int(sample[i] * ((param_info['high'] - param_info['low']) / param_info['step'])) * param_info['step']
                else:
                    params[param_name] = param_info['low'] + sample[i] * (param_info['high'] - param_info['low'])
        return params
 
    def _lhsclassic(self, n: int, samples: int) -> np.ndarray:
        # Generate the intervals
        cut = np.linspace(0, 1, samples + 1)    
    
        # Fill points uniformly in each interval
        u = np.random.rand(samples, n)
        a = cut[:samples]
        b = cut[1:samples + 1]
        rdpoints = np.zeros_like(u)
        for j in range(n):
            rdpoints[:, j] = u[:, j]*(b-a) + a
    
        # Make the random pairings
        H = np.zeros_like(rdpoints)
        for j in range(n):
            order = np.random.permutation(range(samples))
            H[:, j] = rdpoints[order, j]
    
        return H
 
    def save(self, filename: str) -> None:
        r"""
        Saves the LHSIterator object to a file using pickle serialization.
 
        Parameters
        ----------
         - **filename** : str
            The name of the file to save the object to.
        """
        with open(filename, 'wb') as f:
            pickle.dump(self, f)
 
    @staticmethod
    def load(filename: str) -> 'LHSIterator':
        r"""
        Loads an LHSIterator object from a file using pickle deserialization.
 
        Parameters
        ----------
         - **filename** : str
            The name of the file to load the object from.
 
        Returns
        -------
         - **LHSIterator**
            The loaded LHSIterator object.
        """
        with open(filename, 'rb') as f:
            return pickle.load(f)
 
    def toDataFrame(self) -> pd.DataFrame:
        r"""
        Converts the generated LHS samples into a pandas DataFrame.
 
        Returns
        -------
         - **pd.DataFrame**
            A pandas DataFrame containing the generated LHS samples, where each row represents a sample and each column represents a parameter.
        """
        return pd.DataFrame([self._map_sample_to_params(s) for s in self.lhs_samples])