Skip to content

Instantly share code, notes, and snippets.

@pwolfram

pwolfram/interpolate_grids.py

Created January 31, 2019 15:46
Show Gist options
  • Save pwolfram/0f2b83c3dbf492ef06512e95c8d11169 to your computer and use it in GitHub Desktop.
Save pwolfram/0f2b83c3dbf492ef06512e95c8d11169 to your computer and use it in GitHub Desktop.
Download ZIP
Raw
interpolate_grids.py
#!/usr/bin/env python
"""
Interpolation of initial conditions from source grid onto a destination grid.
Currently implementation interpolates
* temperature
* salinity
* layerThickness
Phillip Wolfram, Mark Petersen
01/20/2017
"""
import shutil
import numpy as np
import matplotlib.pyplot as plt
import netCDF4
from scipy.spatial import cKDTree as KDTree
# special timing function
# from http://stackoverflow.com/questions/5478351/python-time-measure-function
from contextlib import contextmanager
import time
@contextmanager
def timeit_context(name):
startTime = time.time()
yield
elapsedTime = time.time() - startTime
print('[{}] finished in {} s'.format(name, int(elapsedTime)))
# general interpolation functions
def get_point_vectors3d(ncdfdata, cartesian): #{{{
"""
inputs:
ncdfdata is an opened netCDF4.Dataset
outputs:
position tuple
"""
maxLevelCell = ncdfdata.variables['maxLevelCell'][:]
npoints = np.sum(maxLevelCell)
# point data
if cartesian:
lonCell = ncdfdata.variables['xCell'][:]
latCell = ncdfdata.variables['yCell'][:]
else:
lonCell = ncdfdata.variables['lonCell'][:]
latCell = ncdfdata.variables['latCell'][:]
# note this is a coarse proxy for the cell center, but it should work for
# this application
vertLevel = ncdfdata.variables['refBottomDepth'][:]
x = np.zeros((npoints))
y = np.zeros((npoints))
z = np.zeros((npoints))
n = 0
for aCell, nLevel in enumerate(maxLevelCell):
# vertical grid
x[n:(n+nLevel)] = lonCell[aCell]
y[n:(n+nLevel)] = latCell[aCell]
z[n:(n+nLevel)] = vertLevel[:nLevel]
# counter
n += nLevel
return np.vstack((x,y,z)).T #}}}
def get_point_vectors2d(ncdfdata,cartesian): #{{{
"""
inputs:
ncdfdata is an opened netCDF4.Dataset
outputs:
position
"""
# point data
if cartesian:
x = ncdfdata.variables['xCell'][:]
y = ncdfdata.variables['yCell'][:]
else:
x = ncdfdata.variables['lonCell'][:]
y = ncdfdata.variables['latCell'][:]
return np.vstack((x,y)).T #}}}
def get_3dcell_data(ncdfdata, datanames=[], ts=0): #{{{
"""
inputs:
ncdfdata is an opened netCDF4.Dataset
datanames is a list of datanames to be interpolated
ts is the time step number
outputs:
position and data
"""
maxLevelCell = ncdfdata.variables['maxLevelCell'][:]
npoints = np.sum(maxLevelCell)
data = np.zeros((len(datanames),npoints))
n = 0
for aCell, nLevel in enumerate(maxLevelCell):
# 3dcell data
for ad, dataname in enumerate(datanames):
data[ad,n:(n+nLevel)] = ncdfdata.variables[dataname][ts, aCell, :nLevel]
# counter
n += nLevel
return data #}}}
def get_2dcell_data(ncdfdata, datanames=[], ts=0): #{{{
"""
inputs:
ncdfdata is an opened netCDF4.Dataset
datanames is a list of datanames to be interpolated
ts is the time step number
outputs:
position and data
"""
nCells = len(ncdfdata.dimensions['nCells'])
data = np.zeros((len(datanames),nCells))
for ad, dataname in enumerate(datanames):
if dataname == 'ssh':
for ac, maxLevel in enumerate(ncdfdata.variables['maxLevelCell']):
data[ad,ac] = -ncdfdata.variables['bottomDepth'][ac] + \
np.sum(ncdfdata.variables['layerThickness'][ts,ac,:maxLevel],axis=-1)
return data #}}}
def set_3dcell_data(ncdfdata, datanames, data, ts=0): #{{{
"""
inputs:
ncdfdata is an opened netCDF4.Dataset
datanames is a list of datanames
data is datavector corresponding to datanames
ts is the time step number
outputs:
in place=> updata data in ncdfdata for datanames
"""
maxLevelCell = ncdfdata.variables['maxLevelCell'][:]
vertLevel = ncdfdata.variables['refBottomDepth'][:]
n = 0
for aCell, nLevel in enumerate(maxLevelCell):
for ad, dataname in enumerate(datanames):
ncdfdata.variables[dataname][ts, aCell, :nLevel] = data[ad,n:(n+nLevel)]
n+= nLevel
return None #}}}
def grid_interp(sgrid, dgrid, ofile, cartesian, interiorscalars=['temperature','salinity']): #{{{
"""
Perform interpolation from a source grid (sgrid) onto a destination grid (pgrid),
writing to an output file (ofile). This interpolation interpolates scalars
list from the source grid onto the destination grid and uses nearest neighbor interpolation.
Interpolated scalars are specified by interiorscalars and the
layerThickness is interpolated on the output grid for z-star coordinates.
Phillip Wolfram
11/18/2015
"""
with timeit_context('Make output file'):
shutil.copyfile(dgrid, ofile)
with timeit_context('Open files'):
sgriddata = netCDF4.Dataset(sgrid,'r')
dgriddata = netCDF4.Dataset(dgrid,'r')
ofiledata = netCDF4.Dataset(ofile,'r+')
with timeit_context('Build vectors of point locations x,y,z and data values'):
cpos3d = get_point_vectors3d(sgriddata,cartesian)
cdata3dcell = get_3dcell_data(sgriddata, interiorscalars)
fpos3d = get_point_vectors3d(dgriddata,cartesian)
cpos2d = get_point_vectors2d(sgriddata,cartesian)
cdata2dcell = get_2dcell_data(sgriddata, ['ssh'])
fpos2d = get_point_vectors2d(dgriddata,cartesian)
with timeit_context('Building trees for source data'):
tree3d = KDTree(cpos3d)
tree2d = KDTree(cpos2d)
with timeit_context('Finding nearest neighbors'):
_, nearestneighs3d = tree3d.query(fpos3d)
_, nearestneighs2d = tree2d.query(fpos2d)
fdata3dcell = cdata3dcell[:,nearestneighs3d]
fssh = cdata2dcell[:,nearestneighs2d]
with timeit_context('Convert ssh to layerThickness'):
rbd = dgriddata.variables['refBottomDepth'][:]
relthick = np.hstack((rbd[0], np.diff(rbd)))
layerThickness = relthick[np.newaxis,:] + (relthick/np.sum(relthick))[np.newaxis,:]*fssh.T
with timeit_context('Saving data'):
set_3dcell_data(ofiledata, interiorscalars, fdata3dcell)
ofiledata.variables['layerThickness'][0,:,:] = layerThickness
with timeit_context('Close files'):
sgriddata.close()
dgriddata.close()
ofiledata.close()
return None #}}}
if __name__ == "__main__":
import argparse
parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawTextHelpFormatter)
parser.add_argument("-s", "--sourcegrid", dest="sgrid", help="Input: Source grid to perform interpolation.")
parser.add_argument("-d", "--destinationgrid", dest="dgrid", help="Input: Destination grid to perform interpolation.")
parser.add_argument("-o", "--outputfile", dest="ofile", help="Output: Interpolated temperature, salinity, and layerThickness on destination grid.")
parser.add_argument("-c", "--cartesian", dest="cartesian", help="If set, use xCell and yCell to interpolate. If not, use latCell and lonCell.", action="store_true")
args = parser.parse_args()
if args.sgrid is None or args.dgrid is None or args.ofile is None:
parser.error('Must fully specify all inputs and outputs')
if args.sgrid == args.ofile:
parser.error('Output grid cannot be the same as the as source input grid.')
if args.dgrid == args.ofile:
parser.error('Output grid cannot be the same as the destination input grid.')
grid_interp(args.sgrid, args.dgrid, args.ofile, args.cartesian)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment