#!/usr/bin/python3
# heatmap.py
# 
# Une méthode pour facilement tracer une fonction à deux variables

import numpy as np
import matplotlib.pyplot as plt


def heatmap(h, x_min=-10, x_max=10, x_res=500, y_min=-10, y_max=10, y_res=500):
# arguments: h = fonction à tracer
#            x_min, x_max, y_min, y_max définit la région
#            x_res, y_res = nombre de points à tracer en direction x et y

    if x_min >= x_max or y_min >= y_max or x_res <= 0 or y_res <= 0:
        raise ValueError
        
    dx = (x_max - x_min) / x_res   
    dy = (y_max - y_min) / y_res
    h_vec = np.vectorize(lambda j, i: h(x_min + i*dx, y_min + j*dy))
    img = np.fromfunction(h_vec, (x_res, y_res))
    
    plt.imshow(img)
    plt.show()