LightningBO 0.1.0__tar.gz

lightningbo-0.1.0/LICENSE

@@ -0,0 +1,208 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+1. Definitions.
+
+   "License" shall mean the terms and conditions for use, reproduction,
+   and distribution as defined by Sections 1 through 9 of this document.
+
+   "Licensor" shall mean the copyright owner or entity authorized by
+   the copyright owner that is granting the License.
+
+   "Legal Entity" shall mean the union of the acting entity and all
+   other entities that control, are controlled by, or are under common
+   control with that entity. For the purposes of this definition,
+   "control" means (i) the power, direct or indirect, to cause the
+   direction or management of such entity, whether by contract or
+   otherwise, or (ii) ownership of fifty percent (50%) or more of the
+   outstanding shares, or (iii) beneficial ownership of such entity.
+
+   "You" (or "Your") shall mean an individual or Legal Entity
+   exercising permissions granted by this License.
+
+   "Source" form shall mean the preferred form for making modifications,
+   including but not limited to software source code, documentation
+   source, and configuration files.
+
+   "Object" form shall mean any form resulting from mechanical
+   transformation or translation of a Source form, including but
+   not limited to compiled object code, generated documentation,
+   and conversions to other media types.
+
+   "Work" shall mean the work of authorship, whether in Source or
+   Object form, made available under the License, as indicated by a
+   copyright notice that is included in or attached to the work
+   (an example is provided in the Appendix below).
+
+   "Derivative Works" shall mean any work, whether in Source or Object
+   form, that is based on (or derived from) the Work and for which the
+   editorial revisions, annotations, elaborations, or other modifications
+   represent, as a whole, an original work of authorship. For the purposes
+   of this License, Derivative Works shall not include works that remain
+   separable from, or merely link (or bind by name) to the interfaces of,
+   the Work and Derivative Works thereof.
+
+   "Contribution" shall mean any work of authorship, including
+   the original version of the Work and any modifications or additions
+   to that Work or Derivative Works thereof, that is intentionally
+   submitted to Licensor for inclusion in the Work by the copyright owner
+   or by an individual or Legal Entity authorized to submit on behalf of
+   the copyright owner. For the purposes of this definition, "submitted"
+   means any form of electronic, verbal, or written communication sent
+   to the Licensor or its representatives, including but not limited to
+   communication on electronic mailing lists, source code control systems,
+   and issue tracking systems that are managed by, or on behalf of, the
+   Licensor for the purpose of discussing and improving the Work, but
+   excluding communication that is conspicuously marked or otherwise
+   designated in writing by the copyright owner as "Not a Contribution."
+
+   "Contributor" shall mean Licensor and any individual or Legal Entity
+   on behalf of whom a Contribution has been received by Licensor and
+   subsequently incorporated within the Work.
+
+2. Grant of Copyright License.
+
+   Subject to the terms and conditions of this License, each Contributor
+   hereby grants to You a perpetual, worldwide, non-exclusive, no-charge,
+   royalty-free, irrevocable copyright license to reproduce, prepare
+   Derivative Works of, publicly display, publicly perform, sublicense,
+   and distribute the Work and such Derivative Works in Source or Object
+   form.
+
+3. Grant of Patent License.
+
+   Subject to the terms and conditions of this License, each Contributor
+   hereby grants to You a perpetual, worldwide, non-exclusive, no-charge,
+   royalty-free, irrevocable (except as stated in this section) patent
+   license to make, have made, use, offer to sell, sell, import, and
+   otherwise transfer the Work, where such license applies only to those
+   patent claims licensable by such Contributor that are necessarily
+   infringed by their Contribution(s) alone or by combination of their
+   Contribution(s) with the Work to which such Contribution(s) was
+   submitted. If You institute patent litigation against any entity
+   (including a cross-claim or counterclaim in a lawsuit) alleging that
+   the Work or a Contribution incorporated within the Work constitutes
+   direct or contributory patent infringement, then any patent licenses
+   granted to You under this License for that Work shall terminate as of
+   the date such litigation is filed.
+
+4. Redistribution.
+
+   You may reproduce and distribute copies of the Work or Derivative
+   Works thereof in any medium, with or without modifications, and in
+   Source or Object form, provided that You meet the following conditions:
+
+   (a) You must give any other recipients of the Work or Derivative Works
+       a copy of this License; and
+
+   (b) You must cause any modified files to carry prominent notices stating
+       that You changed the files; and
+
+   (c) You must retain, in the Source form of any Derivative Works that You
+       distribute, all copyright, patent, trademark, and attribution
+       notices from the Source form of the Work, excluding those notices
+       that do not pertain to any part of the Derivative Works; and
+
+   (d) If the Work includes a "NOTICE" text file as part of its distribution,
+       then any Derivative Works that You distribute must include a readable
+       copy of the attribution notices contained within such NOTICE file,
+       excluding those notices that do not pertain to any part of the
+       Derivative Works, in at least one of the following places: within a
+       NOTICE text file distributed as part of the Derivative Works; within
+       the Source form or documentation, if provided along with the
+       Derivative Works; or, within a display generated by the Derivative
+       Works, if and wherever such third-party notices normally appear. The
+       contents of the NOTICE file are for informational purposes only and
+       do not modify the License. You may add Your own attribution notices
+       within Derivative Works that You distribute, alongside or as an
+       addendum to the NOTICE text from the Work, provided that such
+       additional attribution notices cannot be construed as modifying the
+       License.
+
+   You may add Your own copyright statement to Your modifications and may
+   provide additional or different license terms and conditions for use,
+   reproduction, or distribution of Your modifications, or for any such
+   Derivative Works as a whole, provided Your use, reproduction, and
+   distribution of the Work otherwise complies with the conditions
+   stated in this License.
+
+5. Submission of Contributions.
+
+   Unless You explicitly state otherwise, any Contribution intentionally
+   submitted for inclusion in the Work by You to the Licensor shall be
+   under the terms and conditions of this License, without any additional
+   terms or conditions. Notwithstanding the above, nothing herein shall
+   supersede or modify the terms of any separate license agreement you
+   may have executed with Licensor regarding such Contributions.
+
+6. Trademarks.
+
+   This License does not grant permission to use the trade names,
+   trademarks, service marks, or product names of the Licensor, except as
+   required for describing the origin of the Work and reproducing the
+   content of the NOTICE file.
+
+7. Disclaimer of Warranty.
+
+   Unless required by applicable law or agreed to in writing, Licensor
+   provides the Work (and each Contributor provides its Contributions) on
+   an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+   express or implied, including, without limitation, any warranties or
+   conditions of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR
+   A PARTICULAR PURPOSE. You are solely responsible for determining the
+   appropriateness of using or redistributing the Work and assume any
+   risks associated with Your exercise of permissions under this License.
+
+8. Limitation of Liability.
+
+   In no event and under no legal theory, whether in tort (including
+   negligence), contract, or otherwise, unless required by applicable law
+   (such as deliberate and grossly negligent acts) or agreed to in writing,
+   shall any Contributor be liable to You for damages, including any direct,
+   indirect, special, incidental, or consequential damages of any character
+   arising as a result of this License or out of the use or inability to
+   use the Work (including but not limited to damages for loss of goodwill,
+   work stoppage, computer failure or malfunction, or any and all other
+   commercial damages or losses), even if such Contributor has been advised
+   of the possibility of such damages.
+
+9. Accepting Warranty or Additional Liability.
+
+   While redistributing the Work or Derivative Works thereof, You may choose
+   to offer, and charge a fee for, acceptance of support, warranty, indemnity,
+   or other liability obligations and/or rights consistent with this License.
+   However, in accepting such obligations, You may act only on Your own behalf
+   and on Your sole responsibility, not on behalf of any other Contributor,
+   and only if You agree to indemnify, defend, and hold each Contributor
+   harmless for any liability incurred by, or claims asserted against, such
+   Contributor by reason of your accepting any such warranty or additional
+   liability.
+
+END OF TERMS AND CONDITIONS
+
+APPENDIX: How to apply the Apache License to your work.
+
+   To apply the Apache License to your work, attach the following boilerplate
+   notice, with the fields enclosed by brackets "[]" replaced with your own
+   identifying information. (Don't include the brackets!)  The text should be
+   enclosed in the appropriate comment syntax for the file format. We also
+   recommend that a file or class name and description of purpose be included
+   on the same "printed page" as the copyright notice for easier
+   identification within third-party archives.
+
+Copyright [2025] [s-ohtake]
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.

lightningbo-0.1.0/LightningBO/__init__.py

@@ -0,0 +1,2 @@
+# LightningBO
+__version__ = "0.1.0"

lightningbo-0.1.0/LightningBO/bayesian_core.py

@@ -0,0 +1,460 @@
+"""
+Bayesian Optimization Core Library
+
+
+This module provides core components for Bayesian optimization.
+
+- Parameter space definitions (Continuous, Integer, ParamSpace)
+- Gaussian Process utilities (kernels, prediction, hyperparameter fitting)
+- Acquisition functions (EI, PI, UCB, LCB)
+
+Design philosophy:
+- Flexible abstraction of parameter space for use from Optimizer
+- Fast numerical computation and automatic differentiation via JAX
+- Reusable for both research and production
+
+Main use cases:
+- Black-box function optimization
+- Hyperparameter search
+- Experiment planning and automation
+
+All components are designed to be reusable and extensible for other optimization algorithms.
+"""
+
+from collections import namedtuple
+from functools import partial
+from typing import Any
+
+import jax
+import jax.numpy as jnp
+import optax
+from jax.scipy.linalg import cholesky, solve_triangular
+from jax.scipy.stats import norm
+
+jax.config.update("jax_enable_x64", True)
+from abc import ABC, abstractmethod
+
+class ParameterDomain(ABC):
+    """パラメータ空間の抽象基底クラス"""
+
+    def __init__(self, dtype):
+        self.dtype = dtype
+
+    def __hash__(self):
+        return hash((type(self), self.dtype))
+
+    def __eq__(self, other):
+        return type(self) == type(other) and self.dtype == other.dtype
+
+    @abstractmethod
+    def transform(self, x):
+        pass
+
+    @abstractmethod
+    def sample(self, key, shape):
+        pass
+
+
+
+class Continuous(ParameterDomain):
+    """連続値パラメータ空間"""
+
+    def __init__(self, lower: float, upper: float):
+        if not isinstance(lower, (float, int)):
+            raise TypeError(f"lower must be float or int, got {type(lower)}")
+        if not isinstance(upper, (float, int)):
+            raise TypeError(f"upper must be float or int, got {type(upper)}")
+        if lower > upper:
+            raise ValueError("lower must be <= upper")
+        self.lower = float(lower)
+        self.upper = float(upper)
+        super().__init__(dtype=jnp.float32)
+
+    def __hash__(self):
+        return hash((type(self), self.lower, self.upper))
+
+    def __eq__(self, other):
+        return (
+            isinstance(other, Continuous)
+            and self.lower == other.lower
+            and self.upper == other.upper
+        )
+
+    def transform(self, x):
+        return jnp.clip(x, self.lower, self.upper)
+
+    def sample(self, key, shape):
+        samples = jax.random.uniform(key, shape, minval=self.lower, maxval=self.upper)
+        return self.transform(samples)
+
+
+class Integer(ParameterDomain):
+    """
+    Discrete integer-valued domain with rounding and clipping.
+
+    Represents a parameter that can take integer values within [lower, upper].
+    """
+
+    def __init__(self, lower, upper):
+        """
+        Initialize an integer domain with bounds.
+
+        Args:
+            lower: Lower integer bound (inclusive).
+            upper: Upper integer bound (inclusive).
+        """
+        if not isinstance(lower, int):
+            raise TypeError("Lower bound must be an integer")
+        if not isinstance(upper, int):
+            raise TypeError("Upper bound must be an integer")
+        if lower >= upper:
+            raise ValueError("Lower bound must be less than upper bound")
+
+        self.lower = int(lower)
+        self.upper = int(upper)
+        super().__init__(dtype=jnp.int32)
+
+    def __hash__(self):
+        """Return hash based on bounds."""
+        return hash((self.lower, self.upper))
+
+    def __eq__(self, other):
+        """Check equality based on bounds."""
+        return self.lower == other.lower and self.upper == other.upper
+
+    def transform(self, x: jax.Array):
+        """
+        Round and clip values to the integer domain.
+
+        Args:
+            x: Input values.
+
+        Returns
+        -------
+            Rounded and clipped values as float32.
+        """
+        return jnp.clip(jnp.round(x), self.lower, self.upper).astype(jnp.float32)
+
+    def sample(self, key: jax.Array, shape: tuple):
+        """
+        Sample integers uniformly from the domain.
+
+        Args:
+            key: JAX PRNGKey.
+            shape: Desired output shape.
+
+        Returns
+        -------
+            Sampled values clipped to valid integer range.
+        """
+        samples = jax.random.randint(
+            key, shape, minval=self.lower, maxval=self.upper + 1
+        )
+        return self.transform(samples)
+
+
+
+class ParamSpace:
+    """
+    Parameter space manager for collections of named parameter domains.
+
+    This utility encapsulates logic for sampling, transforming, and handling
+    structured parameter inputs defined by a mapping of variable names to ParameterDomain
+    instances (e.g., Continuous, Integer).
+    """
+
+    def __init__(self, space: dict):
+        self.space = space
+
+    def sample_params(self, key: jax.Array, shape: tuple) -> dict:
+        """Sample parameters from all domains."""
+        keys = jax.random.split(key, len(self.space))
+        return {
+            name: self.space[name].sample(k, shape)
+            for name, k in zip(self.space, keys)
+        }
+
+    def to_array(self, tree: dict) -> jax.Array:
+        """
+        Transform a batch of parameter values into a 2D array suitable for GP input.
+
+        Args:
+            tree: A dictionary of parameter name → array of raw values.
+
+        Returns
+        -------
+            A JAX array of shape (batch_size, num_params) with transformed values.
+        """
+        return jnp.stack([self.space[k].transform(tree[k]) for k in self.space], axis=1)
+
+    def to_dict(self, xs: jax.Array) -> dict:
+        """
+        Convert a stacked parameter matrix back into named parameter trees.
+
+        Args:
+            xs: A 2D JAX array of shape (batch_size, num_params), with each column
+                corresponding to a parameter.
+
+        Returns
+        -------
+            A dictionary mapping parameter names to individual 1D arrays.
+        """
+        return {k: self.space[k].transform(xs[:, i]) for i, k in enumerate(self.space)}
+
+
+# =============================================================================
+# Gaussian Process Components
+# =============================================================================
+
+MASK_VARIANCE = 1e12  # High variance for masked points to not affect the process.
+
+GPParams = namedtuple("GPParams", ["noise", "amplitude", "lengthscale"])
+GPState = namedtuple("GPState", ["params", "momentums", "scales"])
+
+
+def softplus(x):
+    """Softplus activation function."""
+    return jnp.logaddexp(x, 0.0)
+
+
+def exp_quadratic(x1, x2, mask):
+    """Exponential quadratic kernel function."""
+    distance = jnp.sum((x1 - x2) ** 2)
+    return jnp.exp(-distance) * mask
+
+
+def cov(x1, x2, mask1, mask2):
+    """Covariance function with masking."""
+    M = jnp.outer(mask1, mask2)
+    k = exp_quadratic
+    return jax.vmap(jax.vmap(k, in_axes=(None, 0, 0)), in_axes=(0, None, 0))(x1, x2, M)
+
+
+def gaussian_process(
+    params,
+    x: jnp.ndarray,
+    y: jnp.ndarray,
+    mask,
+    xt: Any = None,
+    compute_ml: bool = False,
+) -> Any:
+    """Core Gaussian Process implementation with masking support."""
+    # Number of points in the prior distribution
+    n = x.shape[0]
+
+    noise, amp, ls = jax.tree_util.tree_map(softplus, params)
+
+    ymean = jnp.mean(y, where=mask.astype(bool))
+    y = (y - ymean) * mask
+    x = x / ls
+    K = amp * cov(x, x, mask, mask) + (jnp.eye(n) * (noise + 1e-6))
+    K += jnp.eye(n) * (1.0 - mask.astype(float)) * MASK_VARIANCE
+    L = cholesky(K, lower=True)
+    K_inv_y = solve_triangular(L.T, solve_triangular(L, y, lower=True), lower=False)
+
+    if compute_ml:
+        logp = 0.5 * jnp.dot(y.T, K_inv_y)
+        logp += jnp.sum(jnp.log(jnp.diag(L)))
+        logp -= jnp.sum(1.0 - mask) * 0.5 * jnp.log(MASK_VARIANCE)
+        logp += (jnp.sum(mask) / 2) * jnp.log(2 * jnp.pi)
+        logp += jnp.sum(-0.5 * jnp.log(2 * jnp.pi) - jnp.log(amp) - jnp.log(amp) ** 2)
+        return jnp.sum(logp)
+
+    if xt is None:
+        raise ValueError("xt cannot be None during prediction")
+    xt = xt / ls
+
+    # Compute the covariance with the new point xt
+    mask_t = jnp.ones(len(xt)) == 1
+    K_cross = amp * cov(x, xt, mask, mask_t)
+
+    K_inv_y = K_inv_y * mask  # masking
+    pred_mean = jnp.dot(K_cross.T, K_inv_y) + ymean
+    v = solve_triangular(L, K_cross, lower=True)
+    pred_var = amp * cov(xt, xt, mask_t, mask_t) - v.T @ v
+    pred_std = jnp.sqrt(jnp.maximum(jnp.diag(pred_var), 1e-10))
+    return pred_mean, pred_std
+
+
+marginal_likelihood = partial(gaussian_process, compute_ml=True)
+grad_fun = jax.jit(jax.grad(marginal_likelihood))
+predict = jax.jit(partial(gaussian_process, compute_ml=False))
+
+
+def neg_log_likelihood(params, x, y, mask):
+    """Negative log likelihood for GP hyperparameter optimization."""
+    ll = marginal_likelihood(params, x, y, mask)
+
+    # Weak priors to keep things sane
+    priors = GPParams(-8.0, 1.0, 1.0)
+    log_prior = jax.tree.map(lambda p, m: jnp.sum((p - m) ** 2), params, priors)
+    log_prior = sum(jax.tree.leaves(log_prior))
+    log_posterior = ll - 0.5 * log_prior
+    return -log_posterior
+
+
+def posterior_fit(
+    y: jax.Array,
+    x: jax.Array,
+    mask: jax.Array,
+    params: GPParams,
+    lr: float = 1e-3,
+    trainsteps: int = 100,
+) -> GPParams:
+    """Fit GP hyperparameters using gradient descent."""
+    optimizer = optax.chain(optax.clip_by_global_norm(10.0), optax.adamw(lr))
+    opt_state = optimizer.init(params)
+
+    def train_step(carry, _):
+        params, opt_state = carry
+        grads = jax.grad(neg_log_likelihood)(params, x, y, mask)
+        updates, opt_state = optimizer.update(grads, opt_state, params=params)
+        params = optax.apply_updates(params, updates)
+        return (params, opt_state), None
+
+    (params, _), __ = jax.lax.scan(
+        train_step, (params, opt_state), None, length=trainsteps
+    )
+    return params
+
+
+# =============================================================================
+# Acquisition Functions
+# =============================================================================
+
+
+def expected_improvement(
+    x_pred: jnp.ndarray,
+    xs: jax.Array,
+    ys: jax.Array,
+    mask: jax.Array,
+    gp_params: GPParams,
+    xi: float = 0.01,
+):
+    r"""
+    Compute Expected Improvement (EI) acquisition function.
+
+    Favors points with high improvement over the current best observed value,
+    balancing exploitation and exploration.
+
+    The formula is:
+
+    .. math::
+
+        EI(x) = (\mu(x) - y^* - \xi) \Phi(z) + \sigma(x) \phi(z)
+
+    where:
+
+    .. math::
+
+        z = \frac{\mu(x) - y^* - \xi}{\sigma(x)}
+
+    Args:
+        x_pred: Candidate input locations to evaluate.
+        xs: Observed inputs.
+        ys: Observed function values.
+        mask: Boolean mask indicating valid entries in `ys`.
+        gp_params: Gaussian Process hyperparameters.
+        xi: Exploration-exploitation tradeoff parameter.
+
+    Returns
+    -------
+        EI scores at `x_pred`.
+    """
+    ymax = jnp.max(ys, where=mask.astype(bool), initial=-jnp.inf)
+    mu, std = predict(gp_params, xs, ys, mask, xt=x_pred)
+    a = mu - ymax - xi
+    z = a / (std + 1e-3)
+    ei = a * norm.cdf(z) + std * norm.pdf(z)
+    return ei
+
+
+def probability_improvement(
+    x_pred: jnp.ndarray,
+    xs: jax.Array,
+    ys: jax.Array,
+    mask: jax.Array,
+    gp_params: GPParams,
+    xi: float = 0.01,
+):
+    """
+    Probability of Improvement (PI) acquisition function.
+
+    Estimates the probability that a candidate point will improve
+    over the current best observed value.
+
+    Args:
+        x_pred: Candidate input locations to evaluate.
+        xs: Observed inputs.
+        ys: Observed function values.
+        mask: Boolean mask indicating valid entries in `ys`.
+        gp_params: Gaussian Process hyperparameters.
+        xi: Improvement margin for sensitivity.
+
+    Returns
+    -------
+        PI scores at `x_pred`.
+    """
+    y_max = ys.max()
+    mu, std = predict(gp_params, xs, ys, mask, xt=x_pred)
+    z = (mu - y_max - xi) / std
+    return norm.cdf(z)
+
+
+def upper_confidence_bounds(
+    x_pred: jnp.ndarray,
+    xs: jax.Array,
+    ys: jax.Array,
+    mask: jax.Array,
+    gp_params: GPParams,
+    kappa: float = 0.01,
+):
+    """
+    Upper Confidence Bound (UCB) acquisition function.
+
+    Promotes exploration by favoring points with high predictive uncertainty.
+
+    Args:
+        x_pred: Candidate input locations to evaluate.
+        xs: Observed inputs.
+        ys: Observed function values.
+        mask: Boolean mask indicating valid entries in `ys`.
+        gp_params: Gaussian Process hyperparameters.
+        kappa: Weighting factor for uncertainty.
+
+    Returns
+    -------
+        UCB scores at `x_pred`.
+    """
+    mu, std = predict(gp_params, xs, ys, mask, xt=x_pred)
+    return mu + kappa * std
+
+
+def lower_confidence_bounds(
+    x_pred: jnp.ndarray,
+    xs: jax.Array,
+    ys: jax.Array,
+    mask: jax.Array,
+    gp_params: GPParams,
+    kappa: float = 2.576,
+):
+    """
+    Lower Confidence Bound (LCB) acquisition function.
+
+    Useful for minimization tasks. Encourages sampling in uncertain regions
+    with low predicted values.
+
+    Args:
+        x_pred: Candidate input locations to evaluate.
+        xs: Observed inputs.
+        ys: Observed function values.
+        mask: Boolean mask indicating valid entries in `ys`.
+        gp_params: Gaussian Process hyperparameters.
+        kappa: Weighting factor for uncertainty.
+
+    Returns
+    -------
+        LCB scores at `x_pred`.
+    """
+    mu, std = predict(gp_params, xs, ys, mask, xt=x_pred)
+    return mu - kappa * std

lightningbo-0.1.0/LightningBO/bayesian_optimization.py

@@ -0,0 +1,336 @@
+
+"""
+Bayesian Optimization High-Level Interface
+
+This module provides a high-level API for Bayesian optimization using Gaussian Processes.
+It manages the optimization loop, parameter space, and acquisition functions.
+
+Main features:
+- Optimizer class for black-box function optimization
+- Flexible parameter space definition (Continuous, Integer, etc.)
+- Acquisition functions: EI, PI, UCB, LCB
+- Automatic state management and buffer expansion
+- JAX-based fast computation and differentiation
+
+Typical use cases:
+- Hyperparameter tuning
+- Experimental design
+- Automated search for expensive functions
+
+Designed for extensibility and integration with custom objective functions.
+"""
+
+
+from collections.abc import Callable
+from functools import partial
+from typing import NamedTuple
+
+
+import jax
+from typing import Any
+import jax.numpy as jnp
+import numpy as np
+import optax
+
+from .bayesian_core import (
+    GPParams,
+    ParamSpace,
+    expected_improvement,
+    lower_confidence_bounds,
+    posterior_fit,
+    probability_improvement,
+    upper_confidence_bounds,
+)
+
+
+UtilityFunctionType = Callable[[jnp.ndarray], float]
+
+
+
+class OptimizerState(NamedTuple):
+    """
+    State container for Bayesian optimizer.
+
+    Attributes:
+        params: dict[str, jnp.ndarray] - Padded parameter arrays.
+        ys: jnp.ndarray | np.ndarray - Objective values (with padding).
+        best_score: float - Best observed objective value.
+        best_params: dict[str, float] - Parameter configuration for best_score.
+        mask: jnp.ndarray - Boolean mask for valid entries.
+        gp_params: GPParams - Fitted Gaussian Process parameters.
+    """
+    params: dict[str, jnp.ndarray]
+    ys: jnp.ndarray | np.ndarray
+    best_score: float
+    best_params: dict[str, float]
+    mask: jnp.ndarray
+    gp_params: GPParams
+
+
+
+def _optimize_suggestion(params: jnp.ndarray, fun: Callable[[jnp.ndarray], float], max_iter: int = 10) -> jnp.ndarray:
+    """
+    Local optimization using L-BFGS (Optax).
+
+    Args:
+        params: Initial point in input space.
+        fun: Objective function to maximize (returns scalar, supports autodiff).
+        max_iter: Number of L-BFGS steps.
+
+    Returns:
+        Optimized point after `max_iter` iterations.
+    """
+    opt = optax.lbfgs()
+    value_and_grad_fun = optax.value_and_grad_from_state(lambda x: -fun(x))
+
+    def step(carry, _):
+        params, state = carry
+        value, grad = value_and_grad_fun(params, state=state)
+        updates, state = opt.update(
+            grad, state, params=params, value=value, grad=grad, value_fn=fun
+        )
+        params = optax.apply_updates(params, updates)
+        params = jnp.clip(params, -1e6, 1e6)
+        return (params, state), None
+
+    init_carry = (params, opt.init(params))
+    (final_params, _), __ = jax.lax.scan(step, init_carry, None, length=max_iter)
+    return final_params
+
+
+
+class Optimizer:
+    """
+    Bayesian optimizer using Gaussian Processes and acquisition functions.
+
+    Manages the optimization loop for expensive black-box functions by modeling
+    them with a Gaussian Process and selecting samples via acquisition functions
+    (EI, PI, UCB, LCB).
+    """
+
+    def __init__(self, domain: dict[str, Any], acq: str = "EI", maximize: bool = False):
+        """
+        Initialize the optimizer.
+
+        Args:
+            domain: A dict mapping parameter names to domain objects (e.g., Real, Integer).
+            acq: Acquisition function ('EI', 'PI', 'UCB', 'LCB').
+            maximize: Whether to maximize or minimize the objective.
+        """
+        self.domain = domain
+        self.sign = 1 if maximize else -1
+        self.param_space = ParamSpace(domain)
+
+        acq_map = {
+            "EI": expected_improvement,
+            "PI": probability_improvement,
+            "UCB": upper_confidence_bounds,
+            "LCB": lower_confidence_bounds,
+        }
+        if acq not in acq_map:
+            raise ValueError(f"Acquisition function {acq} is not implemented")
+        self.acq = jax.jit(acq_map[acq])
+
+    def init(self, ys: jax.Array, params: dict, noise_scale: float = -8.0):
+        """
+        Initialize the optimizer state from initial data.
+
+        Args:
+            ys: Objective values for the initial parameters.
+            params: Dict of parameter arrays (same keys as domain).
+            noise_scale: Initial noise scale for GP.
+
+        Returns
+        -------
+            Initialized OptimizerState.
+        """
+        # Create a padded jax array for each parameter and each score.
+        # In order to keep jax compilations at a bay.
+        num_entries = len(ys)
+        pad_value = int(np.ceil(len(ys) / 10) * 10)
+
+        # Convert to jax arrays if they are not already
+        ys = jnp.asarray(ys)
+        ys = self.sign * ys
+        params = jax.tree.map(lambda x: jnp.asarray(x), params)
+
+        # Define padded arrays for the inputs and the outputs
+        mask = jnp.zeros(shape=(pad_value,), dtype=jnp.bool_).at[:num_entries].set(True)
+        ys = jnp.zeros(shape=(pad_value,), dtype=ys.dtype).at[:num_entries].set(ys)
+
+        _params = {}
+        for key, entries in params.items():
+            # Assert that the parameter is in the domain dictionary
+            if key not in self.domain:
+                raise ValueError(f"Parameter {key} is not in the domain")
+
+            # Get dtype from the domain and create a padded array
+            dtype = self.domain[key].dtype
+            values = (
+                jnp.zeros(shape=(pad_value,), dtype=dtype).at[:num_entries].set(entries)
+            )
+            _params[key] = values
+
+        # From the given observation, find the better one (either maxima or minima) and return the
+        # initial optimizer state.
+        best_score = float(jnp.max(ys[mask]))
+        best_params_idx = jnp.argmax(ys[mask])
+        best_params = jax.tree.map(lambda x: x[mask][best_params_idx], _params)
+
+        # Initialize the gaussian processes state
+        gp_params = GPParams(
+            noise=jnp.full((1, 1), 1.0 * noise_scale),
+            amplitude=jnp.zeros((1, 1)),
+            lengthscale=jnp.zeros((1, len(_params))),
+        )
+
+        # Fit to the current observations
+        xs = jnp.stack(
+            [self.domain[key].transform(_params[key]) for key in _params], axis=1
+        )
+        gp_params = posterior_fit(ys, xs, mask=mask, params=gp_params)
+
+        ys = self.sign * ys
+        best_score = self.sign * best_score
+        opt_state = OptimizerState(
+            params=_params,
+            ys=ys,
+            best_score=best_score,
+            best_params=best_params,
+            mask=mask,
+            gp_params=gp_params,
+        )
+
+        return opt_state
+
+    @partial(jax.jit, static_argnames=("self", "size"))
+    def sample(self, key, state, size=10_000):
+        """
+        Sample new parameters using the acquisition function.
+
+        Args:
+            key: JAX PseudoRandom key for random sampling.
+            state: Current optimizer state.
+            size: Number of samples to draw.
+
+        Returns
+        -------
+            Sampled parameters (dict).
+        """
+        # Sample 'size' elements of each distribution.
+        samples = self.param_space.sample_params(key, (size,))
+        xs_samples = self.param_space.to_array(samples)
+
+        # Prepare the data for the Gaussian process prediction.
+        xs = self.param_space.to_array(state.params)
+        ys = self.sign * state.ys
+        mask = state.mask
+        gp_params = state.gp_params
+
+        # Compute the acquisition function values for the sampled points.
+        acq_vals = self.acq(xs_samples, xs, ys, mask, gp_params)
+
+        # Of those, find the best 50 points and optimize them using BFGS.
+        top_idxs = jnp.argsort(acq_vals)[-50:]
+        init_points = xs_samples[top_idxs]
+
+        def acquisition_objective(x):
+            return jnp.squeeze(self.acq(x[None, :], xs, ys, mask, gp_params))
+
+        optimized = jax.vmap(
+            lambda x: _optimize_suggestion(x, acquisition_objective, max_iter=10)
+        )(init_points)
+        opt_vals = self.acq(optimized, xs, ys, mask, gp_params)
+
+        # Return the best point from the optimized points and the sampled points.
+        all_points = jnp.concatenate((optimized, xs_samples), axis=0)
+        all_vals = jnp.concatenate((opt_vals, acq_vals), axis=0)
+        chosen_suggestion = jnp.argmax(all_vals)
+
+        best_params = self.param_space.to_dict(all_points[chosen_suggestion][None])
+        return best_params
+
+    def expand(self, opt_state: OptimizerState):
+        """
+        Expand internal buffers if no space is available.
+
+        Args:
+            opt_state: Current optimizer state.
+
+        Returns
+        -------
+            OptimizerState with expanded storage.
+        """
+        current = jnp.sum(opt_state.mask)
+
+        if current == len(opt_state.mask):
+            pad_value = int(np.ceil(len(opt_state.mask) * 2 / 10) * 10)
+            diff = pad_value - len(opt_state.mask)
+            mask = jnp.pad(opt_state.mask, (0, diff))
+            ys = jnp.pad(opt_state.ys, (0, diff))
+            params = {}
+            for key in opt_state.params:
+                params[key] = jnp.pad(opt_state.params[key], (0, diff))
+        else:
+            mask = opt_state.mask
+            ys = opt_state.ys
+            params = opt_state.params
+
+        opt_state = OptimizerState(
+            params=params,
+            ys=ys,
+            best_score=opt_state.best_score,
+            best_params=opt_state.best_params,
+            mask=mask,
+            gp_params=opt_state.gp_params,
+        )
+        return opt_state
+
+    def fit(self, opt_state, y, new_params):
+        """
+        Update optimizer state with a new observation.
+
+        Args:
+            opt_state: Current optimizer state.
+            y: New objective value.
+            new_params: Parameters that produced y.
+
+        Returns
+        -------
+            Updated OptimizerState.
+        """
+        opt_state = self.expand(opt_state)  # Prompts recompilation
+        opt_state = self._fit(opt_state, y, new_params)
+        return opt_state
+
+    @partial(jax.jit, static_argnums=(0,))
+    def _fit(self, opt_state, y, new_params):
+        """Update optimizer state with new observation (JIT-compiled)."""
+        last_idx = jnp.arange(len(opt_state.mask)) == jnp.argmin(opt_state.mask)
+        mask = jnp.asarray(jnp.where(last_idx, True, opt_state.mask))
+        ys = jnp.where(last_idx, y, opt_state.ys)
+        params = jax.tree_util.tree_map(
+            lambda x, y: jnp.where(last_idx, y, x), opt_state.params, new_params
+        )
+
+        xs = jnp.stack(
+            [self.domain[key].transform(params[key]) for key in params], axis=1
+        )
+        ys = self.sign * ys
+        gp_params = posterior_fit(ys, xs, mask=mask, params=opt_state.gp_params)
+
+        best_score = jnp.max(ys, where=mask, initial=-jnp.inf)
+        best_params_idx = jnp.argmax(jnp.where(mask, ys, -jnp.inf))
+        best_params = jax.tree_util.tree_map(lambda x: x[best_params_idx], params)
+
+        ys = self.sign * ys
+        best_score = self.sign * best_score
+        opt_state = OptimizerState(
+            params=params,
+            ys=ys,
+            best_score=best_score,
+            best_params=best_params,
+            mask=mask,
+            gp_params=gp_params,
+        )
+        return opt_state

lightningbo-0.1.0/LightningBO.egg-info/PKG-INFO

@@ -0,0 +1,15 @@
+Metadata-Version: 2.4
+Name: LightningBO
+Version: 0.1.0
+Summary: Bayesian optimization library
+License-Expression: Apache-2.0
+Requires-Python: >=3.12
+License-File: LICENSE
+Requires-Dist: jax
+Requires-Dist: jaxlib
+Requires-Dist: optax
+Requires-Dist: numpy
+Requires-Dist: absl-py
+Requires-Dist: pytest
+Requires-Dist: matplotlib
+Dynamic: license-file

lightningbo-0.1.0/LightningBO.egg-info/SOURCES.txt

@@ -0,0 +1,13 @@
+LICENSE
+README.md
+pyproject.toml
+LightningBO/__init__.py
+LightningBO/bayesian_core.py
+LightningBO/bayesian_optimization.py
+LightningBO.egg-info/PKG-INFO
+LightningBO.egg-info/SOURCES.txt
+LightningBO.egg-info/dependency_links.txt
+LightningBO.egg-info/requires.txt
+LightningBO.egg-info/top_level.txt
+tests/test_bayesian_core.py
+tests/test_bayesian_optimization.py

lightningbo-0.1.0/LightningBO.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+

lightningbo-0.1.0/LightningBO.egg-info/requires.txt

@@ -0,0 +1,7 @@
+jax
+jaxlib
+optax
+numpy
+absl-py
+pytest
+matplotlib

lightningbo-0.1.0/LightningBO.egg-info/top_level.txt

@@ -0,0 +1 @@
+LightningBO

lightningbo-0.1.0/PKG-INFO

@@ -0,0 +1,15 @@
+Metadata-Version: 2.4
+Name: LightningBO
+Version: 0.1.0
+Summary: Bayesian optimization library
+License-Expression: Apache-2.0
+Requires-Python: >=3.12
+License-File: LICENSE
+Requires-Dist: jax
+Requires-Dist: jaxlib
+Requires-Dist: optax
+Requires-Dist: numpy
+Requires-Dist: absl-py
+Requires-Dist: pytest
+Requires-Dist: matplotlib
+Dynamic: license-file

lightningbo-0.1.0/README.md

@@ -0,0 +1,48 @@
+# LightningBO
+
+LightningBOは、ベイズ最適化（Bayesian Optimization）を簡単に実装・拡張できるPythonライブラリです。
+本ライブラリは、ブラックボックス関数の最適化やハイパーパラメータ探索など、試行回数を抑えつつ効率的に最適解を見つけたい場面で活用できます。
+
+## 特徴
+
+- JAXベースの高速な計算
+- 連続・離散パラメータ両対応
+- 期待値改善（EI）などの獲得関数
+- 拡張性の高い設計
+
+## 使い方例
+
+```python
+import LightningBO.bayesian_optimization as bo
+import LightningBO.bayesian_core as bc
+
+# 最適化したい関数
+def f(x):
+	return x**2 - 4*x + 5
+
+# 探索空間の定義
+domain = {"x": bc.Continuous(-10, 10)}
+
+# Optimizerの作成
+optimizer = bo.Optimizer(domain=domain, maximize=False)
+
+# 初期点
+params = {"x": [0.0, 2.0, 5.0]}
+ys = [f(x) for x in params["x"]]
+state = optimizer.init(ys, params)
+
+# 最適化ループ
+for step in range(20):
+	new_params = optimizer.sample(None, state)
+	y = f(**new_params)
+	state = optimizer.fit(state, y, new_params)
+
+print("最適値:", state.best_score)
+print("最適パラメータ:", state.best_params)
+```
+
+## 主な用途
+
+- 機械学習モデルのハイパーパラメータ最適化
+- 実験・シミュレーションのパラメータ探索
+- その他ブラックボックス関数の最適化

lightningbo-0.1.0/pyproject.toml

@@ -0,0 +1,16 @@
+[project]
+name = "LightningBO"
+version = "0.1.0"
+description = "Bayesian optimization library"
+requires-python = ">=3.12"
+
+license = "Apache-2.0"
+dependencies = [
+	"jax",
+	"jaxlib",
+	"optax",
+	"numpy",
+	"absl-py",
+	"pytest",
+	"matplotlib"
+]

lightningbo-0.1.0/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

lightningbo-0.1.0/tests/test_bayesian_core.py

@@ -0,0 +1,62 @@
+"""Tests for the Bayesian core functionality including GP operations and acquisition functions."""
+
+import jax
+import jax.numpy as jnp
+from absl.testing import absltest, parameterized
+
+from LightningBO.bayesian_core import (
+	GPParams,
+	Integer,
+	Continuous,
+	cov,
+	exp_quadratic,
+	expected_improvement,
+	gaussian_process,
+	lower_confidence_bounds,
+	marginal_likelihood,
+	predict,
+	probability_improvement,
+	upper_confidence_bounds,
+)
+
+class BayesianCoreDomainTest(parameterized.TestCase):
+	"""Test suite for Bayesian domain classes (Continuous, Integer)."""
+
+	@parameterized.parameters(
+		(0.0, 1.0),
+		(-5.5, 5.5),
+	)
+	def test_continuous_transform_clips_within_bounds(self, lower, upper):
+		"""Test that Continuous domain clips values within bounds."""
+		domain = Continuous(lower, upper)
+		x = jnp.array([lower - 1, (lower + upper) / 2, upper + 1])
+		out = domain.transform(x)
+		self.assertTrue(jnp.all(out >= lower) and jnp.all(out <= upper))
+
+	def test_continuous_sample_within_bounds(self):
+		"""Test that Continuous domain samples are within bounds."""
+		domain = Continuous(0.0, 1.0)
+		key = jax.random.PRNGKey(0)
+		shape = (100,)
+		samples = domain.sample(key, shape)
+		self.assertEqual(samples.shape, shape)
+		self.assertTrue(jnp.all(samples >= 0.0) and jnp.all(samples <= 1.0))
+
+	@parameterized.parameters(
+		(0, 5),
+	)
+	def test_integer_transform_clips_within_bounds(self, lower, upper):
+		domain = Integer(lower, upper)
+		x = jnp.array([lower - 1, lower, upper, upper + 1])
+		out = domain.transform(x)
+		self.assertTrue(jnp.all(out >= lower) and jnp.all(out <= upper))
+
+	def test_integer_sample_within_bounds(self):
+		domain = Integer(0, 5)
+		key = jax.random.PRNGKey(0)
+		shape = (100,)
+		samples = domain.sample(key, shape)
+		self.assertEqual(samples.shape, shape)
+		self.assertTrue(jnp.all(samples >= 0) and jnp.all(samples <= 5))
+
+# ...他のテスト...

lightningbo-0.1.0/tests/test_bayesian_optimization.py

@@ -0,0 +1,48 @@
+"""Tests for Bayesian optimization functionality."""
+
+import jax
+import jax.numpy as jnp
+import numpy as np
+from absl.testing import absltest
+
+import LightningBO.bayesian_core as bc
+import LightningBO.bayesian_optimization as bo
+
+KEY = jax.random.PRNGKey(42)
+SEED = 42
+
+class BayesianOptimizationTest(absltest.TestCase):
+	"""Test suite for Bayesian optimization functionality."""
+
+	def test_1D_optim(self):
+		"""Test 1D optimization with sin function."""
+		def f(x):
+			return jnp.sin(x) - 0.1 * x**2
+
+		TARGET = np.max(f(np.linspace(-2, 2, 1000)))
+		domain = {"x": bc.Continuous(-2, 2)}
+		opt = bo.Optimizer(domain=domain, maximize=True)
+		params = {"x": [0.0, 1.0, 2.0]}
+		ys = [f(x) for x in params["x"]]
+		opt_state = opt.init(ys, params)
+		self.assertEqual(opt_state.best_score, np.max(ys))
+		self.assertEqual(opt_state.best_params["x"], params["x"][np.argmax(ys)])
+		self.assertEqual(type(opt_state), bo.OptimizerState)
+		self.assertEqual(type(opt_state.gp_params), bc.GPParams)
+		self.assertEqual(opt_state.params["x"].shape, (10,))
+		self.assertEqual(opt_state.ys.shape, (10,))
+		self.assertEqual(opt_state.mask.shape, (10,))
+		self.assertTrue(np.allclose(opt_state.params["x"][:3], params["x"]))
+		self.assertTrue(np.allclose(opt_state.ys[:3], ys))
+		self.assertTrue(np.allclose(opt_state.mask[:3], [True, True, True]))
+		self.assertTrue(np.allclose(opt_state.mask[3:], [False] * 7))
+		key = jax.random.PRNGKey(SEED)
+		sample_fn = jax.jit(opt.sample)
+		for step in range(100):
+			key = jax.random.fold_in(key, step)
+			new_params = sample_fn(key, opt_state)
+			y = f(**new_params)
+			opt_state = opt.fit(opt_state, y, new_params)
+			if jnp.allclose(opt_state.best_score, TARGET, atol=1e-03):
+				break
+# ...他のテスト...