convert_files.sh

#!/bin/bash
#SBATCH --nodes=1
#SBATCH --time=16:00:00
#SBATCH --job-name=vis
#SBATCH --output=vis.o%j
#SBATCH --error=vis.e%j         
#SBATCH --qos=standard
#SBATCH -A w19_mpascoastal 
#SBATCH --mail-user=pwolfram@lanl.gov
#SBATCH --mail-type=ALL

eval "$(/lustre/scratch4/turquoise/pwolfram/BGC_wrangling/miniconda3/bin/conda shell.bash hook)"

source activate compass_py3.7

paraview_vtk_field_extractor.py \
  -v bottomDepth,latVertex,lonVertex \
  -f /lustre/scratch3/turquoise/maltrud/ACME/input_data/ocn/mpas-o/oRRS30to10v3/mpaso.rst.0026-01-01_00000.GMPAS-IAF_oRRS30to10v3_theta04.Apr07_2018+BGC_IC.Nov30_2017.FESEDx5.nc \
  -x none \
  -o './data'

paraview_vtk_field_extractor.py \
  -v timeMonthly_avg_velocityMeridional,timeMonthly_avg_velocityZonal,timeMonthly_avg_ssh,timeMonthly_avg_layerThickness,timeMonthly_avg_potentialDensity,timeMonthly_avg_dThreshMLD,timeMonthly_avg_vertVelocityTop,timeMonthly_avg_SSHSquared,timeMonthly_avg_ecosys_diag_pH_3D,timeMonthly_avg_activeTracers_temperature,timeMonthly_avg_activeTracers_salinity,timeMonthly_avg_ecosysTracers_PO4,timeMonthly_avg_ecosysTracers_NO3,timeMonthly_avg_ecosysTracers_NH4,timeMonthly_avg_ecosysTracers_DON,timeMonthly_avg_ecosysTracers_spChl,timeMonthly_avg_ecosysTracers_diatChl,timeMonthly_avg_ecosysTracers_diazChl,timeMonthly_avg_ecosysTracers_phaeoChl,timeMonthly_avg_ssh\
  -f '/lustre/scratch4/turquoise//maltrud/ACME/cases/GMPAS-OECO-ODMS-IAF_oRRS30to10v3_grizzly03/run/mpaso.hist.am.timeSeriesStatsMonthly.0046-*.nc' \
  -m /lustre/scratch3/turquoise/maltrud/ACME/input_data/ocn/mpas-o/oRRS30to10v3/mpaso.rst.0026-01-01_00000.GMPAS-IAF_oRRS30to10v3_theta04.Apr07_2018+BGC_IC.Nov30_2017.FESEDx5.nc \
  -x none \
  -d nVertLevels=0:79 nVertLevelsP1=0:80 \
  -o './data'

# build the 3D files
python to3D.vpy

# individually zip it up
for i in TEST-world*; do
  tar czvf $i.tgz $i &
done

to3D.vpy

# Greg Abram, Phillip J. Wolfram, 07/03/2019
from vtk import *
from glob import glob
import numpy as np
from vtk.numpy_interface import dataset_adapter as dsa
from vtk.numpy_interface import algorithms as algs
import sys
import math

args = sys.argv[1:]
DEPTH_SCALE = 300
do_uvw = False
do_latlon = True

variables = []

while len(args) > 0:
  if args[0] == 'd':
    DEPTH_SCALE = float(args[1])  # vertical extent to exagerate veritcal dimension
    args = args[2:]
  elif args[0] == '+uvw':
    do_uvw = True  # uvw is vector zonal, meridiional, vertical vectors
    args = args[1:]
  elif args[0] == '-uvw':
    do_uvw = False
    args = args[1:]
  elif args[0] == '+latlon':
    do_latlon = True  # latlon project
    args = args[1:]
  elif args[0] == '-latlon':
    do_latlon = False
    args = args[1:]
  else:
    variables.append(args[0])
    args = args[1:]

# files coming in from Xylar's post-processing tool
prdr = vtkXMLPolyDataReader()
prdr.SetFileName('data/staticFieldsOnVertices.vtp')

c2p = vtkCellDataToPointData()
c2p.SetInputConnection(prdr.GetOutputPort())
c2p.Update()

pgrid = dsa.WrapDataObject(c2p.GetOutput())

crdr = vtkXMLPolyDataReader()
crdr.SetFileName('data/staticFieldsOnCells.vtp')
crdr.Update()

# create python wrapper for vtk object
cgrid = dsa.WrapDataObject(crdr.GetOutput())

# bottom depth is now assigned to the triangular vertex points
pgrid.PointData.append(cgrid.CellData['bottomDepth'], 'bottomDepth')

nPoints    = pgrid.Points.shape[0]
nTriangles = pgrid.Polygons.shape[0]/4

lat = pgrid.PointData['latVertex']
lon = pgrid.PointData['lonVertex']

print('got latlon')

# radial vectors through each surface point

rxyz = pgrid.Points / np.sqrt((pgrid.Points*pgrid.Points).sum(axis=1)).reshape(len(pgrid.Points),1)

