jaxfolio 0.1.0__py3-none-any.whl

jaxfolio/__init__.py

@@ -0,0 +1,110 @@
+"""jaxfolio — portfolio optimization and options strategies in JAX.
+
+A modern, JAX-powered library covering traditional, learning-based, and
+graph-based portfolio optimizers, a differentiable options toolkit, a
+walk-forward backtester, and dark-themed visualizations.
+
+Quickstart
+----------
+>>> import jaxfolio as jf
+>>> returns = jf.generate_returns(n_assets=8, seed=0)
+>>> result = jf.maximum_sharpe(returns)
+>>> result.top(3)
+"""
+
+from __future__ import annotations
+
+from jaxfolio.custom import CustomStrategy, custom_strategy
+from jaxfolio.data import (
+    generate_prices,
+    generate_returns,
+    load_csv,
+    load_option_chain,
+    load_parquet,
+    load_yfinance,
+    to_returns,
+    train_test_split,
+)
+from jaxfolio.llm import (
+    llm_agent_portfolio,
+    llm_black_litterman,
+    llm_sentiment_portfolio,
+)
+from jaxfolio.optimizers import (
+    black_litterman,
+    deep_sharpe,
+    equal_weight,
+    hierarchical_equal_risk,
+    hierarchical_risk_parity,
+    inverse_volatility,
+    kelly,
+    maximum_diversification,
+    maximum_sharpe,
+    mean_variance,
+    min_cvar,
+    minimum_variance,
+    mst_centrality,
+    online_gradient,
+    risk_parity,
+)
+from jaxfolio.registry import (
+    _register_builtins,
+    get_strategy,
+    list_strategies,
+    register_strategy,
+    strategy_info,
+    unregister,
+)
+from jaxfolio.types import OptimizerConfig, PortfolioResult
+
+# Register the built-in optimizers so they show up in the strategy registry
+# alongside any user-defined strategies.
+_register_builtins()
+
+__version__ = "0.1.0"
+
+__all__ = [
+    "__version__",
+    # types
+    "PortfolioResult",
+    "OptimizerConfig",
+    # registry & custom strategies
+    "register_strategy",
+    "get_strategy",
+    "list_strategies",
+    "strategy_info",
+    "unregister",
+    "custom_strategy",
+    "CustomStrategy",
+    # data
+    "generate_prices",
+    "generate_returns",
+    "load_csv",
+    "load_parquet",
+    "load_yfinance",
+    "load_option_chain",
+    "to_returns",
+    "train_test_split",
+    # classical optimizers
+    "equal_weight",
+    "inverse_volatility",
+    "minimum_variance",
+    "mean_variance",
+    "maximum_sharpe",
+    "maximum_diversification",
+    "risk_parity",
+    "kelly",
+    "min_cvar",
+    "black_litterman",
+    # learning optimizers
+    "deep_sharpe",
+    "online_gradient",
+    # graph optimizers
+    "hierarchical_risk_parity",
+    "hierarchical_equal_risk",
+    "mst_centrality",
+    # LLM (local-model) strategies
+    "llm_black_litterman",
+    "llm_sentiment_portfolio",
+    "llm_agent_portfolio",
+]

jaxfolio/backtest/__init__.py

@@ -0,0 +1,39 @@
+"""Backtesting engine and performance metrics."""
+
+from jaxfolio.backtest.engine import (
+    BacktestResult,
+    backtest,
+    compare,
+    metrics_table,
+)
+from jaxfolio.backtest.metrics import (
+    annualized_return,
+    annualized_volatility,
+    calmar_ratio,
+    conditional_value_at_risk,
+    cumulative_returns,
+    drawdown_series,
+    max_drawdown,
+    sharpe_ratio,
+    sortino_ratio,
+    summary,
+    value_at_risk,
+)
+
+__all__ = [
+    "backtest",
+    "compare",
+    "metrics_table",
+    "BacktestResult",
+    "annualized_return",
+    "annualized_volatility",
+    "sharpe_ratio",
+    "sortino_ratio",
+    "max_drawdown",
+    "calmar_ratio",
+    "value_at_risk",
+    "conditional_value_at_risk",
+    "cumulative_returns",
+    "drawdown_series",
+    "summary",
+]

jaxfolio/backtest/engine.py

