Skip to content

Data & moments

See the Data guide.

Synthetic data

synthetic

Synthetic market data generation.

A correlated geometric-Brownian-motion generator that produces realistic price panels without any network access. Used throughout the tests and examples so everything is reproducible and offline-friendly.

generate_prices 

generate_prices(
    n_assets: int = 12,
    n_days: int = 756,
    *,
    seed: int = 0,
    start_price: float = 100.0,
    mean_annual_return: float = 0.08,
    annual_vol_range: tuple[float, float] = (0.15, 0.45),
    correlation_strength: float = 0.5,
    start_date: str = "2021-01-01",
    tickers: list[str] | None = None,
) -> DataFrame

Generate a panel of correlated daily prices via geometric Brownian motion.

Parameters:

Name Type Description Default
n_assets int

Panel shape. n_days uses ~252 trading days per year.

 12
n_days int

Panel shape. n_days uses ~252 trading days per year.

 12
seed int

RNG seed for reproducibility.

 0
mean_annual_return float

Center of the per-asset drift (drifts are jittered around this).

 0.08
annual_vol_range tuple[float, float]

Range from which per-asset annual volatilities are drawn.

 (0.15, 0.45)
correlation_strength float

Average cross-asset correlation in [0, 1).

 0.5
tickers list[str] | None

Optional explicit column names; otherwise ASSET_00 ... are used.

 None

Returns:

Type Description
DataFrame

Indexed by business day, one column per asset, containing prices.
Source code in src/jaxfolio/data/synthetic.py 
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
def generate_prices(
    n_assets: int = 12,
    n_days: int = 756,
    *,
    seed: int = 0,
    start_price: float = 100.0,
    mean_annual_return: float = 0.08,
    annual_vol_range: tuple[float, float] = (0.15, 0.45),
    correlation_strength: float = 0.5,
    start_date: str = "2021-01-01",
    tickers: list[str] | None = None,
) -> pd.DataFrame:
    """Generate a panel of correlated daily prices via geometric Brownian motion.

    Parameters
    ----------
    n_assets, n_days:
        Panel shape. ``n_days`` uses ~252 trading days per year.
    seed:
        RNG seed for reproducibility.
    mean_annual_return:
        Center of the per-asset drift (drifts are jittered around this).
    annual_vol_range:
        Range from which per-asset annual volatilities are drawn.
    correlation_strength:
        Average cross-asset correlation in ``[0, 1)``.
    tickers:
        Optional explicit column names; otherwise ``ASSET_00`` ... are used.

    Returns
    -------
    pandas.DataFrame
        Indexed by business day, one column per asset, containing prices.
    """
    rng = np.random.default_rng(seed)
    if tickers is None:
        tickers = [f"ASSET_{i:02d}" for i in range(n_assets)]
    elif len(tickers) != n_assets:
        raise ValueError("len(tickers) must equal n_assets")

    dt = 1.0 / 252.0
    vols = rng.uniform(*annual_vol_range, size=n_assets)
    drifts = rng.normal(mean_annual_return, 0.03, size=n_assets)

    corr = _random_correlation(n_assets, rng, correlation_strength)
    chol = np.linalg.cholesky(corr)

    z = rng.standard_normal(size=(n_days, n_assets)) @ chol.T
    # GBM log-returns: (mu - 0.5 sigma^2) dt + sigma sqrt(dt) dW
    log_rets = (drifts - 0.5 * vols**2) * dt + vols * np.sqrt(dt) * z
    log_price = np.log(start_price) + np.cumsum(log_rets, axis=0)
    prices = np.exp(log_price)

    dates = pd.bdate_range(start=start_date, periods=n_days)
    return pd.DataFrame(prices, index=dates, columns=tickers)

generate_returns 

generate_returns(
    n_assets: int = 12, n_days: int = 756, **kwargs
) -> DataFrame

Convenience wrapper returning simple daily returns from synthetic prices.

Source code in src/jaxfolio/data/synthetic.py 
87
88
89
90
def generate_returns(n_assets: int = 12, n_days: int = 756, **kwargs) -> pd.DataFrame:
    """Convenience wrapper returning simple daily returns from synthetic prices."""
    prices = generate_prices(n_assets=n_assets, n_days=n_days, **kwargs)
    return prices.pct_change().dropna()

Loaders

loaders

Data ingestion: CSV, Parquet, and (optionally) live market data via yfinance.

The CSV/Parquet loaders have no third-party dependency beyond pandas. The yfinance loaders are guarded so importing this module never requires the [data] extra to be installed.

