Python vtktools.arr函数代码示例

def calc_mld(files,start,x0=0.0,y0=0.0):
    """ Caclulate density-based MLD from a bunch of VTU files
    """

    mld = []
    times = []
    dates = []
    for file in files:
      
        try:
            os.stat(file)
        except:
            print("No such file: %s" % file)
            sys.exit(1)

        # open vtu and derive the field indices of the edge at (x=0,y=0) ordered by depth
        u=vtktools.vtu(file)
        pos = u.GetLocations()
        ind = get_1d_indices(pos, x0, y0)
    
        # from this we can derive the 1D profile of any field like this:
        depth = vtktools.arr([-pos[i,2] for i in ind])
    
        # handle time for different types of plots
        time = u.GetScalarField('Time')
        times.append(time[0])   # seconds
        dates.append( date2num(start + timedelta(seconds=time[0])) ) # integer datetime
    
        # grab density profile and calculate MLD_den (using 2 different deviation parameters
        d = u.GetScalarField('Density')
        den = vtktools.arr( [d[i] * 1000 for i in ind] )
        mld.append( calc_mld_den(den, depth) ) #den0 = 0.03 is default


    return mld, times, dates

def calc_mld_tke_files(files,start,x0=0.0,y0=0.0):
    """ Caclulate tke-based MLD from a bunch of VTU files
    """

    mld = []
    times = []
    dates = []
    for file in files:
      
        try:
            os.stat(file)
        except:
            print "No such file: %s" % file
            sys.exit(1)

        # open vtu and derive the field indices of the edge at (x=0,y=0) ordered by depth
        u=vtktools.vtu(file)
        pos = u.GetLocations()
        ind = get_1d_indices(pos, x0, y0)
    
        # from this we can derive the 1D profile of any field like this:
        depth = vtktools.arr([-pos[i,2] for i in ind])
    
        # handle time for different types of plots
        time = u.GetScalarField('Time')
        times.append(time[0])   # seconds
        dates.append( date2num(start + timedelta(seconds=time[0])) ) # integer datetime
    
        # grab density profile and calculate MLD
        d = u.GetScalarField('GLSTurbulentKineticEnergy')
        tke = vtktools.arr( [d[i] for i in ind] )
        mld.append( calc_mld_tke(tke, depth) )


    return mld, times, dates

def plot_2d_data(data,depths,time_secs,start_date,file_path,axis_label,finish_date=None,mld_data=None,max_depth=150,interval=3,minimum=None,maximum=None,spacing=None,colour_scale=cm.jet,dates=None):
    """
    """
    # turn given 2d-arrays into numpy arrays (in case they are not already)
    data = vtktools.arr(data)
    time_secs = vtktools.arr(time_secs)
    depths = vtktools.arr(depths)
    
    # convert time profiles in seconds into months
    start = datetime.strptime(start_date, "%Y-%m-%d %H:%M:%S")
    if (dates == None):
        dates = time_secs
        i = 0
        for time in time_secs:
            t = float(time[0].item())
            dates[i,:] = date2num(start + timedelta(seconds=t))
            i += 1
        
    # see if finishing date is given, default to last time given
    if (finish_date != None):
        finish = datetime.strptime(finish_date, "%Y-%m-%d %H:%M:%S")
    else:
        finish = dates[-1][0] 
    
    # define min/max and spacing of data if not given (so we see all of the data)
    if (minimum == None):
        minimum = data.min() 
        minimum = minimum - (0.1*minimum) 
    if (maximum == None):
        maximum = data.max()
        maximum = maximum + (0.1*maximum) 
    if (spacing == None):
        spacing = (maximum - minimum) /256.

    # plot 2d colour graph...
    fig = figure(figsize=(15,8),dpi=90)
    ax = fig.add_axes([.1,.18,.9,.7])
    cs=ax.contour(dates, depths, data, arange(minimum,maximum,spacing),cmap=colour_scale)
    cs=ax.contourf(dates, depths, data, arange(minimum,maximum,spacing),cmap=colour_scale)
    pp=colorbar(cs,format='%.2f')
    if(mld_data!=None):
        ax.plot(dates[:,0],mld_data,'w', alpha=0.7)
    
    dateFmt = mpl.dates.DateFormatter('%m/%Y')
    ax.xaxis.set_major_formatter(dateFmt)
    monthsLoc = mpl.dates.MonthLocator(interval=interval)
    ax.xaxis.set_major_locator(monthsLoc)
    labels = ax.get_xticklabels()
    for label in labels:
        label.set_rotation(30) 
    ax.set_ylim(max_depth, 0)
    ax.set_xlim(start,finish)
    pp.set_label(axis_label)
    xlabel('Date (mm/yyyy)')
    ylabel('Depth (m)')
    
    form = file_path.split('.')[-1].strip()
    savefig(file_path, dpi=90,format=form)
    close(fig)

