Toolkit & custom strategies

The public building blocks for authoring strategies. See the Custom strategies guide.

Toolkit

toolkit

Public building blocks for authoring custom portfolio strategies.

This module is the stable, documented surface for writing your own optimizer. Everything a built-in optimizer uses internally is re-exported here so a custom strategy can be written the same idiomatic way:

from jaxfolio import toolkit as tk def my_strategy(returns): ... mu, cov, names, _ = tk.moments(returns) ... projection = tk.make_projection(long_only=True, weight_bounds=(0.0, 1.0)) ... def objective(w): ... return w @ cov @ w - 0.1 * (w @ mu) # min-variance with a return tilt ... w, _info = tk.solve_projected_gradient(objective, tk.equal_start(len(names)), projection) ... return tk.finalize_result(w, names, "My Strategy", mu=mu, cov=cov)

The heavy lifting (annualized diagnostics, weight cleanup) lives in :func:finalize_result, which the built-in classical and graph optimizers also use, so custom results are indistinguishable from built-in ones downstream (backtester, compare, plots).

PERIODS_PER_YEAR module-attribute

PERIODS_PER_YEAR = 252

Array module-attribute

Array = jnp.ndarray

__all__ module-attribute

__all__ = [
    "as_matrix",
    "moments",
    "mean_returns",
    "sample_covariance",
    "ewma_covariance",
    "ledoit_wolf_covariance",
    "correlation_from_covariance",
    "make_projection",
    "project_simplex",
    "project_box_budget",
    "normalize_weights",
    "softmax_weights",
    "solve_projected_gradient",
    "solve_constrained",
    "select_projection",
    "equal_start",
    "portfolio_return",
    "portfolio_variance",
    "portfolio_volatility",
    "sharpe_ratio",
    "finalize_result",
    "PERIODS_PER_YEAR",
    "PortfolioResult",
    "OptimizerConfig",
]

OptimizerConfig dataclass

OptimizerConfig(
    risk_free_rate: float = 0.0,
    long_only: bool = True,
    weight_bounds: tuple[float, float] | None = None,
    max_iter: int = 2000,
    solver: str = "spg",
    learning_rate: float | None = None,
    tol: float = 1e-07,
    l2_reg: float = 0.0,
)

Configuration shared by the projected-gradient classical optimizers.

Attributes:

Name	Type	Description
risk_free_rate	float	Per-period risk-free rate used by Sharpe-style objectives.
long_only	bool	If True, weights are projected onto the simplex (no shorting).
weight_bounds	tuple[float, float] | None	Explicit (lower, upper) per-asset bounds applied alongside the budget constraint. When left as None the effective bounds follow long_only: (0, 1) when long-only, (-1, 1) when shorting is allowed — so OptimizerConfig(long_only=False) permits shorting without also having to set bounds. Set this to override those defaults.
max_iter	int	Maximum projected-gradient iterations.
solver	str	Which projected-gradient solver to use. "spg" (default) is a spectral projected-gradient method with Barzilai-Borwein step sizes: it needs no learning-rate tuning and converges to the constrained optimum (matching a dedicated QP solver) in far fewer iterations. "adam" keeps the smooth optax-Adam dynamics, which are preferable when differentiating through the optimizer to train an allocation policy.
learning_rate	float | None	Step size for the solver. None (default) lets the solver pick it automatically — a curvature-based (1/L) initial step for "spg", or 1e-2 for "adam". Set a float to override.
tol	float	Convergence tolerance. For "spg" this is the projected-gradient (KKT stationarity) norm; for "adam" it is the weight-update norm.
l2_reg	float	Optional L2 penalty on weights (encourages diversification).

bounds

bounds() -> tuple[float, float]

Resolve the effective per-asset weight bounds.

Honors an explicit weight_bounds; otherwise defaults to (0, 1) for a long-only portfolio and (-1, 1) when shorting is allowed.

Source code in src/jaxfolio/types.py

def bounds(self) -> tuple[float, float]:
    """Resolve the effective per-asset weight bounds.

    Honors an explicit ``weight_bounds``; otherwise defaults to ``(0, 1)``
    for a long-only portfolio and ``(-1, 1)`` when shorting is allowed.
    """
    if self.weight_bounds is not None:
        return self.weight_bounds
    return (0.0, 1.0) if self.long_only else (-1.0, 1.0)

