From 55fe37e59a92ce57b466664ff7271bc0a49a2017 Mon Sep 17 00:00:00 2001
From: porink0424 <daichi2120@gmail.com>
Date: Wed, 10 Apr 2024 18:49:05 +0900
Subject: [PATCH] Add importance props in PlotImportance so that importance can
 be calculated outside

---
 standalone_app/package-lock.json              | 12 +++-
 standalone_app/package.json                   |  1 +
 standalone_app/src/components/StudyDetail.tsx | 68 ++++++++++++++++++-
 3 files changed, 79 insertions(+), 2 deletions(-)

diff --git a/standalone_app/package-lock.json b/standalone_app/package-lock.json
index 2b2040e7..8c4a2211 100644
--- a/standalone_app/package-lock.json
+++ b/standalone_app/package-lock.json
@@ -24,6 +24,7 @@
         "react-router-dom": "^6.22.3"
       },
       "devDependencies": {
+        "@optuna/types": "../tslib/types",
         "@types/plotly.js": "^2.29.2",
         "@types/react": "^18.2.64",
         "@types/react-dom": "^18.2.21",
@@ -54,7 +55,6 @@
         "@mui/material": "^5.15.10",
         "@mui/system": "^5.15.9",
         "@optuna/storage": "../storage/",
-        "optuna": "../../rustlib/pkg",
         "plotly.js-dist-min": "^2.30.1",
         "react": "^18.2.0",
         "react-dom": "^18.2.0",
@@ -93,6 +93,12 @@
         "@optuna/types": "../types/"
       }
     },
+    "../tslib/types": {
+      "name": "@optuna/types",
+      "version": "0.0.1",
+      "dev": true,
+      "license": "MIT"
+    },
     "node_modules/@ampproject/remapping": {
       "version": "2.3.0",
       "resolved": "https://registry.npmjs.org/@ampproject/remapping/-/remapping-2.3.0.tgz",
@@ -1340,6 +1346,10 @@
       "resolved": "../tslib/storage",
       "link": true
     },
+    "node_modules/@optuna/types": {
+      "resolved": "../tslib/types",
+      "link": true
+    },
     "node_modules/@popperjs/core": {
       "version": "2.11.8",
       "resolved": "https://registry.npmjs.org/@popperjs/core/-/core-2.11.8.tgz",
diff --git a/standalone_app/package.json b/standalone_app/package.json
index 9af17831..a9f26a7f 100644
--- a/standalone_app/package.json
+++ b/standalone_app/package.json
@@ -8,6 +8,7 @@
     "build:vscode": "webpack"
   },
   "devDependencies": {
+    "@optuna/types": "../tslib/types",
     "@types/plotly.js": "^2.29.2",
     "@types/react": "^18.2.64",
     "@types/react-dom": "^18.2.21",
diff --git a/standalone_app/src/components/StudyDetail.tsx b/standalone_app/src/components/StudyDetail.tsx
index 7936d857..4c60e0e5 100644
--- a/standalone_app/src/components/StudyDetail.tsx
+++ b/standalone_app/src/components/StudyDetail.tsx
@@ -19,6 +19,8 @@ import {
   PlotIntermediateValues,
   TrialTable,
 } from "@optuna/react"
+import * as Optuna from "@optuna/types"
+import init, { wasm_fanova_calculate } from "optuna"
 import React, { FC, useContext, useState, useEffect } from "react"
 import { Link, useParams } from "react-router-dom"
 import { StorageContext } from "./StorageProvider"
@@ -43,6 +45,68 @@ export const StudyDetail: FC<{
     fetchStudy()
   }, [storage, idxNumber])
 
+  const [importance, setImportance] = useState<Optuna.ParamImportance[][]>([])
+  const filterFunc = (trial: Optuna.Trial, objectiveId: number): boolean => {
+    if (trial.state !== "Complete" && trial.state !== "Pruned") {
+      return false
+    }
+    if (trial.values === undefined) {
+      return false
+    }
+    return (
+      trial.values.length > objectiveId &&
+      trial.values[objectiveId] !== Infinity &&
+      trial.values[objectiveId] !== -Infinity
+    )
+  }
+  // biome-ignore lint/correctness/useExhaustiveDependencies: <explanation>
+  useEffect(() => {
+    async function run_wasm() {
+      if (study === null) {
+        return
+      }
+
+      await init()
+
+      const x: Optuna.ParamImportance[][] = study.directions.map(
+        (_d, objectiveId) => {
+          const filteredTrials = study.trials.filter((t) =>
+            filterFunc(t, objectiveId)
+          )
+          if (filteredTrials.length === 0) {
+            return study.union_search_space.map((s) => {
+              return {
+                name: s.name,
+                importance: 0.5,
+              }
+            })
+          }
+
+          const features = study.intersection_search_space.map((s) =>
+            filteredTrials
+              .map(
+                (t) =>
+                  t.params.find((p) => p.name === s.name) as Optuna.TrialParam
+              )
+              .map((p) => p.param_internal_value)
+          )
+          const values = filteredTrials.map(
+            (t) => t.values?.[objectiveId] as number
+          )
+          // TODO: handle errors thrown by wasm_fanova_calculate
+          const importance = wasm_fanova_calculate(features, values)
+          return study.intersection_search_space.map((s, i) => ({
+            name: s.name,
+            importance: importance[i],
+          }))
+        }
+      )
+      setImportance(x)
+    }
+
+    run_wasm()
+  }, [study])
+
   return (
     <>
       <AppBar position="static">
@@ -116,7 +180,9 @@ export const StudyDetail: FC<{
             <Grid2 xs={6}>
               <Card sx={{ margin: theme.spacing(2) }}>
                 <CardContent>
-                  {!!study && <PlotImportance study={study} />}
+                  {!!study && (
+                    <PlotImportance study={study} importance={importance} />
+                  )}
                 </CardContent>
               </Card>
             </Grid2>