1#!/usr/bin/env python3 2# 3# Reads in the output trajectory data -ts_save_trajectory -ts_trajectory_type visualization 4# Also provides a way to plot the trajectory 5# 6# Make sure $PETSC_DIR/bin is in your PYTHONPATH 7# 8import os 9import PetscBinaryIO 10import numpy as np 11import matplotlib.pyplot as pyplot 12 13def ReadTrajectory(directory): 14 io = PetscBinaryIO.PetscBinaryIO() 15 16 v = [] 17 t = [] 18 19 cnt = 0 20 while 1: 21 try: 22 fh = open(os.path.join(directory,'SA-%06d.bin'%cnt)); 23 cnt = cnt + 1 24 objecttype = io.readObjectType(fh) 25 v.append(io.readVec(fh)) 26 t.append(np.fromfile(fh, dtype=io._scalartype, count=1)[0]) 27 except: 28 break 29 30 names = [] 31 try: 32 fh = open(os.path.join(directory,'variablenames')) 33 nstrings = np.fromfile(fh, dtype=io._inttype, count=1)[0] 34 sizes = np.fromfile(fh, dtype=io._inttype, count=nstrings) 35 for i in range(0,nstrings): 36 s = np.fromfile(fh, dtype=np.byte, count=sizes[i]) 37 names.append("".join(map(chr, s))[0:-1]) 38 except: 39 pass 40 41 return (t,v,names) 42 43def PlotTrajectories(t,v,names,subnames): 44 print(names) 45 sub = [] 46 for s in subnames: 47 sub.append(names.index(s)) 48 w = [] 49 for i in v: 50 w.append(i[sub]) 51 52 pyplot.plot(t,w) 53 pyplot.legend(subnames) 54 pyplot.show() 55 56