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",
    "project_box_budget_vec",
    "project_grouped",
    "Box",
    "Budget",
    "Constraint",
    "GroupCap",
    "GroupFloor",
    "CompiledConstraints",
    "InfeasibleConstraints",
    "compile_constraints",
    "check_feasible",
    "explain",
    "solver_duals",
    "ConstraintReport",
    "SolverDuals",
    "ProjectionDuals",
    "project_box_budget_duals",
    "project_grouped_duals",
    "normalize_weights",
    "softmax_weights",
    "solve_projected_gradient",
    "solve_constrained",
    "select_projection",
    "equal_start",
    "solve_weight_path",
    "available_solvers",
    "resolve_solver",
    "SolverSpec",
    "portfolio_return",
    "portfolio_variance",
    "portfolio_volatility",
    "sharpe_ratio",
    "finalize_result",
    "PERIODS_PER_YEAR",
    "PortfolioResult",
    "OptimizerConfig",
    "TradingCosts",
]

ConstraintReport dataclass

ConstraintReport(
    method: str,
    objective: str,
    sense: str,
    assets: tuple[str, ...],
    weights: ndarray,
    available: bool,
    quality: str,
    reason: str | None,
    stationarity: float | None,
    kkt_residual: float | None,
    step_check: float | None,
    budget_multiplier: float | None,
    budget_shadow_price: float | None,
    budget_note: str | None,
    constraints: tuple[ConstraintDual, ...],
    assets_detail: tuple[AssetAttribution, ...],
    warnings: tuple[str, ...] = (),
    aux_stationarity: float | None = None,
    l2_reg: float = 0.0,
    _index: dict[str, int] = dict(),
)

The result of :func:explain — per-constraint and per-asset attribution.

quality grades everything else in the report:

exact Converged to machine precision; multipliers are trustworthy. approximate Converged loosely, or the step probe disagreed. Numbers are shown with a warning attached. unreliable Not converged, or the objective is non-smooth at the optimum. Shadow prices are suppressed entirely — only primal facts are reported. unavailable This optimizer has no multipliers at all; reason says why.

Four renderings, for four audiences: :meth:explain_text reads the report as prose, :meth:to_table lays the same facts out as fixed-width tables for a terminal or a log, :meth:to_dict / :meth:to_json serialize it whole for a machine, and :meth:to_frame / :meth:constraints_frame hand it to Polars. All five are pure views over the same fields — none of them recompute anything.

binding

binding() -> tuple[ConstraintDual, ...]

Constraints that are active and moving the optimum.

Source code in src/jaxfolio/attribution.py

def binding(self) -> tuple[ConstraintDual, ...]:
    """Constraints that are active *and* moving the optimum."""
    return tuple(c for c in self.constraints if c.status == "binding")

for_asset

for_asset(asset: str) -> AssetAttribution

Attribution for one asset by name.

Source code in src/jaxfolio/attribution.py

def for_asset(self, asset: str) -> AssetAttribution:
    """Attribution for one asset by name."""
    try:
        return self.assets_detail[self._index[asset]]
    except KeyError:
        raise KeyError(f"{asset!r} is not in this portfolio") from None

to_frame

to_frame() -> DataFrame

Per-asset attribution as a Polars frame with an asset column.

Source code in src/jaxfolio/attribution.py

def to_frame(self) -> pl.DataFrame:
    """Per-asset attribution as a Polars frame with an ``asset`` column."""
    import polars as pl

    rows = []
    for a in self.assets_detail:
        rows.append(
            {
                "asset": a.asset,
                "weight": a.weight,
                "status": a.status,
                "reduced_cost": np.nan if a.reduced_cost is None else a.reduced_cost,
                "intrinsic_cost": np.nan if a.intrinsic_cost is None else a.intrinsic_cost,
                "imposed_total": float(sum(v for _, v in a.imposed)),
                "binding": a.primary,
                "binding_kind": a.primary_kind,
                "shadow_price": (
                    np.nan if a.primary_shadow_price is None else a.primary_shadow_price
                ),
                "confidence": a.confidence,
            }
        )
    return pl.DataFrame(rows)

constraints_frame

constraints_frame() -> DataFrame

Per-constraint activity and shadow price, keyed by name.

A different question from :meth:to_frame: which constraint is costing the portfolio the most, rather than what put each asset where it is.

Source code in src/jaxfolio/attribution.py

def constraints_frame(self) -> pl.DataFrame:
    """Per-constraint activity and shadow price, keyed by ``name``.

    A different question from :meth:`to_frame`: which *constraint* is costing
    the portfolio the most, rather than what put each asset where it is.
    """
    import polars as pl

    rows = [
        {
            "name": c.name,
            "kind": c.kind,
            "activity": c.activity,
            "lower": c.lower,
            "upper": c.upper,
            "slack": c.slack,
            "multiplier": np.nan if c.multiplier is None else c.multiplier,
            "shadow_price": np.nan if c.shadow_price is None else c.shadow_price,
            "status": c.status,
            "n_members": len(c.members),
        }
        for c in self.constraints
    ]
    return pl.DataFrame(rows)

explain_text

explain_text(*, top_n: int = 10) -> str

A readable rendering of the report.

Source code in src/jaxfolio/attribution.py

def explain_text(self, *, top_n: int = 10) -> str:
    """A readable rendering of the report."""
    out = [f"Constraint attribution — {self.method} ({len(self.assets)} assets)"]
    if not self.available:
        out.append(f"  unavailable: {self.reason}")
        return "\n".join(out)
    line = f"  quality: {self.quality}"
    if self.stationarity is not None:
        line += f"  ·  stationarity {self.stationarity:.2e}"
    out.append(line)
    if self.budget_shadow_price is not None:
        out.append(f"  cost of capital (budget): {self.budget_shadow_price:+.5f} per unit")
    elif self.budget_note:
        out.append(f"  cost of capital (budget): n/a — {self.budget_note}")
    for w in self.warnings:
        out.append(f"  ! {w}")

    ranked = sorted(self.assets_detail, key=lambda a: (a.status == "interior", -abs(a.weight)))
    for a in ranked[:top_n]:
        out.append("")
        tag = "" if a.status == "interior" else f"   [{a.status.replace('_', ' ')}]"
        out.append(f"Asset {a.asset}: {a.weight:.2%}{tag}")
        if a.reason:
            out.append(f"  {a.reason}")
            continue
        if a.primary is None:
            out.append("  cause: ambiguous — several constraints are tied")
        elif a.primary == "own":
            out.append(f"  cause: own merits ({a.primary_kind})")
        else:
            sp = a.primary_shadow_price
            shown = "n/a" if sp is None else f"{sp:+.5f}"
            out.append(f"  binding: {a.primary} ({a.primary_kind})   shadow price {shown}")
        if a.intrinsic_cost is not None and a.imposed:
            charged = "  ·  ".join(f"charged by {n} {v:+.5f}" for n, v in a.imposed)
            out.append(
                f"  own merit {a.intrinsic_cost:+.5f}  ·  {charged}"
                f"  ·  net {a.reduced_cost:+.5f}"
            )
    return "\n".join(out)

to_table

to_table(
    *, max_assets: int | None = None, notes: bool = True
) -> str

The report as fixed-width tables — for a terminal, a log, or a print-out.

Same facts as :meth:explain_text, laid out as columns instead of prose: a header block, one row per named constraint, one row per asset, and the notes that qualify them. Plain text, and every string jaxfolio itself wrote is folded to ASCII, so the report survives a log file, a CI artifact and a non-UTF-8 console. Asset and constraint names are yours and pass through verbatim — a report that mangles a ticker to fit a charset would be the worse trade.

Parameters:

Name	Type	Description	Default
max_assets	int | None	Show only the first max_assets asset rows, ordered so that the assets sitting on a bound — the ones a constraint may be responsible for — come first. None (the default) prints every asset: a report that silently drops rows is worse than a long one.	None
notes	bool	Include the trailing notes block (units caveat, quality reason, warnings).	True

Examples:

>>> print(jf.explain(result).to_table())

Source code in src/jaxfolio/attribution.py

def to_table(self, *, max_assets: int | None = None, notes: bool = True) -> str:
    """The report as fixed-width tables — for a terminal, a log, or a print-out.

    Same facts as :meth:`explain_text`, laid out as columns instead of prose: a
    header block, one row per named constraint, one row per asset, and the notes
    that qualify them. Plain text, and every string jaxfolio itself wrote is folded
    to ASCII, so the report survives a log file, a CI artifact and a non-UTF-8
    console. Asset and constraint *names* are yours and pass through verbatim — a
    report that mangles a ticker to fit a charset would be the worse trade.

    Parameters
    ----------
    max_assets:
        Show only the first ``max_assets`` asset rows, ordered so that the assets
        sitting on a bound — the ones a constraint may be responsible for — come
        first. ``None`` (the default) prints every asset: a report that silently
        drops rows is worse than a long one.
    notes:
        Include the trailing notes block (units caveat, quality reason, warnings).

    Examples
    --------
    >>> print(jf.explain(result).to_table())  # doctest: +SKIP
    """
    rule = "=" * 78
    out = [rule, f"CONSTRAINT ATTRIBUTION -- {_ascii(self.method)}", rule]

    fields = [("assets", str(len(self.assets)))]
    if self.objective:
        fields.append(("objective", f"{_ascii(self.objective)} ({self.sense})"))
    fields.append(("dual quality", self.quality))
    if self.stationarity is not None:
        fields.append(("stationarity", f"{self.stationarity:.2e}"))
    if self.kkt_residual is not None:
        fields.append(("KKT residual", f"{self.kkt_residual:.2e}"))
    if self.aux_stationarity is not None:
        fields.append(("aux stationarity", f"{self.aux_stationarity:.2e}"))
    extra_notes = []
    if self.budget_shadow_price is not None:
        fields.append(("cost of capital", f"{self.budget_shadow_price:+.3e} per unit"))
    elif self.budget_note:
        # The note is a sentence, not a value — it goes to NOTES where it can wrap,
        # rather than pushing the header block out to 150 columns.
        fields.append(("cost of capital", f"{_MISSING} (see notes)"))
        extra_notes.append(f"cost of capital: {self.budget_note}")
    if self.available:
        binding = [c.name for c in self.binding()]
        summary = f"{len(binding)} of {len(self.constraints)}"
        fields.append(
            ("binding rows", f"{summary}: {', '.join(binding)}" if binding else summary)
        )
    out += _render_fields(fields)

    if not self.available:
        out += ["", *_wrap(f"unavailable: {self.reason}")]

    if self.constraints:
        out += ["", "CONSTRAINTS"]
        out += _render_table(
            [
                "name",
                "kind",
                "activity",
                "bound",
                "slack",
                "multiplier",
                "shadow price",
                "status",
            ],
            [
                [
                    c.name,
                    c.kind,
                    _cell_pct(c.activity),
                    _cell_bound(c.lower, c.upper),
                    _cell_pct(c.slack),
                    _cell_num(c.multiplier),
                    _cell_num(c.shadow_price),
                    c.status,
                ]
                for c in self.constraints
            ],
            "<<>><>><",
        )

    ranked = sorted(self.assets_detail, key=lambda a: (a.status == "interior", -abs(a.weight)))
    shown = ranked if max_assets is None else ranked[:max_assets]
    if shown:
        out += ["", "ASSETS"]
        out += _render_table(
            ["asset", "weight", "status", "own merit", "imposed", "net", "cause", "confidence"],
            [
                [
                    a.asset,
                    _cell_pct(a.weight),
                    a.status.replace("_", " "),
                    _cell_num(a.intrinsic_cost),
                    _cell_num(sum(v for _name, v in a.imposed)) if a.imposed else _MISSING,
                    _cell_num(a.reduced_cost),
                    _cause_cell(a),
                    a.confidence,
                ]
                for a in shown
            ],
            "<><>>><<",
        )
        if len(shown) < len(ranked):
            out.append(f"({len(ranked) - len(shown)} further asset(s) not shown)")

    if notes:
        lines = ["multipliers are per period, in the objective's own units (not annualized)"]
        lines += extra_notes
        if self.reason and self.available:  # an unavailable reason is already printed above
            lines.append(self.reason)
        # No l2_reg note of our own: `warnings` already carries that one, in more
        # detail. Two phrasings of one caveat read as two caveats.
        lines += list(self.warnings)
        out += ["", "NOTES"]
        for line in lines:
            out += _wrap(line, bullet="- ")

    return "\n".join(out)

