Coverage for src/jquantstats/_portfolio_transform.py: 100%
58 statements
« prev ^ index » next coverage.py v7.16.1, created at 2026-09-23 04:11 +0000
1"""Range/lag/smoothing transforms and correlation mixin for Portfolio.
3`PortfolioTransformMixin` groups the methods that derive a *new* Portfolio
4from an existing one (`truncate`, `lag`, `smoothed_holding`) plus the
5`correlation` utility. New portfolios are built through
6``type(self).from_cash_position`` so the transforms inherit the standard
7construction path and return the concrete ``Self`` type.
8"""
10from __future__ import annotations
12from datetime import date, datetime
13from typing import TYPE_CHECKING, Self, cast
15import polars as pl
16import polars.selectors as cs
18from ._portfolio_base import _PortfolioMembers
19from ._truncate import resolve_bounds
22class PortfolioTransformMixin(_PortfolioMembers):
23 """Mixin providing range/lag/smoothing transforms and correlation for Portfolio."""
25 if TYPE_CHECKING:
27 @classmethod
28 def from_cash_position(
29 cls,
30 prices: pl.DataFrame,
31 cash_position: pl.DataFrame,
32 aum: float,
33 cost_per_unit: float = 0.0,
34 cost_bps: float = 0.0,
35 ) -> Self:
36 """Create a Portfolio directly from cash positions aligned with prices."""
37 ...
39 # ── Shared construction helpers ────────────────────────────────────────────
41 def _rebuild(self, cash_position: pl.DataFrame, prices: pl.DataFrame | None = None) -> Self:
42 """Build a new Portfolio of the same concrete type from derived frames.
44 Every transform in this mixin ends the same way: hand the derived
45 frames back to ``from_cash_position`` while carrying ``aum`` and both
46 cost parameters across unchanged. Centralising that here keeps a new
47 construction parameter from having to be threaded through each
48 transform individually.
50 Args:
51 cash_position: The derived cash-position frame.
52 prices: The derived price frame; defaults to the current prices,
53 which the lag and smoothing transforms leave untouched.
55 Returns:
56 A new Portfolio of the same concrete type.
57 """
58 return type(self).from_cash_position(
59 prices=self.prices if prices is None else prices,
60 cash_position=cash_position,
61 aum=self.aum,
62 cost_per_unit=self.cost_per_unit,
63 cost_bps=self.cost_bps,
64 )
66 @property
67 def _numeric_assets(self) -> list[str]:
68 """Names of the numeric asset columns in the cash-position frame.
70 Excludes ``'date'`` and any non-numeric column, so column-wise
71 transforms touch only the asset series.
73 Returns:
74 The numeric asset column names, in frame order.
75 """
76 return [c for c in self.cashposition.columns if c != "date" and self.cashposition[c].dtype.is_numeric()]
78 @staticmethod
79 def _date_range_mask(
80 start: date | datetime | str | int | None,
81 end: date | datetime | str | int | None,
82 ) -> pl.Expr:
83 """Build the inclusive ``[start, end]`` filter over the ``'date'`` column.
85 Args:
86 start: Optional inclusive lower bound; no lower bound when None.
87 end: Optional inclusive upper bound; no upper bound when None.
89 Returns:
90 A boolean Polars expression, ``lit(True)`` when both bounds are None.
91 """
92 cond = pl.lit(True)
93 if start is not None:
94 cond = cond & (pl.col("date") >= pl.lit(start))
95 if end is not None:
96 cond = cond & (pl.col("date") <= pl.lit(end))
97 return cond
99 @staticmethod
100 def _row_slice_bounds(start: int | None, end: int | None, height: int) -> tuple[int, int]:
101 """Resolve integer row bounds into an ``(offset, length)`` slice.
103 Bounds arrive already validated by `resolve_bounds`, so this only
104 substitutes the open-ended defaults.
106 Args:
107 start: Optional inclusive lower row index; 0 when None.
108 end: Optional inclusive upper row index; the last row when None.
109 height: Row count of the frame being sliced.
111 Returns:
112 The ``(offset, length)`` pair to pass to ``DataFrame.slice``.
113 ``length`` is clamped at 0 so an inverted range yields an empty
114 frame rather than a negative slice.
115 """
116 row_start = start if start is not None else 0
117 row_end = end + 1 if end is not None else height
118 return row_start, max(0, row_end - row_start)
120 # ── Transforms ─────────────────────────────────────────────────────────────
122 def truncate(
123 self,
124 start: date | datetime | str | int | None = None,
125 end: date | datetime | str | int | None = None,
126 ) -> Self:
127 """Return a new Portfolio truncated to the inclusive [start, end] range.
129 **The bound type picks the axis.** A ``date``, ``datetime`` or ISO-8601
130 string is compared against the ``'date'`` column; an ``int`` is a 0-based
131 row index and slices positionally. Row indices work on a dated portfolio
132 too — ``truncate(start=10)`` drops the first ten rows — but the two kinds
133 cannot be combined in one call.
135 A portfolio with no ``'date'`` column accepts row indices only.
137 In all cases the ``aum`` value is preserved.
139 Args:
140 start: Optional inclusive lower bound. A ``date``/``datetime``, an
141 ISO-8601 string, or an ``int`` row index; ``int`` only when there
142 is no ``'date'`` column.
143 end: Optional inclusive upper bound. Same type rules as ``start``,
144 and must address the same axis.
146 Returns:
147 A new Portfolio instance with prices and cash positions filtered
148 to the specified range.
150 Raises:
151 IntegerIndexBoundError: When the portfolio has no ``'date'`` column
152 and a bound is not an ``int``.
153 InvalidTruncateBoundError: When a bound is of an unsupported type,
154 or is a string that is not ISO-8601.
155 MixedTruncateBoundsError: When one bound is a row index and the
156 other a date.
157 """
158 mode, lower, upper = resolve_bounds(start, end, temporal="date" in self.prices.columns)
160 if mode == "dates":
161 cond = self._date_range_mask(lower, upper)
162 pr = self.prices.filter(cond)
163 cp = self.cashposition.filter(cond)
164 else:
165 # "none" resolves to a full-width slice, so it needs no separate branch.
166 # The casts record what resolve_bounds guarantees for these modes but
167 # cannot express in its return type.
168 offset, length = self._row_slice_bounds(
169 cast("int | None", lower), cast("int | None", upper), self.prices.height
170 )
171 pr = self.prices.slice(offset, length)
172 cp = self.cashposition.slice(offset, length)
173 return self._rebuild(prices=pr, cash_position=cp)
175 def lag(self, n: int) -> Self:
176 """Return a new Portfolio with cash positions lagged by ``n`` steps.
178 This method shifts the numeric asset columns in the cashposition
179 DataFrame by ``n`` rows, preserving the ``'date'`` column and any
180 non-numeric columns unchanged. Positive ``n`` delays weights (moves
181 them down); negative ``n`` leads them (moves them up); ``n == 0``
182 returns the current portfolio unchanged.
184 Notes:
185 Missing values introduced by the shift are left as nulls;
186 downstream profit computation already guards and treats nulls as
187 zero when multiplying by returns.
189 Args:
190 n: Number of rows to shift (can be negative, zero, or positive).
192 Returns:
193 A new Portfolio instance with lagged cash positions and the same
194 prices/AUM as the original.
195 """
196 if not isinstance(n, int):
197 raise TypeError
198 if n == 0:
199 return self
201 cp_lagged = self.cashposition.with_columns(pl.col(c).shift(n) for c in self._numeric_assets)
202 return self._rebuild(cash_position=cp_lagged)
204 def smoothed_holding(self, n: int) -> Self:
205 """Return a new Portfolio with cash positions smoothed by a rolling mean.
207 Applies a trailing window average over the last ``n`` steps for each
208 numeric asset column (excluding ``'date'``). The window length is
209 ``n + 1`` so that:
211 - n=0 returns the original weights (no smoothing),
212 - n=1 averages the current and previous weights,
213 - n=k averages the current and last k weights.
215 Args:
216 n: Non-negative integer specifying how many previous steps to
217 include.
219 Returns:
220 A new Portfolio with smoothed cash positions and the same
221 prices/AUM.
222 """
223 if not isinstance(n, int):
224 raise TypeError(f"n must be an integer, got {type(n).__name__}") # noqa: TRY003
225 if n < 0:
226 raise ValueError(f"n must be a non-negative integer, got {n}") # noqa: TRY003
227 if n == 0:
228 return self
230 window = n + 1
231 cp_smoothed = self.cashposition.with_columns(
232 pl.col(c).rolling_mean(window_size=window, min_samples=1).alias(c) for c in self._numeric_assets
233 )
234 return self._rebuild(cash_position=cp_smoothed)
236 # ── Utility ────────────────────────────────────────────────────────────────
238 def correlation(self, frame: pl.DataFrame, name: str = "portfolio") -> pl.DataFrame:
239 """Compute a correlation matrix of asset returns plus the portfolio.
241 Computes percentage changes for all numeric columns in ``frame``,
242 appends the portfolio profit series under the provided ``name``, and
243 returns the Pearson correlation matrix across all numeric columns.
245 Args:
246 frame: A Polars DataFrame containing at least the asset price
247 columns (and a date column which will be ignored if
248 non-numeric).
249 name: The column name to use when adding the portfolio profit
250 series to the input frame.
252 Returns:
253 A square Polars DataFrame where each cell is the correlation
254 between a pair of series (values in [-1, 1]).
255 """
256 p = frame.with_columns(cs.by_dtype(pl.Float32, pl.Float64).pct_change())
257 p = p.with_columns(pl.Series(name, self.profit["profit"]))
258 corr_matrix = p.select(cs.numeric()).fill_null(0.0).corr()
259 return corr_matrix