Python vtktools.vtu函数代码示例

def probe_visc():
    # Files and probing resolution:
    vtufile = "Benchmark_Case_1a_1.pvtu"
    npoints = 500

    # First probe viscosity at y = 550e3

    # Setup Coordinates:
    colx               = numpy.array([numpy.linspace(0,500e3,npoints)]).reshape(npoints,1)
    colz               = numpy.zeros((npoints,1))
    coly               = numpy.ones((npoints,1))*550e3
    coordinates        = numpy.concatenate((colx,coly,colz),1) 

    # Open file and probe for Viscosity:
    vtu                = vtktools.vtu(vtufile)
    viscosity_y_550    = vtktools.vtu.ProbeData(vtu,coordinates,'Mantle::Viscosity')[:,0,0]

    # Next probe viscosity at x = 500e3

    # Setup Coordinates:
    coly               = numpy.array([numpy.linspace(0,660e3,npoints)]).reshape(npoints,1)
    colz               = numpy.zeros((npoints,1))
    colx               = numpy.ones((npoints,1))
    coordinates        = numpy.concatenate((colx,coly,colz),1) 

    # Open file and probe for Viscosity:
    vtu                = vtktools.vtu(vtufile)
    viscosity_x_500    = vtktools.vtu.ProbeData(vtu,coordinates,'Mantle::Viscosity')[:,0,0]

    return min(viscosity_y_550), max(viscosity_y_550), min(viscosity_x_500), max(viscosity_x_500)

def Creating_z(Plane,Slope,Sediment,Start_time,End_time,Time_step):

	vtuObject = vtktools.vtu(Plane)
	vtuObject.GetFieldNames()
	gradient = vtuObject.GetScalarField('u')
	ugrid = vtk.vtkUnstructuredGrid()
	gridreader=vtk.vtkXMLUnstructuredGridReader()
	gridreader.SetFileName(Plane)
	gridreader.Update()
	ugrid = gridreader.GetOutput()
	points = ugrid.GetPoints()

	nPoints = ugrid.GetNumberOfPoints()
	for p in range(0,nPoints):
		x = (points.GetPoint(p)[:2] + (gradient[p],))
		points.SetPoint(p,x)
    
	ugrid.Update()
###################################################################################################################
	t = Start_time
	dt = Time_step
	et = End_time
	while t <= et:

		Import = Sediment + str(t) +'000000.vtu'
		NewSave = Sediment + str(t) + '_sed_slope.pvd'
		vtuObjectSed = vtktools.vtu(Import)
		vtuObjectSed.GetFieldNames()
		gradientSed = vtuObjectSed.GetScalarField('u')
		sedgrid = vtk.vtkUnstructuredGrid()
		sedgridreader=vtk.vtkXMLUnstructuredGridReader()
		sedgridreader.SetFileName(Import)
		sedgridreader.Update()
		sedgrid = sedgridreader.GetOutput()
		s = sedgrid.GetPoints()
	
		for p in range(0,nPoints):
			x = ((s.GetPoint(p)[0],) + (s.GetPoint(p)[1],) + ((gradientSed[p]+gradient[p]),))
			s.SetPoint(p,x)

		writer = vtk.vtkUnstructuredGridWriter()
		writer.SetFileName(NewSave)
		writer.SetInput(sedgrid)
		writer.Update()
		writer.Write()
		t += dt
	writer = vtk.vtkUnstructuredGridWriter()
	writer.SetFileName(Slope)
	writer.SetInput(ugrid)
	writer.Update()
	writer.Write()

