def Prepare_Meteorology_Fulldistributed(workspace, wbd, OUTPUT, input_dir,
info, hydrobloks_info):
#Assign other variables
mask = OUTPUT['mask']
mapping_info = {}
#Calculate the fine to coarse scale mapping
for data_var in wbd['files_meteorology']:
#Define the variable name
var = data_var #data_var.split('_')[1]
mapping_info[var] = {}
#Read in the coarse and fine mapping
file_coarse = '%s/%s_ea_coarse.tif' % (workspace, var)
file_fine = '%s/%s_ea_fine.tif' % (workspace, var)
mask_coarse = gdal_tools.read_raster(file_coarse)
mask_fine = gdal_tools.read_raster(file_fine)
nlat = mask_coarse.shape[0]
nlon = mask_coarse.shape[1]
print data_var, "Creating the i and j mapping"
#Create maps of i and j for mapping
coords = {
'i': np.zeros(mask_fine.shape, dtype=np.int),
'j': np.zeros(mask_fine.shape, dtype=np.int)
}
for value in np.unique(mask_coarse):
if value < 0: continue
idx_fine = mask_fine == value
idx_coarse = np.where(mask_coarse == value)
coords['i'][idx_fine] = nlat - idx_coarse[0] - 1 #Careful
coords['j'][idx_fine] = idx_coarse[1]
coords['i'] = coords['i'][mask]
coords['j'] = coords['j'][mask]
print data_var, "Extracting all the data"
#Extract all of the data for that variable
idate = info['time_info']['startdate']
fdate = info['time_info']['enddate']
nt = 24 * ((fdate - idate).days + 1)
var = data_var #data_var.split('_')[1]
date = idate
file = wbd['files_meteorology'][data_var]
fp = nc.Dataset(file)
#Determine the time steps to retrieve
dates = nc.num2date(fp.variables['t'][:],
units='hours since %02d-%02d-%02d 00:00:00' %
(idate.year, idate.month, idate.day))
mask_dates = (dates >= idate) & (dates <= fdate)
data = np.ma.getdata(fp.variables[var][mask_dates])
fp.close()
print data_var, "Assigning the data"
#Create the variable
grp = hydrobloks_info['input_fp'].groups['meteorology']
grp.createVariable(var, 'f4', ('time', 'hsu'))
#Place all the data per time step
tmp = np.zeros(np.sum(mask))
for itime in np.arange(data.shape[0]):
tmp[:] = data[itime, coords['i'], coords['j']]
#Write the timestep
grp.variables[var][itime, :] = tmp[:]
return
def Prepare_Meteorology_Semidistributed(workspace, wbd, OUTPUT, input_dir,
info, hydrobloks_info):
#Define the mapping directory
mapping_info = {}
#Calculate the fine to coarse scale mapping
for data_var in wbd['files_meteorology']:
#Define the variable name
var = data_var #data_var.split('_')[1]
mapping_info[var] = {}
#Read in the coarse and fine mapping
file_coarse = '%s/%s_latlon_coarse.tif' % (workspace, data_var)
file_fine = '%s/%s_ea_fine.tif' % (workspace, data_var)
mask_coarse = gdal_tools.read_raster(file_coarse)
mask_fine = gdal_tools.read_raster(file_fine)
nlat = mask_coarse.shape[0]
nlon = mask_coarse.shape[1]
#Compute the mapping for each hsu
for hsu in np.arange(hydrobloks_info['nhru']):
idx = OUTPUT['hsu_map'] == hsu
#print "Catch:",hydrobloks_info['icatch'], "HRU: ", mask_fine[idx].astype(np.int)
icells = np.unique(
mask_fine[idx][mask_fine[idx] != -9999.0].astype(
np.int)) # Add != -9999 for unique and bicount - Noemi
counts = np.bincount(
mask_fine[idx][mask_fine[idx] != -9999.0].astype(np.int))
coords, pcts = [], []
for icell in icells:
ilat = int(np.floor(icell / mask_coarse.shape[1]))
jlat = icell - ilat * mask_coarse.shape[1]
#ilat = int(mask_coarse.shape[0] - ilat - 1) #CAREFUL
pct = float(counts[icell]) / float(np.sum(counts))
coords.append([ilat, jlat])
pcts.append(pct)
pcts = np.array(pcts)
coords = list(np.array(coords).T)
mapping_info[var][hsu] = {'pcts': pcts, 'coords': coords}
#Iterate through variable creating forcing product per HSU
idate = info['time_info']['startdate']
fdate = info['time_info']['enddate']
dt = info['time_info']['dt']
nt = int(3600 * 24 / dt) * ((fdate - idate).days + 1)
#Create structured array
meteorology = {}
for data_var in wbd['files_meteorology']:
meteorology[data_var] = np.zeros((nt, hydrobloks_info['nhru']))
#Load data into structured array
for data_var in wbd['files_meteorology']:
var = data_var #data_var.split('_')[1]
date = idate
file = wbd['files_meteorology'][data_var]
fp = nc.Dataset(file)
#fp = h5py.File(file)
#Determine the time steps to retrieve
fidate = ' '.join(fp.variables['t'].units.split(' ')[2::])
#dates = nc.num2date(fp.variables['t'][:],units='hours since %02d-%02d-%02d 00:00:00' % (idate.year,idate.month,idate.day))
dates = nc.num2date(fp.variables['t'][:],
units='hours since %s' % fidate)
#dates = nc.num2date(fp['t'][:],units='hours since %02d-%02d-%02d 00:00:00' % (idate.year,idate.month,idate.day))
mask_dates = (dates >= idate) & (dates <= fdate)
data = np.ma.getdata(fp.variables[var][mask_dates, :, :])
fp.close()
#Assing to hsus
for hsu in mapping_info[var]:
#print data_var,data, data.shape, hsu,mapping_info[var][hsu]['pcts'],mapping_info[var][hsu]['coords'],
pcts = mapping_info[var][hsu]['pcts']
coords = mapping_info[var][hsu]['coords']
coords[0][coords[0] >= data.shape[1]] = data.shape[1] - 1
coords[1][coords[1] >= data.shape[2]] = data.shape[2] - 1
tmp = data[:, coords[0], coords[1]]
#m1 = tmp < -999
#m2 = tmp > -999
tmp = pcts * tmp
meteorology[data_var][:, hsu] = np.sum(tmp, axis=1)
#print data_var,np.unique(meteorology[data_var][:,:])
#Write the meteorology to the netcdf file (single chunk for now...)
grp = hydrobloks_info['input_fp'].groups['meteorology']
grp.createVariable(var, 'f4', ('time', 'hsu'))
grp.variables[data_var][:] = meteorology[data_var][:]
#Add time information
dates = []
date = idate
while date <= fdate:
dates.append(date)
date = date + datetime.timedelta(seconds=dt)
dates = np.array(dates)
var = grp.createVariable('time', 'f8', ('time', ))
var.units = 'hours since %4d-01-01' % idate.year
var.calendar = 'standard'
dates = nc.date2num(dates, units=var.units, calendar=var.calendar)
var[:] = dates[:]
return
def Prepare_Model_Input_Data(hydrobloks_info):
#Prepare the info dictionary
info = {}
#Define the start/end dates
info['time_info'] = {}
info['time_info']['startdate'] = hydrobloks_info['idate']
info['time_info']['enddate'] = hydrobloks_info['fdate']
info['time_info']['dt'] = hydrobloks_info['dt']
#Define the workspace
workspace = hydrobloks_info['workspace']
#Define the model input data directory
input_dir = workspace #'%s/input' % workspace
#Read in the metadata
#file = '%s/workspace_info.pck' % workspace
#wbd = pickle.load(open(file))
#Create the dictionary to hold all of the data
output = {}
#Create the Latin Hypercube (Clustering)
nhru = hydrobloks_info['nhru']
#ncores = hydrobloks_info['ncores']
icatch = hydrobloks_info['icatch']
#Get metadata
md = gdal_tools.retrieve_metadata('%s/mask_latlon.tif' % workspace)
#Prepare the input file
wbd = {}
wbd['bbox'] = {
'minlat': md['miny'],
'maxlat': md['maxy'],
'minlon': md['minx'],
'maxlon': md['maxx'],
'res': abs(md['resx'])
}
wbd['files'] = {
'WLTSMC': '%s/theta1500_ea.tif' % workspace,
'TEXTURE_CLASS': '%s/texture_class_ea.tif' % workspace,
'cslope': '%s/cslope_ea.tif' % workspace,
'MAXSMC': '%s/thetas_ea.tif' % workspace,
'BB': '%s/bb_ea.tif' % workspace,
'DRYSMC': '%s/thetar_ea.tif' % workspace,
'fdir': '%s/fdir_ea.tif' % workspace,
'QTZ': '%s/qtz_ea.tif' % workspace,
'SATDW': '%s/dsat_ea.tif' % workspace,
'REFSMC': '%s/theta33_ea.tif' % workspace,
'mask': '%s/mask_ea.tif' % workspace,
'SATPSI': '%s/psisat_ea.tif' % workspace,
'lc': '%s/lc_ea.tif' % workspace,
'carea': '%s/carea_ea.tif' % workspace,
'ti': '%s/ti_ea.tif' % workspace,
'ndvi': '%s/ndvi_ea.tif' % workspace,
'F11': '%s/f11_ea.tif' % workspace,
'SATDK': '%s/ksat_ea.tif' % workspace,
'dem': '%s/dem_ea.tif' % workspace,
'demns': '%s/demns_ea.tif' % workspace,
'sand': '%s/sand_ea.tif' % workspace,
'clay': '%s/clay_ea.tif' % workspace,
'silt': '%s/silt_ea.tif' % workspace,
'om': '%s/om_ea.tif' % workspace,
'bare30': '%s/bare30_ea.tif' % workspace,
'water30': '%s/water30_ea.tif' % workspace,
'tree30': '%s/tree30_ea.tif' % workspace,
}
wbd['files_meteorology'] = {
'lwdown': '%s/lwdown.nc' % workspace,
'swdown': '%s/swdown.nc' % workspace,
'tair': '%s/tair.nc' % workspace,
'precip': '%s/precip.nc' % workspace,
'psurf': '%s/psurf.nc' % workspace,
'wind': '%s/wind.nc' % workspace,
'spfh': '%s/spfh.nc' % workspace,
}
#Create the clusters and their connections
output = Create_Clusters_And_Connections(workspace, wbd, output, input_dir,
nhru, info, hydrobloks_info)
#Extract the meteorological forcing
print "Preparing the meteorology"
if hydrobloks_info['model_type'] == 'semi':
Prepare_Meteorology_Semidistributed(workspace, wbd, output, input_dir,
info, hydrobloks_info)
elif hydrobloks_info['model_type'] == 'full':
Prepare_Meteorology_Fulldistributed(workspace, wbd, output, input_dir,
info, hydrobloks_info)
#Write out the files to the netcdf file
fp = hydrobloks_info['input_fp']
data = output
#Write out the metadata
grp = fp.createGroup('metadata')
grp.latitude = (wbd['bbox']['minlat'] + wbd['bbox']['maxlat']) / 2
lon = (wbd['bbox']['minlon'] + wbd['bbox']['maxlon']) / 2
if lon < 0: lon += 360
grp.longitude = lon
metadata = gdal_tools.retrieve_metadata(wbd['files']['mask'])
grp.dx = metadata['resx']
#grp.longitude = (360.0+(wbd['bbox']['minlon'] + wbd['bbox']['maxlon'])/2)
#Write the HRU mapping
#CONUS conus_albers metadata
metadata['nodata'] = -9999.0
#Save the conus_albers metadata
grp = fp.createGroup('conus_albers_mapping')
grp.createDimension('nx', metadata['nx'])
grp.createDimension('ny', metadata['ny'])
hmca = grp.createVariable('hmca', 'f4', ('ny', 'nx'))
hmca.gt = metadata['gt']
hmca.projection = metadata['projection']
hmca.description = 'HSU mapping (conus albers)'
hmca.nodata = metadata['nodata']
#Save the conus albers mapping
hsu_map = np.copy(output['hsu_map'])
hsu_map[np.isnan(hsu_map) == 1] = metadata['nodata']
hmca[:] = hsu_map
#Write out the mapping
file_ca = '%s/hsu_mapping_ea.tif' % workspace
gdal_tools.write_raster(file_ca, metadata, hsu_map)
#Map the mapping to regular lat/lon
file_ll = '%s/hsu_mapping_latlon.tif' % workspace
os.system('rm -f %s' % file_ll)
res = wbd['bbox']['res']
minlat = wbd['bbox']['minlat']
minlon = wbd['bbox']['minlon']
maxlat = wbd['bbox']['maxlat']
maxlon = wbd['bbox']['maxlon']
log = '%s/log.txt' % workspace
os.system(
'gdalwarp -tr %.16f %.16f -dstnodata %.16f -t_srs \'+proj=longlat +ellps=WGS84 +datum=WGS84 +no_defs \' -te %.16f %.16f %.16f %.16f %s %s >> %s 2>&1'
% (res, res, metadata['nodata'], minlon, minlat, maxlon, maxlat,
file_ca, file_ll, log))
#Write a map for the catchment id
file_icatch = '%s/icatch_latlon.tif' % workspace
metadata = gdal_tools.retrieve_metadata(file_ll)
metadata['nodata'] = -9999.0
tmp = gdal_tools.read_raster(file_ll)
tmp[tmp >= 0] = hydrobloks_info['icatch']
gdal_tools.write_raster(file_icatch, metadata, tmp)
#Retrieve the lat/lon metadata
#metadata = gdal_tools.retrieve_metadata(file_ll)
#metadata['nodata'] = -9999.0
#Save the lat/lon metadata
#grp = fp.createGroup('latlon_mapping')
#grp.createDimension('nlon',metadata['nx'])
#grp.createDimension('nlat',metadata['ny'])
#hmll = grp.createVariable('hmll','f4',('nlat','nlon'))
#hmll.gt = metadata['gt']
#hmll.projection = metadata['projection']
#hmll.description = 'HSU mapping (regular lat/lon)'
#hmll.nodata = metadata['nodata']
#Save the lat/lon mapping
#hsu_map = np.copy(gdal_tools.read_raster(file_ll))
#hsu_map[np.isnan(hsu_map) == 1] = metadata['nodata']
#hmll[:] = hsu_map
#Write the flow matrix
flow_matrix = output['flow_matrix']
nconnections = flow_matrix.data.size
grp = fp.createGroup('flow_matrix')
grp.createDimension('connections_columns', flow_matrix.indices.size)
grp.createDimension('connections_rows', flow_matrix.indptr.size)
grp.createVariable('data', 'f4', ('connections_columns', ))
grp.createVariable('indices', 'f4', ('connections_columns', ))
grp.createVariable('indptr', 'f4', ('connections_rows', ))
grp.variables['data'][:] = flow_matrix.data
grp.variables['indices'][:] = flow_matrix.indices
grp.variables['indptr'][:] = flow_matrix.indptr
#Write the connection matrices
#width
wmatrix = output['cmatrix']['width']
nconnections = wmatrix.data.size
grp = fp.createGroup('wmatrix')
grp.createDimension('connections_columns', wmatrix.indices.size)
grp.createDimension('connections_rows', wmatrix.indptr.size)
grp.createVariable('data', 'f4', ('connections_columns', ))
grp.createVariable('indices', 'f4', ('connections_columns', ))
grp.createVariable('indptr', 'f4', ('connections_rows', ))
grp.variables['data'][:] = wmatrix.data
grp.variables['indices'][:] = wmatrix.indices
grp.variables['indptr'][:] = wmatrix.indptr
#Write the outlet information
outlet = output['outlet']
grp = fp.createGroup('outlet')
full = grp.createGroup('full')
full.createDimension('cell', outlet['full']['hru_org'].size)
full.createVariable('i', 'i4', ('cell', ))
full.createVariable('j', 'i4', ('cell', ))
full.createVariable('hru_org', 'i4', ('cell', ))
full.createVariable('hru_dst', 'i4', ('cell', ))
full.createVariable('d8', 'i4', ('cell', ))
full.variables['i'][:] = outlet['full']['i']
full.variables['j'][:] = outlet['full']['j']
full.variables['hru_org'][:] = outlet['full']['hru_org']
full.variables['hru_dst'][:] = outlet['full']['hru_dst']
full.variables['d8'][:] = outlet['full']['d8']
summary = grp.createGroup('summary')
summary.createDimension('hru', outlet['summary']['hru_org'].size)
summary.createVariable('hru_org', 'i4', ('hru', ))
summary.createVariable('hru_dst', 'i4', ('hru', ))
summary.createVariable('counts', 'i4', ('hru', ))
summary.variables['hru_org'][:] = outlet['summary']['hru_org']
summary.variables['hru_dst'][:] = outlet['summary']['hru_dst']
summary.variables['counts'][:] = outlet['summary']['counts']
#outlet = {'full':{'i':outlet_icoord,'j':outlet_jcoord,'hru_org':outlet_hru_org,'hru_dst':outlet_hru_dst,'d8':outlet_d8},
# 'summary':{'hru_org':outlet_hru_org_summary,'hru_dst':outlet_hru_dst_summary,'counts':counts}}
#Write the model parameters
grp = fp.createGroup('parameters')
vars = [
'slope', 'area_pct', 'land_cover', 'channel', 'dem',
'soil_texture_class', 'ti', 'carea', 'area', 'BB', 'F11', 'SATPSI',
'SATDW', 'QTZ', 'WLTSMC', 'MAXSMC', 'DRYSMC', 'REFSMC', 'SATDK',
'mannings', 'm', 'psoil', 'pksat', 'sdmax'
]
for var in vars:
grp.createVariable(var, 'f4', ('hsu', ))
grp.variables[var][:] = data['hsu'][var]
#Write other metadata
#grp = fp.createGroup('metadata')
#grp.outlet_hsu = data['outlet']['hsu']
#Remove info from output
del output['hsu']
#Add in the catchment info
output['wbd'] = wbd
#Close the file
fp.close()
return output
def Create_and_Curate_Covariates(wbd):
covariates = {}
#Read in and curate all the covariates
for file in wbd['files']:
covariates[file] = gdal_tools.read_raster(wbd['files'][file])
if file == 'cslope':
mask = covariates[file] == 0.0
covariates[file][mask] = 0.000001
#Create lat/lon grids
lats = np.linspace(wbd['bbox']['minlat'] + wbd['bbox']['res'] / 2,
wbd['bbox']['maxlat'] - wbd['bbox']['res'] / 2,
covariates['ti'].shape[0])
lons = np.linspace(wbd['bbox']['minlon'] + wbd['bbox']['res'] / 2,
wbd['bbox']['maxlon'] - wbd['bbox']['res'] / 2,
covariates['ti'].shape[1])
#Constrain the lat/lon grid by the effective resolution (4km...)
nres = int(np.floor(2000.0 / 30.0))
for ilat in xrange(0, lats.size, nres):
lats[ilat:ilat + nres] = np.mean(lats[ilat:ilat + nres])
for ilon in xrange(0, lons.size, nres):
lons[ilon:ilon + nres] = np.mean(lons[ilon:ilon + nres])
#Need to fix so that it doesn't suck up all the clustering:
lats, lons = np.meshgrid(lats, lons)
covariates['lats'] = lats.T
covariates['lons'] = lons.T
#Add sti
sti = covariates['ti'] - np.log(0.1 * covariates['SATDK']) + np.log(
np.mean(0.1 * covariates['SATDK']))
covariates['sti'] = sti
#Define the mask
mask = np.copy(covariates['mask'])
mask[mask >= 0] = 1
mask[mask < 0] = 0
mask = mask.astype(np.bool)
#Set all nans to the mean
for var in covariates:
mask1 = (np.isinf(covariates[var]) == 0) & (np.isnan(covariates[var])
== 0)
mask0 = (np.isinf(covariates[var]) == 1) | (np.isnan(covariates[var])
== 1)
if var in ['fdir', 'nlcd', 'TEXTURE_CLASS', 'lc']:
if len(covariates[var][mask1]) < 1:
print "Full of Nan's Error: 1", var, covariates[var], wbd[
'files'][var] # Noemi insert
covariates[var][
mask0] = -9999.0 #stats.mode(covariates[var][mask1])[0][0]
else:
covariates[var][mask0] = -9999.0 #np.mean(covariates[var][mask1])
#Set everything that is -9999 to the mean
for var in covariates:
if var in ['fdir', 'nlcd', 'TEXTURE_CLASS', 'lc']:
if len(covariates[var][covariates[var] != -9999.0]) < 1:
print "Full of Nan's Error: ", var, covariates[var], wbd[
'files'][var] # Noemi insert
covariates[var][covariates[var] == -9999.0] = stats.mode(
covariates[var][covariates[var] != -9999.0])[0][0]
else:
covariates[var][covariates[var] == -9999.0] = np.mean(
covariates[var][covariates[var] != -9999.0])
#Set everything outside of the mask to -9999
for var in covariates:
covariates[var][mask <= 0] = -9999.0
return (covariates, mask)
import matplotlib.pyplot as mpl
import geospatialtools.gdal_tools as geo
import h5py as h5
def mask(array):
return np.ma.masked_array(array,array==-9999)
def show(array):
p = mpl.imshow(array)
def shape(array):
temp = array.reshape(array.size)
trim = temp[temp != -9999]
return trim[:,np.newaxis]
sand = geo.read_raster('clusters/sand_mean_0_5.tif')
silt = geo.read_raster('clusters/silt_mean_0_5.tif')
clay = geo.read_raster('clusters/clay_mean_0_5.tif')
#sand,silt,clay = mask(sand),mask(silt),mask(clay)
x = sand.reshape(sand.size)
sand1 = np.copy(sand)
silt1 = np.copy(silt)
clay1 = np.copy(clay)
sand,silt,clay = shape(sand),shape(silt),shape(clay)
size = min(sand.shape[0],silt.shape[0],clay.shape[0])
soil = np.concatenate([sand[0:size],silt[0:size],clay[0:size]],axis=1)
model = sk.KMeans(n_clusters=10)
np.random.seed(0)
def Prepare_Meteorology_Semidistributed(workspace,wbd,OUTPUT,input_dir,info,hydrobloks_info):
#Define the mapping directory
mapping_dir = '%s/mapping' % workspace
mapping_info = {}
#Calculate the fine to coarse scale mapping
for data_var in wbd['files_meteorology']:
#Define the variable name
var = data_var#data_var.split('_')[1]
mapping_info[var] = {}
#Read in the coarse and fine mapping
file_coarse = '%s/%s_coarse.tif' % (mapping_dir,data_var)
file_fine = '%s/%s_fine.tif' % (mapping_dir,data_var)
mask_coarse = gdal_tools.read_raster(file_coarse)
mask_fine = gdal_tools.read_raster(file_fine)
nlat = mask_coarse.shape[0]
nlon = mask_coarse.shape[1]
#Compute the mapping for each hsu
for hsu in np.arange(hydrobloks_info['nclusters']):
idx = OUTPUT['hsu_map'] == hsu
icells = np.unique(mask_fine[idx].astype(np.int))
counts = np.bincount(mask_fine[idx].astype(np.int))
coords,pcts = [],[]
for icell in icells:
ilat = int(np.floor(icell/mask_coarse.shape[1]))
jlat = icell - ilat*mask_coarse.shape[1]
#ilat = int(mask_coarse.shape[0] - ilat - 1) #CAREFUL
pct = float(counts[icell])/float(np.sum(counts))
coords.append([ilat,jlat])
pcts.append(pct)
pcts = np.array(pcts)
coords = list(np.array(coords).T)
mapping_info[var][hsu] = {'pcts':pcts,'coords':coords}
#Iterate through variable creating forcing product per HSU
idate = info['time_info']['startdate']
fdate = info['time_info']['enddate']
nt = 24*((fdate - idate).days+1)
#Create structured array
meteorology = {}
for data_var in wbd['files_meteorology']:
meteorology[data_var] = np.zeros((nt,hydrobloks_info['nclusters']))
#Load data into structured array
for data_var in wbd['files_meteorology']:
var = data_var#data_var.split('_')[1]
date = idate
file = wbd['files_meteorology'][data_var]
fp = nc.Dataset(file)
#fp = h5py.File(file)
#Determine the time steps to retrieve
fidate = ' '.join(fp.variables['t'].units.split(' ')[2::])
#dates = nc.num2date(fp.variables['t'][:],units='hours since %02d-%02d-%02d 00:00:00' % (idate.year,idate.month,idate.day))
dates = nc.num2date(fp.variables['t'][:],units='hours since %s' % fidate)
#dates = nc.num2date(fp['t'][:],units='hours since %02d-%02d-%02d 00:00:00' % (idate.year,idate.month,idate.day))
mask_dates = (dates >= idate) & (dates <= fdate)
data = np.ma.getdata(fp.variables[var][mask_dates,:,:])
fp.close()
#Assing to hsus
for hsu in mapping_info[var]:
pcts = mapping_info[var][hsu]['pcts']
coords = mapping_info[var][hsu]['coords']
coords[0][coords[0] >= data.shape[1]] = data.shape[1] - 1
coords[1][coords[1] >= data.shape[2]] = data.shape[2] - 1
tmp = data[:,coords[0],coords[1]]
m1 = tmp < -999
m2 = tmp > -999
tmp = pcts*tmp
#Combine stage iv and nldas here
if data_var not in ['apcpsfc',]:tmp[m1] = np.mean(tmp[m2])
meteorology[data_var][:,hsu] = np.sum(tmp,axis=1)
#print data_var,np.unique(meteorology[data_var][:,:])
#Write the meteorology to the netcdf file (single chunk for now...)
grp = hydrobloks_info['input_fp'].groups['meteorology']
grp.createVariable(var,'f4',('time','hsu'))
grp.variables[data_var][:] = meteorology[data_var][:]
return
def Prepare_Meteorology_Fulldistributed(workspace,wbd,OUTPUT,input_dir,info,hydrobloks_info):
#Assign other variables
mask = OUTPUT['mask']
#Define the mapping directory
mapping_dir = '%s/mapping' % workspace
mapping_info = {}
#Calculate the fine to coarse scale mapping
for data_var in wbd['files_meteorology']:
#Define the variable name
var = data_var#data_var.split('_')[1]
mapping_info[var] = {}
#Read in the coarse and fine mapping
file_coarse = '%s/%s_coarse.tif' % (mapping_dir,data_var)
file_fine = '%s/%s_fine.tif' % (mapping_dir,data_var)
mask_coarse = gdal_tools.read_raster(file_coarse)
mask_fine = gdal_tools.read_raster(file_fine)
nlat = mask_coarse.shape[0]
nlon = mask_coarse.shape[1]
print data_var,"Creating the i and j mapping"
#Create maps of i and j for mapping
coords = {'i':np.zeros(mask_fine.shape,dtype=np.int),
'j':np.zeros(mask_fine.shape,dtype=np.int)}
for value in np.unique(mask_coarse):
if value < 0:continue
idx_fine = mask_fine == value
idx_coarse = np.where(mask_coarse == value)
coords['i'][idx_fine] = nlat - idx_coarse[0] - 1 #Careful
coords['j'][idx_fine] = idx_coarse[1]
coords['i'] = coords['i'][mask]
coords['j'] = coords['j'][mask]
print data_var,"Extracting all the data"
#Extract all of the data for that variable
idate = info['time_info']['startdate']
fdate = info['time_info']['enddate']
nt = 24*((fdate - idate).days+1)
var = data_var#data_var.split('_')[1]
date = idate
file = wbd['files_meteorology'][data_var]
fp = nc.Dataset(file)
#Determine the time steps to retrieve
dates = nc.num2date(fp.variables['t'][:],units='hours since %02d-%02d-%02d 00:00:00' % (idate.year,idate.month,idate.day))
mask_dates = (dates >= idate) & (dates <= fdate)
data = np.ma.getdata(fp.variables[var][mask_dates])
fp.close()
print data_var,"Assigning the data"
#Create the variable
grp = hydrobloks_info['input_fp'].groups['meteorology']
grp.createVariable(var,'f4',('time','hsu'))
#Place all the data per time step
tmp = np.zeros(np.sum(mask))
for itime in np.arange(data.shape[0]):
tmp[:] = data[itime,coords['i'],coords['j']]
#Write the timestep
grp.variables[var][itime,:] = tmp[:]
return
def Create_and_Curate_Covariates(wbd):
covariates = {}
#Read in and curate all the covariates
root = wbd['files']['MAXSMC'][0:-11]
wbd['files']['sand'] = '%s/dssurgo/sandtotal_r.tif' % root
wbd['files']['clay'] = '%s/dssurgo/claytotal_r.tif' % root
for file in wbd['files']:
covariates[file] = gdal_tools.read_raster(wbd['files'][file])
if file == 'cslope':
mask = covariates[file] == 0.0
covariates[file][mask] = 0.000001
#Map the NLCD to IGBP
NLCD2NOAH = {11:17,12:15,21:10,22:10,23:10,24:13,31:16,41:4,42:1,43:5,51:6,52:6,71:10,72:10,73:19,74:19,81:10,82:12,90:11,95:11}
tmp = np.copy(covariates['nlcd'])
for lc in np.unique(covariates['nlcd']):
if lc < 0:continue
tmp[covariates['nlcd'] == lc] = NLCD2NOAH[lc]
covariates['nlcd'][:] = tmp[:]
#Curate the NDVI to match NLCD
for lc in np.unique(covariates['nlcd']):
masklc = covariates['nlcd'] == lc
covariates['ndvi'][masklc] = np.mean(covariates['ndvi'][masklc])
#Create lat/lon grids
lats = np.linspace(wbd['bbox']['minlat']+wbd['bbox']['res']/2,wbd['bbox']['maxlat']-wbd['bbox']['res']/2,covariates['ti'].shape[0])
lons = np.linspace(wbd['bbox']['minlon']+wbd['bbox']['res']/2,wbd['bbox']['maxlon']-wbd['bbox']['res']/2,covariates['ti'].shape[1])
#Constrain the lat/lon grid by the effective resolution (4km...)
nres = int(np.floor(2000.0/30.0))
for ilat in xrange(0,lats.size,nres):
lats[ilat:ilat+nres] = np.mean(lats[ilat:ilat+nres])
for ilon in xrange(0,lons.size,nres):
lons[ilon:ilon+nres] = np.mean(lons[ilon:ilon+nres])
#Need to fix so that it doesn't suck up all the clustering:
lats, lons = np.meshgrid(lats, lons)
covariates['lats'] = lats.T
covariates['lons'] = lons.T
#Add sti
sti = covariates['ti'] - np.log(0.1*covariates['SATDK']) + np.log(np.mean(0.1*covariates['SATDK']))
covariates['sti'] = sti
#Define the mask
mask = np.copy(covariates['mask'])
mask[mask > 0] = 1
mask[mask < 0] = 0
mask = mask.astype(np.bool)
#Set all nans to the mean
for var in covariates:
mask1 = (np.isinf(covariates[var]) == 0) & (np.isnan(covariates[var]) == 0)
mask0 = (np.isinf(covariates[var]) == 1) | (np.isnan(covariates[var]) == 1)
if var in ['fdir','nlcd']:
covariates[var][mask0] = stats.mode(covariates[var][mask1])[0][0]
else:
covariates[var][mask0] = np.mean(covariates[var][mask1])
#Set everything that is -9999 to the mean
for var in covariates:
if var in ['fdir','nlcd','TEXTURE_CLASS']:
covariates[var][covariates[var] == -9999.0] = stats.mode(covariates[var][covariates[var] != -9999.0])[0][0]
else:
covariates[var][covariates[var] == -9999.0] = np.mean(covariates[var][covariates[var] != -9999.0])
#Set everything outside of the mask to -9999
for var in covariates:
covariates[var][mask <= 0] = -9999.0
return (covariates,mask)
def Prepare_Model_Input_Data(hydrobloks_info):
#Prepare the info dictionary
info = {}
#Define the start/end dates
info['time_info'] = {}
info['time_info']['startdate'] = hydrobloks_info['idate']
info['time_info']['enddate'] = hydrobloks_info['fdate']
#Define the workspace
workspace = hydrobloks_info['workspace']
#Define the model input data directory
input_dir = workspace#'%s/input' % workspace
#Read in the metadata
file = '%s/workspace_info.pck' % workspace
wbd = pickle.load(open(file))
#Create the dictionary to hold all of the data
output = {}
#Create the Latin Hypercube (Clustering)
nclusters = hydrobloks_info['nclusters']
ncores = hydrobloks_info['ncores']
icatch = hydrobloks_info['icatch']
#Prepare the input file
wbd['files'] = {
'WLTSMC':'%s/WLTSMC.tif' % workspace,
'TEXTURE_CLASS':'%s/TEXTURE_CLASS.tif' % workspace,
'cslope':'%s/cslope.tif' % workspace,
'MAXSMC':'%s/MAXSMC.tif' % workspace,
'BB':'%s/BB.tif' % workspace,
'DRYSMC':'%s/DRYSMC.tif' % workspace,
'fdir':'%s/fdir.tif' % workspace,
'QTZ':'%s/QTZ.tif' % workspace,
'SATDW':'%s/SATDW.tif' % workspace,
'REFSMC':'%s/REFSMC.tif' % workspace,
'mask':'%s/mask.tif' % workspace,
'channels':'%s/channels.tif' % workspace,
#'SATDW':'%s/SATDW.tif' % workspace,
#'REFSMC':'%s/REFSMC.tif' % workspace,
'SATPSI':'%s/SATPSI.tif' % workspace,
'nlcd':'%s/nlcd.tif' % workspace,
'carea':'%s/carea.tif' % workspace,
'ti':'%s/ti.tif' % workspace,
'ndvi':'%s/ndvi.tif' % workspace,
'F11':'%s/F11.tif' % workspace,
'SATDK':'%s/SATDK.tif' % workspace,
'dem':'%s/dem.tif' % workspace,
'demns':'%s/demns.tif' % workspace,
'strahler':'%s/strahler.tif' % workspace,
#'qbase':'%s/qbase.tif' % workspace
}
wbd['files_meteorology'] = {
'dlwrf':'%s/nldas/dlwrf/dlwrf.nc' % workspace,
'dswrf':'%s/nldas/dswrf/dswrf.nc' % workspace,
'tair':'%s/nldas/tair/tair.nc' % workspace,
'prec':'%s/nldas/prec/prec.nc' % workspace,
'pres':'%s/nldas/pres/pres.nc' % workspace,
'wind':'%s/nldas/wind/wind.nc' % workspace,
'rh':'%s/nldas/rh/rh.nc' % workspace,
'apcpsfc':'%s/stageiv/apcpsfc/apcpsfc.nc' % workspace,
}
#Create the clusters and their connections
output = Create_Clusters_And_Connections(workspace,wbd,output,input_dir,nclusters,ncores,info,hydrobloks_info)
#Extract the meteorological forcing
print "Preparing the meteorology"
if hydrobloks_info['model_type'] == 'semi':
Prepare_Meteorology_Semidistributed(workspace,wbd,output,input_dir,info,hydrobloks_info)
elif hydrobloks_info['model_type'] == 'full':
Prepare_Meteorology_Fulldistributed(workspace,wbd,output,input_dir,info,hydrobloks_info)
#Write out the files to the netcdf file
fp = hydrobloks_info['input_fp']
data = output
#Write out the metadata
grp = fp.createGroup('metadata')
grp.latitude = (wbd['bbox']['minlat'] + wbd['bbox']['maxlat'])/2
grp.longitude = (360.0+(wbd['bbox']['minlon'] + wbd['bbox']['maxlon'])/2)
#Write the HRU mapping
#CONUS conus_albers metadata
metadata = gdal_tools.retrieve_metadata(wbd['files']['mask'])
metadata['nodata'] = -9999.0
#Save the conus_albers metadata
grp = fp.createGroup('conus_albers_mapping')
grp.createDimension('nx',metadata['nx'])
grp.createDimension('ny',metadata['ny'])
hmca = grp.createVariable('hmca','f4',('ny','nx'))
hmca.gt = metadata['gt']
hmca.projection = metadata['projection']
hmca.description = 'HSU mapping (conus albers)'
hmca.nodata = metadata['nodata']
#Save the conus albers mapping
hsu_map = np.copy(output['hsu_map'])
hsu_map[np.isnan(hsu_map) == 1] = metadata['nodata']
hmca[:] = hsu_map
if hydrobloks_info['create_mask_flag'] == True:
#Write out the mapping
file_ca = '%s/hsu_mapping_conus_albers.tif' % workspace
gdal_tools.write_raster(file_ca,metadata,hsu_map)
#Map the mapping to regular lat/lon
file_ll = '%s/hsu_mapping_latlon.tif' % workspace
os.system('rm -f %s' % file_ll)
res = wbd['bbox']['res']
minlat = wbd['bbox']['minlat']
minlon = wbd['bbox']['minlon']
maxlat = wbd['bbox']['maxlat']
maxlon = wbd['bbox']['maxlon']
log = '%s/log.txt' % workspace
os.system('gdalwarp -tr %.16f %.16f -dstnodata %.16f -t_srs EPSG:4326 -s_srs EPSG:102039 -te %.16f %.16f %.16f %.16f %s %s >> %s 2>&1' % (res,res,metadata['nodata'],minlon,minlat,maxlon,maxlat,file_ca,file_ll,log))
#Write a map for the catchment id
file_icatch = '%s/icatch_latlon.tif' % workspace
metadata = gdal_tools.retrieve_metadata(file_ll)
metadata['nodata'] = -9999.0
tmp = gdal_tools.read_raster(file_ll)
tmp[tmp >= 0] = hydrobloks_info['icatch']
gdal_tools.write_raster(file_icatch,metadata,tmp)
#Add the lat/lon mapping
#Retrieve the lat/lon metadata
metadata = gdal_tools.retrieve_metadata(file_ll)
metadata['nodata'] = -9999.0
#Save the lat/lon metadata
grp = fp.createGroup('latlon_mapping')
grp.createDimension('nlon',metadata['nx'])
grp.createDimension('nlat',metadata['ny'])
hmll = grp.createVariable('hmll','f4',('nlat','nlon'))
hmll.gt = metadata['gt']
hmll.projection = metadata['projection']
hmll.description = 'HSU mapping (regular lat/lon)'
hmll.nodata = metadata['nodata']
#Save the lat/lon mapping
hsu_map = np.copy(gdal_tools.read_raster(file_ll))
hsu_map[np.isnan(hsu_map) == 1] = metadata['nodata']
hmll[:] = hsu_map
#Write the flow matrix
flow_matrix = output['flow_matrix']
nconnections = flow_matrix.data.size
grp = fp.createGroup('flow_matrix')
grp.createDimension('connections_columns',flow_matrix.indices.size)
grp.createDimension('connections_rows',flow_matrix.indptr.size)
grp.createVariable('data','f4',('connections_columns',))
grp.createVariable('indices','f4',('connections_columns',))
grp.createVariable('indptr','f4',('connections_rows',))
grp.variables['data'][:] = flow_matrix.data
grp.variables['indices'][:] = flow_matrix.indices
grp.variables['indptr'][:] = flow_matrix.indptr
#Write the outlet information
outlet = output['outlet']
grp = fp.createGroup('outlet')
full = grp.createGroup('full')
full.createDimension('cell',outlet['full']['hru_org'].size)
full.createVariable('i','i4',('cell',))
full.createVariable('j','i4',('cell',))
full.createVariable('hru_org','i4',('cell',))
full.createVariable('hru_dst','i4',('cell',))
full.createVariable('d8','i4',('cell',))
full.variables['i'][:] = outlet['full']['i']
full.variables['j'][:] = outlet['full']['j']
full.variables['hru_org'][:] = outlet['full']['hru_org']
full.variables['hru_dst'][:] = outlet['full']['hru_dst']
full.variables['d8'][:] = outlet['full']['d8']
summary = grp.createGroup('summary')
summary.createDimension('hru',outlet['summary']['hru_org'].size)
summary.createVariable('hru_org','i4',('hru',))
summary.createVariable('hru_dst','i4',('hru',))
summary.createVariable('counts','i4',('hru',))
summary.variables['hru_org'][:] = outlet['summary']['hru_org']
summary.variables['hru_dst'][:] = outlet['summary']['hru_dst']
summary.variables['counts'][:] = outlet['summary']['counts']
#outlet = {'full':{'i':outlet_icoord,'j':outlet_jcoord,'hru_org':outlet_hru_org,'hru_dst':outlet_hru_dst,'d8':outlet_d8},
# 'summary':{'hru_org':outlet_hru_org_summary,'hru_dst':outlet_hru_dst_summary,'counts':counts}}
#Write the model parameters
grp = fp.createGroup('parameters')
vars = ['slope','area_pct','land_cover','channel',
'dem','soil_texture_class','ti','carea','area',
'BB','F11','SATPSI','SATDW','QTZ',
'WLTSMC','MAXSMC','DRYSMC','REFSMC','SATDK',
'mannings','m','psoil','pksat','sdmax']
for var in vars:
grp.createVariable(var,'f4',('hsu',))
grp.variables[var][:] = data['hsu'][var]
#Write other metadata
#grp = fp.createGroup('metadata')
#grp.outlet_hsu = data['outlet']['hsu']
#Remove info from output
del output['hsu']
#Add in the catchment info
output['wbd'] = wbd
#Close the file
fp.close()
return output