if do_uvw:
  # compute zonal/meridonal frame of reference for each point on the surface
  # computes the position vectors

  rx = rxyz[:,0]
  ry = rxyz[:,1]
  rz = rxyz[:,2]

  theta = np.arctan2(ry, rx)
  d = np.sqrt(rx*rx + ry*ry)
  phi = np.arctan2(rz, d)

  ijhat = rx[:,np.newaxis]*[1.0, 0.0, 0.0] + ry[:,np.newaxis]*[0.0, 1.0, 0.0]
  d = np.linalg.norm(ijhat, axis=1)
  ijhat = ijhat / d[:,np.newaxis]

  # zonal, meridional, vertical 3D coordinate vector on the sphere
  U = np.column_stack((-np.sin(theta), np.cos(theta), [0]*len(theta)))
  V = -np.sin(phi)[:,np.newaxis]*ijhat + np.cos(phi)[:,np.newaxis]*[0.0, 0.0, 1.0]
  W = rxyz

  print('uvw ready')
else:
  print('not doing uvw')

# and points at sea level

nPoints = len(pgrid.Points)

timesteps = glob('data/time_series/fieldsOnCells.*.vtp')
#timesteps = glob('data/time_series/timeDependentFieldsOnCells.*.vtp')
#timesteps = [timesteps[0]]


