jsonify-optuna 0.1.0__py3-none-any.whl

jsonify_optuna/__init__.py

@@ -0,0 +1,5 @@
+from jsonify_optuna.jsonify_optuna import jsonify
+from jsonify_optuna.jsonify_optuna import load_study
+
+
+__all__ = ["jsonify", "load_study"]

jsonify_optuna/jsonify_optuna.py

@@ -0,0 +1,109 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import optuna
+from optuna.trial import TrialState
+
+
+if TYPE_CHECKING:
+    from typing import Any
+    from typing import Literal
+    from typing import TypedDict
+
+    from optuna.distributions import CategoricalChoiceType
+
+    class NumericalDistributionType(TypedDict):
+        low: int | float
+        high: int | float
+        step: int | float | None
+        log: bool
+
+    class CategoricalDistributionType(TypedDict):
+        choices: list[CategoricalChoiceType]
+
+    class TrialType(TypedDict):
+        state: Literal["running", "waiting", "completepruned", "fail"]
+        values: list[float]
+        params: dict[str, int | float | CategoricalChoiceType]
+        user_attrs: dict[str, Any]
+        intermediate_values: list[float]
+        distributions: dict[str, NumericalDistributionType | CategoricalDistributionType]
+
+    class StudyType:
+        trials: list[TrialType]
+        best_trial_indices: list[int]
+        directions: list[Literal["minimize", "maximize"]]
+        user_attrs: dict[str, Any]
+        metric_names: list[str]
+
+
+def jsonify(
+    study: optuna.Study, *, states: list[TrialState] | None = None, deepcopy: bool = True
+) -> StudyType:
+    states = states or [state for state in TrialState]
+    study_json = {
+        "directions": [dire.name.lower() for dire in study.directions],
+        "user_attrs": study.user_attrs,
+        "metric_names": study.metric_names,
+    }
+    best_trial_numbers = set(t.number for t in study.best_trials)
+    best_trial_indices = []
+    results = []
+    for t in study.get_trials(deepcopy=deepcopy):
+        if t.state not in states:
+            continue
+        if t.number in best_trial_numbers:
+            best_trial_indices.append(len(results))
+        dists = {}
+        for p, d in t.distributions.items():
+            if isinstance(d, optuna.distributions.CategoricalDistribution):
+                dists[p] = {"choices": d.choices}
+                continue
+            assert isinstance(
+                d, (optuna.distributions.IntDistribution, optuna.distributions.FloatDistribution)
+            )
+            dists[p] = {"low": d.low, "high": d.high, "step": d.step, "log": d.log}
+        row = {
+            "state": t.state.name.lower(),
+            "values": t.values,
+            "params": t.params,
+            "user_attrs": t.user_attrs,
+            "intermediate_values": t.intermediate_values,
+            "distributions": dists,
+        }
+        results.append(row)
+    study_json |= {"trials": results, "best_trial_indices": best_trial_indices}
+    return study_json
+
+
+def load_study(
+    *,
+    study_name: str | None = None,
+    storage: optuna.storages.BaseStorage | None = None,
+    journal_path: str | None = None,
+    rdb_url: str | None = None,
+) -> optuna.Study:
+    if storage is not None:
+        if rdb_url is not None or journal_path is not None:
+            raise ValueError(
+                f"storage is provided but either {rdb_url=} or {journal_path=} is also specified."
+            )
+    elif rdb_url is not None:
+        if journal_path is not None:
+            raise ValueError(f"Specify only one of {rdb_url=} or {journal_path=}.")
+        storage = optuna.storages.RDBStorage(rdb_url)
+    elif journal_path:
+        journal_file = optuna.storages.journal.JournalFileBackend(journal_path)
+        storage = optuna.storages.JournalStorage(journal_file)
+    else:
+        raise ValueError("No storage was specified.")
+    assert isinstance(storage, optuna.storages.BaseStorage), "MyPy Redefinition"
+    study_names = [study.study_name for study in storage.get_all_studies()]
+    if study_name is None and len(study_names) == 1:
+        return optuna.load_study(storage=storage, study_name=study_names[0])
+    if study_name is None:
+        raise ValueError(f"Specify `study_name` from {study_names=}.")
+    if study_name not in study_names:
+        raise ValueError(f"Specify `study_name` from {study_names=} but got {study_name=}.")
+    return optuna.load_study(storage=storage, study_name=study_name)

jsonify_optuna-0.1.0.dist-info/METADATA

