title: How to plot the heat map for a given function in Python
body: <p>I have a function called <code>func(mu, gamma)</code> . For each combination of <code>mu</code> and <code>gamma</code>, the function will return a value, let's call it <code>return_value</code>. </p>
  
  <p>Now I have set range for <code>mu</code> and <code>gamma</code>:</p>
  
  <pre><code>mu = np.linspace(0,1,100)
  gamma = np.linspace(0,1,100)
  </code></pre>
  
  <p>Now we have 1e4 combinations and each combinations corresponds to a <code>return_value</code>. I want to plot a heatmap for <code>return_value</code>.</p>
  
  <p>I have tried to use <code>pcolor</code> in Python. However, from the example in the documentation:</p>
  
  <pre><code>import matplotlib.pyplot as plt
  import numpy as np
  
  # make these smaller to increase the resolution
  dx, dy = 0.15, 0.05
  
  # generate 2 2d grids for the x &amp; y bounds
  y, x = np.mgrid[slice(-3, 3 + dy, dy),
                  slice(-3, 3 + dx, dx)]
  z = (1 - x / 2. + x ** 5 + y ** 3) * np.exp(-x ** 2 - y ** 2)
  # x and y are bounds, so z should be the value *inside* those bounds.
  # Therefore, remove the last value from the z array.
  z = z[:-1, :-1]
  z_min, z_max = -np.abs(z).max(), np.abs(z).max()
  </code></pre>
  
  <p>because the defined function <code>fun</code> in my script can not take array as input, it does not work and I get this message if I follow the example:</p>
  
  <pre><code>ValueError: The truth value of an array with more than one element is ambiguous. Use a.any() or a.all()
  </code></pre>
  
  <p>Below is my code for <code>func</code>:</p>
  
  <pre><code>def fun(mu, gamma2):
      Isolation_Ratio = []
      kappa1 = gamma2
      kappa2 = gamma2                         
      gamma1 = gamma2
      g0 = gamma2 + kappa2 + gamma1 + kappa1
      gammag = kappa1/2. + gamma1/2.
      gamma = gamma2/2. + kappa2/2.
      for ii in range(len(rangedeltaw)):
  
          deltaw = rangedeltaw[ii]
  
          Forward_delta = forward_delta(mu, deltaw)
          Backward_delta = backward_delta(mu, deltaw)
  
          forward_root1, forward_root2, forward_root3 = forward_root(mu, deltaw)
          test_D, backward_root1, backward_root2, backward_root3 = backward_root(mu, deltaw)
  
          Root1.append(backward_root1)
          Root2.append(backward_root2)
          Root3.append(backward_root3)
  
          root1.append(forward_root1)
          root2.append(forward_root2)
          root3.append(forward_root3)
          if Forward_delta &gt;= 0 and Backward_delta &gt;= 0:
              a2sq = [max([forward_root1.real, forward_root2.real, forward_root3.real])]
              b1sq = [max([backward_root1.real, backward_root2.real, backward_root3.real])]
              A2sq.append(max([forward_root1.real, forward_root2.real, forward_root3.real]))
              B1sq.append(max([backward_root1.real, backward_root2.real, backward_root3.real]))
              for ii in range(len(a2sq)):
                  for jj in range(len(b1sq)):
  
                      Isolation_Ratio.append(kappa2*a2sq[ii]/(kappa1*b1sq[jj]))
          elif Forward_delta &gt;= 0 and Backward_delta &lt; 0:
              a2sq = [max([forward_root1.real, forward_root2.real, forward_root3.real])]
              b1sq = [backward_root1.real]
              A2sq.append(max([forward_root1.real, forward_root2.real, forward_root3.real]))
              B1sq.append(backward_root1.real)
              for ii in range(len(a2sq)):
                  for jj in range(len(b1sq)):
                      Isolation_Ratio.append(kappa2*a2sq[ii]/(kappa1*b1sq[jj]))
  
          elif Forward_delta &lt; 0 and Backward_delta &gt;= 0:
              a2sq = [forward_root1.real]
              b1sq = [max([backward_root1.real, backward_root2.real, backward_root3.real])]
              A2sq.append(forward_root1.real)
              B1sq.append(max([backward_root1.real, backward_root2.real, backward_root3.real]))
              for ii in range(len(a2sq)):
                  for jj in range(len(b1sq)):
                      Isolation_Ratio.append(kappa2*a2sq[ii]/(kappa1*b1sq[jj]))
  
          else:   
              A2sq.append(forward_root1.real)
              B1sq.append(backward_root1.real)
              Isolation_Ratio.append(kappa2*forward_root1.real/(kappa1*backward_root1.real))   
  
  
  
      x = Isolation_RangeDeltaw
      y = Isolation_Ratio
  
  
      return max(y)
  </code></pre>