def run_module():
    
    # read in configuration file to execute run
    print("Reading configuration from [%s]" % sys.argv[1])
    
    with open(sys.argv[1]) as f:
        cfg = eval(f.read())
    
    # ensure output path exists
    if not os.path.isdir(cfg['output_dir']): 
        os.mkdir(cfg['output_dir'])

    # configure diagnostics        
    init_diagnostics(os.path.join(cfg['output_dir'], 'moisture_model_v1_diagnostics.txt'))

    # Error covariance matrix condition number in kriging
    diagnostics().configure_tag("skdm_cov_cond", False, True, True)

    # Assimilation parameters
    diagnostics().configure_tag("assim_K0", False, True, True)
    diagnostics().configure_tag("assim_K1", True, True, True)
    diagnostics().configure_tag("assim_data", False, False, True)

    diagnostics().configure_tag("obs_residual_var", True, True, True)

    diagnostics().configure_tag("fm10_model_residual_var", True, True, True)
    diagnostics().configure_tag("fm10_model_var", False, True, True)
    diagnostics().configure_tag("fm10_kriging_var", False, True, True)

    ### Load and preprocess WRF model data

    # load WRF data
    wrf_data = WRFModelData(cfg['input_file'], tz_name = 'US/Mountain')
    
    # read in spatial and temporal extent of WRF variables
    lat, lon = wrf_data.get_lats(), wrf_data.get_lons()
    tm = wrf_data.get_gmt_times()
    Nt = cfg['Nt'] if cfg['Nt'] is not None else len(tm)
    dom_shape = lat.shape

    # retrieve the rain variable
    rain = wrf_data['RAIN']

    # moisture equilibria are now computed from averaged Q,P,T at beginning and end of period
    Ed, Ew = wrf_data.get_moisture_equilibria()

    ### Load observation data from the stations

    # load station data from files
    with open(os.path.join(cfg['station_data_dir'], cfg['station_list_file']), 'r') as f:
        si_list = f.read().split('\n')

    si_list = filter(lambda x: len(x) > 0, map(string.strip, si_list))

    # for each station id, load the station
    stations = []
    for sinfo in si_list:
        code = sinfo.split(',')[0]
        mws = MesoWestStation(sinfo, wrf_data)
        for suffix in [ '_1', '_2', '_3', '_4', '_5', '_6', '_7' ]:
            mws.load_station_data(os.path.join(cfg['station_data_dir'], '%s%s.xls' % (code, suffix)))
        stations.append(mws)

    print('Loaded %d stations.' % len(stations))
    
    # check stations for nans
    stations = filter(MesoWestStation.data_ok, stations)
    print('Have %d stations with complete data.' % len(stations))

    # set the measurement variance of the stations
    for s in stations:
        s.set_measurement_variance('fm10', cfg['fm10_meas_var'])

    # build the observation data
    obs_data_fm10 = build_observation_data(stations, 'fm10', wrf_data, tm)

    ### Initialize model and visualization

    # find maximum moisture overall to set up visualization
    maxE = 0.5
    
    # construct initial conditions from timestep 1 (because Ed/Ew at zero are zero)
    E = 0.5 * (Ed[1,:,:] + Ew[1,:,:])
    
    # set up parameters
    Q = np.eye(9) * cfg['Q']
    P0 = np.eye(9) * cfg['P0']
    dt = (tm[1] - tm[0]).seconds
    print("INFO: Computed timestep from WRF is is %g seconds." % dt)
    K = np.zeros_like(E)
    V = np.zeros_like(E)
    mV = np.zeros_like(E)
    predicted_field = np.zeros_like(E)
    mresV = np.zeros_like(E)
    Kf_fn = np.zeros_like(E)
    Vf_fn = np.zeros_like(E)
    mid = np.zeros_like(E)
    Kg = np.zeros((dom_shape[0], dom_shape[1], 9))
    cV12 = np.zeros_like(E)
    
    # moisture state and observation residual variance estimators
    mod_re = OnlineVarianceEstimator(np.zeros_like(E), np.ones_like(E) * 0.05, 1)
    obs_re = OnlineVarianceEstimator(np.zeros((len(stations),)), np.ones(len(stations),) * 0.05, 1)
    
    # initialize the mean field model (default fit is 1.0 of equilibrium before new information comes in)
    mfm = MeanFieldModel(cfg['lock_gamma'])

    # construct model grid using standard fuel parameters
    Tk = np.array([1.0, 10.0, 100.0]) * 3600
    models = np.zeros(dom_shape, dtype = np.object)
    models_na = np.zeros_like(models)
    for p in np.ndindex(dom_shape): 
        models[p] = CellMoistureModel((lat[p], lon[p]), 3, E[p], Tk, P0 = P0)
        models_na[p] = CellMoistureModel((lat[p], lon[p]), 3, E[p], Tk, P0 = P0)

    m = None
    plt.figure(figsize = (12, 8))
    
    ###  Run model for each WRF timestep and assimilate data when available
    for t in range(1, Nt):
        model_time = tm[t]
        print("INFO: time: %s, step: %d" % (str(model_time), t))

        # run the model update
        for p in np.ndindex(dom_shape):
            i, j = p
            models[p].advance_model(Ed[t-1, i, j], Ew[t-1, i, j], rain[t-1, i, j], dt, Q)
            models_na[p].advance_model(Ed[t-1, i, j], Ew[t-1, i, j], rain[t-1, i, j], dt, Q)
            
        # prepare visualization data
        f = np.zeros((dom_shape[0], dom_shape[1], 3))
        f_na = np.zeros((dom_shape[0], dom_shape[1], 3))
        for p in np.ndindex(dom_shape):
            f[p[0], p[1], :] = models[p].get_state()[:3]
            f_na[p[0], p[1], :] = models_na[p].get_state()[:3]
            P = models[p].get_state_covar()
            cV12[p] = P[0,1]
            mV[p] = P[1,1]
            mid[p] = models[p].get_model_ids()[1]

        diagnostics().push("fm10_model_var", (t, np.mean(mV)))

        # run Kriging on each observed fuel type
        Kf = []
        Vf = []
        fn = []
        for obs_data, fuel_ndx in [ (obs_data_fm10, 1) ]:

            # run the kriging subsystem and the Kalman update only if we have observations
            if model_time in obs_data:

                # retrieve observations for current time
                obs_t = obs_data[model_time]


                # fit the current estimation of the moisture field to the data 
                base_field = f[:,:,fuel_ndx]
                mfm.fit_to_data(base_field, obs_data[model_time])
                
                # find differences (residuals) between observed measurements and nearest grid points
                # use this to update observation residual standard deviation 
                obs_vals = np.array([o.get_value() for o in obs_data[model_time]])
                mod_vals = np.array([base_field[o.get_nearest_grid_point()] for o in obs_data[model_time]])
                mod_na_vals = np.array([f_na[:,:,fuel_ndx][o.get_nearest_grid_point()] for o in obs_data[model_time]])
                obs_re.update_with(obs_vals - mod_vals)
                diagnostics().push("obs_residual_var", (t, np.mean(obs_re.get_variance())))
            
                # predict the moisture field using observed fuel type
                predicted_field = mfm.predict_field(base_field)

                # update the model residual estimator and get current best estimate of variance
                mod_re.update_with(f[:,:,fuel_ndx] - predicted_field)
                mresV = mod_re.get_variance()
                diagnostics().push("fm10_model_residual_var", (t, np.mean(mresV)))

                # krige observations to grid points
                Kf_fn, Vf_fn = trend_surface_model_kriging(obs_data[model_time], wrf_data, predicted_field)

                krig_vals = np.array([Kf_fn[o.get_nearest_grid_point()] for o in obs_data[model_time]])                
                diagnostics().push("assim_data", (t, fuel_ndx, obs_vals, krig_vals, mod_vals, mod_na_vals))
                plot_model_snapshot(cfg, tm, t, fuel_ndx, obs_vals, krig_vals, mod_vals, mod_na_vals)

                diagnostics().push("fm10_kriging_var", (t, np.mean(Vf_fn)))

                # append to storage for kriged fields in this time instant
                Kf.append(Kf_fn)
                Vf.append(Vf_fn)
                fn.append(fuel_ndx)


        # if there were any observations, run the kalman update step
        if len(fn) > 0:
            Nobs = len(fn)
            # run the kalman update in each model independently
            # gather the standard deviations of the moisture fuel after the Kalman update
            for p in np.ndindex(dom_shape):
                O = np.zeros((Nobs,))
                V = np.zeros((Nobs, Nobs))
                
                # construct observations for this position
                for i in range(Nobs):
                    O[i] = Kf[i][p]
                    V[i,i] = Vf[i][p]
                
                # execute the Kalman update
                Kp = models[p].kalman_update(O, V, fn)
                Kg[p[0], p[1], :] = Kp[:, 0]

            # push new diagnostic outputs
            diagnostics().push("assim_K0", (t, np.mean(Kg[:,:,0])))
            diagnostics().push("assim_K1", (t, np.mean(Kg[:,:,1])))

        # prepare visualization data        
        f = np.zeros((dom_shape[0], dom_shape[1], 3))
        for p in np.ndindex(dom_shape):
            f[p[0], p[1], :] = models[p].get_state()[:3]
            
        plt.clf()
        plt.subplot(3,3,1)
        render_spatial_field_fast(m, lon, lat, f[:,:,0], '1-hr fuel')
        plt.clim([0.0, maxE])
        plt.axis('off')
        plt.colorbar()
        plt.subplot(3,3,2)
        render_spatial_field_fast(m, lon, lat, f[:,:,1], '10-hr fuel')
        plt.clim([0.0, maxE])        
        plt.axis('off')
        plt.colorbar()
        plt.subplot(3,3,3)
        render_spatial_field_fast(m, lon, lat, f_na[:,:,1], '10hr fuel - no assim')
        plt.clim([0.0, maxE])
        plt.axis('off')
        plt.colorbar()
        plt.subplot(3,3,4)
        render_spatial_field_fast(m, lon, lat, Kg[:,:,0], 'Kalman gain for 1-hr fuel')  
        plt.clim([0.0, 3.0])        
        plt.axis('off')
        plt.colorbar()
        plt.subplot(3,3,5)
        render_spatial_field_fast(m, lon, lat, Kg[:,:,1], 'Kalman gain for 10-hr fuel')       
        plt.clim([0.0, 1.0])        
        plt.axis('off')
        plt.colorbar()
        plt.subplot(3,3,6)
	render_spatial_field_fast(m, lon, lat, Kf_fn, 'Kriging field')
	plt.clim([0.0, maxE])
        plt.axis('off')
	plt.colorbar()
        plt.subplot(3,3,7)
        render_spatial_field_fast(m, lon, lat, mid, 'Model ids')
        plt.clim([0.0, 5.0])
        plt.axis('off')
        plt.colorbar()
        plt.subplot(3,3,8)
        render_spatial_field_fast(m, lon, lat, Vf_fn, 'Kriging variance')
        plt.clim([0.0, np.max(Vf_fn)])
        plt.axis('off')
        plt.colorbar()
        plt.subplot(3,3,9)
        render_spatial_field_fast(m, lon, lat, mresV, 'Model res. variance')
        plt.clim([0.0, np.max(mresV)])
        plt.axis('off')
        plt.colorbar()
        
        plt.savefig(os.path.join(cfg['output_dir'], 'moisture_model_t%03d.png' % t))


    # store the diagnostics in a binary file
    diagnostics().dump_store(os.path.join(cfg['output_dir'], 'diagnostics.bin'))
    
    # make a plot of gammas
    plt.figure()
    plt.plot(diagnostics().pull('mfm_gamma'), 'bo-')
    plt.title('Mean field model - gamma')
    plt.savefig(os.path.join(cfg['output_dir'], 'plot_gamma.png'))

    plt.figure()
    plt.plot(diagnostics().pull('skdm_cov_cond'))
    plt.title('Condition number of covariance matrix')
    plt.savefig(os.path.join(cfg['output_dir'], 'plot_sigma_cond.png'))

    # make a plot for each substation
    plt.figure()
    D = diagnostics().pull("assim_data")
    for i in range(len(stations)):
        plt.clf()
        # get data for the i-th station
        t_i = [ o[0] for o in D]
        obs_i = [ o[2][i] for o in D]
        krig_i = [ o[3][i] for o in D]
        mod_i = [ o[4][i] for o in D]
        mod_na_i = [ o[5][i] for o in D]
        mx = max(max(obs_i), max(mod_i), max(krig_i), max(mod_i))
        plt.plot(t_i, obs_i, 'ro')
        plt.plot(t_i, krig_i, 'bo-')
        plt.plot(t_i, mod_i, 'kx-', linewidth = 1.5)
        plt.plot(t_i, mod_na_i, 'mx-')
        plt.ylim([0.0, 1.1 * mx])
        plt.legend(['Obs.', 'Kriged', 'Model', 'NoAssim'])
        plt.title('Station observations fit to model and kriging field')
        plt.savefig(os.path.join(cfg['output_dir'], 'station%02d.png' % (i+1)))
        
    plt.figure()
    plt.plot([d[0] for d in diagnostics().pull("assim_K1")],
             [d[1] for d in diagnostics().pull("assim_K1")], 'ro-')
    plt.title('Average Kalman gain')
    plt.savefig(os.path.join(cfg['output_dir'], 'plot_kalman_gain_10hr.png'))
    
    plt.figure()
    plt.plot([d[0] for d in diagnostics().pull("assim_K0")],
             [d[1] for d in diagnostics().pull("assim_K0")], 'ro-')
    plt.title('Average Kalman gain')
    plt.savefig(os.path.join(cfg['output_dir'], 'plot_kalman_gain_1hr.png'))

    plt.figure()
    plt.plot([d[0] for d in diagnostics().pull("fm10_model_var")],
             [d[1] for d in diagnostics().pull("fm10_model_var")], 'ro-')
    plt.title('Average fm10 model variance')
    plt.savefig(os.path.join(cfg['output_dir'], 'plot_fm10_model_variance.png'))

    plt.figure()
    plt.plot([d[0] for d in diagnostics().pull("fm10_model_residual_var")],
             [d[1] for d in diagnostics().pull("fm10_model_residual_var")], 'ro-')
    plt.title('Average fm10 model residual variance')
    plt.savefig(os.path.join(cfg['output_dir'], 'plot_fm10_model_residual_variance.png'))

    plt.figure()
    plt.plot([d[0] for d in diagnostics().pull("fm10_kriging_var")],
             [d[1] for d in diagnostics().pull("fm10_kriging_var")], 'ro-')
    plt.title('Kriging field variance')
    plt.savefig(os.path.join(cfg['output_dir'], 'plot_kriging_variance.png'))

    plt.figure()
    plt.plot([d[0] for d in diagnostics().pull("obs_residual_var")],
             [d[1] for d in diagnostics().pull("obs_residual_var")], 'ro-')
    plt.title('Observation residual variance')
    plt.savefig(os.path.join(cfg['output_dir'], 'plot_observation_residual_variance.png'))
    
    plt.figure()
    plt.plot(diagnostics().pull("mfm_mape"), 'ro-', linewidth = 2)
    plt.title('Mean absolute prediction error of station data')
    plt.savefig(os.path.join(cfg['output_dir'], 'plot_station_mape.png'))
def run_module():
    
    # read in configuration file to execute run
    print("Reading configuration from [%s]" % sys.argv[1])
    
    with open(sys.argv[1]) as f:
        cfg = eval(f.read())
    
    # ensure output path exists
    if not os.path.isdir(cfg['output_dir']): 
        os.mkdir(cfg['output_dir'])
        
    # configure diagnostics        
    init_diagnostics(os.path.join(cfg['output_dir'], 'moisture_model_v1_diagnostics.txt'))
    diagnostics().configure_tag("skdm_obs_res", False, True, True)
    diagnostics().configure_tag("skdm_cov_cond", False, True, True)

    diagnostics().configure_tag("assim_mV", False, False, True)
    diagnostics().configure_tag("assim_K0", False, False, True)
    diagnostics().configure_tag("assim_K1", False, False, True)
    diagnostics().configure_tag("assim_data", False, False, True)
    diagnostics().configure_tag("assim_mresV", False, False, True)

    diagnostics().configure_tag("kriging_variance", False, False, True)
    diagnostics().configure_tag("kriging_obs_res_var", False, False, True)

    print("INFO: input file is [%s]." % cfg['input_file'])
    wrf_data = WRFModelData(cfg['input_file'], tz_name = 'US/Pacific')
    
    # read in vars
    lat, lon = wrf_data.get_lats(), wrf_data.get_lons()
    tm = wrf_data.get_local_times()
    rain = wrf_data['RAIN']
    Ed, Ew = wrf_data.get_moisture_equilibria()
    
    # find maximum moisture overall to set up visualization
