Bottom Page

Thread Rating:
  • 1 Vote(s) - 1 Average
  • 1
  • 2
  • 3
  • 4
  • 5
 Change the colour of a plot
#1
Hi, 

some days ago, I received the task to plot something (I can give you more details if needed) and it worked out quite well, but I have some trouble with the colours. I just can't get them to look like I want them... The code to generate these plots is the following (I think you can ignore the first 150 lines, only functions and stuff that is needed for the plots which are generated in line 156-167):


from numpy import *
from numpy.linalg import eig, norm
from scipy.special import hermite, gamma
from matplotlib.pyplot import *

def arnoldi(A, v0, k):
    r"""Arnoldi algorithm to compute the Krylov approximation :math:`H` of a matrix :math:`A`.

    :param A: The matrix :math:`A` of shape :math:`N \times N` to approximate.
    :param v0: The initial vector :math:`v_0` of length :math:`N`.
    :param k: The number :math:`k` of Krylov steps performed.
    :return: A tuple :math:`(V, H)` where :math:`V` is the large matrix of shape
             :math:`N \times (k+1)` containing the orthogonal vectors and :math:`H` is the
             small matrix of shape :math:`k \times k` containing the Krylov approximation
             of :math:`A`.
    """
    r, c = A.shape
    V = zeros((r, k+1), dtype=complexfloating)
    H = zeros((k+1, k), dtype=complexfloating)

    V[:,0] = v0.reshape(-1) / norm(v0)

    for i in range(1, k+1):
        vi = dot(A, V[:,i-1])
        for j in range(i):
            H[j,i-1] = dot(conjugate(V[:,j]), vi)
            vi -= H[j,i-1] * V[:,j]
        H[i,i-1] = norm(vi)
        V[:,i] = vi / H[i,i-1]

    return V, H[:-1,:]

def Part3():
    # Teil 3: Unteraufgaben g), h), i)

    # Harmonischer Osillator
    v = lambda x, y: 0.5*(x**2 + y**2)

    # Henon-Heiles
    a = 2.0
    b = 0.4
    v = lambda x, y: 0.5*a*(x**2 + y**2) + b*(x**2*y - y**3/3.0)

    N = 32
    x, h = linspace(-3, 3, N, retstep=True)
    X,Y = meshgrid(x,x)
    x = X.reshape(-1)
    y = Y.reshape(-1)

    # Just K for later use
    K = 6


    # Unteraufgabe g)
    
    # Laplace Part
    D = zeros((N*N, N*N), dtype=floating)
    for r in range(0, N):
        for c in range(0, N):
            # X direction stencil
            if c == 0:
                # Forward
                D[N*r+c, N*r+c] += 1.0
                D[N*r+c, N*r+(c+1)] += -2.0
                D[N*r+c, N*r+(c+2)] += 1.0
            elif c == N-1:
                # Backward
                D[N*r+c, N*r+(c-2)] += 1.0
                D[N*r+c, N*r+(c-1)] += -2.0
                D[N*r+c, N*r+c] += 1.0
            else:
                # Central
                D[N*r+c, N*r+(c-1)] += 1.0
                D[N*r+c, N*r+c] += -2.0
                D[N*r+c, N*r+(c+1)] += 1.0

            # Y direction stencil
            if r == 0:
                # Forward
                D[N*r+c, N*r+c] += 1.0
                D[N*r+c, N*(r+1)+c] += -2.0
                D[N*r+c, N*(r+2)+c] += 1.0
            elif r == N-1:
                # Backward
                D[N*r+c, N*(r-2)+c] += 1.0
                D[N*r+c, N*(r-1)+c] += -2.0
                D[N*r+c, N*r+c] += 1.0
            else:
                # Central
                D[N*r+c, N*(r-1)+c] += 1.0
                D[N*r+c, N*r+c] += -2.0
                D[N*r+c, N*(r+1)+c] += 1.0

    D /= h**2
    # Potential Part
    V = diag(v(x,y))
    # Hamiltonian
    H = -0.5*D + V

    figure()
    matshow(H[:50,:50])
    colorbar()
    savefig("discrete_hamiltonian.pdf")
    
    
    # Unteraufgabe h)

    ew, ev = eig(H)
    # Sort the eigenvalues
    I = argsort(ew)
    ew = ew[I]
    Psi = ev[:,I]


    fig = figure(figsize=(12,8))
    for k in range(K):
        # Plot first K eigenstates
        psik = Psi[:,k].reshape((N,N))

        fig.add_subplot(2, 3, k+1)
        ax = fig.gca()
        ax.set_aspect('equal')

        ax.contour(X, Y, v(X,Y), colors="gray", levels=linspace(0, 20, 15))
        ax.contourf(X, Y, abs(psik), levels=linspace(0, 0.15, 40))

        #ax.grid(True)
    fig.savefig("henon_eigenfunctions.pdf")


if __name__ == "__main__":
    Part3()

I would like to give you some pictures but I cant (first post), so I will try and describe how it should look like. If you run the programm you will get one figure with six subplots, each has a "blue" background and in the middle a green part. The background should be black an the middle in some kind of neo-red (like glowing).. I just cant figure out how to change that.

Hope someone can help with this.
Thank you for the helo and sorry for the missing pictures, 
Sito
Quote
#2
Just use some colormap as a parameter for contourf.
    ax.contourf(X, Y, abs(psik), cmap="gist_heat", levels=linspace(0, 0.15, 40))
Quote

Top Page

Possibly Related Threads...
Thread Author Replies Views Last Post
  How to create a new plot in a figure right below the first one after click event? codexx 0 28 Yesterday, 04:45 PM
Last Post: codexx
  How do I make my plot show up? ThomasM4 0 267 Dec-09-2018, 05:28 AM
Last Post: ThomasM4
  Tips on surface plot HW question Cwcox 1 331 Nov-14-2018, 11:45 PM
Last Post: Larz60+
  How to customize x axis in matplotlib.pyplot for a scatter plot? wlsa 9 671 Nov-10-2018, 01:32 AM
Last Post: wlsa
  Adding text to plot Pythcoronas 1 404 Sep-13-2018, 12:27 PM
Last Post: baby_quant
  plot the mean in a bar diagram sussii 3 537 May-23-2018, 11:27 PM
Last Post: wavic
  How to plot date series in matplotlib? StrybolData 2 2,301 Jan-25-2018, 07:13 PM
Last Post: StrybolData
  Removing data in a plot ulrich48155 3 1,405 Jun-19-2017, 06:31 PM
Last Post: zivoni
  How to plot a Graph Kamo 2 1,476 May-12-2017, 12:22 PM
Last Post: zivoni
  matplotlib : Raster Plot adithyakrish 0 1,193 Apr-24-2017, 10:46 AM
Last Post: adithyakrish

Forum Jump:


Users browsing this thread: 1 Guest(s)