to_dict

to_dict() -> dict[str, Any]

The whole report as a JSON-serializable dict.

Nothing is dropped and nothing is rounded: every constraint, every asset and every diagnostic, with plain Python scalars throughout (no numpy types). A number the report declines to stand behind stays None rather than becoming 0.0 — the distinction is the point of the report, so it must survive serialization. schema is there so a consumer can tell versions apart if this grows.

The digest on result.metadata is the small, always-present cousin of this: five keys for a backtest log, where this is the whole picture.

Source code in src/jaxfolio/attribution.py

def to_dict(self) -> dict[str, Any]:
    """The whole report as a JSON-serializable dict.

    Nothing is dropped and nothing is rounded: every constraint, every asset and
    every diagnostic, with plain Python scalars throughout (no numpy types). A
    number the report declines to stand behind stays ``None`` rather than
    becoming ``0.0`` — the distinction is the point of the report, so it must
    survive serialization. ``schema`` is there so a consumer can tell versions
    apart if this grows.

    The digest on ``result.metadata`` is the small, always-present cousin of
    this: five keys for a backtest log, where this is the whole picture.
    """
    return {
        "schema": "jaxfolio.constraint_report/1",
        "method": self.method,
        "objective": self.objective or None,
        "sense": self.sense,
        "available": self.available,
        "quality": self.quality,
        "reason": self.reason,
        "diagnostics": {
            "kkt_residual": _json_float(self.kkt_residual),
            "stationarity": _json_float(self.stationarity),
            "step_check": _json_float(self.step_check),
            "aux_stationarity": _json_float(self.aux_stationarity),
            "l2_reg": _json_float(self.l2_reg),
        },
        "budget": {
            "multiplier": _json_float(self.budget_multiplier),
            "shadow_price": _json_float(self.budget_shadow_price),
            "note": self.budget_note,
        },
        "constraints": [
            {
                "name": c.name,
                "kind": c.kind,
                "activity": _json_float(c.activity),
                "lower": _json_float(c.lower),
                "upper": _json_float(c.upper),
                "slack": _json_float(c.slack),
                "multiplier": _json_float(c.multiplier),
                "shadow_price": _json_float(c.shadow_price),
                "status": c.status,
                "members": list(c.members),
            }
            for c in self.constraints
        ],
        "assets": [
            {
                "asset": a.asset,
                "weight": _json_float(a.weight),
                "status": a.status,
                "reduced_cost": _json_float(a.reduced_cost),
                "intrinsic_cost": _json_float(a.intrinsic_cost),
                "imposed": [
                    {"constraint": name, "charge": _json_float(value)}
                    for name, value in a.imposed
                ],
                "cause": a.primary,
                "cause_kind": a.primary_kind,
                "cause_shadow_price": _json_float(a.primary_shadow_price),
                "confidence": a.confidence,
                "reason": a.reason,
            }
            for a in self.assets_detail
        ],
        "warnings": list(self.warnings),
    }

to_json

to_json(*, indent: int | None = 2) -> str

The report as a JSON string. indent=None for a single compact line.

Source code in src/jaxfolio/attribution.py

def to_json(self, *, indent: int | None = 2) -> str:
    """The report as a JSON string. ``indent=None`` for a single compact line."""
    return json.dumps(self.to_dict(), indent=indent, allow_nan=False)

SolverDuals dataclass

SolverDuals(
    weights: ndarray,
    gradient: ndarray,
    lower: ndarray,
    upper: ndarray,
    budget: float,
    row_names: tuple[str, ...],
    row_lower: tuple[float, ...],
    row_upper: tuple[float, ...],
    group_of: tuple[int, ...],
    budget_multiplier: float,
    row_multipliers: tuple[float, ...],
    reduced_costs: ndarray,
    row_identified: tuple[bool, ...],
    budget_identified: bool,
    kkt_residual: float,
    gradient_norm: float,
    step: float,
    step_check: float,
    aux_stationarity: float | None = None,
    objective: str = "",
    sense: str = "minimize",
    smooth: bool = True,
    convex: bool = True,
    l2_reg: float = 0.0,
    solver_kind: str = "spg",
    tol: float = 1e-07,
)

First-order information captured at a solved portfolio.

Attached to :attr:jaxfolio.types.PortfolioResult.attribution. Deliberately numeric and self-contained — plain numpy arrays and tuples, no JAX arrays, no closures, no reference to the returns panel — so a result stays lightweight and picklable, and so :func:explain is pure post-processing that never re-solves.

Attributes:

Name	Type	Description
weights, gradient		The solution and the objective gradient there.
lower, upper		Effective per-asset bounds in force.
budget, row_names, row_lower, row_upper, group_of		The constraint set, as resolved by :func:jaxfolio.constraints.compile_constraints.
budget_multiplier, row_multipliers, reduced_costs		The recovered multipliers, in objective units per period. Non-negative for a binding cap; see the module docstring for the sign convention.
row_identified, budget_identified		Whether each multiplier is pinned by the data at all. See :class:jaxfolio.constraints.structured.ProjectionDuals.
kkt_residual, gradient_norm, step, step_check		Convergence evidence. kkt_residual is ||d|| where d is the feasible-descent part of the Moreau split; it is a genuine KKT measure for every solver family, including the optax ones whose own residual is only a weight-update norm. step_check re-probes at a tenth of step and reports the disagreement relative to ||grad f||; a large value means the probe left the region where the projection's active set is constant, so the multipliers were read off the wrong face.
aux_stationarity	float | None	For packed [w, tau] problems (CVaR), |d f / d tau| at the solution — the auxiliary variable is unconstrained, so its whole gradient is "descent the constraints could not absorb".
tol	float	The convergence tolerance the solve was asked for. The quality grade is judged against it: "did this solve do what it was told to?" is a fairer and more portable question than any fixed absolute threshold.
objective, sense, smooth, convex		Provenance the report needs to know what it may claim. sense is "minimize" or "maximize" for the objective as the user names it, which sets the reported sign of every shadow price.

stationarity property

stationarity: float

||d|| / ||grad f|| — the scale-free measure the quality gate uses.

summary

summary() -> dict[str, Any]

A small JSON-serializable digest for PortfolioResult.metadata.

Source code in src/jaxfolio/attribution.py

def summary(self) -> dict[str, Any]:
    """A small JSON-serializable digest for ``PortfolioResult.metadata``."""
    binding = [
        name
        for name, mult, ident in zip(
            self.row_names, self.row_multipliers, self.row_identified, strict=True
        )
        if ident and mult != 0.0
    ]
    return {
        "kkt_residual": float(self.kkt_residual),
        "stationarity": self.stationarity,
        "dual_quality": _grade(self),
        "n_binding_constraints": len(binding),
        "binding_constraints": binding,
    }

CompiledConstraints dataclass

CompiledConstraints(
    kind: str,
    assets: tuple[str, ...],
    budget: float,
    lower: tuple[float, ...],
    upper: tuple[float, ...],
    row_names: tuple[str, ...],
    row_lower: tuple[float, ...],
    row_upper: tuple[float, ...],
    group_of: tuple[int, ...],
    specs: tuple[Constraint, ...],
    long_only: bool = True,
    weight_bounds: tuple[float, float] = (0.0, 1.0),
)

A resolved, feasible constraint set ready to hand to the solver.

Everything is stored host-side as tuples so the object stays hashable and picklable; the traced JAX arrays are built on demand by :meth:projection. That also makes :meth:perturb a plain field replacement, which is what lets a what-if re-solve reuse the compiled kernel.

Attributes:

Name	Type	Description
kind	str	Which projection kernel is needed — see the module constants. Part of the jit cache key.
assets	tuple[str, ...]	Asset names, in panel order.
budget	float	Required sum(w).
lower, upper		Effective per-asset bounds, length n_assets, after merging the implicit bounds with every :class:~jaxfolio.constraints.spec.Box.
row_names	tuple[str, ...]	Names of the constraint rows, index-aligned with row_lower / row_upper and with the group ids in group_of.
row_lower, row_upper		Effective row limits, already clamped to what the box alone permits (sum of member bounds). Clamping keeps every value finite — no infinities enter the kernel — and makes "this row is slack" an exact test rather than a tolerance.
group_of	tuple[int, ...]	Length n_assets; group_of[i] is the row index constraining asset i, or len(row_names) for the ungrouped sentinel.
specs	tuple[Constraint, ...]	The original specifications, for reporting and :meth:perturb.
long_only, weight_bounds		The implicit bounds this set was compiled against. Retained so :meth:perturb can recompile faithfully — bounds that came from an OptimizerConfig rather than from a Box spec would otherwise be silently reset to the defaults.

members

members(name: str) -> tuple[str, ...]

Asset names constrained by row name.

Source code in src/jaxfolio/constraints/compile.py

def members(self, name: str) -> tuple[str, ...]:
    """Asset names constrained by row ``name``."""
    k = self.row_names.index(name)
    return tuple(a for a, g in zip(self.assets, self.group_of, strict=True) if g == k)

projection

projection(
    *, aux: bool = False, path: bool = False
) -> tuple[Any, tuple]

Return (projection_fn, pparams) for the cached solver kernel.

aux=True selects the variant that projects only the weight block of a packed [w, tau] vector (the CVaR optimizer); path=True the row-wise variant for a (T, N) weight path. Imported lazily because :mod:jaxfolio.optimizers.base imports this package.

Source code in src/jaxfolio/constraints/compile.py

def projection(self, *, aux: bool = False, path: bool = False) -> tuple[Any, tuple]:
    """Return ``(projection_fn, pparams)`` for the cached solver kernel.

    ``aux=True`` selects the variant that projects only the weight block of a
    packed ``[w, tau]`` vector (the CVaR optimizer); ``path=True`` the
    row-wise variant for a ``(T, N)`` weight path. Imported lazily because
    :mod:`jaxfolio.optimizers.base` imports this package.
    """
    from jaxfolio.optimizers.base import projection_for

    return projection_for(self, aux=aux, path=path)

pparams

pparams() -> tuple

The traced projection parameters for this constraint set.

Bounds and limits are returned as jnp arrays for the vector kinds and as plain Python floats for the two scalar kinds — the latter reproducing byte-for-byte what the pre-existing select_projection produced.

Source code in src/jaxfolio/constraints/compile.py