def meanvelo(file,x,y):

  print "\nRunning velocity profile script on files at times...\n"

  ##### create array of points. Correct for origin not at step.
  pts=[]
  for i in range(len(x)):
    for j in range(len(y)):
      pts.append([x[i]+5.0, y[j], 0.0])

  pts=numpy.array(pts)
  profiles=numpy.zeros([x.size, y.size], float)

  datafile = vtktools.vtu(file)

  ##### Get x-velocity
  uvw = datafile.ProbeData(pts, "AverageVelocity")
  u = uvw[:,0]
  u=u.reshape([x.size,y.size])
  for i in range(len(x)):
    umax = max(u[i,:])
    u[i,:] = u[i,:]/umax
  profiles[:,:] = u

  print "\n...Finished writing data files.\n"
  return profiles

def get_water_depths(filelist, xarray, delta):
  results = []
  for f in filelist:
    try:
      os.stat(f)
    except:
      print "No such file: %s" % f
      sys.exit(1)
    
    y = numpy.arange(delta/2.0,2.0+delta/2.0,delta)[:,numpy.newaxis]
    
    num = int(f.split(".vtu")[0].split('_')[-1])
    vtu = vtktools.vtu(f)
    for name in vtu.GetFieldNames(): 
      if name.endswith("Time"): 
        time = max(vtu.GetScalarRange(name))
        break
    waterdepths = []
    waterdepths.append(num)
    waterdepths.append(time)
    for x in range(len(xarray)):
      coordinates = numpy.concatenate((numpy.ones((len(y),1))*xarray[x], y, numpy.zeros((len(y),1))),1)
      waterdepths.append(sum(vtu.ProbeData(coordinates, "Water::MaterialVolumeFraction"))[0]*delta)
    
    results.append(waterdepths)
  
  results.sort(key=operator.itemgetter(1))
  results = numpy.array(results)
  return results

def GetXandt(filelist):
  time = []
  X_ns = []
  X_fs = []
  for files in filelist:

      data = vtktools.vtu(files) 
      
      time.append(data.GetScalarField("Time")[0])
      
      # Get X
      data.ugrid.GetPointData().SetActiveScalars('Temperature')
      data = data.ugrid
      
      contour = vtk.vtkContourFilter()
      if vtk.vtkVersion.GetVTKMajorVersion() <= 5:
        contour.SetInput(data)
      else:
        contour.SetInputData(data)
      contour.SetValue(0, 0.0)
      contour.Update()
      polydata = contour.GetOutput()

      bounding_box = polydata.GetBounds()
   
      X_ns.append(bounding_box[1])
      X_fs.append(bounding_box[0])

  return time, X_ns, X_fs

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 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 l2(file, numericalfield, analyticalfield):
  ug = vtktools.vtu(file)
  ug.GetFieldNames()
  uv = ug.GetScalarField(numericalfield)
  ex = ug.GetScalarField(analyticalfield)
  pos = ug.GetLocations()
  x = pos[:,0]; y=pos[:,1]; z=pos[:,2]

  NE = ug.ugrid.GetNumberOfCells()
  ML = zeros(size(x), float)
  for ele in range(NE):
    ndglno = ug.GetCellPoints(ele)
    if(size(ndglno)==4):
      t = tetvol(x[ndglno],y[ndglno],z[ndglno])
    elif(size(ndglno)==3):
      t = triarea(x[ndglno],y[ndglno])
    for nod in ndglno:
      ML[nod] = ML[nod]+t/size(ndglno)
  
  err_x = ex- uv
  
  norm_x = 0.0 
  diff = zeros(size(x), float)
  for nod in range(size(x)):
    norm_x = norm_x + ML[nod]*(err_x[nod])**2


  norm_x = sqrt(abs(norm_x))
  return (norm_x)

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 reattachment_length(filelist):

  print "Calculating reattachment point locations using change of x-velocity sign\n"

  nums=[]; results=[]; files = []
  ##### check for no files
  if (len(filelist) == 0):
    print "No files!"
    sys.exit(1)
  for file in filelist:
    try:
      os.stat(file)
    except:
      print "No such file: %s" % file
      sys.exit(1)
    files.append(file)
  sort_nicely(files)

  for file in files:
    ##### Read in data from vtu
    datafile = vtktools.vtu(file)
    ##### Get time for plot:
    t = min(datafile.GetScalarField("Time"))
    print file, ', elapsed time = ', t

    ##### points near bottom surface, 0 < x < 20
    pts=[]; no_pts = 82; offset = 0.01
    x = 5.0
    for i in range(1, no_pts):
      pts.append((x, offset, 0.0))
      x += 0.25

    pts = numpy.array(pts)

    ##### Get x-velocity on bottom boundary
    uvw = datafile.ProbeData(pts, "AverageVelocity")
    u = []
    u = uvw[:,0]
    points = 0.0

    for i in range(len(u)-1):
      ##### Hack to ignore division by zero entries in u.
      ##### All u should be nonzero away from boundary!
      if((u[i] / u[i+1]) < 0. and u[i+1] > 0. and not numpy.isinf(u[i] / u[i+1])):
        ##### interpolate between nodes. Correct for origin not at step.
        p = pts[i][0] + (pts[i+1][0]-pts[i][0]) * (0.0-u[i]) / (u[i+1]-u[i]) -5.0
        print 'p ', p
        ##### Ignore spurious corner points
        if(p>2):
          points = p
          ##### We have our first point on this plane so...
          break
    print "reattachment point found at: ", points

    ##### Append actual reattachment point and time:
    results.append([points,t])

  return results