def botella_p1(NN):
#Botella and Peyret (1998) Table 9. 
  filelist_not_sorted = glob.glob('driven_cavity-%d/*.vtu'%NN)
  vtu_nos_not_sorted = [int(file.split('.vtu')[0].split('_')[-1]) for file in filelist_not_sorted]
  filelist = [filelist_not_sorted[i] for i in numpy.argsort(vtu_nos_not_sorted)]
  file = filelist[-1]
  print file
  try:
    os.stat(file)
  except:
    print "No such file: %s" % file
    sys.exit(1)

  u=vtktools.vtu(file)
  pts=vtktools.arr([
  [0.5, 0.0000, 0.0,  0.110591],
  [0.5, 0.0547, 0.0,  0.109689],
  [0.5, 0.0625, 0.0,  0.109200],
  [0.5, 0.0703, 0.0,  0.108566],
  [0.5, 0.1016, 0.0,  0.104187],
  [0.5, 0.1719, 0.0,  0.081925],
  [0.5, 0.2813, 0.0,  0.040377],
  [0.5, 0.4531, 0.0,  0.004434],
  [0.5, 0.5000, 0.0,  0.000000],
  [0.5, 0.6172, 0.0, -0.000827],
  [0.5, 0.7344, 0.0,  0.012122],
  [0.5, 0.8516, 0.0,  0.034910],
  [0.5, 0.9531, 0.0,  0.050329],
  [0.5, 0.9609, 0.0,  0.050949],
  [0.5, 0.9688, 0.0,  0.051514],
  [0.5, 0.9766, 0.0,  0.052009],
  [0.5, 1.0000, 0.0,  0.052987]])

  velocity = u.ProbeData(pts, "Velocity")
  (ilen, jlen) = velocity.shape
  pressure = u.ProbeData(pts, "Pressure")

  pts0=vtktools.arr([[0.5, 0.5, 0.0]])  # We're going to subtract off the pressure at the centre point
  press0 = u.ProbeData(pts0, "Pressure")

  norm=0.0
  for i in range(ilen):
      diff = pts[i][3] - (pressure[i][0]-press0[0][0])
      norm = norm + diff*diff

  norm = math.sqrt(norm/ilen)
  print "botella_p1_norm:", norm

  return norm

def botella_p2(NN):
#Botella and Peyret (1998) Table 10. 
  filelist_not_sorted = glob.glob('driven_cavity-%d/*.vtu'%NN)
  vtu_nos_not_sorted = [int(file.split('.vtu')[0].split('_')[-1]) for file in filelist_not_sorted]
  filelist = [filelist_not_sorted[i] for i in numpy.argsort(vtu_nos_not_sorted)]
  file = filelist[-1]
  print file
  try:
    os.stat(file)
  except:
    print "No such file: %s" % file
    sys.exit(1)

  u=vtktools.vtu(file)
  pts=vtktools.arr([
  [1.0000, 0.5, 0.0,  0.077455],
  [0.9688, 0.5, 0.0,  0.078837],
  [0.9609, 0.5, 0.0,  0.078685],
  [0.9531, 0.5, 0.0,  0.078148],
  [0.9453, 0.5, 0.0,  0.077154],
  [0.9063, 0.5, 0.0,  0.065816],
  [0.8594, 0.5, 0.0,  0.049029],
  [0.8047, 0.5, 0.0,  0.034552],
  [0.5000, 0.5, 0.0,  0.000000],
  [0.2344, 0.5, 0.0,  0.044848],
  [0.2266, 0.5, 0.0,  0.047260],
  [0.1563, 0.5, 0.0,  0.069511],
  [0.0938, 0.5, 0.0,  0.084386],
  [0.0781, 0.5, 0.0,  0.086716],
  [0.0703, 0.5, 0.0,  0.087653],
  [0.0625, 0.5, 0.0,  0.088445],
  [0.0000, 0.5, 0.0,  0.090477]])

  velocity = u.ProbeData(pts, "Velocity")
  (ilen, jlen) = velocity.shape
  pressure = u.ProbeData(pts, "Pressure")

  pts0=vtktools.arr([[0.5, 0.5, 0.0]])  # We're going to subtract off the pressure at the centre point
  press0 = u.ProbeData(pts0, "Pressure")

  norm=0.0
  for i in range(ilen):
      diff = pts[i][3] - (pressure[i][0]-press0[0][0])
      norm = norm + diff*diff

  norm = math.sqrt(norm/ilen)
  print "botella_p2_norm:", norm

  return norm

