Types & registry

Core types

types

Shared data structures for jaxfolio.

These are lightweight, framework-agnostic containers. Optimizers are configured by an :class:OptimizerConfig (plus a :class:TradingCosts spec for the cost-aware multi-period optimizer) and return a :class:PortfolioResult; the backtester and plotting utilities consume it.

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 | Callable[..., Any] = "spg",
    solver_options: Mapping[str, Any] | None = None,
    learning_rate: float | None = None,
    tol: float = 1e-07,
    l2_reg: float = 0.0,
    constraints: tuple[Any, ...] = (),
    attribution: bool = True,
)

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 | Callable[..., Any]	Which projected-gradient solver to use. Three spellings are accepted: "spg" (default) — 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. the name of any optax optimizer — "adam", "adamw", "sgd", "rmsprop", "lion", ... resolved against optax and optax.contrib. See :func:jaxfolio.solvers.available_solvers. an optax factory callable — optax.adamw, or a module-level function of your own returning a GradientTransformation (the way to use an optax.chain(...) composition). The optax solvers keep smooth fixed-step dynamics, which are preferable when differentiating through the optimizer to train an allocation policy. Pass the factory, not a pre-built optax.adam(1e-2): a pre-built transformation is a new object on every call and would recompile the solver kernel each solve.
solver_options	Mapping[str, Any] | None	Extra keyword arguments forwarded to the optax factory, e.g. {"weight_decay": 1e-3} or {"momentum": 0.9, "nesterov": True}. Only valid for optax solvers — "spg" is tuned via learning_rate and tol. Values must be hashable: they become part of the solver's jit cache key, so each distinct combination costs one compile.
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 any optax optimizer. Set a float to override.
tol	float	Convergence tolerance. For "spg" this is the projected-gradient (KKT stationarity) norm — zero exactly at a KKT point. For the optax solvers it is the weight-update norm, which is only a proxy for optimality: optimizers whose step size does not shrink near the optimum ("sign_sgd", "lion", plain "sgd") may either run to max_iter or stop early when the budget projection cancels a uniform step. Prefer "spg" when you want the exact constrained optimum.
l2_reg	float	Optional L2 penalty on weights (encourages diversification).
attribution	bool	Capture the first-order (KKT) diagnostics needed by :func:jaxfolio.attribution.explain on every solve. Costs one extra gradient and two extra projections, both outside the solver loop. Set to False in tight backtest loops where the explanation is never read.
constraints	tuple[Any, ...]	Named constraint specifications from :mod:jaxfolio.constraints — sector caps, per-asset bound vectors, an explicit budget:: OptimizerConfig(constraints=[ GroupCap("tech", ["AAPL", "MSFT"], max=0.30), Box(lower=0.0, upper=0.10), ]) Unlike long_only / weight_bounds, these carry a name, which is what lets the solver's Lagrange multiplier for each one be reported as an attributable shadow price. Group rows must partition the universe (each asset in at most one group). Empty by default, and an empty set reproduces the previous behaviour exactly — same solver kernel, same numbers. Validation is two-stage, like solver: structural problems (duplicate names, a bad type) are caught here, while feasibility needs the asset universe and is checked by :func:jaxfolio.constraints.compile_constraints at solve time.

with_constraints

with_constraints(constraints: Any) -> OptimizerConfig

Return a copy carrying constraints.

Source code in src/jaxfolio/types.py

def with_constraints(self, constraints: Any) -> OptimizerConfig:
    """Return a copy carrying ``constraints``."""
    from dataclasses import replace

    return replace(self, constraints=tuple(constraints))

solver_spec

solver_spec() -> SolverSpec

Resolve solver / solver_options into a hashable solver key.

Source code in src/jaxfolio/types.py

def solver_spec(self) -> SolverSpec:
    """Resolve ``solver`` / ``solver_options`` into a hashable solver key."""
    from jaxfolio.solvers import resolve_solver  # lazy: avoids an import cycle

    return resolve_solver(self.solver, self.solver_options)

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)

TradingCosts dataclass

TradingCosts(
    spread_bps: float | tuple[float, ...] = 0.0,
    commission_bps: float | tuple[float, ...] = 0.0,
    impact_bps: float | tuple[float, ...] = 0.0,
    turnover_penalty: float = 0.0,
    smoothing: float = 0.1,
)

Per-period trading frictions for the multi-period optimizer.

The three cost coefficients are quoted in basis points of traded notional per unit of |dw|, one-way, per period — the units a trader states them in. Each accepts a scalar (applied to every asset) or a per-asset sequence of length n; sequences are stored as tuples so the dataclass stays hashable.

Costs are per period and never annualized: the objective trades them off against per-period mu and Sigma, so annualizing would double-count.

Attributes:

