Add start/end date selectors for experiment filtering (#370)

Co-authored-by: Remi Lehe <remi.lehe@normalesup.org>

Author: EZoni

dashboard/app.py

@@ -105,7 +105,10 @@ def update(
         ctrl.figure_update(fig)
 
 
-@state.change("experiment")
+@state.change(
+    "experiment",
+    "experiment_date_range",
+)
 def update_on_change_experiment(**kwargs):
     # skip if triggered on server ready (all state variables marked as modified)
     if len(state.modified_keys) == 1:
@@ -361,14 +364,24 @@ def gui_setup():
         # add toolbar components
         with layout.toolbar:
             vuetify.VSpacer()
+            # experiment selector
             vuetify.VSelect(
                 v_model=("experiment",),
                 label="Experiments",
                 items=(experiments,),
                 dense=True,
                 hide_details=True,
                 prepend_icon="mdi-atom",
-                style="max-width: 250px",
+                style="max-width: 250px; margin-right: 14px;",
+            )
+            # date range selector for experiment filtering
+            vuetify.VDateInput(
+                v_model=("experiment_date_range",),
+                label="Date range",
+                multiple="range",
+                dense=True,
+                hide_details=True,
+                style="max-width: 250px; margin-right: 14px;",
             )
         # set up router view
         with layout.content:

dashboard/model_manager.py

@@ -11,7 +11,7 @@
 from lume_model.models.ensemble import NNEnsemble
 from lume_model.models.gp_model import GPModel
 from trame.widgets import vuetify3 as vuetify
-from utils import verify_input_variables, timer, load_config_dict
+from utils import verify_input_variables, timer, load_config_dict, create_date_filter
 from error_manager import add_error
 from sfapi_manager import monitor_sfapi_job
 from state_manager import state
@@ -190,9 +190,13 @@ async def training_kernel(self):
                 client_id=state.sfapi_client_id, secret=state.sfapi_key
             ) as client:
                 perlmutter = await client.compute(Machine.perlmutter)
-                # Upload the config.yaml to nersc
+                # upload the configuration file to NERSC
                 config_dict = load_config_dict(state.experiment)
                 config_dict["simulation_calibration"] = state.simulation_calibration
+                # add date range filter to the configuration dictionary
+                date_filter = create_date_filter(state.experiment_date_range)
+                config_dict["date_filter"] = date_filter
+                # define the target path on NERSC
                 target_path = "/global/cfs/cdirs/m558/superfacility/model_training"
                 [target_path] = await perlmutter.ls(target_path, directory=True)
                 with tempfile.TemporaryDirectory() as temp_dir:

dashboard/state_manager.py

@@ -1,5 +1,6 @@
 from pathlib import Path
 from trame.app import get_server
+from trame.widgets import vuetify3 as vuetify
 
 
 EXPERIMENTS_PATH = Path.cwd().parent / "experiments/"
@@ -8,6 +9,7 @@
 server = get_server(client_type="vue3")
 state = server.state
 ctrl = server.controller
+vuetify.enable_lab()  # Enable Labs components
 
 
 def initialize_state():
@@ -23,6 +25,7 @@ def initialize_state():
     ][0]
     print(f"Setting default experiment to {default_experiment}...")
     state.experiment = default_experiment
+    state.experiment_date_range = []
     # ML model
     state.model_type = "Neural Network (single)"
     state.model_training = False

dashboard/utils.py

@@ -70,11 +70,47 @@ def load_variables(experiment):
     return (input_variables, output_variables, simulation_calibration)
 
 
+def create_date_filter(experiment_date_range):
+    # build date filter if date range is set
+    date_filter = {}
+    if experiment_date_range:
+        start_date = pd.to_datetime(experiment_date_range[0].to_datetime())
+        start_date = start_date.to_pydatetime().replace(hour=0, minute=0, second=0)
+        # VDateInput returns exclusive end date for date ranges:
+        # - subtract 1 day for multi-date ranges with different start/end dates
+        # - do not subtract anything (use end date as is) for single-date ranges
+        end_date = pd.to_datetime(experiment_date_range[-1].to_datetime())
+        end_date_correction = (
+            pd.Timedelta(days=0)
+            if len(experiment_date_range) == 1
+            else pd.Timedelta(days=1)
+        )
+        end_date = end_date - end_date_correction
+        end_date = end_date.to_pydatetime().replace(hour=23, minute=59, second=59)
+        # remove timezone info to match naive datetime in database
+        start_date = (
+            start_date.replace(tzinfo=None) if start_date.tzinfo else start_date
+        )
+        end_date = end_date.replace(tzinfo=None) if end_date.tzinfo else end_date
+        date_filter = {
+            "date": {
+                "$gte": start_date,
+                "$lte": end_date,
+            }
+        }
+        print(f"Filtering data between {start_date.date()} and {end_date.date()}...")
+    return date_filter
+
+
 @timer
 def load_data(db):
     print("Loading data from database...")
+    # create date filter if date range is set
+    date_filter = create_date_filter(state.experiment_date_range)
     # load experiment and simulation data points in dataframes
-    exp_data = pd.DataFrame(db[state.experiment].find({"experiment_flag": 1}))
+    exp_data = pd.DataFrame(
+        db[state.experiment].find({"experiment_flag": 1, **date_filter})
+    )
     sim_data = pd.DataFrame(db[state.experiment].find({"experiment_flag": 0}))
     # Store '_id', 'date' as string
     for key in ["_id", "date"]:

ml/train_model.py

@@ -439,7 +439,8 @@ def write_model(model, model_type, experiment, db):
 
     # Extract experimental and simulation data from the database as pandas dataframe
     db = connect_to_db(config_dict)
-    df_exp = pd.DataFrame(db[experiment].find({"experiment_flag": 1}))
+    date_filter = config_dict.get("date_filter", {})
+    df_exp = pd.DataFrame(db[experiment].find({"experiment_flag": 1, **date_filter}))
     df_sim = pd.DataFrame(db[experiment].find({"experiment_flag": 0}))
 
     # Apply simulation calibration to the simulation data