
import copy

import matplotlib.pyplot as plt
import numpy as np

from lmfit import minimize, fit_report, Minimizer, Parameters


# Let's start with the [example](https://docs.scipy.org/doc/scipy/reference/generated/scipy.optimize.brute.html) given in the documentation of SciPy: "We illustrate the use of brute to seek the global minimum of a function of two variables that is given as the sum of a positive-definite quadratic and two deep “Gaussian-shaped” craters. Specifically, define the objective function f as the sum of three other functions, ``f = f1 + f2 + f3``. We suppose each of these has a signature ``(z, *params), where z = (x, y)``, and params and the functions are as defined below."

# First, we create a set of Parameters where all variables except ``x`` and ``y`` are given fixed values.

# In[3]:


params = Parameters()
params.add_many(
        ('a', 2, False),
        ('b', 3, False),
        ('c', 7, False),
        ('d', 8, False),
        ('e', 9, False),
        ('f', 10, False),
        ('g', 44, False),
        ('h', -1, False),
        ('i', 2, False),
        ('j', 26, False),
        ('k', 1, False),
        ('l', -2, False),
        ('scale', 0.5, False),
        ('x', 0.0, True),
        ('y', 0.0, True))


# Second, create the three functions and the objective function:

# In[4]:


def f1(p):
    par = p.valuesdict()
    return (par['a'] * par['x']**2 + par['b'] * par['x'] * par['y'] +
            par['c'] * par['y']**2 + par['d']*par['x'] + par['e']*par['y'] + par['f'])


def f2(p):
    par = p.valuesdict()
    return (-1.0*par['g']*np.exp(-((par['x']-par['h'])**2 +
            (par['y']-par['i'])**2) / par['scale']))


def f3(p):
    par = p.valuesdict()
    return (-1.0*par['j']*np.exp(-((par['x']-par['k'])**2 +
            (par['y']-par['l'])**2) / par['scale']))


def f(params):
    return f1(params) + f2(params) + f3(params)


#    Just as in the documentation we will do a grid search between ``-4`` and ``4`` and use a stepsize of ``0.25``. The bounds can be set as usual with the ``min`` and ``max`` attributes, and the stepsize is set using ``brute_step``.

# In[5]:


params['x'].set(min=-4, max=4, brute_step=0.25)
params['y'].set(min=-4, max=4, brute_step=0.25)


# Performing a grid search is done with:

# In[6]:


fitter = Minimizer(f, params)
result = fitter.minimize(method='brute')


# , which will increment ``x`` and ``y`` between ``-4`` in increments of ``0.25`` until ``4`` (not inclusive).

# In[7]:


grid_x, grid_y = [np.unique(par.ravel()) for par in result.brute_grid]
print(grid_x)


# The objective function is evaluated on this grid, and the raw output from ``scipy.optimize.brute`` is stored in the MinimizerResult as ``brute_<parname>`` attributes. These attributes are:
#
# ``result.brute_x0`` -- A 1-D array containing the coordinates of a point at which the objective function had its minimum value.

# In[8]:


result.brute_x0


# ``result.brute_fval`` -- Function value at the point x0.

# In[9]:


result.brute_fval


# ``result.brute_grid`` -- Representation of the evaluation grid. It has the same length as x0.

# In[10]:


result.brute_grid


# ``result.brute_Jout`` -- Function values at each point of the evaluation grid, i.e., Jout = func(*grid).

# In[11]:


result.brute_Jout


# **Reassuringly, the obtained results are indentical to using the method in SciPy directly!**

# ## Example: fit a decaying sine wave

# In this example, will explain some of the options
#
# We start off by generating some synthetic data with noise for a decaying sine wave, define an objective function and create a Parameter set.

# In[12]:


# create data to be fitted
x = np.linspace(0, 15, 301)
np.random.seed(7)
data = (5. * np.sin(2*x - 0.1) * np.exp(-x*x*0.025) +
        np.random.normal(size=len(x), scale=0.2))
plt.plot(x, data, 'b')
plt.show()

# In[13]:


# define objective function: returns the array to be minimized
def fcn2min(params, x, data):
    """Model decaying sine wave, subtract data."""
    amp = params['amp']
    shift = params['shift']
    omega = params['omega']
    decay = params['decay']
    model = amp * np.sin(x*omega + shift) * np.exp(-x*x*decay)
    return model - data


# In[14]:


# create a set of Parameters
params = Parameters()
params.add('amp', value=7, min=2.5)
params.add('decay', value=0.05)
params.add('shift', value=0.0, min=-np.pi/2., max=np.pi/2)
params.add('omega', value=3, max=5)


# In contrast to the implementation in SciPy (as shown in the first example), varying parameters do not need to have finite bounds in lmfit. However, in that case they **do** need the ``brute_step`` attribute specified, so let's do that:

# In[15]:


params['amp'].set(brute_step=0.25)
params['decay'].set(brute_step=0.005)
params['omega'].set(brute_step=0.25)