# produces grid to be shared for each timestep and assign values on the mesh
# for each timestep
first = True
for tf in timesteps:

  ts = tf.split('.')[1]

  print( 'timestep: {} {}'.format(ts, tf))

  rdr1 = vtkXMLPolyDataReader()
  rdr1.SetFileName(tf)
  rdr1.Update()

  cgrid = dsa.WrapDataObject(rdr1.GetOutput())

  print('loaded')

  if first:

    first = False

    # get number of layers (checks last two digits from Xylar's vtk output)
    nlayers = -1
    for i in cgrid.CellData.keys():
      if i[:30] == 'timeMonthly_avg_layerThickness':
        l = int(i[-2:])
        if nlayers < l:
          nlayers = l

    # create list of slab thicknesses

    thicknesses = [cgrid.CellData['timeMonthly_avg_layerThickness_%02d' % i] for i in range(nlayers+1)]

    print('{} thicknesses'.format(len(thicknesses)))

    # BottomDepth is positive - eg. units below the sea level.   Slab thicknesses are also
    # positive.  Therefore, the top level is the *negative* bottom depth plus the sum
    # of the layer thicknesses.

    # this is equivalent to the ssh (something that could be checked with an assert)
    toplevel = -pgrid.PointData['bottomDepth'] + np.sum(thicknesses, axis=0)

    # arangement to correspond to interior interpolation of the scalar field
    # (ssh,ssh - dz/2) half top layer
    # ...
    # interior layers
    # ...
    # (dz/2, bottomDepth) half bottom layer

    # To handle the cell-centeredness of the data, we force a layer at the top level, then
    # one 1/2 thickness[0] below that, which would be the center of the first slab.  Then we add
    # a layer at the center of each *subsequent* slab, then a layer at the bottom.   This results
    # in n+2 layers, so for cell-centered variables (that have the same number of vertical layers
    # as the layerThickness variable) we copy the first and last data slices.  For the P1 variables,
    # these are top-centered, so we assign the first slice, then interpolate values for each slab
    # center, then dup the last slab-center slice for the bottom slice.

    # depths[i] are the depths of the i'th slice at each point.   The 0'th slice is at the top
    # level; then there's a slice at the center of each slab, then there's the bottom.   The center
    # of slice k (k > 1) is the center of the previous slab (k-1) + the rest of slab that slice
    # (thickness[k-1]/2) plus the top half of the current slab.


    # (ssh,ssh - dz/2) half top layer (builds out whole layer)
    depths = [toplevel, toplevel - (thicknesses[0]/2)]

    # interior layers (including top of the very bottom layer)
    for i in range(1, len(thicknesses)):
      depths.append(depths[-1] - ((thicknesses[i-1] + thicknesses[i]) / 2))

    # (dz/2, bottomDepth) half bottom layer  (close it out)
    depths.append(-pgrid.PointData['bottomDepth'])

    # diagnostic to check to make sure depths are reasonable
    print('depths at random point on the surface')
    for i,d in enumerate(depths):
      print('{} {}'.format(i, depths[i][10000]))

    # projection with vertical exageration to produce vtk points
    points = []
    for i,depth in enumerate(depths):
      x = pgrid.Points + (DEPTH_SCALE*depth)[:,np.newaxis]*rxyz
      y = np.array(x).astype('f4')
      points.append(y)

    # point values being converted in prep for vtk inputs
    points = np.ascontiguousarray(np.vstack(points)).astype('f4')

    print('points done')

    triangles = pgrid.GetPolygons().reshape((-1,4))[:,[1,2,3]]

    # wedge ("prism") vtk container topology
    prisms = []
    for i in range(len(depths)-1):
      prisms.append(np.hstack(([[6]]*len(triangles), triangles+[i*nPoints]*3, triangles+[(i+1)*nPoints]*3)))

    prisms = np.vstack(prisms).astype('i4')

    nPrisms = prisms.shape[0]
    print('nPrisms {}'.format(nPrisms))

    # build vtk type as wedges
    cellTypes = dsa.VTKArray([VTK_WEDGE]*nPrisms).astype('u1')
    # range doesn't include the last point (e.g., a wedge is defined by 6 points)
    cellLocations = dsa.VTKArray(np.array(range(0, 7*nPrisms, 7)).astype('i4'))

    ogrid = dsa.WrapDataObject(vtkUnstructuredGrid())

    ogrid.SetCells(cellTypes, cellLocations, prisms.flatten())
    ogrid.SetPoints(points)

    print('{} prisms on {} points.  max ID {} ... {}'.format(prisms, len(points), ogrid.GetCells().max(),  len(ogrid.Cells)))

    # appends lat/lon data to produced file
    if do_latlon:
      lat = 90.0 * ((2*lat) / 3.1415926)
      lat = np.concatenate([lat]*len(depths)).astype('f4')
      ogrid.PointData.append(lat, 'lat')
      lon = np.where(lon > 180, lon-360, lon)
      lon = np.concatenate([lon]*len(depths)).astype('f4')
      ogrid.PointData.append(lon, 'lon')
      print('stacked latlon done')
    else:
      print('not doing uvw')


    # vectors to project scalar values into vector
    # multiply velocityZonal, velocityMeridional, etc to project in 3D to build VTK/ParaView vector
    if do_uvw:
      U = np.concatenate([U]*len(depths)).astype('f4')
      V = np.concatenate([V]*len(depths)).astype('f4')
      W = np.concatenate([W]*len(depths)).astype('f4')
      ogrid.PointData.append(U, 'u')
      ogrid.PointData.append(V, 'v')
      ogrid.PointData.append(W, 'w')
      print('uvw done')
    else:
      print('not doing uvw')

  otf = 'TEST-world-%d-%s.vtu' % (DEPTH_SCALE, ts)

  # get variable names
  if len(variables) == 0:
    variables = [i[:-3] for i in cgrid.CellData.keys() if i[-3:] == '_00']

  # for v in variables:
  for v in variables:
    print('starting {}'.format(v))
    vnames = ['%s_%02i' % (v, i) for i in range(len(thicknesses))]
    slices = [np.array(cgrid.CellData[vname]).astype('f4') for vname in vnames]
    if len(slices) == len(thicknesses):

      # Replicate first slab value at top, last at bottom
      data = [slices[0]] + slices + [slices[-1]]

    elif len(slices) == (len(thicknesses)+1):

      # then the data is centered with the *top* of each slab, plus the bottom
      # of the lowest slab.   First slice is correct.  Then, for each *slab*, interpolate
      # between top and bottom values.   Then add in the final slice which,
      # again, is correct

      data = [slices[0]]

      # vertically interpolate a slice at each slab center

      for i in range(len(slices)-1):
        data.append(np.sum([slices[i], slices[i+1]], axis=0) / 2.0)

      # and add the bottom

      data.append(slices[-1])
    else:
      print("unknown data length {} {}".format(v,  len(slices)))
      continue

    data = np.concatenate(data)
    data = np.nan_to_num(data).flatten()  # nans go to zero (should replace with -1E34)
    # needed to do streamlines / particle tracking / particle advection
    # correctly (but note that this won't be a physical value in the output --
    # probably should rename so that this is clear)
    if v == 'timeMonthly_avg_vertVelocityTop':
      data = DEPTH_SCALE * data
    ogrid.PointData.append(data, v)
    print( '{} done {} elements,'.format(v, len(data)))

  print('OGRID =========================================')
  print(ogrid.VTKObject)

  # code that follows is experimental and problem specific (e.g., vis of ice sheets)
  # this helps reduce output  ("clip") in order to produce a smaller file size

  # carve out ice sheets
  threshold1 = vtkThreshold()
  threshold1.SetInputData(ogrid.VTKObject)
  threshold1.SetInputArrayToProcess(1, 0, 0, 0, 'lat')
  #threshold1.ThresholdBetween(-1.48876, -1.261)
  threshold1.ThresholdBetween(0.0579517, 0.96905)
  threshold1.Update()

  print('T1 =========================================')
  print(threshold1.GetOutput())

  threshold2 = vtkThreshold()
  threshold2.SetInputConnection(threshold1.GetOutputPort())
  threshold2.SetInputArrayToProcess(1, 0, 0, 0, 'lon')
  #threshold2.ThresholdBetween(4.84, 5.75)
  threshold2.ThresholdBetween(3.81101, 5.98078)
  threshold2.Update()

  print('T2 =========================================')
  print(threshold2.GetOutput())

  print('writing {}'.format(otf))
  wrtr = vtkXMLUnstructuredGridWriter()
  wrtr.SetFileName(otf)
  wrtr.SetInputData(threshold2.GetOutput())  # need to use clipped output
  wrtr.SetInputData(ogrid.VTKObject)
  wrtr.Write()
  print('{} {} done'.format(ts, tf))

# end the loop: "for tf in timesteps:"