load_csv 

load_csv(
    path: str | Path,
    *,
    date_col: str = "date",
    asset_col: str | None = None,
    price_col: str = "close",
    **read_csv_kwargs,
) -> DataFrame

Load a price panel from CSV.

Two layouts are supported:

  • Wide (default when asset_col is None): one column per asset, plus a date column used as the index.
  • Long/tidy (when asset_col is given): columns date_col, asset_col, price_col are pivoted into a wide panel.
Source code in src/jaxfolio/data/loaders.py 
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
def load_csv(
    path: str | Path,
    *,
    date_col: str = "date",
    asset_col: str | None = None,
    price_col: str = "close",
    **read_csv_kwargs,
) -> pd.DataFrame:
    """Load a price panel from CSV.

    Two layouts are supported:

    * **Wide** (default when ``asset_col`` is ``None``): one column per asset,
      plus a date column used as the index.
    * **Long/tidy** (when ``asset_col`` is given): columns ``date_col``,
      ``asset_col``, ``price_col`` are pivoted into a wide panel.
    """
    df = pd.read_csv(path, **read_csv_kwargs)
    if asset_col is not None:
        return _pivot_to_panel(df, date_col, asset_col, price_col)
    df[date_col] = pd.to_datetime(df[date_col])
    return df.set_index(date_col).sort_index()

load_parquet 

load_parquet(
    path: str | Path,
    *,
    date_col: str = "date",
    asset_col: str | None = None,
    price_col: str = "close",
) -> DataFrame

Load a price panel from Parquet (same layouts as :func:load_csv).

Source code in src/jaxfolio/data/loaders.py 
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
def load_parquet(
    path: str | Path,
    *,
    date_col: str = "date",
    asset_col: str | None = None,
    price_col: str = "close",
) -> pd.DataFrame:
    """Load a price panel from Parquet (same layouts as :func:`load_csv`)."""
    df = pd.read_parquet(path)
    if asset_col is not None:
        return _pivot_to_panel(df, date_col, asset_col, price_col)
    if date_col in df.columns:
        df[date_col] = pd.to_datetime(df[date_col])
        df = df.set_index(date_col)
    return df.sort_index()

load_yfinance 

load_yfinance(
    tickers: list[str] | str,
    *,
    start: str | None = None,
    end: str | None = None,
    period: str | None = "2y",
    interval: str = "1d",
    price_field: str = "Close",
) -> DataFrame

Download a price panel from Yahoo Finance (requires the [data] extra).

Returns a wide date-by-ticker panel of price_field values.

Source code in src/jaxfolio/data/loaders.py 
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
def load_yfinance(
    tickers: list[str] | str,
    *,
    start: str | None = None,
    end: str | None = None,
    period: str | None = "2y",
    interval: str = "1d",
    price_field: str = "Close",
) -> pd.DataFrame:
    """Download a price panel from Yahoo Finance (requires the ``[data]`` extra).

    Returns a wide date-by-ticker panel of ``price_field`` values.
    """
    try:
        import yfinance as yf
    except ImportError as exc:  # pragma: no cover - exercised only without extra
        raise ImportError(
            "load_yfinance requires the optional 'data' extra: "
            "install with `uv sync --extra data` or `pip install jaxfolio[data]`."
        ) from exc

    if isinstance(tickers, str):
        tickers = [t.strip() for t in tickers.split(",") if t.strip()]

    raw = yf.download(
        tickers,
        start=start,
        end=end,
        period=None if start else period,
        interval=interval,
        auto_adjust=True,
        progress=False,
    )
    if isinstance(raw.columns, pd.MultiIndex):
        panel = raw[price_field]
    else:  # single ticker
        panel = raw[[price_field]]
        panel.columns = tickers
    return panel.dropna(how="all").sort_index()

load_option_chain 

load_option_chain(
    ticker: str, *, expiry: str | None = None
) -> DataFrame

Fetch and normalize an option chain from Yahoo Finance ([data] extra).

Returns a tidy frame with columns [type, strike, expiry, last, bid, ask, volume, open_interest, implied_vol] for calls and puts combined. If expiry is None the nearest expiry is used.

Source code in src/jaxfolio/data/loaders.py 
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
def load_option_chain(
    ticker: str,
    *,
    expiry: str | None = None,
) -> pd.DataFrame:
    """Fetch and normalize an option chain from Yahoo Finance (``[data]`` extra).

    Returns a tidy frame with columns ``[type, strike, expiry, last, bid, ask,
    volume, open_interest, implied_vol]`` for calls and puts combined. If
    ``expiry`` is ``None`` the nearest expiry is used.
    """
    try:
        import yfinance as yf
    except ImportError as exc:  # pragma: no cover - exercised only without extra
        raise ImportError(
            "load_option_chain requires the optional 'data' extra: "
            "install with `uv sync --extra data`."
        ) from exc

    tk = yf.Ticker(ticker)
    expiries = tk.options
    if not expiries:
        raise ValueError(f"No listed options found for {ticker!r}")
    chosen = expiry or expiries[0]
    chain = tk.option_chain(chosen)

    def _norm(df: pd.DataFrame, kind: str) -> pd.DataFrame:
        cols = {
            "strike": "strike",
            "lastPrice": "last",
            "bid": "bid",
            "ask": "ask",
            "volume": "volume",
            "openInterest": "open_interest",
            "impliedVolatility": "implied_vol",
        }
        out = df[list(cols)].rename(columns=cols)
        out.insert(0, "type", kind)
        out.insert(2, "expiry", pd.to_datetime(chosen))
        return out

    calls = _norm(chain.calls, "call")
    puts = _norm(chain.puts, "put")
    return pd.concat([calls, puts], ignore_index=True)

Returns & splitting

returns

Return computation, alignment, and train/test splitting.

All functions accept and return pandas objects so they compose cleanly with the loaders, and hand off clean float arrays to the JAX optimizers downstream.

to_returns 

to_returns(
    prices: DataFrame,
    *,
    kind: str = "simple",
    dropna: bool = True,
) -> DataFrame

Convert a price panel to returns.

Parameters:

Name Type Description Default
prices DataFrame

Price panel (rows = dates, columns = assets).

 required
kind str

"simple" for arithmetic returns p_t / p_{t-1} - 1 or "log" for log(p_t / p_{t-1}).

 'simple'
dropna bool

Drop the leading NaN row (and any all-NaN rows) if True.

 True
Source code in src/jaxfolio/data/returns.py 
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
def to_returns(
    prices: pd.DataFrame,
    *,
    kind: str = "simple",
    dropna: bool = True,
) -> pd.DataFrame:
    """Convert a price panel to returns.

    Parameters
    ----------
    prices:
        Price panel (rows = dates, columns = assets).
    kind:
        ``"simple"`` for arithmetic returns ``p_t / p_{t-1} - 1`` or ``"log"``
        for ``log(p_t / p_{t-1})``.
    dropna:
        Drop the leading NaN row (and any all-NaN rows) if ``True``.
    """
    if kind == "simple":
        rets = prices.pct_change()
    elif kind == "log":
        rets = np.log(prices / prices.shift(1))
    else:
        raise ValueError("kind must be 'simple' or 'log'")
    if dropna:
        rets = rets.dropna(how="all")
    return rets

align 

align(
    *frames: DataFrame, how: str = "inner"
) -> list[DataFrame]

Align multiple frames on their common (or union) index.

Useful when combining price series pulled from different sources.

Source code in src/jaxfolio/data/returns.py 
42
43
44
45
46
47
48
49
50
51
52
def align(*frames: pd.DataFrame, how: str = "inner") -> list[pd.DataFrame]:
    """Align multiple frames on their common (or union) index.

    Useful when combining price series pulled from different sources.
    """
    if not frames:
        return []
    index = frames[0].index
    for f in frames[1:]:
        index = index.join(f.index, how=how)
    return [f.reindex(index) for f in frames]

clean_returns 

clean_returns(
    returns: DataFrame,
    *,
    max_missing_frac: float = 0.1,
    fill: str = "zero",
) -> DataFrame

Drop sparse columns and fill remaining gaps.

Columns missing more than max_missing_frac of observations are dropped. Remaining NaNs are filled with 0.0 ("zero") or forward-filled ("ffill").

Source code in src/jaxfolio/data/returns.py 
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
def clean_returns(
    returns: pd.DataFrame,
    *,
    max_missing_frac: float = 0.1,
    fill: str = "zero",
) -> pd.DataFrame:
    """Drop sparse columns and fill remaining gaps.

    Columns missing more than ``max_missing_frac`` of observations are dropped.
    Remaining NaNs are filled with ``0.0`` (``"zero"``) or forward-filled
    (``"ffill"``).
    """
    keep = returns.columns[returns.isna().mean() <= max_missing_frac]
    out = returns[keep].copy()
    if fill == "zero":
        out = out.fillna(0.0)
    elif fill == "ffill":
        out = out.ffill().fillna(0.0)
    else:
        raise ValueError("fill must be 'zero' or 'ffill'")
    return out