def values_per_node(file):
  
  u=vtktools.vtu(file)
  zoo = u.GetScalarField('Zooplankton')
  phyto = u.GetScalarField('Phytoplankton')
  nut = u.GetScalarField('Nutrient')
  det = u.GetScalarField('Detritus')
   
  return phyto, zoo, nut, det

def meanvelo(filelist,xarray,yarray,zarray):

  print "\nRunning velocity profile script on files at times...\n"
  ##### check for no files
  if (len(filelist) < 0):
    print "No files!"
    sys.exit(1)

  ##### create array of points
  pts=[]
  for i in range(len(xarray)):
    for j in range(len(yarray)):
      for k in range(len(zarray)):
        pts.append([xarray[i], yarray[j], zarray[k]])
  pts=numpy.array(pts)

  ##### Create output array of correct shape
  files = 3; filecount = 0
  profiles = numpy.zeros([files, xarray.size, zarray.size], float)

  for file in filelist:
    try:
      os.stat(file)
    except:
      f_log.write("No such file: %s" % files)
      sys.exit(1)

    ##### Only process these 3 datafiles:
    vtu_no = float(file.split('_')[-1].split('.')[0])
    if (vtu_no == 6 or vtu_no == 11 or vtu_no == 33):

      datafile = vtktools.vtu(file)
      t = min(datafile.GetScalarField("Time"))
      print file, ', elapsed time = ', t

      ##### Get x-velocity
      uvw = datafile.ProbeData(pts, "Velocity")
      umax = max(abs(datafile.GetVectorField("Velocity")[:,0]))
      u = uvw[:,0]/umax
      u = u.reshape([xarray.size,yarray.size,zarray.size])

      ##### Spanwise averaging
      usum = numpy.zeros([xarray.size,zarray.size],float)
      usum = numpy.array(usum)
      for i in range(len(yarray)):
        uav = u[:,i,:]
        uav = numpy.array(uav)
        usum += uav
      usum = usum / len(yarray)
      profiles[filecount,:,:] = usum

      ##### reset time to something big to prevent infinite loop
      t = 100.
      filecount += 1

  print "\n...Finished extracting data.\n"
  return profiles

def inf(file, numericalfield, analyticalfield):
  ug = vtktools.vtu(file)
  ug.GetFieldNames()
  uv = ug.GetScalarField(numericalfield)
  ex = ug.GetScalarField(analyticalfield)

  err_x = ex - uv
        
  norm_x = max(abs(err_x))
  return (norm_x)

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