def pparams(self) -> tuple:
    """The traced projection parameters for this constraint set.

    Bounds and limits are returned as ``jnp`` arrays for the vector kinds and
    as plain Python floats for the two scalar kinds — the latter reproducing
    byte-for-byte what the pre-existing ``select_projection`` produced.
    """
    import jax.numpy as jnp

    if self.kind == KIND_SIMPLEX:
        return (self.budget,)
    if self.kind == KIND_BOX:
        return (self.lower[0], self.upper[0], self.budget)
    lower = jnp.asarray(self.lower)
    upper = jnp.asarray(self.upper)
    if self.kind == KIND_BOXVEC:
        return (lower, upper, jnp.asarray(self.budget))
    return (
        lower,
        upper,
        jnp.asarray(self.budget),
        jnp.asarray(self.group_of, dtype=jnp.int32),
        # The sentinel's limits are never read (its multiplier is pinned to
        # zero), but the arrays must be n_rows + 1 long to size the segments.
        jnp.asarray((*self.row_lower, 0.0)),
        jnp.asarray((*self.row_upper, 0.0)),
    )

perturb

perturb(name: str, **changes: float) -> CompiledConstraints

Return a copy with one constraint's limits changed.

Only values move, so the perturbed set keeps the same kind and row count and therefore the same compiled kernel — the basis of a cheap what-if re-solve. Recompiles the specs so the new limits are re-validated and re-clamped, and so an infeasible relaxation is still caught.

Source code in src/jaxfolio/constraints/compile.py

def perturb(self, name: str, **changes: float) -> CompiledConstraints:
    """Return a copy with one constraint's limits changed.

    Only *values* move, so the perturbed set keeps the same ``kind`` and row
    count and therefore the same compiled kernel — the basis of a cheap
    what-if re-solve. Recompiles the specs so the new limits are re-validated
    and re-clamped, and so an infeasible relaxation is still caught.
    """
    if name not in self.row_names and name != "budget":
        known = ", ".join(self.row_names) or "(none)"
        raise KeyError(f"no constraint named {name!r} in this problem; rows are: {known}")
    specs = [replace(s, **changes) if s.name == name else s for s in self.specs]
    out = compile_constraints(
        specs,
        self.assets,
        long_only=self.long_only,
        weight_bounds=self.weight_bounds,
    )
    if out.kind != self.kind or out.n_rows != self.n_rows:
        raise ValueError(
            f"perturbing {name!r} changed the projection kind "
            f"({self.kind} -> {out.kind}); that is a structural change, not a relaxation"
        )
    return out

InfeasibleConstraints

Bases: ValueError

The requested constraint set admits no feasible portfolio.

Raised at compile time, before any moments are estimated or any solve is attempted. Because the supported rows have disjoint supports, the feasibility test is exact — necessary and sufficient — so this is never a false alarm, and conversely a set that compiles can always be satisfied.

Box dataclass

Box(
    lower: float | tuple[float, ...] = 0.0,
    upper: float | tuple[float, ...] = 1.0,
    assets: tuple[str, ...] | tuple[int, ...] | None = None,
    name: str = "box",
)

Bases: Constraint

Per-asset weight bounds lower <= w_i <= upper.

lower and upper are each either a scalar applied to every asset in scope, or a per-asset sequence. assets=None (the default) scopes the box to the whole universe and replaces the bounds implied by OptimizerConfig.long_only / weight_bounds; an assets-scoped box intersects with whatever is already in force, so per-name overrides compose::

Box(lower=0.0, upper=0.10)                    # every asset capped at 10%
Box(upper=0.02, assets=["ILLIQ"], name="illiquid")   # tighter on one name

A per-asset sequence must match the length of assets, or of the whole universe when assets is None.

width

width() -> int | None

Length of the per-asset bounds, or None if both are scalars.

Source code in src/jaxfolio/constraints/spec.py

def width(self) -> int | None:
    """Length of the per-asset bounds, or ``None`` if both are scalars."""
    for b in (self.lower, self.upper):
        if isinstance(b, tuple):
            return len(b)
    return None

Budget dataclass

Budget(total: float = 1.0, name: str = 'budget')

Bases: Constraint

Total invested weight: sum(w) == total.

Always present. Instantiate one explicitly only to override the default fully-invested 1.0 — e.g. Budget(0.0) for a dollar-neutral book.

Constraint

Bases: ABC

Base class for every named constraint.

Deliberately not a dataclass: it declares no fields, so subclasses stay free to order their own (a dataclass base would force every subclass field to come after the base's and to carry a default).

GroupCap dataclass

GroupCap(
    name: str,
    assets: tuple[str, ...] | tuple[int, ...],
    max: float | None = None,
    min: float | None = None,
)

Bases: RowConstraint

A cap and/or floor on the summed weight of a named group of assets.

The canonical sector constraint::

GroupCap("tech", ["AAPL", "MSFT", "NVDA"], max=0.30)   # tech <= 30%
GroupCap("energy", ["XOM", "CVX"], min=0.05)           # energy >= 5%
GroupCap("bonds", ["TLT", "IEF"], min=0.20, max=0.40)  # a band

Groups must be disjoint — no asset may appear in two of them. See :func:jaxfolio.constraints.compile.compile_constraints, which enforces it, and :mod:jaxfolio.constraints.structured for why.

ProjectionDuals

Bases: NamedTuple

A projected point together with the multipliers that produced it.

A NamedTuple so it is a JAX pytree: it flows through jit, vmap and lax control flow with no registration.

Attributes:

Name	Type	Description
weights	Array	The projected point.
budget	Array	Multiplier lam on sum(w) = budget. Signed.
rows	Array	(n_seg,) multipliers theta on the named rows — positive where a cap binds, negative where a floor binds, exactly zero where the row is slack. The last entry is the ungrouped sentinel and is always zero.
bounds	Array	(n,) multipliers beta on the per-asset box — negative at a lower bound, positive at an upper bound, zero in the interior. This is the sign that makes the stationarity identity v - w == budget + rows[gid] + bounds hold as written: writing the box as mu_lo >= 0 on lower - w and mu_hi >= 0 on w - upper, beta = mu_hi - mu_lo. Note the reduced cost of the original portfolio problem carries the opposite sign (positive at a lower bound), because it equals -beta / step; see :mod:jaxfolio.attribution.
identified	Array	(n_seg,) mask: is rows[k] actually pinned by the data? A row's multiplier is identified only when at least one of its members is strictly interior to its box — that member's stationarity equation is the one that separates theta_k from beta_i. When every member sits on a bound, the split of the shadow price between the row and the box bounds is genuinely not determined: any split summing to the same reduced cost is an equally valid KKT certificate. This kernel then reports the minimum-norm choice (theta = 0, everything attributed to the box), which is the reading users mean, but a different QP solver will legitimately report a different split. Verified against CVXPY: every identified row agrees to ~4e-9, while unidentified rows disagree by as much as 0.37.
budget_identified	Array	Is budget pinned by the data? Only if some coordinate is strictly interior and free of a binding row — such a coordinate satisfies v_i - w_i == lam on its own. If every interior coordinate lies inside a binding row, stationarity determines only the sum lam + theta_k, so the returned budget is one admissible value among an interval rather than the shadow price of capital. Callers reporting a budget shadow price must check this first.
feasibility_residual	Array	|sum(w) - budget| left after the bisection and the free-set correction. At the float precision floor when the problem is feasible; non-zero only if it is not.

SolverSpec

Bases: NamedTuple

A hashable, equality-stable solver key (safe as a jit static argument).

Attributes:

Name	Type	Description
kind	str	"spg" for the built-in spectral projected gradient, "optax" for everything else.
factory	Callable[..., Any] | None	The optax factory to call, e.g. optax.adamw. None for "spg".
name	str	Display label used in error messages and diagnostics.
options	tuple[tuple[str, Any], ...]	Extra keyword arguments for the factory, normalized to a sorted tuple of (key, value) pairs so the spec stays hashable.

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)

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

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),
    }

explain

explain(
    result: PortfolioResult,
    *,
    prim_tol: float | None = None,
    dual_tol: float | None = None,
) -> ConstraintReport

Explain a solved portfolio: which constraints bind, and what they cost.

Pure post-processing — reads only the payload captured at solve time and never re-solves. Always returns a well-formed report: optimizers that have no multipliers (closed-form weightings, the graph methods, the learning policies) come back with available=False and a specific reason, with the primal facts still filled in.

Parameters:

Name	Type	Description	Default
prim_tol	float | None	How close to a bound counts as on it. Defaults to 1e-6 scaled by the bound magnitudes.	None
dual_tol	float | None	How large a multiplier counts as non-zero. Defaults to 1e-6 scaled by ||grad f||, so it survives rescaling the objective.	None

Source code in src/jaxfolio/attribution.py

def explain(
    result: PortfolioResult,
    *,
    prim_tol: float | None = None,
    dual_tol: float | None = None,
) -> ConstraintReport:
    """Explain a solved portfolio: which constraints bind, and what they cost.

    Pure post-processing — reads only the payload captured at solve time and never
    re-solves. Always returns a well-formed report: optimizers that have no
    multipliers (closed-form weightings, the graph methods, the learning policies)
    come back with ``available=False`` and a specific ``reason``, with the primal
    facts still filled in.

    Parameters
    ----------
    prim_tol:
        How close to a bound counts as *on* it. Defaults to ``1e-6`` scaled by the
        bound magnitudes.
    dual_tol:
        How large a multiplier counts as non-zero. Defaults to ``1e-6`` scaled by
        ``||grad f||``, so it survives rescaling the objective.
    """
    duals: SolverDuals | None = getattr(result, "attribution", None)
    assets = tuple(result.assets)
    weights = np.asarray(result.weights, dtype=float)
    index = {a: i for i, a in enumerate(assets)}

    if duals is None:
        reason = _NO_DUALS_REASONS.get(
            result.method,
            f"{result.method} does not emit solver diagnostics — it is feasible by "
            "construction rather than optimal subject to constraints, so a zero weight "
            "is caused by the construction, not by a constraint",
        )
        return ConstraintReport(
            method=result.method,
            objective="",
            sense="minimize",
            assets=assets,
            weights=weights,
            available=False,
            quality="unavailable",
            reason=reason,
            stationarity=None,
            kkt_residual=None,
            step_check=None,
            budget_multiplier=None,
            budget_shadow_price=None,
            budget_note=None,
            constraints=(),
            assets_detail=tuple(
                AssetAttribution(
                    asset=a,
                    weight=float(weights[i]),
                    status="unknown",
                    reduced_cost=None,
                    intrinsic_cost=None,
                    imposed=(),
                    primary=None,
                    primary_kind=None,
                    primary_shadow_price=None,
                    confidence="unavailable",
                    reason=reason,
                )
                for i, a in enumerate(assets)
            ),
            _index=index,
        )

    quality = _grade(duals)
    sigma = 1.0 if duals.sense == "maximize" else -1.0
    trust = quality in ("exact", "approximate")

    bound_scale = max(float(np.max(np.abs(duals.upper))), float(np.max(np.abs(duals.lower))), 1e-12)
    ptol = prim_tol if prim_tol is not None else 1e-6 * bound_scale
    dtol = dual_tol if dual_tol is not None else 1e-6 * max(duals.gradient_norm, 1e-12)

    warnings = _collect_warnings(duals, quality)
    constraints = _row_duals(duals, assets, sigma, ptol, dtol, trust)
    detail = _asset_attributions(duals, assets, constraints, ptol, dtol, trust)

    budget_note = None
    budget_sp: float | None = None
    budget_mult: float | None = None
    if not trust:
        budget_note = "suppressed: " + quality
    elif not duals.budget_identified:
        budget_note = (
            "not identified — every interior asset lies inside a binding row, so "
            "stationarity pins only the sum of the budget and row multipliers"
        )
    elif _is_scale_invariant(duals):
        budget_note = (
            "not meaningful: the objective is scale-invariant (degree-0 homogeneous), "
            "so the optimum is unchanged by a change of budget"
        )
    else:
        budget_mult = duals.budget_multiplier
        budget_sp = sigma * duals.budget_multiplier

    return ConstraintReport(
        method=result.method,
        objective=duals.objective,
        sense=duals.sense,
        assets=assets,
        weights=weights,
        available=True,
        quality=quality,
        reason=None if trust else _quality_reason(duals, quality),
        stationarity=duals.stationarity,
        kkt_residual=duals.kkt_residual,
        step_check=duals.step_check,
        budget_multiplier=budget_mult,
        budget_shadow_price=budget_sp,
        budget_note=budget_note,
        constraints=constraints,
        assets_detail=detail,
        warnings=warnings,
        aux_stationarity=duals.aux_stationarity,
        l2_reg=duals.l2_reg,
        _index=index,
    )

