Added assimilate me code

notebooks/create_map.py

@@ -5,10 +5,11 @@
 import json
 import datetime
 import pandas as pd
+from pathlib import Path
 from pygeotile.tile import Tile
 from shapely import geometry
 from bqplot import Lines, Figure, LinearScale, DateScale, Axis, Boxplot
-from ipywidgets import Dropdown, FloatSlider, HBox, VBox, Layout, Label, jslink, Layout, SelectionSlider, Play, Tab, Box
+from ipywidgets import Dropdown, FloatSlider, HBox, VBox, Layout, Label, jslink, Layout, SelectionSlider, Play, Tab, Box, Button
 from ipyleaflet import Map, WidgetControl, LayersControl, ImageOverlay, GeoJSON, Marker, Icon
 from ipywidgets import Image as widgetIMG
 
@@ -17,6 +18,7 @@
 from map_utils import get_lai_gif, get_pixel, get_field_bounds, da_pix, get_lai_color_bar
 from map_utils import debounce
 from wofost_utils import create_ensemble, wofost_parameter_sweep_func, get_era5_gee
+from wofost_utils import ensemble_assimilation
 
 from ipywidgets import Image as ImageWidget
 
@@ -219,19 +221,7 @@
 
 # from ipywidgets import IntSlider
 
-# def read_wofost_data(fname):
-#     f = np.load(fname, allow_pickle=True)
-#     parameters = f.f.parameters
-#     t_axis = f.f.t_axis
-#     samples = f.f.samples
-#     lais = f.f.lais
-#     yields = f.f.yields
-#     DVS = f.f.DVS
-#     print("loading simulations")
-#     doys = [int(datetime.datetime.utcfromtimestamp(i.tolist()/1e9).strftime('%j')) for i in t_axis]
-#     return parameters, t_axis, samples, lais, yields, DVS, doys
 
-# parameters, t_axis, samples, lais, yields, DVS, simu_doys = read_wofost_data('data/wofost_sims_dvs150.npz')
 
 # k_slider1 = IntSlider(min=181, max=224, value=200,        # Opacity is valid in [0,1] range
 #                orientation='horizontal',       # Vertical slider is what we want
@@ -325,6 +315,9 @@
 lat = 8.20
 year = 2021
 
+# Read ensemble
+param_array, sim_times, sim_lai, sim_yields, sim_doys = read_wofost_data()
+
 @debounce(0.2)
 def on_change_wofost_slider(change):
 
@@ -340,14 +333,22 @@ def on_change_wofost_slider(change):
         paras_to_overwrite = [i for i in paras if 'AMAXTB_' not in i]
         for para in paras_to_overwrite:    
             ens_parameters[para] = wofost_sliders_dict[para].value
-        ens_parameters['AMAXTB'] = [0, 55.0, 1.5, wofost_sliders_dict['AMAXTB_150'].value]
-        df = wofost_parameter_sweep_func(year, ens_parameters = ens_parameters.copy(), meteo=meteo_file)
+        ens_parameters['AMAXTB'] = [0, wofost_sliders_dict['AMAXTB_000'].value,
+                                   1.25, wofost_sliders_dict['AMAXTB_125'].value,
+                                   1.50, wofost_sliders_dict['AMAXTB_150'].value,
+                                   2.0, wofost_sliders_dict['AMAXTB_200'].value
+                                   ]
+        print(ens_parameters)
+        df = wofost_parameter_sweep_func(year, 
+                                         ens_parameters = ens_parameters.copy(), 
+                                         meteo=meteo_file)
         dates = df.index
         doys = [int(i.strftime('%j')) for i in dates]
 
         # wofost_lai_fig.x = doys
         # wofost_lai_fig.y = np.array(df.LAI)
-        wofost_out_paras = ['DVS', 'LAI', 'TAGP', 'TWSO', 'TWLV', 'TWST', 'TWRT', 'TRA', 'RD', 'SM', 'WWLOW']
+        wofost_out_paras = ['DVS', 'LAI', 'TAGP', 'TWSO', 'TWLV', 'TWST',
+                            'TWRT', 'TRA', 'RD', 'SM', 'WWLOW']
         for wofost_out_para in wofost_out_paras:
             if wofost_out_para != 'TWSO':
                 wofost_out_dict[wofost_out_para].marks[0].x = doys
@@ -366,11 +367,73 @@ def on_change_wofost_slider(change):
         # lai_fig.axes[0].scale = LinearScale(max=365.0, min=180.0)
 
 prior_df = pd.read_csv('data/par_prior_maize_tropical-C.csv')
