Coverage for src/bartz/stochtree/_preprocess.py: 93%

1# bartz/src/bartz/stochtree/_preprocess.py 

2# 

3# Copyright (c) 2026, The Bartz Contributors 

4# 

5# This file is part of bartz. 

6# 

7# Permission is hereby granted, free of charge, to any person obtaining a copy 

8# of this software and associated documentation files (the "Software"), to deal 

9# in the Software without restriction, including without limitation the rights 

10# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 

11# copies of the Software, and to permit persons to whom the Software is 

12# furnished to do so, subject to the following conditions: 

13# 

14# The above copyright notice and this permission notice shall be included in all 

15# copies or substantial portions of the Software. 

16# 

17# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 

18# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 

19# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 

20# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 

21# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 

22# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 

23# SOFTWARE. 

24 

25"""Auto-preprocessing of covariates for the stochtree-compatible BART interface. 

26 

27Two parallel implementations are provided, `PandasPreprocessor` and 

28`PolarsPreprocessor`, each handling the corresponding dataframe library. Both 

29classes have the same interface:: 

30 

31 pp = PandasPreprocessor() # or PolarsPreprocessor() 

32 varprob = pp.fit(X_train, variable_weights=w) 

33 x_train = pp.transform(X_train) 

34 x_new = pp.transform(X_new) # at prediction time 

35 

36`fit` records the per-column encoding and returns the variable weights expanded 

37to match the new column count (or `None`); `transform` returns the 

38post-processing covariate matrix as a 2-D numpy float32 array (rows=observations, 

39columns=expanded features). 

40 

41Per-column handling: 

42 

43- ordered categorical (pandas ordered `Categorical`): ordinal encoded into a 

44 single integer-valued column, with the declared category order giving the 

45 integer mapping. polars has no ordered categorical dtype; pass an integer 

46 column for ordinal encoding. 

47- unordered categorical (pandas unordered `Categorical`, polars `Enum`): one-hot 

48 encoded into one binary column per declared category. A polars `Enum` 

49 round-trips to a pandas *unordered* `Categorical`, so the two are treated 

50 identically. 

51- boolean: cast to ``{0.0, 1.0}``, single column. 

52- numeric (integer, unsigned, float): pass-through as float. 

53- anything else (strings, ``object``, datetime, polars `Categorical`, etc.): 

54 raises `ValueError`. polars `Categorical` has no reliable per-column category 

55 list (the categories live in a process-wide string cache shared across 

56 columns), so it must be cast to an `Enum` (one-hot) or an integer (ordinal) 

57 first. 

58 

59When a single original column expands into ``k`` output columns (one-hot), the 

60original `variable_weights` entry for that column is split evenly across the 

61``k`` expansions, preserving each original variable's total splitting budget 

62(matching stochtree's `bart.py` behavior). 

63 

64Unknown category values encountered during `transform` raise `ValueError`. 

65""" 

66 

67from collections.abc import Sequence (empty)

68from dataclasses import dataclass (empty)

69from typing import Any, Literal, TypeAlias, overload (empty)

70 

71import numpy as np (empty)

72from jaxtyping import Float32, Shaped (empty)

73 

74# Not `numpy.typing.ArrayLike`: that is a PEP 695 type alias since numpy 2.5, 

75# which jaxtyping cannot subscript. 

76from bartz.mcmcstep._state import ArrayLike (empty)

77 

78# Duck-typed stand-ins for the optional dataframe libraries. bartz does not 

79# depend on pandas or polars at runtime, so we cannot reference their real 

80# classes here; these aliases resolve to `Any` but give the signatures below 

81# legible names. 

82DataFrame: TypeAlias = Any # a pandas or polars DataFrame (empty)

83Series: TypeAlias = Any # a pandas or polars Series (empty)

84PolarsModule: TypeAlias = Any # the polars top-level module (empty)

85 

86_UNSEEN_PREVIEW = 10 (empty)

87 

88ColumnKind: TypeAlias = Literal['numeric', 'bool', 'ordered_cat', 'unordered_cat'] (empty)

89 

90 

91@dataclass(frozen=True) (empty)

92class _ColumnSpec: (empty)

93 """Per-original-column fitted state.""" 

94 

95 kind: ColumnKind 

96 """Encoding to apply to the column.""" (empty)

97 

98 name: str 