train_test_split 

train_test_split(
    returns: DataFrame, *, test_size: float = 0.25
) -> tuple[DataFrame, DataFrame]

Chronological split into in-sample and out-of-sample sets.

test_size is the fraction of the most recent observations held out.

Source code in src/jaxfolio/data/returns.py 
78
79
80
81
82
83
84
85
86
87
88
89
90
91
def train_test_split(
    returns: pd.DataFrame,
    *,
    test_size: float = 0.25,
) -> tuple[pd.DataFrame, pd.DataFrame]:
    """Chronological split into in-sample and out-of-sample sets.

    ``test_size`` is the fraction of the most recent observations held out.
    """
    if not 0.0 < test_size < 1.0:
        raise ValueError("test_size must be in (0, 1)")
    n = len(returns)
    cut = int(round(n * (1.0 - test_size)))
    return returns.iloc[:cut], returns.iloc[cut:]

annualization_factor 

annualization_factor(
    returns: DataFrame | Series, periods_per_year: int = 252
) -> int

Return the annualization factor (kept explicit for clarity in call sites).

Source code in src/jaxfolio/data/returns.py 
94
95
96
def annualization_factor(returns: pd.DataFrame | pd.Series, periods_per_year: int = 252) -> int:
    """Return the annualization factor (kept explicit for clarity in call sites)."""
    return periods_per_year

Moment estimators

estimators

Moment estimators (mean & covariance) implemented in JAX.

Every estimator takes a (T, N) return array and returns JAX arrays, so the downstream optimizers can jit through the whole pipeline. Inputs may be pandas frames or numpy arrays; :func:as_matrix normalizes them.

as_matrix 

as_matrix(
    returns: DataFrame | ndarray | Array,
) -> tuple[Array, list[str]]

Return (matrix, asset_names) from a returns object.

Accepts a pandas DataFrame (columns become asset names), or a raw array (names default to asset_0...). NaNs are replaced with zeros.

Source code in src/jaxfolio/moments/estimators.py 
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
def as_matrix(returns: pd.DataFrame | np.ndarray | Array) -> tuple[Array, list[str]]:
    """Return ``(matrix, asset_names)`` from a returns object.

    Accepts a pandas DataFrame (columns become asset names), or a raw array
    (names default to ``asset_0``...). NaNs are replaced with zeros.
    """
    if isinstance(returns, pd.DataFrame):
        names = [str(c) for c in returns.columns]
        mat = jnp.asarray(returns.to_numpy(dtype=float))
    else:
        mat = jnp.asarray(np.asarray(returns, dtype=float))
        if mat.ndim != 2:
            raise ValueError("returns must be 2-D (T x N)")
        names = [f"asset_{i}" for i in range(mat.shape[1])]
    mat = jnp.nan_to_num(mat, nan=0.0)
    return mat, names

mean_returns 

mean_returns(
    returns: Array, *, periods_per_year: int | None = None
) -> Array

Sample mean of returns; annualized if periods_per_year is given.

Source code in src/jaxfolio/moments/estimators.py 
35
36
37
38
39
40
def mean_returns(returns: Array, *, periods_per_year: int | None = None) -> Array:
    """Sample mean of returns; annualized if ``periods_per_year`` is given."""
    mu = jnp.mean(returns, axis=0)
    if periods_per_year is not None:
        mu = mu * periods_per_year
    return mu

sample_covariance 

sample_covariance(
    returns: Array, *, periods_per_year: int | None = None
) -> Array

Sample covariance matrix (denominator T - 1); optionally annualized.

Requires at least two observations; a single row has no defined covariance.

Source code in src/jaxfolio/moments/estimators.py 
43
44
45
46
47
48
49
50
51
52
53
54
55
def sample_covariance(returns: Array, *, periods_per_year: int | None = None) -> Array:
    """Sample covariance matrix (denominator ``T - 1``); optionally annualized.

    Requires at least two observations; a single row has no defined covariance.
    """
    t = returns.shape[0]
    if t < 2:
        raise ValueError(f"sample_covariance needs at least 2 observations, got {t}")
    demeaned = returns - jnp.mean(returns, axis=0, keepdims=True)
    cov = (demeaned.T @ demeaned) / (t - 1)
    if periods_per_year is not None:
        cov = cov * periods_per_year
    return cov

ewma_covariance 

ewma_covariance(
    returns: Array,
    *,
    halflife: float = 63.0,
    periods_per_year: int | None = None,
) -> Array