#    maxE = max(np.max(Ed), np.max(Ew)) * 1.2
    maxE = 0.3
    
    # obtain sizes
    Nt = rain.shape[0]
    dom_shape = lat.shape
    
    # load station data from files
    tz = pytz.timezone('US/Pacific')
    stations = [StationAdam() for s in station_list]
    for (s,sname) in zip(stations, station_list):
        s.load_station_data(os.path.join(station_data_dir, sname), tz)
        s.register_to_grid(wrf_data)
        s.set_measurement_variance('fm10', 0.05)
    
    # build the observation data structure indexed by time
    obs_data_fm10 = build_observation_data(stations, 'fm10', wrf_data, tm)
    
    # construct initial conditions
    E = 0.5 * (Ed[1,:,:] + Ew[1,:,:])
    
    # set up parameters
    Q = np.eye(9) * 0.001
    P0 = np.eye(9) * 0.01
    dt = 10.0 * 60
    K = np.zeros_like(E)
    V = np.zeros_like(E)
    mV = np.zeros_like(E)
    predicted_field = np.zeros_like(E)
    mresV = np.zeros_like(E)
    Kf_fn = np.zeros_like(E)
    Vf_fn = np.zeros_like(E)
    mid = np.zeros_like(E)
    Kg = np.zeros((dom_shape[0], dom_shape[1], 9))
    cV12 = np.zeros_like(E)
    
    # initialize the mean field model (default fit is 1.0 of equilibrium before new information comes in)
    mfm = MeanFieldModel(cfg['lock_gamma'])

    # construct model grid using standard fuel parameters
    Tk = np.array([1.0, 10.0, 100.0]) * 3600
    models = np.zeros(dom_shape, dtype = np.object)
    models_na = np.zeros_like(models)
    for pos in np.ndindex(dom_shape): 
        models[pos] = CellMoistureModel((lat[pos], lon[pos]), 3, E[pos], Tk, P0 = P0)
        models_na[pos] = CellMoistureModel((lat[pos], lon[pos]), 3, E[pos], Tk, P0 = P0)

    m = None

    plt.figure(figsize = (12, 8))
    
    # run model
    for t in range(1, Nt):
        model_time = tm[t]
        print("Time: %s, step: %d" % (str(model_time), t))

        # pre-compute equilibrium moisture to save a lot of time
        E = 0.5 * (Ed[t,:,:] + Ew[t,:,:])
        
        # run the model update
        for pos in np.ndindex(dom_shape):
            i, j = pos
            models[pos].advance_model(Ed[t, i, j], Ew[t, i, j], rain[t, i, j], dt, Q)
            models_na[pos].advance_model(Ed[t, i, j], Ew[t, i, j], rain[t, i, j], dt, Q)
            
        # prepare visualization data        
        f = np.zeros((dom_shape[0], dom_shape[1], 3))
        f_na = np.zeros((dom_shape[0], dom_shape[1], 3))
        for p in np.ndindex(dom_shape):
            f[p[0], p[1], :] = models[p].get_state()[:3]
            f_na[p[0], p[1], :] = models_na[p].get_state()[:3]
            mV[pos] = models[p].get_state_covar()[1,1]
            cV12[pos] = models[p].get_state_covar()[0,1]
            mid[p] = models[p].get_model_ids()[1]
            

        # run Kriging on each observed fuel type
        Kf = []
        Vf = []
        fn = []
        for obs_data, fuel_ndx in [ (obs_data_fm10, 1) ]:

            if model_time in obs_data:

                # fit the current estimation of the moisture field to the data 
                base_field = f[:,:,fuel_ndx]
                mfm.fit_to_data(base_field, obs_data[model_time])
                
                # find differences (residuals) between observed measurements and nearest grid points
                # use this to update observation residual standard deviation 
                obs_vals = np.array([o.get_value() for o in obs_data[model_time]])
                mod_vals = np.array([f[:,:,fuel_ndx][o.get_nearest_grid_point()] for o in obs_data[model_time]])
                mod_na_vals = np.array([f_na[:,:,fuel_ndx][o.get_nearest_grid_point()] for o in obs_data[model_time]])
                obs_re.update_with(obs_vals - mod_vals)
                diagnostics().push("kriging_obs_res_var", (t, np.mean(obs_re.get_variance())))
            
                # retrieve the variance of the model field
                mresV = mod_re.get_variance()

                # krige data to observations
                if cfg['kriging_strategy'] == 'uk':
                    Kf_fn, Vf_fn, gamma, mape = universal_kriging_data_to_model(obs_data[model_time],
                                                                          obs_re.get_variance() ** 0.5,
                                                                          base_field,
                                                                          wrf_data,
                                                                          mresV ** 0.5, t)
                    # replace the stored gamma with the uk computed gamma
                    diagnostics().pull("mfm_gamma")[-1] = gamma
                    diagnostics().pull("mfm_mape")[-1] = mape
                    print("uk: replaced mfm_gamma %g, mfm_mape %g" % (gamma, mape))

                    # update the residuals estimator with the current
                    mod_re.update_with(gamma * f[:,:,fuel_ndx] - Kf_fn)

                elif cfg['kriging_strategy'] == 'tsm':
                    # predict the moisture field using observed fuel type
                    predicted_field = mfm.predict_field(base_field)

                    # run the tsm kriging estimator
                    Kf_fn, Vf_fn = trend_surface_model_kriging(obs_data[model_time], wrf_data, predicted_field)

                    # update the model residual estimator and get current best estimate of variance
                    mod_re.update_with(f[:,:,fuel_ndx] - predicted_field)

                else:
                    raise ValueError('Invalid kriging strategy [%s] in configuration.' % cfg['kriiging_strategy'])

                krig_vals = np.array([Kf_fn[o.get_nearest_grid_point()] for o in obs_data[model_time]])                
                diagnostics().push("assim_data", (t, fuel_ndx, obs_vals, krig_vals, mod_vals, mod_na_vals))
                plot_model_snapshot(cfg, tm, t, fuel_ndx, obs_vals, krig_vals, mod_vals, mod_na_vals)

                # append to storage for kriged fields in this time instant
                Kf.append(Kf_fn)
                Vf.append(Vf_fn)
                fn.append(fuel_ndx)

        # if there were any observations, run the kalman update step
        if len(fn) > 0:
            Nobs = len(fn)
            # run the kalman update in each model independently
            # gather the standard deviations of the moisture fuel after the Kalman update
            for pos in np.ndindex(dom_shape):
                O = np.zeros((Nobs,))
                V = np.zeros((Nobs, Nobs))
                
                # construct observations for this position
                for i in range(Nobs):
                    O[i] = Kf[i][pos]
                    V[i,i] = Vf[i][pos]
                
                # execute the Kalman update
                Kij = models[pos].kalman_update(O, V, fn)
                Kg[pos[0], pos[1], :] = Kij[:, 0]


        # prepare visualization data        
        f = np.zeros((dom_shape[0], dom_shape[1], 3))
        for p in np.ndindex(dom_shape):
            f[p[0], p[1], :] = models[p].get_state()[:3]
            
        plt.clf()
        plt.subplot(3,3,1)
        render_spatial_field_fast(m, lon, lat, f[:,:,0], '1-hr fuel')
        plt.clim([0.0, maxE])
        plt.axis('off')
        plt.colorbar()
        plt.subplot(3,3,2)
        render_spatial_field_fast(m, lon, lat, f[:,:,1], '10-hr fuel')
        plt.clim([0.0, maxE])        
        plt.axis('off')
        plt.colorbar()
        plt.subplot(3,3,3)
        render_spatial_field_fast(m, lon, lat, f_na[:,:,1], '10hr fuel - no assim')
        plt.clim([0.0, maxE])
        plt.axis('off')
        plt.colorbar()
        plt.subplot(3,3,4)
        render_spatial_field_fast(m, lon, lat, Kg[:,:,0], 'Kalman gain, fm1')  
        plt.clim([0.0, 3.0])
        plt.axis('off')
        plt.colorbar()
        plt.subplot(3,3,5)
        render_spatial_field_fast(m, lon, lat, Kg[:,:,1], 'Kalman gain, fm10')       
        plt.clim([0.0, 1.0])        
        plt.axis('off')
        plt.colorbar()
        plt.subplot(3,3,6)
	render_spatial_field_fast(m, lon, lat, Kf_fn, 'Kriging field')
	plt.clim([0.0, maxE])
        plt.axis('off')
	plt.colorbar()
        plt.subplot(3,3,7)
        render_spatial_field_fast(m, lon, lat, mid, 'Model ids')
        plt.clim([0.0, 5.0])
        plt.axis('off')
        plt.colorbar()
        plt.subplot(3,3,8)
        render_spatial_field_fast(m, lon, lat, Vf_fn, 'Kriging var')
        plt.clim([0.0, np.max(Vf_fn)])
        plt.axis('off')
        plt.colorbar()
        plt.subplot(3,3,9)
        render_spatial_field_fast(m, lon, lat, mresV, 'fm10 model var')
        plt.clim([0.0, np.max(mresV)])
        plt.axis('off')
        plt.colorbar()
        
        plt.savefig(os.path.join(cfg['output_dir'], 'moisture_model_t%03d.png' % t))

        # push new diagnostic outputs
        diagnostics().push("assim_K0", (t, np.mean(Kg[:,:,0])))
        diagnostics().push("assim_K1", (t, np.mean(Kg[:,:,1])))
        diagnostics().push("assim_mV", (t, np.mean(mV)))
        diagnostics().push("assim_mresV", (t, np.mean(mresV)))
        diagnostics().push("kriging_variance", (t, np.mean(Vf_fn)))

        
    # store the gamma coefficients
    with open(os.path.join(cfg['output_dir'], 'gamma.txt'), 'w') as f:
        f.write(str(diagnostics().pull('mfm_gamma')))
        
    # make a plot of gammas
    plt.figure()
    plt.plot(diagnostics().pull('mfm_gamma'))
    plt.title('Mean field model - gamma')
    plt.savefig(os.path.join(cfg['output_dir'], 'plot_gamma.png'))

    plt.figure()
    plt.plot(diagnostics().pull('skdm_cov_cond'))
    plt.title('Condition number of covariance matrix')
    plt.savefig(os.path.join(cfg['output_dir'], 'plot_sigma_cond.png'))

    # make 
    # make a plot for each substation
    plt.figure()
    D = diagnostics().pull("assim_data")
    for i in range(len(stations)):
        plt.clf()
        # get data for the i-th station
        t_i = [ o[0] for o in D]
        obs_i = [ o[2][i] for o in D]
        krig_i = [ o[3][i] for o in D]
        mod_i = [ o[4][i] for o in D]
        mod_na_i = [ o[5][i] for o in D]
        mx = max(max(obs_i), max(mod_i), max(krig_i), max(mod_i))
        plt.plot(t_i, obs_i, 'ro')
        plt.plot(t_i, krig_i, 'bo-')
        plt.plot(t_i, mod_i, 'kx-', linewidth = 1.5)
        plt.plot(t_i, mod_na_i, 'mx-')
        plt.ylim([0.0, 1.1 * mx])
        plt.legend(['Obs.', 'Kriged', 'Model', 'NoAssim'])
        plt.title('Station observations fit to model and kriging field')
        plt.savefig(os.path.join(cfg['output_dir'], 'station%02d.png' % (i+1)))
    
    
    plt.figure()
    plt.plot([d[0] for d in diagnostics().pull("assim_K1")],
             [d[1] for d in diagnostics().pull("assim_K1")], 'ro-')
    plt.title('Average Kalman gain')
    plt.savefig(os.path.join(cfg['output_dir'], 'plot_kalman_gain_10hr.png'))
    
    plt.figure()
    plt.plot([d[0] for d in diagnostics().pull("assim_K0")],
             [d[1] for d in diagnostics().pull("assim_K0")], 'ro-')
    plt.title('Average Kalman gain')
    plt.savefig(os.path.join(cfg['output_dir'], 'plot_kalman_gain_1hr.png'))

    plt.figure()
    plt.plot([d[0] for d in diagnostics().pull("assim_mV")],
             [d[1] for d in diagnostics().pull("assim_mV")], 'ro-')
    plt.title('Average model variance')
    plt.savefig(os.path.join(cfg['output_dir'], 'plot_fm10_model_variance.png'))

    plt.figure()
    plt.plot([d[0] for d in diagnostics().pull("assim_mresV")],
             [d[1] for d in diagnostics().pull("assim_mresV")], 'ro-')
    plt.title('Average fm10 residual variance')
    plt.savefig(os.path.join(cfg['output_dir'], 'plot_fm10_model_residual_variance.png'))

    plt.figure()
    plt.plot([d[0] for d in diagnostics().pull("kriging_variance")],
             [d[1] for d in diagnostics().pull("kriging_variance")], 'ro-')
    plt.title('Kriging field variance')
    plt.savefig(os.path.join(cfg['output_dir'], 'plot_kriging_variance.png'))

    plt.figure()
    plt.plot([d[0] for d in diagnostics().pull("kriging_obs_res_var")],
             [d[1] for d in diagnostics().pull("kriging_obs_res_var")], 'ro-')
    plt.title('Observation residual variance')
    plt.savefig(os.path.join(cfg['output_dir'], 'plot_observation_residual_variance.png'))
    
    plt.figure()
    plt.plot(diagnostics().pull("mfm_mape"), 'ro-', linewidth = 2)
    plt.title('Mean absolute prediction error of station data')
    plt.savefig(os.path.join(cfg['output_dir'], 'plot_station_mape.png'))

    diagnostics().dump_store(os.path.join(cfg['output_dir'], 'diagnostics.bin'))
    
    # as a last step encode all the frames as video
    os.system("cd %s; avconv -qscale 1 -r 20 -b 9600 -i moisture_model_t%%03d.png video.mp4" % cfg['output_dir'])