99 """Original column name (for error messages).""" (empty)

100 

101 categories: tuple[Any, ...] | None = None (empty)

102 """Declared category list for ordered_cat / unordered_cat.""" (empty)

103 

104 @property (empty)

105 def width(self) -> int: (empty)

106 """Number of output columns this spec produces.""" 

107 if self.kind == 'unordered_cat': (empty)

108 assert self.categories is not None (empty)

109 return len(self.categories) (empty)

110 return 1 (empty)

111 

112 

113def _unseen_error(name: str, unseen: Sequence[Any], known: Sequence[Any]) -> ValueError: (empty)

114 """Build the error for category values absent from the fitted list.""" 

115 uniq = sorted({repr(v) for v in unseen}) (empty)

116 msg = ( (empty)

117 f'column {name!r}: {len(unseen)} value(s) at transform time are not in' 

118 f' the fitted category list; unseen sample: {uniq[:_UNSEEN_PREVIEW]};' 

119 f' known categories: {list(known)[:_UNSEEN_PREVIEW]}' 

120 ) 

121 return ValueError(msg) (empty)

122 

123 

124def _unsupported_dtype_error(name: str, dtype: object) -> ValueError: (empty)

125 """Build the error for a column whose dtype has no supported encoding.""" 

126 msg = ( (empty)

127 f'column {name!r} has unsupported dtype {dtype!r}; supported types are' 

128 ' numeric, boolean, pandas ordered/unordered Categorical, and polars' 

129 ' Enum. Convert strings, objects, datetimes, etc. to one of these (e.g.' 

130 ' an explicit Categorical / Enum) before fitting.' 

131 ) 

132 return ValueError(msg) (empty)

133 

134 

135def _polars_categorical_error(name: str) -> ValueError: (empty)

136 """Build the error rejecting a polars `Categorical` column.""" 

137 msg = ( (empty)

138 f'column {name!r} is a polars Categorical, which has no reliable' 

139 ' per-column category list (the categories live in a process-wide' 

140 ' string cache shared across columns). Cast it to a polars Enum with an' 

141 ' explicit category list (pl.Enum([...])) for one-hot encoding, or to an' 

142 ' integer column for ordinal encoding.' 

143 ) 

144 return ValueError(msg) (empty)

145 

146 

147def _ordinal_encode( (empty)

148 values: Shaped[np.ndarray, ' n'], categories: Sequence[Any], name: str 

149) -> Float32[np.ndarray, 'n 1']: 

150 """Map `values` to integer positions in `categories`; raise on unseen.""" 

151 table = {c: i for i, c in enumerate(categories)} (empty)

152 out = np.empty(len(values), dtype=np.float32) (empty)

153 unseen: list[Any] = [] (empty)

154 for i, v in enumerate(values): (empty)

155 code = table.get(v, -1) (empty)

156 if code < 0: 156 ↛ 157line 156 didn't jump to line 157 because the condition on line 156 was never true(empty)

157 unseen.append(v) 

158 else: 

159 out[i] = code (empty)

160 if unseen: 160 ↛ 161line 160 didn't jump to line 161 because the condition on line 160 was never true(empty)

161 raise _unseen_error(name, unseen, categories) 

162 return out[:, None] (empty)

163 

164 

165def _one_hot_encode( (empty)

166 values: Shaped[np.ndarray, ' n'], categories: Sequence[Any], name: str 

167) -> Float32[np.ndarray, 'n k']: 

168 """Build a ``(n, k)`` one-hot matrix using `categories` order; raise on unseen.""" 

169 table = {c: i for i, c in enumerate(categories)} (empty)

170 n = len(values) (empty)

171 k = len(categories) (empty)

172 out = np.zeros((n, k), dtype=np.float32) (empty)

173 unseen: list[Any] = [] (empty)

174 for i, v in enumerate(values): (empty)

175 code = table.get(v, -1) (empty)

176 if code < 0: (empty)

177 unseen.append(v) (empty)

178 else: 

179 out[i, code] = 1.0 (empty)

180 if unseen: (empty)

181 raise _unseen_error(name, unseen, categories) (empty)

182 return out (empty)

183 

184 

185def _polars_one_hot( (empty)

186 pl: PolarsModule, series: Series, categories: Sequence[Any], name: str 

187) -> Float32[np.ndarray, 'n k']: 

