Data & moments

See the Data guide.

Synthetic data

synthetic

Reproducible synthetic market data as Polars frames.

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 correlated daily GBM prices with an explicit date column.

Source code in src/jaxfolio/data/synthetic.py

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,
) -> pl.DataFrame:
    """Generate correlated daily GBM prices with an explicit ``date`` column."""
    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)
    chol = np.linalg.cholesky(_random_correlation(n_assets, rng, correlation_strength))
    z = rng.standard_normal(size=(n_days, n_assets)) @ chol.T
    log_rets = (drifts - 0.5 * vols**2) * dt + vols * np.sqrt(dt) * z
    prices = np.exp(np.log(start_price) + np.cumsum(log_rets, axis=0))

    start = np.datetime64(date.fromisoformat(start_date))
    dates = np.busday_offset(start, np.arange(n_days), roll="forward").astype("datetime64[D]")
    return pl.DataFrame({"date": dates, **dict(zip(tickers, prices.T, strict=True))})

generate_returns

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

Convenience wrapper returning simple daily returns.

Source code in src/jaxfolio/data/synthetic.py

def generate_returns(n_assets: int = 12, n_days: int = 756, **kwargs) -> pl.DataFrame:
    """Convenience wrapper returning simple daily returns."""
    return to_returns(generate_prices(n_assets=n_assets, n_days=n_days, **kwargs))

Loaders

loaders

Data ingestion into explicit-date Polars frames.

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 wide or long/tidy CSV price panel.

Source code in src/jaxfolio/data/loaders.py

def load_csv(
    path: str | Path,
    *,
    date_col: str = "date",
    asset_col: str | None = None,
    price_col: str = "close",
    **read_csv_kwargs,
) -> pl.DataFrame:
    """Load a wide or long/tidy CSV price panel."""
    df = pl.read_csv(path, **read_csv_kwargs)
    if asset_col is not None:
        return _pivot_to_panel(df, date_col, asset_col, price_col)
    return _parse_date(df, date_col).sort(date_col)

load_parquet

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

Load a wide or long/tidy Parquet price panel.

Source code in src/jaxfolio/data/loaders.py

def load_parquet(
    path: str | Path,
    *,
    date_col: str = "date",
    asset_col: str | None = None,
    price_col: str = "close",
) -> pl.DataFrame:
    """Load a wide or long/tidy Parquet price panel."""
    df = pl.read_parquet(path)
    if asset_col is not None:
        return _pivot_to_panel(df, date_col, asset_col, price_col)
    return _parse_date(df, date_col).sort(date_col) if date_col in df.columns else df

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 wide price panel from Yahoo Finance as Polars.

Source code in src/jaxfolio/data/loaders.py

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",
) -> pl.DataFrame:
    """Download a wide price panel from Yahoo Finance as Polars."""
    try:
        import yfinance as yf
    except ImportError as exc:  # pragma: no cover
        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,
    )
    # yfinance uses hierarchical columns for multiple tickers. Duck-type that
    # shape so the provider's frame library stays an implementation detail.
    panel = raw[price_field] if getattr(raw.columns, "nlevels", 1) > 1 else raw[[price_field]]
    out = _provider_panel_to_polars(panel, tickers)
    assets = [c for c in out.columns if c != "date"]
    missing = pl.col(assets).is_null() | pl.col(assets).is_nan()
    return out.filter(~pl.all_horizontal(missing))

load_option_chain

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

Fetch and normalize a Yahoo Finance option chain as a tidy Polars frame.

Source code in src/jaxfolio/data/loaders.py

def load_option_chain(ticker: str, *, expiry: str | None = None) -> pl.DataFrame:
    """Fetch and normalize a Yahoo Finance option chain as a tidy Polars frame."""
    try:
        import yfinance as yf
    except ImportError as exc:  # pragma: no cover
        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, kind: str) -> pl.DataFrame:
        source = {
            "strike": "strike",
            "lastPrice": "last",
            "bid": "bid",
            "ask": "ask",
            "volume": "volume",
            "openInterest": "open_interest",
            "impliedVolatility": "implied_vol",
        }
        values = {dest: df[src].to_numpy() for src, dest in source.items()}
        return (
            pl.DataFrame(values)
            .with_columns(
                pl.lit(kind).alias("type"), pl.lit(date.fromisoformat(chosen)).alias("expiry")
            )
            .select(
                "type",
                "strike",
                "expiry",
                "last",
                "bid",
                "ask",
                "volume",
                "open_interest",
                "implied_vol",
            )
        )

    return pl.concat([_norm(chain.calls, "call"), _norm(chain.puts, "put")])

Returns & splitting

returns

Return computation, alignment, cleaning, and train/test splitting.

Data panels are Polars frames. A temporal column (conventionally date) is kept as an ordinary, explicit column; every other numeric column is an asset.

date_column

date_column(frame: DataFrame) -> str | None

Return the panel's temporal key, preferring the conventional name.

Source code in src/jaxfolio/data/returns.py

def date_column(frame: pl.DataFrame) -> str | None:
    """Return the panel's temporal key, preferring the conventional name."""
    if "date" in frame.columns:
        return "date"
    return next(
        (name for name, dtype in frame.schema.items() if dtype.is_temporal()),
        None,
    )

asset_columns

asset_columns(frame: DataFrame) -> list[str]