@@ -0,0 +1,102 @@
+Metadata-Version: 2.4
+Name: jsonify-optuna
+Version: 0.1.0
+Summary: Convert an Optuna Study into a JSON-serializable dict.
+Requires-Python: >=3.9
+Requires-Dist: optuna>=4.9.0
+Description-Content-Type: text/markdown
+
+# jsonify-optuna
+
+Convert an [Optuna](https://github.com/optuna/optuna) Study into a plain, JSON-serializable Python dict.
+
+## Installation
+
+```bash
+pip install jsonify-optuna
+```
+
+## Usage
+
+Convert an in-memory study to a JSON-friendly dict:
+
+```python
+import optuna
+from jsonify_optuna import jsonify
+from optuna.trial import TrialState
+
+def objective(trial):
+    x = trial.suggest_float("x", -10.0, 10.0)
+    return x ** 2
+
+study = optuna.create_study(direction="minimize")
+study.optimize(objective, n_trials=20)
+result = jsonify(study, states=[TrialState.COMPLETE])
+# This includes all the trials.
+result = jsonify(study)
+```
+
+If you have a study stored in a backend, you can load it as follows:
+
+```python
+from jsonify_optuna import load_study
+
+# From an RDB (SQLite, PostgreSQL, etc.)
+study = load_study(rdb_url="sqlite:///example.db", study_name="my_study")
+
+# From a journal file
+study = load_study(journal_path="journal.log", study_name="my_study")
+
+# From a storage object directly
+study = load_study(storage=my_storage, study_name="my_study")
+```
+
+> [!NOTE]
+> When the storage contains exactly one study, `study_name` can be omitted.
+
+## API Reference
+
+### `jsonify(study, *, states=None, deepcopy=True)`
+
+| Parameter | Type | Description |
+|-----------|------|-------------|
+| `study` | `optuna.Study` | The study to convert. |
+| `states` | `list[TrialState] \| None` | Filter trials by state. `None` includes all states. |
+| `deepcopy` | `bool` | Whether to deep-copy trials from storage. Default `True`. |
+
+Returns a dict with:
+
+| Key | Type | Description |
+|-----|------|-------------|
+| `directions` | `list[str]` | `"minimize"` or `"maximize"` per objective. |
+| `user_attrs` | `dict` | Study-level user attributes. |
+| `metric_names` | `list[str] \| None` | Metric names if set, else `None`. |
+| `trials` | `list[dict]` | Trial dicts (see below). |
+| `best_trial_indices` | `list[int]` | Indices into `trials` for Pareto-front trials. |
+
+Each trial dict contains:
+
+| Key | Type | Description |
+|-----|------|-------------|
+| `state` | `str` | `"complete"`, `"pruned"`, `"fail"`, `"running"`, or `"waiting"`. |
+| `values` | `list[float] \| None` | Objective values. `None` if not complete. |
+| `params` | `dict` | Parameter name-value pairs. |
+| `user_attrs` | `dict` | Trial-level user attributes. |
+| `intermediate_values` | `dict[int, float]` | Step-value pairs from `trial.report()`. |
+| `distributions` | `dict` | Parameter distributions (see below). |
+
+Distribution formats:
+
+- **Int/Float**: `{"low": ..., "high": ..., "step": ..., "log": ...}`
+- **Categorical**: `{"choices": [...]}`
+
+### `load_study(*, study_name=None, storage=None, journal_path=None, rdb_url=None)`
+
+| Parameter | Type | Description |
+|-----------|------|-------------|
+| `study_name` | `str \| None` | Study name. Optional when storage has exactly one study. |
+| `storage` | `BaseStorage \| None` | An Optuna storage object. |
+| `journal_path` | `str \| None` | Path to a journal log file. |
+| `rdb_url` | `str \| None` | SQLAlchemy-style database URL. |
+
+Exactly one of `storage`, `journal_path`, or `rdb_url` must be provided.

jsonify_optuna-0.1.0.dist-info/RECORD

@@ -0,0 +1,5 @@
+jsonify_optuna/__init__.py,sha256=E41S1z9of4nzR_wRJEGkHaY0B8XLyJhDjMWxP7vcVH0,141
+jsonify_optuna/jsonify_optuna.py,sha256=yW5V-46_ba-lnIlq1AW1-30eVOtFLe_uW185s0QaBRU,4078
+jsonify_optuna-0.1.0.dist-info/METADATA,sha256=dbvGicgEcwcDMSO4_6KmBaMUksRg0PsxKk99G_ByVgI,3340
+jsonify_optuna-0.1.0.dist-info/WHEEL,sha256=lCkmxWfQsSc9CfIClYeavTdQeEX2toPqufh9gI35EQA,87
+jsonify_optuna-0.1.0.dist-info/RECORD,,

jsonify_optuna-0.1.0.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: hatchling 1.31.0
+Root-Is-Purelib: true
+Tag: py3-none-any