Name	Type	Description
spread_bps	float | tuple[float, ...]	Bid-ask half-spread, in bps, charged linearly on |dw|. A round trip therefore costs 2 * spread_bps.
commission_bps	float | tuple[float, ...]	Proportional commission / fee, in bps, also linear on |dw|. Kept separate from spread_bps only for reporting clarity — the objective sums the two.
impact_bps	float | tuple[float, ...]	Quadratic market-impact coefficient, in bps per unit of dw**2. Models the convex part of execution cost (temporary impact), which is what makes splitting a large trade across periods strictly cheaper than trading it all at once.
turnover_penalty	float	An extra L1 penalty on |dw| in decimal units (not bps), for soft turnover control that is not a real cash cost. Reported separately from total_cost so P&L attribution stays honest.
smoothing	float	Starting (widest) relative Huber half-width used to make the L1 cost differentiable. The absolute width is smoothing * budget / n, i.e. a fraction of the natural trade unit, which makes the resulting bias independent of n. Must be strictly positive: at 0 the gradient is undefined at dw == 0, which is precisely where the optimum sits when costs bind. The optimizer walks a geometrically shrinking ladder from this width and keeps whichever stage is best under the exact objective, so this is a starting point for a search rather than a value the answer hinges on — which is why it is a tuning knob and not part of the API contract.

Notes

Unit bridge to the backtester: TradingCosts(spread_bps=10.0) prices the same linear cost that backtest(transaction_cost=0.0010) charges. The engine keyword is decimal for historical reasons; this class is bps.

Turnover control here is soft — costs are priced, not capped. See the multi-period guide for why a hard turnover budget is not expressible in this solver, and what to do instead.

Examples:

>>> TradingCosts(spread_bps=5.0, commission_bps=1.0).linear_decimal(3)
array([0.0006, 0.0006, 0.0006])

is_active

is_active() -> bool

True if any cost coefficient is non-zero (reporting/metadata only).

Source code in src/jaxfolio/types.py

def is_active(self) -> bool:
    """True if any cost coefficient is non-zero (reporting/metadata only)."""
    return (
        any(
            np.any(np.asarray(getattr(self, name), dtype=float) != 0.0)
            for name in ("spread_bps", "commission_bps", "impact_bps")
        )
        or self.turnover_penalty != 0.0
    )

linear_decimal

linear_decimal(n: int) -> ndarray

Total linear cost (spread_bps + commission_bps) / 1e4, length n.

Source code in src/jaxfolio/types.py

def linear_decimal(self, n: int) -> np.ndarray:
    """Total linear cost ``(spread_bps + commission_bps) / 1e4``, length ``n``."""
    return (self._broadcast("spread_bps", n) + self._broadcast("commission_bps", n)) / 1e4

quadratic_decimal

quadratic_decimal(n: int) -> ndarray

Quadratic impact coefficient impact_bps / 1e4, length n.

Source code in src/jaxfolio/types.py

def quadratic_decimal(self, n: int) -> np.ndarray:
    """Quadratic impact coefficient ``impact_bps / 1e4``, length ``n``."""
    return self._broadcast("impact_bps", n) / 1e4

as_params

as_params(n: int) -> dict

Traced objective parameters for n assets.

Returns {"c_lin": (n,), "c_quad": (n,)} as JAX arrays. Both are traced rather than static, so changing a cost never recompiles the solver kernel. turnover_penalty is folded into c_lin because it enters the objective identically; the split is preserved for reporting by :meth:linear_decimal.

Source code in src/jaxfolio/types.py

def as_params(self, n: int) -> dict:
    """Traced objective parameters for ``n`` assets.

    Returns ``{"c_lin": (n,), "c_quad": (n,)}`` as JAX arrays. Both are
    *traced* rather than static, so changing a cost never recompiles the
    solver kernel. ``turnover_penalty`` is folded into ``c_lin`` because it
    enters the objective identically; the split is preserved for reporting by
    :meth:`linear_decimal`.
    """
    import jax.numpy as jnp  # lazy: keeps this module importable without jax

    c_lin = self.linear_decimal(n) + self.turnover_penalty
    return {
        "c_lin": jnp.asarray(c_lin, dtype=float),
        "c_quad": jnp.asarray(self.quadratic_decimal(n), dtype=float),
    }

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(),
    trajectory: ndarray | None = None,
    attribution: Any | None = None,
)

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, ...). Values are kept JSON-serializable by convention; bulk arrays belong in a dedicated field like trajectory.
trajectory	ndarray | None	Optional (T, n_assets) planned weight path for multi-period optimizers, where row t is the portfolio to hold in period t and row 0 equals weights. None for every single-period optimizer. :func:jaxfolio.backtest.backtest executes this path step by step when it is present.
attribution	Any | None	Optional :class:jaxfolio.attribution.SolverDuals — the first-order diagnostics captured at the solve, from which :func:jaxfolio.attribution.explain reconstructs which constraints bind and what they cost. None for optimizers that solve no constrained program (equal-weight, risk parity, the graph and learning methods), and whenever OptimizerConfig.attribution is False. explain never requires it: a missing payload yields a well-formed report that says so. Kept out of metadata because it holds bulk arrays; a small JSON-serializable digest of it is placed in metadata.

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