solver_duals

solver_duals(
    objective,
    obj_params,
    compiled: CompiledConstraints,
    weights,
    *,
    objective_name: str,
    sense: str = "minimize",
    smooth: bool = True,
    convex: bool = True,
    l2_reg: float = 0.0,
    solver_kind: str = "spg",
    tol: float = 1e-07,
    aux_dim: int = 0,
) -> SolverDuals

Recover the KKT multipliers of a solved portfolio problem.

Runs after the solve, outside any lax control flow, at the cost of one gradient evaluation and two projections. Solver-agnostic by construction, which is what lets it retrofit a real stationarity measure onto the optax solvers (whose own info["residual"] is a weight-update norm, explicitly not a KKT test).

objective(z, obj_params) is the module-level objective that was minimized; weights is the returned iterate. For a packed [w, tau] variable pass aux_dim=1 and the full packed vector — the projection's normal cone factorizes as N_C(w) x {0}, so the weight-block multipliers are unaffected and tau's stationarity comes back separately.

Source code in src/jaxfolio/attribution.py

def solver_duals(
    objective,
    obj_params,
    compiled: CompiledConstraints,
    weights,
    *,
    objective_name: str,
    sense: str = "minimize",
    smooth: bool = True,
    convex: bool = True,
    l2_reg: float = 0.0,
    solver_kind: str = "spg",
    tol: float = 1e-7,
    aux_dim: int = 0,
) -> SolverDuals:
    """Recover the KKT multipliers of a solved portfolio problem.

    Runs *after* the solve, outside any ``lax`` control flow, at the cost of one
    gradient evaluation and two projections. Solver-agnostic by construction,
    which is what lets it retrofit a real stationarity measure onto the optax
    solvers (whose own ``info["residual"]`` is a weight-update norm, explicitly
    not a KKT test).

    ``objective(z, obj_params)`` is the module-level objective that was minimized;
    ``weights`` is the returned iterate. For a packed ``[w, tau]`` variable pass
    ``aux_dim=1`` and the full packed vector — the projection's normal cone
    factorizes as ``N_C(w) x {0}``, so the weight-block multipliers are unaffected
    and ``tau``'s stationarity comes back separately.
    """
    import jax
    import jax.numpy as jnp

    from jaxfolio.constraints.structured import (
        project_box_budget_duals,
        project_grouped_duals,
    )

    z = jnp.asarray(weights)
    grad = jax.grad(lambda x: objective(x, obj_params))(z)
    w = z[:-aux_dim] if aux_dim else z
    g = grad[:-aux_dim] if aux_dim else grad

    lower = jnp.asarray(compiled.lower)
    upper = jnp.asarray(compiled.upper)
    budget = jnp.asarray(compiled.budget)
    if compiled.kind == "grouped":
        gid = jnp.asarray(compiled.group_of, dtype=jnp.int32)
        g_lo = jnp.asarray((*compiled.row_lower, 0.0))
        g_hi = jnp.asarray((*compiled.row_upper, 0.0))

        def project(v):
            return project_grouped_duals(v, lower, upper, budget, gid, g_lo, g_hi)
    else:
        # Every non-grouped kind — including the simplex, whose feasible set is
        # exactly the box [0, budget] intersected with the budget hyperplane — is
        # served by the box kernel, so the caller sees one uniform payload.
        def project(v):
            return project_box_budget_duals(v, lower, upper, budget)

    # Probe at a displacement scaled to the problem, then again ten times smaller.
    # C is polyhedral, so d_s is constant for all small enough s; disagreement
    # between the two means the larger probe left that region and the multipliers
    # are being read off the wrong face.
    scale = jnp.maximum(jnp.max(jnp.abs(w)), 1e-8)
    gnorm_inf = jnp.maximum(jnp.max(jnp.abs(g)), 1e-30)
    step = _PROBE_FRAC * scale / gnorm_inf

    def descent(s):
        return (project(w - s * g).weights - w) / s

    d_main = descent(step)
    d_fine = descent(step * 0.1)
    # Normalize by the gradient, NOT by ||d_fine||. Both descents have the units of
    # g, and at a converged solution both are essentially zero — dividing by
    # ||d_fine|| there divides noise by noise and reports a meaningless 1.0. Against
    # ||g|| the measure stays interpretable at both ends: ~0 when the two probes
    # agree (including when both are at the float floor), and O(1) when the larger
    # step has left the region where the active set is constant.
    gnorm = jnp.linalg.norm(g)
    step_check = jnp.linalg.norm(d_main - d_fine) / (gnorm + 1e-30)

    duals = project(w - step * g)
    n_rows = compiled.n_rows
    rows = np.asarray(duals.rows)[:n_rows] / float(step) if n_rows else np.zeros(0)
    # rc_i = g_i + lambda + theta_{k(i)} = -beta_i / step. See the module docstring;
    # the sign flip is why `bounds` is negative at a lower bound while a reduced
    # cost is positive there.
    reduced = -np.asarray(duals.bounds) / float(step)

    return SolverDuals(
        weights=np.asarray(w, dtype=float),
        gradient=np.asarray(g, dtype=float),
        lower=np.asarray(compiled.lower, dtype=float),
        upper=np.asarray(compiled.upper, dtype=float),
        budget=float(compiled.budget),
        row_names=tuple(compiled.row_names),
        row_lower=tuple(compiled.row_lower),
        row_upper=tuple(compiled.row_upper),
        group_of=tuple(compiled.group_of),
        budget_multiplier=float(duals.budget) / float(step),
        row_multipliers=tuple(float(x) for x in rows),
        reduced_costs=reduced,
        row_identified=tuple(bool(x) for x in np.asarray(duals.identified)[:n_rows]),
        budget_identified=bool(duals.budget_identified),
        kkt_residual=float(jnp.linalg.norm(d_main)),
        gradient_norm=float(gnorm),
        step=float(step),
        step_check=float(step_check),
        aux_stationarity=float(jnp.abs(grad[-1])) if aux_dim else None,
        objective=objective_name,
        sense=sense,
        smooth=smooth,
        convex=convex,
        l2_reg=float(l2_reg),
        solver_kind=solver_kind,
        tol=float(tol),
    )

check_feasible

check_feasible(
    constraints: Sequence[Constraint] | None,
    assets: Sequence[str],
    *,
    long_only: bool = True,
    weight_bounds: tuple[float, float] = (0.0, 1.0),
) -> None

Validate a constraint set without solving anything.

Raises the same errors :func:compile_constraints would. Useful for validating user input up front — e.g. in a UI — since the feasibility test is exact and costs nothing.

Source code in src/jaxfolio/constraints/compile.py

def check_feasible(
    constraints: Sequence[Constraint] | None,
    assets: Sequence[str],
    *,
    long_only: bool = True,
    weight_bounds: tuple[float, float] = (0.0, 1.0),
) -> None:
    """Validate a constraint set without solving anything.

    Raises the same errors :func:`compile_constraints` would. Useful for
    validating user input up front — e.g. in a UI — since the feasibility test is
    exact and costs nothing.
    """
    compile_constraints(constraints, assets, long_only=long_only, weight_bounds=weight_bounds)

compile_constraints

compile_constraints(
    constraints: Sequence[Constraint] | None,
    assets: Sequence[str],
    *,
    long_only: bool = True,
    weight_bounds: tuple[float, float] = (0.0, 1.0),
) -> CompiledConstraints

Resolve constraint specifications against an asset universe.

Parameters:

Name	Type	Description	Default
constraints	Sequence[Constraint] | None	Specifications from :mod:jaxfolio.constraints.spec. Empty or None reproduces the pre-existing behaviour exactly (see the module docstring).	required
assets	Sequence[str]	Asset names in panel order, as returned by :func:jaxfolio.results.moments.	required
long_only	bool	The implicit bounds, normally config.long_only and config.bounds(). An unscoped :class:~jaxfolio.constraints.spec.Box replaces them; a scoped one intersects with them.	True
weight_bounds	bool	The implicit bounds, normally config.long_only and config.bounds(). An unscoped :class:~jaxfolio.constraints.spec.Box replaces them; a scoped one intersects with them.	True

Raises:

Type	Description
InfeasibleConstraints	The set admits no feasible portfolio. Exact, not heuristic.
ValueError	Two rows share an asset (see :mod:jaxfolio.constraints.structured), two constraints share a name, or more than one budget is given.
KeyError	A constraint names an asset that is not in the panel.

Source code in src/jaxfolio/constraints/compile.py

def compile_constraints(
    constraints: Sequence[Constraint] | None,
    assets: Sequence[str],
    *,
    long_only: bool = True,
    weight_bounds: tuple[float, float] = (0.0, 1.0),
) -> CompiledConstraints:
    """Resolve constraint specifications against an asset universe.

    Parameters
    ----------
    constraints:
        Specifications from :mod:`jaxfolio.constraints.spec`. Empty or ``None``
        reproduces the pre-existing behaviour exactly (see the module docstring).
    assets:
        Asset names in panel order, as returned by
        :func:`jaxfolio.results.moments`.
    long_only, weight_bounds:
        The implicit bounds, normally ``config.long_only`` and
        ``config.bounds()``. An unscoped :class:`~jaxfolio.constraints.spec.Box`
        replaces them; a scoped one intersects with them.

    Raises
    ------
    InfeasibleConstraints
        The set admits no feasible portfolio. Exact, not heuristic.
    ValueError
        Two rows share an asset (see :mod:`jaxfolio.constraints.structured`), two
        constraints share a name, or more than one budget is given.
    KeyError
        A constraint names an asset that is not in the panel.
    """
    names = tuple(str(a) for a in assets)
    n = len(names)
    if n == 0:
        raise ValueError("compile_constraints: the asset universe is empty")
    index = {a: i for i, a in enumerate(names)}
    specs = tuple(constraints or ())

    for spec in specs:
        if not isinstance(spec, Constraint):
            raise TypeError(
                "constraints must be jaxfolio.constraints specifications, got "
                f"{type(spec).__name__}"
            )
    seen: dict[str, Constraint] = {}
    for spec in specs:
        if spec.name in seen:
            raise ValueError(
                f"two constraints are both named {spec.name!r}; names must be unique so "
                "attribution can refer to them unambiguously"
            )
        seen[spec.name] = spec

    budgets = [s for s in specs if isinstance(s, Budget)]
    if len(budgets) > 1:
        raise ValueError(
            f"{len(budgets)} Budget constraints given ({', '.join(b.name for b in budgets)}); "
            "a portfolio has exactly one budget"
        )
    budget = float(budgets[0].total) if budgets else 1.0

    lower, upper = _merge_boxes(specs, names, index, n, long_only, weight_bounds)
    rows = [s for s in specs if isinstance(s, RowConstraint)]
    row_names, row_lower, row_upper, group_of = _build_rows(rows, index, n, lower, upper, names)

    _check_feasible_arrays(
        lower, upper, budget, row_names, row_lower, row_upper, group_of, names, rows
    )

    return CompiledConstraints(
        kind=_select_kind(lower, upper, row_names, long_only, budget),
        assets=names,
        budget=budget,
        lower=tuple(float(x) for x in lower),
        upper=tuple(float(x) for x in upper),
        row_names=tuple(row_names),
        row_lower=tuple(float(x) for x in row_lower),
        row_upper=tuple(float(x) for x in row_upper),
        group_of=tuple(int(g) for g in group_of),
        specs=specs,
        long_only=bool(long_only),
        weight_bounds=(float(weight_bounds[0]), float(weight_bounds[1])),
    )

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)