def plot_1d_comparison(data_dict,style_dict,time_dict,start_date,finish_date,file_path,axis_label,interval=3):
    """ 
    """
    start_time = date2num(datetime.strptime(start_date, "%Y-%m-%d %H:%M:%S"))
    finish_time = date2num(datetime.strptime(finish_date, "%Y-%m-%d %H:%M:%S"))
        
    # plot 1d graph...
    fig = figure(figsize=(15,8),dpi=90)
    ax = fig.add_axes([.05,.12,.9,.85])
    max_value = 0.0
    for key, data_arr in data_dict.iteritems():
        ax.plot(time_dict[key],data_arr,style_dict[key], label=key)
        data_arr = vtktools.arr(data_arr)
        if data_arr.max() > max_value:
            max_value = data_arr.max()
    max_value += max_value * 0.1
    
    dateFmt = mpl.dates.DateFormatter('%m/%Y')
    ax.xaxis.set_major_formatter(dateFmt)
    monthsLoc = mpl.dates.MonthLocator(interval=interval)
    ax.xaxis.set_major_locator(monthsLoc)
    labels = ax.get_xticklabels()
    for label in labels:
        label.set_rotation(30) 
    ax.set_ylim(max_value, 0)
    ax.set_xlim(start_time,finish_time)
    xlabel('Date (mm/yyyy)')
    ylabel(axis_label)
    legend(loc=0)
    
    form = file_path.split('.')[-1].strip()
    savefig(file_path, dpi=90,format=form)
    close(fig)

def MLD(filelist):
  x0 = 0.
  tke0 = 1.0e-5
  last_mld = 0
  
  times = []
  depths = []
  Dm = []
  for file in filelist:
     try:
       os.stat(file)
     except:
       print "No such file: %s" % file
       sys.exit(1)
     
     u=vtktools.vtu(file)
     time = u.GetScalarField('Time')
     tt = time[0]
     kk = u.GetScalarField('GLSTurbulentKineticEnergy')
     pos = u.GetLocations()
     # ignore first 4 hours of simulaiton
     if (tt < 14400):
       continue

     xyzkk = []
     for i in range(0,len(kk)):
       if( abs(pos[i,0] - x0) < 0.1 ):
         xyzkk.append((pos[i,0],-pos[i,1],pos[i,2],(kk[i])))

     xyzkkarr = vtktools.arr(xyzkk)
     III = argsort(xyzkkarr[:,1])
     xyzkkarrsort = xyzkkarr[III,:]
     # march down the column, grabbing the last value above tk0 and the first 
     # one less than tke0. Interpolate between to get the MLD
     kea = 1000
     keb = 0
     zza = 0
     zzb = 0
     for values in xyzkkarrsort:
        if (values[3] > tke0):
            kea = values[3]
            zza = -values[1]
        if (values[3] < tke0):
            keb = values[3]
            zzb = -values[1]
            break

     # the MLD is somewhere between these two values - let's estimate half way!
     mld = (zzb+zza)/2.
     if (last_mld == mld):
        continue

     times.append(tt/3600)
     depths.append(-1.0*mld)
     last_mld = mld
     Dm.append(1.05*0.00988211768*(1.0/sqrt(0.01))*sqrt(tt))


  return times, depths, Dm