188 """Validate via cast to `pl.Enum(categories)` and one-hot via polars APIs. 

189 

190 Polars's `Enum` cast natively raises on any value not in `categories`, and 

191 `to_physical` returns the integer codes in the declared-category order. The 

192 `np.eye` index is the only numpy bit and is just an identity-matrix lookup; 

193 the categorical bookkeeping itself stays inside polars. 

194 """ 

195 cats = list(categories) (empty)

196 try: (empty)

197 coded = series.cast(pl.Enum(cats)) (empty)

198 except pl.exceptions.InvalidOperationError as exc: (empty)

199 # Identify the actual unseen values for a friendly error. Cast to String 

200 # first: the input column may itself be an Enum with different categories, 

201 # which would make a direct is_in(cats) fail trying to coerce the list. 

202 known = set(cats) (empty)

203 unseen = sorted( (empty)

204 { 

205 v 

206 for v in series.cast(pl.String).to_list() 

207 if v not in known and v is not None 

208 } 

209 ) 

210 raise _unseen_error(name, unseen, cats) from exc (empty)

211 if coded.null_count(): 211 ↛ 212line 211 didn't jump to line 212 because the condition on line 211 was never true(empty)

212 msg = f'column {name!r}: null values are not supported in categorical columns' 

213 raise ValueError(msg) 

214 codes = coded.to_physical().to_numpy() (empty)

215 return np.eye(len(cats), dtype=np.float32)[codes] (empty)

216 

217 

218def _expand_variable_weights( (empty)

219 weights: Shaped[ArrayLike, '...'], original_var_indices: Sequence[int], n_orig: int 

220) -> Float32[np.ndarray, ' p']: 

221 """Split each original weight evenly across its one-hot expansions.""" 

222 w = np.asarray(weights, dtype=np.float32) (empty)

223 if w.shape != (n_orig,): (empty)

224 msg = ( (empty)

225 f'variable_weights must have shape ({n_orig},) matching the number' 

226 f' of original columns; got {w.shape}' 

227 ) 

228 raise ValueError(msg) (empty)

229 if not original_var_indices: 229 ↛ 230line 229 didn't jump to line 230 because the condition on line 229 was never true(empty)

230 return np.empty((0,), dtype=np.float32) 

231 counts = np.bincount(np.asarray(original_var_indices), minlength=n_orig) (empty)

232 return np.array([w[j] / counts[j] for j in original_var_indices], dtype=np.float32) (empty)

233 

234 

235def _stack( (empty)

236 cols: Sequence[Float32[np.ndarray, 'n _']], n_rows: int 

237) -> Float32[np.ndarray, 'n p']: 

238 if not cols: 238 ↛ 239line 238 didn't jump to line 239 because the condition on line 238 was never true(empty)

239 return np.empty((n_rows, 0), dtype=np.float32) 

240 return np.concatenate(cols, axis=1) (empty)

241 

242 

243class _PreprocessorBase: (empty)

244 """Common state for `PandasPreprocessor` and `PolarsPreprocessor`.""" 

245 

246 _library: str = '' (empty)

247 """Top-level module prefix of the supported dataframe library.""" (empty)

248 

249 _fitted: bool = False (empty)

250 _specs: Sequence[_ColumnSpec] = () (empty)

251 _original_var_indices: Sequence[int] = () (empty)

252 

253 @property (empty)

254 def fitted(self) -> bool: (empty)

255 """Whether `fit` has been called.""" 

256 return self._fitted 

257 

258 @property (empty)

259 def n_original_columns(self) -> int: (empty)

260 """Number of columns in the dataframe given to `fit`.""" 

261 return len(self._specs) 

262 

263 @property (empty)

264 def n_processed_columns(self) -> int: (empty)

265 """Number of columns in the matrix returned by `transform`.""" 

266 return len(self._original_var_indices) (empty)

267 

268 @property (empty)

269 def original_var_indices(self) -> tuple[int, ...]: (empty)

270 """For each output column, the index of the original column it came from.""" 

271 return tuple(self._original_var_indices) (empty)

272 

273 @overload (empty)

274 def fit( (empty)

275 self, X: DataFrame, *, variable_weights: Shaped[ArrayLike, '...'] (empty)

276 ) -> Float32[np.ndarray, ' p']: ... (empty)

277 

278 @overload (empty)

