Coverage for src/jquantstats/_portfolio_transform.py: 100%
61 statements
« prev ^ index » next coverage.py v7.15.3, created at 2026-08-06 04: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
15import polars as pl
16import polars.selectors as cs
18from ._portfolio_base import _PortfolioMembers
19from .exceptions import IntegerIndexBoundError
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(
101 start: date | datetime | str | int | None,
102 end: date | datetime | str | int | None,
103 height: int,
104 ) -> tuple[int, int]:
105 """Resolve integer row bounds into a ``(offset, length)`` slice.
107 Used when the portfolio has no ``'date'`` column and truncation falls
108 back to 0-based row indexing.
110 Args:
111 start: Optional inclusive lower row index; 0 when None.
112 end: Optional inclusive upper row index; the last row when None.
113 height: Row count of the frame being sliced.
115 Returns:
116 The ``(offset, length)`` pair to pass to ``DataFrame.slice``.
117 ``length`` is clamped at 0 so an inverted range yields an empty
118 frame rather than a negative slice.
120 Raises:
121 IntegerIndexBoundError: When a supplied bound is not an integer.
122 """
123 if start is not None and not isinstance(start, int):
124 raise IntegerIndexBoundError("start", type(start).__name__)
125 if end is not None and not isinstance(end, int):
126 raise IntegerIndexBoundError("end", type(end).__name__)
127 row_start = int(start) if start is not None else 0
128 row_end = int(end) + 1 if end is not None else height
129 return row_start, max(0, row_end - row_start)
131 # ── Transforms ─────────────────────────────────────────────────────────────
133 def truncate(
134 self,
135 start: date | datetime | str | int | None = None,
136 end: date | datetime | str | int | None = None,
137 ) -> Self:
138 """Return a new Portfolio truncated to the inclusive [start, end] range.
140 When a ``'date'`` column is present in both prices and cash positions,
141 truncation is performed by comparing the ``'date'`` column against
142 ``start`` and ``end`` (which should be date/datetime values or strings
143 parseable by Polars).
145 When the ``'date'`` column is absent, integer-based row slicing is
146 used instead. In this case ``start`` and ``end`` must be non-negative
147 integers representing 0-based row indices. Passing non-integer bounds
148 to an integer-indexed portfolio raises `TypeError`.
150 In all cases the ``aum`` value is preserved.
152 Args:
153 start: Optional lower bound (inclusive). A date/datetime or
154 Polars-parseable string when a ``'date'`` column exists; a
155 non-negative int row index when the data has no ``'date'``
156 column.
157 end: Optional upper bound (inclusive). Same type rules as
158 ``start``.
160 Returns:
161 A new Portfolio instance with prices and cash positions filtered
162 to the specified range.
164 Raises:
165 TypeError: When the portfolio has no ``'date'`` column and a
166 non-integer bound is supplied.
167 """
168 if "date" in self.prices.columns:
169 cond = self._date_range_mask(start, end)
170 pr = self.prices.filter(cond)
171 cp = self.cashposition.filter(cond)
172 else:
173 offset, length = self._row_slice_bounds(start, end, self.prices.height)
174 pr = self.prices.slice(offset, length)
175 cp = self.cashposition.slice(offset, length)
176 return self._rebuild(prices=pr, cash_position=cp)
178 def lag(self, n: int) -> Self:
179 """Return a new Portfolio with cash positions lagged by ``n`` steps.
181 This method shifts the numeric asset columns in the cashposition
182 DataFrame by ``n`` rows, preserving the ``'date'`` column and any
183 non-numeric columns unchanged. Positive ``n`` delays weights (moves
184 them down); negative ``n`` leads them (moves them up); ``n == 0``
185 returns the current portfolio unchanged.
187 Notes:
188 Missing values introduced by the shift are left as nulls;
189 downstream profit computation already guards and treats nulls as
190 zero when multiplying by returns.
192 Args:
193 n: Number of rows to shift (can be negative, zero, or positive).
195 Returns:
196 A new Portfolio instance with lagged cash positions and the same
197 prices/AUM as the original.
198 """
199 if not isinstance(n, int):
200 raise TypeError
201 if n == 0:
202 return self
204 cp_lagged = self.cashposition.with_columns(pl.col(c).shift(n) for c in self._numeric_assets)
205 return self._rebuild(cash_position=cp_lagged)
207 def smoothed_holding(self, n: int) -> Self:
208 """Return a new Portfolio with cash positions smoothed by a rolling mean.
210 Applies a trailing window average over the last ``n`` steps for each
211 numeric asset column (excluding ``'date'``). The window length is
212 ``n + 1`` so that:
214 - n=0 returns the original weights (no smoothing),
215 - n=1 averages the current and previous weights,
216 - n=k averages the current and last k weights.
218 Args:
219 n: Non-negative integer specifying how many previous steps to
220 include.
222 Returns:
223 A new Portfolio with smoothed cash positions and the same
224 prices/AUM.
225 """
226 if not isinstance(n, int):
227 raise TypeError(f"n must be an integer, got {type(n).__name__}") # noqa: TRY003
228 if n < 0:
229 raise ValueError(f"n must be a non-negative integer, got {n}") # noqa: TRY003
230 if n == 0:
231 return self
233 window = n + 1
234 cp_smoothed = self.cashposition.with_columns(
235 pl.col(c).rolling_mean(window_size=window, min_samples=1).alias(c) for c in self._numeric_assets
236 )
237 return self._rebuild(cash_position=cp_smoothed)
239 # ── Utility ────────────────────────────────────────────────────────────────
241 def correlation(self, frame: pl.DataFrame, name: str = "portfolio") -> pl.DataFrame:
242 """Compute a correlation matrix of asset returns plus the portfolio.
244 Computes percentage changes for all numeric columns in ``frame``,
245 appends the portfolio profit series under the provided ``name``, and
246 returns the Pearson correlation matrix across all numeric columns.
248 Args:
249 frame: A Polars DataFrame containing at least the asset price
250 columns (and a date column which will be ignored if
251 non-numeric).
252 name: The column name to use when adding the portfolio profit
253 series to the input frame.
255 Returns:
256 A square Polars DataFrame where each cell is the correlation
257 between a pair of series (values in [-1, 1]).
258 """
259 p = frame.with_columns(cs.by_dtype(pl.Float32, pl.Float64).pct_change())
260 p = p.with_columns(pl.Series(name, self.profit["profit"]))
261 corr_matrix = p.select(cs.numeric()).fill_null(0.0).corr()
262 return corr_matrix