Return numeric asset columns, excluding temporal metadata.

Source code in src/jaxfolio/data/returns.py

def asset_columns(frame: pl.DataFrame) -> list[str]:
    """Return numeric asset columns, excluding temporal metadata."""
    date_col = date_column(frame)
    return [
        name
        for name, dtype in frame.schema.items()
        if name != date_col and dtype.is_numeric() and dtype != pl.Boolean
    ]

to_returns

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

Convert a wide price panel to simple or logarithmic returns.

Source code in src/jaxfolio/data/returns.py

def to_returns(
    prices: pl.DataFrame,
    *,
    kind: str = "simple",
    dropna: bool = True,
) -> pl.DataFrame:
    """Convert a wide price panel to simple or logarithmic returns."""
    assets = asset_columns(prices)
    if kind == "simple":
        exprs = [(pl.col(c) / pl.col(c).shift(1) - 1.0).alias(c) for c in assets]
    elif kind == "log":
        exprs = [(pl.col(c) / pl.col(c).shift(1)).log().alias(c) for c in assets]
    else:
        raise ValueError("kind must be 'simple' or 'log'")
    out = prices.with_columns(exprs)
    if dropna:
        missing = pl.col(assets).is_null() | pl.col(assets).is_nan()
        out = out.filter(~pl.all_horizontal(missing))
    return out

align

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

Align frames on their common (inner) or combined (outer) dates.

Source code in src/jaxfolio/data/returns.py

def align(*frames: pl.DataFrame, how: str = "inner") -> list[pl.DataFrame]:
    """Align frames on their common (``inner``) or combined (``outer``) dates."""
    if not frames:
        return []
    if how not in {"inner", "outer"}:
        raise ValueError("how must be 'inner' or 'outer'")
    keys = [date_column(frame) for frame in frames]
    if any(key is None for key in keys):
        raise ValueError("align requires a date or datetime column in every frame")
    key = keys[0]
    assert key is not None
    normalized = [f.rename({k: key}) if k != key else f for f, k in zip(frames, keys, strict=True)]
    dates = normalized[0].select(key)
    for frame in normalized[1:]:
        join_kind = "full" if how == "outer" else "inner"
        dates = dates.join(frame.select(key), on=key, how=join_kind, coalesce=True)
    dates = dates.unique().sort(key)
    return [dates.join(frame, on=key, how="left") for frame in normalized]

clean_returns

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

Drop sparse asset columns and fill remaining null/NaN values.

Source code in src/jaxfolio/data/returns.py

def clean_returns(
    returns: pl.DataFrame,
    *,
    max_missing_frac: float = 0.1,
    fill: str = "zero",
) -> pl.DataFrame:
    """Drop sparse asset columns and fill remaining null/NaN values."""
    assets = asset_columns(returns)
    n = max(returns.height, 1)
    missing = returns.select(
        [((pl.col(c).is_null() | pl.col(c).is_nan()).sum() / n).alias(c) for c in assets]
    ).row(0, named=True)
    keep = [c for c in assets if missing[c] <= max_missing_frac]
    date_col = date_column(returns)
    out = returns.select(([date_col] if date_col else []) + keep)
    if fill == "zero":
        return out.with_columns(pl.col(keep).fill_null(0.0).fill_nan(0.0))
    if fill == "ffill":
        return out.with_columns(
            pl.col(keep).fill_nan(None).fill_null(strategy="forward").fill_null(0.0)
        )
    raise ValueError("fill must be 'zero' or 'ffill'")

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

Source code in src/jaxfolio/data/returns.py

def train_test_split(
    returns: pl.DataFrame,
    *,
    test_size: float = 0.25,
) -> tuple[pl.DataFrame, pl.DataFrame]:
    """Chronological split into in-sample and out-of-sample frames."""
    if not 0.0 < test_size < 1.0:
        raise ValueError("test_size must be in (0, 1)")
    cut = int(round(returns.height * (1.0 - test_size)))
    return returns.slice(0, cut), returns.slice(cut)

annualization_factor

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

Return the explicit annualization factor used by callers.

Source code in src/jaxfolio/data/returns.py

def annualization_factor(returns: pl.DataFrame | pl.Series, periods_per_year: int = 252) -> int:
    """Return the explicit annualization factor used by callers."""
    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 Polars 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 Polars DataFrame (numeric columns become asset names), or a raw array (names default to asset_0...). Date/datetime columns are metadata and are excluded from the optimizer matrix. Nulls and NaNs become zeros.

Source code in src/jaxfolio/moments/estimators.py

def as_matrix(returns: pl.DataFrame | np.ndarray | Array) -> tuple[Array, list[str]]:
    """Return ``(matrix, asset_names)`` from a returns object.

    Accepts a Polars DataFrame (numeric columns become asset names), or a raw
    array (names default to ``asset_0``...). Date/datetime columns are metadata
    and are excluded from the optimizer matrix. Nulls and NaNs become zeros.
    """
    if isinstance(returns, pl.DataFrame):
        names = [
            name
            for name, dtype in returns.schema.items()
            if dtype.is_numeric() and dtype != pl.Boolean
        ]
        if not names:
            raise ValueError("returns must contain at least one numeric asset column")
        mat = jnp.asarray(returns.select(names).fill_null(0.0).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

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

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

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

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

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

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

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

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

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)