GroupFloor

GroupFloor(
    name: str,
    assets: Sequence[str] | Sequence[int],
    min: float,
) -> GroupCap

A named floor on a group's summed weight — sugar for GroupCap(min=...).

Returns a :class:GroupCap rather than a distinct class so that a group can never carry a floor and a cap that disagree, and so the compiler and the attribution report have exactly one row type to reason about.

Source code in src/jaxfolio/constraints/spec.py

def GroupFloor(  # noqa: N802 - a spec constructor, named like the class it returns
    name: str,
    assets: Sequence[str] | Sequence[int],
    min: float,  # noqa: A002 - mirrors GroupCap's public field name
) -> GroupCap:
    """A named floor on a group's summed weight — sugar for ``GroupCap(min=...)``.

    Returns a :class:`GroupCap` rather than a distinct class so that a group can
    never carry a floor and a cap that disagree, and so the compiler and the
    attribution report have exactly one row type to reason about.
    """
    return GroupCap(name, tuple(assets), min=min)

project_box_budget_duals

project_box_budget_duals(
    v: Array,
    lower: Array,
    upper: Array,
    budget: Array,
    *,
    max_iter: int = 50,
) -> ProjectionDuals

Box + budget projection with its multipliers.

The row-free case of :func:project_grouped_duals: rows is a single all-zero sentinel entry and identified is [False], so a caller can treat both kernels uniformly.

Source code in src/jaxfolio/constraints/structured.py

def project_box_budget_duals(
    v: Array,
    lower: Array,
    upper: Array,
    budget: Array,
    *,
    max_iter: int = 50,
) -> ProjectionDuals:
    """Box + budget projection with its multipliers.

    The row-free case of :func:`project_grouped_duals`: ``rows`` is a single
    all-zero sentinel entry and ``identified`` is ``[False]``, so a caller can
    treat both kernels uniformly.
    """
    lo_tau = jnp.min(v - upper) - 1.0
    hi_tau = jnp.max(v - lower) + 1.0

    def body(_, bounds):
        lo, hi = bounds
        mid = 0.5 * (lo + hi)
        too_big = jnp.sum(jnp.clip(v - mid, lower, upper)) > budget
        return (jnp.where(too_big, mid, lo), jnp.where(too_big, hi, mid))

    lo_tau, hi_tau = jax.lax.fori_loop(0, max_iter, body, (lo_tau, hi_tau))
    lam = 0.5 * (lo_tau + hi_tau)
    x = v - lam
    w = jnp.clip(x, lower, upper)
    w, resid = _absorb_residual(w, lower, upper, budget, free=jnp.ones_like(w, dtype=bool))
    # beta_i = x_i - w_i: negative where the lower bound binds, positive at the
    # upper bound, zero in the interior. Recomputed from the final w so it stays
    # consistent with the returned primal point.
    beta = jnp.where(_interior(w, lower, upper), 0.0, x - w)
    return ProjectionDuals(
        weights=w,
        budget=lam,
        rows=jnp.zeros(1, dtype=w.dtype),
        bounds=beta,
        identified=jnp.zeros(1, dtype=bool),
        # With no rows, any strictly-interior coordinate pins lam on its own.
        budget_identified=jnp.any(_interior(w, lower, upper)),
        feasibility_residual=resid,
    )

project_box_budget_vec

project_box_budget_vec(
    v: Array,
    lower: Array,
    upper: Array,
    budget: Array,
    *,
    max_iter: int = 50,
) -> Array

Project v onto {w : lower <= w <= upper, sum(w) = budget}.

The per-asset-bounds generalization of :func:jaxfolio.constraints.projections.project_box_budget: lower and upper are (n,) arrays rather than scalars. Same method — bisect the budget multiplier tau, since w(tau) = clip(v - tau, lower, upper) is monotone decreasing in tau — and the same feasibility requirement, sum(lower) <= budget <= sum(upper).

Source code in src/jaxfolio/constraints/structured.py

def project_box_budget_vec(
    v: Array,
    lower: Array,
    upper: Array,
    budget: Array,
    *,
    max_iter: int = 50,
) -> Array:
    """Project ``v`` onto ``{w : lower <= w <= upper, sum(w) = budget}``.

    The per-asset-bounds generalization of
    :func:`jaxfolio.constraints.projections.project_box_budget`: ``lower`` and
    ``upper`` are ``(n,)`` arrays rather than scalars. Same method — bisect the
    budget multiplier ``tau``, since ``w(tau) = clip(v - tau, lower, upper)`` is
    monotone decreasing in ``tau`` — and the same feasibility requirement,
    ``sum(lower) <= budget <= sum(upper)``.
    """
    return project_box_budget_duals(v, lower, upper, budget, max_iter=max_iter).weights

project_grouped

project_grouped(
    v: Array,
    lower: Array,
    upper: Array,
    budget: Array,
    gid: Array,
    g_lower: Array,
    g_upper: Array,
    *,
    outer_iter: int = _OUTER_ITER,
    inner_iter: int = _INNER_ITER,
) -> Array

Project onto box + budget + disjoint group rows.

See :func:project_grouped_duals for the algorithm and the meaning of every argument; this is the primal-only wrapper used inside the solver loop.

Source code in src/jaxfolio/constraints/structured.py

def project_grouped(
    v: Array,
    lower: Array,
    upper: Array,
    budget: Array,
    gid: Array,
    g_lower: Array,
    g_upper: Array,
    *,
    outer_iter: int = _OUTER_ITER,
    inner_iter: int = _INNER_ITER,
) -> Array:
    """Project onto box + budget + disjoint group rows.

    See :func:`project_grouped_duals` for the algorithm and the meaning of every
    argument; this is the primal-only wrapper used inside the solver loop.
    """
    return project_grouped_duals(
        v,
        lower,
        upper,
        budget,
        gid,
        g_lower,
        g_upper,
        outer_iter=outer_iter,
        inner_iter=inner_iter,
    ).weights

project_grouped_duals

project_grouped_duals(
    v: Array,
    lower: Array,
    upper: Array,
    budget: Array,
    gid: Array,
    g_lower: Array,
    g_upper: Array,
    *,
    outer_iter: int = _OUTER_ITER,
    inner_iter: int = _INNER_ITER,
) -> ProjectionDuals

Project onto box + budget + disjoint group rows, returning multipliers.

Parameters:

Name	Type	Description	Default
v	Array	(n,) point to project.	required
lower	Array	(n,) per-asset bounds.	required
upper	Array	(n,) per-asset bounds.	required
budget	Array	Scalar; sum(w) == budget.	required
gid	Array	(n,) int32 segment id per asset. Index n_seg - 1 is the ungrouped sentinel: assets assigned to it are subject to no row, and its multiplier is pinned to zero. Always allocated, possibly empty.	required
g_lower	Array	(n_seg,) row bounds. The sentinel's entries are ignored.	required
g_upper	Array	(n_seg,) row bounds. The sentinel's entries are ignored.	required

Returns:

Type	Description
class:`ProjectionDuals`. Validated against CVXPY over 40 random problems
(``n`` in 6..40, 2..4 rows, long-only and long-short): weights agree to
2.1e-8, the budget multiplier to 2.2e-9, and every *identified* row
multiplier to 3.9e-9. See ``ProjectionDuals.identified`` for the rows whose
multiplier is not pinned by the data.

Notes

The inner target is clip(freesum_k(lam), L_k, U_k) — not an early exit when the row is already satisfied. That distinction is load-bearing: the clipped-target form is what makes the dual decomposition unique, so lam and theta cannot drift into offsetting each other. It also makes "row is slack" an exact test (target == freesum), which is how theta_k is forced to a clean zero instead of the arbitrary value a degenerate bracket would otherwise return.

The outer bracket needs no theta term even though w depends on theta: the inner solve pins each group's sum to clip(freesum_k(lam), L_k, U_k) regardless of how large theta_k gets, so the total is a function of lam alone — and a non-increasing one, since every freesum_k and every ungrouped coordinate is non-increasing in lam and clip is monotone.

Source code in src/jaxfolio/constraints/structured.py

def project_grouped_duals(
    v: Array,
    lower: Array,
    upper: Array,
    budget: Array,
    gid: Array,
    g_lower: Array,
    g_upper: Array,
    *,
    outer_iter: int = _OUTER_ITER,
    inner_iter: int = _INNER_ITER,
) -> ProjectionDuals:
    """Project onto box + budget + disjoint group rows, returning multipliers.

    Parameters
    ----------
    v:
        ``(n,)`` point to project.
    lower, upper:
        ``(n,)`` per-asset bounds.
    budget:
        Scalar; ``sum(w) == budget``.
    gid:
        ``(n,)`` int32 segment id per asset. Index ``n_seg - 1`` is the
        **ungrouped sentinel**: assets assigned to it are subject to no row, and
        its multiplier is pinned to zero. Always allocated, possibly empty.
    g_lower, g_upper:
        ``(n_seg,)`` row bounds. The sentinel's entries are ignored.

    Returns
    -------
    :class:`ProjectionDuals`. Validated against CVXPY over 40 random problems
    (``n`` in 6..40, 2..4 rows, long-only and long-short): weights agree to
    2.1e-8, the budget multiplier to 2.2e-9, and every *identified* row
    multiplier to 3.9e-9. See ``ProjectionDuals.identified`` for the rows whose
    multiplier is not pinned by the data.

    Notes
    -----
    The inner target is ``clip(freesum_k(lam), L_k, U_k)`` — *not* an early exit
    when the row is already satisfied. That distinction is load-bearing: the
    clipped-target form is what makes the dual decomposition unique, so ``lam``
    and ``theta`` cannot drift into offsetting each other. It also makes "row is
    slack" an exact test (``target == freesum``), which is how ``theta_k`` is
    forced to a clean zero instead of the arbitrary value a degenerate bracket
    would otherwise return.

    The outer bracket needs no ``theta`` term even though ``w`` depends on
    ``theta``: the inner solve pins each group's sum to
    ``clip(freesum_k(lam), L_k, U_k)`` regardless of how large ``theta_k`` gets,
    so the total is a function of ``lam`` alone — and a non-increasing one, since
    every ``freesum_k`` and every ungrouped coordinate is non-increasing in
    ``lam`` and ``clip`` is monotone.
    """
    n_seg = g_lower.shape[0]
    active = jnp.arange(n_seg) < n_seg - 1

    def sums(lam):
        x = v - lam
        free = jnp.clip(x, lower, upper)
        freesum = jax.ops.segment_sum(free, gid, num_segments=n_seg)
        target = jnp.where(active, jnp.clip(freesum, g_lower, g_upper), freesum)
        theta = _solve_theta(x, lower, upper, gid, target, active, n_seg, inner_iter)
        return x, freesum, target, theta

    def total(lam):
        _x, _freesum, target, _theta = sums(lam)
        return jnp.sum(target)

    lo = jnp.min(v - upper) - 1.0
    hi = jnp.max(v - lower) + 1.0

    def body(_, bounds):
        lo, hi = bounds
        mid = 0.5 * (lo + hi)
        too_big = total(mid) > budget
        return (jnp.where(too_big, mid, lo), jnp.where(too_big, hi, mid))

    lo, hi = jax.lax.fori_loop(0, outer_iter, body, (lo, hi))
    lam = 0.5 * (lo + hi)

    x, freesum, target, theta = sums(lam)
    # Exact test: clip() returns freesum unchanged when the row is satisfied, so
    # equality here means "not binding" and pins theta to zero, discarding the
    # arbitrary value a saturated (degenerate) bracket would have produced.
    slack = jnp.logical_or(~active, target == freesum)
    theta = jnp.where(slack, 0.0, theta)

    shifted = x - theta[gid]
    w = jnp.clip(shifted, lower, upper)
    w, resid = _absorb_residual(w, lower, upper, budget, free=slack[gid])
    beta = jnp.where(_interior(w, lower, upper), 0.0, shifted - w)
    # A row's multiplier is pinned by the data only if some member is strictly
    # interior to its box: that member's stationarity equation is the one that
    # separates theta_k from beta. With every member on a bound, only the sum
    # theta_k + beta_i is determined and the split is arbitrary.
    interior = _interior(w, lower, upper)
    identified = jnp.logical_and(
        active, jax.ops.segment_sum(interior.astype(w.dtype), gid, num_segments=n_seg) > 0
    )
    return ProjectionDuals(
        weights=w,
        budget=lam,
        rows=theta,
        bounds=beta,
        identified=identified,
        # lam is pinned alone only by an interior coordinate whose row is slack;
        # inside a binding row, stationarity fixes only lam + theta_k.
        budget_identified=jnp.any(jnp.logical_and(interior, slack[gid])),
        feasibility_residual=resid,
    )

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

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,
    *,
    constraints: Sequence[Any] = (),
    assets: Sequence[str] | None = None,
) -> Callable[[Array], Array]

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