def plot_tracer(vtu_file, line=None):
    
    u=vtktools.vtu(vtu_file)
    pos = u.GetLocations()
    time = u.GetScalarField('Time')[0]
    t = u.GetScalarField('Tracer')
    ind = get_1d_indices(pos)
    distance = vtktools.arr([pos[i,0] for i in ind])
    tracer = vtktools.arr( [t[i] for i in ind] )

    if (line):
        line.set_ydata(tracer)  # update the data
        draw()
    else:
        line, = plot(distance,tracer,color='k')
        
    return line

def MLD(filelist):
  x0 = 0.
  tke0 = 1.0e-5
  last_mld = 0
  
  times = []
  depths = []
  for file in filelist:
     try:
       os.stat(file)
     except:
       print("No such file: %s" % file)
       sys.exit(1)
     
     u=vtktools.vtu(file)
     time = u.GetScalarField('Time')
     tt = time[0]
     kk = u.GetScalarField('GLSTurbulentKineticEnergy')
     pos = u.GetLocations()
     if (tt < 100):
       continue

     xyzkk = []
     for i in range(0,len(kk)):
       if( abs(pos[i,0] - x0) < 0.1 ):
         xyzkk.append((pos[i,0],-pos[i,1],pos[i,2],(kk[i])))

     xyzkkarr = vtktools.arr(xyzkk)
     III = argsort(xyzkkarr[:,1])
     xyzkkarrsort = xyzkkarr[III,:]
     # march down the column, grabbing the last value above tk0 and the first 
     # one less than tke0. Interpolate between to get the MLD
     kea = 1000
     keb = 0
     zza = 0
     zzb = 0
     for values in xyzkkarrsort:
        if (values[3] > tke0):
            kea = values[3]
            zza = -values[1]
        if (values[3] < tke0):
            keb = values[3]
            zzb = -values[1]
            break

     mld = zza
     if (last_mld == mld):
        continue

     times.append(tt/3600)
     depths.append(-1.0*mld)
     last_mld = mld

  return times, depths

def getDistanceMeshDensity(file):
  v = vtktools.vtu(file)
  l = [0.0] * v.ugrid.GetNumberOfPoints()
  a = vtktools.arr(l)
  for i in range(v.ugrid.GetNumberOfPoints()):
    neighbours = v.GetPointPoints(i)
    sum = 0.0
    for neighbour in neighbours:
      sum = sum + v.GetDistance(i, neighbour)
    a[i] = sum / len(neighbours)
  return a

def getData(file, xmin=float('nan'), xmax=float('nan'), ymin=float('nan'), ymax=float('nan'), step_x=1,step_y=1):
  u=vtktools.vtu(file)

  print numpy.isnan(xmin), numpy.isnan(xmax), numpy.isnan(ymin), numpy.isnan(ymax)

  if numpy.isnan(xmin): 
    xmin=u.ugrid.GetBounds()[0]
    print 'xmin = ', xmin
  if numpy.isnan(ymin):
    ymin=u.ugrid.GetBounds()[2]
    print 'ymin = ', ymin
  if numpy.isnan(xmax):
    xmax=u.ugrid.GetBounds()[1]
    print 'xmax = ', xmax
  if numpy.isnan(ymax):
    ymax=u.ugrid.GetBounds()[3]
    print 'ymax = ', ymax
  
  Xlist = numpy.arange(xmin,xmax,step_x)# x coordinates
  Ylist = numpy.arange(ymin,ymax,step_y)# y coordinates
  [X0,Y0] = scipy.meshgrid(Xlist,Ylist)
  X0=X0.transpose()
  Y0=Y0.transpose()
  print Xlist.shape, Ylist.shape, X0.shape, Y0.shape
  Z0 = 0.0*Y0 # This is 2d so z is an array of zeros.
  X = numpy.reshape(X0,(numpy.size(X0),))
  Y = numpy.reshape(Y0,(numpy.size(Y0),))
  Z = numpy.reshape(Z0,(numpy.size(Z0),))
  pts = zip(X,Y,Z)
  pts = vtktools.arr(pts)
  # create arrays of velocity and temperature values at the desired points
  sol = u.ProbeData(pts, 'solution')
  print sol.shape, Xlist.shape, Ylist.shape
  #temperature_structured = u.ProbeData(pts, 'Temperature')
  numpy.savetxt("pts.dat", zip(X,Y,sol))

  sol = sol.reshape((numpy.size(Xlist),numpy.size(Ylist)))
  x=[]
  y=[]
  z=[]
  for i in range(len(Xlist)):
    for j in range(len(Ylist)):
      #print i,j
      x.append(X0[i,j])
      y.append(Y0[i,j])
      z.append(sol[i,j])
  numpy.savetxt("pts2.dat", numpy.array(zip(x, y, z)))

  #data = scipy.interpolate.RectBivariateSpline(Xlist,Ylist,sol)
  
  # return Xlist, Ylist, sol
  return Xlist,Ylist,sol