-paras = ['TDWI', 'TSUM1', 'RGRLAI', 'SDOY', 'SPAN', 'AMAXTB_150']
+paras = prior_df["#PARAM_CODE"].str.strip().tolist()
 all_paras, para_mins, para_maxs = np.array(prior_df.loc[:, ['#PARAM_CODE', 'Min', 'Max']]).T
 para_inds = [all_paras.tolist().index(i) for i in paras]
 wofost_sliders = []
 
+
+def read_wofost_data():
+    """Reads ensemble from JASMIN"""
+    lat, lon = my_map.center
+    lat, lon = (lat // 0.1) * 0.1, (lon // 0.1) * 0.1
+    print(lat, lon, year)
+    f = create_ensemble(lat, lon, year)
+    max_lai = np.nanmax(f.f.LAI, axis=1)
+    y = f.f.Yields.astype(float)
+    lai = f.f.LAI.astype(float)
+    param_names = paras
+
+    param_array = np.array([f[i]
+                            for i in param_names]).astype(float)
+
+    pred_yield = max_lai * 1500 - 700
+    passer = np.abs(pred_yield - y) < 100.
+    sim_yields = y[passer]
+    sim_lai = lai[passer, :]
+    param_array = param_array[:, passer]
+    sim_times = f.f.sim_times
+
+    doys = [int(x.strftime("%j")) for x in sim_times]
+    return param_array, sim_times, sim_lai, sim_yields, doys
+
+
+
+
+# Button needs to be centered
+assimilate_me_button = Button(
+    description='Assimilate Me!',
+    disabled=False,
+    button_style='danger', 
+    tooltip='Click me',
+    icon='check',
+    layout = Layout(display='flex',
+                    flex_flow='horizontal',
+                    align_items='center',
+                    justify_content="center",
+                    width='30%'))
+
+def assimilate_me(b):
+    # Needs obs LAI & obs LAI times
+    global pix_lai, doys
+    print(pix_lai)
+    t_axis = np.array([datetime.datetime.strptime(f"{year}/{x}", "%Y/%j").date()
+              for x in doys])
+    est_yield, est_yield_sd, parameters, _, _ , _, lai_fitted_ensembles = ensemble_assimilation(
+        param_array, sim_times, sim_lai, sim_yields, pix_lai, t_axis)
+    # est_yield: mean estimated yield for this LAI set of observations
+    # est_yield_sd: standard deviation for the yield estimate
+    # parameters: list with parameters of selected ensemble members
+    # lai_fitted_ensembles: list with selected LAI simulations
+    # TOOD:
+    # 1. Update sliders with ?mean/median parameters?
+    # 2. Probably run wofost with solution parameters (or just plot ensemble?)
+    # 3. Update plots of LAI and TWSO
+
+# Link assimilate button to function
+assimilate_me_button.on_click(assimilate_me)    
+
+
 for i in range(len(paras)):
     para_ind = para_inds[i]
     para_name = all_paras[para_ind]
@@ -480,11 +543,9 @@ def update_wofost_fig_val(wofost_out_para, x, y, xmin = 180, xmax = 330):
     # description="Field ID:",
 )
 
-
-
 wofost_output_dropdowns = HBox([wofost_output_dropdown1, wofost_output_dropdown2], layout = Layout(display='flex',
                                   flex_flow='horizontal',
-                                  align_items='flex-start',
+                                  align_items='center',
                                   width='100%'))
 
 
@@ -538,8 +599,10 @@ def on_change_dropdown2(change):
                                   width='400px'))
     tab.children = [panel_box, wofost_box]
 
+
+
 wofost_out_panel = HBox([left_output, right_output])
-wofost_widgets = wofost_sliders + [wofost_out_panel]
+wofost_widgets = wofost_sliders + [assimilate_me_button, wofost_out_panel]
 wofost_output_dropdown1.observe(on_change_dropdown1, 'value')
 wofost_output_dropdown2.observe(on_change_dropdown2, 'value')
 
@@ -1093,16 +1156,21 @@ def handle_interaction(**kwargs):
             pix_cab, pix_lai = mean_bios
 
             max_ind = np.argmax(pix_lai)
-
             field_max = lais[:, max_ind].max()
             field_min = lais[:, max_ind].min() 
 
             lai_ratio = (pix_lai.max() - field_min) / (field_max - field_min)
             lai_ratio = (pix_lai.max() - 0) / (field_max - 0)
             # print(field_max, field_min)
             # print(lai_ratio)
+
 
