bunobee 0.0.4__tar.gz

bunobee-0.0.4/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Edwin Ng
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

bunobee-0.0.4/PKG-INFO

@@ -0,0 +1,45 @@
+Metadata-Version: 2.4
+Name: bunobee
+Version: 0.0.4
+Summary: The time-series engine for forecasting with JAX and NumPyro
+Author-email: Edwin Ng <edwinnglabs@gmail.com>
+License-Expression: MIT
+Project-URL: Homepage, https://github.com/edwinnglabs/bunobee
+Project-URL: Repository, https://github.com/edwinnglabs/bunobee
+Project-URL: Issues, https://github.com/edwinnglabs/bunobee/issues
+Project-URL: Documentation, https://github.com/edwinnglabs/bunobee#readme
+Keywords: time-series,forecasting,bayesian,jax,numpyro,kalman-filter
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Science/Research
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Classifier: Topic :: Scientific/Engineering
+Requires-Python: >=3.11
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: numpyro>=0.20.0
+Requires-Dist: jax<0.10.0,>=0.8.0
+Requires-Dist: pandas>=2.0
+Requires-Dist: arviz
+Requires-Dist: xarray
+Provides-Extra: test
+Requires-Dist: pytest>=8; extra == "test"
+Requires-Dist: black==26.1.0; extra == "test"
+Provides-Extra: notebook
+Requires-Dist: ipywidgets; extra == "notebook"
+Dynamic: license-file
+
+<div align="center">
+  <img src="./bunobee.png" alt="Bunobee Logo" width="320" />
+</div>
+
+# Bunobee
+
+Bunobee is an experimental Bayesian time-series library built on [JAX](https://github.com/google/jax) and [NumPyro](https://github.com/pyro-ppl/numpyro). It currently supports the following model families:
+
+- **Damped Local Trend (DLT)** — an exponential smoothing model with a damped slope component
+- **State Space Models (SSP)** — structured state space models with Kalman filtering and smoothing
+
+**Fun fact:** "bunobee" is a word invented by my kids during their childhood to describe the action of the role they were playing.
+

bunobee-0.0.4/README.md

@@ -0,0 +1,13 @@
+<div align="center">
+  <img src="./bunobee.png" alt="Bunobee Logo" width="320" />
+</div>
+
+# Bunobee
+
+Bunobee is an experimental Bayesian time-series library built on [JAX](https://github.com/google/jax) and [NumPyro](https://github.com/pyro-ppl/numpyro). It currently supports the following model families:
+
+- **Damped Local Trend (DLT)** — an exponential smoothing model with a damped slope component
+- **State Space Models (SSP)** — structured state space models with Kalman filtering and smoothing
+
+**Fun fact:** "bunobee" is a word invented by my kids during their childhood to describe the action of the role they were playing.
+

bunobee-0.0.4/pyproject.toml

@@ -0,0 +1,57 @@
+[project]
+name = "bunobee"
+version = "0.0.4"
+description = "The time-series engine for forecasting with JAX and NumPyro"
+authors = [{name = "Edwin Ng", email = "edwinnglabs@gmail.com"}]
+readme = "README.md"
+requires-python = ">=3.11"
+license = "MIT"
+license-files = ["LICENSE"]
+keywords = ["time-series", "forecasting", "bayesian", "jax", "numpyro", "kalman-filter"]
+classifiers = [
+    "Development Status :: 3 - Alpha",
+    "Intended Audience :: Science/Research",
+    "Programming Language :: Python :: 3.11",
+    "Programming Language :: Python :: 3.12",
+    "Programming Language :: Python :: 3.13",
+    "Topic :: Scientific/Engineering",
+]
+dependencies = [
+    "numpyro>=0.20.0",
+    "jax>=0.8.0,<0.10.0",
+    "pandas>=2.0",
+    "arviz",
+    "xarray",
+]
+
+[project.optional-dependencies]
+test = [
+    "pytest>=8",
+    "black==26.1.0",
+]
+notebook = [
+    "ipywidgets",
+]
+
+[project.urls]
+Homepage = "https://github.com/edwinnglabs/bunobee"
+Repository = "https://github.com/edwinnglabs/bunobee"
+Issues = "https://github.com/edwinnglabs/bunobee/issues"
+Documentation = "https://github.com/edwinnglabs/bunobee#readme"
+
+[build-system]
+requires = ["setuptools"]
+build-backend = "setuptools.build_meta"
+
+[tool.setuptools.packages.find]
+where = ["src"]
+
+[tool.black]
+line-length = 120
+target-version = ['py311']
+include = '\.pyi?$'
+extend-exclude = ""
+
+[tool.pytest.ini_options]
+log_cli = true
+log_cli_level = "INFO"

bunobee-0.0.4/setup.cfg

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

bunobee-0.0.4/src/bunobee/__init__.py

@@ -0,0 +1,4 @@
+from importlib.metadata import version
+
+name = "bunobee"
+__version__ = version("bunobee")

bunobee-0.0.4/src/bunobee/hyper_tune.py

@@ -0,0 +1,180 @@
+import jax
+import jax.numpy as jnp
+import numpy as np
+from jax import vmap
+import logging
+import xarray as xr
+
+from typing import Dict, Tuple
+
+from .utils import flatten_front_dim
+from .models.dlt import run_dlt_model
+
+logger = logging.getLogger("wunku")
+
+
+def slice_trend_single(trend_comp: jnp.ndarray, h: int) -> jnp.ndarray:
+    """Slices a single trend into overlapping windows of size h.
+
+    Args
+    ----
+    trend_comp: array representing the trend component with shape (n_steps, )
+    h: Size of the window to slice the trend into.
+
+    Returns
+    -------
+    A 2D array of shape (n_windows, h) where each row is a window of size h where
+    num_windows = n_steps - h + 1.
+    """
+    n_steps = trend_comp.shape[0]
+    num_windows = n_steps - h + 1
+
+    def get_window(t_idx):
+        return jax.lax.dynamic_slice(trend_comp, start_indices=(t_idx,), slice_sizes=(h,))
+
+    t_indices = jnp.arange(num_windows)
+    # (n_windows, h)
+    return vmap(get_window)(t_indices)
+
+
+def slice_trend(trend_comp, h):
+    """Slices multiple trends into overlapping windows of size h.
+    Args
+    ----
+    trend_comp: array representing multiple trend components with shape (N_samples, n_steps)
+    h: Size of the window to slice the trends into.
+
+    Returns
+    -------
+    A 3D array of shape (N_samples, n_windows, h) where each row is a window of size h
+    and num_windows = n_steps - h + 1 for each trend.
+    """
+    # trends: shape (N_samples, T)
+    return vmap(lambda trend_comp: slice_trend_single(trend_comp, h))(trend_comp)
+
+
+def generate_forecast_span_samples(idata: xr.Dataset, h: int) -> jnp.ndarray:
+    # how to do a sliding window prediction?
+    dlt_comp = flatten_front_dim(idata["dlt_comp"].to_numpy(), n=2)
+    reg_comp = flatten_front_dim(idata["reg_comp"].to_numpy(), n=2)
+
+    logger.debug("dlt_comp shape:", dlt_comp.shape)
+    logger.debug("reg_comp shape:", reg_comp.shape)
+
+    dlt_comp_slice = dlt_comp[:, : -(h - 1), None]
+    reg_comp_slice = slice_trend(reg_comp, h=h)
+
+    # (n_samples, n_windows, h)
+    yhat_span = dlt_comp_slice + reg_comp_slice
+    return yhat_span
+
+
+def compute_log_likelihood(yhat_span, sigma, y):
+    """Computes the log likelihood of the forecasted span given the observations.
+
+    Args
+    ----
+    yhat_span: Forecasted span of shape (n_samples, n_windows, h).
+    sigma: Standard deviation of the noise, shape (n_samples,)
+    y_slice: Observations for the forecasted span, shape (n_windows, h).
+
+    Returns
+    -------
+    A 3D array of log probabilities of shape (n_samples, n_windows, h).
+    """
+    h = yhat_span.shape[-1]
+    # reshape for broadcasting
+    # (n_samples, 1, 1)
+    sigma_broadcast = sigma[:, None, None]
+    # (n_windos, h)
+    y_slice = slice_trend_single(y, h=h)
+    # (1, n_windows, h)
+    y_slice_broadcast = y_slice[None, :, :]
+
+    # compute squared error
+    sq_error = (y_slice_broadcast - yhat_span) ** 2
+
+    # compute log likelihood per (s, t, h)
+    log_prob = -0.5 * jnp.log(2 * jnp.pi) - jnp.log(sigma_broadcast) - 0.5 * sq_error / (sigma_broadcast**2)
+
+    return log_prob
+
+
+def compute_wbic(loglk: jnp.ndarray) -> float:
+    """Compute Weighted Bayesian Information Criterion (WBIC) from log likelihood.
+    Args
+    ----
+    loglk: array-like, log likelihood per samples per obs; should be with values of shape (n_samples, n_steps, h)
+    """
+    # in original paper, they use sum but it leads to unstable scale due to different size of datasets
+    # we use mean instead
+    loglk_per_sample = jnp.nanmean(loglk, axis=(-1, -2))
+    # (n_steps * h, )
+    nobs = loglk.shape[-1] * loglk.shape[-2]
+    beta = 1.0 / jnp.log(nobs)
+    wbic = -(1.0 / beta) * jnp.nanmean(loglk_per_sample)
+    wbic = float(wbic)
+    return wbic
+
+
+def run_dlt_model_and_compute_wbic(params: Tuple, data: Dict[str, np.ndarray], h: int = 12) -> float:
+    lev_sm, slp_sm, theta = params
+    y = data["y"]
+    x_seas = data["x_seas"]
+    x_glb_trend = data["x_glb_trend"]
+    logger.info(f"Trying lev_sm={lev_sm:.4f}, slp_sm={slp_sm:.4f}, theta={theta:.4f}")
+
+    idata = run_dlt_model(
+        lev_sm=lev_sm,
+        slp_sm=slp_sm,
+        theta=theta,
+        x_seas=x_seas,
+        x_glb_trend=x_glb_trend,
+        y=y,
+    )
+
+    yhat_span = generate_forecast_span_samples(idata=idata, h=h)
+    loglk = compute_log_likelihood(
+        yhat_span=yhat_span,
+        sigma=flatten_front_dim(idata["sigma"].to_numpy(), n=2),
+        y=y,
+    )
+    wbic = compute_wbic(loglk)
+
+    print(f"WBIC: {wbic:.4f}")
+    return wbic
+
+
+def hyper_tuning_dlt_with_wbic(
+    data: Dict[str, np.ndarray],
+    # forecast horizon
+    h: int = 12,
+    n_calls: int = 15,
+    random_state: int = 42,
+):
+    print("Starting hyperparameter tuning for DLT model using WBIC...")
+    # print args
+    print(f"h: {h}, n_calls: {n_calls}, random_state: {random_state}")
+    # try import skopt
+    try:
+        from skopt import gp_minimize
+        from skopt.space import Real
+    except ImportError:
+        raise ImportError("Please install scikit-optimize to run hyperparameter tuning.")
+
+    # Define the hyperparam space
+    search_space = [
+        Real(0.0001, 0.1, prior="log-uniform", name="lev_sm"),
+        Real(0.001, 0.1, prior="log-uniform", name="slp_sm"),
+        Real(0, 1.0, prior="uniform", name="theta"),
+    ]
+
+    # Run Bayesian Optimization
+    result = gp_minimize(
+        func=lambda params: run_dlt_model_and_compute_wbic(params=params, data=data, h=h),
+        dimensions=search_space,
+        n_calls=n_calls,
+        random_state=random_state,
+    )
+
+    return result

bunobee-0.0.4/src/bunobee/models/__init__.py

@@ -0,0 +1 @@
+