constraints accepts named specifications from :mod:jaxfolio.constraints.spec (sector caps, per-asset bound vectors). They need assets — the panel's asset names in order — to resolve names to columns; pass integer-indexed specs to skip that. The closure form re-traces per call anyway, so unlike :func:select_projection there is no cache concern here.

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

def make_projection(
    long_only: bool,
    weight_bounds: tuple[float, float],
    budget: float = 1.0,
    *,
    constraints: Sequence[Any] = (),
    assets: Sequence[str] | None = None,
) -> Callable[[Array], Array]:
    """Return the ``w -> w`` projection matching the requested constraint set.

    ``constraints`` accepts named specifications from
    :mod:`jaxfolio.constraints.spec` (sector caps, per-asset bound vectors). They
    need ``assets`` — the panel's asset names in order — to resolve names to
    columns; pass integer-indexed specs to skip that. The closure form re-traces
    per call anyway, so unlike :func:`select_projection` there is no cache concern
    here.
    """
    if constraints:
        from jaxfolio.constraints.compile import compile_constraints

        if assets is None:
            raise ValueError(
                "make_projection: named constraints need `assets` (the panel's asset names) "
                "to resolve names to columns"
            )
        compiled = compile_constraints(
            constraints, assets, long_only=long_only, weight_bounds=weight_bounds
        )
        fn, pparams = projection_for(compiled)
        return lambda w: fn(w, pparams)
    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,
    path: bool = False,
    constraints: Sequence[Any] = (),
    assets: Sequence[str] | None = None,
)

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). path=True selects the row-wise variant for a (T, N) weight path (used by the multi-period optimizer). pparams is a tuple of Python floats — passed as a traced argument to the kernel, and identical for all three variants.

constraints accepts named specifications from :mod:jaxfolio.constraints.spec, which need assets to resolve names to columns. With no named constraints this returns exactly what it always did: the same function object and the same Python-float tuple, so no pre-existing solve changes kernel, cache entry, or dtype.

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

