From c5f0eec50d4d1073423ad80635b3075c9fc2b546 Mon Sep 17 00:00:00 2001
From: Jeroen Overschie <jeroenoverschie@gmail.com>
Date: Tue, 22 Jun 2021 16:58:17 +0200
Subject: [PATCH] [WIP] Plotting charts through API

---
 fseval/callbacks/wandb.py                     | 50 +++++++++++++++++--
 .../rank_and_validate/rank_and_validate.py    |  7 +++
 2 files changed, 54 insertions(+), 3 deletions(-)

diff --git a/fseval/callbacks/wandb.py b/fseval/callbacks/wandb.py
index b2c72da..2896713 100644
--- a/fseval/callbacks/wandb.py
+++ b/fseval/callbacks/wandb.py
@@ -4,11 +4,12 @@
 from logging import Logger, getLogger
 from typing import Dict, Optional, cast
 
-import wandb
-from omegaconf import DictConfig, OmegaConf
-
 from fseval.types import Callback
 from fseval.utils.dict_utils import dict_flatten, dict_merge
+from omegaconf import DictConfig, OmegaConf
+
+import wandb
+from wandb.viz import CustomChart, custom_chart_panel_config
 
 
 class WandbCallback(Callback):
@@ -90,3 +91,46 @@ def upload_table(self, df, name):
         logs = {}
         logs[name] = table
         wandb.log(logs)
+
+    def add_panel(
+        self,
+        viz_id: str,
+        panel_name: str,
+        table_key: str,
+        fields: Dict = {},
+        string_fields: Dict = {},
+        panel_config_callback=lambda panel_config: panel_config,
+    ):
+        """Adds a custom chart panel to the current wandb run. This function uses
+        internal wandb functions, so might be prone to changes in their code. The
+        function is a mixup of the following modules / functions:
+
+        - `wandb.viz`: has `CustomChart` and `custom_chart_panel_config` functions.
+            see https://github.com/wandb/client/blob/master/wandb/viz.py
+        - `wandb.sdk.wandb_run.Run`: has `_add_panel` and `_backend` functions. This
+        function mainly replicates whatever `_history_callback` is doing.
+            see https://github.com/wandb/client/blob/master/wandb/sdk/wandb_run.py
+        """
+
+        assert wandb.run is not None, "no wandb run in progress. wandb.run is None."
+
+        # create custom chart. is just a data holder class for its attributes.
+        custom_chart = CustomChart(
+            viz_id=viz_id,
+            table=None,
+            fields=fields,
+            string_fields=string_fields,
+        )
+
+        # create custom chart config.
+        # Function `custom_chart_panel_config(custom_chart, key, table_key)` has a
+        # useless attribute, `key`.
+        panel_config = custom_chart_panel_config(custom_chart, None, table_key)
+        panel_config = panel_config_callback(panel_config)
+
+        # add chart to current run.
+        wandb.run._add_panel(panel_name, "Vega2", panel_config)
+
+        # "publish" chart to backend
+        if wandb.run._backend:
+            wandb.run._backend.interface.publish_history({}, wandb.run.step)
diff --git a/fseval/pipelines/rank_and_validate/rank_and_validate.py b/fseval/pipelines/rank_and_validate/rank_and_validate.py
index 5a2477a..6939376 100644
--- a/fseval/pipelines/rank_and_validate/rank_and_validate.py
+++ b/fseval/pipelines/rank_and_validate/rank_and_validate.py
@@ -255,4 +255,11 @@ def score(self, X, y, **kwargs):
         if wandb_callback:
             self.logger.info(f"Tables uploaded {tc.green('✓')}")
 
+        # Upload charts
+        # create chart:
+        if wandb_callback:
+            # use wandb_callback.add_panel
+            # "dunnkers/fseval/feature-importances-all-bootstraps-with-ticks"
+            ...
+
         return summary