279 def fit( (empty)

280 self, X: DataFrame, *, variable_weights: None = None (empty)

281 ) -> Float32[np.ndarray, ' p'] | None: ... (empty)

282 

283 def fit( (empty)

284 self, X: DataFrame, *, variable_weights: Shaped[ArrayLike, '...'] | None = None 

285 ) -> Float32[np.ndarray, ' p'] | None: 

286 """Record the per-column encoding and return the expanded variable weights. 

287 

288 Returns `None` when no weights are supplied and no column expands into 

289 several output columns, so the caller can fall back to the native 

290 uniform-weights path; otherwise returns the weights split across each 

291 original column's one-hot expansion. 

292 """ 

293 self._check_library(X) (empty)

294 specs: list[_ColumnSpec] = [] (empty)

295 original_var_indices: list[int] = [] (empty)

296 for orig_idx in range(X.shape[1]): (empty)

297 name, series = self._get_column(X, orig_idx) (empty)

298 spec = self._fit_column(series, str(name)) (empty)

299 specs.append(spec) (empty)

300 original_var_indices.extend([orig_idx] * spec.width) (empty)

301 self._specs = tuple(specs) (empty)

302 self._original_var_indices = tuple(original_var_indices) (empty)

303 self._fitted = True (empty)

304 expanded = len(set(original_var_indices)) != len(original_var_indices) (empty)

305 if variable_weights is None: (empty)

306 if not expanded: (empty)

307 return None (empty)

308 variable_weights = np.full(len(specs), 1.0 / len(specs)) (empty)

309 return _expand_variable_weights( (empty)

310 variable_weights, self._original_var_indices, len(self._specs) 

311 ) 

312 

313 def transform(self, X: DataFrame) -> Float32[np.ndarray, 'n p']: (empty)

314 """Apply the fitted transformation to a new dataframe.""" 

315 self._check_fitted() (empty)

316 self._check_library(X) (empty)

317 self._check_n_columns(X.shape[1]) (empty)

318 cols = [ (empty)

319 self._transform_column(self._get_column(X, orig_idx)[1], spec) 

320 for orig_idx, spec in enumerate(self._specs) 

321 ] 

322 return _stack(cols, X.shape[0]) (empty)

323 

324 def _check_fitted(self) -> None: (empty)

325 if not self._fitted: 325 ↛ 326line 325 didn't jump to line 326 because the condition on line 325 was never true(empty)

326 msg = 'preprocessor has not been fitted yet; call fit first' 

327 raise RuntimeError(msg) 

328 

329 def _check_n_columns(self, n_cols: int) -> None: (empty)

330 if n_cols != len(self._specs): (empty)

331 msg = ( (empty)

332 f'transform input has {n_cols} columns; preprocessor was fitted' 

333 f' on {len(self._specs)} columns' 

334 ) 

335 raise ValueError(msg) (empty)

336 

337 def _check_library(self, X: DataFrame) -> None: (empty)

338 module = type(X).__module__ (empty)

339 if not module.startswith(self._library): (empty)

340 msg = ( (empty)

341 f'this preprocessor handles {self._library} dataframes, but got' 

342 f' an object from {module!r}; fit and transform must use the same' 

343 ' dataframe library' 

344 ) 

345 raise TypeError(msg) (empty)

346 

347 @staticmethod (empty)

348 def _get_column(X: DataFrame, orig_idx: int) -> tuple[Any, Series]: (empty)

349 """Return the ``(name, series)`` of the column at position `orig_idx`.""" 

350 raise NotImplementedError 

351 

352 @staticmethod (empty)

353 def _fit_column(series: Series, name: str) -> _ColumnSpec: (empty)

354 """Inspect a column's dtype and return its encoding spec.""" 

355 raise NotImplementedError 

356 

357 @staticmethod (empty)

358 def _transform_column( (empty)

359 series: Series, spec: _ColumnSpec 

360 ) -> Float32[np.ndarray, 'n _']: 

361 """Encode a single column according to its fitted spec.""" 

362 raise NotImplementedError 

363 

364 

365class PandasPreprocessor(_PreprocessorBase): (empty)

366 """Stochtree-style covariate preprocessor for `pandas.DataFrame` inputs.""" 

367 

368 _library = 'pandas' (empty)

369 

370 @staticmethod (empty)

