Merge pull request #878 from porink0424/feat/metric-names

Add `metric_names` into `Study` type

tslib/storage/src/journal.ts

@@ -29,6 +29,13 @@ interface JournalOpDeleteStudy extends JournalOpBase {
   study_id: number
 }
 
+interface JournalOpSetStudySystemAttr extends JournalOpBase {
+  study_id: number
+  system_attr: {
+    "study:metric_names": string[]
+  }
+}
+
 interface JournalOpCreateTrial extends JournalOpBase {
   study_id: number
   datetime_start?: string
@@ -179,6 +186,14 @@ class JournalStorage {
     this.studies = this.studies.filter((item) => item.id !== log.study_id)
   }
 
+  public applyStudySystemAttr(log: JournalOpSetStudySystemAttr): void {
+    const thisStudy = this.studies.find((item) => item.id === log.study_id)
+    if (thisStudy === undefined) {
+      return
+    }
+    thisStudy.metric_names = log.system_attr["study:metric_names"]
+  }
+
   public applyCreateTrial(log: JournalOpCreateTrial): void {
     const thisStudy = this.studies.find((item) => item.id === log.study_id)
     if (thisStudy === undefined) {
@@ -393,7 +408,9 @@ const loadJournalStorage = (
         // Unsupported
         break
       case JournalOperation.SET_STUDY_SYSTEM_ATTR:
-        // Unsupported
+        journalStorage.applyStudySystemAttr(
+          parsedLog as JournalOpSetStudySystemAttr
+        )
         break
       case JournalOperation.CREATE_TRIAL:
         journalStorage.applyCreateTrial(parsedLog as JournalOpCreateTrial)

tslib/storage/src/sqlite.ts

@@ -144,6 +144,11 @@ const getStudy = (
     trials: [],
   }
 
+  const studySystemAttrs = getStudySystemAttributes(db, summary.id)
+  if (studySystemAttrs !== undefined) {
+    study.metric_names = studySystemAttrs.metric_names
+  }
+
   let intersection_search_space: Set<Optuna.SearchSpaceItem> = new Set()
   study.trials = getTrials(db, summary.id, schemaVersion)
   for (const trial of study.trials) {
@@ -367,6 +372,20 @@ const parseDistributionJSON = (t: string): Optuna.Distribution => {
   }
 }
 
+const getStudySystemAttributes = (db: SQLite3DB, studyId: number) => {
+  let attrs: { metric_names: string[] } | undefined
+  db.exec({
+    sql: `SELECT key, value_json FROM study_system_attributes WHERE study_id = ${studyId} AND key = 'dashboard:objective_names'`,
+    // biome-ignore lint/suspicious/noExplicitAny: <explanation>
+    callback: (vals: any[]) => {
+      attrs = {
+        metric_names: JSON.parse(vals[1]),
+      }
+    },
+  })
+  return attrs
+}
+
 const getTrialUserAttributes = (
   db: SQLite3DB,
   trialId: number

tslib/storage/test/generate_assets.py

@@ -78,6 +78,24 @@ def create_optuna_storage(storage: BaseStorage) -> None:
 
     study.optimize(objective_single_nan_report, n_trials=100)
 
+    # Multi-objective study with metric names
+    study = optuna.create_study(
+        study_name="multi-objective-metric-names",
+        storage=storage,
+        directions=["minimize", "minimize"],
+    )
+    print(f"Generating {study.study_name} for {type(storage).__name__}...")
+    study.set_metric_names(["value1", "value2"])
+
+    def objective_multi(trial: optuna.Trial) -> tuple[float, float]:
+        x = trial.suggest_float("x", 0, 5)
+        y = trial.suggest_float("y", 0, 3)
+        v0 = 4 * x**2 + 4 * y**2
+        v1 = (x - 5) ** 2 + (y - 5) ** 2
+        return v0, v1
+
+    study.optimize(objective_multi, n_trials=50)
+
 
 if __name__ == "__main__":
     remove_assets()

tslib/storage/test/journal.test.mjs

@@ -52,8 +52,13 @@ describe("Test Journal File Storage", async () => {
     )
   })
 
+  it("Check metric_names function", () => {
+    const study = studies.find((s) => s.name === "multi-objective-metric-names")
+    assert.deepStrictEqual(study.metric_names, ["value1", "value2"])
+  })
+
   it("Check the number of studies", () => {
-    const N_STUDIES = 4
+    const N_STUDIES = 5
     assert.strictEqual(studies.length, N_STUDIES)
   })
 })

tslib/types/src/index.ts

@@ -60,6 +60,7 @@ export type Study = {
   union_user_attrs: AttributeSpec[]
   datetime_start?: Date
   trials: Trial[]
+  metric_names?: string[]
 }
 
 export type Trial = {