<filename>gpim/gpbayes/boptim.py
"""
boptim.py
===========

Utility functions for the Gaussian process-based
Bayesian optimization for selecting the next query points in
images and image-like data.

Author: <NAME> (email: <EMAIL>)
"""

import types
import copy
import torch
import numpy as np
from scipy import spatial
from gpim.gpreg import gpr
from gpim.gpbayes import acqfunc
from gpim import gprutils


class boptimizer:
    """
    Gaussian process-based Bayesian optimization for selecting
    next point(s) to measure/evaluate. The Bayesian optimization strategy
    consists of: i) defining prior and posterior distributions over
    the objective (target) function :math:`f` using GP; ii) using the posterior
    to derive an acquistion function :math:`\\alpha (x)`; iii) using the
    acquisition function to derive the next query point according to
    :math:`x_{next}=argmax(\\alpha (x))`; iv) evaluating :math:`f` in
    :math:`x_{next}` and updating the posterior.

    Args:
        X_seed (ndarray):
            Initial seed of (sparse) grid indices with dimensions
            :math:`c \\times N \\times M` or :math:`c \\times N \\times M \\times L`
            where *c* is equal to the number of coordinates
            (for example, for *xyz* coordinates, *c* = 3). Missing coordinates are NaNs.
            Can be potentially larger than 3D but this has not beed tested yet.
        y_seed (ndarray):
            Initial seed with sparse "observations" with dimensions
            :math:`N \\times M` or :math:`N \\times M \\times L`.
            Typically, for 2D image *N* and *M* are image height and width,
            whereas for 3D hyperspectral data *N* and *M* are spatial dimensions
            and *L* is a spectorcopic dimension (e.g. voltage or wavelength).
            Can be potentially larger than 3D but this has not beed tested yet.
        X_full (ndarray):
            Full grid indices (for prediction with a trained GP model)
            with dimensions :math:`N \\times M` or :math:`N \\times M \\times L`.
            Can be potentially larger than 3D but this has not beed tested yet.
        target_function (python function):
            Target (or objective) function. Takes a list of indices and
            returns a function value as a float.
        acquisition_function (str or python function):
            Acquisition function choise.'cb' is confidence bound, 'ei' is
            expected improvement, 'poi' is a probability of improvement.
            One can also pass a custom function, which takes GP model,
            full grid and sparse grid as parameters and returns
            acquisition function values with GP prediction (mean + sd)
        exploration_steps (int):
            Number of exploration-exploitation steps
            (the exploitation-exploration trade-off is
            determined by the acquisition function)
        batch_size (int):
                Number of query points in one batch.
                Returns a single query point when
                batch_update parameter is set to False (default)
        batch_update (bool):
            Filters the query points based on the kernel lengthscale.
            Returns a batch of points when set to True.
            The number of points in the batch may be different from
            batch_size as points are filtered based on the lengthscale.
            In this case the random indices from the initial batch
            will be automatically added to the output.
            Defaults to False.
        kernel (str):
            Kernel type ('RBF', 'Matern52', 'RationalQuadratic').
            Defaults to 'RBF'.
        lengthscale (list of int or list of 2 lists with int):
            Determines lower (1st value or 1st list) and upper (2nd value or 2nd list)
            bounds for kernel lengthscales. For list with two integers,
            the kernel will have only one lenghtscale, even if the dataset
            is multi-dimensional. For lists of two lists, the number of elements
            in each list must be equal to the dataset dimensionality.
        sparse (bool):
            Uses inducing points-based sparse GP regression when set to True.
            False by default.
        indpoints (int):
            Number of inducing points for SparseGPRegression.
            Defaults to total_number_of_points // 10.
        learning_rate (float):
            Learning rate for GP model training
        gp_iterations (int):
            Number of SVI training iteratons for GP model
        seed (int):
            for reproducibility
        **alpha (float or int):
            alpha coefficient in the 'confidence bound' acquisition function
            (Default: 0)
        **beta (float or int):
            beta coefficient in the 'confidence bound' acquisition function
            (Default: 1)
        **xi (float):
            xi coefficient in 'expected improvement'
            and 'probability of improvement' acquisition functions
            (Default: 0.01)
        **use_gpu (bool):
            Uses GPU hardware accelerator when set to 'True'.
            Notice that for large datasets training model without GPU
            is extremely slow. At the same time, for small and (very) sparse data
            there may be no much benefit.
        **precision (str):
            "single" ot "double" floating point precision
        **jitter (float):
            Float between 1e-4 and 1e-6 for numerical stability
        **isotropic (bool):
            Use single lenghtscale for all dimensions in GP (Default: False)
        **mask (ndarray):
            Mask of ones and NaNs (NaNs are values that are not counted when
            searching for acquisition function maximum).
        **dscale (int):
            Distance parameter used in boptimizer.checkvalues.
            It is used in conjuction with 'gamma' and 'memory'
            parameters to select the next query point using the information
            about the location of previous points. It can be understood as a
            short-term memory, where the point at step :math:`s` is chosen no closer
            than :math:`d` to the point selected at step :math:`s-1`, :math:`\\gamma d`
            to the point selected at step :math:`s-2`, :math:`\\gamma^2 d`
            at step :math:`s-3`, and so on. Defaults to 0.
        **batch_dscale (int):
            Distance paramater used in boptimizer.update_points. It is used to
            return a batch of points for a given step, which are no closer
            to each other than batch_dscale value.
            Defauts to kernel average lenghtscale at this step. When used
            in combination with dscale, it will select the *first* queiry
            point in the batch based on the information about the location
            of previous points. THe remianing points will be selected based
            on kernel lengthscale value without taking into account
            the information about the already visited points.
        **batch_out_max (int):
            maximum number of points to output when batch_update=True.
            Note that the number of output points can be smaller than
            this value.
        **gamma (float):
            gamma coefficient, value between 0 and 1.
            Used in boptimizer.checkvalues together with a 'dscale' parameter
            to determine how close the next query point can be to
            the previously selected points.
        **memory (int):
            Number of previously selected points to account for when using 'dscale'.
        **exit_strategy (0 or 1):
            Exit strategy for boptimizer.checkvalues when none
            of the points satisfies the imposed selection criteria.
            1 means that a random value will be chosen, while 0
            means that the last point in the list will be chosen
            (the length of the list is defined by 'batch_size' parameter)
        **extent(list of lists):
            Define bounds for multi-dimensional data. For example, for 2D data,
            the extent parameter is [[xmin, xmax], [ymin, ymax]]
        **save_checkpoints (bool):
            save results to disk after each iteration
        **filename (str):
            name of a file to save checkpoints to
        **verbose (int):
            Level of verbosity (0, 1, or 2)
    """
    def __init__(self,
                 X_seed,
                 y_seed,
                 X_full,
                 target_function,
                 acquisition_function='cb',
                 exploration_steps=10,
                 batch_size=100,
                 batch_update=False,
                 kernel='RBF',
                 lengthscale=None,
                 sparse=False,
                 indpoints=None,
                 gp_iterations=1000,
                 seed=0,
                 **kwargs):
        """
        Initiates Bayesian optimizer parameters
        """
        self.verbose = kwargs.get("verbose", 1)
        self.use_gpu = kwargs.get("use_gpu", False)
        learning_rate = kwargs.get("learning_rate", 5e-2)
        jitter = kwargs.get("jitter", 1.0e-6)
        isotropic = kwargs.get("isotropic", False)
        self.precision = kwargs.get("precision", "double")

        if self.use_gpu and torch.cuda.is_available():
            if self.precision == "single":
                torch.set_default_tensor_type(torch.cuda.FloatTensor)
            else:
                torch.set_default_tensor_type(torch.cuda.DoubleTensor)
        else:
            if self.precision == "single":
                torch.set_default_tensor_type(torch.FloatTensor)
            else:
                torch.set_default_tensor_type(torch.DoubleTensor)

        self.surrogate_model = gpr.reconstructor(
            X_seed, y_seed, X_full, kernel, lengthscale, sparse, indpoints,
            learning_rate, gp_iterations, self.use_gpu, self.verbose, seed,
            isotropic=isotropic, precision=self.precision, jitter=jitter)

        self.X_sparse = X_seed.copy()
        self.y_sparse = y_seed.copy()
        self.X_full = X_full

        self.target_function = target_function
        self.acquisition_function = acquisition_function
        self.exploration_steps = exploration_steps
        self.batch_update = batch_update
        self.batch_size = batch_size
        self.simulate_measurement = kwargs.get("simulate_measurement", False)
        if self.simulate_measurement:
            self.y_true = kwargs.get("y_true")
            if self.y_true is None:
                raise AssertionError(
                    "To simulate measurements, add ground truth ('y_true)")
        self.extent = kwargs.get("extent", None)
        self.alpha, self.beta = kwargs.get("alpha", 0), kwargs.get("beta", 1)
        self.xi = kwargs.get("xi", 0.01)
        self.dscale = kwargs.get("dscale", None)
        self.batch_dscale = kwargs.get("batch_dscale", None)
        self.batch_out_max = kwargs.get("batch_out_max", 10)
        self.gamma = kwargs.get("gamma", 0.8)
        self.points_mem = kwargs.get("memory", 10)
        self.exit_strategy = kwargs.get("exit_strategy", 1)
        self.mask = kwargs.get("mask", None)
        self.save_checkpoints = kwargs.get("save_checkpoints", False)
        self.filename = kwargs.get("filename", "./boptim_results")
        self.indices_all, self.vals_all = [], []
        self.target_func_vals, self.gp_predictions = [y_seed.copy()], []

    def update_posterior(self):
        """
        Updates GP posterior
        """
        X_sparse_new, y_sparse_new = gprutils.prepare_training_data(
            self.X_sparse, self.y_sparse, precision=self.precision)
        if self.use_gpu and torch.cuda.is_available():
            torch.cuda.empty_cache()
            X_sparse_new, y_sparse_new = X_sparse_new.cuda(), y_sparse_new.cuda()
        self.surrogate_model.model.X = X_sparse_new
        self.surrogate_model.model.y = y_sparse_new
        self.surrogate_model.train(verbose=self.verbose)
        return

    def evaluate_function(self, indices, y_measured=None):
        """
        Evaluates target function in the new point(s)
        """
        indices = [indices] if not self.batch_update else indices
        if self.simulate_measurement:
            for idx in indices:
                self.y_sparse[tuple(idx)] = self.y_true[tuple(idx)]
        elif y_measured is not None:
            for idx in indices:
                self.y_sparse[tuple(idx)] = y_measured[tuple(idx)]
        else:
            for idx in indices:
                if self.extent is not None:
                    _idx = []
                    for i, e in zip(idx, self.extent):
                        _idx.append(i + e[0])
                    _idx = tuple(_idx)
                else:
                    _idx = tuple(idx)
                self.y_sparse[tuple(idx)] = self.target_function(_idx)
        self.X_sparse = gprutils.get_sparse_grid(self.y_sparse, self.extent)
        self.target_func_vals.append(self.y_sparse.copy())
        return

    def next_point(self):
        """
        Calculates next query point(s)
        """
        if self.verbose:
            print("Computing acquisition function...")
        if self.acquisition_function == 'cb':
            acq, pred = acqfunc.confidence_bound(
                self.surrogate_model, self.X_full,
                alpha=self.alpha, beta=self.beta)
        elif self.acquisition_function == 'ei':
            acq, pred = acqfunc.expected_improvement(
                self.surrogate_model, self.X_full,
                self.X_sparse, xi=self.xi)
        elif self.acquisition_function == 'poi':
            acq, pred = acqfunc.probability_of_improvement(
                self.surrogate_model, self.X_full,
                self.X_sparse, xi=self.xi)
        elif isinstance(self.acquisition_function, types.FunctionType):
            acq, pred = self.acquisition_function(
                self.surrogate_model, self.X_full, self.X_sparse)
        else:
            raise NotImplementedError(
                "Choose between 'cb', 'ei', and 'poi' acquisition functions or define your own")
        self.gp_predictions.append(pred)
        if self.mask is None:
            indices_list = np.unravel_index(np.argsort(acq.ravel()), acq.shape)
            vals_list = acq[indices_list][::-1][:self.batch_size].tolist()
            indices_list = np.dstack(indices_list)[0][::-1][:self.batch_size].tolist()
        else:
            acq = self.mask*acq
            indices_list = np.unravel_index(np.argsort(acq.ravel()), acq.shape)
            vals_list = acq[indices_list]
            vals_list = vals_list[~np.isnan(vals_list)][::-1]
            indices_list = np.dstack(indices_list)[0]
            indices_list = indices_list[:len(vals_list)][::-1]
            vals_list = vals_list[:self.batch_size].tolist()
            indices_list = indices_list[:self.batch_size].tolist()
        if not self.batch_update:
            return vals_list, indices_list
        if self.batch_dscale is None:
            batch_dscale_ = self.surrogate_model.model.kernel.lengthscale.mean().item()
        else:
            batch_dscale_ = self.batch_dscale
        vals_list, indices_list = self.update_points(
            vals_list, indices_list, batch_dscale_)
        return vals_list, indices_list

    def update_points(self, acqfunc_values, indices, dscale):
        """
        Takes arrays of query points (indices and values) corresponding to
        first *n* max values of the acquisition function and returns a batch
        two lists with updated query points (indices and values) whose
        separation distance is determined by kernel lengthscale.
        Notice that the updated lists will contain less points
        than the original ones.

        Args:
            acqfunc_values (ndarray):
                (*N*,) numpy array with values of acquisition function
            indices (ndarray):
                (*N*, *c*) numpy array with corresponding indices,
                where c ia a number of dimensions of the dataset
            dscale (float):
                kernel lengthscale

        Returns:
            Tuple with computed indices and corresponding values
        """
        ind, val = self.checkvalues(indices, acqfunc_values)
        new_start_idx = np.where(np.array(acqfunc_values) == val)[0][0]
        acqfunc_values = np.array(acqfunc_values)[new_start_idx:]
        indices = np.vstack(indices)[new_start_idx:]
        acqfunc_values_ = copy.deepcopy(acqfunc_values)
        minval = acqfunc_values.min()
        new_max = acqfunc_values.max()
        new_max_id = np.argmax(acqfunc_values)
        max_val_all, max_id_all = [], []
        ck = indices[new_max_id]
        tree = spatial.cKDTree(indices)
        while new_max > minval - 1:
            max_val_all.append(new_max)
            max_id_all.append(new_max_id)
            nn_indices = tree.query_ball_point(ck, dscale)
            acqfunc_values[nn_indices] = minval - 1
            new_max = acqfunc_values.max()
            new_max_id = np.argmax(acqfunc_values)
            ck = indices[new_max_id]
        max_val_all = max_val_all[:self.batch_out_max]
        indices_ = indices[max_id_all].tolist()[:self.batch_out_max]
        if len(indices_) < self.batch_out_max:
            if self.verbose == 2:
                print("Adding {} random indices".format(
                    self.batch_out_max - len(indices_)))
            idx_random = np.random.randint(
                0, len(acqfunc_values), self.batch_out_max - len(indices_))
            indices_.extend(indices[idx_random].tolist())
            max_val_all.extend(acqfunc_values_[idx_random].tolist())
        return max_val_all, indices_

    def checkvalues(self, idx_list, val_list):
        """
        Checks if a current point was already used or if the euclidian
        distances from the current point to the previous points
        are above a value controlled by parameter 'alpha',
        and if so then it selects the next-in-line point.

        Args:
            idx_list (list of lists with ints):
                Indices for max values of acquisition function;
                the list is ordered (max -> min)
            val_list (list with floats):
                Standard deviation values for each index in idx_list

        Returns:
            The first element in the input list (idx_list)
            if no previous occurences found.
            Otherwise, returns the next/closest value from the list.
        """

        def dist(idx):
            idx_prev = self.indices_all[-self.points_mem:]
            # Calculate distances between current point and previous n points
            d_all = [np.linalg.norm(np.array(idx) - np.array(i)) for i in idx_prev]
            # Calculate weighting coefficient for each distance
            dscale_all = [dscale_*self.gamma**i for i in range(len(idx_prev))]
            # Check if each distance satisfies the imposed criteria
            bool_ = 0 in [d > l for (d, l) in zip(d_all[::-1], dscale_all)]
            return bool_

        dscale_ = 0 if self.dscale is None else self.dscale
        _idx = 0
        if self.verbose == 2:
            print('Acquisition function max value {} at {}'.format(
                val_list[_idx], idx_list[_idx]))
        if len(self.indices_all) == 0:
            return idx_list[_idx], val_list[_idx]
        while (1 in [1 for a in self.indices_all if a == idx_list[_idx]]
                or dist(idx_list[_idx])):
            if self.verbose == 2:
                print("Finding the next max point...")
            _idx = _idx + 1
            if _idx == len(idx_list):
                _idx = np.random.randint(0, len(idx_list)) if self.exit_strategy else -1
                if self.verbose == 2:
                    print('Index out of list. Exiting with acquisition function value {} at {}'.format(
                        val_list[_idx], idx_list[_idx]))
                break
            if self.verbose == 2:
                print('Acquisition function max value {} at {}'.format(
                    val_list[_idx], idx_list[_idx]))
        return idx_list[_idx], val_list[_idx]

    def single_step(self, *args):
        """
        Single Bayesian optimization step
        """
        e = args[0]
        if self.verbose:
            print("\nExploration step {} / {}".format(
                e+1, self.exploration_steps))
        # train with seeded data
        if e == 0:
            self.surrogate_model.train()
        # calculate acquisition function and get next query points
        vals, inds = self.next_point()
        if not self.batch_update:
            inds, vals = self.checkvalues(inds, vals)
        # evaluate function
        self.evaluate_function(inds)
        # update posterior
        self.update_posterior()
        # store indices and values
        if isinstance(vals, float):
            self.indices_all.append(inds)
            self.vals_all.append(vals)
        else:
            self.indices_all.extend(inds)
            self.vals_all.extend(vals)
        return

    def run(self):
        """
        Run GP-based Bayesian optimization loop
        """
        for i in range(self.exploration_steps):
            self.single_step(i)
            if self.save_checkpoints:
                self.save_results()
        self.save_results()
        if self.verbose:
            print("\nExploration completed")
        return

    def save_results(self, *args):
        """
        Save indermediary and/or final results
        """
        try:
            filename = args[0]
        except IndexError:
            filename = self.filename
        results = {}
        results['gp_pred'] = self.gp_predictions
        results['func_val'] = self.target_func_vals
        results['inds_all'] = np.array(self.indices_all)
        results['vals_all'] = np.array(self.vals_all)
        np.save(filename+".npy", results)