def erturk_v(NN):
#Erturk et al 2005. Table VII
  filelist_not_sorted = glob.glob('driven_cavity-%d/*.vtu'%NN)
  vtu_nos_not_sorted = [int(file.split('.vtu')[0].split('_')[-1]) for file in filelist_not_sorted]
  filelist = [filelist_not_sorted[i] for i in numpy.argsort(vtu_nos_not_sorted)]
  file = filelist[-1]
  print file
  try:
    os.stat(file)
  except:
    print "No such file: %s" % file
    sys.exit(1)

  u=vtktools.vtu(file)
  pts=vtktools.arr([
  [0.000, 0.5, 0.0,  0.0000],
  [0.015, 0.5, 0.0,  0.1019],
  [0.030, 0.5, 0.0,  0.1792],
  [0.045, 0.5, 0.0,  0.2349],
  [0.060, 0.5, 0.0,  0.2746],
  [0.075, 0.5, 0.0,  0.3041],
  [0.090, 0.5, 0.0,  0.3273],
  [0.105, 0.5, 0.0,  0.3460],
  [0.120, 0.5, 0.0,  0.3605],
  [0.135, 0.5, 0.0,  0.3705],
  [0.150, 0.5, 0.0,  0.3756],
  [0.500, 0.5, 0.0,  0.0258],
  [0.850, 0.5, 0.0, -0.4028],
  [0.865, 0.5, 0.0, -0.4407],
  [0.880, 0.5, 0.0, -0.4803],
  [0.895, 0.5, 0.0, -0.5132],
  [0.910, 0.5, 0.0, -0.5263],
  [0.925, 0.5, 0.0, -0.5052],
  [0.940, 0.5, 0.0, -0.4417],
  [0.955, 0.5, 0.0, -0.3400],
  [0.970, 0.5, 0.0, -0.2173],
  [0.985, 0.5, 0.0, -0.0973],
  [1.000, 0.5, 0.0,  0.0000]])
  
  velocity = u.ProbeData(pts, "Velocity")
  (ilen, jlen) = velocity.shape
  norm=0.0
  for i in range(ilen):
      diff = pts[i][3] - velocity[i][1]
      norm = norm + diff*diff
  
  norm = math.sqrt(norm/ilen)
  print "erturk_v_norm:", norm

  return norm