def run_module():

    # read in configuration file to execute run
    print("Reading configuration from [%s]" % sys.argv[1])

    with open(sys.argv[1]) as f:
        cfg = eval(f.read())

    # ensure output path exists
    if not os.path.isdir(cfg['output_dir']):
        os.mkdir(cfg['output_dir'])

    # configure diagnostics
    init_diagnostics(
        os.path.join(cfg['output_dir'], 'moisture_model_v1_diagnostics.txt'))

    # Error covariance matrix condition number in kriging
    diagnostics().configure_tag("skdm_cov_cond", False, True, True)

    # Assimilation parameters
    diagnostics().configure_tag("assim_K0", False, True, True)
    diagnostics().configure_tag("assim_K1", True, True, True)
    diagnostics().configure_tag("assim_data", False, False, True)

    diagnostics().configure_tag("fm10_model_var", False, True, True)
    diagnostics().configure_tag("fm10_kriging_var", False, True, True)

    ### Load and preprocess WRF model data

    # load WRF data
    wrf_data = WRFModelData(cfg['input_file'], tz_name='US/Mountain')

    # read in spatial and temporal extent of WRF variables
    lat, lon = wrf_data.get_lats(), wrf_data.get_lons()
    tm = wrf_data.get_gmt_times()
    Nt = cfg['Nt'] if cfg.has_key('Nt') and cfg['Nt'] is not None else len(tm)
    dom_shape = lat.shape

    # retrieve the rain variable
    rain = wrf_data['RAIN']

    # moisture equilibria are now computed from averaged Q,P,T at beginning and end of period
    Ed, Ew = wrf_data.get_moisture_equilibria()

    ### Load observation data from the stations

    # load station data from files
    with open(os.path.join(cfg['station_data_dir'], cfg['station_list_file']),
              'r') as f:
        si_list = f.read().split('\n')

    si_list = filter(lambda x: len(x) > 0 and x[0] != '#',
                     map(string.strip, si_list))

    # for each station id, load the station
    stations = []
    for code in si_list:
        mws = MesoWestStation(code)
        mws.load_station_info(
            os.path.join(cfg["station_data_dir"], "%s.info" % code))
        mws.register_to_grid(wrf_data)
        mws.load_station_data(
            os.path.join(cfg["station_data_dir"], "%s.obs" % code))
        stations.append(mws)

    print('Loaded %d stations.' % len(stations))

    # check stations for nans
    stations = filter(MesoWestStation.data_ok, stations)
    print('Have %d stations with complete data.' % len(stations))

    # build the observation data
    #    obs_data_fm10 = build_observation_data(stations, 'FM')
    obs_data_fm10 = {}

    ### Initialize model and visualization

    # find maximum moisture overall to set up visualization
    maxE = 0.5

    # construct initial conditions from timestep 1 (because Ed/Ew at zero are zero)
    E = 0.5 * (Ed[1, :, :] + Ew[1, :, :])

    # set up parameters
    Q = np.eye(9) * cfg['Q']
    P0 = np.eye(9) * cfg['P0']
    dt = (tm[1] - tm[0]).seconds
    print("INFO: Computed timestep from WRF is is %g seconds." % dt)
    K = np.zeros_like(E)
    V = np.zeros_like(E)
    mV = np.zeros_like(E)
    predicted_field = np.zeros_like(E)
    mresV = np.zeros_like(E)
    Kf_fn = np.zeros_like(E)
    Vf_fn = np.zeros_like(E)
    mid = np.zeros_like(E)
    Kg = np.zeros((dom_shape[0], dom_shape[1], 9))
    cV12 = np.zeros_like(E)

    # initialize the mean field model (default fit is 1.0 of equilibrium before new information comes in)
    mfm = MeanFieldModel(cfg['lock_gamma'])

    # construct model grid using standard fuel parameters
    Tk = np.array([1.0, 10.0, 100.0]) * 3600
    models = np.zeros(dom_shape, dtype=np.object)
    models_na = np.zeros_like(models)
    for p in np.ndindex(dom_shape):
        models[p] = CellMoistureModel((lat[p], lon[p]), 3, E[p], Tk, P0=P0)
        models_na[p] = CellMoistureModel((lat[p], lon[p]), 3, E[p], Tk, P0=P0)

    m = None
    plt.figure(figsize=(12, 8))

    ###  Run model for each WRF timestep and assimilate data when available
    for t in range(1, Nt):
        model_time = tm[t]
        print("INFO: time: %s, step: %d" % (str(model_time), t))

        # run the model update
        for p in np.ndindex(dom_shape):
            i, j = p
            models[p].advance_model(Ed[t - 1, i, j], Ew[t - 1, i, j],
                                    rain[t - 1, i, j], dt, Q)
            models_na[p].advance_model(Ed[t - 1, i, j], Ew[t - 1, i, j],
                                       rain[t - 1, i, j], dt, Q)

        # prepare visualization data
        f = np.zeros((dom_shape[0], dom_shape[1], 3))
        f_na = np.zeros((dom_shape[0], dom_shape[1], 3))
        for p in np.ndindex(dom_shape):
            f[p[0], p[1], :] = models[p].get_state()[:3]
            f_na[p[0], p[1], :] = models_na[p].get_state()[:3]
            P = models[p].get_state_covar()
            cV12[p] = P[0, 1]
            mV[p] = P[1, 1]
            mid[p] = models[p].get_model_ids()[1]

        diagnostics().push("fm10_model_var", (t, np.mean(mV)))

        # run Kriging on each observed fuel type
        Kf = []
        Vf = []
        fn = []
        for obs_data, fuel_ndx in [(obs_data_fm10, 1)]:

            # run the kriging subsystem and the Kalman update only if we have observations
            if model_time in obs_data:

                # retrieve observations for current time
                obs_t = obs_data[model_time]

                # fit the current estimation of the moisture field to the data
                base_field = f[:, :, fuel_ndx]
                mfm.fit_to_data(base_field, obs_data[model_time])

                # find differences (residuals) between observed measurements and nearest grid points
                # use this to update observation residual standard deviation
                obs_vals = np.array(
                    [o.get_value() for o in obs_data[model_time]])
                mod_vals = np.array([
                    base_field[o.get_nearest_grid_point()]
                    for o in obs_data[model_time]
                ])
                mod_na_vals = np.array([
                    f_na[:, :, fuel_ndx][o.get_nearest_grid_point()]
                    for o in obs_data[model_time]
                ])

                # predict the moisture field using observed fuel type
                predicted_field = mfm.predict_field(base_field)

                # krige observations to grid points
                Kf_fn, Vf_fn = trend_surface_model_kriging(
                    obs_data[model_time], wrf_data, predicted_field)

                krig_vals = np.array([
                    Kf_fn[o.get_nearest_grid_point()]
                    for o in obs_data[model_time]
                ])
                diagnostics().push(
                    "assim_data",
                    (t, fuel_ndx, obs_vals, krig_vals, mod_vals, mod_na_vals))
                plot_model_snapshot(cfg, tm, t, fuel_ndx, obs_vals, krig_vals,
                                    mod_vals, mod_na_vals)

                diagnostics().push("fm10_kriging_var", (t, np.mean(Vf_fn)))

                # append to storage for kriged fields in this time instant
                Kf.append(Kf_fn)
                Vf.append(Vf_fn)
                fn.append(fuel_ndx)

        # if there were any observations, run the kalman update step
        if len(fn) > 0:
            Nobs = len(fn)
            # run the kalman update in each model independently
            # gather the standard deviations of the moisture fuel after the Kalman update
            for p in np.ndindex(dom_shape):
                O = np.zeros((Nobs, ))
                V = np.zeros((Nobs, Nobs))

                # construct observations for this position
                for i in range(Nobs):
                    O[i] = Kf[i][p]
                    V[i, i] = Vf[i][p]

                # execute the Kalman update
                Kp = models[p].kalman_update(O, V, fn)
                Kg[p[0], p[1], :] = Kp[:, 0]

            # push new diagnostic outputs
            diagnostics().push("assim_K0", (t, np.mean(Kg[:, :, 0])))
            diagnostics().push("assim_K1", (t, np.mean(Kg[:, :, 1])))

        # prepare visualization data
        f = np.zeros((dom_shape[0], dom_shape[1], 3))
        for p in np.ndindex(dom_shape):
            f[p[0], p[1], :] = models[p].get_state()[:3]

        plt.clf()
        plt.subplot(3, 3, 1)
        render_spatial_field_fast(m, lon, lat, f[:, :, 0], '1-hr fuel')
        plt.clim([0.0, maxE])
        plt.axis('off')
        plt.colorbar()
        plt.subplot(3, 3, 2)
        render_spatial_field_fast(m, lon, lat, f[:, :, 1], '10-hr fuel')
        plt.clim([0.0, maxE])
        plt.axis('off')
        plt.colorbar()
        plt.subplot(3, 3, 3)
        render_spatial_field_fast(m, lon, lat, f_na[:, :, 1],
                                  '10hr fuel - no assim')
        plt.clim([0.0, maxE])
        plt.axis('off')
        plt.colorbar()
        plt.subplot(3, 3, 4)
        render_spatial_field_fast(m, lon, lat, Kg[:, :, 0],
                                  'Kalman gain for 1-hr fuel')
        plt.clim([0.0, 3.0])
        plt.axis('off')
        plt.colorbar()
        plt.subplot(3, 3, 5)
        render_spatial_field_fast(m, lon, lat, Kg[:, :, 1],
                                  'Kalman gain for 10-hr fuel')
        plt.clim([0.0, 1.0])
        plt.axis('off')
        plt.colorbar()
        plt.subplot(3, 3, 6)
        render_spatial_field_fast(m, lon, lat, Kf_fn, 'Kriging field')
        plt.clim([0.0, maxE])
        plt.axis('off')
        plt.colorbar()
        plt.subplot(3, 3, 7)
        render_spatial_field_fast(m, lon, lat, mid, 'Model ids')
        plt.clim([0.0, 5.0])
        plt.axis('off')
        plt.colorbar()
        plt.subplot(3, 3, 8)
        render_spatial_field_fast(m, lon, lat, Vf_fn, 'Kriging variance')
        plt.clim([0.0, np.max(Vf_fn)])
        plt.axis('off')
        plt.colorbar()
        #        plt.subplot(3,3,9)
        #        render_spatial_field_fast(m, lon, lat, mresV, 'Model variance')
        #        plt.clim([0.0, np.max(mresV)])
        #       plt.axis('off')
        #        plt.colorbar()

        plt.savefig(
            os.path.join(cfg['output_dir'], 'moisture_model_t%03d.png' % t))

    # store the diagnostics in a binary file
    diagnostics().dump_store(os.path.join(cfg['output_dir'],
                                          'diagnostics.bin'))

    # make a plot of gammas
    plt.figure()
    plt.plot(diagnostics().pull('mfm_gamma'), 'bo-')
    plt.title('Mean field model - gamma')
    plt.savefig(os.path.join(cfg['output_dir'], 'plot_gamma.png'))

    plt.figure()
    plt.plot(diagnostics().pull('skdm_cov_cond'))
    plt.title('Condition number of covariance matrix')
    plt.savefig(os.path.join(cfg['output_dir'], 'plot_sigma_cond.png'))

    # make a plot for each substation
    plt.figure()
    D = diagnostics().pull("assim_data")
    for i in range(len(stations)):
        plt.clf()
        # get data for the i-th station
        t_i = [o[0] for o in D]
        obs_i = [o[2][i] for o in D]
        krig_i = [o[3][i] for o in D]
        mod_i = [o[4][i] for o in D]
        mod_na_i = [o[5][i] for o in D]
        mx = max(max(obs_i), max(mod_i), max(krig_i), max(mod_i))
        plt.plot(t_i, obs_i, 'ro')
        plt.plot(t_i, krig_i, 'bo-')
        plt.plot(t_i, mod_i, 'kx-', linewidth=1.5)
        plt.plot(t_i, mod_na_i, 'mx-')
        plt.ylim([0.0, 1.1 * mx])
        plt.legend(['Obs.', 'Kriged', 'Model', 'NoAssim'])
        plt.title('Station observations fit to model and kriging field')
        plt.savefig(
            os.path.join(cfg['output_dir'], 'station%02d.png' % (i + 1)))

    plt.figure()
    plt.plot([d[0] for d in diagnostics().pull("assim_K1")],
             [d[1] for d in diagnostics().pull("assim_K1")], 'ro-')
    plt.title('Average Kalman gain')
    plt.savefig(os.path.join(cfg['output_dir'], 'plot_kalman_gain_10hr.png'))

    plt.figure()
    plt.plot([d[0] for d in diagnostics().pull("assim_K0")],
             [d[1] for d in diagnostics().pull("assim_K0")], 'ro-')
    plt.title('Average Kalman gain')
    plt.savefig(os.path.join(cfg['output_dir'], 'plot_kalman_gain_1hr.png'))

    plt.figure()
    plt.plot([d[0] for d in diagnostics().pull("fm10_model_var")],
             [d[1] for d in diagnostics().pull("fm10_model_var")], 'ro-')
    plt.title('Average fm10 model variance')
    plt.savefig(os.path.join(cfg['output_dir'],
                             'plot_fm10_model_variance.png'))

    plt.figure()
    plt.plot([d[0] for d in diagnostics().pull("fm10_model_residual_var")],
             [d[1] for d in diagnostics().pull("fm10_model_residual_var")],
             'ro-')
    plt.title('Average fm10 model residual variance')
    plt.savefig(
        os.path.join(cfg['output_dir'],
                     'plot_fm10_model_residual_variance.png'))

    plt.figure()
    plt.plot([d[0] for d in diagnostics().pull("fm10_kriging_var")],
             [d[1] for d in diagnostics().pull("fm10_kriging_var")], 'ro-')
    plt.title('Kriging field variance')
    plt.savefig(os.path.join(cfg['output_dir'], 'plot_kriging_variance.png'))

    plt.figure()
    plt.plot([d[0] for d in diagnostics().pull("obs_residual_var")],
             [d[1] for d in diagnostics().pull("obs_residual_var")], 'ro-')
    plt.title('Observation residual variance')
    plt.savefig(
        os.path.join(cfg['output_dir'],
                     'plot_observation_residual_variance.png'))

    plt.figure()
    plt.plot(diagnostics().pull("mfm_mape"), 'ro-', linewidth=2)
    plt.title('Mean absolute prediction error of station data')
    plt.savefig(os.path.join(cfg['output_dir'], 'plot_station_mape.png'))