371 def _get_column(X: DataFrame, orig_idx: int) -> tuple[Any, Series]: (empty)

372 return X.columns[orig_idx], X.iloc[:, orig_idx] (empty)

373 

374 @staticmethod (empty)

375 def _fit_column(series: Series, name: str) -> _ColumnSpec: (empty)

376 import pandas as pd # noqa: PLC0415 # optional runtime dependency (empty)

377 

378 dt = series.dtype (empty)

379 if isinstance(dt, pd.CategoricalDtype): (empty)

380 cats = tuple(dt.categories) (empty)

381 kind: ColumnKind = 'ordered_cat' if dt.ordered else 'unordered_cat' (empty)

382 return _ColumnSpec(kind, name, categories=cats) (empty)

383 if pd.api.types.is_bool_dtype(dt): (empty)

384 return _ColumnSpec('bool', name) (empty)

385 if pd.api.types.is_numeric_dtype(dt): (empty)

386 return _ColumnSpec('numeric', name) (empty)

387 raise _unsupported_dtype_error(name, dt) (empty)

388 

389 @staticmethod (empty)

390 def _transform_column( (empty)

391 series: Series, spec: _ColumnSpec 

392 ) -> Float32[np.ndarray, 'n _']: 

393 if spec.kind == 'ordered_cat': (empty)

394 assert spec.categories is not None (empty)

395 return _ordinal_encode(series.to_numpy(), spec.categories, spec.name) (empty)

396 if spec.kind == 'unordered_cat': (empty)

397 assert spec.categories is not None (empty)

398 return _one_hot_encode(series.to_numpy(), spec.categories, spec.name) (empty)

399 return series.to_numpy(dtype=np.float32)[:, None] (empty)

400 

401 

402class PolarsPreprocessor(_PreprocessorBase): (empty)

403 """Stochtree-style covariate preprocessor for `polars.DataFrame` inputs.""" 

404 

405 _library = 'polars' (empty)

406 

407 @staticmethod (empty)

408 def _get_column(X: DataFrame, orig_idx: int) -> tuple[Any, Series]: (empty)

409 name = X.columns[orig_idx] (empty)

410 return name, X[name] (empty)

411 

412 @staticmethod (empty)

413 def _fit_column(series: Series, name: str) -> _ColumnSpec: (empty)

414 import polars as pl # noqa: PLC0415 # optional runtime dependency (empty)

415 

416 dt = series.dtype (empty)

417 if isinstance(dt, pl.Enum): (empty)

418 # A polars Enum round-trips to a pandas *unordered* Categorical, so 

419 # we treat it as unordered (one-hot). For ordinal encoding, pass an 

420 # integer column. 

421 return _ColumnSpec( (empty)

422 'unordered_cat', name, categories=tuple(dt.categories.to_list()) 

423 ) 

424 if isinstance(dt, pl.Categorical): (empty)

425 raise _polars_categorical_error(name) (empty)

426 if dt == pl.Boolean: (empty)

427 return _ColumnSpec('bool', name) (empty)

428 if dt.is_numeric(): (empty)

429 return _ColumnSpec('numeric', name) (empty)

430 raise _unsupported_dtype_error(name, dt) (empty)

431 

432 @staticmethod (empty)

433 def _transform_column( (empty)

434 series: Series, spec: _ColumnSpec 

435 ) -> Float32[np.ndarray, 'n _']: 

436 import polars as pl # noqa: PLC0415 # optional runtime dependency (empty)

437 

438 if spec.kind == 'unordered_cat': (empty)

439 assert spec.categories is not None (empty)

440 return _polars_one_hot(pl, series, spec.categories, spec.name) (empty)

441 return series.cast(pl.Float32).to_numpy()[:, None] (empty)

442 

443 

444def make_preprocessor(X: object) -> _PreprocessorBase | None: (empty)

445 """Return a preprocessor matched to `X`'s library, or `None` if `X` is not a DataFrame. 

446 

447 Dispatches by inspecting ``type(X).__module__`` to avoid hard imports of 

448 pandas/polars. 

449 """ 

450 mod = type(X).__module__ (empty)

451 if mod.startswith('polars'): (empty)

452 return PolarsPreprocessor() (empty)

453 if mod.startswith('pandas'): (empty)

454 return PandasPreprocessor() (empty)

455 return None (empty)