In [19]:

    

import math
import numpy as np
from numpy.polynomial.hermite import Hermite
from matplotlib import pylab as plt

from mpl_toolkits.mplot3d.axes3d import Axes3D
%matplotlib notebook

def wf(n, x):
    shape = []
    if n == 0:
        shape = [1]
    else:
        for i in range(n):
            shape.append(0)
        shape.append(1)
    h_n = Hermite(shape)
    exp_factor = np.exp(-np.square(x)/2)
    #some_factor = (-1)**n/(np.sqrt(2**n *math.factorial(n)*np.sqrt(math.pi)))
    some_factor = 1/(np.sqrt(2**n * math.factorial(n)*np.sqrt(math.pi)))
    res = some_factor*exp_factor*h_n(x)
    return res.reshape(x.size,)


def grid_vect(a,b,N):
    x = np.linspace(a, b, N, endpoint=True)
    h = x[1] - x[0]
    return np.mgrid[a:b+h:h, a:b+h:h].reshape(2,-1).T


def wf_nd(n_list, x):
    Z = wf(n_list[0], x[:,0])
    i = 0
    for n in n_list:
        if (i>=1):
            Z *= wf(n, x[:,i])
        i += 1
    return Z



def show_wf(n, x):
    plt.title('Output:')
    plt.grid(True)
    plt.plot(x[0], wf(n,x), "r--")
    
def show_2d_wf(n_list, xy):
    fig = plt.figure(figsize=(14,6))
    ax = fig.add_subplot(1, 2, 1, projection='3d')
    xg,yg = np.meshgrid(xy[:,0], xy[:,1])
    ax.plot_wireframe(xy[:,0], xy[:,1], Z)


    

In [20]:

    

a = -5
b = 5
N = 50
xy = grid_vect(a,b,N)
n_list = (0,1)
Z = wf_nd(n_list,xy)
Z = Z.reshape(1,Z.size)


    

In [21]:

    

fig = plt.figure(figsize=(14,6))
ax = fig.add_subplot(1, 2, 1, projection='3d')
xg,yg = np.meshgrid(xy[:,0], xy[:,1])
ax.plot_wireframe(xy[:,0], xy[:,1], Z)


    

    















    












    
Out[21]:





<mpl_toolkits.mplot3d.art3d.Line3DCollection at 0x1992e0a5358>

In [18]:

    

Z.shape


    

    
Out[18]:





(1, 2601)

In [ ]:

    




    

In [17]:

    

xy[:,1].shape


    

    
Out[17]:





(2601,)

In [ ]:

    




    

In [ ]:

    




    

In [ ]:

    




    

In [ ]:

    




    

In [ ]:

    




    

In [ ]:

    




    

In [ ]:

    




    

In [ ]:

    




    

In [ ]: