diff --git a/.github/FUNDING.yml b/.github/FUNDING.yml new file mode 100644 index 00000000..c2c7a342 --- /dev/null +++ b/.github/FUNDING.yml @@ -0,0 +1 @@ +github: optuna diff --git a/docs/tutorials/hitl.rst b/docs/tutorials/hitl.rst index 6d31ae55..e06709c8 100644 --- a/docs/tutorials/hitl.rst +++ b/docs/tutorials/hitl.rst @@ -1,11 +1,13 @@ -Tutorial: Human-in-the-loop Optimization -======================================== +.. _tutorial-hitl-objective-form-widgets: + +Tutorial: Human-in-the-loop Optimization using Objective Form Widgets +===================================================================== .. image:: ./images/hitl1.png In tasks involving image generation, natural language, or speech synthesis, evaluating results mechanically can be tough, and human evaluation becomes crucial. Until now, managing such tasks with Optuna has been challenging. However, the introduction of Optuna Dashboard enables humans and optimization algorithms to work interactively and execute the optimization process. -In this tutorial, we will explain how to optimize hyperparameters to generate a simple image using Optuna Dashboard. While the tutorial focuses on a simple task, the same approach can be applied to for instance optimize more complex images, natural language, and speech. +In this tutorial, we will explain how to optimize hyperparameters to generate a simple image using Optuna Dashboard. While the tutorial focuses on a simple task, the same approach can be applied to for instance optimize more complex images, natural language, and speech. The tutorial is organized as follows: @@ -93,11 +95,11 @@ Given the above system, we carry out HITL optimization as follows: Environment setup ^^^^^^^^^^^^^^^^^ -To run `the script `_ used in this tutorial, you need to install two libraries: +To run `the script `_ used in this tutorial, you need to install following libraries: .. code-block:: console - $ pip install "optuna>=3.2.0" "optuna-dashboard>=0.10.0" pillow + $ pip install "optuna>=3.3.0" "optuna-dashboard>=0.12.0" pillow You will use SQLite for the storage backend in this tutorial. Ensure that the following library is installed: @@ -179,7 +181,9 @@ Let’s walk through the script we used for the optimization. .. code-block:: python :linenos: - def suggest_and_generate_image(study: optuna.Study, artifact_backend: FileSystemBackend) -> None: + def suggest_and_generate_image( + study: optuna.Study, artifact_store: FileSystemArtifactStore + ) -> None: # 1. Ask new parameters trial = study.ask() r = trial.suggest_int("r", 0, 255) @@ -192,7 +196,7 @@ Let’s walk through the script we used for the optimization. image.save(image_path) # 3. Upload Artifact - artifact_id = upload_artifact(artifact_backend, trial, image_path) + artifact_id = upload_artifact(trial, image_path, artifact_store) artifact_path = get_artifact_path(trial, artifact_id) # 4. Save Note @@ -205,12 +209,12 @@ Let’s walk through the script we used for the optimization. ) save_note(trial, note) -In the ``suggest_and_generate_image`` function, a new Trial is obtained and new hyperparameters are suggested for that Trial. Based on those hyperparameters, an RGB image is generated as an artifact. The generated image is then uploaded to the Artifact Storage of the Optuna Dashboard, and the image is also displayed in the Dashboard's Note. For more information on how to use the Note feature, please refer to the API Reference of :func:`~optuna_dashboard.save_note`. +In the ``suggest_and_generate_image`` function, a new Trial is obtained and new hyperparameters are suggested for that Trial. Based on those hyperparameters, an RGB image is generated as an artifact. The generated image is then uploaded to the Artifact Store of the Optuna, and the image is also displayed in the Dashboard's Note. For more information on how to use the Note feature, please refer to the API Reference of :func:`~optuna_dashboard.save_note`. .. code-block:: python :linenos: - def start_optimization(artifact_backend: FileSystemBackend) -> NoReturn: + def start_optimization(artifact_store: FileSystemArtifactStore) -> NoReturn: # 1. Create Study study = optuna.create_study( study_name="Human-in-the-loop Optimization", @@ -218,10 +222,10 @@ In the ``suggest_and_generate_image`` function, a new Trial is obtained and new sampler=optuna.samplers.TPESampler(constant_liar=True, n_startup_trials=5), load_if_exists=True, ) - + # 2. Set an objective name study.set_metric_names(["Looks like sunset color?"]) - + # 3. Register ChoiceWidget register_objective_form_widgets( study, @@ -234,15 +238,14 @@ In the ``suggest_and_generate_image`` function, a new Trial is obtained and new ], ) - # 4. Start Optimization + # 4. Start Human-in-the-loop Optimization n_batch = 4 while True: running_trials = study.get_trials(deepcopy=False, states=(TrialState.RUNNING,)) if len(running_trials) >= n_batch: time.sleep(1) # Avoid busy-loop continue - suggest_and_generate_image(study, artifact_backend) - + suggest_and_generate_image(study, artifact_store) The function ``start_optimization`` defines our loop for HITL optimization to generate an image resembling a sunset color. @@ -256,10 +259,10 @@ The function ``start_optimization`` defines our loop for HITL optimization to ge def main() -> NoReturn: tmp_path = os.path.join(os.path.dirname(__file__), "tmp") - + # 1. Create Artifact Store artifact_path = os.path.join(os.path.dirname(__file__), "artifact") - artifact_backend = FileSystemBackend(base_path=artifact_path) + artifact_store = FileSystemArtifactStore(artifact_path) if not os.path.exists(artifact_path): os.mkdir(artifact_path) @@ -268,11 +271,11 @@ The function ``start_optimization`` defines our loop for HITL optimization to ge os.mkdir(tmp_path) # 2. Run optimize loop - start_optimization(artifact_backend) + start_optimization(artifact_store) In the ``main`` function, at first, the locations of the Artifact Store is set. -* At #1, the :class:`~optuna_dashboard.FileSystemBackend` is created, which is one of the Artifact Storage options used in the Optuna Dashboard. Artifact Storage is used to store artifacts (data, files, etc.) generated during Optuna trials. For more information, please refer to the API Reference. +* At #1, the `FileSystemArtifactStore `_ is created, which is one of the Artifact Store options used in the Optuna. Artifact Store is used to store artifacts (data, files, etc.) generated during Optuna trials. For more information, please refer to the API Reference. * At #2, `start_optimization()` function, which is described above, is called. After that, two folders are created, artifact and tmp, and then ``start_optimization`` function is called to start the HITL optimization using Optuna. diff --git a/docs/tutorials/images/preferential-optimization/anim.gif b/docs/tutorials/images/preferential-optimization/anim.gif new file mode 100644 index 00000000..ba15e2b5 Binary files /dev/null and b/docs/tutorials/images/preferential-optimization/anim.gif differ diff --git a/docs/tutorials/images/preferential-optimization/system-architecture.png b/docs/tutorials/images/preferential-optimization/system-architecture.png new file mode 100644 index 00000000..73dfd8d6 Binary files /dev/null and b/docs/tutorials/images/preferential-optimization/system-architecture.png differ diff --git a/docs/tutorials/index.rst b/docs/tutorials/index.rst index 5af1ba40..f396e740 100644 --- a/docs/tutorials/index.rst +++ b/docs/tutorials/index.rst @@ -5,3 +5,4 @@ Tutorials :maxdepth: 1 hitl + preferential-optimization diff --git a/docs/tutorials/preferential-optimization.rst b/docs/tutorials/preferential-optimization.rst new file mode 100644 index 00000000..706eabe8 --- /dev/null +++ b/docs/tutorials/preferential-optimization.rst @@ -0,0 +1,135 @@ +Tutorial: Preferential Optimization +=================================== + +What is Preferential Optimization? +---------------------------------- + +Preferential optimization is the way to optimize hyperparameters based on human preferences, +specifically by determining which trial is better when given a pair to compare. +Compared to the `human-in-the-loop optimization utilizing objective form widgets `_, +which relies on absolute evaluations, preferential optimization significantly reduces fluctuations in the evaluators' criteria, +ensuring more consistent results. + +In this tutorial, we will interactively optimize RGB values between 0 and 255 to generate a color that resembles the "sunset hue", which is the same problem setting as `this tutorial `_. +Hence, familiarizing yourself with the tutorial on objective form widgets beforehand might offer a smoother understanding. + +How to Run Preferential Optimization +------------------------------------ + +In preferential optimization, we run two programs simultaneously: `generator.py`_ which executes parameter sampling or image generation, +and the Optuna Dashboard which provides a user interface for human evaluation. + +.. figure:: ./images/preferential-optimization/system-architecture.png + :alt: System Architecture + :align: center + :width: 800px + +To start, ensure you have the necessary packages installed. You can do this by running the following command in your terminal: + +.. code-block:: console + + $ pip install "optuna>=3.3.0" "optuna-dashboard>=0.13.0b1" pillow botorch + +Run a Python script below which you copied from `generator.py`_. + +.. code-block:: console + + $ python generator.py + +Then run a following command to launch Optuna Dashboard in a separate process. + +.. code-block:: console + + $ optuna-dashboard sqlite:///example.db --artifact-dir ./artifact + +In the command, the storage is set to ``sqlite:///example.db`` to persist Optuna's trial history. +To store the artifacts (output images), ``--artifact-dir ./artifact`` is specified. + +.. code-block:: console + + Listening on http://127.0.0.1:8080/ + Hit Ctrl-C to quit. + +When you run the command, you will see a message like the one above. +Please open `http://127.0.0.1:8080/dashboard/ `_ in your browser, then you can see the Optuna Dashboard as follows: + +.. figure:: ./images/preferential-optimization/anim.gif + :alt: GIF animation for preferential optimization + :align: center + :width: 800px + + Selecting the least sunset-like color from four trials to report human preferences. + + +Script Explanation +------------------ + +Here, we specify the SQLite database URL and setup the artifact store, a filesystem to store images generated during the trial. + +.. code-block:: python + :linenos: + + STORAGE_URL = "sqlite:///example.db" + artifact_path = os.path.join(os.path.dirname(__file__), "artifact") + artifact_store = FileSystemArtifactStore(base_path=artifact_path) + os.makedirs(artifact_path, exist_ok=True) + +Within the ``main()`` function, we initialize the study with necessary parameters, including specifying the preferential sampler. +ote that the ``Study`` and ``Sampler`` instantiated here are different from the conventional Optuna's ``Study``` and the ``Sampler``. +Preferential optimization relies solely on the comparison results between trials, and there are no absolute evaluation values for each trial. +Therefore, it is necessary to create dedicated ``Study`` and ``Sampler`` objects. + +.. code-block:: python + :linenos: + + from optuna_dashboard.preferential import create_study + from optuna_dashboard.preferential.samplers.gp import PreferentialGPSampler + + study = create_study( + n_generate=5, + study_name="Preferential Optimization", + storage=STORAGE_URL, + sampler=PreferentialGPSampler(), + load_if_exists=True, + ) + +Then, we create a loop that continuously checks if new trials should be generated, awaiting human evaluation if not. +Within the while loop, new trials are generated if the condition :meth:`~optuna_dashboard.preferential.PreferentialStudy.should_generate` returns ``True``. +For each trial, RGB values are sampled, and an image is generated with these values. +The image is saved temporarily, uploaded to artifact store, and then saved a Markdown note using :func:`~optuna_dashboard.save_note`. + +.. code-block:: python + :linenos: + + while True: + # If study.should_generate() returns False, the generator waits for human evaluation. + if not study.should_generate(): + time.sleep(0.1) # Avoid busy-loop + continue + + trial = study.ask() + # Ask new parameters + r = trial.suggest_int("r", 0, 255) + g = trial.suggest_int("g", 0, 255) + b = trial.suggest_int("b", 0, 255) + + # Generate an image + image_path = os.path.join(tmpdir, f"sample-{trial.number}.png") + image = Image.new("RGB", (320, 240), color=(r, g, b)) + image.save(image_path) + + # Upload to Artifact store + artifact_id = upload_artifact(trial, image_path, artifact_store) + trial.set_user_attr("artifact_id", artifact_id) + print("RGB:", (r, g, b)) + + # Save a Markdown note + note = textwrap.dedent( + f"""\ + ![generated-image]({get_artifact_path(trial, artifact_id)}) + + (R, G, B) = ({r}, {g}, {b}) + """ + ) + +.. _generator.py: https://github.com/optuna/optuna-dashboard/blob/main/examples/preferential-optimization/generator.py diff --git a/examples/hitl/main.py b/examples/hitl/main.py index 1557ae21..7844bc7a 100644 --- a/examples/hitl/main.py +++ b/examples/hitl/main.py @@ -4,17 +4,19 @@ import time from typing import NoReturn import optuna +from optuna.artifacts import FileSystemArtifactStore +from optuna.artifacts import upload_artifact from optuna.trial import TrialState from optuna_dashboard import ChoiceWidget from optuna_dashboard import register_objective_form_widgets from optuna_dashboard import save_note from optuna_dashboard.artifact import get_artifact_path -from optuna_dashboard.artifact import upload_artifact -from optuna_dashboard.artifact.file_system import FileSystemBackend from PIL import Image -def suggest_and_generate_image(study: optuna.Study, artifact_backend: FileSystemBackend) -> None: +def suggest_and_generate_image( + study: optuna.Study, artifact_store: FileSystemArtifactStore +) -> None: # 1. Ask new parameters trial = study.ask() r = trial.suggest_int("r", 0, 255) @@ -27,7 +29,7 @@ def suggest_and_generate_image(study: optuna.Study, artifact_backend: FileSystem image.save(image_path) # 3. Upload Artifact - artifact_id = upload_artifact(artifact_backend, trial, image_path) + artifact_id = upload_artifact(trial, image_path, artifact_store) artifact_path = get_artifact_path(trial, artifact_id) # 4. Save Note @@ -41,7 +43,7 @@ def suggest_and_generate_image(study: optuna.Study, artifact_backend: FileSystem save_note(trial, note) -def start_optimization(artifact_backend: FileSystemBackend) -> NoReturn: +def start_optimization(artifact_store: FileSystemArtifactStore) -> NoReturn: # 1. Create Study study = optuna.create_study( study_name="Human-in-the-loop Optimization", @@ -72,7 +74,7 @@ def start_optimization(artifact_backend: FileSystemBackend) -> NoReturn: if len(running_trials) >= n_batch: time.sleep(1) # Avoid busy-loop continue - suggest_and_generate_image(study, artifact_backend) + suggest_and_generate_image(study, artifact_store) def main() -> NoReturn: @@ -80,7 +82,7 @@ def main() -> NoReturn: # 1. Create Artifact Store artifact_path = os.path.join(os.path.dirname(__file__), "artifact") - artifact_backend = FileSystemBackend(base_path=artifact_path) + artifact_store = FileSystemArtifactStore(artifact_path) if not os.path.exists(artifact_path): os.mkdir(artifact_path) @@ -89,7 +91,7 @@ def main() -> NoReturn: os.mkdir(tmp_path) # 2. Run optimize loop - start_optimization(artifact_backend) + start_optimization(artifact_store) if __name__ == "__main__": diff --git a/examples/preferential-optimization/generator.py b/examples/preferential-optimization/generator.py index d8f7d3bd..66be1765 100644 --- a/examples/preferential-optimization/generator.py +++ b/examples/preferential-optimization/generator.py @@ -6,10 +6,10 @@ import textwrap import time from typing import NoReturn +from optuna.artifacts import FileSystemArtifactStore +from optuna.artifacts import upload_artifact from optuna_dashboard import save_note from optuna_dashboard.artifact import get_artifact_path -from optuna_dashboard.artifact import upload_artifact -from optuna_dashboard.artifact.file_system import FileSystemBackend from optuna_dashboard.preferential import create_study from optuna_dashboard.preferential.samplers.gp import PreferentialGPSampler from PIL import Image @@ -17,7 +17,7 @@ from PIL import Image STORAGE_URL = "sqlite:///example.db" artifact_path = os.path.join(os.path.dirname(__file__), "artifact") -artifact_backend = FileSystemBackend(base_path=artifact_path) +artifact_store = FileSystemArtifactStore(base_path=artifact_path) os.makedirs(artifact_path, exist_ok=True) @@ -32,7 +32,7 @@ def main() -> NoReturn: with tempfile.TemporaryDirectory() as tmpdir: while True: - # If n_comparison "best" trials (that are not reported bad) exists, + # If study.should_generate() returns False, # the generator waits for human evaluation. if not study.should_generate(): time.sleep(0.1) # Avoid busy-loop @@ -50,7 +50,7 @@ def main() -> NoReturn: image.save(image_path) # 3. Upload Artifact - artifact_id = upload_artifact(artifact_backend, trial, image_path) + artifact_id = upload_artifact(trial, image_path, artifact_store) trial.set_user_attr("artifact_id", artifact_id) print("RGB:", (r, g, b)) diff --git a/optuna_dashboard/preferential/samplers/gp.py b/optuna_dashboard/preferential/samplers/gp.py index b90a1de4..8349b1a4 100644 --- a/optuna_dashboard/preferential/samplers/gp.py +++ b/optuna_dashboard/preferential/samplers/gp.py @@ -1,155 +1,41 @@ from __future__ import annotations import math -from math import erfc from typing import Any +from typing import Callable -from botorch.acquisition.analytic import LogExpectedImprovement -from botorch.models.gpytorch import GPyTorchModel -from botorch.optim import optimize_acqf +import botorch.acquisition.analytic +import botorch.models.model +import botorch.optim +import botorch.posteriors.gpytorch import gpytorch.constraints import gpytorch.kernels -import gpytorch.likelihoods.gaussian_likelihood -from gpytorch.likelihoods.gaussian_likelihood import GaussianLikelihood -from gpytorch.likelihoods.gaussian_likelihood import Interval from gpytorch.likelihoods.gaussian_likelihood import Prior -from gpytorch.models.exact_gp import ExactGP -import gpytorch.module -from linear_operator.operators import DiagLinearOperator -from linear_operator.operators import LinearOperator -from linear_operator.utils.errors import NotPSDError import numpy as np import optuna -from optuna import distributions -from optuna import Study -from optuna._transform import _SearchSpaceTransform -from optuna.distributions import BaseDistribution -from optuna.search_space import IntersectionSearchSpace -from optuna.trial import FrozenTrial -import pyro -import pyro.infer.autoguide -import pyro.infer.mcmc -from scipy.special import erfcinv +import optuna._transform import torch from torch import Tensor from .._system_attrs import get_preferences -class _WeightedGaussianLikelihood(GaussianLikelihood): - def __init__( - self, - weights: torch.Tensor | None = None, - noise_prior: Prior | None = None, - noise_constraint: Interval | None = None, - batch_shape: torch.Size = torch.Size(), - **kwargs: Any, - ) -> None: - super().__init__( - noise_prior=noise_prior, - noise_constraint=noise_constraint, - batch_shape=batch_shape, - **kwargs, - ) - self.weights = weights - - def _shaped_noise_covar( - self, base_shape: torch.Size, *params: Any, **kwargs: Any - ) -> Tensor | LinearOperator: - assert self.weights is not None - assert base_shape[-1] == self.weights.shape[-1] - return DiagLinearOperator(1.0 / self.weights) * super()._shaped_noise_covar( - base_shape, *params, **kwargs - ) - - -def _sample_y( - preferences: np.ndarray, - cov_X_X: np.ndarray, - obs_noise_var: float, - cycles: int, - initial_sample: np.ndarray, - rng: np.random.RandomState, -) -> np.ndarray: - # TODO: Refactor and write tests for this function. - - N = cov_X_X.shape[0] - M = len(preferences) - cov_X_X = cov_X_X + np.eye(N) * 1e-6 # Add jitter - cov_X_X_chol = np.linalg.cholesky(cov_X_X) - cov_X_X_inv = np.linalg.inv(cov_X_X) - - # (sI + A K A^T)^-1 = s^-1 I - s^-2 A(K^-1 + s^-1 A^T A)^-1 A^T - - schur = cov_X_X_inv.copy() - np.add.at(schur, (preferences[:, 0], preferences[:, 0]), 1.0 / (2 * obs_noise_var)) - np.add.at(schur, (preferences[:, 1], preferences[:, 1]), 1.0 / (2 * obs_noise_var)) - np.add.at(schur, (preferences[:, 0], preferences[:, 1]), -1.0 / (2 * obs_noise_var)) - np.add.at(schur, (preferences[:, 1], preferences[:, 0]), -1.0 / (2 * obs_noise_var)) - idx_M = np.arange(M) - - schur_inv = np.linalg.inv(schur) - - cov_diff_inv = schur_inv[:, preferences[:, 0]] - schur_inv[:, preferences[:, 1]] - cov_diff_inv = cov_diff_inv[preferences[:, 0], :] - cov_diff_inv[preferences[:, 1], :] - cov_diff_inv *= -1 / (2 * obs_noise_var) ** 2 - cov_diff_inv[idx_M, idx_M] += 1.0 / (2 * obs_noise_var) - - diffs = _orthants_MVN_Gibbs_sampling( - cov_diff_inv, - cycles=cycles, - initial_sample=initial_sample[:, 0] - initial_sample[:, 1], - rng=rng, - )[-1] - - random_ys = (cov_X_X_chol @ rng.randn(N))[preferences] + np.sqrt(obs_noise_var) * rng.randn( - M, 2 - ) - errors = diffs - (random_ys[:, 0] - random_ys[:, 1]) - cov_diff_inv_errors = cov_diff_inv @ errors - - AT_cov_diff_inv_errors = np.zeros((N,)) - np.add.at(AT_cov_diff_inv_errors, preferences[:, 0], cov_diff_inv_errors) - np.add.at(AT_cov_diff_inv_errors, preferences[:, 1], -cov_diff_inv_errors) - - return ( - random_ys - + (cov_X_X @ AT_cov_diff_inv_errors)[preferences] - + obs_noise_var * np.array([[1, -1]]) * cov_diff_inv_errors[:, None] - ) - - -_SQRT2 = math.sqrt(2) - - -def _orthants_MVN_Gibbs_sampling( - cov_inv: np.ndarray, - cycles: int, - initial_sample: np.ndarray, - rng: np.random.RandomState, -) -> np.ndarray: +def _orthants_MVN_Gibbs_sampling(cov_inv: Tensor, cycles: int, initial_sample: Tensor) -> Tensor: dim = cov_inv.shape[0] assert cov_inv.shape == (dim, dim) - if initial_sample is None: - sample_chain = np.zeros(dim) - else: - sample_chain = initial_sample + sample_chain = initial_sample + conditional_std = torch.rsqrt(torch.diag(cov_inv)) + scaled_cov_inv = cov_inv / torch.diag(cov_inv)[:, None] - conditional_std = 1 / np.sqrt(np.diag(cov_inv)) - - scaled_cov_inv = cov_inv / np.c_[np.diag(cov_inv)] - - out = np.empty((cycles + 1, dim)) + out = torch.empty((cycles + 1, dim), dtype=torch.float64) out[0, :] = sample_chain for i in range(cycles): for j in range(dim): conditional_mean = sample_chain[j] - scaled_cov_inv[j] @ sample_chain sample_chain[j] = ( - _one_side_trunc_norm_sampling( - lower=-conditional_mean / conditional_std[j], rng=rng - ) + _one_side_trunc_norm_sampling(lower=-conditional_mean / conditional_std[j]) * conditional_std[j] + conditional_mean ) @@ -158,144 +44,234 @@ def _orthants_MVN_Gibbs_sampling( return out -def _one_side_trunc_norm_sampling(lower: float, rng: np.random.RandomState) -> float: - return erfcinv(rng.rand() * erfc(lower / _SQRT2)) * _SQRT2 +def _one_side_trunc_norm_sampling(lower: Tensor) -> Tensor: + if lower > 4.0: + r = torch.clamp_min(torch.rand(torch.Size(()), dtype=torch.float64), min=1e-300) + return (lower * lower - 2 * r.log()).sqrt() + else: + SQRT2 = math.sqrt(2) + r = torch.rand(torch.Size(()), dtype=torch.float64) * torch.erfc(lower / SQRT2) + while 1 - r == 1: + r = torch.rand(torch.Size(()), dtype=torch.float64) * torch.erfc(lower / SQRT2) + return torch.erfinv(1 - r) * SQRT2 -class _PreferentialGP(GPyTorchModel, ExactGP): - _num_outputs = 1 +_orthants_MVN_Gibbs_sampling_jit = torch.jit.script(_orthants_MVN_Gibbs_sampling) + +def _compute_cov_diff_diff_inv(preferences: Tensor, cov_x_x: Tensor, noise_var: Tensor) -> Tensor: + N = cov_x_x.shape[0] + M = preferences.shape[0] + + # (sI + A K A^T)^-1 = s^-1 I - s^-2 A(K^-1 + s^-1 A^T A)^-1 A^T + # (K^-1 + s^-1 A^T A)^-1 = K (I + s^-1 A^T A K)^-1 (To avoid computing K^-1) + + I_plus_sinv_AT_A_K = torch.eye(N, dtype=torch.float64) + A_K = cov_x_x[preferences[:, 0], :] - cov_x_x[preferences[:, 1], :] + I_plus_sinv_AT_A_K.index_add_(0, preferences[:, 0], A_K * (1 / noise_var)) + I_plus_sinv_AT_A_K.index_add_(0, preferences[:, 1], A_K * (-1 / noise_var)) + schur_inv: Tensor = torch.linalg.solve(I_plus_sinv_AT_A_K, cov_x_x, left=False) + cov_diff_diff_inv = schur_inv[:, preferences[:, 0]] - schur_inv[:, preferences[:, 1]] + cov_diff_diff_inv = ( + cov_diff_diff_inv[preferences[:, 0], :] - cov_diff_diff_inv[preferences[:, 1], :] + ) + cov_diff_diff_inv *= -1 / noise_var**2 + idx_M = torch.arange(M) + cov_diff_diff_inv[idx_M, idx_M] += 1.0 / noise_var + + return cov_diff_diff_inv + + +class _SampledGP(botorch.models.model.Model): def __init__( self, - kernel: gpytorch.kernels.Kernel, - noise_prior: Prior | None = None, - noise_constraint: Interval | None = None, + kernel_func: Callable[[Tensor, Tensor], Tensor], + x: Tensor, + preferences: Tensor, + noise_var: Tensor, + diff: Tensor, ) -> None: - GPyTorchModel.__init__(self) - likelihood = _WeightedGaussianLikelihood( - noise_prior=noise_prior, noise_constraint=noise_constraint + super().__init__() + self.kernel_func = kernel_func + self.x = x + self.preferences = preferences + self.diff = diff + self.noise_var = noise_var + self._cov_diff_diff_inv = _compute_cov_diff_diff_inv( + preferences=preferences, + cov_x_x=self.kernel_func(x, x), + noise_var=noise_var, ) - ExactGP.__init__(self, train_inputs=None, train_targets=None, likelihood=likelihood) - self.covar_module = kernel - self._last_params: dict[str, torch.Tensor] | None = None - self._last_mcmc_step_size: float | None = None + def posterior( + self, + X: Tensor, + output_indices: list[int] | None = None, + observation_noise: bool = False, + posterior_transform: Any | None = None, + **kwargs: Any, + ) -> botorch.posteriors.gpytorch.GPyTorchPosterior: + assert posterior_transform is None + assert output_indices is None + assert self.x.shape[-1] == X.shape[-1] - def _pyro_model(self, train_x: torch.Tensor, train_y: torch.Tensor) -> None: - # with gpytorch.settings.fast_computations(False, False, False): - sampled_model = self.pyro_sample_from_prior() + x_expanded = self.x.expand(X.shape[:-2] + (self.x.shape[-2], X.shape[-1])) - ys = sampled_model.likelihood(sampled_model.forward(train_x)) + cov_X_x = self.kernel_func(X, x_expanded) + cov_X_diff = cov_X_x[..., self.preferences[:, 0]] - cov_X_x[..., self.preferences[:, 1]] - pyro.sample("y", ys, obs=train_y) + mean = cov_X_diff @ (self._cov_diff_diff_inv @ self.diff) + cov = self.kernel_func(X, X) - cov_X_diff @ self._cov_diff_diff_inv @ cov_X_diff.transpose( + -1, -2 + ) + if observation_noise: + idx = torch.arange(cov.shape[-1]) + cov[..., idx, idx] += self.noise_var - def fit_mcmc( - self, X: torch.Tensor, preferences: torch.Tensor, cycles: int, rng: np.random.RandomState - ) -> None: + return botorch.posteriors.gpytorch.GPyTorchPosterior( + distribution=gpytorch.distributions.MultivariateNormal( + mean=mean, + covariance_matrix=cov, + ) + ) + + @property + def batch_shape(self) -> torch.Size: + return torch.Size() + + @property + def num_outputs(self) -> int: + return 1 + + +def _truncnorm_mean_var_logz(alpha: Tensor) -> tuple[Tensor, Tensor, Tensor]: + SQRT_HALF = math.sqrt(0.5) + SQRT_HALF_PI = math.sqrt(0.5 * math.pi) + logz = torch.special.log_ndtr(-alpha) + mean = 1 / (SQRT_HALF_PI * torch.special.erfcx(alpha * SQRT_HALF)) + var = 1 - mean * (mean - alpha) + return mean, var, logz + + +def _orthants_MVN_EP( + cov0: Tensor, preferences: Tensor, noise_var: Tensor, cycles: int +) -> tuple[Tensor, Tensor, Tensor]: + N = cov0.shape[0] + M = preferences.shape[0] + mu = torch.zeros(N, dtype=cov0.dtype) + cov = cov0.clone() + virtual_obs_a = [torch.tensor(0.0, dtype=cov0.dtype) for _ in range(M)] + virtual_obs_b = [torch.tensor(0.0, dtype=cov0.dtype) for _ in range(M)] + log_zs = torch.zeros(M, dtype=cov0.dtype) + + for _ in range(cycles): + for i in range(M): + pref_i = preferences[i, :] + mean1 = mu[pref_i[0]] - mu[pref_i[1]] + Sxy = cov[pref_i[0]] - cov[pref_i[1]] + var1 = Sxy[pref_i[0]] - Sxy[pref_i[1]] + + r0 = (1 - var1 * virtual_obs_a[i]).reciprocal() + var0 = var1 * r0 + mean0 = (mean1 + var1 * virtual_obs_b[i]) * r0 + + obs_var = var0 + noise_var + obs_sigma = torch.sqrt(obs_var) + alpha = -mean0 / torch.clamp_min(obs_sigma, min=1e-20) + mean_norm, var_norm, logz = _truncnorm_mean_var_logz(alpha) + + kalman_factor = var0 / torch.clamp_min(obs_var, min=1e-20) + mean2 = mean0 + obs_sigma * mean_norm * kalman_factor + var2 = kalman_factor * (noise_var + var_norm * var0) + + var1_var2_inv = torch.clamp_min(var1 * var2, min=1e-20).reciprocal() + db = (mean1 * var2 - mean2 * var1) * var1_var2_inv + da = (var1 - var2) * var1_var2_inv + virtual_obs_b[i] = virtual_obs_b[i] + db + virtual_obs_a[i] = virtual_obs_a[i] + da + + dr = (1 + var1 * da).reciprocal() + mu = mu - Sxy * ((db + mean1 * da) * dr) + cov = cov - (Sxy[:, None] * (da * dr)) @ Sxy[None, :] + log_zs[i] = logz + return mu, cov, torch.sum(log_zs) + + +_orthants_MVN_EP_jit = torch.jit.script(_orthants_MVN_EP) + + +class _PreferentialGP: + def __init__(self, kernel: gpytorch.kernels.Kernel, noise_prior: Prior, dims: int) -> None: + self.kernel = kernel + self.noise_prior = noise_prior + self.dims = dims + + self.diff = torch.empty((0,), dtype=torch.float64, requires_grad=False) + self.log_noise = torch.nn.Parameter( + torch.tensor(0.0, dtype=torch.float64), requires_grad=True + ) + + def fit_params_EP(self, X: Tensor, preferences: Tensor) -> None: if len(preferences) == 0: - # Skip actual MCMC computation - self.set_train_data( - inputs=torch.empty((0, X.shape[-1])), - targets=torch.empty((0,)), - strict=False, - ) - self.likelihood.weights = torch.empty((0,)) - else: - dtype = torch.float64 + return + tolerance = 1e-3 + max_iter = 100 - cnt = torch.bincount(preferences.reshape(-1)) - mask = cnt > 0 - train_x = X[mask] - weights = cnt[mask] + optim = torch.optim.LBFGS([*self.kernel.parameters(), self.log_noise]) - assert isinstance(self.likelihood, _WeightedGaussianLikelihood) - self.likelihood.weights = weights + last_params = [p.detach().clone() for p in optim.param_groups[0]["params"]] + for _ in range(max_iter): - preferences_np = preferences.detach().numpy() + def closure() -> Tensor: + optim.zero_grad() + noise = self.log_noise.exp() + cov0 = self.kernel.forward(X, X).to_dense() + _, _, logz = _orthants_MVN_EP_jit(cov0, preferences, noise, cycles=2) - all_ys_np = np.zeros((len(preferences), 2)) - train_y = torch.zeros( - ( - len( - train_x, - ) - ), - dtype=dtype, + loss = -logz - self.noise_prior.log_prob(noise) + for _, _, prior, param, _ in self.kernel.named_priors(): + loss = loss - prior.log_prob(param(self.kernel)).sum() + + loss.backward() + return loss + + optim.step(closure) + + # Check for convergence + params = optim.param_groups[0]["params"] + for p_old, p_new in zip(last_params, params): + if torch.max(torch.abs(p_old - p_new)) > tolerance: + break + else: + break + last_params = [p.detach().clone() for p in params] + + def sample_gp(self, x: Tensor, preferences: Tensor) -> _SampledGP: + self.fit_params_EP(x, preferences) + + with torch.no_grad(): + cov_diff_diff_inv = _compute_cov_diff_diff_inv( + preferences=preferences, + cov_x_x=self.kernel(x, x).to_dense(), + noise_var=self.log_noise.exp(), ) - nuts = pyro.infer.mcmc.NUTS( - model=self._pyro_model, - init_strategy=pyro.infer.autoguide.init_to_sample, - step_size=self._last_mcmc_step_size or 1.0, + original_diff_size = len(self.diff) + self.diff.resize_(len(preferences)) + self.diff[original_diff_size:] = 0.0 + + self.diff = _orthants_MVN_Gibbs_sampling_jit( + cov_inv=cov_diff_diff_inv, + initial_sample=self.diff, + cycles=20, + )[-1] + return _SampledGP( + kernel_func=lambda x1, x2: self.kernel(x1, x2).to_dense(), + x=x, + preferences=preferences, + noise_var=self.log_noise.exp(), + diff=self.diff, ) - warmup_steps = max(0, cycles - 2) - nuts.setup(warmup_steps=warmup_steps, train_x=train_x, train_y=train_y) - - raw_params = self._last_params or nuts.initial_params - for i in range(cycles): - params = { - name: nuts.transforms[name].inv(value) for name, value in raw_params.items() - } - _set_params(self, params) - self.set_train_data(train_x, train_y, strict=False) - all_ys_np = _sample_y( - preferences=preferences_np, - cov_X_X=self.covar_module(train_x).detach().numpy(), - obs_noise_var=float(self.likelihood.noise_covar.noise), - cycles=10, - initial_sample=all_ys_np, - rng=rng, - ) - ys_sum_np = np.zeros((len(X),)) - np.add.at(ys_sum_np, preferences_np.reshape(-1), all_ys_np.reshape(-1)) - ys_sum = torch.from_numpy(ys_sum_np) - train_y[:] = ys_sum[mask] / cnt[mask] - nuts.clear_cache() - try: - raw_params = nuts.sample(raw_params) - except NotPSDError: - nuts.cleanup() - nuts = pyro.infer.mcmc.NUTS( - model=self._pyro_model, - init_strategy=pyro.infer.autoguide.init_to_sample, - step_size=self._last_mcmc_step_size or 1.0, - ) - nuts.setup(warmup_steps=warmup_steps, train_x=train_x, train_y=train_y) - raw_params = nuts.initial_params - - params = {name: nuts.transforms[name].inv(value) for name, value in raw_params.items()} - self.set_train_data(train_x, train_y, strict=False) - _set_params(self, params) - - self._last_params = raw_params - self._last_mcmc_step_size = nuts.step_size - nuts.cleanup() - - def forward(self, x: torch.Tensor) -> gpytorch.distributions.MultivariateNormal: - mean_module = gpytorch.means.ZeroMean() - return gpytorch.distributions.MultivariateNormal( - mean_module(x), - self.covar_module(x), - ) - - -def _set_params( - module: gpytorch.Module, - params_dict: dict[str, torch.Tensor], - memo: set | None = None, - prefix: str = "", -) -> None: - if memo is None: - memo = set() - if hasattr(module, "_priors"): - for name, (prior, closure, setting_closure) in module._priors.items(): - if prior is not None and prior not in memo: - memo.add(prior) - setting_closure(module, params_dict[prefix + ("." if prefix else "") + name]) - - for mname, module_ in module.named_children(): - submodule_prefix = prefix + ("." if prefix else "") + mname - _set_params(module_, params_dict, memo=memo, prefix=submodule_prefix) class PreferentialGPSampler(optuna.samplers.BaseSampler): @@ -306,18 +282,16 @@ class PreferentialGPSampler(optuna.samplers.BaseSampler): noise_prior: Prior | None = None, independent_sampler: optuna.samplers.BaseSampler | None = None, seed: int | None = None, - device: torch.device | None = None, ) -> None: - self._rng = np.random.RandomState(seed) - self._search_space = IntersectionSearchSpace() - self.kernel = kernel - self.noise_prior = noise_prior - self.independent_sampler = independent_sampler or optuna.samplers.RandomSampler( - seed=self._rng.randint(2**32), - ) - self.device = device or torch.device("cpu") + self.noise_prior = noise_prior or gpytorch.priors.GammaPrior(5.0, 50.0) + self._rng = np.random.RandomState(seed) + self.independent_sampler = independent_sampler or optuna.samplers.RandomSampler( + seed=self._rng.randint(2**32) + ) + + self._search_space = optuna.search_space.IntersectionSearchSpace() self._gp: _PreferentialGP | None = None def reseed_rng(self) -> None: @@ -325,75 +299,64 @@ class PreferentialGPSampler(optuna.samplers.BaseSampler): self._rng = np.random.RandomState() def infer_relative_search_space( - self, study: Study, trial: FrozenTrial - ) -> dict[str, BaseDistribution]: + self, study: optuna.Study, trial: optuna.trial.FrozenTrial + ) -> dict[str, optuna.distributions.BaseDistribution]: return self._search_space.calculate(study) def sample_relative( self, - study: Study, - trial: FrozenTrial, - search_space: dict[str, BaseDistribution], + study: optuna.Study, + trial: optuna.trial.FrozenTrial, + search_space: dict[str, optuna.distributions.BaseDistribution], ) -> dict[str, Any]: + preferences = get_preferences(study.system_attrs) + if len(preferences) == 0: + return {} + + trials = study.get_trials(deepcopy=False) + trials_with_preference = list({t for (b, w) in preferences for t in (b, w)}) + ids = {t: i for i, t in enumerate(trials_with_preference)} + + trans = optuna._transform._SearchSpaceTransform( + search_space, transform_log=True, transform_step=True, transform_0_1=True + ) + params = torch.tensor( + np.array([trans.transform(trials[t].params) for t in trials_with_preference]), + dtype=torch.float64, + ) + pref_ids = torch.tensor([[ids[b], ids[w]] for b, w in preferences], dtype=torch.int32) with torch.random.fork_rng(): torch.manual_seed(self._rng.randint(2**32)) - pyro.set_rng_seed(self._rng.randint(2**32)) - if len(search_space) == 0: - return {} - - preferences = get_preferences(study.system_attrs) - trials = study.get_trials(deepcopy=False) - if len(preferences) == 0: - return {} - - trans = _SearchSpaceTransform( - search_space, transform_log=True, transform_step=True, transform_0_1=True - ) - dims = len(trans.bounds) self._gp = self._gp or _PreferentialGP( kernel=self.kernel or gpytorch.kernels.MaternKernel( - nu=2.5, - ard_num_dims=dims, - lengthscale_prior=gpytorch.priors.GammaPrior(3.0, 6.0), - lengthscale_constraint=gpytorch.constraints.Positive(), + nu=1.5, + ard_num_dims=len(trans.bounds), + lengthscale_prior=gpytorch.priors.GammaPrior(5.0, 10.0), + lengthscale_constraint=gpytorch.constraints.GreaterThan( + 0.0, + transform=torch.exp, + inv_transform=torch.log, + ), ), - noise_prior=self.noise_prior or gpytorch.priors.GammaPrior(1.1, 2.0), - noise_constraint=gpytorch.constraints.Positive(), + noise_prior=self.noise_prior, + dims=len(trans.bounds), ) + if self._gp.dims != len(trans.bounds): + raise NotImplementedError( + "The search space has changed. " + "Dynamic search space is not supported in PreferentialGPSampler." + ) - ids: dict[int, int] = {} - params: list[torch.Tensor] = [] - pref_ids: list[tuple[int, int]] = [] - - for better, worse in preferences: - for t in (better, worse): - if t not in ids: - ids[t] = len(ids) - params.append(trans.transform(trials[t].params)) - pref_ids.append((ids[better], ids[worse])) - dtype = torch.float64 - - params_torch = torch.tensor(np.array(params), dtype=dtype, device=self.device) - pref_ids_torch = torch.tensor( - np.array(pref_ids), - dtype=torch.int32, - device=self.device, - ) - self._gp.fit_mcmc(params_torch, pref_ids_torch, cycles=10, rng=self._rng) - self._gp.eval() - scores = self._gp(params_torch).mean - - best_f = torch.max(scores) - - acqf = LogExpectedImprovement( - model=self._gp, - best_f=best_f, + sampled_gp = self._gp.sample_gp(params, pref_ids) + acqf = botorch.acquisition.analytic.LogExpectedImprovement( + model=sampled_gp, + best_f=torch.max(sampled_gp.posterior(params[:, None, :]).mean), ) # TODO: Make it possible to apply it on categorical variables - candidates, _ = optimize_acqf( + candidates, _ = botorch.optim.optimize_acqf( acq_function=acqf, bounds=torch.from_numpy(trans.bounds.T), q=1, @@ -407,10 +370,10 @@ class PreferentialGPSampler(optuna.samplers.BaseSampler): def sample_independent( self, - study: Study, - trial: FrozenTrial, + study: optuna.Study, + trial: optuna.trial.FrozenTrial, param_name: str, - param_distribution: distributions.BaseDistribution, + param_distribution: optuna.distributions.BaseDistribution, ) -> Any: return self.independent_sampler.sample_independent( study, trial, param_name, param_distribution diff --git a/optuna_dashboard/ts/components/ArtifactCardMedia.tsx b/optuna_dashboard/ts/components/ArtifactCardMedia.tsx new file mode 100644 index 00000000..a994aa41 --- /dev/null +++ b/optuna_dashboard/ts/components/ArtifactCardMedia.tsx @@ -0,0 +1,41 @@ +import React, { FC } from "react" +import { + ThreejsArtifactViewer, + isThreejsArtifact, +} from "./ThreejsArtifactViewer" +import InsertDriveFileIcon from "@mui/icons-material/InsertDriveFile" +import { CardMedia } from "@mui/material" + +export const ArtifactCardMedia: FC<{ + artifact: Artifact + urlPath: string + height: string +}> = ({ artifact, urlPath, height }) => { + if (isThreejsArtifact(artifact)) { + return ( + + ) + } else if (artifact.mimetype.startsWith("audio")) { + return ( + + ) + } else if (artifact.mimetype.startsWith("image")) { + return ( + + ) + } + return +} diff --git a/optuna_dashboard/ts/components/ThreejsArtifactViewer.tsx b/optuna_dashboard/ts/components/ThreejsArtifactViewer.tsx index 536f1d91..a88b9d2c 100644 --- a/optuna_dashboard/ts/components/ThreejsArtifactViewer.tsx +++ b/optuna_dashboard/ts/components/ThreejsArtifactViewer.tsx @@ -1,10 +1,17 @@ import * as THREE from "three" -import React, { useEffect, useState } from "react" +import React, { useEffect, useState, ReactNode } from "react" import { Canvas } from "@react-three/fiber" import { GizmoHelper, GizmoViewport, OrbitControls } from "@react-three/drei" import { STLLoader } from "three/examples/jsm/loaders/STLLoader" import { Rhino3dmLoader } from "three/examples/jsm/loaders/3DMLoader" import { PerspectiveCamera } from "three" +import { Modal, Box } from "@mui/material" + +export const isThreejsArtifact = (artifact: Artifact): boolean => { + return ( + artifact.filename.endsWith(".stl") || artifact.filename.endsWith(".3dm") + ) +} interface ThreejsArtifactViewerProps { src: string @@ -109,3 +116,48 @@ export const ThreejsArtifactViewer: React.FC = ( ) } + +export const useThreejsArtifactModal = (): [ + (path: string, artifact: Artifact) => void, + () => ReactNode +] => { + const [open, setOpen] = useState(false) + const [target, setTarget] = useState<[string, Artifact | null]>(["", null]) + + const openModal = (artifactUrlPath: string, artifact: Artifact) => { + setTarget([artifactUrlPath, artifact]) + setOpen(true) + } + + const renderDeleteStudyDialog = () => { + return ( + { + setOpen(false) + setTarget(["", null]) + }} + > + + + + + ) + } + return [openModal, renderDeleteStudyDialog] +} diff --git a/optuna_dashboard/ts/components/TrialArtifactCards.tsx b/optuna_dashboard/ts/components/TrialArtifactCards.tsx new file mode 100644 index 00000000..3e15f7ed --- /dev/null +++ b/optuna_dashboard/ts/components/TrialArtifactCards.tsx @@ -0,0 +1,233 @@ +import React, { + ChangeEventHandler, + DragEventHandler, + FC, + MouseEventHandler, + useRef, + useState, +} from "react" +import { + Typography, + Box, + useTheme, + IconButton, + Card, + CardContent, + CardActionArea, +} from "@mui/material" +import UploadFileIcon from "@mui/icons-material/UploadFile" +import DownloadIcon from "@mui/icons-material/Download" +import DeleteIcon from "@mui/icons-material/Delete" +import FullscreenIcon from "@mui/icons-material/Fullscreen" + +import { actionCreator } from "../action" +import { useDeleteArtifactDialog } from "./DeleteArtifactDialog" +import { + useThreejsArtifactModal, + isThreejsArtifact, +} from "./ThreejsArtifactViewer" +import { ArtifactCardMedia } from "./ArtifactCardMedia" + +export const TrialArtifactCards: FC<{ trial: Trial }> = ({ trial }) => { + const theme = useTheme() + const [openDeleteArtifactDialog, renderDeleteArtifactDialog] = + useDeleteArtifactDialog() + const [openThreejsArtifactModal, renderThreejsArtifactModal] = + useThreejsArtifactModal() + + const width = "200px" + const height = "150px" + + return ( + <> + + Artifacts + + + {trial.artifacts.map((artifact) => { + const urlPath = `/artifacts/${trial.study_id}/${trial.trial_id}/${artifact.artifact_id}` + return ( + + + + + {artifact.filename} + + {isThreejsArtifact(artifact) ? ( + { + openThreejsArtifactModal(urlPath, artifact) + }} + > + + + ) : null} + { + openDeleteArtifactDialog( + trial.study_id, + trial.trial_id, + artifact + ) + }} + > + + + + + + + + ) + })} + + + {renderDeleteArtifactDialog()} + {renderThreejsArtifactModal()} + + ) +} + +const TrialArtifactUploader: FC<{ + trial: Trial + width: string + height: string +}> = ({ trial, width, height }) => { + const theme = useTheme() + const action = actionCreator() + const [dragOver, setDragOver] = useState(false) + + if (trial.state !== "Running" && trial.state !== "Waiting") { + return null + } + const inputRef = useRef(null) + const handleClick: MouseEventHandler = () => { + if (!inputRef || !inputRef.current) { + return + } + inputRef.current.click() + } + const handleOnChange: ChangeEventHandler = (e) => { + const files = e.target.files + if (files === null) { + return + } + action.uploadArtifact(trial.study_id, trial.trial_id, files[0]) + } + const handleDrop: DragEventHandler = (e) => { + e.stopPropagation() + e.preventDefault() + const files = e.dataTransfer.files + setDragOver(false) + for (let i = 0; i < files.length; i++) { + action.uploadArtifact(trial.study_id, trial.trial_id, files[i]) + } + } + const handleDragOver: DragEventHandler = (e) => { + e.stopPropagation() + e.preventDefault() + e.dataTransfer.dropEffect = "copy" + setDragOver(true) + } + const handleDragLeave: DragEventHandler = (e) => { + e.stopPropagation() + e.preventDefault() + e.dataTransfer.dropEffect = "copy" + setDragOver(false) + } + return ( + + + + + + Upload a New File + + Drag your file here or click to browse. + + + + + ) +} diff --git a/optuna_dashboard/ts/components/TrialList.tsx b/optuna_dashboard/ts/components/TrialList.tsx index 4aceb619..6922003a 100644 --- a/optuna_dashboard/ts/components/TrialList.tsx +++ b/optuna_dashboard/ts/components/TrialList.tsx @@ -1,13 +1,4 @@ -import React, { - ChangeEventHandler, - DragEventHandler, - FC, - MouseEventHandler, - ReactNode, - useMemo, - useRef, - useState, -} from "react" +import React, { FC, ReactNode, useMemo } from "react" import { Typography, Box, @@ -16,11 +7,6 @@ import { IconButton, Menu, MenuItem, - Card, - CardContent, - CardMedia, - CardActionArea, - Modal, } from "@mui/material" import Chip from "@mui/material/Chip" import Divider from "@mui/material/Divider" @@ -32,11 +18,6 @@ import ListSubheader from "@mui/material/ListSubheader" import FilterListIcon from "@mui/icons-material/FilterList" import CheckBoxOutlineBlankIcon from "@mui/icons-material/CheckBoxOutlineBlank" import CheckBoxIcon from "@mui/icons-material/CheckBox" -import UploadFileIcon from "@mui/icons-material/UploadFile" -import DownloadIcon from "@mui/icons-material/Download" -import DeleteIcon from "@mui/icons-material/Delete" -import FullscreenIcon from "@mui/icons-material/Fullscreen" -import InsertDriveFileIcon from "@mui/icons-material/InsertDriveFile" import StopCircleIcon from "@mui/icons-material/StopCircle" import { TrialNote } from "./Note" @@ -45,9 +26,8 @@ import ListItemIcon from "@mui/material/ListItemIcon" import { useRecoilValue } from "recoil" import { artifactIsAvailable } from "../state" import { actionCreator } from "../action" -import { useDeleteArtifactDialog } from "./DeleteArtifactDialog" import { TrialFormWidgets } from "./TrialFormWidgets" -import { ThreejsArtifactViewer } from "./ThreejsArtifactViewer" +import { TrialArtifactCards } from "./TrialArtifactCards" const states: TrialState[] = [ "Complete", @@ -319,454 +299,11 @@ export const TrialListDetail: FC<{ value !== null ? renderInfo(key, value) : null )} - {artifactEnabled && } + {artifactEnabled && } ) } -const TrialArtifact: FC<{ trial: Trial }> = ({ trial }) => { - const theme = useTheme() - const action = actionCreator() - const [openDeleteArtifactDialog, renderDeleteArtifactDialog] = - useDeleteArtifactDialog() - const [dragOver, setDragOver] = useState(false) - const [open3dModelViewer, setOpen3dModelViewer] = useState<{ - [key: string]: boolean - }>({}) - - const width = "200px" - const height = "150px" - - const inputRef = useRef(null) - const handleClick: MouseEventHandler = () => { - if (!inputRef || !inputRef.current) { - return - } - inputRef.current.click() - } - const handleOnChange: ChangeEventHandler = (e) => { - const files = e.target.files - if (files === null) { - return - } - action.uploadArtifact(trial.study_id, trial.trial_id, files[0]) - } - const handleDrop: DragEventHandler = (e) => { - e.stopPropagation() - e.preventDefault() - const files = e.dataTransfer.files - setDragOver(false) - for (let i = 0; i < files.length; i++) { - action.uploadArtifact(trial.study_id, trial.trial_id, files[i]) - } - } - const handleDragOver: DragEventHandler = (e) => { - e.stopPropagation() - e.preventDefault() - e.dataTransfer.dropEffect = "copy" - setDragOver(true) - } - const handleDragLeave: DragEventHandler = (e) => { - e.stopPropagation() - e.preventDefault() - e.dataTransfer.dropEffect = "copy" - setDragOver(false) - } - - return ( - <> - - Artifacts - - - {trial.artifacts.map((a) => { - if (a.mimetype.startsWith("image")) { - return ( - - - - - {a.filename} - - { - openDeleteArtifactDialog( - trial.study_id, - trial.trial_id, - a - ) - }} - > - - - - - - - - ) - } else if ( - a.filename.endsWith(".stl") || - a.filename.endsWith(".3dm") - ) { - return ( - - - - - - - {a.filename} - - { - setOpen3dModelViewer(() => { - const obj = { ...open3dModelViewer } - obj[a.artifact_id] = true - return obj - }) - }} - > - - - { - setOpen3dModelViewer(() => { - const obj = { ...open3dModelViewer } - obj[a.artifact_id] = false - return obj - }) - }} - > - - - - - { - openDeleteArtifactDialog( - trial.study_id, - trial.trial_id, - a - ) - }} - > - - - - - - - - ) - } else if (a.mimetype.startsWith("audio")) { - return ( - - - - - - - {a.filename} - - { - openDeleteArtifactDialog( - trial.study_id, - trial.trial_id, - a - ) - }} - > - - - - - - - - ) - } else { - return ( - - - - - - - {a.filename} - - { - openDeleteArtifactDialog( - trial.study_id, - trial.trial_id, - a - ) - }} - > - - - - - - - - ) - } - })} - {trial.state === "Running" || trial.state === "Waiting" ? ( - - - - - - Upload a New File - - Drag your file here or click to browse. - - - - - ) : null} - - {renderDeleteArtifactDialog()} - - ) -} - const getTrialListLink = ( studyId: number, exclude: TrialState[],