<filename>gpim/gpbayes/boptim.py
"""
boptim.py
===========

Utility functions for the Gaussian process-based
Bayesian optimization for selecting the next query points in
images and image-like data.

Author: <NAME> (email: <EMAIL>)
"""

import types
import copy
import torch
import numpy as np
from scipy import spatial
from gpim.gpreg import gpr
from gpim.gpbayes import acqfunc
from gpim import gprutils


class boptimizer:
    """
    Gaussian process-based Bayesian optimization for selecting
    next point(s) to measure/evaluate. The Bayesian optimization strategy
    consists of: i) defining prior and posterior distributions over
    the objective (target) function :math:`f` using GP; ii) using the posterior
    to derive an acquistion function :math:`\\alpha (x)`; iii) using the
    acquisition function to derive the next query point according to
    :math:`x_{next}=argmax(\\alpha (x))`; iv) evaluating :math:`f` in
    :math:`x_{next}` and updating the posterior.

    Args:
        X_seed (ndarray):
            Initial seed of (sparse) grid indices with dimensions
            :math:`c \\times N \\times M` or :math:`c \\times N \\times M \\times L`
            where *c* is equal to the number of coordinates
            (for example, for *xyz* coordinates, *c* = 3). Missing coordinates are NaNs.
            Can be potentially larger than 3D but this has not beed tested yet.
        y_seed (ndarray):
            Initial seed with sparse "observations" with dimensions
            :math:`N \\times M` or :math:`N \\times M \\times L`.
            Typically, for 2D image *N* and *M* are image height and width,
            whereas for 3D hyperspectral data *N* and *M* are spatial dimensions
            and *L* is a spectorcopic dimension (e.g. voltage or wavelength).
            Can be potentially larger than 3D but this has not beed tested yet.
        X_full (ndarray):
            Full grid indices (for prediction with a trained GP model)
            with dimensions :math:`N \\times M` or :math:`N \\times M \\times L`.
            Can be potentially larger than 3D but this has not beed tested yet.
        target_function (python function):
            Target (or objective) function. Takes a list of indices and
            returns a function value as a float.
        acquisition_function (str or python function):
            Acquisition function choise.'cb' is confidence bound, 'ei' is
            expected improvement, 'poi' is a probability of improvement.
            One can also pass a custom function, which takes GP model,
            full grid and sparse grid as parameters and returns
            acquisition function values with GP prediction (mean + sd)
        exploration_steps (int):
            Number of exploration-exploitation steps
            (the exploitation-exploration trade-off is
            determined by the acquisition function)
        batch_size (int):
                Number of query points in one batch.
                Returns a single query point when
                batch_update parameter is set to False (default)
        batch_update (bool):
            Filters the query points based on the kernel lengthscale.
            Returns a batch of points when set to True.
            The number of points in the batch may be different from
            batch_size as points are filtered based on the lengthscale.
            In this case the random indices from the initial batch
            will be automatically added to the output.
            Defaults to False.
        kernel (str):
            Kernel type ('RBF', 'Matern52', 'RationalQuadratic').
            Defaults to 'RBF'.
        lengthscale (list of int or list of 2 lists with int):
            Determines lower (1st value or 1st list) and upper (2nd value or 2nd list)
            bounds for kernel lengthscales. For list with two integers,
            the kernel will have only one lenghtscale, even if the dataset
            is multi-dimensional. For lists of two lists, the number of elements
            in each list must be equal to the dataset dimensionality.
        sparse (bool):
            Uses inducing points-based sparse GP regression when set to True.
            False by default.
        indpoints (int):
            Number of inducing points for SparseGPRegression.
            Defaults to total_number_of_points // 10.
        learning_rate (float):
            Learning rate for GP model training
        gp_iterations (int):
            Number of SVI training iteratons for GP model
        seed (int):
            for reproducibility
        **alpha (float or int):
            alpha coefficient in the 'confidence bound' acquisition function
            (Default: 0)
        **beta (float or int):
            beta coefficient in the 'confidence bound' acquisition function
            (Default: 1)
        **xi (float):
            xi coefficient in 'expected improvement'
            and 'probability of improvement' acquisition functions
            (Default: 0.01)
        **use_gpu (bool):
            Uses GPU hardware accelerator when set to 'True'.
            Notice that for large datasets training model without GPU
            is extremely slow. At the same time, for small and (very) sparse data
            there may be no much benefit.
        **precision (str):
            "single" ot "double" floating point precision
        **jitter (float):
            Float between 1e-4 and 1e-6 for numerical stability
        **isotropic (bool):
            Use single lenghtscale for all dimensions in GP (Default: False)
        **mask (ndarray):
            Mask of ones and NaNs (NaNs are values that are not counted when
            searching for acquisition function maximum).
        **dscale (int):
            Distance parameter used in boptimizer.checkvalues.
            It is used in conjuction with 'gamma' and 'memory'
            parameters to select the next query point using the information
            about the location of previous points. It can be understood as a
            short-term memory, where the point at step :math:`s` is chosen no closer
            than :math:`d` to the point selected at step :math:`s-1`, :math:`\\gamma d`
            to the point selected at step :math:`s-2`, :math:`\\gamma^2 d`
            at step :math:`s-3`, and so on. Defaults to 0.
        **batch_dscale (int):
            Distance paramater used in boptimizer.update_points. It is used to
            return a batch of points for a given step, which are no closer
            to each other than batch_dscale value.
            Defauts to kernel average lenghtscale at this step. When used
            in combination with dscale, it will select the *first* queiry
            point in the batch based on the information about the location
            of previous points. THe remianing points will be selected based
            on kernel lengthscale value without taking into account
            the information about the already visited points.
        **batch_out_max (int):
            maximum number of points to output when batch_update=True.
            Note that the number of output points can be smaller than
            this value.
        **gamma (float):
            gamma coefficient, value between 0 and 1.
            Used in boptimizer.checkvalues together with a 'dscale' parameter
            to determine how close the next query point can be to
            the previously selected points.
        **memory (int):
            Number of previously selected points to account for when using 'dscale'.
        **exit_strategy (0 or 1):
            Exit strategy for boptimizer.checkvalues when none
            of the points satisfies the imposed selection criteria.
            1 means that a random value will be chosen, while 0
            means that the last point in the list will be chosen
            (the length of the list is defined by 'batch_size' parameter)
        **extent(list of lists):
            Define bounds for multi-dimensional data. For example, for 2D data,
            the extent parameter is [[xmin, xmax], [ymin, ymax]]
        **save_checkpoints (bool):
            save results to disk after each iteration
        **filename (str):
            name of a file to save checkpoints to
        **verbose (int):
            Level of verbosity (0, 1, or 2)
    """
    def __init__(self,
                 X_seed,
                 y_seed,
                 X_full,
                 target_function,
                 acquisition_function='cb',
                 exploration_steps=10,
                 batch_size=100,
                 batch_update=False,
                 kernel='RBF',
                 lengthscale=None,
                 sparse=False,
                 indpoints=None,
                 gp_iterations=1000,
                 seed=0,
                 **kwargs):
        """
        Initiates Bayesian optimizer parameters
        """
        self.verbose = kwargs.get("verbose", 1)
        self.use_gpu = kwargs.get("use_gpu", False)
        learning_rate = kwargs.get("learning_rate", 5e-2)
        jitter = kwargs.get("jitter", 1.0e-6)
        isotropic = kwargs.get("isotropic", False)
        self.precision = kwargs.get("precision", "double")

        if self.use_gpu and torch.cuda.is_available():
            if self.precision == "single":
                torch.set_default_tensor_type(torch.cuda.FloatTensor)
            else:
                torch.set_default_tensor_type(torch.cuda.DoubleTensor)
        else:
            if self.precision == "single":
                torch.set_default_tensor_type(torch.FloatTensor)
            else:
                torch.set_default_tensor_type(torch.DoubleTensor)

        self.surrogate_model = gpr.reconstructor(
            X_seed, y_seed, X_full, kernel, lengthscale, sparse, indpoints,
            learning_rate, gp_iterations, self.use_gpu, self.verbose, seed,
            isotropic=isotropic, precision=self.precision, jitter=jitter)

        self.X_sparse = X_seed.copy()
        self.y_sparse = y_seed.copy()
        self.X_full = X_full

        self.target_function = target_function
        self.acquisition_function = acquisition_function
        self.exploration_steps = exploration_steps
        self.batch_update = batch_update
        self.batch_size = batch_size
        self.simulate_measurement = kwargs.get("simulate_measurement", False)
        if self.simulate_measurement:
            self.y_true = kwargs.get("y_true")
            if self.y_true is None:
                raise AssertionError(
                    "To simulate measurements, add ground truth ('y_true)")
        self.extent = kwargs.get("extent", None)
        self.alpha, self.beta = kwargs.get("alpha", 0), kwargs.get("beta", 1)
        self.xi = kwargs.get("xi", 0.01)
        self.dscale = kwargs.get("dscale", None)
        self.batch_dscale = kwargs.get("batch_dscale", None)
        self.batch_out_max = kwargs.get("batch_out_max", 10)
        self.gamma = kwargs.get("gamma", 0.8)
        self.points_mem = kwargs.get("memory", 10)
        self.exit_strategy = kwargs.get("exit_strategy", 1)
        self.mask = kwargs.get("mask", None)
        self.save_checkpoints = kwargs.get("save_checkpoints", False)
        self.filename = kwargs.get("filename", "./boptim_results")
        self.indices_all, self.vals_all = [], []
        self.target_func_vals, self.gp_predictions = [y_seed.copy()], []

    def update_posterior(self):
        """
        Updates GP posterior
        """
        X_sparse_new, y_sparse_new = gprutils.prepare_training_data(
            self.X_sparse, self.y_sparse, precision=self.precision)
        if self.use_gpu and torch.cuda.is_available():
            torch.cuda.empty_cache()
            X_sparse_new, y_sparse_new = X_sparse_new.cuda(), y_sparse_new.cuda()
        self.surrogate_model.model.X = X_sparse_new
        self.surrogate_model.model.y = y_sparse_new
        self.surrogate_model.train(verbose=self.verbose)
        return

    def evaluate_function(self, indices, y_measured=None):
        """
        Evaluates target function in the new point(s)
        """
        indices = [indices] if not self.batch_update else indices
        if self.simulate_measurement:
            for idx in indices:
                self.y_sparse[tuple(idx)] = self.y_true[tuple(idx)]
        elif y_measured is not None:
            for idx in indices:
                self.y_sparse[tuple(idx)] = y_measured[tuple(idx)]
        else:
            for idx in indices:
                if self.extent is not None:
                    _idx = []
                    for i, e in zip(idx, self.extent):
                        _idx.append(i + e[0])
                    _idx = tuple(_idx)
                else:
                    _idx = tuple(idx)
                self.y_sparse[tuple(idx)] = self.target_function(_idx)
        self.X_sparse = gprutils.get_sparse_grid(self.y_sparse, self.extent)
        self.target_func_vals.append(self.y_sparse.copy())
        return

    def next_point(self):
        """
        Calculates next query point(s)
        """
        if self.verbose:
            print("Computing acquisition function...")
        if self.acquisition_function == 'cb':
            acq, pred = acqfunc.confidence_bound(
                self.surrogate_model, self.X_full,
                alpha=self.alpha, beta=self.beta)
        elif self.acquisition_function == 'ei':
            acq, pred = acqfunc.expected_improvement(
                self.surrogate_model, self.X_full,
                self.X_sparse, xi=self.xi)
        elif self.acquisition_function == 'poi':
            acq, pred = acqfunc.probability_of_improvement(
                self.surrogate_model, self.X_full,
                self.X_sparse, xi=self.xi)
        elif isinstance(self.acquisition_function, types.FunctionType):
            acq, pred = self.acquisition_function(
                self.surrogate_model, self.X_full, self.X_sparse)
        else:
            raise NotImplementedError(
                "Choose between 'cb', 'ei', and 'poi' acquisition functions or define your own")
        self.gp_predictions.append(pred)
        if self.mask is None:
            indices_list = np.unravel_index(np.argsort(acq.ravel()), acq.shape)
            vals_list = acq[indices_list][::-1][:self.batch_size].tolist()
            indices_list = np.dstack(indices_list)[0][::-1][:self.batch_size].tolist()
        else:
            acq = self.mask*acq
            indices_list = np.unravel_index(np.argsort(acq.ravel()), acq.shape)
            vals_list = acq[indices_list]
            vals_list = vals_list[~np.isnan(vals_list)][::-1]
            indices_list = np.dstack(indices_list)[0]
            indices_list = indices_list[:len(vals_list)][::-1]
            vals_list = vals_list[:self.batch_size].tolist()
            indices_list = indices_list[:self.batch_size].tolist()
        if not self.batch_update:
            return vals_list, indices_list
        if self.batch_dscale is None:
            batch_dscale_ = self.surrogate_model.model.kernel.lengthscale.mean().item()
        else:
            batch_dscale_ = self.batch_dscale
        vals_list, indices_list = self.update_points(
            vals_list, indices_list, batch_dscale_)
        return vals_list, indices_list

    def update_points(self, acqfunc_values, indices, dscale):
        """
        Takes arrays of query points (indices and values) corresponding to
        first *n* max values of the acquisition function and returns a batch
        two lists with updated query points (indices and values) whose
        separation distance is determined by kernel lengthscale.
        Notice that the updated lists will contain less points
        than the original ones.

        Args:
            acqfunc_values (ndarray):
                (*N*,) numpy array with values of acquisition function
            indices (ndarray):
                (*N*, *c*) numpy array with corresponding indices,
                where c ia a number of dimensions of the dataset
            dscale (float):
                kernel lengthscale

        Returns:
            Tuple with computed indices and corresponding values
        """
        ind, val = self.checkvalues(indices, acqfunc_values)
        new_start_idx = np.where(np.array(acqfunc_values) == val)[0][0]
        acqfunc_values = np.array(acqfunc_values)[new_start_idx:]
        indices = np.vstack(indices)[new_start_idx:]
        acqfunc_values_ = copy.deepcopy(acqfunc_values)
        minval = acqfunc_values.min()
        new_max = acqfunc_values.max()
        new_max_id = np.argmax(acqfunc_values)
        max_val_all, max_id_all = [], []
        ck = indices[new_max_id]
        tree = spatial.cKDTree(indices)
        while new_max > minval - 1:
            max_val_all.append(new_max)
            max_id_all.append(new_max_id)
            nn_indices = tree.query_ball_point(ck, dscale)
            acqfunc_values[nn_indices] = minval - 1
            new_max = acqfunc_values.max()
            new_max_id = np.argmax(acqfunc_values)
            ck = indices[new_max_id]
        max_val_all = max_val_all[:self.batch_out_max]
        indices_ = indices[max_id_all].tolist()[:self.batch_out_max]
        if len(indices_) < self.batch_out_max:
            if self.verbose == 2:
                print("Adding {} random indices".format(
                    self.batch_out_max - len(indices_)))
            idx_random = np.random.randint(
                0, len(acqfunc_values), self.batch_out_max - len(indices_))
            indices_.extend(indices[idx_random].tolist())
            max_val_all.extend(acqfunc_values_[idx_random].tolist())
        return max_val_all, indices_

    def checkvalues(self, idx_list, val_list):
        """
        Checks if a current point was already used or if the euclidian
        distances from the current point to the previous points
        are above a value controlled by parameter 'alpha',
        and if so then it selects the next-in-line point.

        Args:
            idx_list (list of lists with ints):
                Indices for max values of acquisition function;
                the list is ordered (max -> min)
            val_list (list with floats):
                Standard deviation values for each index in idx_list

        Returns:
            The first element in the input list (idx_list)
            if no previous occurences found.
            Otherwise, returns the next/closest value from the list.
        """

        def dist(idx):
            idx_prev = self.indices_all[-self.points_mem:]
            # Calculate distances between current point and previous n points
            d_all = [np.linalg.norm(np.array(idx) - np.array(i)) for i in idx_prev]
            # Calculate weighting coefficient for each distance
            dscale_all = [dscale_*self.gamma**i for i in range(len(idx_prev))]
            # Check if each distance satisfies the imposed criteria
            bool_ = 0 in [d > l for (d, l) in zip(d_all[::-1], dscale_all)]
            return bool_

        dscale_ = 0 if self.dscale is None else self.dscale
        _idx = 0
        if self.verbose == 2:
            print('Acquisition function max value {} at {}'.format(
                val_list[_idx], idx_list[_idx]))
        if len(self.indices_all) == 0:
            return idx_list[_idx], val_list[_idx]
        while (1 in [1 for a in self.indices_all if a == idx_list[_idx]]
                or dist(idx_list[_idx])):
            if self.verbose == 2:
                print("Finding the next max point...")
            _idx = _idx + 1
            if _idx == len(idx_list):
                _idx = np.random.randint(0, len(idx_list)) if self.exit_strategy else -1
                if self.verbose == 2:
                    print('Index out of list. Exiting with acquisition function value {} at {}'.format(
                        val_list[_idx], idx_list[_idx]))
                break
            if self.verbose == 2:
                print('Acquisition function max value {} at {}'.format(
                    val_list[_idx], idx_list[_idx]))
        return idx_list[_idx], val_list[_idx]

    def single_step(self, *args):
        """
        Single Bayesian optimization step
        """
        e = args[0]
        if self.verbose:
            print("\nExploration step {} / {}".format(
                e+1, self.exploration_steps))
        # train with seeded data
        if e == 0:
            self.surrogate_model.train()
        # calculate acquisition function and get next query points
        vals, inds = self.next_point()
        if not self.batch_update:
            inds, vals = self.checkvalues(inds, vals)
        # evaluate function
        self.evaluate_function(inds)
        # update posterior
        self.update_posterior()
        # store indices and values
        if isinstance(vals, float):
            self.indices_all.append(inds)
            self.vals_all.append(vals)
        else:
            self.indices_all.extend(inds)
            self.vals_all.extend(vals)
        return

    def run(self):
        """
        Run GP-based Bayesian optimization loop
        """
        for i in range(self.exploration_steps):
            self.single_step(i)
            if self.save_checkpoints:
                self.save_results()
        self.save_results()
        if self.verbose:
            print("\nExploration completed")
        return

    def save_results(self, *args):
        """
        Save indermediary and/or final results
        """
        try:
            filename = args[0]
        except IndexError:
            filename = self.filename
        results = {}
        results['gp_pred'] = self.gp_predictions
        results['func_val'] = self.target_func_vals
        results['inds_all'] = np.array(self.indices_all)
        results['vals_all'] = np.array(self.vals_all)
        np.save(filename+".npy", results)