-            icon = Icon(icon_url='https://leafletjs.com/examples/custom-icons/leaf-green.png', icon_size=[(38*lai_ratio), int(95*lai_ratio)], icon_anchor=[22,94])
+            try:
+                icon = Icon(icon_url='data/leaf-green.png',
+                            icon_size=[(38*lai_ratio), int(95*lai_ratio)], icon_anchor=[22,94])
+            except ValueError:
+                icon = Icon(icon_url='https://leafletjs.com/SlavaUkraini/examples//custom-icons/leaf-green.png',
+                            icon_size=[(38*lai_ratio), int(95*lai_ratio)], icon_anchor=[22,94])
             maize_icon = Icon(icon_url='https://gws-access.jasmin.ac.uk/public/nceo_ard/Ghana/Ghana_workshop2022/imgs/maize.png', 
                         icon_size=[36.5*lai_ratio, 98.5*lai_ratio], 
                         icon_anchor=[36.5/2*lai_ratio, 98.5*lai_ratio])

notebooks/data/par_prior_maize_tropical-C.csv

@@ -1,11 +1,9 @@
 #PARAM_CODE,Variation,PARAM_XVALUE,PARAM_YVALUE,Min,Max,StdDev,Scale,Distribution
-TDWI,S,-1,2,0.4,7,1,1,NA
-TSUM1,S,-1,1825,1350,2060,20,1,NA
-TSUM2,S,-1,1210,1030,1390,10,1,NA
-LAIEM,S,-1,0.02,0.003,0.005,0.001,1,NA
-RGRLAI,S,-1,0.02,0.005,0.006,0.1,1,NA
-WAV,S,-1,20,1,45,10,1,NA
-SDOY,S,-1,175,150,200,10,1,Gaussian
-SPAN,S,-1,35,24,48,4,1,Gaussian
-AMAXTB_000,YA,0,55,52,86.4,5.7,1,Gaussian
-AMAXTB_150,YA,1.5,53,44,78,5.7,1,Gaussian
+TDWI,S,-1,20,2,30,5,1,NA
+SDOY,S,-1,195,175,210,10,1,NA
+SPAN,S,-1,35,30,45,8,1,Gaussian
+CVO,S,-1,0.5,0.4,0.56,0.15,1,Gaussian
+AMAXTB_000,YA,0,65,55,85,20,1,Gaussian
+AMAXTB_125,YA,1.25,45,40,55,10,1,Gaussian
+AMAXTB_150,YA,1.5,20,10,30,10,1,Gaussian
+AMAXTB_200,YA,2.0,2.0,0,5,2,1,Gaussian

notebooks/python/wofost_utils.py

