def smooth(da_w): da, w = da_w mid = int(da.shape[0]/2) if (da.shape[-1]==0) | (w.shape[-1]==0): return da[mid], w[mid] data = np.array(smoothn(da, s=10., smoothOrder=1., axis=0, TolZ=0.001, verbose=False, isrobust=True, W = w))[[0, 3],] return data[0][mid], data[1][mid]
def do_aot_tcwv(aot_bands, tcwv_bands, cams_dir, obs_time, sun_ang_name, view_ang_name, dem, tcwv_name=None, aot_name=None): toas = [ reproject_data(str(band), tcwv_bands[0], dstNodata=0, resample=5, xRes=120, yRes=120).data for band in aot_bands ] toas = np.array(toas) / 10000. mask = np.all(toas >= 0.0001, axis=0) time_ind = np.abs((obs_time.hour + obs_time.minute / 60. + obs_time.second / 3600.) - np.arange(0, 25, 3)).argmin() prior_uncs = 0.1 prior_scale = 46.698 prior_f = cams_dir + '/'.join([ datetime.datetime.strftime(obs_time, '%Y_%m_%d'), datetime.datetime.strftime(obs_time, '%Y_%m_%d') + '_gtco3.tif' ]) var_g = gdal.Open(prior_f) prior_g = reproject_data(prior_f, tcwv_bands[0], dstNodata=0, resample=1, xRes=120, yRes=120).g g = var_g.GetRasterBand(int(time_ind + 1)) offset = g.GetOffset() scale = g.GetScale() tco3 = prior_g.GetRasterBand( int(time_ind + 1)).ReadAsArray() * scale + offset tco3[:] = np.nanmean(tco3) * prior_scale saa, sza = reproject_data(str(sun_ang_name), tcwv_bands[0], dstNodata=0, resample=1, xRes=120, yRes=120).data / 100. vaa, vza = reproject_data(str(view_ang_name), tcwv_bands[0], dstNodata=0, resample=1, xRes=120, yRes=120).data / 100. raa = vaa - saa sza = np.cos(np.deg2rad(sza)) vza = np.cos(np.deg2rad(vza)) raa = np.cos(np.deg2rad(raa)) ele = reproject_data( dem, tcwv_bands[0], dstNodata=0, resample=1, xRes=120, yRes=120).data / 10000. X = np.vstack([toas[[9, 8]], np.array([sza, vza, raa, tco3, ele])]).reshape(7, -1).T iso_tcwv = Two_NN(np_model_file=file_path + '/emus/S2_TCWV.npz') tcwv = iso_tcwv.predict(X)[0].reshape(toas[0].shape).astype(float) bad = (tcwv < 0) | (tcwv > 8) | (~mask) | np.isnan(tcwv) tcwv[bad] = np.nanmedian(tcwv[~bad]) bands_min = [ 0.11572497, 0.08986528, 0.07280412, 0.05007033, 0.06228712, 0.06849915, 0.07015634, 0.06554198, 0.06809415, 0.02089167, 0.00080242, 0.04392302, 0.03012728 ] bands_max = [ 0.32072931, 0.30801709, 0.32420026, 0.38214899, 0.3951169, 0.41631785, 0.44391895, 0.42444658, 0.46161492, 0.19927658, 0.00870671, 0.39533241, 0.32505772 ] for _, toa in enumerate(toas): mask = mask & (toa >= bands_min[_]) & (toa <= bands_max[_]) X = np.vstack([toas, np.array([sza, vza, raa, tco3, ele])]).reshape(18, -1).T gbm = joblib.load(file_path + '/emus/lgb.pkl') aot = gbm.predict(X).reshape(toas[0].shape).astype(float) aot = np.exp(-1 * aot) shape = toas[0].shape if mask.sum() > 3: aot_min, aot_median, aot_max = np.nanpercentile(aot[mask], [5, 50, 95]) good_aot = (aot >= aot_min) & (aot <= aot_max) indx, indy = np.where(good_aot.reshape(shape)) myInterpolator = NearestNDInterpolator((indx, indy), aot[good_aot]) grid_x, grid_y = np.mgrid[0:shape[0]:1, 0:shape[1]:1, ] aot = myInterpolator(grid_x, grid_y) aot = smoothn(aot, isrobust=True, TolZ=1e-6, W=100 * ((aot >= aot_min) & (aot <= aot_max)), s=10, MaxIter=1000)[0] else: aot = np.nanmedian(aot) * np.ones(shape) if tcwv_name is not None: g = gdal.Open(tcwv_bands[0]) ySize, xSize = tcwv.shape dst = gdal.GetDriverByName('GTiff').Create( tcwv_name, xSize, ySize, 1, gdal.GDT_UInt16, options=["TILED=YES", "COMPRESS=DEFLATE"]) dst.SetGeoTransform(g.GetGeoTransform()) dst.SetProjection(g.GetProjection()) dst.GetRasterBand(1).WriteArray((tcwv * 1000).astype(int)) dst.GetRasterBand(1).SetNoDataValue(65535) dst = None g = None if aot_name is not None: g = gdal.Open(aot_bands[0]) ySize, xSize = aot.shape dst = gdal.GetDriverByName('GTiff').Create( aot_name, xSize, ySize, 1, gdal.GDT_UInt16, options=["TILED=YES", "COMPRESS=DEFLATE"]) dst.SetGeoTransform(g.GetGeoTransform()) dst.SetProjection(g.GetProjection()) dst.GetRasterBand(1).WriteArray((aot * 1000).astype(int)) dst.GetRasterBand(1).SetNoDataValue(65535) dst = None g = None return aot, tcwv