# sensitivity.py
#
import numpy as np
import matplotlib.pyplot as pl
import malaria_control_start_function as mcs
from matplotlib import cm

# Define the ranges for the sensitivity
ivo = 0.01
bro = 0.24
ivs = np.arange(0.001, 0.5, (0.5-0.001)/20)
brs = np.arange(0.01, 0.99, (0.99-0.01)/20)
#assert(len(ivs) == len(brs))

# Define the array
S = np.zeros([len(ivs),len(brs)])

# Loop through each of the values to a sensitivity array
for i in range(len(brs)):
	for j in range(len(ivs)):
		P = mcs.malaria_control_orig(ivs[j], brs[i], 200, False)
		S[j][i] = P[0]

# Plot

fig = pl.figure()
ax = fig.gca(projection='3d')
X,Y = np.meshgrid(ivs, brs)
surf = ax.plot_surface(X,Y,S.T, rstride=1, cstride=1, cmap=cm.coolwarm,
                       linewidth=0, antialiased=False)
ax.set_xlabel('Immunity Rate')
ax.set_ylabel('Bite Rate')
ax.set_zlabel('Population')
ax.set_title('Sensitivity Plot')
pl.show()