@@ -424,7 +424,7 @@ def create_ensemble(
     lat,
     lon,
     year,
-    en_size=20000,
+    en_size=50,
     param_file="data/par_prior_maize_tropical-C.csv",
     cropfile="data/MAIZGA-C.CAB",
     soil="data/ec4.new",
@@ -521,11 +521,11 @@ def ensemble_assimilation(
     sim_yields,
     obs_lai,
     obs_lai_time,
-    sigma_lai=None,
+    sigma_lai=0.05,
     obs_yield = None,
     sigma_yield=1.,
-    sel_n_best=50,
-    cost_prior = 0.
+    sel_n_best=20,
+    fit_tail_end=False
 ):
     """A function that performs ensemble assimilation. Requires:
     * parameters (n_params, n_ens) set of model parameters
@@ -541,80 +541,67 @@ def ensemble_assimilation(
     * sel_n_best (int) Select the best N simulations.
     * cost_prior (n_ens) You can also add a prior cost to each
     ensemble member.
+
+    Returns
+    -------
+    est_yield: estimated yield
+    est_yield_sd: estiamted yield std dev
+    parameters: list of parameters for selectede ensemble members
+    obs_dates: dates of the observations
+    obs_lai_sub: lai measurements
+    work_sim_times: simulation times that match observations
+    fit_lai: Simulated LAI predictions to match observations.
     """
     n_par, n_ens = parameters.shape
     cost_lai = np.zeros(n_ens)
+    sim_lai = sim_lai.astype(float)
+    sim_lai[np.isnan(sim_lai)] = 0.
     if obs_yield is not None:
         cost_yield = np.zeros_like(cost_lai)
 
-    #####obs_dates = obs_lai_time
-    #####for iens in range(n_ens):
-        #####sim_timex = pd.date_range(sim_times[iens][0],
-                                  #####sim_times[iens][1]
-                                  #####)
-        #####passer = obs_lai_time <= sim_timex.max()
-        #####obs_dates = obs_lai_time[passer]
-        #####obs_lai_sub = obs_lai[passer]
-        #####doys = (obs_dates -
-                #####sim_timex.to_pydatetime()[0].date())
-
-        #####time_index = [x.days for x in doys]
-        #####work_sim_times = sim_timex[time_index]
-        #####diffs = np.array(sim_lai[iens])[time_index] - obs_lai_sub.squeeze()
-        #####cost_lai[iens] = np.nansum(diffs*diffs)
     # Get observations that match the simulation period
     passer = obs_lai_time<= sim_times.max()
     obs_dates = obs_lai_time[passer]
     obs_lai = obs_lai[passer]
+    if fit_tail_end:
+        iloc = np.argmax(obs_lai)-5
+        jloc = np.argmax(obs_lai)+5
+
+
+        obs_dates = obs_dates[iloc:jloc]
+        obs_lai_sub = obs_lai[iloc:jloc]
+    else:
+        obs_lai_sub = obs_lai
     # We need a pointer that matches times in the
     # dense simulations array to the observations
     doys = obs_dates - sim_times[0]
     time_index = [x.days for x in doys]
 
-    diffs = sim_lai[:, time_index] - obs_lai.squeeze()
-    if sigma_lai is None:
-        cost_lai = np.nansum(diffs*diffs, axis=1)
-    else:
-        cost_lai = np.nansum(-0.5*diffs**2/(sigma_lai*sigma_lai),
-                             axis=1)
+    work_sim_times = sim_times[time_index]
+
+    # work sims has the simulations on the same time axis
+    # as the observations now.
+
+    diffs = sim_lai[:, time_index] - obs_lai_sub.squeeze()
 
-    posterior = cost_prior + cost_lai
+    cost_lai = np.nansum(diffs*diffs, axis=1)
+
+    posterior = cost_lai
     if obs_yield is not None:
         cost_yield= 0.5 * ((sim_yields - obs_yield) ** 2 / (sigma_yield ** 2))
         posterior += cost_yield
 
     ilocs = posterior.argsort()[:sel_n_best] # best 20?
 
+
+    eposterior = np.exp(-posterior.astype(float))
     sim_yields = np.array(sim_yields)
-    parameters=np.array(parameters)
+    parameters = np.array(parameters)
     est_yield = np.nanmean(sim_yields[ilocs])
+    #est_yield = np.average(sim_yields, weights=eposterior)
     est_yield_sd = np.nanstd(sim_yields[ilocs])
-    parameters = np.nanmean(parameters[:, ilocs],
-                            axis=1)
-
-    return est_yield, est_yield_sd, parameters
-
-
-########def wofost_parameter_sweep():
-########widgets.interact_manual(wofost_parameter_sweep_func,
-########field_code=widgets.Dropdown(options=all_fields),
-########crop_start_date=widgets.fixed(dt.date(2021, 7, 25)),
-########crop_end_date=widgets.fixed(dt.date(2021, 12, 1)),
-########span=widgets.FloatSlider(value=35.0, min=10, max=50),
-########cvo = widgets.FloatSlider(value=0.54, min=0.1, max=0.9, step=0.02),
-########cvl = widgets.FloatSlider(value=0.68, min=0.1, max=0.9, step=0.02),
-########tdwi=widgets.FloatSlider(value=2.0,min=0.5, max=20),
-########amaxtb_000=widgets.FloatSlider(value=55.0,min=0, max=90),
-########amaxtb_150=widgets.FloatSlider(value=55.0,min=0, max=90),
-########tsum1=widgets.FloatSlider(value=850.0,min=100, max=1800),
-########tsum2=widgets.FloatSlider(value=850.0,min=100, max=1800),
-########tsumem=widgets.FloatSlider(value=56,min=10, max=100),
-########rgrlai=widgets.FloatSlider(value=0.005, min=0.001, max=0.3, step=0.001,
-########readout_format='.3f',),
-########meteo=widgets.fixed("AgERA5_Togo_Tamale_2021_2022.csv"),
-########soil=widgets.fixed("ec4.new"),
-########wav=widgets.FloatSlider(value=5, min=0, max=100),
-########co2=widgets.fixed(400),
-########rdmsol=widgets.fixed(100.),
-########potential=widgets.Checkbox(value=True, description='Potential mode',
-########icon='check'))
+    parameters = parameters[:, ilocs]#np.nanmean(parameters[:, ilocs],
+                            #axis=1)
+    fit_lai = sim_lai[:, time_index].astype(float)[ilocs, :]
+    return (est_yield, est_yield_sd, parameters, 
+            obs_dates, obs_lai_sub, work_sim_times, fit_lai)