def select_projection(
    long_only: bool,
    weight_bounds: tuple[float, float],
    *,
    aux: bool = False,
    path: bool = False,
    constraints: Sequence[Any] = (),
    assets: Sequence[str] | None = None,
):
    """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). ``path=True``
    selects the row-wise variant for a ``(T, N)`` weight path (used by the
    multi-period optimizer). ``pparams`` is a tuple of Python floats — passed as
    a traced argument to the kernel, and identical for all three variants.

    ``constraints`` accepts named specifications from
    :mod:`jaxfolio.constraints.spec`, which need ``assets`` to resolve names to
    columns. With no named constraints this returns exactly what it always did:
    the same function object and the same Python-float tuple, so no pre-existing
    solve changes kernel, cache entry, or dtype.
    """
    if constraints:
        from jaxfolio.constraints.compile import compile_constraints

        if assets is None:
            raise ValueError(
                "select_projection: named constraints need `assets` (the panel's asset "
                "names) to resolve names to columns"
            )
        compiled = compile_constraints(
            constraints, assets, long_only=long_only, weight_bounds=weight_bounds
        )
        return projection_for(compiled, aux=aux, path=path)
    _shape_key(aux, path)  # validate the aux/path combination
    lo, hi = weight_bounds
    simplex = long_only and lo <= 0.0 and hi >= 1.0
    if simplex:
        if path:
            return _proj_simplex_path, (1.0,)
        return (_proj_simplex_aux if aux else _proj_simplex), (1.0,)
    bounds = (float(lo), float(hi), 1.0)
    if path:
        return _proj_box_path, bounds
    return (_proj_box_aux if aux else _proj_box), bounds

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 | Callable[..., Any] | SolverSpec = "spg",
    learning_rate: float | None = None,
    max_iter: int = 2000,
    tol: float = 1e-07,
    solver_options: Mapping[str, Any] | None = None,
) -> 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. solver is "spg", any optax optimizer name or factory, or a pre-resolved :class:~jaxfolio.solvers.SolverSpec. Returns (weights, info) with info["iterations"] and info["residual"] (the solver's convergence measure at the final iterate).

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

def solve_constrained(
    objective,
    obj_params,
    w0: Array,
    projection,
    proj_params,
    *,
    solver: str | Callable[..., Any] | SolverSpec = "spg",
    learning_rate: float | None = None,
    max_iter: int = 2000,
    tol: float = 1e-7,
    solver_options: Mapping[str, Any] | None = None,
) -> 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.
    ``solver`` is ``"spg"``, any optax optimizer name or factory, or a
    pre-resolved :class:`~jaxfolio.solvers.SolverSpec`.
    Returns ``(weights, info)`` with ``info["iterations"]`` and
    ``info["residual"]`` (the solver's convergence measure at the final iterate).
    """
    w, iters, resid = _solve_cached(
        objective,
        projection,
        obj_params,
        proj_params,
        w0,
        jnp.asarray(tol),
        solver=resolve_solver(solver, solver_options),
        learning_rate=learning_rate,
        max_iter=max_iter,
    )
    return w, {"iterations": iters, "residual": resid}

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 | Callable[..., Any] | SolverSpec = "spg",
    solver_options: Mapping[str, Any] | None = None,
) -> 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. solver is "spg", any optax optimizer name or factory, or a pre-resolved :class:~jaxfolio.solvers.SolverSpec. Returns (weights, info) with info["iterations"] and info["residual"].

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 | Callable[..., Any] | SolverSpec = "spg",
    solver_options: Mapping[str, Any] | None = None,
) -> 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`. ``solver`` is ``"spg"``, any optax optimizer name
    or factory, or a pre-resolved :class:`~jaxfolio.solvers.SolverSpec`.
    Returns ``(weights, info)`` with ``info["iterations"]`` and
    ``info["residual"]``.
    """
    w, iters, resid = _run_solver(
        objective,
        projection,
        w0,
        solver=solver,
        learning_rate=learning_rate,
        max_iter=max_iter,
        tol=tol,
        solver_options=solver_options,
    )
    return w, {"iterations": iters, "residual": resid}

solve_weight_path

solve_weight_path(
    mu_path,
    cov,
    w_prev,
    *,
    risk_aversion: float = 1.0,
    c_lin=0.0,
    c_quad=0.0,
    l2_reg: float = 0.0,
    trade_eps_rel: float = 0.1,
    refine: int = DEFAULT_REFINE,
    long_only: bool = True,
    weight_bounds: tuple[float, float] = (0.0, 1.0),
    solver: Any = "spg",
    learning_rate: float | None = None,
    max_iter: int = 2000,
    tol: float = 1e-06,
    solver_options: Mapping[str, Any] | None = None,
) -> tuple[Array, dict[str, Any]]

Solve the multi-period mean-variance problem for a whole weight path.

This is the primitive behind :func:multi_period_mean_variance, exposed for callers who already hold forecasts and want the raw path back.

Parameters:

Name	Type	Description	Default
mu_path		Expected returns per period, shape (T, N).	required
cov		Covariance, either (N, N) (shared across periods) or (T, N, N).	required
w_prev		Currently held weights, shape (N,). Need not sum to one — the zero vector is a valid "flat book", and the cost of establishing a position from it is priced correctly.	required
risk_aversion	float	Coefficient gamma on the per-period variance term.	1.0
c_lin		Linear and quadratic trade-cost coefficients in decimal units, broadcast against (T, N). Scalars, (N,) and (T, N) all work.	0.0
c_quad		Linear and quadratic trade-cost coefficients in decimal units, broadcast against (T, N). Scalars, (N,) and (T, N) all work.	0.0
trade_eps_rel	float	Starting (widest) Huber half-width, relative to budget / N.	0.1
refine	int	Number of geometric refinements of the Huber width beyond the first stage, each shrinking it by a factor of three. Every stage warm-starts from the previous one and reuses the same compiled kernel, because the width is a traced parameter — so the ladder costs iterations, never compilations. The returned path is the stage that scored best under the exact (un-smoothed) objective, not simply the last one: accuracy is not monotone in the width, so selecting is what makes the result reliable.	DEFAULT_REFINE
tol	float	Convergence tolerance. Defaults to 1e-6 rather than the package-wide 1e-7 because the float32 projected-gradient norm floors around 1e-6 on this problem, so a tighter value only buys wasted iterations.	1e-06

Returns:

Type	Description
tuple	(W, info) where W has shape (T, N) and info carries iteration counts, the per-stage widths, the final residual, and both the smoothed and exact objective values.

Notes

A high c_lin relative to the Huber width legitimately leaves the projected-gradient residual on a plateau far above tol while the answer is exactly right (the correct action is "do not trade", and the residual carries the un-actioned cost gradient). Read info["converged"] as a diagnostic, not a failure.

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

def solve_weight_path(
    mu_path,
    cov,
    w_prev,
    *,
    risk_aversion: float = 1.0,
    c_lin=0.0,
    c_quad=0.0,
    l2_reg: float = 0.0,
    trade_eps_rel: float = 1e-1,
    refine: int = DEFAULT_REFINE,
    long_only: bool = True,
    weight_bounds: tuple[float, float] = (0.0, 1.0),
    solver: Any = "spg",
    learning_rate: float | None = None,
    max_iter: int = 2000,
    tol: float = 1e-6,
    solver_options: Mapping[str, Any] | None = None,
) -> tuple[Array, dict[str, Any]]:
    """Solve the multi-period mean-variance problem for a whole weight path.

    This is the primitive behind :func:`multi_period_mean_variance`, exposed for
    callers who already hold forecasts and want the raw path back.

    Parameters
    ----------
    mu_path:
        Expected returns per period, shape ``(T, N)``.
    cov:
        Covariance, either ``(N, N)`` (shared across periods) or ``(T, N, N)``.
    w_prev:
        Currently held weights, shape ``(N,)``. Need **not** sum to one — the
        zero vector is a valid "flat book", and the cost of establishing a
        position from it is priced correctly.
    risk_aversion:
        Coefficient ``gamma`` on the per-period variance term.
    c_lin, c_quad:
        Linear and quadratic trade-cost coefficients in decimal units, broadcast
        against ``(T, N)``. Scalars, ``(N,)`` and ``(T, N)`` all work.
    trade_eps_rel:
        Starting (widest) Huber half-width, relative to ``budget / N``.
    refine:
        Number of geometric refinements of the Huber width beyond the first stage,
        each shrinking it by a factor of three. Every stage warm-starts from the
        previous one and reuses the *same* compiled kernel, because the width is a
        traced parameter — so the ladder costs iterations, never compilations.
        The returned path is the stage that scored best under the **exact**
        (un-smoothed) objective, not simply the last one: accuracy is not monotone
        in the width, so selecting is what makes the result reliable.
    tol:
        Convergence tolerance. Defaults to ``1e-6`` rather than the package-wide
        ``1e-7`` because the float32 projected-gradient norm floors around
        ``1e-6`` on this problem, so a tighter value only buys wasted iterations.

    Returns
    -------
    tuple
        ``(W, info)`` where ``W`` has shape ``(T, N)`` and ``info`` carries
        iteration counts, the per-stage widths, the final residual, and both the
        smoothed and exact objective values.

    Notes
    -----
    A high ``c_lin`` relative to the Huber width legitimately leaves the
    projected-gradient residual on a plateau far above ``tol`` while the answer is
    exactly right (the correct action is "do not trade", and the residual carries
    the un-actioned cost gradient). Read ``info["converged"]`` as a diagnostic,
    not a failure.
    """
    mu_path = jnp.asarray(mu_path, dtype=float)
    cov = jnp.asarray(cov, dtype=float)
    if mu_path.ndim != 2:
        raise ValueError(f"mu_path must be 2-D (T, N), got shape {mu_path.shape}")
    horizon, n = mu_path.shape
    if cov.ndim == 2:
        objective = _mp_shared_cov
        if cov.shape != (n, n):
            raise ValueError(
                f"cov shape {cov.shape} does not match {n} assets; expected ({n}, {n})"
            )
    elif cov.ndim == 3:
        objective = _mp_path_cov
        if cov.shape != (horizon, n, n):
            raise ValueError(
                f"cov shape {cov.shape} does not match the path; expected ({horizon}, {n}, {n})"
            )
    else:
        raise ValueError(f"cov must be (N, N) or (T, N, N), got shape {cov.shape}")
    if refine < 0:
        raise ValueError(f"refine must be non-negative, got {refine}")
    if trade_eps_rel <= 0.0:
        raise ValueError(f"trade_eps_rel must be strictly positive, got {trade_eps_rel}")

    projection, pparams = select_projection(long_only, weight_bounds, path=True)
    budget = pparams[-1]

    # ``dtype=`` is required, not decorative: jnp.full(n, 1/n) is weak-typed and
    # jnp.tile preserves that, but the solver's output is strongly typed — so a
    # weak-typed start would make stage 2 a *different* jit cache key and compile
    # the kernel twice for what is one problem.
    w_prev_j = jnp.asarray(w_prev, dtype=mu_path.dtype).reshape(-1)
    if w_prev_j.shape != (n,):
        raise ValueError(
            f"w_prev shape {w_prev_j.shape} does not match {n} assets; expected ({n},)"
        )
    # Project the warm start onto the feasible set before anything reads it. The
    # solver would project internally anyway, but the stage-selection below also
    # scores this iterate as a candidate — and an *infeasible* point can score
    # lower than every feasible one, which would let it win and be returned. That
    # matters concretely whenever ``w_prev`` is not a valid portfolio: a partially
    # invested book (the backtester hands over exactly that before its first
    # rebalance) or holdings that violate the configured bounds.
    W = projection(jnp.tile(w_prev_j, (horizon, 1)), pparams)

    params = {
        "mu_path": mu_path,
        "cov": cov,
        "w_prev": w_prev_j,
        "risk_aversion": jnp.asarray(risk_aversion, dtype=mu_path.dtype),
        "l2": jnp.asarray(l2_reg, dtype=mu_path.dtype),
        "c_lin": jnp.broadcast_to(jnp.asarray(c_lin, dtype=mu_path.dtype), (n,))
        if jnp.ndim(c_lin) <= 1
        else jnp.asarray(c_lin, dtype=mu_path.dtype),
        "c_quad": jnp.broadcast_to(jnp.asarray(c_quad, dtype=mu_path.dtype), (n,))
        if jnp.ndim(c_quad) <= 1
        else jnp.asarray(c_quad, dtype=mu_path.dtype),
        "trade_eps": jnp.asarray(1.0, dtype=mu_path.dtype),
    }

    # Walk a ladder of shrinking Huber widths, warm-starting each stage from the
    # previous one, and keep the iterate that is best under the *exact* objective.
    #
    # Selecting rather than simply taking the last stage is what makes this
    # reliable. A single width cannot serve every problem: when the linear cost is
    # comparable to the whole mean-variance curvature (common at 50-100bps), a
    # narrow band makes the smoothed problem so stiff that the projected-gradient
    # method crawls, while a wide band biases the answer. Worse, accuracy is *not*
    # monotone in the width, so trusting the narrowest stage is measurably wrong —
    # it was observed 8% to 160% off the true optimum. Because the un-smoothed
    # objective is cheap to evaluate, the ladder can be treated as a search over
    # smoothing levels judged by the real criterion, which bounds the error by the
    # best stage instead of the last. That took the worst case across a grid of
    # panel sizes and cost regimes from ~163% down to ~1e-4.
    eps = float(trade_eps_rel) * budget / n
    stage_iters: list[int] = []
    stage_eps: list[float] = []
    stage_exact: list[float] = []
    best_W = W
    # ``W`` is feasible here (projected above), so this is a legitimate candidate:
    # holding still is often genuinely optimal once trading is expensive enough.
    best_exact = float(_exact_objective(W, {**params, "trade_eps": jnp.asarray(eps)}))
    best_stage = -1  # -1 => the (projected) no-trade path was never beaten
    best_residual = float("inf")
    for stage in range(refine + 1):
        params["trade_eps"] = jnp.asarray(eps, dtype=mu_path.dtype)
        W, info = solve_constrained(
            objective,
            params,
            W,
            projection,
            pparams,
            solver=solver,
            learning_rate=learning_rate,
            max_iter=max_iter,
            tol=tol,
            solver_options=solver_options,
        )
        stage_iters.append(int(info["iterations"]))
        stage_eps.append(eps)
        value = float(_exact_objective(W, params))
        stage_exact.append(value)
        if value < best_exact:
            best_exact, best_W, best_stage = value, W, stage
            best_residual = float(info["residual"])
        eps *= _REFINE_FACTOR

    W = best_W
    # Report the smoothed objective at the width that actually produced ``W``.
    params["trade_eps"] = jnp.asarray(
        stage_eps[best_stage] if best_stage >= 0 else stage_eps[0], dtype=mu_path.dtype
    )
    smoothed = float(objective(W, params))
    exact = best_exact
    residual = best_residual
    d = np.asarray(W - _path_prev(W, w_prev_j))
    turnover_path = np.abs(d).sum(axis=1)
    c_lin_np = np.asarray(params["c_lin"])
    c_quad_np = np.asarray(params["c_quad"])

    info = {
        "iterations": int(sum(stage_iters)),
        "stage_iterations": tuple(stage_iters),
        "stage_trade_eps": tuple(stage_eps),
        "stage_objective_exact": tuple(stage_exact),
        "selected_stage": int(best_stage),
        "trade_eps": float(stage_eps[best_stage] if best_stage >= 0 else stage_eps[0]),
        "residual": float(residual),
        "converged": bool(float(residual) <= tol),
        "turnover_path": turnover_path,
        "total_turnover": float(turnover_path.sum()),
        "total_cost": float((c_lin_np * np.abs(d)).sum() + (c_quad_np * d**2).sum()),
        "objective_smoothed": smoothed,
        "objective_exact": exact,
        "smoothing_bias": exact - smoothed,
    }
    return W, info

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,
    trajectory=None,
    attribution=None,
) -> 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.

trajectory is an optional (T, n) planned weight path for multi-period optimizers. It is attached verbatim — not cleaned by clean_eps, so the path keeps the solver's exact iterate while weights is tidied for display.

attribution is an optional :class:jaxfolio.attribution.SolverDuals. It is attached verbatim and a small JSON-serializable digest of it is merged into metadata, so consumers that only read metadata still see the convergence and binding-constraint summary.

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,
    trajectory=None,
    attribution=None,
) -> 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.

    ``trajectory`` is an optional ``(T, n)`` planned weight path for multi-period
    optimizers. It is attached verbatim — *not* cleaned by ``clean_eps``, so the
    path keeps the solver's exact iterate while ``weights`` is tidied for display.

    ``attribution`` is an optional :class:`jaxfolio.attribution.SolverDuals`. It is
    attached verbatim and a small JSON-serializable digest of it is merged into
    ``metadata``, so consumers that only read ``metadata`` still see the
    convergence and binding-constraint summary.
    """
    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

    meta = dict(metadata or {})
    if attribution is not None:
        meta.update(attribution.summary())

    return PortfolioResult(
        weights=w,
        assets=list(assets),
        method=method,
        expected_return=ann_mu,
        volatility=ann_vol,
        sharpe=sharpe,
        metadata=meta,
        trajectory=trajectory,
        attribution=attribution,
    )

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

available_solvers

available_solvers() -> tuple[str, ...]

Return the solver names that can be passed as solver=....

("spg",) followed by the sorted names of every optax and optax.contrib factory that takes a learning_rate. This is a discovery aid for humans and for error messages — resolution itself is more permissive, so optimizers with their own adaptive step (dowg, cocob) still work even though they are not listed here.

Source code in src/jaxfolio/solvers.py

def available_solvers() -> tuple[str, ...]:
    """Return the solver names that can be passed as ``solver=...``.

    ``("spg",)`` followed by the sorted names of every optax and ``optax.contrib``
    factory that takes a ``learning_rate``. This is a discovery aid for humans and
    for error messages — resolution itself is more permissive, so optimizers with
    their own adaptive step (``dowg``, ``cocob``) still work even though they are
    not listed here.
    """
    names: set[str] = set()
    for module in (optax, optax.contrib):
        for name in dir(module):
            if name.startswith("_") or name.startswith("scale_by_"):
                continue
            candidate = getattr(module, name, None)
            if not callable(candidate) or inspect.isclass(candidate):
                continue
            try:
                params = inspect.signature(candidate).parameters
            except (ValueError, TypeError):
                continue
            if "learning_rate" in params:
                names.add(name)
    return (SPG, *sorted(names))

resolve_solver

resolve_solver(
    solver: str | Callable[..., Any] | SolverSpec = SPG,
    solver_options: Mapping[str, Any] | None = None,
) -> SolverSpec

Normalize a solver spelling into a :class:SolverSpec.

Accepts "spg", an optax optimizer name, an optax factory callable, or an already-resolved :class:SolverSpec (the function is idempotent, so callers can normalize defensively). Raises :class:ValueError for an unknown name and :class:TypeError for an unusable object.

Source code in src/jaxfolio/solvers.py

def resolve_solver(
    solver: str | Callable[..., Any] | SolverSpec = SPG,
    solver_options: Mapping[str, Any] | None = None,
) -> SolverSpec:
    """Normalize a solver spelling into a :class:`SolverSpec`.

    Accepts ``"spg"``, an optax optimizer name, an optax factory callable, or an
    already-resolved :class:`SolverSpec` (the function is idempotent, so callers
    can normalize defensively). Raises :class:`ValueError` for an unknown name and
    :class:`TypeError` for an unusable object.
    """
    if isinstance(solver, SolverSpec):
        return solver

    options = _normalize_options(solver_options)

    if solver == SPG:
        if options:
            raise ValueError(
                "solver_options apply to optax solvers only; 'spg' is tuned by "
                "`learning_rate` (its initial Barzilai-Borwein step) and `tol`"
            )
        return _SPG_SPEC

    if isinstance(solver, str):
        factory = _lookup(solver)
        return SolverSpec("optax", factory, solver, options)

    # ``optax.nadam``/``nadamw`` are functools.partial objects, so gate on
    # callable() rather than inspect.isfunction().
    if callable(solver):
        name = getattr(solver, "__name__", None) or repr(solver)
        return SolverSpec("optax", solver, name, options)

    if hasattr(solver, "init") and hasattr(solver, "update"):
        raise TypeError(
            "pass the optax *factory* rather than a pre-built GradientTransformation "
            "(e.g. solver=optax.adamw with solver_options={'weight_decay': 1e-3}, not "
            "solver=optax.adamw(1e-3)). A pre-built transformation is a fresh object on "
            "every call, so it would recompile the solver kernel on every solve. For an "
            "optax.chain(...) composition, wrap it in a module-level factory function "
            "`def my_solver(learning_rate): ...` and pass that."
        )

    raise TypeError(
        f"solver must be 'spg', an optax optimizer name, or an optax factory callable; "
        f"got {type(solver).__name__}"
    )

Solver selection

solvers

Solver selection: the built-in SPG method plus any optax optimizer.

The shared projected-gradient solver in :mod:jaxfolio.optimizers.base accepts a solver spec rather than a hard-coded name. A spec is produced by :func:resolve_solver, which understands three spellings:

• "spg" — the built-in spectral projected-gradient method (default).
• any optax optimizer by name — "adam", "adamw", "sgd", "rmsprop", "lion", ... Resolved against :mod:optax and then :mod:optax.contrib, so anything the installed optax ships is reachable.
• an optax factory callable — optax.adamw, optax.contrib.adopt, or a module-level function of your own that returns a GradientTransformation (the escape hatch for optax.chain(...) compositions).

Hyperparameters go in solver_options ({"weight_decay": 1e-3}), not in a pre-built transformation — see the cache note below.

Cache stability

The solver spec is a static argument of the jit-cached kernel, so it has to be hashable and compare equal across separately constructed instances. Name strings and module-level factories satisfy that: optax.adamw is the same object every time you look it up, and :class:SolverSpec is a plain NamedTuple of hashable fields. A pre-built transformation (optax.adam(1e-2)) does not: it is a fresh tuple of closures on every call, hashed by identity, so passing one would recompile the kernel on every single solve. :func:resolve_solver therefore rejects pre-built transformations and points at the factory spelling.

SolverSpec

Bases: NamedTuple

A hashable, equality-stable solver key (safe as a jit static argument).

Attributes:

Name	Type	Description
kind	str	"spg" for the built-in spectral projected gradient, "optax" for everything else.
factory	Callable[..., Any] | None	The optax factory to call, e.g. optax.adamw. None for "spg".
name	str	Display label used in error messages and diagnostics.
options	tuple[tuple[str, Any], ...]	Extra keyword arguments for the factory, normalized to a sorted tuple of (key, value) pairs so the spec stays hashable.

resolve_solver

resolve_solver(
    solver: str | Callable[..., Any] | SolverSpec = SPG,
    solver_options: Mapping[str, Any] | None = None,
) -> SolverSpec

Normalize a solver spelling into a :class:SolverSpec.

Accepts "spg", an optax optimizer name, an optax factory callable, or an already-resolved :class:SolverSpec (the function is idempotent, so callers can normalize defensively). Raises :class:ValueError for an unknown name and :class:TypeError for an unusable object.

Source code in src/jaxfolio/solvers.py

def resolve_solver(
    solver: str | Callable[..., Any] | SolverSpec = SPG,
    solver_options: Mapping[str, Any] | None = None,
) -> SolverSpec:
    """Normalize a solver spelling into a :class:`SolverSpec`.

    Accepts ``"spg"``, an optax optimizer name, an optax factory callable, or an
    already-resolved :class:`SolverSpec` (the function is idempotent, so callers
    can normalize defensively). Raises :class:`ValueError` for an unknown name and
    :class:`TypeError` for an unusable object.
    """
    if isinstance(solver, SolverSpec):
        return solver

    options = _normalize_options(solver_options)

    if solver == SPG:
        if options:
            raise ValueError(
                "solver_options apply to optax solvers only; 'spg' is tuned by "
                "`learning_rate` (its initial Barzilai-Borwein step) and `tol`"
            )
        return _SPG_SPEC

    if isinstance(solver, str):
        factory = _lookup(solver)
        return SolverSpec("optax", factory, solver, options)

    # ``optax.nadam``/``nadamw`` are functools.partial objects, so gate on
    # callable() rather than inspect.isfunction().
    if callable(solver):
        name = getattr(solver, "__name__", None) or repr(solver)
        return SolverSpec("optax", solver, name, options)

    if hasattr(solver, "init") and hasattr(solver, "update"):
        raise TypeError(
            "pass the optax *factory* rather than a pre-built GradientTransformation "
            "(e.g. solver=optax.adamw with solver_options={'weight_decay': 1e-3}, not "
            "solver=optax.adamw(1e-3)). A pre-built transformation is a fresh object on "
            "every call, so it would recompile the solver kernel on every solve. For an "
            "optax.chain(...) composition, wrap it in a module-level factory function "
            "`def my_solver(learning_rate): ...` and pass that."
        )

    raise TypeError(
        f"solver must be 'spg', an optax optimizer name, or an optax factory callable; "
        f"got {type(solver).__name__}"
    )

build_optimizer

build_optimizer(
    spec: SolverSpec, learning_rate: float | None = None
) -> GradientTransformation

Instantiate the optax transformation described by spec.

learning_rate=None falls back to :data:DEFAULT_LEARNING_RATE (1e-2) — optax factories have no default of their own. Factories that take no learning rate (optax.contrib.dowg, cocob, ...) are called without one, and an explicitly supplied rate is then an error rather than silently ignored.

Source code in src/jaxfolio/solvers.py

def build_optimizer(
    spec: SolverSpec,
    learning_rate: float | None = None,
) -> optax.GradientTransformation:
    """Instantiate the optax transformation described by ``spec``.

    ``learning_rate=None`` falls back to :data:`DEFAULT_LEARNING_RATE` (``1e-2``)
    — optax factories have no default of their own. Factories that take no
    learning rate (``optax.contrib.dowg``, ``cocob``, ...) are called without one,
    and an explicitly supplied rate is then an error rather than silently ignored.
    """
    if spec.kind != "optax" or spec.factory is None:
        raise ValueError(f"cannot build an optax optimizer from solver kind {spec.kind!r}")

    kwargs: dict[str, Any] = dict(spec.options)
    try:
        params = inspect.signature(spec.factory).parameters
        takes_lr = "learning_rate" in params
    except (ValueError, TypeError):  # builtins / C-implemented callables
        takes_lr = True

    if takes_lr:
        # Always by keyword: contrib factories order their arguments differently.
        kwargs["learning_rate"] = DEFAULT_LEARNING_RATE if learning_rate is None else learning_rate
    elif learning_rate is not None:
        raise ValueError(f"solver {spec.name!r} does not take a learning_rate; leave it as None")

    try:
        tx = spec.factory(**kwargs)
    except TypeError as exc:
        given = dict(spec.options)
        raise ValueError(f"solver {spec.name!r} rejected solver_options {given}: {exc}") from exc

    if not (hasattr(tx, "init") and hasattr(tx, "update")):
        raise ValueError(
            f"solver {spec.name!r} did not return an optax GradientTransformation "
            f"(got {type(tx).__name__}); it is probably not an optimizer"
        )

    _reject_line_search(tx, spec.name)
    return tx

available_solvers

available_solvers() -> tuple[str, ...]

Return the solver names that can be passed as solver=....

("spg",) followed by the sorted names of every optax and optax.contrib factory that takes a learning_rate. This is a discovery aid for humans and for error messages — resolution itself is more permissive, so optimizers with their own adaptive step (dowg, cocob) still work even though they are not listed here.

Source code in src/jaxfolio/solvers.py

def available_solvers() -> tuple[str, ...]:
    """Return the solver names that can be passed as ``solver=...``.

    ``("spg",)`` followed by the sorted names of every optax and ``optax.contrib``
    factory that takes a ``learning_rate``. This is a discovery aid for humans and
    for error messages — resolution itself is more permissive, so optimizers with
    their own adaptive step (``dowg``, ``cocob``) still work even though they are
    not listed here.
    """
    names: set[str] = set()
    for module in (optax, optax.contrib):
        for name in dir(module):
            if name.startswith("_") or name.startswith("scale_by_"):
                continue
            candidate = getattr(module, name, None)
            if not callable(candidate) or inspect.isclass(candidate):
                continue
            try:
                params = inspect.signature(candidate).parameters
            except (ValueError, TypeError):
                continue
            if "learning_rate" in params:
                names.add(name)
    return (SPG, *sorted(names))

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)
        # A Polars panel carries its temporal key as an explicit column.
        # User weight functions see only investable assets, matching the
        # pre-Polars contract where dates lived outside DataFrame columns.
        asset_frame = returns.select(names) if hasattr(returns, "select") else returns
        w = _coerce_weights(weight_fn(asset_frame), 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, including config.solver / config.solver_options — so a custom objective can be minimized with any optax optimizer.

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,
    including ``config.solver`` / ``config.solver_options`` — so a custom
    objective can be minimized with any optax optimizer.
    """
    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(),
            constraints=cfg.constraints,
            assets=names,
        )
        w, info = tk.solve_projected_gradient(
            lambda w: objective_fn(w, ctx),
            tk.equal_start(len(names)),
            projection,
            solver=cfg.solver_spec(),
            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,
    trajectory=None,
    attribution=None,
) -> 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.