@@ -0,0 +1,141 @@
+"""A vectorized walk-forward backtester.
+
+The backtester is optimizer-agnostic: it takes any callable
+``optimizer(returns_window) -> PortfolioResult`` and rebalances on a fixed
+schedule, applying linear transaction costs on turnover. It returns a
+:class:`BacktestResult` holding the strategy return series, weight history, and a
+metrics summary — ready for the plotting utilities.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Callable
+from dataclasses import dataclass, field
+
+import numpy as np
+import pandas as pd
+
+from jaxfolio.backtest import metrics as M
+from jaxfolio.types import PortfolioResult
+
+Optimizer = Callable[[pd.DataFrame], PortfolioResult]
+
+
+@dataclass
+class BacktestResult:
+    """Output of a backtest run."""
+
+    name: str
+    returns: pd.Series
+    weights: pd.DataFrame
+    turnover: pd.Series
+    metrics: dict[str, float] = field(default_factory=dict)
+
+    @property
+    def equity_curve(self) -> pd.Series:
+        return pd.Series(M.cumulative_returns(self.returns), index=self.returns.index)
+
+    @property
+    def drawdown(self) -> pd.Series:
+        return pd.Series(M.drawdown_series(self.returns), index=self.returns.index)
+
+
+def backtest(
+    returns: pd.DataFrame,
+    optimizer: Optimizer,
+    *,
+    name: str | None = None,
+    lookback: int = 252,
+    rebalance_every: int = 21,
+    transaction_cost: float = 0.0010,
+    periods_per_year: int = 252,
+    risk_free: float = 0.0,
+) -> BacktestResult:
+    """Walk-forward backtest of a single optimizer.
+
+    Parameters
+    ----------
+    returns:
+        Asset return panel (rows = periods, columns = assets).
+    optimizer:
+        Callable mapping a trailing return window to a :class:`PortfolioResult`.
+    lookback:
+        Number of trailing periods handed to the optimizer at each rebalance.
+    rebalance_every:
+        Rebalance frequency in periods (21 ≈ monthly for daily data).
+    transaction_cost:
+        Proportional cost per unit of turnover (one-way), e.g. ``0.0010`` = 10bps.
+
+    Returns
+    -------
+    BacktestResult
+        Net-of-cost strategy returns, weight history, turnover, and metrics.
+    """
+    assets = list(returns.columns)
+    n = len(assets)
+    dates = returns.index
+    ret_vals = returns.to_numpy()
+
+    weights = np.zeros(n)
+    weight_hist: dict[pd.Timestamp, np.ndarray] = {}
+    turnover_hist: dict[pd.Timestamp, float] = {}
+    strat_returns = np.zeros(len(returns))
+
+    for t in range(lookback, len(returns)):
+        # Rebalance at the cadence; otherwise let weights drift with returns.
+        if (t - lookback) % rebalance_every == 0:
+            window = returns.iloc[t - lookback : t]
+            result = optimizer(window)
+            target = _align_weights(result, assets)
+            cost = transaction_cost * np.abs(target - weights).sum()
+            turnover_hist[dates[t]] = float(np.abs(target - weights).sum())
+            weights = target
+        else:
+            cost = 0.0
+
+        period_ret = ret_vals[t]
+        gross = float(np.dot(weights, period_ret))
+        strat_returns[t] = gross - cost
+        weight_hist[dates[t]] = weights.copy()
+
+        # Drift weights with realized returns (buy-and-hold between rebalances).
+        grown = weights * (1.0 + period_ret)
+        total = grown.sum()
+        if total > 0:
+            weights = grown / total
+
+    ret_series = pd.Series(strat_returns, index=dates, name=name or "strategy")
+    ret_series = ret_series.iloc[lookback:]
+    w_df = pd.DataFrame.from_dict(weight_hist, orient="index", columns=assets)
+    to_series = pd.Series(turnover_hist, name="turnover")
+
+    summary = M.summary(ret_series, risk_free=risk_free, periods_per_year=periods_per_year)
+    summary["avg_turnover"] = float(to_series.mean()) if len(to_series) else 0.0
+    return BacktestResult(
+        name=name or "strategy",
+        returns=ret_series,
+        weights=w_df,
+        turnover=to_series,
+        metrics=summary,
+    )
+
+
+def compare(
+    returns: pd.DataFrame,
+    optimizers: dict[str, Optimizer],
+    **kwargs,
+) -> dict[str, BacktestResult]:
+    """Backtest several optimizers on the same data and return a name->result map."""
+    return {name: backtest(returns, opt, name=name, **kwargs) for name, opt in optimizers.items()}
+
+
+def metrics_table(results: dict[str, BacktestResult]) -> pd.DataFrame:
+    """Assemble a tidy metrics comparison table across backtest results."""
+    rows = {name: res.metrics for name, res in results.items()}
+    return pd.DataFrame(rows).T
+
+
+def _align_weights(result: PortfolioResult, assets: list[str]) -> np.ndarray:
+    """Reorder/pad a result's weights to the backtest's asset ordering."""
+    lookup = dict(zip(result.assets, result.weights, strict=True))
+    return np.array([lookup.get(a, 0.0) for a in assets], dtype=float)

jaxfolio/backtest/metrics.py

@@ -0,0 +1,133 @@
+"""Performance and risk metrics for return series.
+
+All functions accept a 1-D array-like of periodic returns and return plain
+floats. Annualization uses ``periods_per_year`` (252 for daily by default).
+"""
+
+from __future__ import annotations
+
+import numpy as np
+import pandas as pd
+
+_PPY = 252
+
+
+def _to_array(returns) -> np.ndarray:
+    if isinstance(returns, pd.Series | pd.DataFrame):
+        returns = returns.to_numpy()
+    return np.asarray(returns, dtype=float).reshape(-1)
+
+
+def annualized_return(returns, periods_per_year: int = _PPY) -> float:
+    """Geometric annualized return (CAGR) of a periodic return series."""
+    r = _to_array(returns)
+    if r.size == 0:
+        return 0.0
+    growth = np.prod(1.0 + r)
+    years = r.size / periods_per_year
+    if growth <= 0:
+        return -1.0
+    return float(growth ** (1.0 / years) - 1.0)
+
+
+def annualized_volatility(returns, periods_per_year: int = _PPY) -> float:
+    """Annualized standard deviation of returns."""
+    r = _to_array(returns)
+    return float(np.std(r, ddof=1) * np.sqrt(periods_per_year)) if r.size > 1 else 0.0
+
+
+def sharpe_ratio(returns, risk_free: float = 0.0, periods_per_year: int = _PPY) -> float:
+    """Annualized Sharpe ratio (``risk_free`` is a per-period rate)."""
+    r = _to_array(returns) - risk_free
+    sd = np.std(r, ddof=1)
+    if sd == 0:
+        return 0.0
+    return float(np.mean(r) / sd * np.sqrt(periods_per_year))
+
+
+def sortino_ratio(returns, risk_free: float = 0.0, periods_per_year: int = _PPY) -> float:
+    """Annualized Sortino ratio (downside-deviation denominator).
+
+    With no downside (zero denominator), returns ``+inf`` for positive mean
+    excess return, ``-inf`` for negative, and ``0`` for a flat series — so a
+    downside-free strategy ranks above one with drawdowns rather than tying at 0.
+    """
+    r = _to_array(returns) - risk_free
+    downside = r[r < 0]
+    dd = np.sqrt(np.mean(downside**2)) if downside.size else 0.0
+    mean = float(np.mean(r)) if r.size else 0.0
+    if dd == 0:
+        return float(np.sign(mean) * np.inf) if mean != 0 else 0.0
+    return float(mean / dd * np.sqrt(periods_per_year))
+
+
+def cumulative_returns(returns) -> np.ndarray:
+    """Cumulative wealth path starting from 1.0 (equity curve)."""
+    r = _to_array(returns)
+    return np.cumprod(1.0 + r)
+
+
+def drawdown_series(returns) -> np.ndarray:
+    """Drawdown at each point: ``equity / running_peak - 1`` (<= 0)."""
+    equity = cumulative_returns(returns)
+    peak = np.maximum.accumulate(equity)
+    return equity / peak - 1.0
+
+
+def max_drawdown(returns) -> float:
+    """Maximum peak-to-trough drawdown (a negative number)."""
+    dd = drawdown_series(returns)
+    return float(dd.min()) if dd.size else 0.0
+
+
+def calmar_ratio(returns, periods_per_year: int = _PPY) -> float:
+    """Annualized return divided by the absolute max drawdown.
+
+    With no drawdown (zero denominator), returns ``+inf`` for a positive
+    annualized return, ``-inf`` for negative, and ``0`` for flat — so a
+    drawdown-free strategy is not mis-ranked as the worst performer.
+    """
+    ann = annualized_return(returns, periods_per_year)
+    mdd = abs(max_drawdown(returns))
+    if mdd == 0:
+        return float(np.sign(ann) * np.inf) if ann != 0 else 0.0
+    return float(ann / mdd)
+
+
+def value_at_risk(returns, alpha: float = 0.95) -> float:
+    """Historical Value-at-Risk at confidence ``alpha`` (a positive loss)."""
+    r = _to_array(returns)
+    if r.size == 0:
+        return 0.0
+    return float(-np.quantile(r, 1.0 - alpha))
+
+
+def conditional_value_at_risk(returns, alpha: float = 0.95) -> float:
+    """Historical CVaR / expected shortfall at confidence ``alpha``."""
+    r = _to_array(returns)
+    if r.size == 0:
+        return 0.0
+    var = -value_at_risk(r, alpha)
+    tail = r[r <= var]
+    return float(-tail.mean()) if tail.size else float(-var)
+
+
+def hit_rate(returns) -> float:
+    """Fraction of periods with a positive return."""
+    r = _to_array(returns)
+    return float(np.mean(r > 0)) if r.size else 0.0
+
+
+def summary(returns, risk_free: float = 0.0, periods_per_year: int = _PPY) -> dict[str, float]:
+    """Bundle the headline metrics into a single dict for reporting/plots."""
+    return {
+        "annual_return": annualized_return(returns, periods_per_year),
+        "annual_volatility": annualized_volatility(returns, periods_per_year),
+        "sharpe": sharpe_ratio(returns, risk_free, periods_per_year),
+        "sortino": sortino_ratio(returns, risk_free, periods_per_year),
+        "max_drawdown": max_drawdown(returns),
+        "calmar": calmar_ratio(returns, periods_per_year),
+        "var_95": value_at_risk(returns, 0.95),
+        "cvar_95": conditional_value_at_risk(returns, 0.95),
+        "hit_rate": hit_rate(returns),
+    }

jaxfolio/constraints/__init__.py

@@ -0,0 +1,15 @@
+"""Constraint projections for weight vectors."""
+
+from jaxfolio.constraints.projections import (
+    normalize_weights,
+    project_box_budget,
+    project_simplex,
+    softmax_weights,
+)
+
+__all__ = [
+    "project_simplex",
+    "project_box_budget",
+    "normalize_weights",
+    "softmax_weights",
+]

jaxfolio/constraints/projections.py

@@ -0,0 +1,87 @@
+"""Constraint projections used by the projected-gradient optimizers.
+
+All projections are pure JAX and jit/vmap-safe. The two workhorses are:
+
+* :func:`project_simplex` — Euclidean projection onto the probability simplex
+  ``{w : w >= 0, sum(w) = 1}`` (long-only, fully-invested).
+* :func:`project_box_budget` — projection onto a box ``[lo, hi]`` intersected
+  with the budget hyperplane ``sum(w) = budget`` (allows bounded shorting).
+"""
+
+from __future__ import annotations
+
+import jax
+import jax.numpy as jnp
+
+Array = jnp.ndarray
+
+
+def project_simplex(v: Array, budget: float = 1.0) -> Array:
+    """Euclidean projection of ``v`` onto the scaled simplex summing to ``budget``.
+
+    Implements the classic Duchi et al. (2008) sort-based algorithm, which is
+    exact and differentiable almost everywhere.
+    """
+    n = v.shape[0]
+    u = jnp.sort(v)[::-1]
+    cssv = jnp.cumsum(u) - budget
+    ind = jnp.arange(1, n + 1)
+    cond = u - cssv / ind > 0
+    # rho = number of positive-threshold coordinates.
+    rho = jnp.sum(cond)
+    theta = cssv[rho - 1] / rho
+    return jnp.maximum(v - theta, 0.0)
+
+
+def project_box_budget(
+    v: Array,
+    lower: float = 0.0,
+    upper: float = 1.0,
+    budget: float = 1.0,
+    *,
+    max_iter: int = 50,
+    tol: float = 1e-10,
+) -> Array:
+    """Project ``v`` onto ``{w : lower <= w <= upper, sum(w) = budget}``.
+
+    Solves for the Lagrange multiplier ``tau`` on the budget constraint by
+    bisection: ``w(tau) = clip(v - tau, lower, upper)`` is monotone decreasing in
+    ``tau``, so we bisect until ``sum(w(tau)) == budget``. Feasible whenever
+    ``n*lower <= budget <= n*upper``.
+    """
+    n = v.shape[0]
+
+    def sum_at(tau: Array) -> Array:
+        return jnp.sum(jnp.clip(v - tau, lower, upper))
+
+    # Bracket tau: sum decreases as tau grows.
+    lo = jnp.min(v) - upper
+    hi = jnp.max(v) - lower
+
+    def body(_, bounds):
+        lo, hi = bounds
+        mid = 0.5 * (lo + hi)
+        s = sum_at(mid)
+        # If sum too big, need larger tau -> move lo up; else move hi down.
+        too_big = s > budget
+        lo = jnp.where(too_big, mid, lo)
+        hi = jnp.where(too_big, hi, mid)
+        return (lo, hi)
+
+    lo, hi = jax.lax.fori_loop(0, max_iter, body, (lo, hi))
+    tau = 0.5 * (lo + hi)
+    w = jnp.clip(v - tau, lower, upper)
+    # Correct tiny residual so weights sum exactly to budget.
+    w = w + (budget - jnp.sum(w)) / n
+    return w
+
+
+def normalize_weights(w: Array, budget: float = 1.0) -> Array:
+    """Rescale weights to sum to ``budget`` (assumes a non-zero sum)."""
+    s = jnp.sum(w)
+    return jnp.where(jnp.abs(s) > 1e-12, w * (budget / s), jnp.full_like(w, budget / w.shape[0]))
+
+
+def softmax_weights(logits: Array, budget: float = 1.0) -> Array:
+    """Map unconstrained logits to long-only weights via softmax (sums to budget)."""
+    return budget * jax.nn.softmax(logits)