1

I'm trying to plot a wave function over one dimension but it has real and imaginary parts, so I did a 3D plot animation of it. This is a screenshot:

Wave function

The main thing I would like to do is to spread it along the x-axis (which now is vertical) so it doesn't look squeezed. Also, it would be nice to set it up in a set of 3 RGB axes that intersect at the point (0,0,0). In the documentation I couldn't find any straight forward way to do this. I'm attaching the part of the code I'm using to animate it:

import matplotlib.pyplot as plt
import matplotlib.animation as animation
import mpl_toolkits.mplot3d.axes3d as p3

fig = plt.figure()
ax = fig.gca(projection='3d')
line, = ax.plot(REAL[0,:],IMAG[0,:],x,"r",linewidth=0.5)

def animacio(i):
    ax.collections.clear()
    line.set_data(REAL[i,:],IMAG[i,:])
    line.set_3d_properties(x, 'z')
    return line,

ani = animation.FuncAnimation(fig,animacio,interval=50, frames=Nt,repeat=True)
nom = 'Evolució_'
ani.save(str(nom)+'['+str(V0)+','+str(L)+','+str(l)+','+str(xi)+','+str(sigmax)+','+str(T)+']'+'.mp4', writer="ffmpeg", dpi=300)
plt.show()

print('Animation saved as: '+str(nom)+'['+str(V0)+','+str(L)+','+str(l)+','+str(xi)+','+str(sigmax)+','+str(T)+']'+'.mp4')
Trenton McKinney
  • 56,955
  • 33
  • 144
  • 158
Isidre Mas Magre
  • 149
  • 7
  • PS: Using f-strings (python 3.6), you could write `nom_complet = f'{nom}[{V0},{L},{l},{xi},{sigmax},{T}].mp4'` – JohanC Jan 05 '20 at 11:57
  • Could you share the code to generate a similar curve? It helps to provide an answer that really works for your case. – JohanC Jan 05 '20 at 11:58
  • Thank you for the tip on the name! Sure I can share it! But the code is a bit long, shall I edit the question to add It? – Isidre Mas Magre Jan 05 '20 at 12:28
  • Maybe [this post](https://stackoverflow.com/questions/30223161/matplotlib-mplot3d-how-to-increase-the-size-of-an-axis-stretch-in-a-3d-plo) also helps? [This article](https://geoffboeing.com/2015/04/animated-3d-plots-python/) provides some ideas about 3D animation. – JohanC Jan 05 '20 at 13:53
  • [This post](https://stackoverflow.com/questions/15617207/line-colour-of-3d-parametric-curve-in-pythons-matplotlib-pyplot/15617576#15617576) handles multi coloring – JohanC Jan 05 '20 at 13:55

1 Answers1

0

You can add colored lines to the plot, just by giving start and end points and assigning a color. The limits for the 'up'-axis can be set by ax.set_zlim. I created a demo curve that roughly resembles yours.

import numpy as np
from mpl_toolkits import mplot3d
import matplotlib.pyplot as plt

x = np.linspace(-10, 10, 1000)
y = np.sin(10*x)/(x*x+1)
z = np.cos(10*x)/(x*x+1)
ax = plt.axes(projection='3d')
ax.plot3D([0,0], [0,0], [-10,10], color='crimson')
ax.plot3D([0,0], [-1,1], [0,0], color='limegreen')
ax.plot3D([-1,1], [0,0], [0,0], color='dodgerblue')
line, = ax.plot3D(y, z, x, color='blueviolet')
ax.set_zlim(-1, 1)
plt.show()

At the left the plot without limiting, at the right with limits: resulting images

To get a more elongated view, you could use something like:

plt.gcf().set_size_inches(4, 12)
JohanC
  • 71,591
  • 8
  • 33
  • 66
  • Very nice on resembling the curve, I was hoping however to spread It while showing the whole range on the x axis, in a way that the volume of the plot would become more like a rectangular box, do you know if that's possible with matplotlib? – Isidre Mas Magre Jan 05 '20 at 13:00
  • Either changing the surrounding window manually, or programming it with `plt.gcf().set_size_inches(4, 12) ` – JohanC Jan 05 '20 at 13:33
  • This sort of works if I get the camera angle right, but what it really does is elongate the image as you said. By doing this proportionality on the other axes is lost. – Isidre Mas Magre Jan 05 '20 at 13:41
  • 1
    [This post](https://stackoverflow.com/questions/30223161/matplotlib-mplot3d-how-to-increase-the-size-of-an-axis-stretch-in-a-3d-plo) you commented on the main question in one of the answers solves smoothly the problem by adding a single line in the code: `ax.get_proj = lambda: np.dot(Axes3D.get_proj(ax), np.diag([0.2, 0.2, 1, 1]))` – Isidre Mas Magre Jan 05 '20 at 14:52
  • 1
    Thanks! Or `ax.get_proj = lambda: np.dot(Axes3D.get_proj(ax), np.diag([0.3, 0.3, 1.5, 1]))` to not reduce the size too much. – JohanC Jan 05 '20 at 15:04