Creating 1D array where elements are the sum of a 2D numpy array of stored functions

Question

I apologize if this has been asked before. I'm new to python and programming in general so please point me in the right direction if it has been asked. I'm using python 3.7.

I have a 2D numpy array where each element is a stored function. I want to add the functions in each column to get a 1D array where the elements of the 1D array are a single function. I'm not sure why the np.sum() function doesn't work to do this. I get a 1D array but the functions are only from the first column of the "npwavefxns" array.

i.e.

[[X₀₀,X₀₁,...,X_0n]

[X₁₀, X₁₁,...,X_1n]

...

[X_n0,X_n1,...[X_nn]]

-> [[X₀₀+X₁₀+...+X_n0, X₀₁+X₁₁+...+X_n1, X_0n+X_1n+...+X_nn]]

The np.sum() function seems to work for integers, so I'm not sure why it doesn't work when the elements are a function. A sample of my code is given below. If this code works correctly I suspect to get these 4 plots when "4" basis functions are used.

from scipy import mat, matrix, integrate
from scipy.integrate import quad
import numpy as np
from numpy import linalg as la
import os
from matplotlib import pyplot as plt

# Defining variables and functions
MP=float(9.10938356e-31)    #mass of electron in kg
WL=float(1e-10) #length of well in meters
CON=float(1.60218e-19)  #constant height in joules
Hb = float(1.054571726e-34) #reduced planck's constant in J s
NB=int(input("Number of basis functions ")) #define number of basis sets to be used

#####Potential energy of initial state#####
PE=[]
for j in range(1,NB+1):
    alist=[]
    for k in range(1,NB+1):
        F1=integrate.quad(lambda x:((2/WL)*np.sin((k*np.pi*x)/WL)*
        ((-CON)*np.sin(np.pi*x/WL))*np.sin((j*np.pi*x)/WL)),0,WL)[0]
        if F1 < -1e-25:
            F1=F1
        elif F1 > 1e-25:
            F1=F1
        else:
            F1=0
        alist.append(F1)
    PE.append(alist)
PEarray=np.asarray(PE)

#####Kinetic Energy of initial state#####
KE=[]
for j in range(1,NB+1):
    blist=[]
    for k in range(1,NB+1):
        F2=integrate.quad(lambda x:(((((Hb**2)*(j**2)*(np.pi**2))/(MP*(WL**3)))*
        ((np.sin(j*np.pi*x/WL))*(np.sin(k*np.pi*x/WL))))),0,WL)[0]
        if F2 < -1e-25:
            F2=F2
        elif F2 > 1e-25:
            F2=F2
        else:
            F2=0
        blist.append(F2)
    KE.append(blist)
KEarray=np.asarray(KE)

#####Adding PE and KE to give the full hamiltonian of the initial state#####
#####Then convert the list to a numpy array#####
sum=[0]*NB
for i in range(NB):
    sum[i]=[0]*NB
for y in range(len(PEarray)):
    for z in range(len(PEarray[0])):
        sum[y][z]=PEarray[y][z]+KEarray[y][z]
npsum=np.asarray(sum)

EVal, EVec=la.eigh(npsum)

wavefxns=[]
for j in range(1,NB+1):
    clist=[]
    for k in range(1,NB+1):
        F3 = (lambda x: ((EVec.item(k-1, j-1))*
        (np.sin((((k+1)*np.pi)/WL)*x))))
        clist.append(F3)
    wavefxns.append(clist)
npwavefxns=np.asarray(wavefxns)

EQS=[]
for j in range(0,NB):
    F4=np.sum(npwavefxns.item(j))
    EQS.append(F4)
npEQS=np.asarray(EQS)

for j in range(0,NB):
    wfxn1=(lambda x: ((npEQS.item(j))(x)))
    plt.xlabel("Box length")
    plt.ylabel("energy")
    x = np.linspace(0,WL,500)
    plt.plot(x, wfxn1(x), '--m')
    plt.show()

So here's the thing: numpy arrays are best used for storing numbers, and maybe strings. The arrays are flexible enough that you can store functions (and other arbitrary objects), but in general, you probably shouldn't. For example, you can't sum over an array of functions: normally you would get an error like `TypeError: unsupported operand type(s) for +: 'function' and 'function'`. The reason you aren't getting that error with your code is that the line `F4=np.sum(npwavefxns.item(j))` is only taking the sum of a single function, which will just return that function unchanged. — tel, Nov 13 '18 at 00:04
Ah I see. If I stored them in a nested list (is that the right term?) would I still end up with the unsupported operand error? — jc990, Nov 13 '18 at 03:09
You would, since at some point in their evaluation both `sum([f1, f2, f2])` and `np.array([f1, f2, f3]).sum()` will end up calling the basic addition operator. This is the equivalent of `f1 + f2`, which isn't supported for functions. I have to say, your idea of lazy evaluation of a sum of lambda functions is an interesting one. It's just not one that you're going to be able to easily implement in Python. Your best bet will be to evaluate the functions as you go and store only values. Alternatively, you can try to tweak what you have so that functions and input come together at the end. — tel, Nov 13 '18 at 03:22

tel · Accepted Answer · 2018-11-13T06:27:09.493

Okay, so as far as I can tell, you were running into two problems. One of them was the sum you were trying to take over an array of functions (and everything else related to npwavefxns).

The other problem was a real stinker, and turned out to lead back to the lambda that you were assigning to F3. The short version is that you were using the loop variables j and k in your lambda. lambdas "capture" variables so you can use them when you call the lambda later on. The catch is that the value of those variables can change, as did the value of j and k on every iteration of the loop. By the time you were actually calling those lambdas, they all ended up using the exact same value of j and k (which in this case was the final value they each had on the last loop).

I fixed the lambda issue using a technique called a closure (see this thread for more details). The short explanation is that it allows you to explicitly capture the current value of a variable for later use.

Anyhow, here's a complete working example of your code. Everything above the wavefxns=[] line has been left untouched:

from scipy import mat, matrix, integrate
from scipy.integrate import quad
import numpy as np
from numpy import linalg as la
import os
from matplotlib import pyplot as plt

# Defining variables and functions
MP=float(9.10938356e-31)    #mass of electron in kg
WL=float(1e-10) #length of well in meters
CON=float(1.60218e-19)  #constant height in joules
Hb = float(1.054571726e-34) #reduced planck's constant in J s
NB=int(input("Number of basis functions ")) #define number of basis sets to be used

#####Potential energy of initial state#####
PE=[]
for j in range(1,NB+1):
    alist=[]
    for k in range(1,NB+1):
        F1=integrate.quad(lambda x:((2/WL)*np.sin((k*np.pi*x)/WL)*
        ((-CON)*np.sin(np.pi*x/WL))*np.sin((j*np.pi*x)/WL)),0,WL)[0]
        if F1 < -1e-25:
            F1=F1
        elif F1 > 1e-25:
            F1=F1
        else:
            F1=0
        alist.append(F1)
    PE.append(alist)
PEarray=np.asarray(PE)

#####Kinetic Energy of initial state#####
KE=[]
for j in range(1,NB+1):
    blist=[]
    for k in range(1,NB+1):
        F2=integrate.quad(lambda x:(((((Hb**2)*(j**2)*(np.pi**2))/(MP*(WL**3)))*
        ((np.sin(j*np.pi*x/WL))*(np.sin(k*np.pi*x/WL))))),0,WL)[0]
        if F2 < -1e-25:
            F2=F2
        elif F2 > 1e-25:
            F2=F2
        else:
            F2=0
        blist.append(F2)
    KE.append(blist)
KEarray=np.asarray(KE)

#####Adding PE and KE to give the full hamiltonian of the initial state#####
#####Then convert the list to a numpy array#####
sum=[0]*NB
for i in range(NB):
    sum[i]=[0]*NB
for y in range(len(PEarray)):
    for z in range(len(PEarray[0])):
        sum[y][z]=PEarray[y][z]+KEarray[y][z]
npsum=np.asarray(sum)

EVal, EVec=la.eigh(npsum)

wavefxns=[]
for j in range(0,NB):
    clist=[]
    for k in range(0,NB):
        F3 = (lambda a,b: (lambda x: ((EVec.item(b-1, a-1)) * (np.sin((((b+1)*np.pi)/WL)*x)))))(j,k)
        clist.append(F3)
    wavefxns.append(clist)

gridspec_kw = {'wspace': 0, 'hspace': 0}
fig,axarr = plt.subplots(NB, sharex=True, squeeze=False, gridspec_kw=gridspec_kw, figsize=(3,7.5))
fig.subplots_adjust(left=0, bottom=0, right=1, top=1)

for j,ax in zip(range(0,NB), axarr.ravel()):
    wfxn = lambda x: np.sum([wavefxns[i][j](x) for i in range(len(wavefxns))], axis=0)
    if j==(NB - 1):
        ax.set_xlabel("Box length")
    ax.set_ylabel("energy")
    x = np.linspace(0,WL,500)
    ax.plot(x, wfxn(x), '--m')

Which (when you run it and enter 4) produces the output:

I tweaked the plotting routines so that all of the wavefunctions would plot onto one figure (and so I'd only have one figure to copy/paste into this answer).

Thanks! I guess a simpler method for this could be to plot the principle diagonal of the array rather than adding all of the functions together, which should give the same result as long as there isn't a large perturbation from the off diagonal elements. The end goal of this program will be to calculate the integral of the initial wavefunction multiplied by the final wavefunction, which will give a single value for each row. I will probably try to build that into a single loop as you suggested in the previous comment. I just wanted to be able to plot and check the functions on the way. — jc990, Nov 13 '18 at 16:51

Creating 1D array where elements are the sum of a 2D numpy array of stored functions

1 Answers1