q_read_split.py

from clawpack.pyclaw import Solution
import matplotlib as mpl
mpl.use('Agg')
import matplotlib.pyplot as plt
import glob
import os
visdir = './_ntf'

def plane_cut(path='./output',cutplane='z',qn=1,outdir='results',saveq=False,saveplot=False,splitall=False,planecutplot=False):
    if not os.path.exists(os.path.join(path,outdir)):
        os.mkdir(os.path.join(path,outdir))
    for i in range(0,len(glob.glob(os.path.join(path,'*.pkl*')))+1):
        sol = Solution()
        sol.read(i,path=path,file_format='petsc',read_aux=False)
        q = sol.state.get_q_global()
        print q.min(), q.max(),q.shape[0]
        print q.shape
        if splitall:
            print q.shape[0]
            for n in range(0,q.shape[0]):
                q[n,:,:,:].tofile(os.path.join(path,'q'+str(n)+'.'+str(i).zfill(3)))
        if saveq:
            q[qn,:,:,:].tofile(os.path.join(path,'q'+str(qn)+'.'+str(i).zfill(3)))

        if planecutplot:
            if cutplane=='z':
                plt.pcolor(q[qn,:,:,int(q.shape[3]/2.0)])
            elif cutplane=='y':
                plt.pcolor(q[qn,:,int(q.shape[2]/2.0),:])
            elif cutplane=='x':
                plt.pcolor(q[qn,int(q.shape[1]/2.0),:,:])
            plt.colorbar()
            plt.draw()
            plt.savefig(os.path.join(path,outdir,cutplane+'-'+str(i).zfill(3)+'.png'),dpi=240)
            plt.close()
        del sol
        del q

if __name__=="__main__":
    import sys
    from clawpack.pyclaw.util import run_app_from_main
    run_app_from_main(plane_cut)