Experiments
A collection of experiments, code dump etc, that are more for the vibe than for the writing (i.e. these may be heavily written by AI, but are still “cool”).
2026
Speech Diarizer
A fairly vanilla speech diarizer without torch or unusual dependencies. Only
onnxruntime. See https://github.com/NoRaincheck/Speech-Segmentation
1import urllib.request
2import wave
3import numpy as np
4import onnxruntime as ort
5
6model_id = "onnx-community/pyannote-segmentation-3.0"
7
8model_path = "model.onnx"
9urllib.request.urlretrieve(
10 f"https://huggingface.co/{model_id}/resolve/main/onnx/model.onnx",
11 model_path,
12)
13
14url = "https://huggingface.co/datasets/Xenova/transformers.js-docs/resolve/main/mlk.wav"
15audio_path = "mlk.wav"
16urllib.request.urlretrieve(url, audio_path)
17
18with wave.open(audio_path, "rb") as wf:
19 sr = wf.getframerate()
20 frames = wf.readframes(wf.getnframes())
21 audio = np.frombuffer(frames, dtype=np.int16).astype(np.float32) / 32768.0
22
23target_sr = 16000
24step = 270
25
26if sr != target_sr:
27 old_len = len(audio)
28 new_len = int(old_len * target_sr / sr)
29 audio = np.interp(np.linspace(0, old_len - 1, new_len), np.arange(old_len), audio)
30
31session = ort.InferenceSession(model_path)
32logits = session.run(
33 None, {"input_values": audio[np.newaxis, np.newaxis, :].astype(np.float32)}
34)[0]
35
36frame_logits = logits[0]
37exps = np.exp(frame_logits - frame_logits.max(axis=1, keepdims=True))
38probs = exps / exps.sum(axis=1, keepdims=True)
39
40preds = probs.argmax(axis=1)
41confidence = probs.max(axis=1)
42
43segments = []
44current_spk = None
45current_start = None
46max_conf = 0.0
47
48for i, (cls, conf) in enumerate(zip(preds, confidence)):
49 if cls in (1, 2, 3):
50 if cls != current_spk:
51 if current_spk is not None:
52 segments.append(
53 (
54 current_spk,
55 current_start * step / target_sr,
56 i * step / target_sr,
57 float(max_conf),
58 )
59 )
60 current_spk = cls
61 current_start = i
62 max_conf = conf
63 else:
64 max_conf = max(max_conf, conf)
65 else:
66 if current_spk is not None:
67 segments.append(
68 (
69 current_spk,
70 current_start * step / target_sr,
71 i * step / target_sr,
72 float(max_conf),
73 )
74 )
75 current_spk = None
76 current_start = None
77 max_conf = 0.0
78
79if current_spk is not None:
80 segments.append(
81 (
82 current_spk,
83 current_start * step / target_sr,
84 len(preds) * step / target_sr,
85 float(max_conf),
86 )
87 )
88
89for spk_id, start, end, conf in segments:
90 print(f" SPEAKER_{spk_id:02d} {start:7.2f}s - {end:7.2f}s (conf={conf:.3f})")
pre-commit hook for anonymous git commits
1#!/bin/sh
2# Pre-commit hook to check that git config contains $NAME
3NAME='~myname'
4
5# Define color codes
6RED='\033[0;31m'
7GREEN='\033[0;32m'
8YELLOW='\033[1;33m'
9NC='\033[0m' # No Color
10
11# Check if user.name
12USER_NAME=$(git config user.name)
13if ! echo "$USER_NAME" | grep -q "${NAME}"; then
14 echo "${RED}❌ Error: git config user.name does not contain '${NAME}'${NC}"
15 echo "${YELLOW}Current user.name: $USER_NAME${NC}"
16 exit 1
17fi
18
19echo "${GREEN}✅ Git config validation passed: user.name and user.email contain '${NAME}'${NC}"
Full Text Search (sqlite3 FTS5)
Semantic search gets all the attention, but full-text search with BM25 scoring
is fast, dependency-free, and surprisingly effective for local RAG pipelines.
Python’s built-in sqlite3 ships with FTS5 support — no pip installs, no
external databases. Just create a regular table for your documents (plus JSON
metadata), an FTS5 virtual table for indexing, and a few triggers to keep them
in sync.
1"""
2Minimal full-text search client using sqlite3 FTS5 with BM25 scoring.
3
4Stdlib only — no pip installs required. Python 3.12+ (union types).
5Uses JSON functions (SQLite 3.38+) for metadata filtering.
6"""
7
8import json
9import sqlite3
10from typing import Any, Sequence
11
12
13class FTSClient:
14 """
15 A full-text search client backed by sqlite3 FTS5.
16
17 Documents are stored in a regular table (id + content + JSON metadata).
18 An FTS5 virtual table indexes only the ``content`` column. Triggers keep
19 the index in sync automatically.
20
21 BM25 relevance scores are returned via ``bm25(fts_table)`` — lower is better.
22 """
23
24 def __init__(self, db: str = ":memory:"):
25 self.con = sqlite3.connect(db)
26 self._ensure_schema()
27
28 # ------------------------------------------------------------------ schema
29 def _ensure_schema(self) -> None:
30 self.con.executescript("""
31 CREATE TABLE IF NOT EXISTS documents (
32 id TEXT PRIMARY KEY,
33 content TEXT,
34 metadata TEXT -- JSON string stored in the regular table
35 );
36
37 CREATE VIRTUAL TABLE IF NOT EXISTS documents_fts USING fts5(
38 content,
39 tokenize='unicode61'
40 );
41
42 /* Triggers keep the FTS index synced with the main table. */
43 CREATE TRIGGER IF NOT EXISTS documents_ai AFTER INSERT ON documents BEGIN
44 INSERT INTO documents_fts(_rowid_, content) VALUES (new.rowid, new.content);
45 END;
46
47 CREATE TRIGGER IF NOT EXISTS documents_ad AFTER DELETE ON documents BEGIN
48 DELETE FROM documents_fts WHERE rowid = old.rowid;
49 END;
50
51 CREATE TRIGGER IF NOT EXISTS documents_au AFTER UPDATE ON documents BEGIN
52 DELETE FROM documents_fts WHERE rowid = old.rowid;
53 INSERT INTO documents_fts(_rowid_, content) VALUES (new.rowid, new.content);
54 END;
55 """)
56
57 # -------------------------------------------------------------- CRUD helpers
58 def insert(
59 self,
60 data: Sequence[tuple[str, str, dict[str, Any]]],
61 replace: bool = False,
62 ) -> None:
63 """Insert documents.
64
65 Args:
66 data: iterable of ``(id, content, metadata_dict)``.
67 replace: if True, existing rows with the same ``id`` are overwritten.
68
69 Note:
70 We explicitly DELETE before INSERT when ``replace=True`` because
71 SQLite's ``INSERT OR REPLACE`` does not fire AFTER DELETE triggers,
72 which would leave stale FTS entries behind.
73 """
74 serialized = [
75 (doc_id, content, json.dumps(metadata))
76 for doc_id, content, metadata in data
77 ]
78 if replace and serialized:
79 ids = tuple(d[0] for d in serialized)
80 placeholders = ",".join("?" for _ in ids)
81 self.con.execute(
82 f"DELETE FROM documents WHERE id IN ({placeholders})", ids,
83 )
84 sql = "INSERT INTO documents VALUES (?, ?, ?)"
85 self.con.executemany(sql, serialized)
86
87 def delete(self, ids: Sequence[str]) -> None:
88 placeholders = ",".join("?" for _ in ids)
89 self.con.execute(
90 f"DELETE FROM documents WHERE id IN ({placeholders})", ids,
91 )
92
93 # -------------------------------------------------------------- search API
94 def search(
95 self,
96 query: str = "",
97 limit: int = 10,
98 metadata_filter: dict[str, Any] | None = None,
99 ) -> list[tuple[str, str, str, float]]:
100 """
101 Search documents by full-text query (BM25) with optional metadata filter.
102
103 Returns a list of ``(id, content, metadata_json, bm25_score)`` tuples.
104 Scores are lower for better matches.
105 """
106 if not query and not metadata_filter:
107 return []
108
109 meta_clauses: list[str] = []
110 params: list[Any] = []
111 if metadata_filter:
112 for key, val in metadata_filter.items():
113 path = f"$.{key}"
114 # Normalize booleans to lowercase (SQLite JSON stores true/false)
115 norm = str(val).lower() if isinstance(val, bool) else val
116 meta_clauses.append("json_extract(d.metadata, ?) = ?")
117 params.extend([path, norm])
118
119 has_fts = bool(query)
120 has_meta = bool(metadata_filter)
121
122 if has_fts and has_meta:
123 sql = f"""
124 SELECT d.id, d.content, d.metadata,
125 bm25(documents_fts) AS score
126 FROM documents d
127 JOIN documents_fts ON documents_fts._rowid_ = d.rowid
128 WHERE documents_fts MATCH ?
129 AND {' AND '.join(meta_clauses)}
130 ORDER BY score ASC LIMIT ?"""
131 params = [query, *params, limit]
132
133 elif has_meta and not has_fts:
134 sql = f"""
135 SELECT d.id, d.content, d.metadata, NULL AS score
136 FROM documents d
137 WHERE {' AND '.join(meta_clauses)}
138 ORDER BY d.id ASC LIMIT ?"""
139 params = [*params, limit]
140
141 elif has_fts and not has_meta:
142 sql = """
143 SELECT d.id, d.content, d.metadata,
144 bm25(documents_fts) AS score
145 FROM documents d
146 JOIN documents_fts ON documents_fts._rowid_ = d.rowid
147 WHERE documents_fts MATCH ?
148 ORDER BY score ASC LIMIT ?"""
149 params = [query, limit]
150
151 else:
152 return []
153
154 rows = self.con.execute(sql, params).fetchall()
155 return [(rid, content, meta, score) for rid, content, meta, score in rows]
2025
Lightweight SHAP Implementation
December 2025
When performing SHAP, sometimes the requirements to install are overly onerous.
This is a lightweight implementation of Permutation SHAP including helper
functions matching the style of scikit-learn interface.
1"""
2Permutation SHAP with scikit-learn style API.
3
4This module provides permutation-based SHAP feature importance computation
5following scikit-learn's API conventions, with feature selection utilities
6including Boruta.
7
8Example usage:
9
10 >>> from sklearn.ensemble import RandomForestClassifier
11 >>> from sklearn.datasets import make_classification
12 >>> X, y = make_classification(n_samples=100, n_features=10, random_state=42)
13 >>> model = RandomForestClassifier(random_state=42).fit(X, y)
14
15 # Global importance (sklearn permutation_importance style)
16 >>> result = permutation_shap(model, X, n_repeats=5, random_state=42)
17 >>> print(result.importances_mean)
18
19 # Per-instance SHAP values
20 >>> shap_values = permutation_shap_values(model, X[:5], background=X, random_state=42)
21 >>> print(shap_values.shape) # (5, 10) or (5, 10, n_classes)
22
23 # Feature selection
24 >>> selector = SelectFromShap(model, threshold='median', random_state=42)
25 >>> X_selected = selector.fit_transform(X)
26"""
27
28import numpy as np
29from sklearn.base import BaseEstimator, clone, is_classifier
30from sklearn.feature_selection import SelectorMixin
31from sklearn.utils import Bunch, check_random_state
32from sklearn.cluster import KMeans, MiniBatchKMeans
33
34
35def _get_predict_fn(estimator):
36 """
37 Get a prediction function from an estimator.
38
39 For classifiers, uses predict_proba if available, otherwise predict.
40 For regressors, uses predict.
41
42 Parameters
43 ----------
44 estimator : estimator object
45 A fitted estimator.
46
47 Returns
48 -------
49 predict_fn : callable
50 Function that takes X and returns predictions.
51 """
52 if hasattr(estimator, 'predict_proba'):
53 return estimator.predict_proba
54 return estimator.predict
55
56
57def _summarize_data(X, max_samples, random_state):
58 """
59 Summarize data using k-means clustering if max_samples is specified.
60
61 Parameters
62 ----------
63 X : np.ndarray
64 Input data to summarize. Shape: (n_samples, n_features)
65 max_samples : int, float, or None
66 If None, return X unchanged.
67 If int, reduce to that many centroids via k-means.
68 If float (0-1), reduce to that fraction of samples.
69 random_state : RandomState
70 Random state for k-means.
71
72 Returns
73 -------
74 np.ndarray
75 Either the original X or k-means cluster centroids.
76 """
77 if max_samples is None:
78 return X
79
80 # Determine number of clusters
81 if isinstance(max_samples, float):
82 if not 0 < max_samples <= 1:
83 raise ValueError("max_samples as float must be in (0, 1]")
84 n_clusters = max(1, int(max_samples * X.shape[0]))
85 else:
86 n_clusters = int(max_samples)
87
88 # No reduction needed if n_clusters >= n_samples
89 if n_clusters >= X.shape[0]:
90 return X
91
92 if n_clusters < 1:
93 raise ValueError("max_samples must result in at least 1 cluster")
94
95 # Use MiniBatchKMeans for large datasets (faster), KMeans otherwise (more accurate)
96 threshold = 10000
97 KMeansClass = MiniBatchKMeans if X.shape[0] > threshold else KMeans
98 seed = random_state.randint(0, 2**31) if hasattr(random_state, 'randint') else random_state
99 kmeans = KMeansClass(n_clusters=n_clusters, random_state=seed, n_init='auto')
100 kmeans.fit(X)
101 return kmeans.cluster_centers_
102
103
104def _compute_shap_single_instance(x, predict_fn, background, n_permutations, rng):
105 """
106 Compute SHAP values for a single instance using permutation method
107 with antithetic sampling.
108
109 Parameters
110 ----------
111 x : np.ndarray
112 Instance to explain. Shape: (n_features,)
113 predict_fn : callable
114 Prediction function.
115 background : np.ndarray
116 Background data. Shape: (n_background, n_features)
117 n_permutations : int
118 Number of permutations to sample.
119 rng : RandomState
120 Random number generator.
121
122 Returns
123 -------
124 shap_values : np.ndarray
125 SHAP values for each feature.
126 """
127 n_features = x.shape[0]
128 n_background = background.shape[0]
129
130 # Determine output dimensionality
131 sample_pred = predict_fn(background[:1])
132 if sample_pred.ndim == 1:
133 n_outputs = 1
134 contributions = [[] for _ in range(n_features)]
135 else:
136 n_outputs = sample_pred.shape[1]
137 contributions = [[] for _ in range(n_features)]
138
139 for _ in range(n_permutations):
140 # Sample a random permutation of feature indices
141 perm = rng.permutation(n_features)
142
143 # Sample a background instance for absent features
144 bg_idx = rng.randint(n_background)
145 bg_sample = background[bg_idx].copy()
146
147 # Forward pass: add features according to permutation
148 _compute_marginal_contributions(x, perm, bg_sample, predict_fn, contributions)
149
150 # Backward pass (antithetic): use reverse permutation to reduce variance
151 reverse_perm = perm[::-1]
152 _compute_marginal_contributions(x, reverse_perm, bg_sample, predict_fn, contributions)
153
154 # Average all marginal contributions for each feature
155 if n_outputs == 1:
156 shap_values = np.array([np.mean(contribs) for contribs in contributions])
157 else:
158 shap_values = np.array([
159 np.mean(contribs, axis=0) for contribs in contributions
160 ])
161
162 return shap_values
163
164
165def _compute_marginal_contributions(x, permutation, background_sample, predict_fn,
166 contributions):
167 """
168 Compute marginal contributions for a single permutation pass.
169
170 Starting from a coalition containing no features from x (all from background),
171 we add features one at a time according to the permutation order.
172 Each addition gives us a marginal contribution for that feature.
173
174 Parameters
175 ----------
176 x : np.ndarray
177 Instance to explain. Shape: (n_features,)
178 permutation : np.ndarray
179 Order in which to add features. Shape: (n_features,)
180 background_sample : np.ndarray
181 Background sample for absent features. Shape: (n_features,)
182 predict_fn : callable
183 Prediction function.
184 contributions : list of lists
185 Contribution accumulator. contributions[i] is a list of marginal
186 contributions for feature i.
187 """
188 # Start with all features from background (empty coalition w.r.t. x)
189 current_sample = background_sample.copy()
190 prev_pred = predict_fn(current_sample.reshape(1, -1))[0]
191
192 for feat_idx in permutation:
193 # Add this feature to the coalition (use value from x)
194 current_sample[feat_idx] = x[feat_idx]
195 current_pred = predict_fn(current_sample.reshape(1, -1))[0]
196
197 # Marginal contribution: v(S ∪ {i}) - v(S)
198 marginal = current_pred - prev_pred
199 contributions[feat_idx].append(marginal)
200
201 # Reuse this prediction as the "before" for the next feature
202 prev_pred = current_pred
203
204
205def permutation_shap_values(estimator, X, *, background=None, n_permutations=10,
206 max_samples=None, random_state=None):
207 """
208 Compute per-instance SHAP values using permutation method.
209
210 This function computes SHAP values for each instance in X using the
211 permutation method with antithetic sampling for variance reduction.
212
213 Parameters
214 ----------
215 estimator : estimator object
216 A fitted estimator. Must have a predict or predict_proba method.
217 X : array-like of shape (n_samples, n_features)
218 Data for which to compute SHAP values.
219 background : array-like of shape (n_background, n_features), optional
220 Background data used for computing baseline expectations.
221 If None, uses X as the background.
222 n_permutations : int, default=10
223 Number of permutations to sample. Each permutation generates both
224 forward and backward (antithetic) passes.
225 max_samples : int, float, or None, default=None
226 The number of samples to use from background. If background is large,
227 k-means clustering is used to summarize it to centroids.
228 - If None, use the full background unchanged.
229 - If int, reduce background to that many centroids via k-means.
230 - If float (0-1), reduce background to that fraction of samples.
231 random_state : int, RandomState instance, or None, default=None
232 Controls the randomness. Pass an int for reproducible results.
233
234 Returns
235 -------
236 shap_values : np.ndarray
237 SHAP values for each instance and feature.
238 - If estimator outputs 1D predictions: shape (n_samples, n_features)
239 - If estimator outputs 2D predictions: shape (n_samples, n_features, n_outputs)
240
241 Examples
242 --------
243 >>> from sklearn.ensemble import RandomForestClassifier
244 >>> from sklearn.datasets import make_classification
245 >>> X, y = make_classification(n_samples=100, n_features=10, random_state=42)
246 >>> model = RandomForestClassifier(random_state=42).fit(X, y)
247 >>> shap_values = permutation_shap_values(model, X[:5], background=X, random_state=42)
248 >>> shap_values.shape
249 (5, 10, 2)
250 """
251 X = np.asarray(X)
252 rng = check_random_state(random_state)
253
254 if background is None:
255 background = X
256 else:
257 background = np.asarray(background)
258
259 # Summarize background if requested
260 background = _summarize_data(background, max_samples, rng)
261
262 predict_fn = _get_predict_fn(estimator)
263
264 # Compute SHAP values for each instance
265 results = []
266 for i in range(X.shape[0]):
267 shap_vals = _compute_shap_single_instance(
268 X[i], predict_fn, background, n_permutations, rng
269 )
270 results.append(shap_vals)
271
272 return np.array(results)
273
274
275def permutation_shap(estimator, X, y=None, *, background=None, n_repeats=5,
276 n_permutations=10, max_samples=None, random_state=None):
277 """
278 Compute permutation-based SHAP feature importance.
279
280 This function follows the scikit-learn permutation_importance API pattern.
281 It computes SHAP values and aggregates them to global feature importance
282 using mean absolute SHAP values.
283
284 Parameters
285 ----------
286 estimator : estimator object
287 A fitted estimator. Must have a predict or predict_proba method.
288 X : array-like of shape (n_samples, n_features)
289 Data on which to compute feature importance.
290 y : array-like of shape (n_samples,), optional
291 Target values. Not used in computation but accepted for API compatibility.
292 background : array-like of shape (n_background, n_features), optional
293 Background data used for computing baseline expectations.
294 If None, uses X as the background.
295 n_repeats : int, default=5
296 Number of times to repeat the importance computation with different
297 random seeds for variance estimation.
298 n_permutations : int, default=10
299 Number of permutations to sample per repeat for SHAP value estimation.
300 max_samples : int, float, or None, default=None
301 The number of samples to use from background. If background is large,
302 k-means clustering is used to summarize it to centroids.
303 - If None, use the full background unchanged.
304 - If int, reduce background to that many centroids via k-means.
305 - If float (0-1), reduce background to that fraction of samples.
306 random_state : int, RandomState instance, or None, default=None
307 Controls the randomness. Pass an int for reproducible results.
308
309 Returns
310 -------
311 result : Bunch
312 Dictionary-like object with the following attributes:
313
314 importances_mean : np.ndarray of shape (n_features,)
315 Mean of feature importances across repeats.
316 importances_std : np.ndarray of shape (n_features,)
317 Standard deviation of feature importances across repeats.
318 importances : np.ndarray of shape (n_features, n_repeats)
319 Raw feature importances for each repeat.
320
321 Examples
322 --------
323 >>> from sklearn.ensemble import RandomForestClassifier
324 >>> from sklearn.datasets import make_classification
325 >>> X, y = make_classification(n_samples=100, n_features=10, random_state=42)
326 >>> model = RandomForestClassifier(random_state=42).fit(X, y)
327 >>> result = permutation_shap(model, X, n_repeats=5, random_state=42)
328 >>> result.importances_mean.shape
329 (10,)
330 >>> for i in result.importances_mean.argsort()[::-1][:3]:
331 ... print(f"Feature {i}: {result.importances_mean[i]:.4f} +/- {result.importances_std[i]:.4f}")
332 """
333 X = np.asarray(X)
334 rng = check_random_state(random_state)
335 n_features = X.shape[1]
336
337 if background is None:
338 background = X
339 else:
340 background = np.asarray(background)
341
342 # Summarize background if requested (do this once, not per repeat)
343 background = _summarize_data(background, max_samples, rng)
344
345 # Collect importances across repeats
346 importances = np.zeros((n_features, n_repeats))
347
348 for repeat in range(n_repeats):
349 # Compute SHAP values for all instances
350 shap_values = permutation_shap_values(
351 estimator, X,
352 background=background,
353 n_permutations=n_permutations,
354 max_samples=None, # Already summarized
355 random_state=rng
356 )
357
358 # For multi-output, sum absolute values across outputs
359 if shap_values.ndim == 3:
360 # Shape: (n_samples, n_features, n_outputs)
361 # Take mean absolute across samples, sum across outputs
362 feature_importance = np.mean(np.sum(np.abs(shap_values), axis=2), axis=0)
363 else:
364 # Shape: (n_samples, n_features)
365 # Take mean absolute across samples
366 feature_importance = np.mean(np.abs(shap_values), axis=0)
367
368 importances[:, repeat] = feature_importance
369
370 return Bunch(
371 importances_mean=np.mean(importances, axis=1),
372 importances_std=np.std(importances, axis=1),
373 importances=importances
374 )
375
376
377class SelectFromShap(SelectorMixin, BaseEstimator):
378 """
379 Feature selector based on permutation SHAP importance.
380
381 This transformer uses permutation-based SHAP values to select features
382 based on their importance. It supports threshold-based selection and
383 Boruta feature selection.
384
385 Parameters
386 ----------
387 estimator : estimator object
388 A fitted estimator to compute SHAP importance for.
389 threshold : float, str, or None, default=None
390 The threshold value to use for feature selection. Features with
391 importance greater than or equal to the threshold are selected.
392 - If float, features with importance >= threshold are selected.
393 - If "median", uses the median of feature importances.
394 - If "mean", uses the mean of feature importances.
395 - If None and mode='threshold', defaults to "median".
396 Ignored when mode='boruta'.
397 max_features : int or None, default=None
398 The maximum number of features to select. If not None, only the
399 top max_features features are selected. If both threshold and
400 max_features are specified, max_features takes precedence.
401 mode : {'threshold', 'boruta'}, default='threshold'
402 Feature selection mode:
403 - 'threshold': Select features above importance threshold.
404 - 'boruta': Use Boruta algorithm with shadow features.
405 n_permutations : int, default=10
406 Number of permutations for SHAP value estimation.
407 n_repeats : int, default=5
408 Number of repeats for importance variance estimation.
409 max_samples : int, float, or None, default=None
410 Maximum samples for background summarization.
411 boruta_max_iter : int, default=100
412 Maximum iterations for Boruta algorithm. Only used when mode='boruta'.
413 boruta_alpha : float, default=0.05
414 Significance level for Boruta statistical test. Only used when mode='boruta'.
415 random_state : int, RandomState instance, or None, default=None
416 Controls randomness.
417
418 Attributes
419 ----------
420 importances_ : np.ndarray of shape (n_features,)
421 Feature importances (mean absolute SHAP values).
422 importances_std_ : np.ndarray of shape (n_features,)
423 Standard deviation of feature importances.
424 support_ : np.ndarray of shape (n_features,)
425 Boolean mask of selected features.
426 n_features_in_ : int
427 Number of features seen during fit.
428 ranking_ : np.ndarray of shape (n_features,)
429 Feature ranking (1 = selected, higher = less important).
430 Only available when mode='boruta'.
431
432 Examples
433 --------
434 >>> from sklearn.ensemble import RandomForestClassifier
435 >>> from sklearn.datasets import make_classification
436 >>> X, y = make_classification(n_samples=100, n_features=20,
437 ... n_informative=5, random_state=42)
438 >>> model = RandomForestClassifier(random_state=42).fit(X, y)
439
440 # Threshold-based selection
441 >>> selector = SelectFromShap(model, threshold='median', random_state=42)
442 >>> X_selected = selector.fit_transform(X)
443 >>> X_selected.shape[1] # About half the features
444
445 # Top-k selection
446 >>> selector = SelectFromShap(model, max_features=5, random_state=42)
447 >>> X_selected = selector.fit_transform(X)
448 >>> X_selected.shape[1]
449 5
450
451 # Boruta selection
452 >>> selector = SelectFromShap(model, mode='boruta', random_state=42)
453 >>> X_selected = selector.fit_transform(X)
454 """
455
456 def __init__(self, estimator, *, threshold=None, max_features=None,
457 mode='threshold', n_permutations=10, n_repeats=5,
458 max_samples=None, boruta_max_iter=100, boruta_alpha=0.05,
459 random_state=None):
460 self.estimator = estimator
461 self.threshold = threshold
462 self.max_features = max_features
463 self.mode = mode
464 self.n_permutations = n_permutations
465 self.n_repeats = n_repeats
466 self.max_samples = max_samples
467 self.boruta_max_iter = boruta_max_iter
468 self.boruta_alpha = boruta_alpha
469 self.random_state = random_state
470
471 def fit(self, X, y=None):
472 """
473 Fit the feature selector.
474
475 Parameters
476 ----------
477 X : array-like of shape (n_samples, n_features)
478 Training data.
479 y : array-like of shape (n_samples,), optional
480 Target values. Not used but accepted for pipeline compatibility.
481
482 Returns
483 -------
484 self : object
485 Returns self.
486 """
487 X = np.asarray(X)
488 self.n_features_in_ = X.shape[1]
489
490 if self.mode == 'boruta':
491 self._fit_boruta(X, y)
492 else:
493 self._fit_threshold(X, y)
494
495 return self
496
497 def _fit_threshold(self, X, y):
498 """Fit using threshold-based selection."""
499 rng = check_random_state(self.random_state)
500
501 # Compute SHAP importance
502 result = permutation_shap(
503 self.estimator, X, y,
504 background=X,
505 n_repeats=self.n_repeats,
506 n_permutations=self.n_permutations,
507 max_samples=self.max_samples,
508 random_state=rng
509 )
510
511 self.importances_ = result.importances_mean
512 self.importances_std_ = result.importances_std
513
514 # Determine support mask
515 if self.max_features is not None:
516 # Select top max_features
517 n_select = min(self.max_features, self.n_features_in_)
518 top_indices = np.argsort(self.importances_)[::-1][:n_select]
519 self.support_ = np.zeros(self.n_features_in_, dtype=bool)
520 self.support_[top_indices] = True
521 else:
522 # Use threshold
523 threshold = self.threshold
524 if threshold is None or threshold == 'median':
525 threshold = np.median(self.importances_)
526 elif threshold == 'mean':
527 threshold = np.mean(self.importances_)
528
529 self.support_ = self.importances_ >= threshold
530
531 # Ensure at least one feature is selected
532 if not np.any(self.support_):
533 best_idx = np.argmax(self.importances_)
534 self.support_[best_idx] = True
535
536 def _fit_boruta(self, X, y):
537 """
538 Fit using Boruta algorithm.
539
540 The Boruta algorithm:
541 1. Create shadow features (shuffled copies of all original features)
542 2. Train model on real + shadow features
543 3. Compute SHAP importance for all features
544 4. Compare each real feature's importance to max shadow importance
545 5. Use statistical test to classify features as confirmed/rejected/tentative
546 6. Repeat until all features are classified or max_iter reached
547 """
548 rng = check_random_state(self.random_state)
549 n_features = X.shape[1]
550
551 # Track feature status: 0=tentative, 1=confirmed, -1=rejected
552 status = np.zeros(n_features, dtype=int)
553
554 # Track hits (times feature beat max shadow)
555 hits = np.zeros(n_features, dtype=int)
556
557 # Track total trials
558 n_trials = 0
559
560 for iteration in range(self.boruta_max_iter):
561 # Check if all features are classified
562 if np.all(status != 0):
563 break
564
565 # Get tentative feature indices
566 tentative_mask = (status == 0)
567
568 # Create shadow features by shuffling each column independently
569 X_shadow = X.copy()
570 for j in range(n_features):
571 rng.shuffle(X_shadow[:, j])
572
573 # Combine real and shadow features
574 X_combined = np.hstack([X, X_shadow])
575
576 # Clone and refit estimator on combined data
577 combined_estimator = clone(self.estimator)
578 combined_estimator.fit(X_combined, y)
579
580 # Compute SHAP importance for combined features
581 result = permutation_shap(
582 combined_estimator, X_combined, y,
583 background=X_combined,
584 n_repeats=1, # Single repeat per iteration
585 n_permutations=self.n_permutations,
586 max_samples=self.max_samples,
587 random_state=rng
588 )
589
590 importances = result.importances_mean
591 real_importances = importances[:n_features]
592 shadow_importances = importances[n_features:]
593
594 # Max shadow importance (threshold)
595 max_shadow = np.max(shadow_importances)
596
597 # Update hits for tentative features
598 for j in range(n_features):
599 if status[j] == 0: # Tentative
600 if real_importances[j] > max_shadow:
601 hits[j] += 1
602
603 n_trials += 1
604
605 # Statistical test using binomial distribution
606 # Under null hypothesis, P(beat max shadow) = 0.5
607 # Use two-tailed test with Bonferroni correction
608 if n_trials >= 5: # Need minimum trials for stable test
609 for j in range(n_features):
610 if status[j] == 0: # Tentative
611 # Binomial test: probability of getting this many hits by chance
612 # Using normal approximation for binomial
613 p_value = self._binomial_test(hits[j], n_trials)
614
615 if p_value < self.boruta_alpha / n_features: # Bonferroni correction
616 if hits[j] > n_trials / 2:
617 status[j] = 1 # Confirmed
618 else:
619 status[j] = -1 # Rejected
620
621 # Final classification: treat remaining tentative as rejected
622 self.support_ = (status == 1)
623
624 # If no features confirmed, select the one with most hits
625 if not np.any(self.support_):
626 best_idx = np.argmax(hits)
627 self.support_[best_idx] = True
628
629 # Compute final importances on original features
630 result = permutation_shap(
631 self.estimator, X, y,
632 background=X,
633 n_repeats=self.n_repeats,
634 n_permutations=self.n_permutations,
635 max_samples=self.max_samples,
636 random_state=rng
637 )
638
639 self.importances_ = result.importances_mean
640 self.importances_std_ = result.importances_std
641
642 # Create ranking: 1 for confirmed, 2 for tentative, 3 for rejected
643 self.ranking_ = np.where(status == 1, 1, np.where(status == 0, 2, 3))
644
645 def _binomial_test(self, k, n, p=0.5):
646 """
647 Two-tailed binomial test using normal approximation.
648
649 Tests whether the number of successes k in n trials is
650 significantly different from expected under p=0.5.
651
652 Parameters
653 ----------
654 k : int
655 Number of successes (hits).
656 n : int
657 Number of trials.
658 p : float, default=0.5
659 Null hypothesis probability.
660
661 Returns
662 -------
663 p_value : float
664 Two-tailed p-value.
665 """
666 # Normal approximation to binomial
667 mean = n * p
668 std = np.sqrt(n * p * (1 - p))
669
670 if std == 0:
671 return 1.0
672
673 # Z-score
674 z = abs(k - mean) / std
675
676 # Two-tailed p-value using normal CDF approximation
677 # Using error function approximation
678 p_value = 2 * (1 - self._norm_cdf(z))
679
680 return p_value
681
682 def _norm_cdf(self, x):
683 """
684 Standard normal CDF approximation.
685
686 Uses the error function approximation.
687 """
688 # Approximation using the error function
689 # CDF(x) = 0.5 * (1 + erf(x / sqrt(2)))
690 return 0.5 * (1 + self._erf(x / np.sqrt(2)))
691
692 def _erf(self, x):
693 """
694 Error function approximation.
695
696 Abramowitz and Stegun approximation (max error ~ 1.5e-7).
697 """
698 # Constants
699 a1 = 0.254829592
700 a2 = -0.284496736
701 a3 = 1.421413741
702 a4 = -1.453152027
703 a5 = 1.061405429
704 p = 0.3275911
705
706 sign = np.sign(x)
707 x = np.abs(x)
708
709 t = 1.0 / (1.0 + p * x)
710 y = 1.0 - (((((a5 * t + a4) * t) + a3) * t + a2) * t + a1) * t * np.exp(-x * x)
711
712 return sign * y
713
714 def _get_support_mask(self):
715 """Get the boolean mask of selected features."""
716 return self.support_
Tree Ensemble Hashing
August 2025
I’ve been thinking about how to encode decision trees (on and off). This is a vibe-coded approach that tries to encode the path of every instance based on what leaf of the decision tree(s) the instance lands on. The rough idea is that it also gets ‘weight’ from every parent node (weighted based on distance to leaf) and is encoded by feature hashing, which allows it to be aggregated by all trees. It somewhat works. I think this kind of approach may be useful for scenarios where online learning with just a linear head is used, especially when feature engineering (from the trees) is done in an online manner, for example via Mondrian Trees.
1import mmh3
2import numpy as np
3from typing import Dict, List, Tuple, Optional, Union, Any
4from onnx import ModelProto
5import logging
6
7logger = logging.getLogger(__name__)
8
9
10class TreeEnsembleHash:
11 """
12 Implements the tree ensemble hashing algorithm described in the README.
13
14 The algorithm works as follows:
15 1. Extract all nodes, splits (thresholds), and paths from an ONNX tree ensemble model
16 2. For each input sample, traverse the trees to find which leaves are visited
17 3. Hash only the visited leaves based on {node_id} {operation} {threshold}
18 4. Compute weights based on inverse distance to leaf nodes using 1/d(x,y)^p, 0<p<1
19 5. Aggregate all hashed leaf features to create the final representation
20 """
21
22 def __init__(self, onnx_model: Union[ModelProto, Tuple[ModelProto, Any]], hash_dim: int = 128, distance_power: float = 0.5):
23 """
24 Initialize the TreeEnsembleHash with an ONNX model.
25
26 Args:
27 onnx_model: ONNX model containing tree ensemble operators, or tuple (ModelProto, Topology)
28 hash_dim: Dimension for murmurhash3 output
29 distance_power: Power parameter for distance weighting (0 < p < 1)
30 """
31 if not 0 < distance_power < 1:
32 raise ValueError("distance_power must be between 0 and 1")
33 if hash_dim <= 0:
34 raise ValueError("hash_dim must be positive")
35
36 self.hash_dim = hash_dim
37 self.distance_power = distance_power
38
39 # Extract and store tree ensemble information
40 self.tree_ensembles = self._extract_tree_ensemble_info(onnx_model)
41 if not self.tree_ensembles:
42 raise ValueError("No tree ensemble operators found in the model")
43
44 # Pre-compute distances for all trees
45 self.tree_distances = []
46 for tree_info in self.tree_ensembles:
47 distances = self._compute_node_distances(tree_info)
48 self.tree_distances.append(distances)
49
50 def _extract_tree_ensemble_info(self, model: Union[ModelProto, Tuple[ModelProto, Any]]) -> List[Dict]:
51 """
52 Extract tree ensemble information from ONNX model.
53
54 Args:
55 model: ONNX model containing tree ensemble operators, or tuple (ModelProto, Topology)
56
57 Returns:
58 List of dictionaries containing tree ensemble information
59 """
60 # Handle case where convert_sklearn returns (ModelProto, Topology)
61 if isinstance(model, tuple):
62 model = model[0]
63
64 tree_ensembles = []
65
66 for node in model.graph.node:
67 # Check both the op_type and domain for tree ensemble operators
68 if (node.op_type in ['TreeEnsembleClassifier', 'TreeEnsembleRegressor'] or
69 (hasattr(node, 'domain') and node.domain == 'ai.onnx.ml' and
70 node.op_type in ['TreeEnsembleClassifier', 'TreeEnsembleRegressor'])):
71 # Extract attributes
72 attrs = {attr.name: attr for attr in node.attribute}
73
74 # Get node information
75 node_ids = self._get_attribute_values(attrs, 'nodes_nodeids')
76 feature_ids = self._get_attribute_values(attrs, 'nodes_featureids')
77 values = self._get_attribute_values(attrs, 'nodes_values')
78 modes = self._get_attribute_values(attrs, 'nodes_modes')
79 true_node_ids = self._get_attribute_values(attrs, 'nodes_truenodeids')
80 false_node_ids = self._get_attribute_values(attrs, 'nodes_falsenodeids')
81 tree_ids = self._get_attribute_values(attrs, 'nodes_treeids')
82
83 if all(v is not None for v in [node_ids, feature_ids, values, modes, true_node_ids, false_node_ids, tree_ids]):
84 tree_ensembles.append({
85 'node_ids': node_ids,
86 'feature_ids': feature_ids,
87 'values': values,
88 'modes': modes,
89 'true_node_ids': true_node_ids,
90 'false_node_ids': false_node_ids,
91 'tree_ids': tree_ids
92 })
93
94 return tree_ensembles
95
96 def _get_attribute_values(self, attrs: Dict, name: str) -> Optional[List]:
97 """Extract attribute values from ONNX node attributes."""
98 if name not in attrs:
99 logger.debug(f"Attribute {name} not found in {list(attrs.keys())}")
100 return None
101
102 attr = attrs[name]
103 logger.debug(f"Attribute {name} type: {attr.type}")
104
105 try:
106 if attr.type == 1: # FLOAT
107 result = list(attr.floats)
108 return result
109 elif attr.type == 2: # INT
110 result = list(attr.ints)
111 return result
112 elif attr.type == 3: # STRING
113 result = list(attr.strings)
114 return result
115 elif attr.type == 6: # FLOATS (alternative representation)
116 result = list(attr.floats)
117 return result
118 elif attr.type == 7: # INTS (alternative representation)
119 result = list(attr.ints)
120 return result
121 elif attr.type == 8: # STRINGS (alternative representation)
122 result = list(attr.strings)
123 return result
124 else:
125 return None
126 except Exception as e:
127 return None
128
129 def _compute_node_distances(self, tree_info: Dict) -> Dict[int, Union[int, float]]:
130 """
131 Compute the distance from each node to its nearest leaf node.
132
133 Args:
134 tree_info: Dictionary containing tree structure information
135
136 Returns:
137 Dictionary mapping node_id to distance to nearest leaf
138 """
139 node_ids = tree_info['node_ids']
140 modes = tree_info['modes']
141 true_node_ids = tree_info['true_node_ids']
142 false_node_ids = tree_info['false_node_ids']
143
144 # Create adjacency list
145 adjacency = {}
146 for i, node_id in enumerate(node_ids):
147 if modes[i] != b'LEAF' and modes[i] != 'LEAF':
148 adjacency[node_id] = []
149 if true_node_ids[i] != 0:
150 adjacency[node_id].append(true_node_ids[i])
151 if false_node_ids[i] != 0:
152 adjacency[node_id].append(false_node_ids[i])
153
154 # BFS to compute distances from each node to nearest leaf
155 distances = {}
156 for node_id in node_ids:
157 if node_id not in distances:
158 distances[node_id] = self._bfs_distance_to_leaf(node_id, adjacency, modes, node_ids)
159
160 return distances
161
162 def _bfs_distance_to_leaf(self, start_node: int, adjacency: Dict, modes: List, node_ids: List) -> Union[int, float]:
163 """Compute distance from start_node to nearest leaf using BFS."""
164 if start_node not in adjacency:
165 return 0 # This is a leaf node
166
167 visited = set()
168 queue = [(start_node, 0)]
169
170 while queue:
171 current_node, distance = queue.pop(0)
172
173 if current_node in visited:
174 continue
175
176 visited.add(current_node)
177
178 # Find the index of current_node in node_ids
179 try:
180 idx = node_ids.index(current_node)
181 if modes[idx] == b'LEAF' or modes[idx] == 'LEAF':
182 return distance
183 except ValueError:
184 continue
185
186 # Add neighbors to queue
187 if current_node in adjacency:
188 for neighbor in adjacency[current_node]:
189 if neighbor != 0: # 0 indicates no child
190 queue.append((neighbor, distance + 1))
191
192 return float('inf') # No path to leaf found
193
194 def _traverse_tree(self, X: np.ndarray, tree_info: Dict, tree_idx: int) -> List[Tuple[int, str, float]]:
195 """
196 Traverse a single tree based on input data to find all visited nodes along the path.
197
198 Args:
199 X: Input data (single sample)
200 tree_info: Dictionary containing tree information
201 tree_idx: Index of the tree in the ensemble
202
203 Returns:
204 List of tuples (node_id, operation, threshold) for all nodes along the path
205 """
206 node_ids = tree_info['node_ids']
207 feature_ids = tree_info['feature_ids']
208 values = tree_info['values']
209 modes = tree_info['modes']
210 true_node_ids = tree_info['true_node_ids']
211 false_node_ids = tree_info['false_node_ids']
212 tree_ids = tree_info['tree_ids']
213
214 # Find nodes belonging to this specific tree
215 tree_mask = [i for i, tid in enumerate(tree_ids) if tid == tree_idx]
216
217 if not tree_mask:
218 return []
219
220 # Create a mapping from node_id to index for this tree
221 node_to_idx = {node_ids[i]: i for i in tree_mask}
222
223 # Find the root node - it should be the one that's not referenced as a child
224 all_child_nodes = set()
225 for i in tree_mask:
226 if true_node_ids[i] != 0:
227 all_child_nodes.add(true_node_ids[i])
228 if false_node_ids[i] != 0:
229 all_child_nodes.add(false_node_ids[i])
230
231 # Root node is the one that's not a child of any other node
232 root_candidates = [node_ids[i] for i in tree_mask if node_ids[i] not in all_child_nodes]
233
234 if not root_candidates:
235 # Fallback: use first node in tree
236 root_node_id = node_ids[tree_mask[0]]
237 else:
238 root_node_id = root_candidates[0]
239
240 current_node_id = root_node_id
241 visited_nodes = []
242
243 while True:
244 if current_node_id not in node_to_idx:
245 break
246
247 idx = node_to_idx[current_node_id]
248 mode = modes[idx]
249 feature_id = feature_ids[idx] if feature_ids else 0
250 value = values[idx] if values else 0.0
251
252 # Add this node to visited nodes (including its operation and threshold)
253 visited_nodes.append((current_node_id, mode, value))
254
255 # If we've reached a leaf, stop traversing
256 if mode == b'LEAF' or mode == 'LEAF':
257 break
258
259 # If it's a branch node, traverse based on feature value
260 if mode == b'BRANCH_LEQ' or mode == 'BRANCH_LEQ':
261 # Get the feature value for this sample
262 if feature_id < X.shape[0]:
263 feature_value = X[feature_id]
264
265 # Decide which path to take
266 if feature_value <= value:
267 next_node_id = true_node_ids[idx]
268 else:
269 next_node_id = false_node_ids[idx]
270
271 # Move to next node
272 if next_node_id != 0: # 0 indicates no child
273 current_node_id = next_node_id
274 else:
275 break
276 else:
277 # Feature ID out of bounds, stop
278 break
279 else:
280 # Unknown mode, stop
281 break
282
283 return visited_nodes
284
285 def _hash_node(self, node_id: int, operation: str, threshold: float) -> int:
286 """
287 Hash a node using murmurhash3 based on {node_id} {operation} {threshold}.
288
289 Args:
290 node_id: Node identifier
291 operation: Operation type (BRANCH_LEQ, LEAF, etc.)
292 threshold: Threshold value for the split
293
294 Returns:
295 Hash value
296 """
297 node_string = f"{node_id} {operation} {threshold}"
298 return mmh3.hash(node_string)
299
300 def _compute_node_weight(self, distance: Union[int, float]) -> float:
301 """
302 Compute weight based on inverse distance to leaf node.
303
304 Args:
305 distance: Distance to leaf node
306
307 Returns:
308 Weight value using 1/d(x,y)^p formula
309 """
310 if distance == 0:
311 return 1.0
312 if distance == float('inf'):
313 return 0.0 # Infinite distance gets zero weight
314 return 1.0 / (distance ** self.distance_power)
315
316 def _create_hash_vector_from_nodes(self, visited_nodes: List[Tuple[int, str, float]], tree_info: Dict, distances: Dict[int, Union[int, float]]) -> np.ndarray:
317 """
318 Create hash vector from all visited nodes along the path.
319
320 Args:
321 visited_nodes: List of tuples (node_id, operation, threshold) for visited nodes
322 tree_info: Dictionary containing tree information
323 distances: Dictionary mapping node_id to distance to leaf
324
325 Returns:
326 Hash vector of dimension hash_dim
327 """
328 hash_vector = np.zeros(self.hash_dim)
329
330 for node_id, operation, threshold in visited_nodes:
331 # Hash the node using its ID, operation, and threshold
332 node_hash = self._hash_node(node_id, operation, threshold)
333
334 # Compute weight based on distance to leaf
335 distance = distances.get(node_id, 0)
336 weight = self._compute_node_weight(distance)
337
338 # Add weighted hash to vector
339 hash_idx = abs(node_hash) % self.hash_dim
340 hash_vector[hash_idx] += weight
341
342 return hash_vector
343
344 def transform(self, X: np.ndarray) -> np.ndarray:
345 """
346 Transform input data to hash vector representation by traversing trees.
347
348 Args:
349 X: Input data matrix of shape (n_samples, n_features)
350
351 Returns:
352 Hash vector representation of shape (n_samples, hash_dim)
353 """
354 if X.ndim == 1:
355 X = X.reshape(1, -1)
356
357 n_samples = X.shape[0]
358 result = np.zeros((n_samples, self.hash_dim))
359
360 for sample_idx in range(n_samples):
361 sample = X[sample_idx]
362
363 # Aggregate hash vectors from all trees
364 sample_vector = np.zeros(self.hash_dim)
365
366 for tree_idx, tree_info in enumerate(self.tree_ensembles):
367 # Find unique tree IDs in this tree ensemble
368 unique_tree_ids = list(set(tree_info['tree_ids']))
369
370 for tree_id in unique_tree_ids:
371 # Traverse this specific tree to find visited nodes along the path
372 visited_nodes = self._traverse_tree(sample, tree_info, tree_id)
373
374 # Create hash vector from visited nodes
375 tree_vector = self._create_hash_vector_from_nodes(
376 visited_nodes, tree_info, self.tree_distances[tree_idx]
377 )
378
379 # Add to sample vector
380 sample_vector += tree_vector
381
382 result[sample_idx] = sample_vector
383
384 return result
385
386 def fit_transform(self, X: np.ndarray) -> np.ndarray:
387 """
388 Alias for transform method to maintain sklearn-like interface.
389
390 Args:
391 X: Input data matrix
392
393 Returns:
394 Hash vector representation
395 """
396 return self.transform(X)
397
398
399def create_tree_ensemble_hash(
400 onnx_model: Union[ModelProto, Tuple[ModelProto, Any]],
401 X: np.ndarray,
402 hash_dim: int = 128,
403 distance_power: float = 0.5
404) -> np.ndarray:
405 """
406 Convenience function to create tree ensemble hash vector.
407
408 Args:
409 onnx_model: ONNX model containing tree ensemble operators
410 X: Input data matrix
411 hash_dim: Dimension for murmurhash3 output
412 distance_power: Power parameter for distance weighting (0 < p < 1)
413
414 Returns:
415 Hash vector representation of the tree ensemble
416 """
417 hasher = TreeEnsembleHash(onnx_model, hash_dim=hash_dim, distance_power=distance_power)
418 return hasher.transform(X)
Using Local LLMs as an Arbiter/Classifier
June 2025
Using LLMs for doing arbitration or classification is nothing new. Generally I
found that “smaller” local LLMs struggled or led to overly optimistic results. I
think ‘smaller’ models are finally performant enough (to some extent at least)
where coming up with binary outcomes in a structured manner is “good enough”. In
general as of writing things still aren’t good enough. This is reflected in
things like the
Goose Blogpost
where 32gb of RAM is still generally recommended. Here, I use the very good
gemma-3-12b model to accomplish this – I generally found models smaller/older
than this failed to properly perform structured outputs in a meaningful manner.
The below example is based on Tricube Tales which is a TTRPG released under CC BY 3.0 license.
1import os
2from pathlib import Path
3from pprint import pprint
4from typing import Literal
5
6from llama_cpp import Llama
7from llama_cpp.llama_tokenizer import LlamaHFTokenizer
8from pydantic import BaseModel, model_validator
9
10model_id = str(
11 (
12 Path(os.environ.get("HOME", "~")) / "dev/google_gemma-3-12b-it-qat-Q4_0.gguf"
13 ).absolute()
14)
15
16prompt = """
17You are "The Arbiter", an impartial rules adjudicator for the RPG Tricube Tales.
18Your sole job is to determine which dice modifiers apply to a proposed action. Do not narrate outcomes or roll dice.
19
20Your judgments follow this logic:
21
22- Award +1 die if the action clearly aligns with the character's trait.
23 - The action must directly reflect the strengths of that trait:
24 - Brawny: feats of strength, endurance, or physical toughness
25 - Agile: speed, balance, dexterity, or evasive finesse
26 - Crafty: cleverness, deception, planning, or technical skill
27
28- Apply -1 die if the character's quirk meaningfully impedes the action.
29- Apply -1 die if the action is inherently challenging for most people in the setting, or the character lacks relevant expertise for the task
30
31These modifiers are cumulative. The resulting diceModifier may be: +1, 0, -1, or -2.
32
33Return your judgment using the following strict JSON schema:
34
35```json
36{
37 "trait_bonus": true | false,
38 "quirk_penalty": true | false,
39 "difficult": true | false,
40 "rationale": "concise human-readable explanation (30-70 words)"
41}
42"""
43
44
45class Arbiter(BaseModel):
46 trait_bonus: bool
47 quirk_penalty: bool
48 difficult: bool
49 rationale: str
50
51
52class DiceModifierArbiter(Arbiter):
53 dice_modifier: int | None = None
54
55 @model_validator(mode="after")
56 def validate_dice_modifier(self):
57 modifier = 0
58 if self.trait_bonus:
59 modifier += 1
60 if self.quirk_penalty:
61 modifier -= 1
62 if self.difficult:
63 modifier -= 1
64 self.dice_modifier = modifier
65 return self
66
67
68def create_chat_completion_with_schema(
69 messages: list[dict], model: Llama, schema: type[BaseModel], **kwargs
70) -> type[BaseModel]:
71 output = model.create_chat_completion(
72 messages=messages,
73 response_format={
74 "type": "json_object",
75 "schema": schema.model_json_schema(),
76 },
77 **kwargs,
78 )
79 return schema.model_validate_json(output["choices"][0]["message"]["content"])
80
81
82def determine_dice_modifier(context: str) -> DiceModifierArbiter:
83 llama_model = Llama(
84 model_path=model_id,
85 tokenizer=LlamaHFTokenizer.from_pretrained("gemma-3-12b"),
86 n_ctx=2048,
87 )
88 result = create_chat_completion_with_schema(
89 [
90 {"role": "system", "content": prompt},
91 {"role": "user", "content": context},
92 ],
93 llama_model,
94 Arbiter,
95 ) # type: ignore
96 return DiceModifierArbiter.model_validate(result.model_dump())
97
98
99if __name__ == "__main__":
100 result = []
101 result.append(
102 determine_dice_modifier("A brawny barbarian swings his sword at a goblin.")
103 )
104 result.append(
105 determine_dice_modifier("A brawny barbarian draws his bow at a goblin.")
106 )
107 result.append(
108 determine_dice_modifier("A brawny barbarian attempts to bargin with a vendor")
109 )
110 result.append(
111 determine_dice_modifier(
112 "Would a greedy, agile, thief successfully flee from a situation or get distracted with a treasure chest?"
113 )
114 )
115 pprint(result)
Output:
1 DiceModifierArbiter(trait_bonus=True, quirk_penalty=False, difficult=False, rationale="Swinging a sword is a direct application of brawny strength. There's no indication of a quirk impeding the action, and while combat is challenging, it's not inherently difficult for most warriors.", dice_modifier=1)
2
3 DiceModifierArbiter(trait_bonus=False, quirk_penalty=False, difficult=False, rationale="Drawing a bow is not inherently a feat of strength, endurance, or toughness. While barbarians might be proficient, it doesn't automatically grant a bonus. There's no indication of a quirk impeding the action, and drawing a bow isn't exceptionally difficult.", dice_modifier=0)
4
5 DiceModifierArbiter(trait_bonus=False, quirk_penalty=False, difficult=True, rationale="While a barbarian might have some persuasive ability, bargaining is not a direct application of Brawny. It is also inherently difficult for most people, requiring social finesse and negotiation skills, which are not typically a barbarian's forte.", dice_modifier=-1)
6
7 DiceModifierArbiter(trait_bonus=True, quirk_penalty=True, difficult=True, rationale="The thief's agility grants a bonus for fleeing. However, their greed introduces a penalty as they'd be tempted to grab treasure, hindering their escape. Fleeing a dangerous situation is generally difficult, warranting a further penalty.", dice_modifier=-1)
Using Qwen3-8B model also works relative well with this output:
1 DiceModifierArbiter(trait_bonus=True, quirk_penalty=False, difficult=False, rationale='Swinging a sword is a strength-based action, aligning with Brawny. No quirk or difficulty mentioned.', dice_modifier=1)
2
3 DiceModifierArbiter(trait_bonus=False, quirk_penalty=False, difficult=False, rationale='Drawing a bow is not a strength-based action. Brawny traits relate to melee, not ranged combat.', dice_modifier=0)
4
5 DiceModifierArbiter(trait_bonus=False, quirk_penalty=False, difficult=True, rationale='Brawny characters are physically strong but not necessarily skilled in negotiation. Bargaining is an inherently difficult task without relevant expertise.', dice_modifier=-1)
6
7 DiceModifierArbiter(trait_bonus=False, quirk_penalty=True, difficult=False, rationale="The thief's greed (quirk) would likely distract them from fleeing, outweighing their agility.", dice_modifier=-1)]
stdlib Numpy for when you don’t need all of Numpy
May 2025
Sometimes you just want the convenience of numpy-like notation for slicing, and setting values.
1from copy import deepcopy
2import json
3import re
4
5
6class Array:
7 def __init__(self, data):
8 # Accepts list, list of lists, or another Array
9 if isinstance(data, Array):
10 self.data = self._deepcopy(data.data)
11 else:
12 self.data = self._deepcopy(data)
13 self._update_shape()
14
15 def _deepcopy(self, data):
16 if isinstance(data, list):
17 if data and isinstance(data[0], list):
18 return [row[:] for row in data]
19 else:
20 return data[:]
21 return data
22
23 def _update_shape(self):
24 if isinstance(self.data, list):
25 if self.data and isinstance(self.data[0], list):
26 self.shape = (len(self.data), len(self.data[0]))
27 else:
28 self.shape = (len(self.data),)
29 else:
30 self.shape = ()
31
32 def __getitem__(self, key):
33 if len(self.shape) == 2:
34 if isinstance(key, tuple):
35 row_key, col_key = key
36 rows = self._process_slice(row_key, axis=0)
37 cols = self._process_slice(col_key, axis=1)
38 result = [[self.data[i][j] for j in cols] for i in rows]
39 if len(result) == 1 and len(result[0]) == 1:
40 return result[0][0]
41 if len(result) == 1:
42 return Array(result[0])
43 if len(result[0]) == 1:
44 return Array([row[0] for row in result])
45 return Array(result)
46 else:
47 # Single row
48 rows = self._process_slice(key, axis=0)
49 result = [self.data[i][:] for i in rows]
50 if len(result) == 1:
51 return Array(result[0])
52 return Array(result)
53 elif len(self.shape) == 1:
54 if isinstance(key, int):
55 idx = key
56 if idx < 0:
57 idx += self.shape[0]
58 if idx < 0 or idx >= self.shape[0]:
59 raise IndexError("index out of range")
60 return self.data[idx]
61 elif isinstance(key, slice) or isinstance(key, list):
62 idxs = self._process_slice(key, axis=0)
63 result = [self.data[i] for i in idxs]
64 return Array(result)
65 else:
66 raise TypeError(f"Invalid index type: {type(key)}")
67 else:
68 raise IndexError("Array is empty or has invalid shape.")
69
70 def __setitem__(self, key, value):
71 if len(self.shape) == 2:
72 if isinstance(key, tuple):
73 row_key, col_key = key
74 rows = self._process_slice(row_key, axis=0)
75 cols = self._process_slice(col_key, axis=1)
76 if isinstance(value, Array):
77 value = value.data
78 # Broadcasting for scalar assignment
79 if not isinstance(value, list) or (
80 value and not isinstance(value[0], list)
81 ):
82 value = [[value for _ in cols] for _ in rows]
83 elif value and not isinstance(value[0], list):
84 value = [value] * len(rows)
85 for i, row in enumerate(rows):
86 for j, col in enumerate(cols):
87 self.data[row][col] = value[i][j]
88 else:
89 rows = self._process_slice(key, axis=0)
90 if isinstance(value, Array):
91 value = value.data
92 if value and not isinstance(value[0], list):
93 value = [value] * len(rows)
94 for i, row in enumerate(rows):
95 self.data[row] = value[i]
96 elif len(self.shape) == 1:
97 if isinstance(key, int):
98 idx = key
99 if idx < 0:
100 idx += self.shape[0]
101 if idx < 0 or idx >= self.shape[0]:
102 raise IndexError("index out of range")
103 self.data[idx] = value
104 elif isinstance(key, slice) or isinstance(key, list):
105 idxs = self._process_slice(key, axis=0)
106 if isinstance(value, Array):
107 value = value.data
108 # Broadcasting for scalar assignment
109 if not isinstance(value, list):
110 value = [value] * len(idxs)
111 if len(value) != len(idxs):
112 raise ValueError("could not broadcast input array to shape")
113 for i, idx in enumerate(idxs):
114 self.data[idx] = value[i]
115 else:
116 raise TypeError(f"Invalid index type: {type(key)}")
117 else:
118 raise IndexError("Array is empty or has invalid shape.")
119
120 def _process_slice(self, key, axis):
121 n = self.shape[axis]
122 if isinstance(key, int):
123 if key < 0:
124 key += n
125 return [key]
126 elif isinstance(key, slice):
127 return list(range(*key.indices(n)))
128 elif isinstance(key, list):
129 return [(k + n if k < 0 else k) for k in key]
130 else:
131 raise TypeError(f"Invalid index type: {type(key)}")
132 def __repr__(self):
133 data_copy = deepcopy(self.data)
134 max_string_length = 0
135 is_numeric_data = True
136
137 def process_value(val):
138 nonlocal is_numeric_data
139 str_val = json.dumps(val)
140 if not isinstance(val, (int, float)):
141 is_numeric_data = False
142 return str_val
143
144 def format_value(str_val):
145 if is_numeric_data:
146 return str_val.rjust(max_string_length).replace(" ", "_")
147 return str_val.ljust(max_string_length)
148
149 if len(self.shape) == 2:
150 # Process 2D array values
151 for i, row in enumerate(data_copy):
152 for j, col in enumerate(row):
153 data_copy[i][j] = process_value(col)
154 max_string_length = max(max_string_length, len(data_copy[i][j]))
155
156 # Format values
157 for i, row in enumerate(data_copy):
158 for j, col in enumerate(row):
159 data_copy[i][j] = format_value(data_copy[i][j])
160
161 # Format output
162 output = json.dumps(data_copy, indent=2)
163 output = output.replace('"', "")
164 output = re.sub(r"\n\s*(?=[^\s\[]|])", " ", output)
165 if is_numeric_data:
166 output = output.replace("_", " ")
167 output = re.sub(r"\]\s*\]", "]\n]", output)
168
169 else: # 1D array
170 # Process 1D array values
171 for i, val in enumerate(data_copy):
172 data_copy[i] = process_value(val)
173 max_string_length = max(max_string_length, len(data_copy[i]))
174
175 # Format values
176 for i, val in enumerate(data_copy):
177 data_copy[i] = format_value(data_copy[i])
178
179 # Format output
180 output = json.dumps(data_copy)
181 output = output.replace('"', "")
182 if is_numeric_data:
183 output = output.replace("_", " ")
184
185 return output
186
187
188def array(data):
189 return Array(data)
190
191
192def zeros(shape):
193 if isinstance(shape, int):
194 return Array([0 for _ in range(shape)])
195 elif len(shape) == 1:
196 return Array([0 for _ in range(shape[0])])
197 elif len(shape) == 2:
198 rows, cols = shape
199 return Array([[0 for _ in range(cols)] for _ in range(rows)])
200 else:
201 raise ValueError("Only 1D and 2D arrays are supported.")
202
203
204def ones(shape):
205 if isinstance(shape, int):
206 return Array([1 for _ in range(shape)])
207 elif len(shape) == 1:
208 return Array([1 for _ in range(shape[0])])
209 elif len(shape) == 2:
210 rows, cols = shape
211 return Array([[1 for _ in range(cols)] for _ in range(rows)])
212 else:
213 raise ValueError("Only 1D and 2D arrays are supported.")
Prompting ChatGPT to Generate Prompts
March 2025
Okay, so this isn’t that “new” though sometimes its hard to find a “winner”. This is one that I like for writing code reviews.
Prompt
Write me a chatgpt prompt for reviewing code PRs specifically off git diff. that does the following:
- Summarizes the changes
- Code quality and adherence to best practices
- Potential bugs or edge cases
- Performance optimizations
- Readability and maintainability
- Any security concerns
Suggest improvements and explain your reasoning for each suggestion. Keep each element short and succint
At the end of the review it should also give a rating on the estimated risk of merge
The generated output:
Here’s a solid ChatGPT prompt for reviewing a PR using a git diff:
Prompt:
You are a highly skilled code reviewer analyzing a git diff. Your task is to review the changes and provide a structured response covering:
1. Summary of Changes: Briefly describe what was modified.
2. Code Quality & Best Practices: Note any deviations from coding standards.
3. Potential Bugs & Edge Cases: Highlight any potential issues.
4. Performance Optimizations: Suggest improvements for efficiency.
5. Readability & Maintainability: Evaluate how easy it is to understand and extend.
6. Security Concerns: Identify vulnerabilities or security risks.
7. Suggested Improvements: Provide specific recommendations and reasoning.
At the end, provide an Estimated Merge Risk rating (Low, Medium, High) based on potential issues and complexity. Keep feedback concise but actionable.
Input: Provide the git diff for analysis.
What is interesting is how near (or far?) it is from curated offerings. For example this is the one from continue.dev
<system>
You will be acting as a senior software engineer performing a code review for a colleague.
</system>
You will follow the guidelines for giving a great code review outlined below:
{{{ url "https://google.github.io/eng-practices/review/reviewer/looking-for.html" }}}
---
Here is the proposed code changes you will be reviewing:
{{{ diff }}}
---
Do not include a greeting. Immediately begin reviewing the changes.
For each file, decide if you need to provide any feedback on the changes.
If so, outline the feedback using one or two sentences.
If a code change is required, then mention the original code, and
then propose a code change to fix it.
Do not add any other text after the suggestion.
If you have no feedback on a file, do not add a comment for that file.
Lastly, provide a one to two summary of your feedback at the end.
Here are some examples.
<example>
### filename.js
The name of this variable is unclear.
Original:
```js
const x = getAllUsers();
```
Suggestion:
```js
const allUsers = getAllUsers();
```
</example>
<example>
### filename.js
This code is overly complex.
Original:
```py
class AgeCalculator:
def __init__(self, birth_year):
self.birth_year = birth_year
def calculate_age(self, current_year):
age = current_year - self.birth_year
return self._validate_and_format_age(age)
def _validate_and_format_age(self, age):
if age < 0:
raise ValueError("Invalid age calculated")
return f"User is {age} years old"
def get_user_age(birth_year, current_year):
calculator = AgeCalculator(birth_year)
return calculator.calculate_age(current_year)
```
Suggestion:
```python
def get_user_age(birth_year, current_year):
return current_year - birth_year
```
</example>
<example>
### Summary
Overall, these changes appear to be minor improvements to the
project structure and code cleanliness.
</example>
Here is the additional input from the code author:
<input>
{{ input }}
</input>
Think through your feedback step by step before replying.
IsoForest with k-NN
February 2025
Following on the general theme of using random-ness to infer or generate predictions, I’ve always been very curious about isolation forests, more specifically given the innate measure of “similarity” or “distance” can we use this in the supervised learning scenario via k-NN?
Below is some code which implements this completely using stdlib Python with no
dependencies on numpy. Written with heavy assistance from LLMs.
1import random
2import math
3from collections import Counter
4
5# ----------------------
6# 1) Isolation Forest
7# ----------------------
8
9
10class IsolationTree:
11 """
12 A single isolation tree.
13 """
14
15 def __init__(self, max_depth, seed=None):
16 self.max_depth = max_depth
17 # We'll store a seed to replicate the random picks
18 # If None, we use the global random generator
19 self.seed = seed if seed is not None else random.randint(0, 2**31 - 1)
20
21 self.split_feature = None
22 self.split_value = None
23 self.left = None
24 self.right = None
25 self.size = 0 # number of samples in this node (for leaf nodes)
26
27 def fit(self, data, depth=0, rng=None):
28 """
29 Recursively build the isolation tree.
30
31 :param data: List of data points (each a list of features).
32 :param depth: Current depth of the tree.
33 :param rng: A random.Random instance. If None, we build one using self.seed.
34 """
35 if rng is None:
36 rng = random.Random(self.seed) # create a local RNG
37
38 n = len(data)
39 self.size = n
40
41 # Base case: stop if max depth reached or not enough samples
42 if depth >= self.max_depth or n <= 1:
43 return
44
45 # Number of features
46 num_features = len(data[0])
47
48 # Randomly pick a feature to split on
49 self.split_feature = rng.randint(0, num_features - 1)
50
51 # Get min and max of that feature
52 feature_values = [point[self.split_feature] for point in data]
53 min_val = min(feature_values)
54 max_val = max(feature_values)
55
56 # If all points are the same on this feature, stop splitting
57 if min_val == max_val:
58 self.split_feature = None
59 return
60
61 # Randomly pick a split value
62 split_val = rng.uniform(min_val, max_val)
63 self.split_value = split_val
64
65 # Partition data into left and right
66 left_data = []
67 right_data = []
68 for point in data:
69 if point[self.split_feature] < split_val:
70 left_data.append(point)
71 else:
72 right_data.append(point)
73
74 # Create child nodes
75 self.left = IsolationTree(self.max_depth)
76 self.right = IsolationTree(self.max_depth)
77
78 # Recursively fit child nodes (reuse the same rng for consistency)
79 self.left.fit(left_data, depth + 1, rng)
80 self.right.fit(right_data, depth + 1, rng)
81
82 def path_length(self, x, depth=0):
83 """
84 Compute the path length of a sample x in this isolation tree.
85 Includes the c_factor for leaf nodes.
86 """
87 # If leaf node or no further split
88 if self.split_feature is None or self.left is None or self.right is None:
89 return depth + c_factor(self.size)
90
91 # Check which branch x goes to
92 if x[self.split_feature] < self.split_value:
93 return self.left.path_length(x, depth + 1)
94 else:
95 return self.right.path_length(x, depth + 1)
96
97 def shared_depth(self, x, y, depth=0):
98 """
99 Compute how many levels x and y share in the same branch
100 before diverging (or reaching a leaf).
101 """
102 # If we are at a leaf node or no further split
103 if self.split_feature is None or self.left is None or self.right is None:
104 return depth
105
106 x_left = x[self.split_feature] < self.split_value
107 y_left = y[self.split_feature] < self.split_value
108
109 if x_left != y_left:
110 return depth
111
112 if x_left:
113 return self.left.shared_depth(x, y, depth + 1)
114 else:
115 return self.right.shared_depth(x, y, depth + 1)
116
117
118def harmonic_number(n):
119 """
120 Return the nth harmonic number H_n = 1 + 1/2 + 1/3 + ... + 1/n
121 """
122 return sum(1.0 / i for i in range(1, n + 1))
123
124
125def c_factor(n):
126 """
127 Normalization factor commonly used in isolation forest scoring:
128 c(n) = 2 * H_{n-1} - 2*(n-1)/n
129 For large n, c(n) ~ 2 * ln(n-1) + gamma - 2*(n-1)/n
130 """
131 if n <= 1:
132 return 0
133 return 2.0 * harmonic_number(n - 1) - 2.0 * (n - 1) / n
134
135
136class IsolationForest:
137 """
138 An Isolation Forest: collection of IsolationTrees.
139 """
140
141 def __init__(self, n_estimators=10, max_samples=256, max_depth=8, random_seed=None):
142 """
143 :param n_estimators: Number of isolation trees.
144 :param max_samples: Subsample size for each tree.
145 :param max_depth: Maximum depth of each tree.
146 :param random_seed: Optional integer seed for reproducibility.
147 """
148 self.n_estimators = n_estimators
149 self.max_samples = max_samples
150 self.max_depth = max_depth
151 self.trees = []
152
153 # For partial_fit, we want a random generator we can use each time
154 # to pick which trees to replace, etc.
155 if random_seed is not None:
156 self._rng = random.Random(random_seed)
157 else:
158 self._rng = random.Random()
159
160 def fit(self, data):
161 """
162 Train the isolation forest on the given dataset.
163
164 :param data: List of data points (each a list of features).
165 """
166 self.trees = []
167 n = len(data)
168
169 for _ in range(self.n_estimators):
170 # Subsample the data
171 if n > self.max_samples:
172 subset = random.sample(data, self.max_samples)
173 else:
174 subset = data[:]
175
176 tree = IsolationTree(self.max_depth)
177 tree.fit(subset)
178 self.trees.append(tree)
179
180 def path_length(self, x):
181 """
182 Compute the average path length of x across all trees.
183
184 :param x: A single data point (list of features).
185 :return: Average path length as a float.
186 """
187 total_path = 0.0
188 for tree in self.trees:
189 total_path += tree.path_length(x)
190 return total_path / float(len(self.trees))
191
192 def anomaly_score(self, x):
193 """
194 Compute the anomaly (outlier) score for a single sample x.
195
196 A common formula is:
197 score(x) = 2^(- E(path_length) / c_factor(max_samples))
198 """
199 avg_path = self.path_length(x)
200 # Classical isolation score:
201 return math.pow(2, -avg_path / c_factor(self.max_samples))
202
203 def predict(self, data, threshold=0.5):
204 """
205 Classify data points as -1 (outlier) or 1 (inlier) based on a threshold.
206 """
207 preds = []
208 for x in data:
209 score = self.anomaly_score(x)
210 preds.append(-1 if score > threshold else 1)
211 return preds
212
213 def similarity_score(self, x, y):
214 """
215 Compute a similarity score between two points x and y, based on
216 shared depth across all trees (normalized by max_depth).
217
218 The higher the score, the more 'similar' the points are
219 in the sense of how the forest partitions the space.
220 """
221 total_shared_depth = 0
222 for tree in self.trees:
223 total_shared_depth += tree.shared_depth(x, y, depth=0)
224
225 avg_shared_depth = total_shared_depth / float(len(self.trees))
226 # Normalize by the maximum possible depth
227 similarity = avg_shared_depth / float(self.max_depth)
228 return similarity
229
230 def partial_fit(self, data, n_replace=1):
231 """
232 *Heuristic* partial fit that:
233 1) Takes the new minibatch 'data'
234 2) Randomly picks 'n_replace' trees to remove
235 3) Builds 'n_replace' new trees on the minibatch (up to max_samples)
236 reusing seeds from the removed trees if desired
237 """
238 if not self.trees:
239 # If forest was empty, just do a fresh fit with n_estimators
240 self.fit(data)
241 return
242
243 # Bound n_replace by the current number of trees
244 n_replace = min(n_replace, len(self.trees))
245
246 # 1) Randomly pick which trees to remove
247 remove_indices = self._rng.sample(range(len(self.trees)), n_replace)
248 remove_indices = set(remove_indices)
249
250 # 2) For each tree to remove, gather its seed to reuse
251 removed_seeds = []
252 kept_trees = []
253 for i, tree in enumerate(self.trees):
254 if i in remove_indices:
255 removed_seeds.append(tree.seed)
256 else:
257 kept_trees.append(tree)
258
259 # 3) Build n_replace new trees on the minibatch
260 # If minibatch larger than max_samples, sample
261 n_data = len(data)
262 if n_data > self.max_samples:
263 subset = self._rng.sample(data, self.max_samples)
264 else:
265 subset = data[:]
266
267 new_trees = []
268 for seed in removed_seeds:
269 # We create a new tree reusing the old tree's seed
270 new_tree = IsolationTree(self.max_depth, seed=seed)
271 new_tree.fit(subset)
272 new_trees.append(new_tree)
273
274 # 4) Update self.trees
275 self.trees = kept_trees + new_trees
276
277 def __len__(self):
278 """
279 Number of trees in the forest.
280 """
281 return len(self.trees)
282
283
284# -------------------------------------
285# 2) IsoKNNClassifier (Meta-Learner)
286# -------------------------------------
287
288
289class IsoKNNClassifier:
290 """
291 A classifier that uses Isolation Forest to define a similarity metric,
292 then applies a k-Nearest Neighbors approach based on highest similarity.
293 """
294
295 def __init__(self, k=3, iso_forest_params=None):
296 """
297 :param k: Number of neighbors to consider.
298 :param iso_forest_params: Dictionary of parameters for IsolationForest.
299 """
300 self.k = k
301 if iso_forest_params is None:
302 iso_forest_params = {}
303 self.iso_forest = IsolationForest(**iso_forest_params)
304 self.X_train = []
305 self.y_train = []
306
307 def fit(self, X, y):
308 """
309 Fit the IsoKNN model: store training data and build the Isolation Forest.
310
311 :param X: List of feature vectors.
312 :param y: List of corresponding labels (same order).
313 """
314 # Make sure X and y are in a usable structure
315 self.X_train = X
316 self.y_train = y
317
318 # Train the Isolation Forest on X
319 self.iso_forest.fit(X)
320
321 def predict(self, X_test):
322 """
323 Predict labels for a list of test points using the IsoForest-based kNN.
324
325 :param X_test: List of feature vectors for testing.
326 :return: List of predicted labels.
327 """
328 predictions = []
329 for x in X_test:
330 # 1. Compute similarity to each training point
331 similarities = []
332 for i, x_train in enumerate(self.X_train):
333 sim = self.iso_forest.similarity_score(x, x_train)
334 similarities.append((sim, self.y_train[i]))
335
336 # 2. Sort by similarity in descending order
337 similarities.sort(key=lambda tup: tup[0], reverse=True)
338
339 # 3. Get top-k neighbors
340 top_k = similarities[: self.k]
341
342 # 4. Majority vote among neighbors
343 labels = [label for (_, label) in top_k]
344 label_count = Counter(labels)
345 # Pick the label with the highest count; break ties arbitrarily
346 majority_label = label_count.most_common(1)[0][0]
347 predictions.append(majority_label)
348
349 return predictions
350
351
352# ----------------
353# Example Usage
354# ----------------
355
356
357def main():
358 # Synthetic training data
359 X_train = [
360 [0.1, 0.2], # label 0
361 [0.2, 0.1], # label 0
362 [0.15, 0.18], # label 0
363 [5.0, 6.0], # label 1
364 [5.2, 5.9], # label 1
365 [10.0, 10.0], # label 2
366 ]
367 y_train = [0, 0, 0, 1, 1, 2]
368
369 # Define test points
370 X_test = [
371 [0.15, 0.17], # Close to first cluster => likely label 0
372 [5.1, 6.1], # Close to second cluster => likely label 1
373 [10.1, 9.9], # Close to the outlier => likely label 2
374 [4.9, 5.8], # Possibly near cluster (5.0, 6.0) => label 1
375 ]
376
377 # Create and fit IsoKNNClassifier
378 # We'll use an Isolation Forest with a few trees for demonstration.
379 iso_knn = IsoKNNClassifier(
380 k=2,
381 iso_forest_params={
382 "n_estimators": 5,
383 "max_samples": 4,
384 "max_depth": 5,
385 "random_seed": 42,
386 },
387 )
388 iso_knn.fit(X_train, y_train)
389
390 # Predict
391 preds = iso_knn.predict(X_test)
392 for x, label in zip(X_test, preds):
393 print(f"Test point {x} => predicted label: {label}")
394
395
396def demo_partial_fit():
397 """
398 Minimal demonstration of partial_fit usage.
399 """
400 # Original data
401 data_initial = [
402 [0.1, 0.2],
403 [0.2, 0.1],
404 [0.15, 0.18],
405 [5.0, 6.0],
406 [5.2, 5.9],
407 [4.9, 6.2],
408 ]
409
410 # New data (minibatch) to incorporate
411 data_minibatch = [[10.0, 10.0], [10.1, 9.9]]
412
413 # Build an initial forest
414 iso_f = IsolationForest(n_estimators=4, max_depth=5, max_samples=3, random_seed=42)
415 iso_f.fit(data_initial)
416 print("Initial number of trees =", len(iso_f))
417
418 # Check anomaly scores before partial_fit
419 outlier_score_before = iso_f.anomaly_score([10.0, 10.0])
420
421 # partial_fit with new data, replacing 2 trees
422 iso_f.partial_fit(data_minibatch, n_replace=2)
423 print("Number of trees after partial_fit =", len(iso_f))
424
425 # Check anomaly scores after partial_fit
426 outlier_score_after = iso_f.anomaly_score([10.0, 10.0])
427
428 print(f"Outlier score (before) = {outlier_score_before:.3f}")
429 print(f"Outlier score (after) = {outlier_score_after:.3f}")
430
431
432if __name__ == "__main__":
433 main()
434 print("== Partial fit demo ==")
435 demo_partial_fit()
Commentary from the future: this approach is very similar to Mondrian Trees, in which the random partitions are induced by a Mondrian Process. I should include a “fast approximation” to Mondrian processes and expand it to multi-dimensions using standard lib in the future.
Brownian Bridges
February 2025
I got LLM via prompting to recreate how it may have generated the simple kriging diagram on the wikipedia page. I intentionally got it to do it via brownian bridges rather than the “typical” guassian process way simply because I was interested in brownian bridges. I may have made changes for ease of implementation (I was thinking about re-implementing it using stdlib or in another language as a close approximation), though the “idea” remains. This technique maybe useful if one wants to do something better than sampling over a grid search.
To extend to multiple dimensions, we can make use of tree-parzen like approach to sample per region (chosen heuristically) and presume univariate relationships (which was the approach taken in the original paper, though later approaches used multivariate KDEs). For the purposes of a simplistic implementation, simplifying assumptions will suffice. For example, it could be that the tree-structure restricts sampling to a fixed number of points per a region to determine whether or not a split is necessary.
1# %%
2import numpy as np
3from numpy.random import default_rng
4from numpy.random import default_rng
5import matplotlib.pyplot as plt
6
7def brownian_realisation(x_obs, y_obs, x_pred, sigma=1.0, mean=0.0, rng=None):
8 """
9 Generate one realization of the Gaussian process (Brownian motion
10 with variance scale sigma^2) at the locations x_pred,
11 conditioned on the data (x_obs, y_obs).
12
13 Parameters
14 ----------
15 x_obs : array-like
16 Observed locations (1D).
17 y_obs : array-like
18 Observed values corresponding to x_obs.
19 x_pred: array-like
20 Locations where we want a realization.
21 sigma : float, optional
22 Standard deviation parameter for Brownian motion. Cov = sigma^2 * min(t1, t2).
23 mean : float or array-like, optional
24 Global mean (scalar or an array the same length as x_obs/x_pred).
25 rng : np.random.Generator or None
26 Random generator for reproducibility. If None, a new default generator is used.
27
28 Returns
29 -------
30 y_pred_samp : ndarray
31 One random sample (realization) of the GP at x_pred.
32 """
33 if rng is None:
34 rng = default_rng()
35
36 # Convert to arrays
37 x_obs = np.asarray(x_obs, dtype=float)
38 y_obs = np.asarray(y_obs, dtype=float)
39 x_pred = np.asarray(x_pred, dtype=float)
40
41 # Sort so everything is in ascending order (not strictly required, but good practice)
42 # (If your x_obs, x_pred come sorted, this step can be skipped or adapted.)
43 idx_obs_sort = np.argsort(x_obs)
44 x_obs = x_obs[idx_obs_sort]
45 y_obs = y_obs[idx_obs_sort]
46
47 # offset by the initial point as it is assumed to be 0
48 y_obs_initial = y_obs[0]
49 y_obs = y_obs - y_obs_initial
50
51 idx_pred_sort = np.argsort(x_pred)
52 x_pred_sorted = x_pred[idx_pred_sort]
53
54 # Build entire (obs+pred) set of x's
55 x_all = np.concatenate([x_obs, x_pred_sorted])
56
57 # Build full covariance
58 K_all = brownian_covariance(x_all, sigma=sigma)
59
60 # Partition covariance
61 n_obs = len(x_obs)
62 n_pred = len(x_pred_sorted)
63
64 K_oo = K_all[:n_obs, :n_obs] # Cov among observed points
65 K_op = K_all[:n_obs, n_obs:] # Cov between obs and pred
66 K_po = K_all[n_obs:, :n_obs] # Cov between pred and obs
67 K_pp = K_all[n_obs:, n_obs:] # Cov among pred points
68
69 # Mean vectors
70 # We'll assume 'mean' is a scalar. If it's not, adapt below for shape
71 m_o = mean * np.ones(n_obs)
72 m_p = mean * np.ones(n_pred)
73
74 # Compute the conditional mean
75 # y_obs - m_o is the 'centered' observed data
76 K_oo_inv = np.linalg.pinv(K_oo)
77 cond_mean = m_p + K_po @ K_oo_inv @ (y_obs - m_o)
78
79 # Compute the conditional covariance
80 cond_cov = K_pp - K_po @ K_oo_inv @ K_op
81
82 # Draw one realization from N(cond_mean, cond_cov)
83 y_pred_sample = rng.multivariate_normal(mean=cond_mean, cov=cond_cov)
84
85 # Reorder to match the original order of x_pred
86 y_pred_samp_original_order = np.empty_like(y_pred_sample)
87 y_pred_samp_original_order[idx_pred_sort] = y_pred_sample
88
89 return y_pred_samp_original_order + y_obs_initial
90
91
92def brownian_bridge_realisation(
93 x_obs,
94 y_obs,
95 x_pred,
96 sigma=1.0,
97 mean=0.0,
98 rng=None
99):
100 """
101 Generate a 1D Brownian motion sample that is 'pinned' to observed data
102 at certain locations (times) WITHOUT doing the full covariance-matrix inversion.
103
104 This uses the Brownian-bridge property on each segment between consecutive
105 observed points. In 1D, this produces the exact conditional sample:
106
107 Y(x_obs[i]) = y_obs[i] for all observed i,
108
109 and at any intermediate x_pred in that segment, Y is drawn from a
110 Brownian bridge. This avoids the NxN matrix inversion from the usual
111 'brownian_realisation' approach.
112
113 Parameters
114 ----------
115 x_obs : array-like
116 The observed (pinned) locations (1D), shape (n_obs,).
117 y_obs : array-like
118 The observed values at x_obs, shape (n_obs,).
119 x_pred: array-like
120 The locations (times) where we want to simulate the Brownian motion.
121 sigma : float, optional
122 Brownian motion scale parameter. Covariance is sigma^2 * min(t1, t2).
123 Equivalently, increments have variance sigma^2 * (t_{i+1} - t_i).
124 mean : float, optional
125 Global constant mean for the process. If not 0, we effectively
126 do Brownian bridging on (y_obs - mean) and then add 'mean' at the end.
127 rng : np.random.Generator or None
128 NumPy random generator for reproducibility. If None, use default_rng().
129
130 Returns
131 -------
132 y_pred : ndarray
133 A single realization of the Brownian motion at x_pred, shaped like x_pred.
134 """
135
136 if rng is None:
137 rng = default_rng()
138
139 # Convert inputs to arrays
140 x_obs = np.asarray(x_obs, dtype=float)
141 y_obs = np.asarray(y_obs, dtype=float)
142 x_pred = np.asarray(x_pred, dtype=float)
143
144 # 1) Sort the observed data by x_obs
145 sort_obs_idx = np.argsort(x_obs)
146 x_obs_sorted = x_obs[sort_obs_idx]
147 y_obs_sorted = y_obs[sort_obs_idx]
148
149 # 2) Sort x_pred
150 sort_pred_idx = np.argsort(x_pred)
151 x_pred_sorted = x_pred[sort_pred_idx]
152
153 # 3) Merge (x_obs_sorted) and (x_pred_sorted) into a single sorted array x_all
154 # We'll keep all unique points but preserve order.
155 x_all = np.unique(np.concatenate([x_obs_sorted, x_pred_sorted]))
156
157 # We'll create an array to hold the entire path: Y(x_all)
158 # For pinned points (observed), we fix Y. For the rest, we fill by bridging.
159 y_all = np.empty_like(x_all, dtype=float)
160
161 # A helper dictionary for pinned values: pinned_dict[x_obs[i]] = y_obs[i]
162 # We'll store 'centered' pinned values by subtracting the global mean
163 pinned_dict = {}
164 for xo, yo in zip(x_obs_sorted, y_obs_sorted):
165 pinned_dict[xo] = yo - mean
166
167 # Mark which of x_all are pinned vs. not
168 pinned_mask = np.array([ (x in pinned_dict) for x in x_all ])
169
170 # Fill in pinned points in y_all
171 y_all[:] = np.nan
172 for i, xval in enumerate(x_all):
173 if pinned_mask[i]:
174 y_all[i] = pinned_dict[xval]
175
176 # Indices of pinned points
177 pinned_indices = np.where(pinned_mask)[0]
178
179 # If there are no pinned points, or only one pinned point,
180 # you might define boundary conditions or do an unconditional random walk.
181 # For a minimal demonstration, let's handle segments only if we have >= 2 pins.
182 if len(pinned_indices) == 0:
183 # No data at all: unconditionally simulate from 0 up to max(x_all)
184 # For a pure Brownian motion with mean=0, we can do increments in ascending x_all.
185 # Then add 'mean' at the end.
186 y_all[0] = 0.0
187 for i in range(1, len(x_all)):
188 dt = x_all[i] - x_all[i-1] # time step
189 # Increments ~ Normal(0, sigma^2 * dt)
190 incr = rng.normal(loc=0.0, scale=sigma*np.sqrt(dt))
191 incr = rng.uniform(low=-sigma*np.sqrt(dt), high=sigma*np.sqrt(dt))
192 y_all[i] = y_all[i-1] + incr
193 y_all += mean
194
195 else:
196 # We'll do piecewise bridging from one pinned index to the next
197 # If pinned_indices doesn't include the first or last index,
198 # we might do an unconditional portion before the first pin or after the last pin.
199
200 # Handle any region before the first pinned index (unconditional from x_all[0] to that pinned point)
201 first_pin = pinned_indices[0]
202 if first_pin > 0:
203 # From i=0 up to i=first_pin
204 # We'll do a "bridge" from (x_all[0], random start) to (x_all[first_pin], pinned_value)
205 # If you want a known start, set y_all[0] = 0 or something. We'll just do it unconditionally
206 # until we reach the pinned point, so let's pick y_all[0] = pinned_value as a hack, or 0.
207 # We'll do an unconditional path and then shift so that at i=first_pin we get pinned.
208 #
209 # Simpler approach:
210 # y_all[0] = pinned_value_of_the_first_pin to avoid a big jump,
211 # or
212 # start from 0.0 and see where we end up, then shift everything.
213 # We'll do something straightforward for demonstration:
214 y_all[0] = y_all[first_pin] # pin the start to the same y so variance is only from increments
215
216 for i in range(1, first_pin+1):
217 dt = x_all[i] - x_all[i-1]
218 incr = rng.normal(loc=0.0, scale=sigma*np.sqrt(dt))
219 incr = rng.uniform(low=-sigma*np.sqrt(dt), high=sigma*np.sqrt(dt))
220 y_all[i] = y_all[i-1] + incr
221
222 # Now forcibly set the pinned index to pinned value
223 y_all[first_pin] = pinned_dict[x_all[first_pin]]
224
225 # Now do bridging on each segment between consecutive pinned points
226 for seg_start, seg_end in zip(pinned_indices, pinned_indices[1:]):
227 # We have x_left, y_left and x_right, y_right
228 x_left = x_all[seg_start]
229 y_left = y_all[seg_start] # pinned value
230 x_right = x_all[seg_end]
231 y_right = y_all[seg_end] # pinned value
232
233 # The interior points are seg_start+1 .. seg_end-1
234 # We'll fill them in a forward pass using the Brownian bridge increments
235 for i in range(seg_start+1, seg_end):
236 t_i = x_all[i-1]
237 y_i = y_all[i-1]
238 t_next = x_all[i]
239
240 # Markov property for Brownian bridge:
241 # Y(t_next) | Y(t_i)=y_i, Y(x_right)=y_right is normal with:
242 # mean = y_i + ((t_next - t_i)/(x_right - t_i)) * (y_right - y_i)
243 # var = sigma^2 * (t_next - t_i) * (x_right - t_next) / (x_right - t_i)
244 #
245 dt = (t_next - t_i)
246 denom = (x_right - t_i)
247
248 # Mean increment (drift toward y_right)
249 mean_inc = (dt/denom) * (y_right - y_i)
250
251 # Variance of that increment
252 var_inc = sigma**2 * dt * (x_right - t_next) / denom
253 sd_inc = np.sqrt(var_inc)
254
255 # Sample from Normal( mean_inc, sd_inc^2 )
256 inc = rng.normal(loc=mean_inc, scale=sd_inc)
257 inc = rng.uniform(low=mean_inc-sd_inc, high=mean_inc + sd_inc)
258
259 y_all[i] = y_i + inc
260
261 # do we forcibly set end point to pinned?
262 # y_all[seg_end] = y_right
263
264 # Handle any region after the last pinned index
265 last_pin = pinned_indices[-1]
266 if last_pin < len(x_all) - 1:
267 # unconditional from that pinned point to the end
268 for i in range(last_pin+1, len(x_all)):
269 dt = x_all[i] - x_all[i-1]
270 incr = rng.normal(loc=0.0, scale=sigma*np.sqrt(dt))
271 incr = rng.uniform(low=-sigma*np.sqrt(dt), high=sigma*np.sqrt(dt))
272 y_all[i] = y_all[i-1] + incr
273
274 # Finally, add 'mean' back in, since we subtracted it from y_obs
275 y_all += mean
276
277 # 4) We only need to return values at x_pred
278 # So we'll make a dict from x_all -> y_all
279 val_dict = {x: y for (x, y) in zip(x_all, y_all)}
280
281 # 5) Extract the results in the original order of x_pred
282 y_pred_out = np.array([val_dict[x] for x in x_pred_sorted])
283
284 # Re-map to the original (unsorted) x_pred order
285 y_pred = np.empty_like(y_pred_out)
286 y_pred[sort_pred_idx] = y_pred_out
287
288 return y_pred
289
290
291# %%
292def simple_kriging(x_obs, y_obs, x_pred, n_reals=1000, sigma=1.0, mean=0.0, rng=None, use_random=False):
293 """
294 Perform simple kriging under Brownian motion assumption by Monte Carlo.
295
296 1. Draw multiple conditional realizations from the GP.
297 2. Compute empirical mean and credible intervals from those realizations.
298
299 Parameters
300 ----------
301 x_obs : array-like
302 Observed locations.
303 y_obs : array-like
304 Observed values.
305 x_pred: array-like
306 Prediction locations.
307 n_reals : int
308 Number of Monte Carlo realizations to draw.
309 sigma : float
310 Standard deviation parameter for Brownian motion.
311 mean : float
312 Global mean assumption for "simple" kriging.
313 rng : np.random.Generator or None
314 Random generator for reproducibility.
315
316 Returns
317 -------
318 pred_mean : ndarray
319 The empirical mean of the GP at x_pred over all realizations.
320 pred_lower : ndarray
321 The lower bound of the 95% credible interval (2.5th percentile).
322 pred_upper : ndarray
323 The upper bound of the 95% credible interval (97.5th percentile).
324 all_realizations : ndarray
325 All realizations of shape (n_reals, len(x_pred)).
326 """
327 if rng is None:
328 rng = default_rng()
329
330 x_pred = np.asarray(x_pred, dtype=float)
331 all_realizations = []
332
333 # Draw n_reals times
334 for _ in range(n_reals):
335 sample = brownian_bridge_realisation(
336 x_obs, y_obs, x_pred, sigma=sigma, mean=mean, rng=rng
337 )
338 all_realizations.append(sample)
339
340 all_realizations = np.array(all_realizations) # shape (n_reals, n_pred)
341
342 # Empirical stats across realizations
343 pred_mean = np.mean(all_realizations, axis=0)
344 pred_lower = np.percentile(all_realizations, 2.5, axis=0)
345 pred_upper = np.percentile(all_realizations, 97.5, axis=0)
346
347 return pred_mean, pred_lower, pred_upper, all_realizations
348
349
350# %%
351# Observed data
352x_obs = np.array([0.0, 2.0, 3.0, 7.0, 8.0])
353y_obs = np.array([1.4, 1.2, 0.9, 2.5, 2.2])
354
355# Points where we want to predict
356x_pred = np.linspace(0, 8, 50) # 50 points from 0 to 8
357
358# Perform simple kriging via Brownian motion realisations
359pred_mean, pred_lower, pred_upper, reals = simple_kriging(
360 x_obs, y_obs, x_pred,
361 n_reals=2000,
362 sigma=1.0, # if you believe the process has variance sigma^2
363 mean=0.0 # simple kriging with known mean 0
364)
365
366# Print some results
367# print("Prediction Mean:", pred_mean)
368# print("95% Credible Interval Lower:", pred_lower)
369# print("95% Credible Interval Upper:", pred_upper)
370
371# Plotting (optional) -- requires matplotlib
372import matplotlib.pyplot as plt
373
374plt.figure(figsize=(8,5))
375# Observed data
376plt.plot(x_obs, y_obs, 'ko', label='Observations')
377# Mean
378plt.plot(x_pred, pred_mean, 'b-', label='Posterior mean')
379
380# print some realisations
381plt.plot(x_pred, reals[0, :], '-', label='A single realization')
382
383# 95% Credible Interval
384plt.fill_between(x_pred, pred_lower, pred_upper, color='blue', alpha=0.2, label='95% CI')
385plt.legend()
386plt.title("Simple Kriging under Brownian Motion Assumption")
387plt.xlabel("x")
388plt.ylabel("y")
389plt.show()
