def test_create_merged_weight_file_unstructured(self):
self.remove_dir = False
ufile = self.get_temporary_file_path('ugrid.nc')
resolution = 10.
self.fixture_regular_ugrid_file(ufile, resolution, crs=Spherical())
src_grid = RequestDataset(ufile, driver=DriverNetcdfUGRID, grid_abstraction='point').get().grid
self.assertEqual(src_grid.abstraction, 'point')
dst_grid = self.get_gridxy_global(resolution=20., crs=Spherical())
dst_path = self.get_temporary_file_path('dst.nc')
dst_grid.parent.write(dst_path)
gs = GridSplitter(src_grid, dst_grid, (3, 3), check_contains=False, src_grid_resolution=10.,
paths=self.fixture_paths)
gs.write_subsets()
# Load the grid splitter index file ----------------------------------------------------------------------------
index_path = gs.create_full_path_from_template('index_file')
ifile = RequestDataset(uri=index_path).get()
ifile.load()
gidx = ifile[GridSplitterConstants.IndexFile.NAME_INDEX_VARIABLE].attrs
source_filename = ifile[gidx[GridSplitterConstants.IndexFile.NAME_SOURCE_VARIABLE]]
sv = source_filename.join_string_value()
destination_filename = ifile[gidx[GridSplitterConstants.IndexFile.NAME_DESTINATION_VARIABLE]]
dv = destination_filename.join_string_value()
# Create weight files for each subset --------------------------------------------------------------------------
for ii, sfn in enumerate(sv):
esp = os.path.join(self.current_dir_output, sfn)
edp = os.path.join(self.current_dir_output, dv[ii])
ewp = gs.create_full_path_from_template('wgt_template', index=ii + 1)
cmd = ['ESMF_RegridWeightGen', '-s', esp, '--src_type', 'UGRID', '--src_meshname',
VariableName.UGRID_HOST_VARIABLE, '-d', edp, '--dst_type', 'GRIDSPEC', '-w', ewp, '--method',
'conserve', '-r', '--no_log']
subprocess.check_call(cmd)
# Merge weight files -------------------------------------------------------------------------------------------
mwf = self.get_temporary_file_path('merged_weight_file.nc')
gs.create_merged_weight_file(mwf)
# Generate a global weight file using ESMF ---------------------------------------------------------------------
global_weights_filename = self.get_temporary_file_path('global_weights.nc')
cmd = ['ESMF_RegridWeightGen', '-s', ufile, '--src_type', 'UGRID', '-d', dst_path, '--dst_type',
'GRIDSPEC', '-w', global_weights_filename, '--method', 'conserve', '--weight-only', '--no_log',
'--src_meshname', VariableName.UGRID_HOST_VARIABLE]
subprocess.check_call(cmd)
# Test merged and global weight files are equivalent -----------------------------------------------------------
self.assertWeightFilesEquivalent(global_weights_filename, mwf)
def smm(index_path, wd=None, data_variables='auto'):
"""
Run a chunked sparse matrix multiplication.
:param str index_path: Path to chunked operation's index file.
:param str wd: Path to the directory containing the chunk files.
:param list data_variables: Optional list of data variables. Otherwise, auto-discovery is used.
"""
# Assume the current working directory
if wd is None:
wd = ''
# Turn on auto-discovery
if data_variables == 'auto':
v = None
else:
v = data_variables
# Get the attribute host
index_field = RequestDataset(index_path).get()
gs_index_v = index_field[GridChunkerConstants.IndexFile.NAME_INDEX_VARIABLE]
# Collect the source filenames
src_filenames = gs_index_v.attrs[GridChunkerConstants.IndexFile.NAME_SOURCE_VARIABLE]
src_filenames = index_field[src_filenames]
src_filenames = src_filenames.join_string_value()
# Collect the destination file names
dst_filenames = gs_index_v.attrs[GridChunkerConstants.IndexFile.NAME_DESTINATION_VARIABLE]
dst_filenames = index_field[dst_filenames]
dst_filenames = dst_filenames.join_string_value()
# Collect the weight filenames
wgt_filenames = gs_index_v.attrs[GridChunkerConstants.IndexFile.NAME_WEIGHTS_VARIABLE]
wgt_filenames = index_field[wgt_filenames]
wgt_filenames = wgt_filenames.join_string_value()
# Main loop for executing the SMM
for ii in range(src_filenames.size):
ocgis_lh(msg="Running SMM {} of {}".format(ii + 1, src_filenames.size), level=10, logger='regrid.base')
src_path = os.path.join(wd, src_filenames[ii])
src_field = RequestDataset(src_path, variable=v, decomp_type=DecompositionType.ESMF).create_field()
src_field.load()
dst_path = os.path.join(wd, dst_filenames[ii])
dst_field = RequestDataset(dst_path, decomp_type=DecompositionType.ESMF).create_field()
dst_field.load()
from ocgis.regrid.base import RegridOperation
from ESMF.api.constants import RegridMethod
# HACK: Note that weight filenames are stored with their full path.
# HACK: Always use a bilinear regridding method to allows for the identity sparse matrix in the case of
# equal grids.
regrid_options = {'split': False, 'weights_in': wgt_filenames[ii], 'regrid_method': RegridMethod.BILINEAR}
ro = RegridOperation(src_field, dst_field, regrid_options=regrid_options)
regridded = ro.execute()
regridded.write(dst_path)
def test_create_merged_weight_file(self):
path_src = self.get_temporary_file_path('src.nc')
path_dst = self.get_temporary_file_path('dst.nc')
src_grid = create_gridxy_global(resolution=30.0, wrapped=False, crs=Spherical())
dst_grid = create_gridxy_global(resolution=35.0, wrapped=False, crs=Spherical())
src_grid.write(path_src)
dst_grid.write(path_dst)
# Split source and destination grids ---------------------------------------------------------------------------
gs = GridSplitter(src_grid, dst_grid, (2, 2), check_contains=False, allow_masked=True, paths=self.fixture_paths)
gs.write_subsets()
# Load the grid splitter index file ----------------------------------------------------------------------------
index_filename = gs.create_full_path_from_template('index_file')
ifile = RequestDataset(uri=index_filename).get()
ifile.load()
gidx = ifile[GridSplitterConstants.IndexFile.NAME_INDEX_VARIABLE].attrs
source_filename = ifile[gidx[GridSplitterConstants.IndexFile.NAME_SOURCE_VARIABLE]]
sv = source_filename.join_string_value()
destination_filename = ifile[gidx[GridSplitterConstants.IndexFile.NAME_DESTINATION_VARIABLE]]
dv = destination_filename.join_string_value()
# Create weight files for each subset --------------------------------------------------------------------------
for ii, sfn in enumerate(sv):
esp = os.path.join(self.current_dir_output, sfn)
edp = os.path.join(self.current_dir_output, dv[ii])
ewp = gs.create_full_path_from_template('wgt_template', index=ii + 1)
cmd = ['ESMF_RegridWeightGen', '-s', esp, '--src_type', 'GRIDSPEC', '-d', edp, '--dst_type',
'GRIDSPEC', '-w', ewp, '--method', 'conserve', '-r', '--no_log']
subprocess.check_call(cmd)
# Merge weight files -------------------------------------------------------------------------------------------
merged_weight_filename = self.get_temporary_file_path('merged_weights.nc')
gs.create_merged_weight_file(merged_weight_filename)
# Generate a global weight file using ESMF ---------------------------------------------------------------------
global_weights_filename = self.get_temporary_file_path('global_weights.nc')
cmd = ['ESMF_RegridWeightGen', '-s', path_src, '--src_type', 'GRIDSPEC', '-d', path_dst, '--dst_type',
'GRIDSPEC', '-w', global_weights_filename, '--method', 'conserve', '--weight-only', '--no_log']
subprocess.check_call(cmd)
# Test merged and global weight files are equivalent -----------------------------------------------------------
self.assertWeightFilesEquivalent(global_weights_filename, merged_weight_filename)
def smm(index_path, wd=None, data_variables='auto'):
"""
Run a chunked sparse matrix multiplication.
:param str index_path: Path to chunked operation's index file.
:param str wd: Path to the directory containing the chunk files.
:param list data_variables: Optional list of data variables. Otherwise, auto-discovery is used.
"""
# Assume the current working directory
if wd is None:
wd = ''
# Turn on auto-discovery
if data_variables == 'auto':
v = None
else:
v = data_variables
# Get the attribute host
index_field = RequestDataset(index_path).get()
gs_index_v = index_field[
GridChunkerConstants.IndexFile.NAME_INDEX_VARIABLE]
# Collect the source filenames
src_filenames = gs_index_v.attrs[
GridChunkerConstants.IndexFile.NAME_SOURCE_VARIABLE]
src_filenames = index_field[src_filenames]
src_filenames = src_filenames.join_string_value()
# Collect the destination file names
dst_filenames = gs_index_v.attrs[
GridChunkerConstants.IndexFile.NAME_DESTINATION_VARIABLE]
dst_filenames = index_field[dst_filenames]
dst_filenames = dst_filenames.join_string_value()
# Collect the weight filenames
wgt_filenames = gs_index_v.attrs[
GridChunkerConstants.IndexFile.NAME_WEIGHTS_VARIABLE]
wgt_filenames = index_field[wgt_filenames]
wgt_filenames = wgt_filenames.join_string_value()
# Main loop for executing the SMM
for ii in range(src_filenames.size):
ocgis_lh(msg="Running SMM {} of {}".format(ii + 1, src_filenames.size),
level=10,
logger='regrid.base')
src_path = os.path.join(wd, src_filenames[ii])
src_field = RequestDataset(
src_path, variable=v,
decomp_type=DecompositionType.ESMF).create_field()
src_field.load()
dst_path = os.path.join(wd, dst_filenames[ii])
dst_field = RequestDataset(
dst_path, decomp_type=DecompositionType.ESMF).create_field()
dst_field.load()
from ocgis.regrid.base import RegridOperation
from ESMF.api.constants import RegridMethod
# HACK: Note that weight filenames are stored with their full path.
# HACK: Always use a bilinear regridding method to allows for the identity sparse matrix in the case of
# equal grids.
regrid_options = {
'split': False,
'weights_in': wgt_filenames[ii],
'regrid_method': RegridMethod.BILINEAR
}
ro = RegridOperation(src_field,
dst_field,
regrid_options=regrid_options)
regridded = ro.execute()
regridded.write(dst_path)