trajectory is an optional (T, n) planned weight path for multi-period optimizers. It is attached verbatim — not cleaned by clean_eps, so the path keeps the solver's exact iterate while weights is tidied for display.

attribution is an optional :class:jaxfolio.attribution.SolverDuals. It is attached verbatim and a small JSON-serializable digest of it is merged into metadata, so consumers that only read metadata still see the convergence and binding-constraint summary.

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,
    trajectory=None,
    attribution=None,
) -> 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.

    ``trajectory`` is an optional ``(T, n)`` planned weight path for multi-period
    optimizers. It is attached verbatim — *not* cleaned by ``clean_eps``, so the
    path keeps the solver's exact iterate while ``weights`` is tidied for display.

    ``attribution`` is an optional :class:`jaxfolio.attribution.SolverDuals`. It is
    attached verbatim and a small JSON-serializable digest of it is merged into
    ``metadata``, so consumers that only read ``metadata`` still see the
    convergence and binding-constraint summary.
    """
    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

    meta = dict(metadata or {})
    if attribution is not None:
        meta.update(attribution.summary())

    return PortfolioResult(
        weights=w,
        assets=list(assets),
        method=method,
        expected_return=ann_mu,
        volatility=ann_vol,
        sharpe=sharpe,
        metadata=meta,
        trajectory=trajectory,
        attribution=attribution,
    )