Exponentially-weighted covariance matrix.

Recent observations receive more weight; halflife is in periods (default ~one quarter of trading days).

Source code in src/jaxfolio/moments/estimators.py 
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
def ewma_covariance(
    returns: Array,
    *,
    halflife: float = 63.0,
    periods_per_year: int | None = None,
) -> Array:
    """Exponentially-weighted covariance matrix.

    Recent observations receive more weight; ``halflife`` is in periods (default
    ~one quarter of trading days).
    """
    t = returns.shape[0]
    decay = jnp.log(2.0) / halflife
    ages = jnp.arange(t)[::-1]  # most recent row -> age 0
    weights = jnp.exp(-decay * ages)
    weights = weights / jnp.sum(weights)
    wmean = jnp.sum(weights[:, None] * returns, axis=0)
    demeaned = returns - wmean
    cov = (demeaned * weights[:, None]).T @ demeaned
    # Symmetrize against floating point drift.
    cov = 0.5 * (cov + cov.T)
    if periods_per_year is not None:
        cov = cov * periods_per_year
    return cov

ledoit_wolf_covariance 

ledoit_wolf_covariance(
    returns: Array, *, periods_per_year: int | None = None
) -> tuple[Array, float]

Ledoit-Wolf shrinkage toward a scaled-identity target.

Returns (shrunk_cov, shrinkage_intensity). The shrinkage intensity is estimated analytically (Ledoit & Wolf, 2004) and clipped to [0, 1].

Source code in src/jaxfolio/moments/estimators.py 
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
def ledoit_wolf_covariance(
    returns: Array,
    *,
    periods_per_year: int | None = None,
) -> tuple[Array, float]:
    """Ledoit-Wolf shrinkage toward a scaled-identity target.

    Returns ``(shrunk_cov, shrinkage_intensity)``. The shrinkage intensity is
    estimated analytically (Ledoit & Wolf, 2004) and clipped to ``[0, 1]``.
    """
    t, n = returns.shape
    x = returns - jnp.mean(returns, axis=0, keepdims=True)
    sample = (x.T @ x) / t

    mu = jnp.trace(sample) / n
    target = mu * jnp.eye(n)

    # pi: sum of asymptotic variances of the sample covariance entries.
    x2 = x**2
    phi_mat = (x2.T @ x2) / t - sample**2
    pi_hat = jnp.sum(phi_mat)

    # rho: for an identity-scaled target the off-diagonal correction is zero;
    # only the diagonal contributes.
    rho_hat = jnp.sum(jnp.diag(phi_mat))

    gamma_hat = jnp.sum((sample - target) ** 2)
    # When the sample already equals the target (e.g. n == 1) gamma_hat is 0 and
    # the shrinkage is undefined; fall back to no shrinkage rather than NaN.
    safe_gamma = jnp.where(gamma_hat > 0, gamma_hat, 1.0)
    kappa = jnp.where(gamma_hat > 0, (pi_hat - rho_hat) / safe_gamma, 0.0)
    shrinkage = jnp.clip(kappa / t, 0.0, 1.0)

    shrunk = shrinkage * target + (1.0 - shrinkage) * sample
    if periods_per_year is not None:
        shrunk = shrunk * periods_per_year
    return shrunk, float(shrinkage)

correlation_from_covariance 

correlation_from_covariance(cov: Array) -> Array

Convert a covariance matrix to a correlation matrix.

Source code in src/jaxfolio/moments/estimators.py 
123
124
125
126
127
def correlation_from_covariance(cov: Array) -> Array:
    """Convert a covariance matrix to a correlation matrix."""
    d = jnp.sqrt(jnp.clip(jnp.diag(cov), 1e-18, None))
    corr = cov / jnp.outer(d, d)
    return jnp.clip(corr, -1.0, 1.0)

Constraint projections

projections

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).

project_simplex 

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.

Source code in src/jaxfolio/constraints/projections.py 
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
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)

project_box_budget 

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.

Source code in src/jaxfolio/constraints/projections.py 
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
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

normalize_weights 

normalize_weights(w: Array, budget: float = 1.0) -> Array

Rescale weights to sum to budget (assumes a non-zero sum).

Source code in src/jaxfolio/constraints/projections.py 
79
80
81
82
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]))

softmax_weights 

softmax_weights(
    logits: Array, budget: float = 1.0
) -> Array

Map unconstrained logits to long-only weights via softmax (sums to budget).

Source code in src/jaxfolio/constraints/projections.py 
85
86
87
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)