def erturk_u(NN):
#Erturk et al 2005. Table VI
  filelist_not_sorted = glob.glob('driven_cavity-%d/*.vtu'%NN)
  vtu_nos_not_sorted = [int(file.split('.vtu')[0].split('_')[-1]) for file in filelist_not_sorted]
  filelist = [filelist_not_sorted[i] for i in numpy.argsort(vtu_nos_not_sorted)]
  file = filelist[-1]
  print file
  try:
    os.stat(file)
  except:
    print "No such file: %s" % file
    sys.exit(1)

  u=vtktools.vtu(file)
  pts=vtktools.arr([[0.5, 0.000, 0.000, 0.0000],
  [0.5, 0.000, 0.000,  0.0000],
  [0.5, 0.020, 0.000, -0.0757],
  [0.5, 0.040, 0.000, -0.1392],
  [0.5, 0.060, 0.000, -0.1951],
  [0.5, 0.080, 0.000, -0.2472],
  [0.5, 0.100, 0.000, -0.2960],
  [0.5, 0.120, 0.000, -0.3381],
  [0.5, 0.140, 0.000, -0.3690],
  [0.5, 0.160, 0.000, -0.3854],
  [0.5, 0.180, 0.000, -0.3869],
  [0.5, 0.200, 0.000, -0.3756],
  [0.5, 0.500, 0.000, -0.0620],
  [0.5, 0.900, 0.000,  0.3838],
  [0.5, 0.910, 0.000,  0.3913],
  [0.5, 0.920, 0.000,  0.3993],
  [0.5, 0.930, 0.000,  0.4101],
  [0.5, 0.940, 0.000,  0.4276],
  [0.5, 0.950, 0.000,  0.4582],
  [0.5, 0.960, 0.000,  0.5102],
  [0.5, 0.970, 0.000,  0.5917],
  [0.5, 0.980, 0.000,  0.7065],
  [0.5, 0.990, 0.000,  0.8486],
  [0.5, 1.000, 0.000,  1.0000]])
  
  velocity = u.ProbeData(pts, "Velocity")
  (ilen, jlen) = velocity.shape
  norm=0.0
  for i in range(ilen):
      diff = pts[i][3] - velocity[i][0]
      norm = norm + diff*diff
  
  norm = math.sqrt(norm/ilen)
  print "erturk_u_norm:", norm

  return norm

def getElementMeshDensity(file):
  v = vtktools.vtu(file)
  l = [0.0] * v.ugrid.GetNumberOfPoints()
  a = vtktools.arr(l)
  c = v.ugrid.GetCell(1)

  for i in range(v.ugrid.GetNumberOfPoints()):
    eles = v.GetPointCells(i)
    sum = 0.0
    for ele in eles:
      points = v.ugrid.GetCell(ele).GetPoints().GetData()
      sum = sum + c.ComputeVolume(points.GetTuple3(1), points.GetTuple3(2),
                                  points.GetTuple3(3), points.GetTuple3(4))
    a[i] = sum / len(eles)
  return a

def calc_mld(files,mld_depths):
    for file in files:
        try:
            os.stat(file)
        except:
            print "No such file: %s" % file
            sys.exit(1)
        num = int(file.split(".vtu")[0].split('_')[-1])
        u=vtktools.vtu(file)
        pos = u.GetLocations()
        time = u.GetScalarField('Time')
        tt = time[0]

        # grab the data e need from the VTUs and shove in an array
        den = u.GetScalarField('Density')
        xyz_data = []
        for i in range(0,len(den)):
            if (x0-0.1 < pos[i,0] < x0+0.1 and y0-0.1 < pos[i,1] < y0+0.1):
                xyz_data.append((pos[i,0],pos[i,1],-pos[i,2]+z0,1000*den[i]))
       

        # sorted the array based on depth
        xyz_data = vtktools.arr(xyz_data)
        III = argsort(xyz_data[:,2])
        xyz_data_sorted = xyz_data[III,:]

        # Surface values
        sden = xyz_data_sorted[0,3]

        # grab any values where the UML condition is met for temperature, density and TKE
        uml_den = ((xyz_data_sorted[:,3]) <= (sden+den0)).nonzero()

        # ...on density
        if( (size(uml_den) > 0 ) ):
            LL = uml_den[-1][-1]
            zz = xyz_data_sorted[:,2]
            if (LL+1 < size(zz)):
                zza = zz[LL+1]
                kea = xyz_data_sorted[LL+1,3]
            else:
                zza = zz[LL]
                kea = xyz_data_sorted[LL,3]
            zzb = zz[LL]
            keb = xyz_data_sorted[LL,3]
            tt = tt/(24*60*60)
            time = start_datetime + timedelta(days=tt)
            key = '%04d/%02d/%02d' % (time.year, time.month, time.day)
            mld_depths[key] = (zza-(zzb-zza)*(((sden+den0)-kea)/(kea-keb)))