PortfolioResult dataclass

PortfolioResult(
    weights: ndarray,
    assets: list[str],
    method: str,
    expected_return: float | None = None,
    volatility: float | None = None,
    sharpe: float | None = None,
    metadata: dict[str, Any] = dict(),
)

The output of an optimizer: weights plus diagnostics.

Attributes:

Name	Type	Description
weights	ndarray	Portfolio weights as a 1-D numpy array (sums to 1).
assets	list[str]	Asset names aligned with weights.
method	str	Human-readable optimizer name.
expected_return	float | None	Annualized (or per-period, if unscaled) expected portfolio return.
volatility	float | None	Portfolio volatility on the same scale as expected_return.
sharpe	float | None	Sharpe ratio implied by expected_return / volatility.
metadata	dict[str, Any]	Free-form diagnostics (iterations, converged flag, cluster labels, ...).

as_dict

as_dict() -> dict[str, float]

Return {asset: weight} sorted by descending absolute weight.

Source code in src/jaxfolio/types.py

def as_dict(self) -> dict[str, float]:
    """Return ``{asset: weight}`` sorted by descending absolute weight."""
    pairs = zip(self.assets, self.weights.tolist(), strict=True)
    return dict(sorted(pairs, key=lambda kv: abs(kv[1]), reverse=True))

top

top(n: int = 10) -> dict[str, float]

Return the n largest holdings by absolute weight.

Source code in src/jaxfolio/types.py

def top(self, n: int = 10) -> dict[str, float]:
    """Return the ``n`` largest holdings by absolute weight."""
    return dict(list(self.as_dict().items())[:n])

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

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

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)

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)

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

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

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)

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)

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

make_projection

make_projection(
    long_only: bool,
    weight_bounds: tuple[float, float],
    budget: float = 1.0,
) -> Callable[[Array], Array]

Return the w -> w projection matching the requested constraint set.

Source code in src/jaxfolio/optimizers/base.py

def make_projection(
    long_only: bool,
    weight_bounds: tuple[float, float],
    budget: float = 1.0,
) -> Callable[[Array], Array]:
    """Return the ``w -> w`` projection matching the requested constraint set."""
    lo, hi = weight_bounds
    if long_only and lo <= 0.0 and hi >= 1.0:
        # Plain probability simplex is cheaper and exact.
        return lambda w: project_simplex(w, budget)
    return lambda w: project_box_budget(w, lo, hi, budget)

portfolio_return

portfolio_return(weights: Array, mu: Array) -> Array

Expected portfolio return w . mu.

Source code in src/jaxfolio/optimizers/base.py

@jax.jit
def portfolio_return(weights: Array, mu: Array) -> Array:
    """Expected portfolio return ``w . mu``."""
    return jnp.dot(weights, mu)

portfolio_variance

portfolio_variance(weights: Array, cov: Array) -> Array

Portfolio variance w' Sigma w.

Source code in src/jaxfolio/optimizers/base.py

@jax.jit
def portfolio_variance(weights: Array, cov: Array) -> Array:
    """Portfolio variance ``w' Sigma w``."""
    return weights @ cov @ weights

portfolio_volatility

portfolio_volatility(weights: Array, cov: Array) -> Array

