diff --git a/CHANGELOG.md b/CHANGELOG.md index 145b0cb76..05487c26c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,15 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/), and this project adheres to [PEP 440](https://www.python.org/dev/peps/pep-0440/) and uses [Semantic Versioning](https://semver.org/spec/v2.0.0.html). +## [0.5.0] +### Added +* A `--input-bucket-prefix` argument to `calcDelaysGUNW` which will allow RAiDER to process ARIA GUNW products under one prefix and upload the final products to another prefix provided by the `--bucket-prefix` argument. +### Fixed +* [613](https://github.com/dbekaert/RAiDER/issues/613) - ensure NASA Earthdata credentials for downloading orbits from ASF +* [634](https://github.com/dbekaert/RAiDER/issues/634) - download orbits from ASF before trying ESA +* [630](https://github.com/dbekaert/RAiDER/pull/630) - use correct model name so (hrrr-ak) in azimuth_timing_grid +* [620](https://github.com/dbekaert/RAiDER/issues/620) - Fix MERRA-2 access because API was updated + ## [0.4.7] ### Fixed diff --git a/test/__init__.py b/test/__init__.py index 113a8326c..0faf8b712 100644 --- a/test/__init__.py +++ b/test/__init__.py @@ -16,7 +16,7 @@ WM_DIR = os.path.join(TEST_DIR, 'weather_files') ORB_DIR = os.path.join(TEST_DIR, 'orbit_files') -WM = 'GMAO' +WM = 'MERRA2' @contextmanager def pushd(dir): diff --git a/test/test_GUNW.py b/test/test_GUNW.py index a57152d2f..6eb2a39f1 100644 --- a/test/test_GUNW.py +++ b/test/test_GUNW.py @@ -100,7 +100,7 @@ def test_GUNW_hyp3_metadata_update(test_gunw_json_path, test_gunw_json_schema_pa # We only need to make sure the json file is passed, the netcdf file name will not have # any impact on subsequent testing - mocker.patch("RAiDER.aws.get_s3_file", side_effect=['foo.nc', temp_json_path]) + mocker.patch("RAiDER.aws.get_s3_file", side_effect=['foo.nc', temp_json_path, 'foo.png']) mocker.patch("RAiDER.aws.upload_file_to_s3") mocker.patch("RAiDER.aria.prepFromGUNW.main", return_value=['my_path_cfg', 'my_wavelength']) mocker.patch('RAiDER.aria.prepFromGUNW.check_hrrr_dataset_availablity_for_s1_azimuth_time_interpolation', @@ -111,7 +111,7 @@ def test_GUNW_hyp3_metadata_update(test_gunw_json_path, test_gunw_json_schema_pa iargs = ['--weather-model', 'HRES', '--bucket', 'myBucket', - '--bucket-prefix', 'myPrefix'] + '--bucket-prefix', 'myPrefix',] calcDelaysGUNW(iargs) metadata = json.loads(temp_json_path.read_text()) @@ -123,6 +123,7 @@ def test_GUNW_hyp3_metadata_update(test_gunw_json_path, test_gunw_json_schema_pa assert aws.get_s3_file.mock_calls == [ unittest.mock.call('myBucket', 'myPrefix', '.nc'), unittest.mock.call('myBucket', 'myPrefix', '.json'), + unittest.mock.call('myBucket', 'myPrefix', '.png'), ] RAiDER.aria.prepFromGUNW.main.assert_called_once() @@ -138,6 +139,60 @@ def test_GUNW_hyp3_metadata_update(test_gunw_json_path, test_gunw_json_schema_pa assert aws.upload_file_to_s3.mock_calls == [ unittest.mock.call('foo.nc', 'myBucket', 'myPrefix'), unittest.mock.call(temp_json_path, 'myBucket', 'myPrefix'), + unittest.mock.call('foo.png', 'myBucket', 'myPrefix'), + ] + + +def test_GUNW_hyp3_metadata_update_different_prefix(test_gunw_json_path, test_gunw_json_schema_path, tmp_path, mocker): + """This test performs the GUNW entrypoint using a different input bucket/prefix than the output bucket/prefix. + Only updates the json. Monkey patches the upload/download to/from s3 and the actual computation. + """ + temp_json_path = tmp_path / 'temp.json' + shutil.copy(test_gunw_json_path, temp_json_path) + + # We only need to make sure the json file is passed, the netcdf file name will not have + # any impact on subsequent testing + mocker.patch("RAiDER.aws.get_s3_file", side_effect=['foo.nc', temp_json_path, 'foo.png']) + mocker.patch("RAiDER.aws.upload_file_to_s3") + mocker.patch("RAiDER.aria.prepFromGUNW.main", return_value=['my_path_cfg', 'my_wavelength']) + mocker.patch('RAiDER.aria.prepFromGUNW.check_hrrr_dataset_availablity_for_s1_azimuth_time_interpolation', + side_effect=[True]) + mocker.patch("RAiDER.aria.prepFromGUNW.check_weather_model_availability", return_value=True) + mocker.patch("RAiDER.cli.raider.calcDelays", return_value=['file1', 'file2']) + mocker.patch("RAiDER.aria.calcGUNW.tropo_gunw_slc") + + iargs = ['--weather-model', 'HRES', + '--bucket', 'myBucket', + '--bucket-prefix', 'myOutputPrefix', + '--input-bucket-prefix', 'myInputPrefix',] + calcDelaysGUNW(iargs) + + metadata = json.loads(temp_json_path.read_text()) + schema = json.loads(test_gunw_json_schema_path.read_text()) + + assert metadata['metadata']['weather_model'] == ['HRES'] + assert (jsonschema.validate(instance=metadata, schema=schema) is None) + + assert aws.get_s3_file.mock_calls == [ + unittest.mock.call('myBucket', 'myInputPrefix', '.nc'), + unittest.mock.call('myBucket', 'myInputPrefix', '.json'), + unittest.mock.call('myBucket', 'myInputPrefix', '.png'), + ] + + RAiDER.aria.prepFromGUNW.main.assert_called_once() + + raider.calcDelays.assert_called_once_with(['my_path_cfg']) + + RAiDER.aria.calcGUNW.tropo_gunw_slc.assert_called_once_with( + ['file1', 'file2'], + 'foo.nc', + 'my_wavelength', + ) + + assert aws.upload_file_to_s3.mock_calls == [ + unittest.mock.call('foo.nc', 'myBucket', 'myOutputPrefix'), + unittest.mock.call(temp_json_path, 'myBucket', 'myOutputPrefix'), + unittest.mock.call('foo.png', 'myBucket', 'myOutputPrefix'), ] @@ -277,7 +332,7 @@ def test_azimuth_timing_interp_against_center_time_interp(weather_model_name: st assert np.nanmax(abs_diff_mm) < 1 -@pytest.mark.parametrize('weather_model_name', ['GMAO', 'HRRR', 'HRES', 'ERA5', 'ERA5']) +@pytest.mark.parametrize('weather_model_name', ['HRRR', 'HRES', 'ERA5', 'ERA5T']) def test_check_weather_model_availability(test_gunw_path_factory, weather_model_name, mocker): # Should be True for all weather models # S1-GUNW-D-R-071-tops-20200130_20200124-135156-34956N_32979N-PP-913f-v2_0_4.nc @@ -288,12 +343,12 @@ def test_check_weather_model_availability(test_gunw_path_factory, weather_model_ mocker.patch("RAiDER.aria.prepFromGUNW.get_acq_time_from_slc_id", side_effect=[pd.Timestamp('2015-01-01'), pd.Timestamp('2014-01-01')]) cond = check_weather_model_availability(test_gunw_path, weather_model_name) - if weather_model_name in ['HRRR', 'GMAO']: + if weather_model_name in ['HRRR', 'MERRA2']: cond = not cond assert cond -@pytest.mark.parametrize('weather_model_name', ['GMAO', 'HRRR']) +@pytest.mark.parametrize('weather_model_name', ['HRRR']) def test_check_weather_model_availability_over_alaska(test_gunw_path_factory, weather_model_name, mocker): # Should be True for all weather models # S1-GUNW-D-R-059-tops-20230320_20220418-180300-00179W_00051N-PP-c92e-v2_0_6.nc @@ -309,7 +364,7 @@ def test_check_weather_model_availability_over_alaska(test_gunw_path_factory, we assert cond -@pytest.mark.parametrize('weather_model_name', ['HRRR', 'GMAO']) +@pytest.mark.parametrize('weather_model_name', ['HRRR']) @pytest.mark.parametrize('location', ['california-t71', 'alaska']) def test_weather_model_availability_integration_using_valid_range(location, test_gunw_path_factory, diff --git a/test/test_s1_orbits.py b/test/test_s1_orbits.py index 424b333a2..681f1e444 100644 --- a/test/test_s1_orbits.py +++ b/test/test_s1_orbits.py @@ -12,34 +12,63 @@ class EmptyNetrc(): def __init__(self, netrc_file): self.netrc_file = netrc_file self.hosts = {} + def __str__(self): return str(self.hosts) - # No .netrc, no ESA CDSE env variables + # No .netrc, no ESA CDSE or Earthdata env variables + with monkeypatch.context() as mp: + mp.setattr(netrc, 'netrc', EmptyNetrc, raising=False) + mp.delenv('ESA_USERNAME', raising=False) + mp.delenv('ESA_PASSWORD', raising=False) + mp.delenv('EARTHDATA_USERNAME', raising=False) + mp.delenv('EARTHDATA_PASSWORD', raising=False) + with pytest.raises(ValueError): + s1_orbits.ensure_orbit_credentials() + + # No .netrc or Earthdata env vars, set ESA CDSE env variables + with monkeypatch.context() as mp: + mp.setattr(netrc, 'netrc', EmptyNetrc, raising=False) + mp.setenv('ESA_USERNAME', 'foo') + mp.setenv('ESA_PASSWORD', 'bar') + mp.delenv('EARTHDATA_USERNAME', raising=False) + mp.delenv('EARTHDATA_PASSWORD', raising=False) + with pytest.raises(ValueError): + s1_orbits.ensure_orbit_credentials() + + # No .netrc or ESA CDSE env vars, set Earthdata env variables with monkeypatch.context() as mp: mp.setattr(netrc, 'netrc', EmptyNetrc, raising=False) mp.delenv('ESA_USERNAME', raising=False) mp.delenv('ESA_PASSWORD', raising=False) + mp.setenv('EARTHDATA_USERNAME', 'fizz') + mp.setenv('EARTHDATA_PASSWORD', 'buzz') with pytest.raises(ValueError): s1_orbits.ensure_orbit_credentials() - # No .netrc, set ESA CDSE env variables + # No .netrc, set Earthdata and ESA CDSE env variables with monkeypatch.context() as mp: mp.setattr(netrc, 'netrc', EmptyNetrc, raising=False) mp.setenv('ESA_USERNAME', 'foo') mp.setenv('ESA_PASSWORD', 'bar') + mp.setenv('EARTHDATA_USERNAME', 'fizz') + mp.setenv('EARTHDATA_PASSWORD', 'buzz') mp.setattr(Path, 'write_text', lambda self, write_text: write_text) written_credentials = s1_orbits.ensure_orbit_credentials() - assert written_credentials == str({s1_orbits.ESA_CDSE_HOST: ('foo', None, 'bar')}) + assert written_credentials == str({ + s1_orbits.ESA_CDSE_HOST: ('foo', None, 'bar'), + s1_orbits.NASA_EDL_HOST: ('fizz', None, 'buzz') + }) class NoCDSENetrc(): def __init__(self, netrc_file): self.netrc_file = netrc_file - self.hosts = {'fizz.buzz.org': ('foo', None, 'bar')} + self.hosts = {s1_orbits.NASA_EDL_HOST: ('fizz', None, 'buzz')} + def __str__(self): return str(self.hosts) - # No CDSE in .netrc, no ESA CDSE env variables + # No CDSE in .netrc or ESA CDSE env variables, Earthdata in .netrc with monkeypatch.context() as mp: mp.setattr(netrc, 'netrc', NoCDSENetrc, raising=False) mp.delenv('ESA_USERNAME', raising=False) @@ -47,66 +76,122 @@ def __str__(self): with pytest.raises(ValueError): s1_orbits.ensure_orbit_credentials() - # No CDSE in .netrc, set ESA CDSE env variables + # No CDSE in .netrc, set ESA CDSE env variables, Earthdata in .netrc with monkeypatch.context() as mp: mp.setattr(netrc, 'netrc', NoCDSENetrc, raising=False) mp.setenv('ESA_USERNAME', 'foo') mp.setenv('ESA_PASSWORD', 'bar') mp.setattr(Path, 'write_text', lambda self, write_text: write_text) written_credentials = s1_orbits.ensure_orbit_credentials() - assert written_credentials == str({'fizz.buzz.org': ('foo', None, 'bar'), s1_orbits.ESA_CDSE_HOST: ('foo', None, 'bar')}) + assert written_credentials == str({ + s1_orbits.NASA_EDL_HOST: ('fizz', None, 'buzz'), + s1_orbits.ESA_CDSE_HOST: ('foo', None, 'bar'), + }) - class CDSENetrc(): + class NoEarthdataNetrc(): def __init__(self, netrc_file): self.netrc_file = netrc_file self.hosts = {s1_orbits.ESA_CDSE_HOST: ('foo', None, 'bar')} + def __str__(self): return str(self.hosts) - # cdse in .netrc, no ESA CDSE env variables + # cdse in .netrc, no ESA CDSE env variables, Earthdata env variables with monkeypatch.context() as mp: - mp.setattr(netrc, 'netrc', CDSENetrc, raising=False) + mp.setattr(netrc, 'netrc', NoEarthdataNetrc, raising=False) mp.delenv('ESA_USERNAME', raising=False) mp.delenv('ESA_PASSWORD', raising=False) + mp.setenv('EARTHDATA_USERNAME', 'fizz') + mp.setenv('EARTHDATA_PASSWORD', 'buzz') + mp.setattr(Path, 'write_text', lambda self, write_text: write_text) + written_credentials = s1_orbits.ensure_orbit_credentials() + assert written_credentials == str({ + s1_orbits.ESA_CDSE_HOST: ('foo', None, 'bar'), + s1_orbits.NASA_EDL_HOST: ('fizz', None, 'buzz'), + }) + + class CDSEAndEarthdataNetrc(): + def __init__(self, netrc_file): + self.netrc_file = netrc_file + self.hosts = { + s1_orbits.ESA_CDSE_HOST: ('foo', None, 'bar'), + s1_orbits.NASA_EDL_HOST: ('fizz', None, 'buzz') + } + + def __str__(self): + return str(self.hosts) + + # cdse and Earthdata in netrc, no env variables + with monkeypatch.context() as mp: + mp.setattr(netrc, 'netrc', CDSEAndEarthdataNetrc, raising=False) + written_credentials = s1_orbits.ensure_orbit_credentials() + assert written_credentials is None + + with monkeypatch.context() as mp: + mp.setattr(netrc, 'netrc', CDSEAndEarthdataNetrc, raising=False) + mp.delenv('ESA_USERNAME', raising=False) + mp.delenv('ESA_PASSWORD', raising=False) + mp.setenv('EARTHDATA_USERNAME', 'fizz') + mp.setenv('EARTHDATA_PASSWORD', 'buzz') written_credentials = s1_orbits.ensure_orbit_credentials() assert written_credentials is None - # cdse in .netrc, set ESA CDSE env variables with monkeypatch.context() as mp: - mp.setattr(netrc, 'netrc', CDSENetrc, raising=False) + mp.setattr(netrc, 'netrc', CDSEAndEarthdataNetrc, raising=False) mp.setenv('ESA_USERNAME', 'foo') mp.setenv('ESA_PASSWORD', 'bar') + mp.setenv('EARTHDATA_USERNAME', 'fizz') + mp.setenv('EARTHDATA_PASSWORD', 'buzz') + written_credentials = s1_orbits.ensure_orbit_credentials() + assert written_credentials is None + + with monkeypatch.context() as mp: + mp.setattr(netrc, 'netrc', CDSEAndEarthdataNetrc, raising=False) + mp.setenv('ESA_USERNAME', 'foo') + mp.setenv('ESA_PASSWORD', 'bar') + mp.delenv('ESA_USERNAME', raising=False) + mp.delenv('ESA_PASSWORD', raising=False) written_credentials = s1_orbits.ensure_orbit_credentials() assert written_credentials is None def test_get_orbits_from_slc_ids(mocker): side_effect = [ - [Path('foo.txt')], - [Path('bar.txt'), Path('fiz.txt')], + [Path('foo_start.txt'), Path('foo_stop.txt')], + [Path('bar_start.txt'), Path('bar_end.txt'), + Path('fiz_start.txt'), Path('fiz_end')], ] mocker.patch('eof.download.download_eofs', - side_effect=side_effect) + side_effect=side_effect[0]) orbit_files = s1_orbits.get_orbits_from_slc_ids( ['S1A_IW_SLC__1SSV_20150621T120220_20150621T120232_006471_008934_72D8'] ) - assert orbit_files == [Path('foo.txt')] - assert eof.download.download_eofs.call_count == 1 - eof.download.download_eofs.assert_called_with( - [ '20150621T120220', '20150621T120232'], - ['S1A'] * 2, - save_dir=str(Path.cwd()) - ) + assert orbit_files == side_effect[0] + assert eof.download.download_eofs.call_count == 2 + for dt in '20150621T120220 20150621T120232'.split(): + eof.download.download_eofs.assert_any_call( + [dt], + ['S1A'], + save_dir=str(Path.cwd()), + force_asf=True + ) + + mocker.patch('eof.download.download_eofs', + side_effect=side_effect[1]) orbit_files = s1_orbits.get_orbits_from_slc_ids( ['S1B_IW_SLC__1SDV_20201115T162313_20201115T162340_024278_02E29D_5C54', 'S1A_IW_SLC__1SDV_20201203T162353_20201203T162420_035524_042744_6D5C'] ) - assert orbit_files == [Path('bar.txt'), Path('fiz.txt')] - assert eof.download.download_eofs.call_count == 2 - eof.download.download_eofs.assert_called_with( - ['20201115T162313', '20201203T162353', '20201115T162340', '20201203T162420'], - ['S1B', 'S1A'] * 2, - save_dir=str(Path.cwd()) - ) + assert orbit_files == side_effect[1] + assert eof.download.download_eofs.call_count == 4 + missions = 'S1B S1B S1A S1A'.split() + dts = '20201115T162313 20201115T162340 20201203T162353 20201203T162420'.split() + for dt, mission in zip(dts, missions): + eof.download.download_eofs.assert_any_call( + [dt], + [mission], + save_dir=str(Path.cwd()), + force_asf=True + ) diff --git a/tools/RAiDER/aria/prepFromGUNW.py b/tools/RAiDER/aria/prepFromGUNW.py index 4561abe4f..52f10afdb 100644 --- a/tools/RAiDER/aria/prepFromGUNW.py +++ b/tools/RAiDER/aria/prepFromGUNW.py @@ -8,7 +8,6 @@ import os from datetime import datetime import numpy as np -import eof.download import xarray as xr import rasterio import pandas as pd @@ -23,10 +22,10 @@ from RAiDER.models import credentials from RAiDER.models.hrrr import HRRR_CONUS_COVERAGE_POLYGON, AK_GEO, check_hrrr_dataset_availability from RAiDER.s1_azimuth_timing import get_times_for_azimuth_interpolation -from RAiDER.s1_orbits import ensure_orbit_credentials +from RAiDER.s1_orbits import download_eofs ## cube spacing in degrees for each model -DCT_POSTING = {'HRRR': 0.05, 'HRES': 0.10, 'GMAO': 0.10, 'ERA5': 0.10, 'ERA5T': 0.10} +DCT_POSTING = {'HRRR': 0.05, 'HRES': 0.10, 'GMAO': 0.10, 'ERA5': 0.10, 'ERA5T': 0.10, 'MERRA2': 0.1} def _get_acq_time_from_gunw_id(gunw_id: str, reference_or_secondary: str) -> datetime: @@ -95,7 +94,7 @@ def check_weather_model_availability(gunw_path: str, ---------- gunw_path : str weather_model_name : str - Should be one of 'HRRR', 'HRES', 'ERA5', 'ERA5T', 'GMAO'. + Should be one of 'HRRR', 'HRES', 'ERA5', 'ERA5T', 'GMAO', 'MERRA2'. Returns ------- bool: @@ -276,8 +275,7 @@ def get_orbit_file(self): sat = slc.split('_')[0] dt = datetime.strptime(f'{self.dates[0]}T{self.mid_time}', '%Y%m%dT%H:%M:%S') - ensure_orbit_credentials() - path_orb = eof.download.download_eofs([dt], [sat], save_dir=orbit_dir) + path_orb = download_eofs([dt], [sat], str(orbit_dir)) return [str(o) for o in path_orb] diff --git a/tools/RAiDER/cli/raider.py b/tools/RAiDER/cli/raider.py index 4034d8da4..facf34eab 100644 --- a/tools/RAiDER/cli/raider.py +++ b/tools/RAiDER/cli/raider.py @@ -461,7 +461,15 @@ def calcDelaysGUNW(iargs: list[str] = None) -> xr.Dataset: p.add_argument( '--bucket-prefix', default='', - help='S3 bucket prefix containing ARIA GUNW NetCDF file. Will be ignored if the --file argument is provided.' + help='S3 bucket prefix which may contain an ARIA GUNW NetCDF file to calculate delays for and which the final ' + 'ARIA GUNW NetCDF file will be upload to. Will be ignored if the --file argument is provided.' + ) + + p.add_argument( + '--input-bucket-prefix', + help='S3 bucket prefix that contains an ARIA GUNW NetCDF file to calculate delays for. ' + 'If not provided, will look in --bucket-prefix for an ARIA GUNW NetCDF file. ' + 'Will be ignored if the --file argument is provided.' ) p.add_argument( @@ -501,6 +509,9 @@ def calcDelaysGUNW(iargs: list[str] = None) -> xr.Dataset: iargs = p.parse_args(iargs) + if not iargs.input_bucket_prefix: + iargs.input_bucket_prefix = iargs.bucket_prefix + if iargs.interpolate_time not in ['none', 'center_time', 'azimuth_time_grid']: raise ValueError('interpolate_time arg must be in [\'none\', \'center_time\', \'azimuth_time_grid\']') @@ -520,7 +531,7 @@ def calcDelaysGUNW(iargs: list[str] = None) -> xr.Dataset: if not iargs.file and iargs.bucket: # only use GUNW ID for checking if HRRR available - iargs.file = aws.get_s3_file(iargs.bucket, iargs.bucket_prefix, '.nc') + iargs.file = aws.get_s3_file(iargs.bucket, iargs.input_bucket_prefix, '.nc') if iargs.weather_model == 'HRRR' and (iargs.interpolate_time == 'azimuth_time_grid'): file_name_str = str(iargs.file) gunw_nc_name = file_name_str.split('/')[-1] @@ -536,11 +547,16 @@ def calcDelaysGUNW(iargs: list[str] = None) -> xr.Dataset: # we include this within this portion of the control flow. print('Nothing to do because outside of weather model range') return - json_file_path = aws.get_s3_file(iargs.bucket, iargs.bucket_prefix, '.json') + json_file_path = aws.get_s3_file(iargs.bucket, iargs.input_bucket_prefix, '.json') json_data = json.load(open(json_file_path)) json_data['metadata'].setdefault('weather_model', []).append(iargs.weather_model) json.dump(json_data, open(json_file_path, 'w')) + # also get browse image -- if RAiDER is running in its own HyP3 job, the browse image will be needed for ingest + browse_file_path = aws.get_s3_file(iargs.bucket, iargs.input_bucket_prefix, '.png') + + + elif not iargs.file: raise ValueError('Either argument --file or --bucket must be provided') @@ -563,6 +579,7 @@ def calcDelaysGUNW(iargs: list[str] = None) -> xr.Dataset: if iargs.bucket: aws.upload_file_to_s3(iargs.file, iargs.bucket, iargs.bucket_prefix) aws.upload_file_to_s3(json_file_path, iargs.bucket, iargs.bucket_prefix) + aws.upload_file_to_s3(browse_file_path, iargs.bucket, iargs.bucket_prefix) return ds @@ -597,13 +614,13 @@ def combineZTDFiles(): def getWeatherFile(wfiles, times, t, model, interp_method='none'): ''' # Time interpolation - # + # # Need to handle various cases, including if the exact weather model time is - # requested, or if one or more datetimes are not available from the weather + # requested, or if one or more datetimes are not available from the weather # model data provider - ''' - - # time interpolation method: number of expected files + ''' + + # time interpolation method: number of expected files EXPECTED_NUM_FILES = {'none': 1, 'center_time': 2, 'azimuth_time_grid': 3} Nfiles = len(wfiles) @@ -613,8 +630,8 @@ def getWeatherFile(wfiles, times, t, model, interp_method='none'): Nfiles_expected = EXPECTED_NUM_FILES[interp_method] except KeyError: raise ValueError('getWeatherFile: interp_method {} is not known'.format(interp_method)) - - Nmatch = (Nfiles_expected == Nfiles) + + Nmatch = (Nfiles_expected == Nfiles) Tmatch = (Nfiles == Ntimes) # Case 1: no files downloaded @@ -622,7 +639,7 @@ def getWeatherFile(wfiles, times, t, model, interp_method='none'): logger.error('No weather model data was successfully processed.') return None - # Case 2 - nearest weather model time is requested and retrieved + # Case 2 - nearest weather model time is requested and retrieved if (interp_method == 'none'): weather_model_file = wfiles[0] @@ -631,8 +648,8 @@ def getWeatherFile(wfiles, times, t, model, interp_method='none'): if Nmatch: # Case 3: two weather files downloaded weather_model_file = combine_weather_files( - wfiles, - t, + wfiles, + t, model, interp_method='center_time' ) @@ -644,17 +661,17 @@ def getWeatherFile(wfiles, times, t, model, interp_method='none'): weather_model_file = wfiles[0] else: raise WrongNumberOfFiles(Nfiles_expected, Nfiles) - + elif (interp_method) == 'azimuth_time_grid': if Nmatch or Tmatch: # Case 6: all files downloaded weather_model_file = combine_weather_files( - wfiles, + wfiles, t, model, interp_method='azimuth_time_grid' ) - else: + else: raise WrongNumberOfFiles(Nfiles_expected, Nfiles) # Case 7 - Anything else errors out @@ -662,7 +679,7 @@ def getWeatherFile(wfiles, times, t, model, interp_method='none'): N = len(wfiles) raise NotImplementedError(f'The {interp_method} with {N} retrieved weather model files was not well posed ' 'for the current workflow.') - + return weather_model_file @@ -678,7 +695,7 @@ def combine_weather_files(wfiles, t, model, interp_method='center_time'): times = [] for ds in datasets: times.append(datetime.datetime.strptime(ds.attrs['datetime'], '%Y_%m_%dT%H_%M_%S')) - + if len(times)==0: raise NoWeatherModelData() @@ -699,8 +716,8 @@ def combine_weather_files(wfiles, t, model, interp_method='center_time'): # Give the weighted combination a new file name weather_model_file = os.path.join( os.path.dirname(wfiles[0]), - os.path.basename(wfiles[0]).split('_')[0] + '_' + - t.strftime('%Y_%m_%dT%H_%M_%S') + STYLE[interp_method] + + os.path.basename(wfiles[0]).split('_')[0] + '_' + + t.strftime('%Y_%m_%dT%H_%M_%S') + STYLE[interp_method] + '_'.join(wfiles[0].split('_')[-4:]), ) @@ -720,7 +737,7 @@ def combine_files_using_azimuth_time(wfiles, t, times): times = [] for ds in datasets: times.append(datetime.datetime.strptime(ds.attrs['datetime'], '%Y_%m_%dT%H_%M_%S')) - + model = datasets[0].attrs['model_name'] time_grid = get_time_grid_for_aztime_interp(datasets, times, t, model) @@ -750,13 +767,13 @@ def get_weights_time_interp(times, t): '''Calculate weights for time interpolation using simple inverse linear weighting''' date1,date2 = times wgts = [ 1 - get_dt(t, date1) / get_dt(date2, date1), 1 - get_dt(date2, t) / get_dt(date2, date1)] - + try: assert np.isclose(np.sum(wgts), 1) except AssertionError: logger.error('Time interpolation weights do not sum to one; something is off with query datetime: %s', t) return None - + return wgts @@ -767,7 +784,7 @@ def get_time_grid_for_aztime_interp(datasets, t, model): # the subsequent s1 azimuth time grid requires dimension # inputs to all have same dimensions and either be # 1d or 3d. - AZ_TIME_ALLOWED_MODELS = ['hrrr', 'hrrrak'] + AZ_TIME_ALLOWED_MODELS = 'hrrr hrrrak hrrr-ak'.split() if model.lower() in AZ_TIME_ALLOWED_MODELS: lat_2d = datasets[0].latitude.data @@ -785,5 +802,5 @@ def get_time_grid_for_aztime_interp(datasets, t, model): time_grid = get_s1_azimuth_time_grid(lon, lat, hgt, t) # This is the acq time from loop if np.any(np.isnan(time_grid)): raise ValueError('The Time Grid return nans meaning no orbit was downloaded.') - + return time_grid diff --git a/tools/RAiDER/losreader.py b/tools/RAiDER/losreader.py index 426305f0b..3faaa7ef5 100644 --- a/tools/RAiDER/losreader.py +++ b/tools/RAiDER/losreader.py @@ -378,11 +378,11 @@ def filter_ESA_orbit_file_p(path: str) -> bool: svs = read_shelve(los_file) except BaseException: raise ValueError( - 'get_sv: I cannot parse the statevector file {}'.format(los_file) + f'get_sv: I cannot parse the statevector file {los_file}' ) except: - raise ValueError('get_sv: I cannot parse the statevector file {}'.format(los_file)) - + raise ValueError(f'get_sv: I cannot parse the statevector file {los_file}') + if ref_time: idx = cut_times(svs[0], ref_time, pad=pad) diff --git a/tools/RAiDER/models/__init__.py b/tools/RAiDER/models/__init__.py index 1d9307bda..79666992e 100644 --- a/tools/RAiDER/models/__init__.py +++ b/tools/RAiDER/models/__init__.py @@ -3,5 +3,7 @@ from .gmao import GMAO from .hres import HRES from .hrrr import HRRR, HRRRAK +from .merra2 import MERRA2 +from .ncmr import NCMR -__all__ = ['HRRR', 'HRRRAK', 'GMAO', 'ERA5', 'ERA5T', 'HRES'] +__all__ = ['HRRR', 'HRRRAK', 'GMAO', 'ERA5', 'ERA5T', 'HRES', 'MERRA2'] diff --git a/tools/RAiDER/models/allowed.py b/tools/RAiDER/models/allowed.py index e5be23652..9d1b1d1ea 100755 --- a/tools/RAiDER/models/allowed.py +++ b/tools/RAiDER/models/allowed.py @@ -3,5 +3,7 @@ 'ERA5T', 'HRRR', 'GMAO', - 'HRES' + 'HRES', + 'MERRA2', + 'NCMR', ] diff --git a/tools/RAiDER/models/gmao.py b/tools/RAiDER/models/gmao.py index 04ab1f507..bf23ec58c 100755 --- a/tools/RAiDER/models/gmao.py +++ b/tools/RAiDER/models/gmao.py @@ -9,7 +9,7 @@ from RAiDER.models.weatherModel import WeatherModel, TIME_RES from RAiDER.logger import logger -from RAiDER.utilFcns import writeWeatherVars2NETCDF4, round_date, requests_retry_session +from RAiDER.utilFcns import writeWeatherVarsXarray, round_date, requests_retry_session from RAiDER.models.model_levels import ( LEVELS_137_HEIGHTS, ) @@ -145,7 +145,7 @@ def _fetch(self, out): try: # Note that lat/lon gets written twice for GMAO because they are the same as y/x - writeWeatherVars2NETCDF4(self, lats, lons, h, qv, p, t, outName=out) + writeWeatherVarsXarray(lat, lon, h, q, p, t, dt, crs, outName=None, NoDataValue=None, chunk=(1, 91, 144)) except Exception: logger.exception("Unable to save weathermodel to file") diff --git a/tools/RAiDER/models/merra2.py b/tools/RAiDER/models/merra2.py index 85a3a7fc4..e4aec16dc 100755 --- a/tools/RAiDER/models/merra2.py +++ b/tools/RAiDER/models/merra2.py @@ -1,4 +1,6 @@ import os +import xarray + import datetime as dt import numpy as np import pydap.cas.urs @@ -8,7 +10,7 @@ from RAiDER.models.weatherModel import WeatherModel from RAiDER.logger import logger -from RAiDER.utilFcns import writeWeatherVars2NETCDF4, read_EarthData_loginInfo +from RAiDER.utilFcns import writeWeatherVarsXarray, read_EarthData_loginInfo from RAiDER.models.model_levels import ( LEVELS_137_HEIGHTS, ) @@ -94,6 +96,8 @@ def _fetch(self, out): self._lon_res ) + lon, lat = np.meshgrid(lons, lats) + if time.year < 1992: url_sub = 100 elif time.year < 2001: @@ -107,39 +111,22 @@ def _fetch(self, out): DT = time - T0 time_ind = int(DT.total_seconds() / 3600.0 / 3.0) - ml_min = 0 - ml_max = 71 - # Earthdata credentials earthdata_usr, earthdata_pwd = read_EarthData_loginInfo(EARTHDATA_RC) # open the dataset and pull the data - url = 'https://goldsmr5.gesdisc.eosdis.nasa.gov:443/opendap/MERRA2/M2I3NVASM.5.12.4/' + time.strftime('%Y/%m') + '/MERRA2_' + str(url_sub) + '.inst3_3d_asm_Nv.' + time.strftime('%Y%m%d') + '.nc4' + url = 'https://goldsmr5.gesdisc.eosdis.nasa.gov/opendap/MERRA2/M2T3NVASM.5.12.4/' + time.strftime('%Y/%m') + '/MERRA2_' + str(url_sub) + '.tavg3_3d_asm_Nv.' + time.strftime('%Y%m%d') + '.nc4' + session = pydap.cas.urs.setup_session(earthdata_usr, earthdata_pwd, check_url=url) - ds = pydap.client.open_url(url, session=session) + stream = pydap.client.open_url(url, session=session) - ############# The MERRA-2 server changes the pydap data retrieval format frequently between these two formats; so better to retain both of them rather than only using either one of them ############# - try: - q = ds['QV'].data[0][time_ind, ml_min:(ml_max + 1), lat_min_ind:(lat_max_ind + 1), lon_min_ind:(lon_max_ind + 1)][0] - p = ds['PL'].data[0][time_ind, ml_min:(ml_max + 1), lat_min_ind:(lat_max_ind + 1), lon_min_ind:(lon_max_ind + 1)][0] - t = ds['T'].data[0][time_ind, ml_min:(ml_max + 1), lat_min_ind:(lat_max_ind + 1), lon_min_ind:(lon_max_ind + 1)][0] - h = ds['H'].data[0][time_ind, ml_min:(ml_max + 1), lat_min_ind:(lat_max_ind + 1), lon_min_ind:(lon_max_ind + 1)][0] - except IndexError: - q = ds['QV'].data[time_ind, ml_min:(ml_max + 1), lat_min_ind:(lat_max_ind + 1), lon_min_ind:(lon_max_ind + 1)][0] - p = ds['PL'].data[time_ind, ml_min:(ml_max + 1), lat_min_ind:(lat_max_ind + 1), lon_min_ind:(lon_max_ind + 1)][0] - t = ds['T'].data[time_ind, ml_min:(ml_max + 1), lat_min_ind:(lat_max_ind + 1), lon_min_ind:(lon_max_ind + 1)][0] - h = ds['H'].data[time_ind, ml_min:(ml_max + 1), lat_min_ind:(lat_max_ind + 1), lon_min_ind:(lon_max_ind + 1)][0] - except AttributeError: - q = ds['QV'][time_ind, ml_min:(ml_max + 1), lat_min_ind:(lat_max_ind + 1), lon_min_ind:(lon_max_ind + 1)][0] - p = ds['PL'][time_ind, ml_min:(ml_max + 1), lat_min_ind:(lat_max_ind + 1), lon_min_ind:(lon_max_ind + 1)][0] - t = ds['T'][time_ind, ml_min:(ml_max + 1), lat_min_ind:(lat_max_ind + 1), lon_min_ind:(lon_max_ind + 1)][0] - h = ds['H'][time_ind, ml_min:(ml_max + 1), lat_min_ind:(lat_max_ind + 1), lon_min_ind:(lon_max_ind + 1)][0] - except BaseException: - logger.exception("MERRA-2: Unable to read weathermodel data") - ######################################################################################################################## + q = stream['QV'][0,:,lat_min_ind:lat_max_ind + 1, lon_min_ind:lon_max_ind + 1].data.squeeze() + p = stream['PL'][0,:,lat_min_ind:lat_max_ind + 1, lon_min_ind:lon_max_ind + 1].data.squeeze() + t = stream['T'][0,:,lat_min_ind:lat_max_ind + 1, lon_min_ind:lon_max_ind + 1].data.squeeze() + h = stream['H'][0,:,lat_min_ind:lat_max_ind + 1, lon_min_ind:lon_max_ind + 1].data.squeeze() try: - writeWeatherVars2NETCDF4(self, lats, lons, h, q, p, t, outName=out) + writeWeatherVarsXarray(lat, lon, h, q, p, t, time, self._proj, outName=out) except Exception as e: logger.debug(e) logger.exception("MERRA-2: Unable to save weathermodel to file") @@ -161,26 +148,19 @@ def _load_model_level(self, filename): ''' # adding the import here should become absolute when transition to netcdf - from netCDF4 import Dataset - with Dataset(filename, mode='r') as f: - lons = np.array(f.variables['x'][:]) - lats = np.array(f.variables['y'][:]) - h = np.array(f.variables['H'][:]) - q = np.array(f.variables['QV'][:]) - p = np.array(f.variables['PL'][:]) - t = np.array(f.variables['T'][:]) - - # restructure the 3-D lat/lon/h in regular grid - _lons = np.broadcast_to(lons[np.newaxis, np.newaxis, :], t.shape) - _lats = np.broadcast_to(lats[np.newaxis, :, np.newaxis], t.shape) + ds = xarray.load_dataset(filename) + lons = ds['longitude'].values + lats = ds['latitude'].values + h = ds['h'].values + q = ds['q'].values + p = ds['p'].values + t = ds['t'].values # Re-structure everything from (heights, lats, lons) to (lons, lats, heights) p = np.transpose(p) q = np.transpose(q) t = np.transpose(t) h = np.transpose(h) - _lats = np.transpose(_lats) - _lons = np.transpose(_lons) # check this # data cube format should be lats,lons,heights @@ -188,24 +168,19 @@ def _load_model_level(self, filename): q = q.swapaxes(0, 1) t = t.swapaxes(0, 1) h = h.swapaxes(0, 1) - _lats = _lats.swapaxes(0, 1) - _lons = _lons.swapaxes(0, 1) # For some reason z is opposite the others p = np.flip(p, axis=2) q = np.flip(q, axis=2) t = np.flip(t, axis=2) h = np.flip(h, axis=2) - _lats = np.flip(_lats, axis=2) - _lons = np.flip(_lons, axis=2) # assign the regular-grid (lat/lon/h) variables - self._p = p self._q = q self._t = t - self._lats = _lats - self._lons = _lons - self._xs = _lons - self._ys = _lats + self._lats = lats + self._lons = lons + self._xs = lons + self._ys = lats self._zs = h diff --git a/tools/RAiDER/models/ncmr.py b/tools/RAiDER/models/ncmr.py index 4f64984aa..01b9f62ad 100755 --- a/tools/RAiDER/models/ncmr.py +++ b/tools/RAiDER/models/ncmr.py @@ -13,9 +13,9 @@ from RAiDER.models.weatherModel import WeatherModel, TIME_RES from RAiDER.logger import logger from RAiDER.utilFcns import ( - writeWeatherVars2NETCDF4, read_NCMR_loginInfo, - show_progress + show_progress, + writeWeatherVarsXarray, ) from RAiDER.models.model_levels import ( LEVELS_137_HEIGHTS, @@ -174,7 +174,7 @@ def _download_ncmr_file(self, out, date_time, bounding_box): ######################################################################################################################## try: - writeWeatherVars2NETCDF4(self, lats, lons, hgt, q, p, t, outName=out) + writeWeatherVarsXarray(lats, lons, hgt, q, p, t, self._time, self._proj, outName=out) except Exception: logger.exception("Unable to save weathermodel to file") diff --git a/tools/RAiDER/models/weatherModel.py b/tools/RAiDER/models/weatherModel.py index 37165f07f..868ed5632 100755 --- a/tools/RAiDER/models/weatherModel.py +++ b/tools/RAiDER/models/weatherModel.py @@ -20,7 +20,7 @@ from RAiDER.models import plotWeather as plots, weatherModel from RAiDER.models.customExceptions import * from RAiDER.utilFcns import ( - robmax, robmin, write2NETCDF4core, calcgeoh, transform_coords, clip_bbox + robmax, robmin, writeWeatherVarsXarray, calcgeoh, transform_coords, clip_bbox ) TIME_RES = {'GMAO': 3, diff --git a/tools/RAiDER/s1_orbits.py b/tools/RAiDER/s1_orbits.py index 18b852f43..92659c3b0 100644 --- a/tools/RAiDER/s1_orbits.py +++ b/tools/RAiDER/s1_orbits.py @@ -3,12 +3,14 @@ import re from pathlib import Path from platform import system -from typing import List, Optional, Tuple +from typing import List, Optional import eof.download +from RAiDER.logger import logger ESA_CDSE_HOST = 'dataspace.copernicus.eu' +NASA_EDL_HOST = 'urs.earthdata.nasa.gov' def _netrc_path() -> Path: @@ -17,11 +19,13 @@ def _netrc_path() -> Path: def ensure_orbit_credentials() -> Optional[int]: - """Ensure credentials exist for ESA's CDSE to download orbits + """Ensure credentials exist for ESA's CDSE and ASF's S1QC to download orbits - This method will prefer to use CDSE credentials from your `~/.netrc` file if they exist, - otherwise will look for ESA_USERNAME and ESA_PASSWORD environment variables and - update or create your `~/.netrc` file. + This method will prefer to use CDSE and NASA Earthdata credentials from your `~/.netrc` file if they exist, + otherwise will look for environment variables and update or create your `~/.netrc` file. The environment variables + used are: + CDSE: ESA_USERNAME, ESA_PASSWORD + S1QC: EARTHDATA_USERNAME, EARTHDATA_PASSWORD Returns `None` if the `~/.netrc` file did not need to be updated and the number of characters written if it did. """ @@ -32,17 +36,29 @@ def ensure_orbit_credentials() -> Optional[int]: netrc_file.touch(mode=0o600) netrc_credentials = netrc.netrc(netrc_file) - if ESA_CDSE_HOST in netrc_credentials.hosts: + if ESA_CDSE_HOST in netrc_credentials.hosts and NASA_EDL_HOST in netrc_credentials.hosts: return - username = os.environ.get('ESA_USERNAME') - password = os.environ.get('ESA_PASSWORD') - if username is None and password is None: - raise ValueError('Credentials are required for fetching orbit data from dataspace.copernicus.eu!\n' - 'Either add your credentials to ~/.netrc or set the ESA_USERNAME and ESA_PASSWORD ' - 'environment variables.') + if ESA_CDSE_HOST not in netrc_credentials.hosts: + username = os.environ.get('ESA_USERNAME') + password = os.environ.get('ESA_PASSWORD') + if username is None or password is None: + raise ValueError('Credentials are required for fetching orbit data from dataspace.copernicus.eu!\n' + 'Either add your credentials to ~/.netrc or set the ESA_USERNAME and ESA_PASSWORD ' + 'environment variables.') + + netrc_credentials.hosts[ESA_CDSE_HOST] = (username, None, password) + + if NASA_EDL_HOST not in netrc_credentials.hosts: + username = os.environ.get('EARTHDATA_USERNAME') + password = os.environ.get('EARTHDATA_PASSWORD') + if username is None or password is None: + raise ValueError(f'Credentials are required for fetching orbit data from s1qc.asf.alaska.edu!\n' + 'Either add your credentials to ~/.netrc or set the EARTHDATA_USERNAME and' + ' EARTHDATA_PASSWORD environment variables.') + + netrc_credentials.hosts[NASA_EDL_HOST] = (username, None, password) - netrc_credentials.hosts[ESA_CDSE_HOST] = (username, None, password) return netrc_file.write_text(str(netrc_credentials)) @@ -53,12 +69,34 @@ def get_orbits_from_slc_ids(slc_ids: List[str], directory=Path.cwd()) -> List[Pa Returns a list of orbit file paths """ - _ = ensure_orbit_credentials() - missions = [slc_id[0:3] for slc_id in slc_ids] start_times = [re.split(r'_+', slc_id)[4] for slc_id in slc_ids] stop_times = [re.split(r'_+', slc_id)[5] for slc_id in slc_ids] - - orb_files = eof.download.download_eofs(start_times + stop_times, missions * 2, save_dir=str(directory)) + + orb_files = download_eofs(start_times + stop_times, missions * 2, str(directory)) + + return orb_files + + +def download_eofs(dts: list, missions: list, save_dir: str): + """Wrapper around sentineleof to first try downloading from ASF and fall back to CDSE""" + _ = ensure_orbit_credentials() + + orb_files = [] + for dt, mission in zip(dts, missions): + dt = dt if isinstance(dt, list) else [dt] + mission = mission if isinstance(mission, list) else [mission] + + try: + orb_file = eof.download.download_eofs(dt, mission, save_dir=save_dir, force_asf=True) + except: + logger.error(f'Could not download orbit from ASF, trying ESA...') + orb_file = eof.download.download_eofs(dt, mission, save_dir=save_dir, force_asf=False) + + orb_file = orb_file[0] if isinstance(orb_file, list) else orb_file + orb_files.append(orb_file) + + if not len(orb_files) == len(dts): + raise Exception(f'Missing {len(dts) - len(orb_files)} orbit files! dts={dts}, orb_files={len(orb_files)}') return orb_files diff --git a/tools/RAiDER/utilFcns.py b/tools/RAiDER/utilFcns.py index dce21673b..a8a7507cc 100755 --- a/tools/RAiDER/utilFcns.py +++ b/tools/RAiDER/utilFcns.py @@ -1,6 +1,7 @@ """Geodesy-related utility functions.""" import os import re +import xarray from datetime import datetime, timedelta from numpy import ndarray @@ -622,199 +623,54 @@ def requests_retry_session(retries=10, session=None): return session -def writeWeatherVars2NETCDF4(self, lat, lon, h, q, p, t, outName=None, NoDataValue=None, chunk=(1, 91, 144), mapping_name='WGS84'): - ''' - By calling the abstract/modular netcdf writer (RAiDER.utilFcns.write2NETCDF4core), write the OpenDAP/PyDAP-retrieved weather model data (GMAO and MERRA-2) to a NETCDF4 file - that can be accessed by external programs. - - The point of doing this is to alleviate some of the memory load of keeping - the full model in memory and make it easier to scale up the program. - ''' - - import netCDF4 - - if outName is None: - outName = os.path.join( - os.getcwd() + '/weather_files', - self._Name + datetime.strftime( - self._time, '_%Y_%m_%d_T%H_%M_%S' - ) + '.nc' - ) - - if NoDataValue is None: - NoDataValue = -9999. - - self._time = getTimeFromFile(outName) - - dimidZ, dimidY, dimidX = t.shape - chunk_lines_Y = np.min([chunk[1], dimidY]) - chunk_lines_X = np.min([chunk[2], dimidX]) - ChunkSize = [1, chunk_lines_Y, chunk_lines_X] - - nc_outfile = netCDF4.Dataset(outName, 'w', clobber=True, format='NETCDF4') - nc_outfile.setncattr('Conventions', 'CF-1.6') - nc_outfile.setncattr('datetime', datetime.strftime(self._time, "%Y_%m_%dT%H_%M_%S")) - nc_outfile.setncattr('date_created', datetime.now().strftime("%Y_%m_%dT%H_%M_%S")) - title = self._Name + ' weather model data' - nc_outfile.setncattr('title', title) - - tran = [lon[0], lon[1] - lon[0], 0.0, lat[0], 0.0, lat[1] - lat[0]] - +def writeWeatherVarsXarray(lat, lon, h, q, p, t, dt, crs, outName=None, NoDataValue=-9999, chunk=(1, 91, 144)): + + # I added datetime as an input to the function and just copied these two lines from merra2 for the attrs_dict + attrs_dict = { + 'datetime': dt.strftime("%Y_%m_%dT%H_%M_%S"), + 'date_created': datetime.now().strftime("%Y_%m_%dT%H_%M_%S"), + 'NoDataValue': NoDataValue, + 'chunksize': chunk, + # 'mapping_name': mapping_name, + } + dimension_dict = { - 'x': {'varname': 'x', - 'datatype': np.dtype('float64'), - 'dimensions': ('x'), - 'length': dimidX, - 'FillValue': None, - 'standard_name': 'longitude', - 'description': 'longitude', - 'dataset': lon, - 'units': 'degrees_east'}, - 'y': {'varname': 'y', - 'datatype': np.dtype('float64'), - 'dimensions': ('y'), - 'length': dimidY, - 'FillValue': None, - 'standard_name': 'latitude', - 'description': 'latitude', - 'dataset': lat, - 'units': 'degrees_north'}, - 'z': {'varname': 'z', - 'datatype': np.dtype('float32'), - 'dimensions': ('z'), - 'length': dimidZ, - 'FillValue': None, - 'standard_name': 'model_layers', - 'description': 'model layers', - 'dataset': np.arange(dimidZ), - 'units': 'layer'} + 'latitude': (('y', 'x'), lat), + 'longitude': (('y', 'x'), lon), } dataset_dict = { - 'h': {'varname': 'H', - 'datatype': np.dtype('float32'), - 'dimensions': ('z', 'y', 'x'), - 'grid_mapping': mapping_name, - 'FillValue': NoDataValue, - 'ChunkSize': ChunkSize, - 'standard_name': 'mid_layer_heights', - 'description': 'mid layer heights', - 'dataset': h, - 'units': 'm'}, - 'q': {'varname': 'QV', - 'datatype': np.dtype('float32'), - 'dimensions': ('z', 'y', 'x'), - 'grid_mapping': mapping_name, - 'FillValue': NoDataValue, - 'ChunkSize': ChunkSize, - 'standard_name': 'specific_humidity', - 'description': 'specific humidity', - 'dataset': q, - 'units': 'kg kg-1'}, - 'p': {'varname': 'PL', - 'datatype': np.dtype('float32'), - 'dimensions': ('z', 'y', 'x'), - 'grid_mapping': mapping_name, - 'FillValue': NoDataValue, - 'ChunkSize': ChunkSize, - 'standard_name': 'mid_level_pressure', - 'description': 'mid level pressure', - 'dataset': p, - 'units': 'Pa'}, - 't': {'varname': 'T', - 'datatype': np.dtype('float32'), - 'dimensions': ('z', 'y', 'x'), - 'grid_mapping': mapping_name, - 'FillValue': NoDataValue, - 'ChunkSize': ChunkSize, - 'standard_name': 'air_temperature', - 'description': 'air temperature', - 'dataset': t, - 'units': 'K'} + 'h': (('z', 'y', 'x'), h), + 'q': (('z', 'y', 'x'), q), + 'p': (('z', 'y', 'x'), p), + 't': (('z', 'y', 'x'), t), } - nc_outfile = write2NETCDF4core(nc_outfile, dimension_dict, dataset_dict, tran, mapping_name='WGS84') - - nc_outfile.sync() # flush data to disk - nc_outfile.close() - - -def write2NETCDF4core(nc_outfile, dimension_dict, dataset_dict, tran, mapping_name='WGS84'): - ''' - The abstract/modular netcdf writer that can be called by a wrapper function to write data to a NETCDF4 file - that can be accessed by external programs. - - The point of doing this is to alleviate some of the memory load of keeping - the full model in memory and make it easier to scale up the program. - ''' - from osgeo import osr - - if mapping_name == 'WGS84': - epsg = 4326 - crs = CRS.from_epsg(epsg) - - grid_mapping = 'WGS84' # need to set this as an attribute for the image variables - else: - crs = CRS.from_wkt(mapping_name) - grid_mapping = 'CRS' - - datatype = np.dtype('S1') - dimensions = () - var = nc_outfile.createVariable( - grid_mapping, - datatype, - dimensions, - fill_value=None - ) - - # variable made, now add attributes - for k, v in crs.to_cf().items(): - var.setncattr(k, v) - - var.setncattr('GeoTransform', ' '.join(str(x) for x in tran)) # note this has pixel size in it - set explicitly above - - for dim in dimension_dict: - nc_outfile.createDimension(dim, dimension_dict[dim]['length']) - varname = dimension_dict[dim]['varname'] - datatype = dimension_dict[dim]['datatype'] - dimensions = dimension_dict[dim]['dimensions'] - FillValue = dimension_dict[dim]['FillValue'] - var = nc_outfile.createVariable(varname, datatype, dimensions, fill_value=FillValue) - var.setncattr('standard_name', dimension_dict[dim]['standard_name']) - var.setncattr('description', dimension_dict[dim]['description']) - var.setncattr('units', dimension_dict[dim]['units']) - var[:] = dimension_dict[dim]['dataset'].astype(dimension_dict[dim]['datatype']) - - for data in dataset_dict: - varname = dataset_dict[data]['varname'] - datatype = dataset_dict[data]['datatype'] - dimensions = dataset_dict[data]['dimensions'] - FillValue = dataset_dict[data]['FillValue'] - ChunkSize = dataset_dict[data]['ChunkSize'] - var = nc_outfile.createVariable( - varname, - datatype, - dimensions, - fill_value=FillValue, - zlib=True, - complevel=2, - shuffle=True, - chunksizes=ChunkSize + ds = xarray.Dataset( + data_vars=dataset_dict, + coords=dimension_dict, + attrs=attrs_dict, ) - # Override with correct name here - var.setncattr('grid_mapping', grid_mapping) # dataset_dict[data]['grid_mapping']) - var.setncattr('standard_name', dataset_dict[data]['standard_name']) - var.setncattr('description', dataset_dict[data]['description']) - if 'units' in dataset_dict[data]: - var.setncattr('units', dataset_dict[data]['units']) - - ndmask = np.isnan(dataset_dict[data]['dataset']) - dataset_dict[data]['dataset'][ndmask] = FillValue - - var[:] = dataset_dict[data]['dataset'].astype(datatype) - - return nc_outfile + + ds['h'].attrs['standard_name'] = 'mid_layer_heights' + ds['p'].attrs['standard_name'] = 'mid_level_pressure' + ds['q'].attrs['standard_name'] = 'specific_humidity' + ds['t'].attrs['standard_name'] = 'air_temperature' + + ds['h'].attrs['units'] = 'm' + ds['p'].attrs['units'] = 'Pa' + ds['q'].attrs['units'] = 'kg kg-1' + ds['t'].attrs['units'] = 'K' + ds["proj"] = int() + for k, v in crs.to_cf().items(): + ds.proj.attrs[k] = v + for var in ds.data_vars: + ds[var].attrs['grid_mapping'] = 'proj' + + ds.to_netcdf(outName) + del ds + def convertLons(inLons): '''Convert lons from 0-360 to -180-180'''