# Our initial parameter set is now defined as shown below and this will determine how the grid is set-up.

# In[16]:


params.pretty_print()


# First, we initialize a Minimizer and perform the grid search:

# In[17]:


fitter = Minimizer(fcn2min, params, fcn_args=(x, data))
result_brute = fitter.minimize(method='brute', Ns=25, keep=25)


# We used two new parameters here: ``Ns`` and ``keep``. The parameter ``Ns`` determines the 'number of grid points along the axes' similarly to its usage in SciPy. Together with ``brute_step``, ``min`` and ``max`` for a Parameter it will dictate how the grid is set-up:
#
# **(1)** finite bounds are specified ("SciPy implementation"): uses ``brute_step`` if present (in the example above) or uses ``Ns`` to generate the grid. The latter scenario that interpolates ``Ns`` points from ``min`` to ``max`` (inclusive), is here shown for the parameter ``shift``:

# In[18]:


par_name = 'shift'
indx_shift = result_brute.var_names.index(par_name)
grid_shift = np.unique(result_brute.brute_grid[indx_shift].ravel())
print("parameter = {}: \nnumber of steps = {}\ngrid = {}".format(par_name, len(grid_shift), grid_shift))


# If finite bounds are not set for a certain parameter then the user **must** specify ``brute_step`` - three more scenarios are considered here:
#
# **(2)** lower bound (min) and brute_step are specified: range = (min, min + Ns * brute_step, brute_step)

# In[19]:


par_name = 'amp'
indx_shift = result_brute.var_names.index(par_name)
grid_shift = np.unique(result_brute.brute_grid[indx_shift].ravel())
print("parameter = {}: \nnumber of steps = {}\ngrid = {}".format(par_name, len(grid_shift), grid_shift))


# **(3)** upper bound (max) and brute_step are specified: range = (max - Ns * brute_step, max, brute_step)

# In[20]:


par_name = 'omega'
indx_shift = result_brute.var_names.index(par_name)
grid_shift = np.unique(result_brute.brute_grid[indx_shift].ravel())
print("parameter = {}: \nnumber of steps = {}\ngrid = {}".format(par_name, len(grid_shift), grid_shift))


# **(4)** numerical value (value) and brute_step are specified: range = (value - (Ns//2) * brute_step, value + (Ns//2) * brute_step, brute_step)

# In[21]:


par_name = 'decay'
indx_shift = result_brute.var_names.index(par_name)
grid_shift = np.unique(result_brute.brute_grid[indx_shift].ravel())
print("parameter = {}: \nnumber of steps = {}\ngrid = {}".format(par_name, len(grid_shift), grid_shift))


# The [``MinimizerResult``](http://lmfit.github.io/lmfit-py/fitting.html#lmfit.minimizer.MinimizerResult) contains all the usual best-fit parameters and fitting statistics. For example, the optimal solution from the grid search is given below together with a plot:

# In[22]:


print(fit_report(result_brute))
plt.plot(x, data, 'b')
plt.plot(x, data + fcn2min(result_brute.params, x, data), 'r--')
plt.show()

# We can see that this fit is already very good, which is what we should expect since our ``brute`` force grid is sampled rather finely and encompasses the "correct" values.
#
# In a more realistic, complicated example the ``brute`` method will be used to get reasonable values for the parameters and perform another minimization (e.g., using ``leastsq``) using those as starting values. That is where the ``keep`` parameter comes into play: it determines the "number of best candidates from the brute force method that are stored in the ``candidates`` attribute". In the example above we store the best-ranking 25 solutions (the default value is ``50`` and storing all the grid points can be accomplished by choosing ``all``). The ``candidates`` attribute contains the parameters and ``chisqr`` from the brute force method as a namedtuple, ``(‘Candidate’, [‘params’, ‘score’])``, sorted on the (lowest) ``chisqr`` value. To access the values for a particular candidate one can use ``result.candidate[#].params`` or ``result.candidate[#].score``, where a lower # represents a better candidate. The ``show_candidates(#)`` uses the ``pretty_print()`` method to show a specific candidate-# or all candidates when no number is specified.
#
# The optimal fit is, as usual, stored in the ``MinimizerResult.params`` attribute and is, therefore, identical to ``result_brute.show_candidates(0)``.

# In[23]:


result_brute.show_candidates(0)


#  In this case, the next-best scoring candidate has already a ``chisqr`` that increased quite a bit:

# In[24]:


result_brute.show_candidates(1)


# and is, therefore, probably not so likely... However, as said above, in most cases you'll want to do another minimization using the solutions from the ``brute`` method as starting values. That can be easily accomplished as shown in the code below, where we now perform a ``leastsq`` minimization starting from the top-25 solutions and accept the solution if the ``chisqr`` is lower than the previously 'optimal' solution:

# In[25]:


best_result = copy.deepcopy(result_brute)

for candidate in result_brute.candidates:
    trial = fitter.minimize(method='leastsq', params=candidate.params)
    if trial.chisqr < best_result.chisqr:
        best_result = trial