boptim.py ===========  Utility functions for the Gaussian process-based Bayesian optimization for selecting the next query points in images and image-like data.  Author: <NAME> (email: <EMAIL>) Gaussian process-based Bayesian optimization for selecting     next point(s) to measure/evaluate. The Bayesian optimization strategy     consists of: i) defining prior and posterior distributions over     the objective (target) function :math:`f` using GP; ii) using the posterior     to derive an acquistion function :math:`\\alpha (x)`; iii) using the     acquisition function to derive the next query point according to     :math:`x_{next}=argmax(\\alpha (x))`; iv) evaluating :math:`f` in     :math:`x_{next}` and updating the posterior.      Args:         X_seed (ndarray):             Initial seed of (sparse) grid indices with dimensions             :math:`c \\times N \\times M` or :math:`c \\times N \\times M \\times L`             where *c* is equal to the number of coordinates             (for example, for *xyz* coordinates, *c* = 3). Missing coordinates are NaNs.             Can be potentially larger than 3D but this has not beed tested yet.         y_seed (ndarray):             Initial seed with sparse "observations" with dimensions             :math:`N \\times M` or :math:`N \\times M \\times L`.             Typically, for 2D image *N* and *M* are image height and width,             whereas for 3D hyperspectral data *N* and *M* are spatial dimensions             and *L* is a spectorcopic dimension (e.g. voltage or wavelength).             Can be potentially larger than 3D but this has not beed tested yet.         X_full (ndarray):             Full grid indices (for prediction with a trained GP model)             with dimensions :math:`N \\times M` or :math:`N \\times M \\times L`.             Can be potentially larger than 3D but this has not beed tested yet.         target_function (python function):             Target (or objective) function. Takes a list of indices and             returns a function value as a float.         acquisition_function (str or python function):             Acquisition function choise.'cb' is confidence bound, 'ei' is             expected improvement, 'poi' is a probability of improvement.             One can also pass a custom function, which takes GP model,             full grid and sparse grid as parameters and returns             acquisition function values with GP prediction (mean + sd)         exploration_steps (int):             Number of exploration-exploitation steps             (the exploitation-exploration trade-off is             determined by the acquisition function)         batch_size (int):                 Number of query points in one batch.                 Returns a single query point when                 batch_update parameter is set to False (default)         batch_update (bool):             Filters the query points based on the kernel lengthscale.             Returns a batch of points when set to True.             The number of points in the batch may be different from             batch_size as points are filtered based on the lengthscale.             In this case the random indices from the initial batch             will be automatically added to the output.             Defaults to False.         kernel (str):             Kernel type ('RBF', 'Matern52', 'RationalQuadratic').             Defaults to 'RBF'.         lengthscale (list of int or list of 2 lists with int):             Determines lower (1st value or 1st list) and upper (2nd value or 2nd list)             bounds for kernel lengthscales. For list with two integers,             the kernel will have only one lenghtscale, even if the dataset             is multi-dimensional. For lists of two lists, the number of elements             in each list must be equal to the dataset dimensionality.         sparse (bool):             Uses inducing points-based sparse GP regression when set to True.             False by default.         indpoints (int):             Number of inducing points for SparseGPRegression.             Defaults to total_number_of_points // 10.         learning_rate (float):             Learning rate for GP model training         gp_iterations (int):             Number of SVI training iteratons for GP model         seed (int):             for reproducibility         **alpha (float or int):             alpha coefficient in the 'confidence bound' acquisition function             (Default: 0)         **beta (float or int):             beta coefficient in the 'confidence bound' acquisition function             (Default: 1)         **xi (float):             xi coefficient in 'expected improvement'             and 'probability of improvement' acquisition functions             (Default: 0.01)         **use_gpu (bool):             Uses GPU hardware accelerator when set to 'True'.             Notice that for large datasets training model without GPU             is extremely slow. At the same time, for small and (very) sparse data             there may be no much benefit.         **precision (str):             "single" ot "double" floating point precision         **jitter (float):             Float between 1e-4 and 1e-6 for numerical stability         **isotropic (bool):             Use single lenghtscale for all dimensions in GP (Default: False)         **mask (ndarray):             Mask of ones and NaNs (NaNs are values that are not counted when             searching for acquisition function maximum).         **dscale (int):             Distance parameter used in boptimizer.checkvalues.             It is used in conjuction with 'gamma' and 'memory'             parameters to select the next query point using the information             about the location of previous points. It can be understood as a             short-term memory, where the point at step :math:`s` is chosen no closer             than :math:`d` to the point selected at step :math:`s-1`, :math:`\\gamma d`             to the point selected at step :math:`s-2`, :math:`\\gamma^2 d`             at step :math:`s-3`, and so on. Defaults to 0.         **batch_dscale (int):             Distance paramater used in boptimizer.update_points. It is used to             return a batch of points for a given step, which are no closer             to each other than batch_dscale value.             Defauts to kernel average lenghtscale at this step. When used             in combination with dscale, it will select the *first* queiry             point in the batch based on the information about the location             of previous points. THe remianing points will be selected based             on kernel lengthscale value without taking into account             the information about the already visited points.         **batch_out_max (int):             maximum number of points to output when batch_update=True.             Note that the number of output points can be smaller than             this value.         **gamma (float):             gamma coefficient, value between 0 and 1.             Used in boptimizer.checkvalues together with a 'dscale' parameter             to determine how close the next query point can be to             the previously selected points.         **memory (int):             Number of previously selected points to account for when using 'dscale'.         **exit_strategy (0 or 1):             Exit strategy for boptimizer.checkvalues when none             of the points satisfies the imposed selection criteria.             1 means that a random value will be chosen, while 0             means that the last point in the list will be chosen             (the length of the list is defined by 'batch_size' parameter)         **extent(list of lists):             Define bounds for multi-dimensional data. For example, for 2D data,             the extent parameter is [[xmin, xmax], [ymin, ymax]]         **save_checkpoints (bool):             save results to disk after each iteration         **filename (str):             name of a file to save checkpoints to         **verbose (int):             Level of verbosity (0, 1, or 2) Initiates Bayesian optimizer parameters Updates GP posterior Evaluates target function in the new point(s) Calculates next query point(s) Takes arrays of query points (indices and values) corresponding to         first *n* max values of the acquisition function and returns a batch         two lists with updated query points (indices and values) whose         separation distance is determined by kernel lengthscale.         Notice that the updated lists will contain less points         than the original ones.          Args:             acqfunc_values (ndarray):                 (*N*,) numpy array with values of acquisition function             indices (ndarray):                 (*N*, *c*) numpy array with corresponding indices,                 where c ia a number of dimensions of the dataset             dscale (float):                 kernel lengthscale          Returns:             Tuple with computed indices and corresponding values Checks if a current point was already used or if the euclidian         distances from the current point to the previous points         are above a value controlled by parameter 'alpha',         and if so then it selects the next-in-line point.          Args:             idx_list (list of lists with ints):                 Indices for max values of acquisition function;                 the list is ordered (max -> min)             val_list (list with floats):                 Standard deviation values for each index in idx_list          Returns:             The first element in the input list (idx_list)             if no previous occurences found.             Otherwise, returns the next/closest value from the list. # Calculate distances between current point and previous n points # Calculate weighting coefficient for each distance # Check if each distance satisfies the imposed criteria Single Bayesian optimization step # train with seeded data # calculate acquisition function and get next query points # evaluate function # update posterior # store indices and values Run GP-based Bayesian optimization loop Save indermediary and/or final results