Coverage for src/jquantstats/_portfolio_transform.py: 100%
67 statements
« prev ^ index » next coverage.py v7.17.1, created at 2026-11-03 15:52 +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
20if TYPE_CHECKING:
21 from ._cost_model import CostModel
22from ._truncate import resolve_bounds
25class PortfolioTransformMixin(_PortfolioMembers):
26 """Mixin providing range/lag/smoothing transforms and correlation for Portfolio."""
28 if TYPE_CHECKING:
30 @classmethod
31 def from_cash_position(
32 cls,
33 prices: pl.DataFrame,
34 cash_position: pl.DataFrame,
35 aum: float,
36 cost_per_unit: float = 0.0,
37 cost_bps: float = 0.0,
38 cost_model: CostModel | None = None,
39 annual_fee: float = 0.0,
40 compounding: bool = False,
41 realized_costs: pl.DataFrame | None = None,
42 ) -> Self:
43 """Create a Portfolio directly from cash positions aligned with prices."""
44 ...
46 # ── Shared construction helpers ────────────────────────────────────────────
48 def _rebuild(
49 self,
50 cash_position: pl.DataFrame,
51 prices: pl.DataFrame | None = None,
52 aum: float | None = None,
53 realized_costs: pl.DataFrame | None = None,
54 ) -> Self:
55 """Build a new Portfolio of the same concrete type from derived frames.
57 Every transform in this mixin ends the same way: hand the derived
58 frames back to ``from_cash_position`` while carrying ``aum``, both
59 cost parameters, the management fee and the ``compounding`` flag
60 across unchanged. Centralising that here keeps a new construction
61 parameter from having to be threaded through each transform
62 individually.
64 Realised costs are *not* carried by default: they were paid on the
65 original trades, and a transform that changes the positions trades
66 differently. Only `truncate`, which keeps the trades, passes them on.
68 Args:
69 cash_position: The derived cash-position frame.
70 prices: The derived price frame; defaults to the current prices,
71 which the lag and smoothing transforms leave untouched.
72 aum: Starting capital of the new portfolio; defaults to the
73 current ``aum``.
74 realized_costs: Realised costs of the new portfolio; defaults to
75 None.
77 Returns:
78 A new Portfolio of the same concrete type.
79 """
80 return type(self).from_cash_position(
81 prices=self.prices if prices is None else prices,
82 cash_position=cash_position,
83 aum=self.aum if aum is None else aum,
84 cost_per_unit=self.cost_per_unit,
85 cost_bps=self.cost_bps,
86 annual_fee=self.annual_fee,
87 compounding=self.compounding,
88 realized_costs=realized_costs,
89 )
91 @property
92 def _numeric_assets(self) -> list[str]:
93 """Names of the numeric asset columns in the cash-position frame.
95 Excludes ``'date'`` and any non-numeric column, so column-wise
96 transforms touch only the asset series.
98 Returns:
99 The numeric asset column names, in frame order.
100 """
101 return [c for c in self.cashposition.columns if c != "date" and self.cashposition[c].dtype.is_numeric()]
103 @staticmethod
104 def _date_range_mask(
105 start: date | datetime | str | int | None,
106 end: date | datetime | str | int | None,
107 ) -> pl.Expr:
108 """Build the inclusive ``[start, end]`` filter over the ``'date'`` column.
110 Args:
111 start: Optional inclusive lower bound; no lower bound when None.
112 end: Optional inclusive upper bound; no upper bound when None.
114 Returns:
115 A boolean Polars expression, ``lit(True)`` when both bounds are None.
116 """
117 cond = pl.lit(True)
118 if start is not None:
119 cond = cond & (pl.col("date") >= pl.lit(start))
120 if end is not None:
121 cond = cond & (pl.col("date") <= pl.lit(end))
122 return cond
124 @staticmethod
125 def _row_slice_bounds(start: int | None, end: int | None, height: int) -> tuple[int, int]:
126 """Resolve integer row bounds into an ``(offset, length)`` slice.
128 Bounds arrive already validated by `resolve_bounds`, so this only
129 substitutes the open-ended defaults.
131 Args:
132 start: Optional inclusive lower row index; 0 when None.
133 end: Optional inclusive upper row index; the last row when None.
134 height: Row count of the frame being sliced.
136 Returns:
137 The ``(offset, length)`` pair to pass to ``DataFrame.slice``.
138 ``length`` is clamped at 0 so an inverted range yields an empty
139 frame rather than a negative slice.
140 """
141 row_start = start if start is not None else 0
142 row_end = end + 1 if end is not None else height
143 return row_start, max(0, row_end - row_start)
145 # ── Transforms ─────────────────────────────────────────────────────────────
147 def truncate(
148 self,
149 start: date | datetime | str | int | None = None,
150 end: date | datetime | str | int | None = None,
151 ) -> Self:
152 """Return a new Portfolio truncated to the inclusive [start, end] range.
154 **The bound type picks the axis.** A ``date``, ``datetime`` or ISO-8601
155 string is compared against the ``'date'`` column; an ``int`` is a 0-based
156 row index and slices positionally. Row indices work on a dated portfolio
157 too — ``truncate(start=10)`` drops the first ten rows — but the two kinds
158 cannot be combined in one call.
160 A portfolio with no ``'date'`` column accepts row indices only.
162 In all cases the ``aum`` value is preserved.
164 Args:
165 start: Optional inclusive lower bound. A ``date``/``datetime``, an
166 ISO-8601 string, or an ``int`` row index; ``int`` only when there
167 is no ``'date'`` column.
168 end: Optional inclusive upper bound. Same type rules as ``start``,
169 and must address the same axis.
171 Returns:
172 A new Portfolio instance with prices and cash positions filtered
173 to the specified range.
175 Raises:
176 IntegerIndexBoundError: When the portfolio has no ``'date'`` column
177 and a bound is not an ``int``.
178 InvalidTruncateBoundError: When a bound is of an unsupported type,
179 or is a string that is not ISO-8601.
180 MixedTruncateBoundsError: When one bound is a row index and the
181 other a date.
182 """
183 mode, lower, upper = resolve_bounds(start, end, temporal="date" in self.prices.columns)
185 rows = self.prices.with_row_index("_row")
186 if mode == "dates":
187 rows = rows.filter(self._date_range_mask(lower, upper))
188 else:
189 # "none" resolves to a full-width slice, so it needs no separate branch.
190 # The casts record what resolve_bounds guarantees for these modes but
191 # cannot express in its return type.
192 offset, length = self._row_slice_bounds(
193 cast("int | None", lower), cast("int | None", upper), self.prices.height
194 )
195 rows = rows.slice(offset, length)
196 kept = rows["_row"]
197 pr = self.prices[kept]
198 cp = self.cashposition[kept]
199 if kept.is_empty():
200 return self._rebuild(prices=pr, cash_position=cp)
202 # A compounding book sized its positions off the NAV it had reached, so
203 # the window starts from that NAV rather than from the original aum.
204 aum = float(self.nav_accumulated["NAV_accumulated"][kept[0]]) if self.compounding else self.aum
205 costs = None
206 if self.realized_costs is not None:
207 # The first kept row's trades happened before the window opens: they
208 # are already inside its starting capital.
209 is_first = pl.int_range(pl.len()) == 0
210 costs = self.realized_costs[kept]
211 costs = costs.with_columns(
212 pl.when(is_first).then(0.0).otherwise(pl.col(c)).alias(c) for c in self.assets if c in costs.columns
213 )
214 return self._rebuild(prices=pr, cash_position=cp, aum=aum, realized_costs=costs)
216 def lag(self, n: int) -> Self:
217 """Return a new Portfolio with cash positions lagged by ``n`` steps.
219 This method shifts the numeric asset columns in the cashposition
220 DataFrame by ``n`` rows, preserving the ``'date'`` column and any
221 non-numeric columns unchanged. Positive ``n`` delays weights (moves
222 them down); negative ``n`` leads them (moves them up); ``n == 0``
223 returns the current portfolio unchanged.
225 Notes:
226 Missing values introduced by the shift are left as nulls;
227 downstream profit computation already guards and treats nulls as
228 zero when multiplying by returns.
230 Args:
231 n: Number of rows to shift (can be negative, zero, or positive).
233 Returns:
234 A new Portfolio instance with lagged cash positions and the same
235 prices/AUM as the original.
236 """
237 if not isinstance(n, int):
238 raise TypeError
239 if n == 0:
240 return self
242 cp_lagged = self.cashposition.with_columns(pl.col(c).shift(n) for c in self._numeric_assets)
243 return self._rebuild(cash_position=cp_lagged)
245 def smoothed_holding(self, n: int) -> Self:
246 """Return a new Portfolio with cash positions smoothed by a rolling mean.
248 Applies a trailing window average over the last ``n`` steps for each
249 numeric asset column (excluding ``'date'``). The window length is
250 ``n + 1`` so that:
252 - n=0 returns the original weights (no smoothing),
253 - n=1 averages the current and previous weights,
254 - n=k averages the current and last k weights.
256 Args:
257 n: Non-negative integer specifying how many previous steps to
258 include.
260 Returns:
261 A new Portfolio with smoothed cash positions and the same
262 prices/AUM.
263 """
264 if not isinstance(n, int):
265 raise TypeError(f"n must be an integer, got {type(n).__name__}") # noqa: TRY003
266 if n < 0:
267 raise ValueError(f"n must be a non-negative integer, got {n}") # noqa: TRY003
268 if n == 0:
269 return self
271 window = n + 1
272 cp_smoothed = self.cashposition.with_columns(
273 pl.col(c).rolling_mean(window_size=window, min_samples=1).alias(c) for c in self._numeric_assets
274 )
275 return self._rebuild(cash_position=cp_smoothed)
277 # ── Utility ────────────────────────────────────────────────────────────────
279 def correlation(self, frame: pl.DataFrame, name: str = "portfolio") -> pl.DataFrame:
280 """Compute a correlation matrix of asset returns plus the portfolio.
282 Computes percentage changes for all numeric columns in ``frame``,
283 appends the portfolio profit series under the provided ``name``, and
284 returns the Pearson correlation matrix across all numeric columns.
286 Args:
287 frame: A Polars DataFrame containing at least the asset price
288 columns (and a date column which will be ignored if
289 non-numeric).
290 name: The column name to use when adding the portfolio profit
291 series to the input frame.
293 Returns:
294 A square Polars DataFrame where each cell is the correlation
295 between a pair of series (values in [-1, 1]).
296 """
297 p = frame.with_columns(cs.by_dtype(pl.Float32, pl.Float64).pct_change())
298 p = p.with_columns(pl.Series(name, self.profit["profit"]))
299 corr_matrix = p.select(cs.numeric()).fill_null(0.0).corr()
300 return corr_matrix