Registry

registry

A lightweight registry for portfolio strategies.

Any callable with the optimizer shape method(returns, ...) -> PortfolioResult can be registered by name and then looked up, listed, and mixed with the built-ins in the backtester. Built-in optimizers register themselves at import time (see :func:_register_builtins), so :func:list_strategies always reflects the full catalog — including the user's own additions.

Usage

from jaxfolio import register_strategy, get_strategy, list_strategies @register_strategy("my_momentum", description="12-1 momentum tilt") ... def my_momentum(returns): ... ... get_strategy("my_momentum") is my_momentum True "my_momentum" in list_strategies() True

StrategyInfo dataclass

StrategyInfo(
    name: str,
    func: Strategy,
    description: str = "",
    builtin: bool = False,
)

A registered strategy: its callable, canonical name, and description.

register_strategy

register_strategy(
    name: str | None = None,
    *,
    description: str = "",
    builtin: bool = False,
    overwrite: bool = False,
) -> Callable[[Strategy], Strategy]

Register a strategy under name (defaults to the function's name).

Usable as a decorator (@register_strategy("x")) or directly (register_strategy("x")(fn)). The wrapped function is returned unchanged, so decorating never alters behavior. Re-registering an existing name raises unless overwrite=True.

Source code in src/jaxfolio/registry.py

def register_strategy(
    name: str | None = None,
    *,
    description: str = "",
    builtin: bool = False,
    overwrite: bool = False,
) -> Callable[[Strategy], Strategy]:
    """Register a strategy under ``name`` (defaults to the function's name).

    Usable as a decorator (``@register_strategy("x")``) or directly
    (``register_strategy("x")(fn)``). The wrapped function is returned unchanged,
    so decorating never alters behavior. Re-registering an existing name raises
    unless ``overwrite=True``.
    """

    def decorator(func: Strategy) -> Strategy:
        key = name or getattr(func, "__name__", None)
        if not key:
            raise ValueError("register_strategy needs a name (no __name__ on the callable)")
        if key in _REGISTRY and not overwrite:
            raise ValueError(
                f"strategy {key!r} is already registered; pass overwrite=True to replace it"
            )
        _REGISTRY[key] = StrategyInfo(name=key, func=func, description=description, builtin=builtin)
        return func

    return decorator

get_strategy

get_strategy(name: str) -> Strategy

Return the registered strategy callable for name.

Source code in src/jaxfolio/registry.py

def get_strategy(name: str) -> Strategy:
    """Return the registered strategy callable for ``name``."""
    try:
        return _REGISTRY[name].func
    except KeyError:
        raise KeyError(f"unknown strategy {name!r}; registered: {sorted(_REGISTRY)}") from None

strategy_info

strategy_info(name: str) -> StrategyInfo

Return the full :class:StrategyInfo record for name.

Source code in src/jaxfolio/registry.py

def strategy_info(name: str) -> StrategyInfo:
    """Return the full :class:`StrategyInfo` record for ``name``."""
    try:
        return _REGISTRY[name]
    except KeyError:
        raise KeyError(f"unknown strategy {name!r}") from None

list_strategies

list_strategies(
    *, builtin_only: bool = False, custom_only: bool = False
) -> list[str]

List registered strategy names (optionally filtered by origin).

Source code in src/jaxfolio/registry.py

def list_strategies(*, builtin_only: bool = False, custom_only: bool = False) -> list[str]:
    """List registered strategy names (optionally filtered by origin)."""
    names = sorted(_REGISTRY)
    if builtin_only:
        return [n for n in names if _REGISTRY[n].builtin]
    if custom_only:
        return [n for n in names if not _REGISTRY[n].builtin]
    return names

unregister

unregister(name: str) -> None

Remove a strategy from the registry (no error if absent).

Source code in src/jaxfolio/registry.py

def unregister(name: str) -> None:
    """Remove a strategy from the registry (no error if absent)."""
    _REGISTRY.pop(name, None)

registry

registry() -> dict[str, StrategyInfo]

Return a shallow copy of the registry mapping (for inspection).

Source code in src/jaxfolio/registry.py

def registry() -> dict[str, StrategyInfo]:
    """Return a shallow copy of the registry mapping (for inspection)."""
    return dict(_REGISTRY)