エネルギーの時間変化をプロット

import glob
import matplotlib.pyplot as plt
import numpy as np
from scipy import interpolate
 
def unesco_density(t,s):
	if s < 0.0: s = 0.0 
	return 999.842594 \
	+6.793952E-02*t\
	-9.095290E-03*t*t\
	+1.001685E-04*t*t*t\
	-1.120083E-06*t*t*t*t\
	+6.536332E-09*t*t*t*t*t\
	+(0.824493-4.0899E-03*t\
	+7.6438E-05*t*t\
	-8.2467E-07*t*t*t\
	+5.3875E-09*t*t*t*t)*s\
	+(-5.72466E-03\
	+1.0277E-04*t\
	-1.6546E-06*t*t)*s**1.5\
	+4.8314E-04*s*s
 
def get_data(file):
  print ("processing... ", file)
  columns = []
  for line in open(file).readlines():
	items = line.rstrip().split()
	if items[0] == "time":
	  time = float(items[1])
	elif items[0] == "dz":
	  dzs = [float(x) for x in items[1:]]
	elif items[0] == "limit_depth":
	  pass
	elif items[0] == "items":
	  pass
	elif items[0] == "column":
	  columns.append({"number"  :int(items[1]),    "k_offset" :int(items[2]),
					  "center_x":float(items[3]),  "center_y" :float(items[4]),
					  "delta_x" : float(items[5]), "delta_y"  :float(items[6]),
					  "bottom_k":int(items[7]),    "surface_k":int(items[8]),
					  "bottom_z":float(items[9]),  "surface_z":float(items[10]),
					  "variables":{}})
	else:
	  columns[int(items[0])]["variables"][items[1]]=[float(x) for x in items[2:]]
  return time, dzs, columns 

def energy_estimation(data):
  total_base_area  = 0.0
  kinetic_energy   = 0.0
  potential_energy = 0.0
  background_pe    = 0.0
  bgpe_components  = []
  rho_0 = unesco_density(6.48,0.0)
  dz = list(map(float, data[1]))
  for column in data[2]:
	area     = float(column["delta_x"] )*float(column["delta_y"])
	btm_z    = float(column["bottom_z"])
	sfc_z    = float(column["surface_z"])
	cell_top = 0.0
	cell_btm = 0.0
	total_base_area += area
	for k in range(int(column["bottom_k"]), int(column["surface_k"])+1):
	  cell_top += dz[k]
	  cell_height = min(cell_top, sfc_z)-max(cell_btm, btm_z)
	  cell_mid = 0.5*(min(cell_top, sfc_z)+max(cell_btm, btm_z))
	  cell_btm = cell_top # for next loop
	  u   = column["variables"]["u"          ][k]
	  w   = column["variables"]["w"          ][k]
	  s   = column["variables"]["salinity"   ][k]
	  t   = column["variables"]["temperature"][k]
	  rho = unesco_density(t,s)
	  vol = area*cell_height
	  bgpe_components.append((rho, vol))
	  kinetic_energy   += vol*0.5*rho_0*(u*u+w*w)
	  potential_energy += vol*rho*9.8*cell_mid

  layer_top = 0.0
  layer_btm = 0.0    
  for component in sorted(bgpe_components)[::-1]:
	rho = component[0]
	vol = component[1]
	layer_top += vol/total_base_area
	layer_mid  = 0.5*(layer_top+layer_btm)
	background_pe += vol*rho*9.8*layer_mid
	layer_btm = layer_top
  return (kinetic_energy, potential_energy-background_pe)

def main():
 
  files = glob.glob("intwave-0*.dat")
  files.sort()
  
  t,ke, pe, te = [],[],[],[]
  for file in files:
	data    = get_data(file)
	time    = data[0]
	columns = data[2] 
	t.append(time)
	energy = energy_estimation(data)
	te_temp = energy[0]+energy[1]
	if files.index(file)==0:
	  ini_te = te_temp

	ke.append(energy[0]/ini_te)
	pe.append(energy[1]/ini_te)
	te.append(te_temp/ini_te)

  print(ke, pe, te)
  plt.plot(t,te, label='total energy')
  plt.plot(t,pe, label='potential energy')
  plt.plot(t,ke, label='kinetic energy')
  plt.legend()
  plt.savefig('energy.png')
  plt.show()

if __name__ == '__main__':
  main()