# From the ``leastsq`` minimization we obtain the following parameters for the most optimal result:

# In[26]:


print(fit_report(best_result))


# As expected the parameters have not changed significantly as they were already very close to the "real" values, which can also be appreciated from the plots below.

# In[27]:


plt.plot(x, data, 'b')
plt.plot(x, data + fcn2min(result_brute.params, x, data), 'r--', label='brute')
plt.plot(x, data + fcn2min(best_result.params, x, data), 'g--', label='brute followed by leastsq')
plt.legend()
plt.show()

# Finally, the results from the ``brute`` force grid-search can be visualized using the rather lengthy Python function below (which might get incorporated  in lmfit at some point).

# In[28]:


def plot_results_brute(result, best_vals=True, varlabels=None,
                       output=None):
    """Visualize the result of the brute force grid search.

    The output file will display the chi-square value per parameter and contour
    plots for all combination of two parameters.

    Inspired by the `corner` package (https://github.com/dfm/corner.py).

    Parameters
    ----------
    result : :class:`~lmfit.minimizer.MinimizerResult`
        Contains the results from the :meth:`brute` method.

    best_vals : bool, optional
        Whether to show the best values from the grid search (default is True).

    varlabels : list, optional
        If None (default), use `result.var_names` as axis labels, otherwise
        use the names specified in `varlabels`.

    output : str, optional
        Name of the output PDF file (default is 'None')
    """
    from matplotlib.colors import LogNorm

    npars = len(result.var_names)
    fig, axes = plt.subplots(npars, npars)

    if not varlabels:
        varlabels = result.var_names
    if best_vals and isinstance(best_vals, bool):
        best_vals = result.params

    for i, par1 in enumerate(result.var_names):
        for j, par2 in enumerate(result.var_names):

            # parameter vs chi2 in case of only one parameter
            if npars == 1:
                axes.plot(result.brute_grid, result.brute_Jout, 'o', ms=3)
                axes.set_ylabel(r'$\chi^{2}$')
                axes.set_xlabel(varlabels[i])
                if best_vals:
                    axes.axvline(best_vals[par1].value, ls='dashed', color='r')

            # parameter vs chi2 profile on top
            elif i == j and j < npars-1:
                if i == 0:
                    axes[0, 0].axis('off')
                ax = axes[i, j+1]
                red_axis = tuple([a for a in range(npars) if a != i])
                ax.plot(np.unique(result.brute_grid[i]),
                        np.minimum.reduce(result.brute_Jout, axis=red_axis),
                        'o', ms=3)
                ax.set_ylabel(r'$\chi^{2}$')
                ax.yaxis.set_label_position("right")
                ax.yaxis.set_ticks_position('right')
                ax.set_xticks([])
                if best_vals:
                    ax.axvline(best_vals[par1].value, ls='dashed', color='r')

            # parameter vs chi2 profile on the left
            elif j == 0 and i > 0:
                ax = axes[i, j]
                red_axis = tuple([a for a in range(npars) if a != i])
                ax.plot(np.minimum.reduce(result.brute_Jout, axis=red_axis),
                        np.unique(result.brute_grid[i]), 'o', ms=3)
                ax.invert_xaxis()
                ax.set_ylabel(varlabels[i])
                if i != npars-1:
                    ax.set_xticks([])
                elif i == npars-1:
                    ax.set_xlabel(r'$\chi^{2}$')
                if best_vals:
                    ax.axhline(best_vals[par1].value, ls='dashed', color='r')

            # contour plots for all combinations of two parameters
            elif j > i:
                ax = axes[j, i+1]
                red_axis = tuple([a for a in range(npars) if a != i and a != j])
                X, Y = np.meshgrid(np.unique(result.brute_grid[i]),
                                   np.unique(result.brute_grid[j]))
                lvls1 = np.linspace(result.brute_Jout.min(),
                                    np.median(result.brute_Jout)/2.0, 7, dtype='int')
                lvls2 = np.linspace(np.median(result.brute_Jout)/2.0,
                                    np.median(result.brute_Jout), 3, dtype='int')
                lvls = np.unique(np.concatenate((lvls1, lvls2)))
                ax.contourf(X.T, Y.T, np.minimum.reduce(result.brute_Jout, axis=red_axis),
                            lvls, norm=LogNorm())
                ax.set_yticks([])
                if best_vals:
                    ax.axvline(best_vals[par1].value, ls='dashed', color='r')
                    ax.axhline(best_vals[par2].value, ls='dashed', color='r')
                    ax.plot(best_vals[par1].value, best_vals[par2].value, 'rs', ms=3)
                if j != npars-1:
                    ax.set_xticks([])
                elif j == npars-1:
                    ax.set_xlabel(varlabels[i])
                if j - i >= 2:
                    axes[i, j].axis('off')

    if output is not None:
        plt.savefig(output)


# and finally, to generated the figure:

# In[29]:


plot_results_brute(result_brute, best_vals=True, varlabels=None)
plt.show()
