流速ベクトルを表示するスクリプト

import numpy as np
import glob
import matplotlib
import matplotlib.pyplot as plt
import matplotlib.gridspec as gridspec
import matplotlib.patches as patches
import copy

# function for reading fantom dat
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":
      key_names =items[1:]
    elif items[0] == "column":
      dic = {"variables":{}}
      for n, name in enumerate(key_names):
        dic[name] = items[n+1]
      columns.append(dic)
    else:
      columns[int(items[0])]["variables"][items[1]]=[float(x) for x in items[2:]]
  return time, dzs, columns

# prepare list 
file_names = glob.glob("./kasumi*.dat")

# setup fig properties
fig  = plt.figure(figsize=(10,10))
gs   = gridspec.GridSpec(2, 2)  
axes = [plt.subplot(gs[0,0]), plt.subplot(gs[0,1]), plt.subplot(gs[1,0]), plt.subplot(gs[1,1])]
gs.update(left=0.1, right=0.9, bottom=0.1, top=0.9, hspace=0.3, wspace=0.2)

for num,fname in enumerate(file_names):
  time, dzs, columns = get_data(fname)

  #plot section
  fig.suptitle('velocity plot', size=20)
  x, y, us, vs , ub, vb = [], [], [], [], [], []
  for column in columns:
    if column["dry_wet"]=="wet":
      x.append(float(column["center_x"]))
      y.append(float(column["center_y"]))
      us.append(float(column["variables"]["u"][int(column["sfc_k"])]))
      vs.append(float(column["variables"]["v"][int(column["sfc_k"])]))
      ub.append(float(column["variables"]["u"][int(column["btm_k"])]))
      vb.append(float(column["variables"]["v"][int(column["btm_k"])]))

  local_range =(10000, 21000, 14000, 23000)
  rect = patches.Rectangle((local_range[0],local_range[1]),
                 local_range[2]-local_range[0],
                 local_range[3]-local_range[1],
                 linewidth=2,edgecolor='red',fill = False)

  axes[0].add_patch(rect)

  axes[0].quiver(x,y,us,vs,scale=10)
  axes[1].quiver(x,y,ub,vb,scale=10)
  axes[0].set_title('surface velocity')
  axes[1].set_title('bottom velocity')

  axes[2].quiver(x,y,us,vs,scale=0.5)
  axes[3].quiver(x,y,ub,vb,scale=0.5)
  axes[2].set_xlim(local_range[0], local_range[2])
  axes[2].set_ylim(local_range[1], local_range[3])
  axes[3].set_xlim(local_range[0], local_range[2])
  axes[3].set_ylim(local_range[1], local_range[3])

  axes[2].set_title('surface velocity (local)')
  axes[3].set_title('bottom velocity (local)')

  plt.savefig("vector_fig_"+fname[-10:-4]+".png")