| 1 | """ |
|---|
| 2 | Filename: io.py |
|---|
| 3 | Author: Shannon Mason, shannon.mason@ecmwf.int |
|---|
| 4 | Description: I/O functions for loading IFS data |
|---|
| 5 | """ |
|---|
| 6 | |
|---|
| 7 | #For loading and handling netCDF data |
|---|
| 8 | import xarray as xr |
|---|
| 9 | |
|---|
| 10 | def load_inputs(input_srcfile): |
|---|
| 11 | """ |
|---|
| 12 | Load ecRAD input files. |
|---|
| 13 | """ |
|---|
| 14 | |
|---|
| 15 | with xr.open_dataset(input_srcfile) as input_dset: |
|---|
| 16 | if 'lat' in input_dset: |
|---|
| 17 | input_dset = input_dset.rename({'lat':'latitude', 'lon':'longitude'}) |
|---|
| 18 | |
|---|
| 19 | input_dset['pressure_fl'] = xr.DataArray( |
|---|
| 20 | input_dset.pressure_hl[:,:-1] + 0.5*input_dset.pressure_hl.diff('half_level'), |
|---|
| 21 | coords={'column':input_dset.column, 'level':input_dset.level}, |
|---|
| 22 | dims=['column', 'level'] |
|---|
| 23 | ) |
|---|
| 24 | |
|---|
| 25 | #input_dset['p'] = input_dset.pressure_fl.median('column') |
|---|
| 26 | #input_dset['phl'] = input_dset.pressure_hl.median('column') |
|---|
| 27 | |
|---|
| 28 | if 'aerosol_mmr' in input_dset.data_vars: |
|---|
| 29 | input_dset['sea_salt'] = input_dset.aerosol_mmr.sel( |
|---|
| 30 | aerosol_type=slice(0,2)).sum('aerosol_type') |
|---|
| 31 | input_dset['dust'] = input_dset.aerosol_mmr.sel( |
|---|
| 32 | aerosol_type=slice(3,5)).sum('aerosol_type') |
|---|
| 33 | input_dset['organics'] = input_dset.aerosol_mmr.sel( |
|---|
| 34 | aerosol_type=slice(6,7)).sum('aerosol_type') |
|---|
| 35 | input_dset['black_carbon'] = input_dset.aerosol_mmr.sel( |
|---|
| 36 | aerosol_type=slice(8,9)).sum('aerosol_type') |
|---|
| 37 | input_dset['sulphate'] = input_dset.aerosol_mmr.sel(aerosol_type=10) |
|---|
| 38 | |
|---|
| 39 | input_dset = input_dset.set_coords(['latitude', 'longitude']) #, 'p', 'phl']) |
|---|
| 40 | |
|---|
| 41 | return input_dset.swap_dims({'column':'latitude'})# 'half_level':'phl', 'level':'p'}) |
|---|
| 42 | |
|---|
| 43 | |
|---|
| 44 | def load_ecRAD(output_srcfile, input_srcfile): |
|---|
| 45 | """ |
|---|
| 46 | Load an ecRAD output file and merge with inputs. |
|---|
| 47 | Calculate net fluxes and heating rates, including necessary pressures at full levels and half |
|---|
| 48 | levels. |
|---|
| 49 | """ |
|---|
| 50 | |
|---|
| 51 | with xr.open_dataset(output_srcfile) as rad_output, xr.open_dataset(input_srcfile) as ifs_input: |
|---|
| 52 | |
|---|
| 53 | if 'lat' in ifs_input: |
|---|
| 54 | ifs_input = ifs_input.rename({'lat':'latitude', 'lon':'longitude'}) |
|---|
| 55 | |
|---|
| 56 | input_dset = rad_output.merge(ifs_input) #, compat='override') |
|---|
| 57 | |
|---|
| 58 | input_dset['pressure_fl'] = xr.DataArray( |
|---|
| 59 | input_dset.pressure_hl[:,:-1] + 0.5*input_dset.pressure_hl.diff('half_level'), |
|---|
| 60 | coords={'column':input_dset.column, 'level':input_dset.level}, |
|---|
| 61 | dims=['column', 'level'] |
|---|
| 62 | ) |
|---|
| 63 | |
|---|
| 64 | #input_dset['p'] = input_dset.pressure_fl.median('column') |
|---|
| 65 | #input_dset['phl'] = input_dset.pressure_hl.median('column') |
|---|
| 66 | |
|---|
| 67 | input_dset = input_dset.set_coords(['latitude', 'longitude']) #, 'p', 'phl']) |
|---|
| 68 | |
|---|
| 69 | input_dset['cloud_radiative_effect_lw'] = ( |
|---|
| 70 | (input_dset.flux_dn_lw - input_dset.flux_dn_lw_clear) -\ |
|---|
| 71 | (input_dset.flux_up_lw - input_dset.flux_up_lw_clear) |
|---|
| 72 | ) |
|---|
| 73 | input_dset['cloud_radiative_effect_sw'] = ( |
|---|
| 74 | (input_dset.flux_dn_sw - input_dset.flux_dn_sw_clear) -\ |
|---|
| 75 | (input_dset.flux_up_sw - input_dset.flux_up_sw_clear) |
|---|
| 76 | ) |
|---|
| 77 | input_dset['flux_net_lw'] = input_dset.flux_dn_lw - input_dset.flux_up_lw |
|---|
| 78 | input_dset['flux_net_sw'] = input_dset.flux_dn_sw - input_dset.flux_up_sw |
|---|
| 79 | |
|---|
| 80 | input_dset['flux_net_lw_clear'] = input_dset.flux_dn_lw_clear - input_dset.flux_up_lw_clear |
|---|
| 81 | input_dset['flux_net_sw_clear'] = input_dset.flux_dn_sw_clear - input_dset.flux_up_sw_clear |
|---|
| 82 | |
|---|
| 83 | c_factor = 24*3600*(9.81/1004.) |
|---|
| 84 | input_dset['heating_rate_lw'] = (-1*c_factor*input_dset.flux_net_lw.diff('half_level')\ |
|---|
| 85 | /input_dset.pressure_hl.diff('half_level')).rename({'half_level':'level'}) |
|---|
| 86 | input_dset['heating_rate_sw'] = (-1*c_factor*input_dset.flux_net_sw.diff('half_level')\ |
|---|
| 87 | /input_dset.pressure_hl.diff('half_level')).rename({'half_level':'level'}) |
|---|
| 88 | |
|---|
| 89 | input_dset['heating_rate_lw_clear'] = \ |
|---|
| 90 | (-1*c_factor*input_dset.flux_net_lw_clear.diff('half_level')\ |
|---|
| 91 | /input_dset.pressure_hl.diff('half_level')).rename({'half_level':'level'}) |
|---|
| 92 | input_dset['heating_rate_sw_clear'] = \ |
|---|
| 93 | (-1*c_factor*input_dset.flux_net_sw_clear.diff('half_level')\ |
|---|
| 94 | /input_dset.pressure_hl.diff('half_level')).rename({'half_level':'level'}) |
|---|
| 95 | |
|---|
| 96 | return input_dset.swap_dims({'column':'latitude'}) # 'half_level':'phl', 'level':'p'}) |
|---|
