PYTHON: Add faster block interpolation using ndimage.zoom.

Signed-off-by: Grzegorz Kowal <grzegorz@amuncode.org>
This commit is contained in:
Grzegorz Kowal 2021-05-26 18:18:09 -03:00
parent b350ecdedd
commit be7e0ea5f2
3 changed files with 25 additions and 9 deletions

View File

@ -137,7 +137,7 @@ def amun_coordinate(fname, iname):
return None
def amun_dataset(fname, vname, shrink=1, interpolation='rebin', progress=False):
def amun_dataset(fname, vname, shrink=1, interpolation='rebin', order=3, progress=False):
'''
Subroutine to read datasets from AMUN HDF5 snapshots.
@ -432,11 +432,11 @@ def amun_dataset(fname, vname, shrink=1, interpolation='rebin', progress=False):
if ndims == 3:
ib, jb, kb = ibeg[0], ibeg[1], ibeg[2]
ie, je, ke = iend[0], iend[1], iend[2]
ret[kb:ke,jb:je,ib:ie] = interpolate(dataset[:,:,:,l], cm, ng, method=method)
ret[kb:ke,jb:je,ib:ie] = interpolate(dataset[:,:,:,l], cm, ng, method=method, order=order)
else:
ib, jb = ibeg[0], ibeg[1]
ie, je = iend[0], iend[1]
ret[jb:je,ib:ie] = interpolate(dataset[0,:,:,l], cm, ng, method=method)
ret[jb:je,ib:ie] = interpolate(dataset[0,:,:,l], cm, ng, method=method, order=order)
nb += 1

View File

@ -895,7 +895,7 @@ class AmunXML:
return ids, dset
def dataset(self, var, extent=None, maxlev=-1, shrink=1, interpolation='rebin', progress=False):
def dataset(self, var, extent=None, maxlev=-1, shrink=1, interpolation='rebin', order=3, progress=False):
if not var in self.variables:
print("Dataset '{}' not available in the snapshot!".format(var))
@ -983,14 +983,14 @@ class AmunXML:
if isinstance(dset, (list, tuple)):
for di in range(3):
if self.attributes['ndims'] == 3:
arr[di][ib[2]:ie[2],ib[1]:ie[1],ib[0]:ie[0]] = interpolate(dset[di][:,:,:,p], cm, self.attributes['nghosts'], method=method)[jb[2]:je[2],jb[1]:je[1],jb[0]:je[0]]
arr[di][ib[2]:ie[2],ib[1]:ie[1],ib[0]:ie[0]] = interpolate(dset[di][:,:,:,p], cm, self.attributes['nghosts'], method=method, order=order)[jb[2]:je[2],jb[1]:je[1],jb[0]:je[0]]
else:
arr[di][ib[1]:ie[1],ib[0]:ie[0]] = interpolate(dset[di][:,:,p], cm, self.attributes['nghosts'], method=method)[jb[1]:je[1],jb[0]:je[0]]
arr[di][ib[1]:ie[1],ib[0]:ie[0]] = interpolate(dset[di][:,:,p], cm, self.attributes['nghosts'], method=method, order=order)[jb[1]:je[1],jb[0]:je[0]]
else:
if self.attributes['ndims'] == 3:
arr[ib[2]:ie[2],ib[1]:ie[1],ib[0]:ie[0]] = interpolate(dset[:,:,:,p], cm, self.attributes['nghosts'], method=method)[jb[2]:je[2],jb[1]:je[1],jb[0]:je[0]]
arr[ib[2]:ie[2],ib[1]:ie[1],ib[0]:ie[0]] = interpolate(dset[:,:,:,p], cm, self.attributes['nghosts'], method=method, order=order)[jb[2]:je[2],jb[1]:je[1],jb[0]:je[0]]
else:
arr[ib[1]:ie[1],ib[0]:ie[0]] = interpolate(dset[:,:,p], cm, self.attributes['nghosts'], method=method)[jb[1]:je[1],jb[0]:je[0]]
arr[ib[1]:ie[1],ib[0]:ie[0]] = interpolate(dset[:,:,p], cm, self.attributes['nghosts'], method=method, order=order)[jb[1]:je[1],jb[0]:je[0]]
if progress:
mfile = 'datablocks_{:06d}.xml'.format(n)

View File

@ -35,6 +35,7 @@
"""
import numpy as np
try:
from scipy.ndimage import zoom
from scipy.interpolate import splrep, splev, interp1d, pchip_interpolate
scipy_available = True
except ImportError:
@ -65,7 +66,7 @@ def rebin(a, newshape):
return(a)
def interpolate(a, newshape, nghosts, method = 'rebin'):
def interpolate(a, newshape, nghosts, method='rebin', order=3):
'''
Subroutine rescales the block by interpolating its values.
'''
@ -93,6 +94,21 @@ def interpolate(a, newshape, nghosts, method = 'rebin'):
return rebin(a[jb:je,ib:ie], newshape)
elif (method == 'zoom'):
fc = int(newshape[1] / (a.shape[1] - 2 * nghosts))
ib, ie = fc * nghosts, - fc * nghosts
jb, je = fc * nghosts, - fc * nghosts
if a.ndim == 3:
kb, ke = fc * nghosts, - fc * nghosts
if (a.shape[0] == 1):
kb = 0
ke = 1
return zoom(a, fc, order=order, grid_mode=True, mode='nearest')[kb:ke,jb:je,ib:ie]
else:
return zoom(a, fc, order=order, grid_mode=True, mode='nearest')[jb:je,ib:ie]
elif (method == 'monotonic' or method == 'pchip'):
dims = np.arange(a.ndim)