From e215b1bf10a95cd19f35aaf860d322d5d1e7f845 Mon Sep 17 00:00:00 2001 From: porink0424 Date: Fri, 21 Jun 2024 18:36:57 +0900 Subject: [PATCH 1/3] Modify the logic of storages --- tslib/storage/src/journal.ts | 72 +++++++++++++++++++++++++++--------- tslib/storage/src/sqlite.ts | 64 +++++++++++++++++++++----------- 2 files changed, 98 insertions(+), 38 deletions(-) diff --git a/tslib/storage/src/journal.ts b/tslib/storage/src/journal.ts index 05c0779d..6ee09281 100644 --- a/tslib/storage/src/journal.ts +++ b/tslib/storage/src/journal.ts @@ -1,6 +1,36 @@ import * as Optuna from "@optuna/types" import { OptunaStorage } from "./storage" +// TODO(porink0424): Refactor to common function with sqlite.ts (current workaround duplicates code due to missing file extensions in tsc build output). +const isDistributionEqual = ( + a: Optuna.Distribution, + b: Optuna.Distribution +) => { + if (a.type !== b.type) { + return false + } + + if (a.type === "IntDistribution" || a.type === "FloatDistribution") { + if (b.type !== "IntDistribution" && b.type !== "FloatDistribution") { + throw new Error("Invalid distribution type") + } + return ( + a.low === b.low && + a.high === b.high && + a.step === b.step && + a.log === b.log + ) + } + if (a.type === "CategoricalDistribution") { + if (b.type !== "CategoricalDistribution") { + throw new Error("Invalid distribution type") + } + return JSON.stringify(a.choices) === JSON.stringify(b.choices) + } + + throw new Error("Invalid distribution type") +} + // JournalStorage enum JournalOperation { CREATE_STUDY = 0, @@ -137,30 +167,42 @@ class JournalStorage { public getStudies(): Optuna.Study[] { for (const study of this.studies) { const unionUserAttrs: Set = new Set() - const unionSearchSpace: Set = new Set() - let intersectionSearchSpace: string[] = [] - const nameToSearchSpaceItem: Map = - new Map() + const unionSearchSpace: Optuna.SearchSpaceItem[] = [] + let intersectionSearchSpace: Optuna.SearchSpaceItem[] = [] study.trials.forEach((trial, index) => { for (const userAttr of trial.user_attrs) { unionUserAttrs.add(userAttr.key) } for (const param of trial.params) { - unionSearchSpace.add(param.name) - if (!nameToSearchSpaceItem.has(param.name)) { - nameToSearchSpaceItem.set(param.name, { + if ( + !unionSearchSpace.some( + (item) => + item.name === param.name && + isDistributionEqual(item.distribution, param.distribution) + ) + ) { + unionSearchSpace.push({ name: param.name, distribution: param.distribution, }) } } if (index === 0) { - intersectionSearchSpace = Array.from(unionSearchSpace) + intersectionSearchSpace = [...unionSearchSpace] } else { - intersectionSearchSpace = intersectionSearchSpace.filter((name) => { - return trial.params.some((param) => param.name === name) - }) + intersectionSearchSpace = intersectionSearchSpace.filter( + (searchSpaceItem) => { + return trial.params.some( + (param) => + param.name === searchSpaceItem.name && + isDistributionEqual( + param.distribution, + searchSpaceItem.distribution + ) + ) + } + ) } }) study.union_user_attrs = Array.from(unionUserAttrs).map((key) => { @@ -169,12 +211,8 @@ class JournalStorage { sortable: false, } }) - study.union_search_space = Array.from(unionSearchSpace).map((name) => { - return nameToSearchSpaceItem.get(name) as Optuna.SearchSpaceItem - }) - study.intersection_search_space = intersectionSearchSpace.map((name) => { - return nameToSearchSpaceItem.get(name) as Optuna.SearchSpaceItem - }) + study.union_search_space = unionSearchSpace + study.intersection_search_space = intersectionSearchSpace } return this.studies diff --git a/tslib/storage/src/sqlite.ts b/tslib/storage/src/sqlite.ts index 11912ac2..1f9bc7a9 100644 --- a/tslib/storage/src/sqlite.ts +++ b/tslib/storage/src/sqlite.ts @@ -3,6 +3,36 @@ import * as Optuna from "@optuna/types" import sqlite3InitModule from "@sqlite.org/sqlite-wasm" import { OptunaStorage } from "./storage" +// TODO(porink0424): Refactor to common function with journal.ts (current workaround duplicates code due to missing file extensions in tsc build output). +const isDistributionEqual = ( + a: Optuna.Distribution, + b: Optuna.Distribution +) => { + if (a.type !== b.type) { + return false + } + + if (a.type === "IntDistribution" || a.type === "FloatDistribution") { + if (b.type !== "IntDistribution" && b.type !== "FloatDistribution") { + throw new Error("Invalid distribution type") + } + return ( + a.low === b.low && + a.high === b.high && + a.step === b.step && + a.log === b.log + ) + } + if (a.type === "CategoricalDistribution") { + if (b.type !== "CategoricalDistribution") { + throw new Error("Invalid distribution type") + } + return JSON.stringify(a.choices) === JSON.stringify(b.choices) + } + + throw new Error("Invalid distribution type") +} + type SQLite3DB = { exec(options: { sql: string @@ -149,7 +179,7 @@ const getStudy = ( study.metric_names = studySystemAttrs.metric_names } - let intersection_search_space: Set = new Set() + let intersectionSearchSpace: Optuna.SearchSpaceItem[] = [] study.trials = getTrials(db, summary.id, schemaVersion) for (const trial of study.trials) { const userAttrs = getTrialUserAttributes(db, trial.trial_id) @@ -165,16 +195,7 @@ const getStudy = ( } const params = getTrialParams(db, trial.trial_id) - const paramNames = new Set() - const paramNameToSearchSpaceItem = new Map() for (const param of params) { - paramNames.add(param.name) - if (paramNameToSearchSpaceItem.has(param.name)) { - paramNameToSearchSpaceItem.set(param.name, { - name: param.name, - distribution: param.distribution, - }) - } if ( study.union_search_space.findIndex((s) => s.name === param.name) === -1 ) { @@ -184,23 +205,24 @@ const getStudy = ( }) } } - if (intersection_search_space.size === 0) { - for (const s of paramNames) { - intersection_search_space.add( - paramNameToSearchSpaceItem.get(s) as Optuna.SearchSpaceItem - ) - } + if (intersectionSearchSpace.length === 0) { + intersectionSearchSpace = params.map((param) => ({ + name: param.name, + distribution: param.distribution, + })) } else { - intersection_search_space = new Set( - Array.from(intersection_search_space).filter((s) => - paramNames.has(s.name) + intersectionSearchSpace = intersectionSearchSpace.filter((item) => { + return params.some( + (param) => + item.name === param.name && + isDistributionEqual(item.distribution, param.distribution) ) - ) + }) } trial.params = params trial.user_attrs = userAttrs } - study.intersection_search_space = Array.from(intersection_search_space) + study.intersection_search_space = intersectionSearchSpace return study } From 747098b6894b88a2445cf03aaf34bc4b2e009f27 Mon Sep 17 00:00:00 2001 From: porink0424 Date: Fri, 21 Jun 2024 18:37:28 +0900 Subject: [PATCH 2/3] Add the test to check the study with dynamic search space --- tslib/storage/test/generate_assets.py | 4 ++-- tslib/storage/test/journal.test.mjs | 12 ++++++++++++ 2 files changed, 14 insertions(+), 2 deletions(-) diff --git a/tslib/storage/test/generate_assets.py b/tslib/storage/test/generate_assets.py index 2e79e141..7799b726 100644 --- a/tslib/storage/test/generate_assets.py +++ b/tslib/storage/test/generate_assets.py @@ -44,9 +44,9 @@ def create_optuna_storage(storage: BaseStorage) -> None: def objective_single_dynamic(trial: optuna.Trial) -> float: category = trial.suggest_categorical("category", ["foo", "bar"]) if category == "foo": - return (trial.suggest_float("x1", 0, 10) - 2) ** 2 + return (trial.suggest_float("x", 0, 10) - 2) ** 2 else: - return -((trial.suggest_float("x2", -10, 0) + 5) ** 2) + return -((trial.suggest_float("x", -10, 0) + 5) ** 2) study.optimize(objective_single_dynamic, n_trials=50) diff --git a/tslib/storage/test/journal.test.mjs b/tslib/storage/test/journal.test.mjs index 1f8f1299..70e4f7a5 100644 --- a/tslib/storage/test/journal.test.mjs +++ b/tslib/storage/test/journal.test.mjs @@ -16,6 +16,18 @@ describe("Test Journal File Storage", async () => { studySummaries.map((_summary, index) => storage.getStudy(index)) ) + it("Check the study with dynamic search space", () => { + const study = studies.find((s) => s.name === "single-objective-dynamic") + assert.deepStrictEqual( + study.union_search_space.map((item) => item.name).sort(), + ["x", "x", "category"].sort() + ) + assert.deepStrictEqual( + study.intersection_search_space.map((item) => item.name).sort(), + ["category"].sort() + ) + }) + it("Check the study including Infinities", () => { const study = studies.find((s) => s.name === "single-inf") study.trials.forEach((trial, index) => { From 29c4e73b94f767a0dfa044eb6572a2bed8e5fa22 Mon Sep 17 00:00:00 2001 From: porink0424 Date: Wed, 26 Jun 2024 10:34:32 +0900 Subject: [PATCH 3/3] Add the test for checking the distribution property as well --- tslib/storage/test/journal.test.mjs | 42 +++++++++++++++++++++++++++++ 1 file changed, 42 insertions(+) diff --git a/tslib/storage/test/journal.test.mjs b/tslib/storage/test/journal.test.mjs index 70e4f7a5..a9bb4c04 100644 --- a/tslib/storage/test/journal.test.mjs +++ b/tslib/storage/test/journal.test.mjs @@ -22,10 +22,52 @@ describe("Test Journal File Storage", async () => { study.union_search_space.map((item) => item.name).sort(), ["x", "x", "category"].sort() ) + assert.strictEqual( + study.union_search_space.some( + (item) => + item.name === "category" && + item.distribution.type === "CategoricalDistribution" && + item.distribution.choices.length === 2 + ), + true + ) + assert.strictEqual( + study.union_search_space.some( + (item) => + item.name === "x" && + item.distribution.type === "FloatDistribution" && + item.distribution.low === 0 && + item.distribution.high === 10 && + item.distribution.step === null && + item.distribution.log === false + ), + true + ) + assert.strictEqual( + study.union_search_space.some( + (item) => + item.name === "x" && + item.distribution.type === "FloatDistribution" && + item.distribution.low === -10 && + item.distribution.high === 0 && + item.distribution.step === null && + item.distribution.log === false + ), + true + ) assert.deepStrictEqual( study.intersection_search_space.map((item) => item.name).sort(), ["category"].sort() ) + assert.strictEqual( + study.intersection_search_space.some( + (item) => + item.name === "category" && + item.distribution.type === "CategoricalDistribution" && + item.distribution.choices.length === 2 + ), + true + ) }) it("Check the study including Infinities", () => {