Portfolio volatility sqrt(w' Sigma w).

Source code in src/jaxfolio/optimizers/base.py

@jax.jit
def portfolio_volatility(weights: Array, cov: Array) -> Array:
    """Portfolio volatility ``sqrt(w' Sigma w)``."""
    return jnp.sqrt(jnp.clip(portfolio_variance(weights, cov), 1e-18, None))

select_projection

select_projection(
    long_only: bool,
    weight_bounds: tuple[float, float],
    *,
    aux: bool = False,
)

Return (projection_fn, pparams) for the cached kernel.

aux=True selects the variant that projects only the weight-block of a packed [w, tau] vector (used by the CVaR optimizer). pparams is a tuple of Python floats — passed as a traced argument to the kernel.

Source code in src/jaxfolio/optimizers/base.py

def select_projection(long_only: bool, weight_bounds: tuple[float, float], *, aux: bool = False):
    """Return ``(projection_fn, pparams)`` for the cached kernel.

    ``aux=True`` selects the variant that projects only the weight-block of a
    packed ``[w, tau]`` vector (used by the CVaR optimizer). ``pparams`` is a
    tuple of Python floats — passed as a traced argument to the kernel.
    """
    lo, hi = weight_bounds
    simplex = long_only and lo <= 0.0 and hi >= 1.0
    if simplex:
        return (_proj_simplex_aux if aux else _proj_simplex), (1.0,)
    return (_proj_box_aux if aux else _proj_box), (float(lo), float(hi), 1.0)

sharpe_ratio

sharpe_ratio(
    weights: Array,
    mu: Array,
    cov: Array,
    risk_free: float = 0.0,
) -> Array

Sharpe ratio of a weight vector given moments.

Source code in src/jaxfolio/optimizers/base.py

@jax.jit
def sharpe_ratio(weights: Array, mu: Array, cov: Array, risk_free: float = 0.0) -> Array:
    """Sharpe ratio of a weight vector given moments."""
    excess = portfolio_return(weights, mu) - risk_free
    return excess / portfolio_volatility(weights, cov)

solve_constrained

solve_constrained(
    objective,
    obj_params,
    w0: Array,
    projection,
    proj_params,
    *,
    solver: str = "spg",
    learning_rate: float | None = None,
    max_iter: int = 2000,
    tol: float = 1e-07,
) -> tuple[Array, dict[str, Array]]

Cached entry point for the built-in optimizers.

objective and projection must be module-level functions (see :mod:jaxfolio.optimizers.classical) so the underlying jit cache hits. Returns (weights, info) with info["iterations"].

Source code in src/jaxfolio/optimizers/base.py

def solve_constrained(
    objective,
    obj_params,
    w0: Array,
    projection,
    proj_params,
    *,
    solver: str = "spg",
    learning_rate: float | None = None,
    max_iter: int = 2000,
    tol: float = 1e-7,
) -> tuple[Array, dict[str, Array]]:
    """Cached entry point for the built-in optimizers.

    ``objective`` and ``projection`` must be *module-level* functions (see
    :mod:`jaxfolio.optimizers.classical`) so the underlying jit cache hits.
    Returns ``(weights, info)`` with ``info["iterations"]``.
    """
    w, iters = _solve_cached(
        objective,
        projection,
        obj_params,
        proj_params,
        w0,
        jnp.asarray(tol),
        solver=solver,
        learning_rate=learning_rate,
        max_iter=max_iter,
    )
    return w, {"iterations": iters}

solve_projected_gradient

solve_projected_gradient(
    objective: Objective,
    w0: Array,
    projection: Callable[[Array], Array],
    *,
    learning_rate: float | None = None,
    max_iter: int = 2000,
    tol: float = 1e-07,
    solver: str = "spg",
) -> tuple[Array, dict[str, Array]]

Minimize objective over the feasible set defined by projection.

Accepts arbitrary Python closures objective(w) -> scalar and projection(w) -> w (the shape used by custom strategies and the toolkit). Because the closure identity changes per call this re-traces each time — for hot loops over a built-in optimizer, prefer the cached :func:solve_constrained. Returns (weights, info).

Source code in src/jaxfolio/optimizers/base.py

def solve_projected_gradient(
    objective: Objective,
    w0: Array,
    projection: Callable[[Array], Array],
    *,
    learning_rate: float | None = None,
    max_iter: int = 2000,
    tol: float = 1e-7,
    solver: str = "spg",
) -> tuple[Array, dict[str, Array]]:
    """Minimize ``objective`` over the feasible set defined by ``projection``.

    Accepts arbitrary Python closures ``objective(w) -> scalar`` and
    ``projection(w) -> w`` (the shape used by custom strategies and the
    ``toolkit``). Because the closure identity changes per call this re-traces
    each time — for hot loops over a built-in optimizer, prefer the cached
    :func:`solve_constrained`. Returns ``(weights, info)``.
    """
    w, iters = _run_solver(
        objective,
        projection,
        w0,
        solver=solver,
        learning_rate=learning_rate,
        max_iter=max_iter,
        tol=tol,
    )
    return w, {"iterations": iters}

equal_start

equal_start(n: int) -> Array

A uniform 1/n starting weight vector for the solver.

Source code in src/jaxfolio/results.py

def equal_start(n: int) -> Array:
    """A uniform ``1/n`` starting weight vector for the solver."""
    return jnp.full(n, 1.0 / n)

finalize_result

finalize_result(
    weights,
    assets: list[str],
    method: str,
    *,
    mu: Array | None = None,
    cov: Array | None = None,
    returns=None,
    risk_free: float = 0.0,
    periods_per_year: int = PERIODS_PER_YEAR,
    metadata: dict | None = None,
    clean_eps: float = 1e-10,
) -> PortfolioResult

Assemble a :class:PortfolioResult with annualized diagnostics.

Provide the moments either directly (mu and cov) or implicitly via returns (sample moments are computed from it). Weights are coerced to a flat float array and values below clean_eps in magnitude are zeroed. This is the single source of truth for the annualized return / volatility / Sharpe reported on every optimizer result — built-in and custom alike.

Source code in src/jaxfolio/results.py

def finalize_result(
    weights,
    assets: list[str],
    method: str,
    *,
    mu: Array | None = None,
    cov: Array | None = None,
    returns=None,
    risk_free: float = 0.0,
    periods_per_year: int = PERIODS_PER_YEAR,
    metadata: dict | None = None,
    clean_eps: float = 1e-10,
) -> PortfolioResult:
    """Assemble a :class:`PortfolioResult` with annualized diagnostics.

    Provide the moments either directly (``mu`` and ``cov``) or implicitly via
    ``returns`` (sample moments are computed from it). Weights are coerced to a
    flat float array and values below ``clean_eps`` in magnitude are zeroed. This
    is the single source of truth for the annualized return / volatility / Sharpe
    reported on every optimizer result — built-in and custom alike.
    """
    w = np.asarray(weights, dtype=float).reshape(-1)
    w = np.where(np.abs(w) < clean_eps, 0.0, w)

    if mu is None or cov is None:
        if returns is None:
            raise ValueError("finalize_result requires either (mu, cov) or returns")
        mat, _ = as_matrix(returns)
        mu = mean_returns(mat) if mu is None else mu
        cov = sample_covariance(mat) if cov is None else cov

    wj = jnp.asarray(w)
    ann_mu = float(jnp.dot(wj, mu) * periods_per_year)
    port_var = float(wj @ cov @ wj)
    ann_vol = float(np.sqrt(max(port_var, 0.0)) * np.sqrt(periods_per_year))
    excess = ann_mu - risk_free * periods_per_year
    sharpe = excess / ann_vol if ann_vol > 0 else None

    return PortfolioResult(
        weights=w,
        assets=list(assets),
        method=method,
        expected_return=ann_mu,
        volatility=ann_vol,
        sharpe=sharpe,
        metadata=metadata or {},
    )

moments

moments(
    returns, cov_estimator=None
) -> tuple[Array, Array, list[str], Array]

Return (mu, cov, asset_names, return_matrix) from a returns object.

cov_estimator optionally overrides the default sample covariance.

Source code in src/jaxfolio/results.py

def moments(returns, cov_estimator=None) -> tuple[Array, Array, list[str], Array]:
    """Return ``(mu, cov, asset_names, return_matrix)`` from a returns object.

    ``cov_estimator`` optionally overrides the default sample covariance.
    """
    mat, names = as_matrix(returns)
    mu = mean_returns(mat)
    cov = sample_covariance(mat) if cov_estimator is None else cov_estimator(mat)
    return mu, cov, names, mat

Custom strategy authoring

custom

Author your own portfolio strategy with minimal boilerplate.

Two entry points, both producing a proper :class:PortfolioResult that works everywhere a built-in optimizer does (backtester, compare, plots):

• :func:custom_strategy / :meth:CustomStrategy.from_weights — wrap a function that maps returns to a weight vector (array) or a {asset: weight} mapping. Weights are validated, optionally renormalized, and annualized diagnostics are attached automatically.
• :meth:CustomStrategy.from_objective — provide a JAX objective f(w, ctx) and let the shared projected-gradient solver optimize it under the configured constraints (this is exactly how the built-in classical optimizers are built).

Both optionally register the resulting strategy by name via :mod:jaxfolio.registry.

CustomStrategy

CustomStrategy(
    name: str, fn: Callable[[object], PortfolioResult]
)

A user-defined strategy exposed with the standard optimizer interface.

Instances are callable as strategy(returns) -> PortfolioResult, so they drop straight into :func:jaxfolio.backtest.compare and the plotting layer.

Source code in src/jaxfolio/custom.py

def __init__(self, name: str, fn: Callable[[object], PortfolioResult]):
    self.name = name
    self._fn = fn

from_weights classmethod

from_weights(
    name: str,
    weight_fn: Callable[[object], object],
    *,
    renormalize: bool = True,
    risk_free: float = 0.0,
    register: bool = False,
    description: str = "",
) -> CustomStrategy

Wrap a returns -> weights function into a full strategy.

weight_fn may return a numpy/JAX array aligned to the asset columns or a {asset: weight} mapping. If renormalize is True (default) and the weights sum to a non-zero value, they are scaled to sum to 1.

Source code in src/jaxfolio/custom.py

@classmethod
def from_weights(
    cls,
    name: str,
    weight_fn: Callable[[object], object],
    *,
    renormalize: bool = True,
    risk_free: float = 0.0,
    register: bool = False,
    description: str = "",
) -> CustomStrategy:
    """Wrap a ``returns -> weights`` function into a full strategy.

    ``weight_fn`` may return a numpy/JAX array aligned to the asset columns or
    a ``{asset: weight}`` mapping. If ``renormalize`` is ``True`` (default)
    and the weights sum to a non-zero value, they are scaled to sum to 1.
    """

    def run(returns) -> PortfolioResult:
        mu, cov, names, _mat = tk.moments(returns)
        w = _coerce_weights(weight_fn(returns), names)
        # Only rescale to a fully-invested book when the net exposure is
        # positive. Dividing by a negative sum would flip every sign and
        # invert a net-short / dollar-neutral stance, so those are left as-is.
        if renormalize and w.sum() > 1e-12:
            w = w / w.sum()
        return tk.finalize_result(
            w,
            names,
            name,
            mu=mu,
            cov=cov,
            risk_free=risk_free,
            metadata={"custom": True, "kind": "weights"},
        )

    strat = cls(name, run)
    if register:
        register_strategy(name, description=description, overwrite=True)(strat)
    return strat

from_objective classmethod

from_objective(
    name: str,
    objective_fn: Callable[[Array, _Context], Array],
    *,
    config: OptimizerConfig | None = None,
    register: bool = False,
    description: str = "",
) -> CustomStrategy

Build a strategy by minimizing a JAX objective under constraints.

objective_fn(w, ctx) must be JAX-differentiable and return a scalar to minimize; ctx is a :class:_Context exposing mu, cov, returns, assets, n. The shared projected-gradient solver and the configured constraint set (long-only / bounds) are reused verbatim.

Source code in src/jaxfolio/custom.py

@classmethod
def from_objective(
    cls,
    name: str,
    objective_fn: Callable[[Array, _Context], Array],
    *,
    config: OptimizerConfig | None = None,
    register: bool = False,
    description: str = "",
) -> CustomStrategy:
    """Build a strategy by minimizing a JAX objective under constraints.

    ``objective_fn(w, ctx)`` must be JAX-differentiable and return a scalar to
    *minimize*; ``ctx`` is a :class:`_Context` exposing ``mu``, ``cov``,
    ``returns``, ``assets``, ``n``. The shared projected-gradient solver and
    the configured constraint set (long-only / bounds) are reused verbatim.
    """
    cfg = config or OptimizerConfig()

    def run(returns) -> PortfolioResult:
        mu, cov, names, mat = tk.moments(returns)
        ctx = _Context(mu=mu, cov=cov, returns=mat, assets=names)
        projection = tk.make_projection(cfg.long_only, cfg.bounds())
        w, info = tk.solve_projected_gradient(
            lambda w: objective_fn(w, ctx),
            tk.equal_start(len(names)),
            projection,
            learning_rate=cfg.learning_rate,
            max_iter=cfg.max_iter,
            tol=cfg.tol,
        )
        return tk.finalize_result(
            w,
            names,
            name,
            mu=mu,
            cov=cov,
            risk_free=cfg.risk_free_rate,
            metadata={
                "custom": True,
                "kind": "objective",
                "iterations": int(info["iterations"]),
            },
        )

    strat = cls(name, run)
    if register:
        register_strategy(name, description=description, overwrite=True)(strat)
    return strat

custom_strategy

custom_strategy(
    name: str,
    weight_fn: Callable[[object], object],
    *,
    renormalize: bool = True,
    risk_free: float = 0.0,
    register: bool = False,
    description: str = "",
) -> CustomStrategy

Shorthand for :meth:CustomStrategy.from_weights (the common case).

Source code in src/jaxfolio/custom.py

def custom_strategy(
    name: str,
    weight_fn: Callable[[object], object],
    *,
    renormalize: bool = True,
    risk_free: float = 0.0,
    register: bool = False,
    description: str = "",
) -> CustomStrategy:
    """Shorthand for :meth:`CustomStrategy.from_weights` (the common case)."""
    return CustomStrategy.from_weights(
        name,
        weight_fn,
        renormalize=renormalize,
        risk_free=risk_free,
        register=register,
        description=description,
    )

Result assembly

results

Low-level result assembly and moment extraction.

Kept dependency-light on purpose: this module imports only the moment estimators (never the optimizer package), so both the optimizers and the higher-level :mod:jaxfolio.toolkit can import it without creating an import cycle.

moments

moments(
    returns, cov_estimator=None
) -> tuple[Array, Array, list[str], Array]

Return (mu, cov, asset_names, return_matrix) from a returns object.

cov_estimator optionally overrides the default sample covariance.

Source code in src/jaxfolio/results.py

def moments(returns, cov_estimator=None) -> tuple[Array, Array, list[str], Array]:
    """Return ``(mu, cov, asset_names, return_matrix)`` from a returns object.

    ``cov_estimator`` optionally overrides the default sample covariance.
    """
    mat, names = as_matrix(returns)
    mu = mean_returns(mat)
    cov = sample_covariance(mat) if cov_estimator is None else cov_estimator(mat)
    return mu, cov, names, mat

equal_start

equal_start(n: int) -> Array

A uniform 1/n starting weight vector for the solver.

Source code in src/jaxfolio/results.py

def equal_start(n: int) -> Array:
    """A uniform ``1/n`` starting weight vector for the solver."""
    return jnp.full(n, 1.0 / n)

finalize_result

finalize_result(
    weights,
    assets: list[str],
    method: str,
    *,
    mu: Array | None = None,
    cov: Array | None = None,
    returns=None,
    risk_free: float = 0.0,
    periods_per_year: int = PERIODS_PER_YEAR,
    metadata: dict | None = None,
    clean_eps: float = 1e-10,
) -> PortfolioResult

Assemble a :class:PortfolioResult with annualized diagnostics.

Provide the moments either directly (mu and cov) or implicitly via returns (sample moments are computed from it). Weights are coerced to a flat float array and values below clean_eps in magnitude are zeroed. This is the single source of truth for the annualized return / volatility / Sharpe reported on every optimizer result — built-in and custom alike.

Source code in src/jaxfolio/results.py

def finalize_result(
    weights,
    assets: list[str],
    method: str,
    *,
    mu: Array | None = None,
    cov: Array | None = None,
    returns=None,
    risk_free: float = 0.0,
    periods_per_year: int = PERIODS_PER_YEAR,
    metadata: dict | None = None,
    clean_eps: float = 1e-10,
) -> PortfolioResult:
    """Assemble a :class:`PortfolioResult` with annualized diagnostics.

    Provide the moments either directly (``mu`` and ``cov``) or implicitly via
    ``returns`` (sample moments are computed from it). Weights are coerced to a
    flat float array and values below ``clean_eps`` in magnitude are zeroed. This
    is the single source of truth for the annualized return / volatility / Sharpe
    reported on every optimizer result — built-in and custom alike.
    """
    w = np.asarray(weights, dtype=float).reshape(-1)
    w = np.where(np.abs(w) < clean_eps, 0.0, w)

    if mu is None or cov is None:
        if returns is None:
            raise ValueError("finalize_result requires either (mu, cov) or returns")
        mat, _ = as_matrix(returns)
        mu = mean_returns(mat) if mu is None else mu
        cov = sample_covariance(mat) if cov is None else cov

    wj = jnp.asarray(w)
    ann_mu = float(jnp.dot(wj, mu) * periods_per_year)
    port_var = float(wj @ cov @ wj)
    ann_vol = float(np.sqrt(max(port_var, 0.0)) * np.sqrt(periods_per_year))
    excess = ann_mu - risk_free * periods_per_year
    sharpe = excess / ann_vol if ann_vol > 0 else None

    return PortfolioResult(
        weights=w,
        assets=list(assets),
        method=method,
        expected_return=ann_mu,
        volatility=ann_vol,
        sharpe=sharpe,
        metadata=metadata or {},
    )