PM-pinou
6cbbc24f14
Break monolithic distance.py into 8 focused sub-modules: - _geo.py: IP distance functions (subnet mask, haversine, geo) - _text.py: String distance (Levenshtein + enum) + bool distance - _numeric.py: Bytes/hex popcount + timestamp FFT distance - _normalize.py: StandardScaler normalization - _svd.py: svd_denoise (TruncatedSVD, returns noise_profile) - _umap.py: umap_reduce (3D→2D fallback) - _cluster_svd.py: cluster_svd_extract (per-cluster SVD) DELETE: run_preprocessing_pipeline (dead code, replaced by ColumnProcessor) All public APIs re-exported via __init__.py. Backward compatible.
123 lines
4.2 KiB
Python
123 lines
4.2 KiB
Python
"""UMAP dimensionality reduction (3D → 2D fallback)."""
|
|
|
|
import logging
|
|
import numpy as np
|
|
import polars as pl
|
|
from sklearn.preprocessing import StandardScaler
|
|
import umap
|
|
|
|
|
|
def umap_reduce(
|
|
lf: pl.LazyFrame,
|
|
feature_cols: list[str],
|
|
n_components: int = 3,
|
|
) -> tuple[pl.LazyFrame, dict]:
|
|
"""UMAP dimensionality reduction with automatic 3D→2D fallback.
|
|
|
|
Trains on a ≤10K-row sample (when *lf* exceeds that limit), then
|
|
batch-transforms the remaining rows in 1K-row batches.
|
|
|
|
If ``n_components=3`` fails (e.g. too few rows or numerical issues),
|
|
falls back to ``n_components=2`` automatically.
|
|
|
|
Parameters
|
|
----------
|
|
lf:
|
|
LazyFrame containing the numeric feature columns.
|
|
feature_cols:
|
|
Numeric column names to use for UMAP projection. Columns not
|
|
present in *lf* are silently skipped.
|
|
n_components:
|
|
Target dimensionality (2 or 3). Default 3 with fallback to 2.
|
|
|
|
Returns
|
|
-------
|
|
tuple[pl.LazyFrame, dict]
|
|
- LazyFrame: input data extended with ``_umap_x``, ``_umap_y``
|
|
(and ``_umap_z`` when ``n_components ≥ 3``).
|
|
- dict: metadata with keys ``n_components`` and ``reducer``.
|
|
"""
|
|
logger = logging.getLogger(__name__)
|
|
MAX_TRAIN = 10000
|
|
BATCH_SIZE = 1000
|
|
|
|
available = [c for c in feature_cols if c in lf.collect_schema().names()]
|
|
if not available:
|
|
return lf, {'n_components': 0, 'reducer': 'umap', 'error': 'no feature columns available'}
|
|
|
|
df = lf.collect(streaming=True)
|
|
n_total = len(df)
|
|
|
|
if n_total < 2:
|
|
return lf, {'n_components': 0, 'reducer': 'umap', 'error': 'too few rows'}
|
|
|
|
mat = df.select(available).to_numpy().astype(np.float64)
|
|
mat = np.nan_to_num(mat, nan=0.0, posinf=0.0, neginf=0.0)
|
|
|
|
scaler = StandardScaler()
|
|
if n_total > MAX_TRAIN:
|
|
rng = np.random.RandomState(42)
|
|
idx_sample = rng.choice(n_total, size=MAX_TRAIN, replace=False)
|
|
mat_sample = mat[idx_sample]
|
|
scaler.fit(mat_sample)
|
|
mat_sample_scaled = np.nan_to_num(scaler.transform(mat_sample), nan=0.0)
|
|
else:
|
|
scaler.fit(mat)
|
|
mat_sample_scaled = np.nan_to_num(scaler.transform(mat), nan=0.0)
|
|
|
|
# ── Train UMAP; fallback 3→2 on failure ──
|
|
reducer = None
|
|
coords = None
|
|
effective_n = n_components
|
|
|
|
# Try requested n_components first, then fall back to 2 (no duplicates)
|
|
attempts = sorted(set((n_components, 2)), reverse=True)
|
|
for attempt_n in attempts:
|
|
if attempt_n != n_components:
|
|
logger.warning(f'UMAP {n_components}D failed, falling back to 2D')
|
|
effective_n = attempt_n
|
|
try:
|
|
reducer = umap.UMAP(
|
|
n_components=effective_n, random_state=42,
|
|
n_neighbors=15, min_dist=0.1, metric='euclidean',
|
|
)
|
|
if n_total > MAX_TRAIN:
|
|
reducer.fit(mat_sample_scaled)
|
|
else:
|
|
coords = reducer.fit_transform(mat_sample_scaled)
|
|
break
|
|
except Exception as e:
|
|
if attempt_n == 2:
|
|
raise
|
|
logger.warning(f'UMAP {attempt_n}D failed: {e}')
|
|
|
|
# ── Batch transform if dataset was larger than training sample ──
|
|
if n_total > MAX_TRAIN and reducer is not None:
|
|
coords_list: list[np.ndarray] = []
|
|
for start in range(0, n_total, BATCH_SIZE):
|
|
end = min(start + BATCH_SIZE, n_total)
|
|
mat_batch = scaler.transform(mat[start:end])
|
|
mat_batch = np.nan_to_num(mat_batch, nan=0.0, posinf=0.0, neginf=0.0)
|
|
coords_list.append(reducer.transform(mat_batch))
|
|
coords = np.vstack(coords_list)
|
|
|
|
if coords is None:
|
|
return lf, {'n_components': 0, 'reducer': 'umap', 'error': 'UMAP produced no output'}
|
|
|
|
# ── Build result LazyFrame with coordinate columns ──
|
|
coord_series = [
|
|
pl.Series('_umap_x', coords[:, 0]),
|
|
pl.Series('_umap_y', coords[:, 1]),
|
|
]
|
|
if effective_n >= 3 and coords.shape[1] >= 3:
|
|
coord_series.append(pl.Series('_umap_z', coords[:, 2]))
|
|
|
|
df_result = df.with_columns(coord_series)
|
|
|
|
metadata: dict = {
|
|
'n_components': effective_n,
|
|
'reducer': 'umap',
|
|
}
|
|
|
|
return df_result.lazy(), metadata
|