#!/usr/bin/env python

# Grid Chi-square
# Module to perform Chi-square analysis for a grid of values.
from __future__ import division, print_function
import os
from joblib import Memory
from joblib import Parallel, delayed
import numpy as np
import matplotlib.pyplot as plt

path = "/home/jneal/Phd/Codes/equanimous-octo-tribble/"  # save path
cachedir = os.path.join(path, "cache")  # save path
memory = Memory(cachedir=cachedir, verbose=0)


# @memory.cache
def chi_squared(observed, expected, error=None):
    """Calculate chi squared.
    Same result as as scipy.stats.chisquare
    """
    if np.any(error):
        chisqr = np.sum((observed - expected) ** 2 / (error ** 2))
    else:
        # chisqr = np.sum((observed-expected)**2)
        chisqr = np.sum((observed - expected)**2 / expected)
        # When divided by exted the result is identical to scipy
    return chisqr


def parallel_chisqr_1D(t, obs, model, iter_1, n_jobs=4):
    grid = Parallel(n_jobs=n_jobs)(delayed(chi_squared)(obs, model(t, a)) for a in iter_1)
    return grid


# @memory.cache
def parallel_chisqr_2D(t, obs, model, iter_1, iter_2, n_jobs=4):
    grid = Parallel(n_jobs=n_jobs)(delayed(chi_squared)(obs, model(t, a, b))
                                   for a in iter_1 for b in iter_2)
    return grid


# @memory.cache
def parallel_chisqr_3D(t, obs, model, iter_1, iter_2, iter_3, n_jobs=4):
    grid = Parallel(n_jobs=n_jobs)(delayed(chi_squared)(obs, model(t, a, b, c))
                                   for a in iter_1 for b in iter_2 for c in iter_3)
    return grid


def parallel_chisqr_4D(t, obs, model, iter_1, iter_2, iter_3, iter_4, n_jobs=4):
    grid = Parallel(n_jobs=n_jobs)(delayed(chi_squared)(obs, model(t, a, b, c, d))
                                   for a in iter_1 for b in iter_2 for c in iter_3 for d in iter_4)
    return grid


# @memory.cache
def model_func(t, a, b, c):
    return a * np.exp(-(t-b)**2 / (2*c**2))


if __name__ == "__main__":
    a = 5
    b = 5
    c = 2
    t = np.linspace(0, 5 * np.pi, 1000)
    iterable_a = np.linspace(0.1, 20, 50)
    iterable_b = np.linspace(0.1, 20, 30)
    iterable_c = np.arange(1, 4)
    simulation = model_func(t, a, b, c)

    chisqr_grid = parallel_chisqr_3D(t, simulation, model_func, iterable_a, iterable_b, iterable_c)

    print(chisqr_grid)
    X, Y, Z = np.meshgrid(iterable_a, iterable_b, iterable_c, indexing="ij")

    min_loc = np.argmin(chisqr_grid)
    print("min location", min_loc)

    a_sol = X.ravel()[min_loc]
    b_sol = Y.ravel()[min_loc]
    c_sol = Z.ravel()[min_loc]

    plt.plot(t, simulation, "x", label="sim")
    plt.plot(t, model_func(t, a_sol, b_sol, c_sol), label="solution")
    plt.legend()
    plt.show()
    print("Found Solution", a_sol, b_sol, c_sol)
    # main()

    plt.contourf(X[:, :, 1], Y[:, :, 1],
                 np.array(np.log10(chisqr_grid)).reshape(X.shape)[:, :, 1], 20)
    plt.show()
    # print("Time to run = {} seconds".format(time.time()-start))

@@ -903,17 +903,22 @@
 isqr_1D(
-t
+iter_1
 , obs, m
@@ -923,22 +923,28 @@
  model, 
-iter_1
+model_params
 , n_jobs
@@ -1015,20 +1015,32 @@
 , model(
-t, a
+a, *model_params
 )) for a
@@ -1112,37 +1112,48 @@
 _2D(
-t, obs, model, iter_1, iter_2
+iter_1, iter_2, obs, model, model_params
 , n_
@@ -1228,23 +1228,35 @@
 , model(
-t, a, b
+a, b, *model_params
 ))%0A     
@@ -1375,31 +1375,16 @@
 isqr_3D(
-t, obs, model, 
 iter_1, 
@@ -1398,16 +1398,45 @@
  iter_3,
+ t, obs, model, model_params,
  n_jobs=
@@ -1506,26 +1506,38 @@
 , model(
-t, 
 a, b, c
+, *model_params
 ))%0A     
@@ -1656,31 +1656,16 @@
 isqr_4D(
-t, obs, model, 
 iter_1, 
@@ -1687,16 +1687,42 @@
  iter_4,
+ obs, model, model_params,
  n_jobs=
@@ -1792,19 +1792,16 @@
 , model(
-t, 
 a, b, c,
@@ -1802,16 +1802,31 @@
  b, c, d
+, *model_params
 ))%0A     
@@ -1969,26 +1969,26 @@
 el_func(
-t, 
 a, b, c
+, t
 ):%0A    r
@@ -2011,23 +2011,31 @@
 (-(t
--
+ - 
 b)
+ 
 **
+ 
 2 / (2
-*c
+ * c 
 **
+ 
 2))%0A
@@ -2282,18 +2282,18 @@
 unc(
-t, 
 a, b, c
+, t
 )%0A%0A 
@@ -2332,35 +2332,8 @@
 _3D(
-t, simulation, model_func, 
 iter
@@ -2362,16 +2362,43 @@
 erable_c
+, simulation, model_func, t
 )%0A%0A    p

Reorder parameters to put the iterables first and the model parameters last