Plotting functions where the x-variable is multiplied by an array

Python Help
1

I’m interested to plot a function in which the (scalar) independent variable is multiplied by an array. The function returns a scalar, so I should still be able to plot it in a 1D graph. My function is a little complicated, but for the sake of example consider the following:

import numpy as np
import matplotlib.pyplot as plt 

def f(x):
     return sum(x*np.array([1,2]))

x = np.linspace(-10, 10, 100)
plt.plot(x, f(x), color='red')
plt.show()

This code doesn’t work, because in the operation x*np.array([1,2]) the variable x is taken as the whole 100-dimensional array (which obviously does not have the same shape as [1,2]).

How do I tell it to take x as a scalar and plot it correctly?

2

So x was [1, 2, 3], then f(x) would be sum((1*1, 1*2)), sum((2*1, 2*2)), sum((3*1, 3*2)), i.e (3, 6, 9)?

If that is the case (ignore me if it isn’t :sweat_smile:) you can fix your code by adding an extra dimension to x. x[..., None] will give [[1],[2],[3]], which you can then multiply by [1, 2].

x = np.array([1, 2, 3])
a = np.array([1, 2])

y = x[..., None] * a
print(y)
# [[1, 2],
#  [2, 4],
#  [3, 6]]
print(y.sum(1))
# [3, 6, 9]

(btw, x[..., None] here is equivalent to x.reshape(len(x), 1), or np.expand_dims(x, axis=1))

You could also do it the other way round i.e. sum(x*np.array([[1],[2]]))

3

Hey wookie. This isn’t exactly what I’m looking for. I’d just like to plot the function with the various values of x given as input to f one at a time. For example, say x=[1,2,3], then the program should plot the points f(1)=1*(1+2)=3, f(2)=2*(1+2)=6, f(3)=3*(1+2)=9. Obviously I’m looking for a more complete graph, so I’ll be calling the function np.linspace to fill x.

4

What about writing your function as

def f(x):
     return x*(np.array([1,2]).sum())

?

5

Hey gkb. This function is merely an example; the code I have is a bit more complicated, but ultimately the problem is the same. I can’t really modify the original function that way, I’m afraid.

6

Just for completeness’ sake: the full code I’m dealing with is the following.

import numpy as np
import math as math
import matplotlib.pyplot as plt

def sigma_x(j,n):
    if j>n:
        raise Exception('j should not exceed n')
    elif j==n:
        return sigma_x(0,n)
    else:
        m=1
        for k in range(j):
            m=np.kron(m,np.eye(2))
        m=np.kron(m,np.array([[0,1],[1,0]]))
        for k in range (j+1,n):
                m=np.kron(m,np.eye(2))
        return m

def sigma_z(j,n):
    if j>n:
        raise Exception('j should not exceed n')
    elif j==n:
        return sigma_z(0,n)
    else:
        m=1
        for k in range(j):
            m=np.kron(m,np.eye(2))
        m=np.kron(m,np.array([[1,0],[0,-1]]))
        for k in range (j+1,n):
                m=np.kron(m,np.eye(2))
        return m

def H(n,J,g,h,bound): #MAIN FUNCTION
    term1=0.
    term2=0.
    term3=0.
    
    if bound=='PBC':
        for k in range(n):
            term1+=np.dot(sigma_z(k,n),sigma_z(k+1,n))
        for k in range(n):
            term2+=sigma_x(k,n)
        for k in range(n):
            term3+=sigma_z(k,n)

        return -J*term1-g*term2-h*term3 #PROBLEM HERE
    
    elif bound=='OBC':
        for k in range(n-1):
            term1+=np.dot(sigma_z(k,n),sigma_z(k+1,n))
        for k in range(n):
            term2+=sigma_x(k,n)
        for k in range(n):
            term3+=sigma_z(k,n)
    
        return -J*term1-g*term2-h*term3 #SAME AS ABOVE
    else:
        raise Exception('Choose PBC or OBC')

def energy(i,n,J,g,h,bound):
    array=np.linalg.eigh(H(n,J,g,h,bound))[0]
    return array[i]

def delta_E(n,J,g,h,bound):
    return energy(1,n,J,g,h,bound)-energy(0,n,J,g,h,bound)

def plot_gap(n,J,bound):
    g=np.linspace(0,1,50)
    E=delta_E(n,J,g,0,bound)
    
    fig = plt.figure()
    ax = fig.add_subplot(1, 1, 1)
    ax.spines['left'].set_position('center')
    ax.spines['bottom'].set_position('zero')
    ax.spines['right'].set_color('none')
    ax.spines['top'].set_color('none')
    ax.xaxis.set_ticks_position('bottom')
    ax.yaxis.set_ticks_position('left')
    
    plt.plot(g,E,'b')
    plt.show()

The problem here is the variable g in the function H, which multiplies term2 (a 2D array). When I’m trying to plot the function delta_E with g=np.linspace(...), I get the issue above.

7

Do I understand correctly that basically you want to vectorize delta_E? I.e. you want to call delta_E with multiple values for g?

In that case, you could do the following:

delta_E_vec = np.vectorize(delta_E, excluded=["n", "J", "h", "bound"])

And then call this new function in your function plot_gap:

def plot_gap(n, J, bound):
    g = np.linspace(0, 1, 10)
    E = delta_E_vec(n, J, g, 0, bound)

    fig = plt.figure()
    ax = fig.add_subplot(1, 1, 1)
    ax.spines["left"].set_position("center")
    ax.spines["bottom"].set_position("zero")
    ax.spines["right"].set_color("none")
    ax.spines["top"].set_color("none")
    ax.xaxis.set_ticks_position("bottom")
    ax.yaxis.set_ticks_position("left")

    plt.plot(g, E, "b")
    plt.show()

Unfortunately, this will use a Python loop to iterate over all values of g. For thousands of values this might be too slow for you.

8

Hey @gkb, sorry for the late reply, forgot to check in on the forum. This does work, thank you very much! Although as you say this can become quickly very computationally expensive. I’m afraid this is hardwired in the nature of the problem, though.

1 Like