update_river_data.py 3.3 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11
import numpy as np
import datetime as dt
import glob as gb
import netCDF4 as nc
import pickle as pk
import sys

import fvcom_river as fr


wrf_forecast_out_dir = sys.argv[1]
12 13
wrf_forecast_out_dir_strfmt = sys.argv[2]
end_date = dt.datetime.strptime(sys.argv[3],'%Y-%m-%d')
14
no_miss_loops = 4
15 16 17

# Load the river model
with open('river_model.pk1','rb') as f:
18
    river_dict = pk.load(f)
19 20 21

start_date = end_date
for this_river in river_dict.values():
22 23
    if hasattr(this_river, 'catchment_precipitation'):
        this_river_update = np.max(this_river.catchment_precipitation[0])
24 25
        if this_river_update < start_date:
            start_date = this_river_update
26

27

28
if start_date == end_date:
29 30 31 32 33 34
    print('Already up to date')

else:    

    missing_dates = np.asarray([start_date + dt.timedelta(days=int(this_ind)) for this_ind in np.arange(0, (end_date - start_date).days + 1)])

35
    for this_missing_loop in np.arange(-1, no_miss_loops):
36 37
        new_missing_dates = []
        for this_date in missing_dates:
38
            if this_missing_loop >= 0:
39 40 41 42
                print('Trying again to fill for {}'.format(this_date))
            else:
                print(this_date)
            this_date_m1 = this_date - dt.timedelta(days=int(this_missing_loop))
43 44
            this_date_str = this_date_m1.strftime(wrf_forecast_out_dir_strfmt)
            potential_files = gb.glob('{}/*{}*/wrfout_d03*'.format(wrf_forecast_out_dir, this_date_str))
45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75

            try:
                this_wrf_nc = nc.Dataset(potential_files[-1], 'r')
                wrf_date_str_raw = this_wrf_nc.variables['Times'][:]
                wrf_date_str = np.asarray([b''.join(this_str) for this_str in wrf_date_str_raw])
                wrf_dt = np.asarray([dt.datetime.strptime(this_str.decode('utf-8'),'%Y-%m-%d_%H:%M:%S') for this_str in wrf_date_str])
                wrf_dt_date = np.asarray([this_dt.date() for this_dt in wrf_dt]) 

                date_match = wrf_dt_date == this_date.date()

                forecast_data = {'times': wrf_dt[date_match], 'RAINNC': this_wrf_nc.variables['RAINNC'][date_match,:,:],
                                    'T2': this_wrf_nc.variables['T2'][date_match,:,:]}
                this_wrf_nc.close()

                for this_river_name, this_river in river_dict.items():
                    if hasattr(this_river, 'addToSeries'):
                        this_rain = np.sum(np.sum(forecast_data['RAINNC']*this_river.wrf_catchment_factors, axis=2), axis=1)
                        this_river.addToSeries('catchment_precipitation', this_rain, forecast_data['times'], override=True)

                        this_temp = np.zeros(len(forecast_data['times']))
                        for i in range(0, len(forecast_data['times'])):
                            this_temp[i] = np.average(forecast_data['T2'][i,:,:], weights=this_river.wrf_catchment_factors)
                        this_river.addToSeries('catchment_temp', this_temp, forecast_data['times'], override=True)
                        if hasattr(this_river, 'river_obj'):
                            this_river.catchment_precipitation = this_river.river_obj.catchment_precipitation
            except:
                new_missing_dates.append(this_date)
            missing_dates = new_missing_dates[:]            

    with open('river_model.pk1','wb') as f:
        pk.dump(river_dict, f, pk.HIGHEST_PROTOCOL)
76 77