schrodinger-simulation

schrodinger's playground
git clone git://popovic.xyz/schrodinger-simulation.git
Log | Files | Refs | README | LICENSE

plotter.py (1033B)


      1 #!/usr/bin/python3.9
      2 
      3 import numpy as np
      4 import matplotlib.pyplot as plt
      5 import os
      6 
      7 
      8 def plotter(U, V, dt, xx, yy, name:str):
      9 
     10     n = len(xx)
     11     path = './.plotcache'
     12     if os.path.exists(path) != True:
     13         os.mkdir(path)
     14 
     15     print('Making gif...')
     16     for i in range(0, len(U)):
     17         levels = np.linspace(0, max(U[i].reshape(n*n)), 200)
     18         fig, ax = plt.subplots(figsize=[7,5])
     19 
     20         c = ax.contourf(xx, yy, U[i], levels=levels, zorder=1,\
     21                         cmap=plt.cm.inferno)
     22 
     23         ax.contour(xx, yy, V.reshape(n, n), extend='both',\
     24                    cmap=plt.cm.binary)
     25 
     26         ax.set_title(f'Time {round(i*dt, 2)}')
     27         fig.colorbar(c)
     28 
     29         plt.savefig(f'{path}/img-{i}.png')
     30         plt.close()
     31 
     32     plt.plot(U[-1][15])
     33     plt.savefig('./inteferenz.png')
     34     plt.close()
     35 
     36     os.system(f'ffmpeg -start_number 0 -i {path}/img-%d.png {name}.gif')
     37     os.system(f'ffmpeg -start_number 0 -i {path}/img-%d.png {name}.mp4')
     38     os.system('rm -rf ' + path)
     39     os.system('rm -rf __pycache__')