From e215b1bf10a95cd19f35aaf860d322d5d1e7f845 Mon Sep 17 00:00:00 2001 From: porink0424 Date: Fri, 21 Jun 2024 18:36:57 +0900 Subject: [PATCH] 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 }