def run_module():

    # read in configuration file to execute run
    print("Reading configuration from [%s]" % sys.argv[1])

    with open(sys.argv[1]) as f:
        cfg = eval(f.read())

    # ensure output path exists
    if not os.path.isdir(cfg['output_dir']):
        os.mkdir(cfg['output_dir'])

    # configure diagnostics
    init_diagnostics(
        os.path.join(cfg['output_dir'], 'moisture_model_v1_diagnostics.txt'))
    diagnostics().configure_tag("skdm_obs_res", False, True, True)
    diagnostics().configure_tag("skdm_cov_cond", False, True, True)

    diagnostics().configure_tag("assim_mV", False, False, True)
    diagnostics().configure_tag("assim_K0", False, False, True)
    diagnostics().configure_tag("assim_K1", False, False, True)
    diagnostics().configure_tag("assim_data", False, False, True)
    diagnostics().configure_tag("assim_mresV", False, False, True)

    diagnostics().configure_tag("kriging_variance", False, False, True)
    diagnostics().configure_tag("kriging_obs_res_var", False, False, True)

    print("INFO: input file is [%s]." % cfg['input_file'])
    wrf_data = WRFModelData(cfg['input_file'], tz_name='US/Pacific')

    # read in vars
    lat, lon = wrf_data.get_lats(), wrf_data.get_lons()
    tm = wrf_data.get_local_times()
    rain = wrf_data['RAIN']
    Ed, Ew = wrf_data.get_moisture_equilibria()

    # find maximum moisture overall to set up visualization
    #    maxE = max(np.max(Ed), np.max(Ew)) * 1.2
    maxE = 0.3

    # obtain sizes
    Nt = rain.shape[0]
    dom_shape = lat.shape

    # load station data from files
    tz = pytz.timezone('US/Pacific')
    stations = [StationAdam() for s in station_list]
    for (s, sname) in zip(stations, station_list):
        s.load_station_data(os.path.join(station_data_dir, sname), tz)
        s.register_to_grid(wrf_data)
        s.set_measurement_variance('fm10', 0.05)

    # build the observation data structure indexed by time
    obs_data_fm10 = build_observation_data(stations, 'fm10', wrf_data, tm)

    # construct initial conditions
    E = 0.5 * (Ed[1, :, :] + Ew[1, :, :])

    # set up parameters
    Q = np.eye(9) * 0.001
    P0 = np.eye(9) * 0.01
    dt = 10.0 * 60
    K = np.zeros_like(E)
    V = np.zeros_like(E)
    mV = np.zeros_like(E)
    predicted_field = np.zeros_like(E)
    mresV = np.zeros_like(E)
    Kf_fn = np.zeros_like(E)
    Vf_fn = np.zeros_like(E)
    mid = np.zeros_like(E)
    Kg = np.zeros((dom_shape[0], dom_shape[1], 9))
    cV12 = np.zeros_like(E)

    # initialize the mean field model (default fit is 1.0 of equilibrium before new information comes in)
    mfm = MeanFieldModel(cfg['lock_gamma'])

    # construct model grid using standard fuel parameters
    Tk = np.array([1.0, 10.0, 100.0]) * 3600
    models = np.zeros(dom_shape, dtype=np.object)
    models_na = np.zeros_like(models)
    for pos in np.ndindex(dom_shape):
        models[pos] = CellMoistureModel((lat[pos], lon[pos]),
                                        3,
                                        E[pos],
                                        Tk,
                                        P0=P0)
        models_na[pos] = CellMoistureModel((lat[pos], lon[pos]),
                                           3,
                                           E[pos],
                                           Tk,
                                           P0=P0)

    m = None

    plt.figure(figsize=(12, 8))

    # run model
    for t in range(1, Nt):
        model_time = tm[t]
        print("Time: %s, step: %d" % (str(model_time), t))

        # pre-compute equilibrium moisture to save a lot of time
        E = 0.5 * (Ed[t, :, :] + Ew[t, :, :])

        # run the model update
        for pos in np.ndindex(dom_shape):
            i, j = pos
            models[pos].advance_model(Ed[t, i, j], Ew[t, i, j], rain[t, i, j],
                                      dt, Q)
            models_na[pos].advance_model(Ed[t, i, j], Ew[t, i, j],
                                         rain[t, i, j], dt, Q)

        # prepare visualization data
        f = np.zeros((dom_shape[0], dom_shape[1], 3))
        f_na = np.zeros((dom_shape[0], dom_shape[1], 3))
        for p in np.ndindex(dom_shape):
            f[p[0], p[1], :] = models[p].get_state()[:3]
            f_na[p[0], p[1], :] = models_na[p].get_state()[:3]
            mV[pos] = models[p].get_state_covar()[1, 1]
            cV12[pos] = models[p].get_state_covar()[0, 1]
            mid[p] = models[p].get_model_ids()[1]

        # run Kriging on each observed fuel type
        Kf = []
        Vf = []
        fn = []
        for obs_data, fuel_ndx in [(obs_data_fm10, 1)]:

            if model_time in obs_data:

                # fit the current estimation of the moisture field to the data
                base_field = f[:, :, fuel_ndx]
                mfm.fit_to_data(base_field, obs_data[model_time])

                # find differences (residuals) between observed measurements and nearest grid points
                # use this to update observation residual standard deviation
                obs_vals = np.array(
                    [o.get_value() for o in obs_data[model_time]])
                mod_vals = np.array([
                    f[:, :, fuel_ndx][o.get_nearest_grid_point()]
                    for o in obs_data[model_time]
                ])
                mod_na_vals = np.array([
                    f_na[:, :, fuel_ndx][o.get_nearest_grid_point()]
                    for o in obs_data[model_time]
                ])
                obs_re.update_with(obs_vals - mod_vals)
                diagnostics().push("kriging_obs_res_var",
                                   (t, np.mean(obs_re.get_variance())))

                # retrieve the variance of the model field
                mresV = mod_re.get_variance()

                # krige data to observations
                if cfg['kriging_strategy'] == 'uk':
                    Kf_fn, Vf_fn, gamma, mape = universal_kriging_data_to_model(
                        obs_data[model_time],
                        obs_re.get_variance()**0.5, base_field, wrf_data,
                        mresV**0.5, t)
                    # replace the stored gamma with the uk computed gamma
                    diagnostics().pull("mfm_gamma")[-1] = gamma
                    diagnostics().pull("mfm_mape")[-1] = mape
                    print("uk: replaced mfm_gamma %g, mfm_mape %g" %
                          (gamma, mape))

                    # update the residuals estimator with the current
                    mod_re.update_with(gamma * f[:, :, fuel_ndx] - Kf_fn)

                elif cfg['kriging_strategy'] == 'tsm':
                    # predict the moisture field using observed fuel type
                    predicted_field = mfm.predict_field(base_field)

                    # run the tsm kriging estimator
                    Kf_fn, Vf_fn = trend_surface_model_kriging(
                        obs_data[model_time], wrf_data, predicted_field)

                    # update the model residual estimator and get current best estimate of variance
                    mod_re.update_with(f[:, :, fuel_ndx] - predicted_field)

                else:
                    raise ValueError(
                        'Invalid kriging strategy [%s] in configuration.' %
                        cfg['kriiging_strategy'])

                krig_vals = np.array([
                    Kf_fn[o.get_nearest_grid_point()]
                    for o in obs_data[model_time]
                ])
                diagnostics().push(
                    "assim_data",
                    (t, fuel_ndx, obs_vals, krig_vals, mod_vals, mod_na_vals))
                plot_model_snapshot(cfg, tm, t, fuel_ndx, obs_vals, krig_vals,
                                    mod_vals, mod_na_vals)

                # append to storage for kriged fields in this time instant
                Kf.append(Kf_fn)
                Vf.append(Vf_fn)
                fn.append(fuel_ndx)

        # if there were any observations, run the kalman update step
        if len(fn) > 0:
            Nobs = len(fn)
            # run the kalman update in each model independently
            # gather the standard deviations of the moisture fuel after the Kalman update
            for pos in np.ndindex(dom_shape):
                O = np.zeros((Nobs, ))
                V = np.zeros((Nobs, Nobs))

                # construct observations for this position
                for i in range(Nobs):
                    O[i] = Kf[i][pos]
                    V[i, i] = Vf[i][pos]

                # execute the Kalman update
                Kij = models[pos].kalman_update(O, V, fn)
                Kg[pos[0], pos[1], :] = Kij[:, 0]

        # prepare visualization data
        f = np.zeros((dom_shape[0], dom_shape[1], 3))
        for p in np.ndindex(dom_shape):
            f[p[0], p[1], :] = models[p].get_state()[:3]

        plt.clf()
        plt.subplot(3, 3, 1)
        render_spatial_field_fast(m, lon, lat, f[:, :, 0], '1-hr fuel')
        plt.clim([0.0, maxE])
        plt.axis('off')
        plt.colorbar()
        plt.subplot(3, 3, 2)
        render_spatial_field_fast(m, lon, lat, f[:, :, 1], '10-hr fuel')
        plt.clim([0.0, maxE])
        plt.axis('off')
        plt.colorbar()
        plt.subplot(3, 3, 3)
        render_spatial_field_fast(m, lon, lat, f_na[:, :, 1],
                                  '10hr fuel - no assim')
        plt.clim([0.0, maxE])
        plt.axis('off')
        plt.colorbar()
        plt.subplot(3, 3, 4)
        render_spatial_field_fast(m, lon, lat, Kg[:, :, 0], 'Kalman gain, fm1')
        plt.clim([0.0, 3.0])
        plt.axis('off')
        plt.colorbar()
        plt.subplot(3, 3, 5)
        render_spatial_field_fast(m, lon, lat, Kg[:, :, 1],
                                  'Kalman gain, fm10')
        plt.clim([0.0, 1.0])
        plt.axis('off')
        plt.colorbar()
        plt.subplot(3, 3, 6)
        render_spatial_field_fast(m, lon, lat, Kf_fn, 'Kriging field')
        plt.clim([0.0, maxE])
        plt.axis('off')
        plt.colorbar()
        plt.subplot(3, 3, 7)
        render_spatial_field_fast(m, lon, lat, mid, 'Model ids')
        plt.clim([0.0, 5.0])
        plt.axis('off')
        plt.colorbar()
        plt.subplot(3, 3, 8)
        render_spatial_field_fast(m, lon, lat, Vf_fn, 'Kriging var')
        plt.clim([0.0, np.max(Vf_fn)])
        plt.axis('off')
        plt.colorbar()
        plt.subplot(3, 3, 9)
        render_spatial_field_fast(m, lon, lat, mresV, 'fm10 model var')
        plt.clim([0.0, np.max(mresV)])
        plt.axis('off')
        plt.colorbar()

        plt.savefig(
            os.path.join(cfg['output_dir'], 'moisture_model_t%03d.png' % t))

        # push new diagnostic outputs
        diagnostics().push("assim_K0", (t, np.mean(Kg[:, :, 0])))
        diagnostics().push("assim_K1", (t, np.mean(Kg[:, :, 1])))
        diagnostics().push("assim_mV", (t, np.mean(mV)))
        diagnostics().push("assim_mresV", (t, np.mean(mresV)))
        diagnostics().push("kriging_variance", (t, np.mean(Vf_fn)))

    # store the gamma coefficients
    with open(os.path.join(cfg['output_dir'], 'gamma.txt'), 'w') as f:
        f.write(str(diagnostics().pull('mfm_gamma')))

    # make a plot of gammas
    plt.figure()
    plt.plot(diagnostics().pull('mfm_gamma'))
    plt.title('Mean field model - gamma')
    plt.savefig(os.path.join(cfg['output_dir'], 'plot_gamma.png'))

    plt.figure()
    plt.plot(diagnostics().pull('skdm_cov_cond'))
    plt.title('Condition number of covariance matrix')
    plt.savefig(os.path.join(cfg['output_dir'], 'plot_sigma_cond.png'))

    # make
    # make a plot for each substation
    plt.figure()
    D = diagnostics().pull("assim_data")
    for i in range(len(stations)):
        plt.clf()
        # get data for the i-th station
        t_i = [o[0] for o in D]
        obs_i = [o[2][i] for o in D]
        krig_i = [o[3][i] for o in D]
        mod_i = [o[4][i] for o in D]
        mod_na_i = [o[5][i] for o in D]
        mx = max(max(obs_i), max(mod_i), max(krig_i), max(mod_i))
        plt.plot(t_i, obs_i, 'ro')
        plt.plot(t_i, krig_i, 'bo-')
        plt.plot(t_i, mod_i, 'kx-', linewidth=1.5)
        plt.plot(t_i, mod_na_i, 'mx-')
        plt.ylim([0.0, 1.1 * mx])
        plt.legend(['Obs.', 'Kriged', 'Model', 'NoAssim'])
        plt.title('Station observations fit to model and kriging field')
        plt.savefig(
            os.path.join(cfg['output_dir'], 'station%02d.png' % (i + 1)))

    plt.figure()
    plt.plot([d[0] for d in diagnostics().pull("assim_K1")],
             [d[1] for d in diagnostics().pull("assim_K1")], 'ro-')
    plt.title('Average Kalman gain')
    plt.savefig(os.path.join(cfg['output_dir'], 'plot_kalman_gain_10hr.png'))

    plt.figure()
    plt.plot([d[0] for d in diagnostics().pull("assim_K0")],
             [d[1] for d in diagnostics().pull("assim_K0")], 'ro-')
    plt.title('Average Kalman gain')
    plt.savefig(os.path.join(cfg['output_dir'], 'plot_kalman_gain_1hr.png'))

    plt.figure()
    plt.plot([d[0] for d in diagnostics().pull("assim_mV")],
             [d[1] for d in diagnostics().pull("assim_mV")], 'ro-')
    plt.title('Average model variance')
    plt.savefig(os.path.join(cfg['output_dir'],
                             'plot_fm10_model_variance.png'))

    plt.figure()
    plt.plot([d[0] for d in diagnostics().pull("assim_mresV")],
             [d[1] for d in diagnostics().pull("assim_mresV")], 'ro-')
    plt.title('Average fm10 residual variance')
    plt.savefig(
        os.path.join(cfg['output_dir'],
                     'plot_fm10_model_residual_variance.png'))

    plt.figure()
    plt.plot([d[0] for d in diagnostics().pull("kriging_variance")],
             [d[1] for d in diagnostics().pull("kriging_variance")], 'ro-')
    plt.title('Kriging field variance')
    plt.savefig(os.path.join(cfg['output_dir'], 'plot_kriging_variance.png'))

    plt.figure()
    plt.plot([d[0] for d in diagnostics().pull("kriging_obs_res_var")],
             [d[1] for d in diagnostics().pull("kriging_obs_res_var")], 'ro-')
    plt.title('Observation residual variance')
    plt.savefig(
        os.path.join(cfg['output_dir'],
                     'plot_observation_residual_variance.png'))

    plt.figure()
    plt.plot(diagnostics().pull("mfm_mape"), 'ro-', linewidth=2)
    plt.title('Mean absolute prediction error of station data')
    plt.savefig(os.path.join(cfg['output_dir'], 'plot_station_mape.png'))

    diagnostics().dump_store(os.path.join(cfg['output_dir'],
                                          'diagnostics.bin'))

    # as a last step encode all the frames as video
    os.system(
        "cd %s; avconv -qscale 1 -r 20 -b 9600 -i moisture_model_t%%03d.png video.mp4"
        % cfg['output_dir'])