From 6cbbc24f14c838b5635baebf3211d3aca9ad5923 Mon Sep 17 00:00:00 2001 From: PM-pinou <2504420230@qq.com> Date: Fri, 24 Jul 2026 13:50:44 +0800 Subject: [PATCH] refactor: split distance.py (1114 lines) into analysis/distance/ package MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 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. --- analysis/distance.py | 1114 ----------------------------- analysis/distance/__init__.py | 37 + analysis/distance/_cluster_svd.py | 87 +++ analysis/distance/_geo.py | 170 +++++ analysis/distance/_normalize.py | 52 ++ analysis/distance/_numeric.py | 278 +++++++ analysis/distance/_svd.py | 116 +++ analysis/distance/_text.py | 164 +++++ analysis/distance/_umap.py | 122 ++++ 9 files changed, 1026 insertions(+), 1114 deletions(-) delete mode 100644 analysis/distance.py create mode 100644 analysis/distance/__init__.py create mode 100644 analysis/distance/_cluster_svd.py create mode 100644 analysis/distance/_geo.py create mode 100644 analysis/distance/_normalize.py create mode 100644 analysis/distance/_numeric.py create mode 100644 analysis/distance/_svd.py create mode 100644 analysis/distance/_text.py create mode 100644 analysis/distance/_umap.py diff --git a/analysis/distance.py b/analysis/distance.py deleted file mode 100644 index 763cfde..0000000 --- a/analysis/distance.py +++ /dev/null @@ -1,1114 +0,0 @@ -""" -IP distance functions for Polars LazyFrames. - -Provides subnet-mask distance and Haversine geographic distance -between pairs of IP columns. All distances are normalised to [0, 1]; -invalid / non-geolocatable IPs produce ``NaN``. - -Typical usage:: - - from analysis.distance import ip_distance - - lf = pl.scan_csv("flows.csv") - lf = ip_distance(lf, ["src_ip", "dst_ip"]) - lf = ip_distance(lf, ["proxy_ip", "origin_ip"], subnet_masks=[16, 24]) -""" - -from __future__ import annotations - -import math -from typing import Optional - -import polars as pl - -from analysis.geoip import ip_to_int, lookup - -# --------------------------------------------------------------------------- -# Earth constant -# --------------------------------------------------------------------------- - -_HALF_EARTH_CIRCUMFERENCE_KM = 20_037.508 # equatorial / 2 - - -# --------------------------------------------------------------------------- -# Core distance helpers -# --------------------------------------------------------------------------- - -def _subnet_mask_distance(ip1: Optional[str], ip2: Optional[str], mask: int) -> float: - """Normalised subnet-mask distance between two IP strings. - - Computes ``abs(int(ip1)/2**mask - int(ip2)/2**mask) / 2**(32-mask)``. - Returns ``NaN`` when either IP is ``None`` or unparseable. - """ - if ip1 is None or ip2 is None: - return float("nan") - if not isinstance(ip1, str): - ip1 = str(ip1) - if not isinstance(ip2, str): - ip2 = str(ip2) - - int1 = ip_to_int(ip1) - int2 = ip_to_int(ip2) - if int1 is None or int2 is None: - return float("nan") - - divisor = 2.0 ** mask - max_val = 2.0 ** (32 - mask) - return abs(int1 / divisor - int2 / divisor) / max_val - - -def _haversine_km(lat1: float, lon1: float, lat2: float, lon2: float) -> float: - """Great-circle distance (km) via the Haversine formula.""" - dlat = math.radians(lat2 - lat1) - dlon = math.radians(lon2 - lon1) - a = ( - math.sin(dlat / 2) ** 2 - + math.cos(math.radians(lat1)) - * math.cos(math.radians(lat2)) - * math.sin(dlon / 2) ** 2 - ) - return 6371.0 * 2.0 * math.asin(math.sqrt(a)) - - -def _geo_distance(ip1: Optional[str], ip2: Optional[str]) -> float: - """Normalised [0, 1] geographic distance between two IP strings. - - Uses ``analysis.geoip.lookup()`` to obtain lat/lon, then Haversine. - Returns ``NaN`` when either IP cannot be geolocated. - """ - if ip1 is None or ip2 is None: - return float("nan") - if not isinstance(ip1, str): - ip1 = str(ip1) - if not isinstance(ip2, str): - ip2 = str(ip2) - - geo1 = lookup(ip1) - geo2 = lookup(ip2) - if geo1 is None or geo2 is None: - return float("nan") - - km = _haversine_km(geo1["lat"], geo1["lon"], geo2["lat"], geo2["lon"]) - return km / _HALF_EARTH_CIRCUMFERENCE_KM - - -# --------------------------------------------------------------------------- -# Public API -# --------------------------------------------------------------------------- - -def ip_distance( - lf: pl.LazyFrame, - ip_columns: list[str], - subnet_masks: Optional[list[int]] = None, -) -> pl.LazyFrame: - """Append IP distance columns to a LazyFrame. - - For each **consecutive pair** of *ip_columns* - (``ip_columns[0]`` vs ``ip_columns[1]``, - ``ip_columns[2]`` vs ``ip_columns[3]``, …) computes: - - * one subnet-mask distance per entry in *subnet_masks* - * one Haversine geographic distance - - All distances are normalised to [0, 1]. Invalid IPs or IPs not - found in the GeoIP database produce ``NaN``. - - Parameters - ---------- - lf: - LazyFrame containing the IP columns. - ip_columns: - List of IP column names (used in consecutive pairs). - subnet_masks: - Subnet masks for network distance (default ``[24, 28]``). - - Returns - ------- - pl.LazyFrame - Input frame with ``_ip_dist_0``, ``_ip_dist_1``, … - columns appended. - - Examples - -------- - >>> lf = pl.LazyFrame({ - ... "src_ip": ["8.8.8.8", "1.1.1.1"], - ... "dst_ip": ["8.8.4.4", "1.0.0.1"], - ... }) - >>> result = ip_distance(lf, ["src_ip", "dst_ip"], subnet_masks=[24]) - >>> result.collect().columns - ['src_ip', 'dst_ip', '_ip_dist_0', '_ip_dist_1'] - """ - if subnet_masks is None: - subnet_masks = [24, 28] - - result = lf - dist_idx = 0 - - # Walk consecutive pairs - for i in range(0, len(ip_columns) - 1, 2): - col_a = ip_columns[i] - col_b = ip_columns[i + 1] - - # -- subnet-mask distances -- - for mask in subnet_masks: - name = f"_ip_dist_{dist_idx}" - result = result.with_columns( - pl.struct([col_a, col_b]) - .map_elements( - lambda s, _mask=mask, _ca=col_a, _cb=col_b: ( - _subnet_mask_distance(s[_ca], s[_cb], _mask) - ), - return_dtype=pl.Float64, - ) - .alias(name) - ) - dist_idx += 1 - - # -- geographic distance -- - name = f"_ip_dist_{dist_idx}" - result = result.with_columns( - pl.struct([col_a, col_b]) - .map_elements( - lambda s, _ca=col_a, _cb=col_b: _geo_distance(s[_ca], s[_cb]), - return_dtype=pl.Float64, - ) - .alias(name) - ) - dist_idx += 1 - - return result - - -# --------------------------------------------------------------------------- -# String distance -# --------------------------------------------------------------------------- - - -def _is_enum_like(s: pl.Series) -> bool: - """Check if ≥50% of rows have values appearing ≥2 times.""" - total = len(s) - if total == 0: - return False - non_null = s.drop_nulls() - if len(non_null) == 0: - return False - vc = non_null.value_counts() - if 'count' not in vc.columns: - return False - duplicated = vc.filter(pl.col('count') >= 2) - if len(duplicated) == 0: - return False - duplicated_count: int = duplicated['count'].sum() # type: ignore[union-attr] - return duplicated_count >= 0.5 * total - - -def _levenshtein_normalized(s1: str | None, s2: str) -> float: - """Normalised Levenshtein distance between *s1* and *s2*. - - Empty or ``None`` values return 1.0 (maximum distance). - """ - if s1 is None or s1 == '': - return 1.0 - if s2 is None or s2 == '': - return 1.0 - s1_str = str(s1) - s2_str = str(s2) - max_len = max(len(s1_str), len(s2_str)) - if max_len == 0: - return 0.0 - import Levenshtein - return Levenshtein.distance(s1_str, s2_str) / max_len - - -def string_distance(lf: pl.LazyFrame, str_columns: list[str]) -> pl.LazyFrame: - """Compute string distance columns. - - For each column in *str_columns*: - - * If **≥50%** of rows have values appearing ≥2 times, the column is - treated as an enum: 0 when the value equals the mode, 1 otherwise. - * Otherwise, the normalised Levenshtein distance to the mode is - computed for each row. - - Empty or ``None`` values always produce the maximum distance (1.0). - - Parameters - ---------- - lf: - Input LazyFrame. - str_columns: - List of string column names to compute distances for. - - Returns - ------- - pl.LazyFrame - Input LazyFrame extended with ``_str_dist_0``, ``_str_dist_1``, - ... columns (one per entry in *str_columns*). - """ - df = lf.collect(streaming=True) - dist_exprs: list[pl.Expr] = [] - - for idx, col in enumerate(str_columns): - if col not in df.columns: - continue - - series = df[col].cast(pl.Utf8) - enum_like = _is_enum_like(series) - - non_null = series.drop_nulls() - if len(non_null) > 0: - mode_val = non_null.value_counts().sort('count', descending=True).row(0)[0] - mode_str = str(mode_val) - else: - mode_str = '' - - col_name = f'_str_dist_{idx}' - - if enum_like: - dist_exprs.append( - pl.when(pl.col(col).is_null()) - .then(pl.lit(1.0)) - .when(pl.col(col).cast(pl.Utf8) == pl.lit(mode_str)) - .then(pl.lit(0.0)) - .otherwise(pl.lit(1.0)) - .alias(col_name) - ) - else: - # Capture in closure for map_elements - _mode = mode_str - dist_exprs.append( - pl.col(col).map_elements( - lambda v: _levenshtein_normalized(v, _mode), - return_dtype=pl.Float64, - ).alias(col_name) - ) - - return lf.with_columns(dist_exprs) - - -# --------------------------------------------------------------------------- -# Boolean distance -# --------------------------------------------------------------------------- - - -def bool_distance(lf: pl.LazyFrame, bool_columns: list[str]) -> pl.LazyFrame: - """Compute boolean distance columns. - - For each column in *bool_columns*, each row gets: - 0 if its value matches the column mode, 1 otherwise. - ``None`` values always produce 1.0 (maximum distance). - - Parameters - ---------- - lf: - Input LazyFrame. - bool_columns: - List of boolean column names to compute distances for. - - Returns - ------- - pl.LazyFrame - Input LazyFrame extended with ``_bool_dist_0``, ``_bool_dist_1``, - ... columns (one per entry in *bool_columns*). - """ - df = lf.collect(streaming=True) - dist_exprs: list[pl.Expr] = [] - - for idx, col in enumerate(bool_columns): - if col not in df.columns: - continue - - series = df[col] - non_null = series.drop_nulls() - col_name = f'_bool_dist_{idx}' - - if len(non_null) == 0: - dist_exprs.append(pl.lit(0.0).alias(col_name)) - else: - mode_val = non_null.value_counts().sort('count', descending=True).row(0)[0] - dist_exprs.append( - pl.when(pl.col(col).is_null()) - .then(pl.lit(1.0)) - .when(pl.col(col) == pl.lit(mode_val)) - .then(pl.lit(0.0)) - .otherwise(pl.lit(1.0)) - .alias(col_name) - ) - - return lf.with_columns(dist_exprs) - - -# --------------------------------------------------------------------------- -# Bytes Hamming distance -# --------------------------------------------------------------------------- - -_HEX_POPCOUNT: dict[str, int] = { - '0': 0, '1': 1, '2': 1, '3': 2, - '4': 1, '5': 2, '6': 2, '7': 3, - '8': 1, '9': 2, - 'a': 2, 'b': 3, 'c': 2, 'd': 3, 'e': 3, 'f': 4, - 'A': 2, 'B': 3, 'C': 2, 'D': 3, 'E': 3, 'F': 4, -} - - -def _hex_popcount_normalized(val: str | None) -> float: - """Return the proportion of 1-bits in a hex byte string (0.0–1.0). - - This is the Hamming distance from the all-zero bit string, - normalised by total bits. Returns ``NaN`` for null or - non-hex input. - """ - if val is None: - return float("nan") - s = str(val).replace(' ', '').replace('\t', '').replace('\n', '') - if not s: - return 0.0 - - total_ones = 0 - for ch in s: - ones = _HEX_POPCOUNT.get(ch) - if ones is None: - return float("nan") # non-hex character - total_ones += ones - - total_bits = len(s) * 4 - return total_ones / total_bits - - -def bytes_distance( - lf: pl.LazyFrame, - byte_columns: list[str], -) -> pl.LazyFrame: - """Append normalised Hamming-distance columns for hex byte columns. - - Each byte column is converted from hex to binary, and the - proportion of 1-bits (Hamming distance from the all-zero - bit vector) is computed as a ``_bytes_dist_N`` Float64 column - normalised to [0, 1]. Non-hex / null values produce ``NaN``. - - Parameters - ---------- - lf: - LazyFrame containing the byte columns. - byte_columns: - Column names whose values are hex byte strings - (with optional whitespace separators, e.g. ``"A3 F2"``). - - Returns - ------- - pl.LazyFrame - Input frame with ``_bytes_dist_0``, ``_bytes_dist_1``, … - appended. Columns not present in *lf* are silently skipped. - - Examples - -------- - >>> df = pl.DataFrame({"payload": ["A3F2", "FFFF", None]}) - >>> result = bytes_distance(df.lazy(), ["payload"]) - >>> result.collect() - shape: (3, 2) - ┌─────────┬────────────────┐ - │ payload ┆ _bytes_dist_0 │ - │ --- ┆ --- │ - │ str ┆ f64 │ - ╞═════════╪════════════════╡ - │ A3F2 ┆ 0.625 │ - │ FFFF ┆ 1.0 │ - │ null ┆ NaN │ - └─────────┴────────────────┘ - """ - existing = set(lf.collect_schema().names()) - result = lf - - for i, col_name in enumerate(byte_columns): - if col_name not in existing: - continue - result = result.with_columns( - pl.col(col_name) - .cast(pl.Utf8) - .map_elements(_hex_popcount_normalized, return_dtype=pl.Float64) - .alias(f'_bytes_dist_{i}') - ) - - return result - - -# --------------------------------------------------------------------------- -# Timestamp FFT phase distance -# --------------------------------------------------------------------------- - -import numpy as np - -# ── threshold: minimum fraction of repeated intervals to use FFT path ── -_FFT_PERIODICITY_THRESHOLD = 0.05 -# ── minimum samples before FFT is attempted ── -_FFT_MIN_SAMPLES = 4 - - -def _parse_timestamp_to_epoch(series: pl.Series) -> pl.Series: - """Convert a timestamp series to Unix epoch seconds (float64). - - Handles int/float (assumed already epoch), pl.Datetime/pl.Date - (ns→seconds), and str/pl.Utf8 (strptime→ns→seconds, with direct - float-cast fallback on parse failure). - """ - dtype = series.dtype - if dtype in (pl.Utf8, pl.String): - try: - parsed = series.str.strptime(pl.Datetime, format=None, strict=False) - except Exception: - return series.cast(pl.Float64) - return parsed.cast(pl.Int64).cast(pl.Float64) / 1_000_000_000.0 - elif dtype in (pl.Datetime, pl.Date): - return series.cast(pl.Int64).cast(pl.Float64) / 1_000_000_000.0 - else: - return series.cast(pl.Float64) - - -def _compute_fft_distances(values: np.ndarray) -> np.ndarray: - """Compute per-row phase distances for a 1-D array. - - Sorts *values*, checks for periodicity, then either: - * FFT path: ``np.fft.fft()`` → ``np.angle()``, mapped back to - original (unsorted) row order via rank lookup. - * Statistical fallback: interval z-scores mapped to row order. - - Returns an ``(n,)`` float64 array, normalised to zero-mean - unit-variance. - """ - n = len(values) - if n < _FFT_MIN_SAMPLES: - return np.full(n, np.nan, dtype=np.float64) - - # Sort once; record rank→position mapping for unsorting later - sorted_order = np.argsort(values) - ranks = np.empty(n, dtype=np.int64) - ranks[sorted_order] = np.arange(n) - sorted_vals = values[sorted_order] - - # ── periodicity check ────────────────────────────────────────── - intervals = np.diff(sorted_vals) - unique_int, counts = np.unique( - np.round(intervals, 6), return_counts=True - ) - repeat_ratio = float(np.sum(counts[counts > 1])) / max(len(intervals), 1) - - if repeat_ratio >= _FFT_PERIODICITY_THRESHOLD: - # ── FFT phase path ───────────────────────────────────────── - centered = sorted_vals - np.mean(sorted_vals) - fft_result = np.fft.fft(centered) - phases = np.angle(fft_result) - dist_values = phases[ranks] - else: - # ── statistical fallback ─────────────────────────────────── - mean_int = float(np.mean(intervals)) - std_int = float(np.std(intervals)) - if std_int > 1e-10: - z_scores = (intervals - mean_int) / std_int - else: - z_scores = np.zeros_like(intervals) - - # Row at sorted position r gets the z-score of interval r-1 - # (preceding interval); row at position 0 reuses last interval. - dist_values = np.empty(n, dtype=np.float64) - for r in range(n): - if r > 0: - dist_values[ranks[r]] = z_scores[r - 1] - else: - dist_values[ranks[0]] = ( - z_scores[-1] if n > 1 else 0.0 - ) - - # Normalise to zero-mean unit-variance - d_std = float(np.std(dist_values)) - d_mean = float(np.mean(dist_values)) - if d_std > 1e-10: - dist_values = (dist_values - d_mean) / d_std - else: - dist_values = dist_values - d_mean - - return dist_values - - -def timestamp_fft_distance( - lf: pl.LazyFrame, ts_columns: list[str] -) -> pl.LazyFrame: - """Append FFT-based phase-distance columns for timestamp columns. - - For each column in *ts_columns* that exists in the schema: - - 1. Parse values to Unix epoch float seconds (handles int, float, - datetime strings, and ``pl.Datetime``). - 2. Collect the column to numpy, sort, then run ``np.fft.fft`` - followed by ``np.angle`` to extract per-row phase values. - 3. If fewer than 5 % of timestamp intervals repeat (weak - periodicity), fall back to statistical features derived - from interval mean, standard deviation, and z-scores. - - Columns are appended as ``_ts_dist_0``, ``_ts_dist_1``, … in the - order *ts_columns* are supplied. Rows whose timestamp is null - receive ``NaN``. - - Parameters - ---------- - lf : pl.LazyFrame - The input LazyFrame with at least one timestamp column. - ts_columns : list[str] - Column names to process. - - Returns - ------- - pl.LazyFrame - A new LazyFrame with ``_ts_dist_N`` columns appended. - - Examples - -------- - >>> import polars as pl - >>> from analysis.distance import timestamp_fft_distance - >>> df = pl.DataFrame({ - ... 'ts': [0, 60, 120, 180, 240, 300], - ... }) - >>> result = timestamp_fft_distance(df.lazy(), ['ts']) - >>> result.collect().columns - ['ts', '_ts_dist_0'] - """ - schema = lf.collect_schema() - available = [c for c in ts_columns if c in schema.names()] - if not available: - return lf - - # Collect all needed columns once - df = lf.select(available).collect(streaming=True) - n_rows = len(df) - - dist_arrays: list[np.ndarray] = [] - - for col in available: - series = df[col] - ts_numeric = _parse_timestamp_to_epoch(series) - - # Identify valid (non-null) rows - valid = ~ts_numeric.is_null() - valid_indices = np.where(valid.to_numpy())[0] - values = ts_numeric.filter(valid).to_numpy().astype(np.float64) - - if len(values) < _FFT_MIN_SAMPLES: - continue - - phase_values = _compute_fft_distances(values) - - # Build full-length array; invalid rows stay NaN - full_dist = np.full(n_rows, np.nan, dtype=np.float64) - full_dist[valid_indices] = phase_values - dist_arrays.append(full_dist) - - if not dist_arrays: - return lf - - # Re-materialise the original frame and bolt on distance columns - result_df = lf.collect(streaming=True) - for i, arr in enumerate(dist_arrays): - result_df = result_df.with_columns( - pl.Series(f'_ts_dist_{i}', arr) - ) - - return result_df.lazy() - - -# --------------------------------------------------------------------------- -# Normalization — StandardScaler (step 2 of preprocessing pipeline) -# --------------------------------------------------------------------------- - -from sklearn.preprocessing import StandardScaler # noqa: E402 - - -def normalize_features(lf: pl.LazyFrame, feature_cols: list[str]) -> pl.LazyFrame: - """Apply StandardScaler normalisation to numeric feature columns. - - Collects *feature_cols* into a numpy array, replaces ``NaN`` with - zero, centres to μ=0 σ=1 via ``sklearn.preprocessing.StandardScaler``, - then returns the result as *lf* extended with ``_norm_*`` prefixed - columns. - - Parameters - ---------- - lf: - LazyFrame containing *feature_cols*. - feature_cols: - Numeric column names to normalise. Columns not present in - *lf* are silently skipped; non-numeric columns raise an error. - - Returns - ------- - pl.LazyFrame - Input frame with additional ``_norm_{col}`` columns. - - Memory - ------ - Peak memory ≈ collected numpy matrix + scaler workspace - (well under the 800 MB budget for typical inputs). - """ - available = [c for c in feature_cols if c in lf.collect_schema().names()] - if not available: - return lf - - df = lf.select(available).collect(streaming=True) - mat = df.to_numpy().astype(np.float64) - mat = np.nan_to_num(mat, nan=0.0) - - scaler = StandardScaler() - scaled = scaler.fit_transform(mat) - - # Build new columns with _norm_ prefix - df_result = lf.collect(streaming=True) - for i, col in enumerate(available): - df_result = df_result.with_columns( - pl.Series(f'_norm_{col}', scaled[:, i]) - ) - - return df_result.lazy() - - -# --------------------------------------------------------------------------- -# SVD noise removal (step 3 of preprocessing pipeline) -# --------------------------------------------------------------------------- - -from sklearn.decomposition import TruncatedSVD # noqa: E402 - - -def svd_denoise( - lf: pl.LazyFrame, - feature_cols: list[str], - variance_threshold: float = 0.95, -) -> tuple[pl.LazyFrame, dict]: - """Remove noise from numeric features via TruncatedSVD. - - 1. Fits TruncatedSVD on *feature_cols*. - 2. Keeps *k* components where cumulative explained variance - ``trace(S[:k]) / trace(S) ≥ variance_threshold`` (default 0.95). - 3. Inverse-transforms to reconstruct the denoised matrix. - 4. **Replaces** the original *feature_cols* with denoised values. - - The noise components (removed by SVD) are returned as a dict for - display, NOT as new feature columns. - - Parameters - ---------- - lf: - LazyFrame containing *feature_cols*. - feature_cols: - Numeric column names to denoise. These columns are replaced - in-place by their denoised reconstructions. - variance_threshold: - Fraction of total variance to preserve (0.0–1.0). - Default 0.95 keeps 95 % of signal and treats 5 % as noise. - - Returns - ------- - tuple[pl.LazyFrame, dict] - ``(denoised_lf, noise_profile)`` where *noise_profile* contains: - - ``kept_components`` : int — number of components kept - - ``total_components`` : int — max possible components - - ``kept_variance`` : float — fraction of variance kept - - ``noise_components`` : list[dict] — per-component noise info - (component index, explained variance, top-5 contributing features) - """ - noise_profile = { - 'kept_components': 0, - 'total_components': 0, - 'kept_variance': 0.0, - 'noise_components': [], - } - - available = [c for c in feature_cols if c in lf.collect_schema().names()] - if not available: - return lf, noise_profile - - df = lf.select(available).collect(streaming=True) - mat = df.to_numpy().astype(np.float64) - mat = np.nan_to_num(mat, nan=0.0) - - n, d = mat.shape - max_components = min(n, d - 1) - if max_components < 1: - return lf, noise_profile - - from sklearn.decomposition import TruncatedSVD - svd = TruncatedSVD(n_components=max_components, random_state=42) - _ = svd.fit_transform(mat) - - total_variance = float(np.sum(svd.explained_variance_)) - if total_variance <= 1e-10: - return lf, noise_profile - - cumsum = np.cumsum(svd.explained_variance_) - for k in range(1, max_components + 1): - if cumsum[k - 1] / total_variance >= variance_threshold: - break - else: - k = max_components - - # Re-fit with chosen k - svd_k = TruncatedSVD(n_components=k, random_state=42) - reduced = svd_k.fit_transform(mat) - reconstructed = svd_k.inverse_transform(reduced) - - # --- Build noise profile (removed components) --- - # The "noise" is the variance in components k..max_components - kept_var_ratio = cumsum[k - 1] / total_variance if k > 0 else 0.0 - noise_profile = { - 'kept_components': k, - 'total_components': max_components, - 'kept_variance': round(float(kept_var_ratio), 4), - } - - # Per-component noise info for removed components (k and beyond) - noise_info = [] - for comp_idx in range(k, min(k + 10, max_components)): - loading = np.abs(svd.components_[comp_idx]) - top5_idx = np.argsort(loading)[-5:][::-1] - top5 = [ - {'feature': available[i], 'strength': round(float(loading[i]), 4)} - for i in top5_idx - ] - noise_info.append({ - 'component': int(comp_idx + 1), - 'explained_var_ratio': round(float(svd.explained_variance_ratio_[comp_idx]) * 100, 3), - 'top_features': top5, - }) - noise_profile['noise_components'] = noise_info - - # --- Replace original columns with denoised values --- - df_result = lf.collect(streaming=True) - for i, col in enumerate(available): - df_result = df_result.with_columns( - pl.Series(col, reconstructed[:, i]) - ) - - return df_result.lazy(), noise_profile - - -# --------------------------------------------------------------------------- -# Master preprocessing pipeline -# --------------------------------------------------------------------------- - - -def run_preprocessing_pipeline( - lf: pl.LazyFrame, - column_types: dict | None = None, -) -> pl.LazyFrame: - """Full preprocessing pipeline for UMAP / clustering. - - Steps (executed sequentially): - 1. **Column type detection** — uses - :func:`analysis.type_classifier.classify_schema` unless - *column_types* is explicitly provided. - 2. **Distance computation** — IP distance, string distance, - bool distance, bytes distance, and timestamp FFT distance - computed per-column-type and appended as ``_ip_dist_*``, - ``_str_dist_*``, ``_bool_dist_*``, ``_bytes_dist_*``, - ``_ts_dist_*`` columns. - 3. **Normalisation** — :func:`normalize_features` (StandardScaler - → μ=0, σ=1). - 4. **SVD denoise** — :func:`svd_denoise` (TruncatedSVD keeping - 95 % variance). - - Parameters - ---------- - lf: - Input LazyFrame with raw data columns. - column_types: - Optional ``{col_name: DataType}`` mapping. When ``None``, - :func:`analysis.type_classifier.classify_schema` is called - to auto-detect types. - - Returns - ------- - pl.LazyFrame - LazyFrame with only the SVD-denoised feature columns - (``_svd__norm__*_dist_*`` etc.) plus original row anchors. - """ - if column_types is None: - from analysis.types import classify_schema - column_types = classify_schema(lf) - - ip_cols: list[str] = [] - str_cols: list[str] = [] - bool_cols: list[str] = [] - byte_cols: list[str] = [] - ts_cols: list[str] = [] - float_cols: list[str] = [] - int_cols: list[str] = [] - - for col_name, ct in column_types.items(): - if ct.name == "IPv4": - ip_cols.append(col_name) - elif ct.name == "字符串": - str_cols.append(col_name) - elif ct.name == "布尔": - bool_cols.append(col_name) - elif ct.name == "字节": - byte_cols.append(col_name) - elif ct.name == "时间戳": - ts_cols.append(col_name) - elif ct.name in ("浮点", "整数"): - # Float/Int are kept as direct numeric features - if ct.name == "浮点": - float_cols.append(col_name) - else: - int_cols.append(col_name) - - result = lf - - # -- Step 2: Distance computation -- - if ip_cols and len(ip_cols) >= 2: - result = ip_distance(result, ip_cols) - if str_cols: - result = string_distance(result, str_cols) - if bool_cols: - result = bool_distance(result, bool_cols) - if byte_cols: - result = bytes_distance(result, byte_cols) - if ts_cols: - result = timestamp_fft_distance(result, ts_cols) - - # Collect numeric feature columns present after distance computation - schema = result.collect_schema() - all_dist_cols = [ - c for c in schema.names() - if (c.startswith('_ip_dist_') or c.startswith('_str_dist_') - or c.startswith('_bool_dist_') or c.startswith('_bytes_dist_') - or c.startswith('_ts_dist_')) - ] - # Also include raw FLOAT/INT columns as direct features - numeric_raw = [ - c for c in (float_cols + int_cols) - if c in schema.names() - ] - - feature_cols = all_dist_cols + numeric_raw - if not feature_cols: - return result - - # -- Step 3: Normalisation -- - result = normalize_features(result, feature_cols) - - # -- Step 4: SVD denoise (REPLACES noisy features with cleaned ones) -- - norm_cols = [f'_norm_{c}' for c in feature_cols] - result, noise_profile = svd_denoise(result, norm_cols) - - return result, noise_profile - - -# --------------------------------------------------------------------------- -# UMAP dimensionality reduction (3D → 2D fallback) -# --------------------------------------------------------------------------- - - -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``. - """ - import numpy as np - from sklearn.preprocessing import StandardScaler - import umap - import logging - - 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 - - -# --------------------------------------------------------------------------- -# Per-cluster SVD feature extraction -# --------------------------------------------------------------------------- - - -def cluster_svd_extract( - lf: pl.LazyFrame, - labels: list[int], - feature_cols: list[str], -) -> dict[int, dict]: - """Per-cluster SVD-based distinguishing feature extraction. - - For each cluster label: - 1. Extracts rows belonging to that cluster. - 2. Fits ``TruncatedSVD(n_components=min(N, len(feature_cols)))``. - 3. The absolute loadings on the first SVD component serve as - *feature_strength* — higher values mean more distinguishing. - - Returns a mapping ``{cluster_label: {...}}`` where each value is:: - - { - 'features': [...], # top-10 feature names by strength - 'explained_variance_ratio': [...], # per-component ratios - 'feature_strength': [...], # per-feature loading magnitude - } - - Parameters - ---------- - lf: - LazyFrame with rows aligned to *labels*. - labels: - Cluster labels (same length as *lf* row count, ``-1`` = noise). - feature_cols: - Numeric column names to analyse. - - Returns - ------- - dict[int, dict] - Per-cluster SVD extraction results. - """ - import numpy as np - from sklearn.decomposition import TruncatedSVD - - available = [c for c in feature_cols if c in lf.collect_schema().names()] - if not available: - return {} - - df = lf.select(available).collect(streaming=True) - mat = df.to_numpy().astype(np.float64) - mat = np.nan_to_num(mat, nan=0.0, posinf=0.0, neginf=0.0) - - if len(labels) != len(mat): - labels = labels[:len(mat)] - - unique_labels = sorted(set(labels)) - result: dict[int, dict] = {} - - for label in unique_labels: - mask = np.array(labels) == label - n = int(mask.sum()) - if n < 2: - continue - - cluster_mat = mat[mask] - n_comp = min(n, len(available)) - svd = TruncatedSVD(n_components=n_comp, random_state=42) - svd.fit(cluster_mat) - - # First component loadings → feature strength - if n_comp > 0 and svd.components_.shape[0] > 0: - feature_strength = np.abs(svd.components_[0]).astype(np.float64) - else: - feature_strength = np.zeros(len(available), dtype=np.float64) - - total_strength = float(feature_strength.sum()) - if total_strength > 1e-10: - feature_strength = feature_strength / total_strength - - top_indices = np.argsort(-feature_strength) - top_k = min(10, len(available)) - - result[int(label)] = { - 'features': [available[i] for i in top_indices[:top_k]], - 'explained_variance_ratio': svd.explained_variance_ratio_.tolist(), - 'feature_strength': feature_strength.tolist(), - } - - return result diff --git a/analysis/distance/__init__.py b/analysis/distance/__init__.py new file mode 100644 index 0000000..f119571 --- /dev/null +++ b/analysis/distance/__init__.py @@ -0,0 +1,37 @@ +"""Distance computation package for TianXuan analysis pipeline. + +Provides distance functions for various data types (IP, string, boolean, +bytes, timestamp), plus normalisation, SVD denoising, UMAP reduction, +and per-cluster SVD feature extraction. + +Public API: + ip_distance — subnet-mask + Haversine geographic distance for IP pairs + string_distance — Levenshtein / enum-like distance for string columns + bool_distance — 0/1 match-to-mode distance for boolean columns + bytes_distance — Hamming (popcount) distance for hex byte columns + timestamp_fft_distance — FFT phase distance for timestamp columns + normalize_features — StandardScaler normalisation (μ=0, σ=1) + svd_denoise — TruncatedSVD noise removal (returns noise_profile) + umap_reduce — UMAP dimensionality reduction (3D→2D fallback) + cluster_svd_extract — Per-cluster SVD distinguishing feature extraction +""" + +from analysis.distance._geo import ip_distance +from analysis.distance._text import bool_distance, string_distance +from analysis.distance._numeric import bytes_distance, timestamp_fft_distance +from analysis.distance._normalize import normalize_features +from analysis.distance._svd import svd_denoise +from analysis.distance._umap import umap_reduce +from analysis.distance._cluster_svd import cluster_svd_extract + +__all__ = [ + "ip_distance", + "string_distance", + "bool_distance", + "bytes_distance", + "timestamp_fft_distance", + "normalize_features", + "svd_denoise", + "umap_reduce", + "cluster_svd_extract", +] diff --git a/analysis/distance/_cluster_svd.py b/analysis/distance/_cluster_svd.py new file mode 100644 index 0000000..d814dd2 --- /dev/null +++ b/analysis/distance/_cluster_svd.py @@ -0,0 +1,87 @@ +"""Per-cluster SVD feature extraction.""" + +import numpy as np +import polars as pl +from sklearn.decomposition import TruncatedSVD + + +def cluster_svd_extract( + lf: pl.LazyFrame, + labels: list[int], + feature_cols: list[str], +) -> dict[int, dict]: + """Per-cluster SVD-based distinguishing feature extraction. + + For each cluster label: + 1. Extracts rows belonging to that cluster. + 2. Fits ``TruncatedSVD(n_components=min(N, len(feature_cols)))``. + 3. The absolute loadings on the first SVD component serve as + *feature_strength* — higher values mean more distinguishing. + + Returns a mapping ``{cluster_label: {...}}`` where each value is:: + + { + 'features': [...], # top-10 feature names by strength + 'explained_variance_ratio': [...], # per-component ratios + 'feature_strength': [...], # per-feature loading magnitude + } + + Parameters + ---------- + lf: + LazyFrame with rows aligned to *labels*. + labels: + Cluster labels (same length as *lf* row count, ``-1`` = noise). + feature_cols: + Numeric column names to analyse. + + Returns + ------- + dict[int, dict] + Per-cluster SVD extraction results. + """ + available = [c for c in feature_cols if c in lf.collect_schema().names()] + if not available: + return {} + + df = lf.select(available).collect(streaming=True) + mat = df.to_numpy().astype(np.float64) + mat = np.nan_to_num(mat, nan=0.0, posinf=0.0, neginf=0.0) + + if len(labels) != len(mat): + labels = labels[:len(mat)] + + unique_labels = sorted(set(labels)) + result: dict[int, dict] = {} + + for label in unique_labels: + mask = np.array(labels) == label + n = int(mask.sum()) + if n < 2: + continue + + cluster_mat = mat[mask] + n_comp = min(n, len(available)) + svd = TruncatedSVD(n_components=n_comp, random_state=42) + svd.fit(cluster_mat) + + # First component loadings → feature strength + if n_comp > 0 and svd.components_.shape[0] > 0: + feature_strength = np.abs(svd.components_[0]).astype(np.float64) + else: + feature_strength = np.zeros(len(available), dtype=np.float64) + + total_strength = float(feature_strength.sum()) + if total_strength > 1e-10: + feature_strength = feature_strength / total_strength + + top_indices = np.argsort(-feature_strength) + top_k = min(10, len(available)) + + result[int(label)] = { + 'features': [available[i] for i in top_indices[:top_k]], + 'explained_variance_ratio': svd.explained_variance_ratio_.tolist(), + 'feature_strength': feature_strength.tolist(), + } + + return result diff --git a/analysis/distance/_geo.py b/analysis/distance/_geo.py new file mode 100644 index 0000000..dac3c50 --- /dev/null +++ b/analysis/distance/_geo.py @@ -0,0 +1,170 @@ +"""IP distance functions for Polars LazyFrames. + +Provides subnet-mask distance and Haversine geographic distance +between pairs of IP columns. All distances are normalised to [0, 1]; +invalid / non-geolocatable IPs produce ``NaN``. +""" + +from __future__ import annotations + +import math +from typing import Optional + +import polars as pl + +from analysis.geoip import ip_to_int, lookup + +# --------------------------------------------------------------------------- +# Earth constant +# --------------------------------------------------------------------------- + +_HALF_EARTH_CIRCUMFERENCE_KM = 20_037.508 # equatorial / 2 + + +# --------------------------------------------------------------------------- +# Core distance helpers +# --------------------------------------------------------------------------- + +def _subnet_mask_distance(ip1: Optional[str], ip2: Optional[str], mask: int) -> float: + """Normalised subnet-mask distance between two IP strings. + + Computes ``abs(int(ip1)/2**mask - int(ip2)/2**mask) / 2**(32-mask)``. + Returns ``NaN`` when either IP is ``None`` or unparseable. + """ + if ip1 is None or ip2 is None: + return float("nan") + if not isinstance(ip1, str): + ip1 = str(ip1) + if not isinstance(ip2, str): + ip2 = str(ip2) + + int1 = ip_to_int(ip1) + int2 = ip_to_int(ip2) + if int1 is None or int2 is None: + return float("nan") + + divisor = 2.0 ** mask + max_val = 2.0 ** (32 - mask) + return abs(int1 / divisor - int2 / divisor) / max_val + + +def _haversine_km(lat1: float, lon1: float, lat2: float, lon2: float) -> float: + """Great-circle distance (km) via the Haversine formula.""" + dlat = math.radians(lat2 - lat1) + dlon = math.radians(lon2 - lon1) + a = ( + math.sin(dlat / 2) ** 2 + + math.cos(math.radians(lat1)) + * math.cos(math.radians(lat2)) + * math.sin(dlon / 2) ** 2 + ) + return 6371.0 * 2.0 * math.asin(math.sqrt(a)) + + +def _geo_distance(ip1: Optional[str], ip2: Optional[str]) -> float: + """Normalised [0, 1] geographic distance between two IP strings. + + Uses ``analysis.geoip.lookup()`` to obtain lat/lon, then Haversine. + Returns ``NaN`` when either IP cannot be geolocated. + """ + if ip1 is None or ip2 is None: + return float("nan") + if not isinstance(ip1, str): + ip1 = str(ip1) + if not isinstance(ip2, str): + ip2 = str(ip2) + + geo1 = lookup(ip1) + geo2 = lookup(ip2) + if geo1 is None or geo2 is None: + return float("nan") + + km = _haversine_km(geo1["lat"], geo1["lon"], geo2["lat"], geo2["lon"]) + return km / _HALF_EARTH_CIRCUMFERENCE_KM + + +# --------------------------------------------------------------------------- +# Public API +# --------------------------------------------------------------------------- + +def ip_distance( + lf: pl.LazyFrame, + ip_columns: list[str], + subnet_masks: Optional[list[int]] = None, +) -> pl.LazyFrame: + """Append IP distance columns to a LazyFrame. + + For each **consecutive pair** of *ip_columns* + (``ip_columns[0]`` vs ``ip_columns[1]``, + ``ip_columns[2]`` vs ``ip_columns[3]``, …) computes: + + * one subnet-mask distance per entry in *subnet_masks* + * one Haversine geographic distance + + All distances are normalised to [0, 1]. Invalid IPs or IPs not + found in the GeoIP database produce ``NaN``. + + Parameters + ---------- + lf: + LazyFrame containing the IP columns. + ip_columns: + List of IP column names (used in consecutive pairs). + subnet_masks: + Subnet masks for network distance (default ``[24, 28]``). + + Returns + ------- + pl.LazyFrame + Input frame with ``_ip_dist_0``, ``_ip_dist_1``, … + columns appended. + + Examples + -------- + >>> lf = pl.LazyFrame({ + ... "src_ip": ["8.8.8.8", "1.1.1.1"], + ... "dst_ip": ["8.8.4.4", "1.0.0.1"], + ... }) + >>> result = ip_distance(lf, ["src_ip", "dst_ip"], subnet_masks=[24]) + >>> result.collect().columns + ['src_ip', 'dst_ip', '_ip_dist_0', '_ip_dist_1'] + """ + if subnet_masks is None: + subnet_masks = [24, 28] + + result = lf + dist_idx = 0 + + # Walk consecutive pairs + for i in range(0, len(ip_columns) - 1, 2): + col_a = ip_columns[i] + col_b = ip_columns[i + 1] + + # -- subnet-mask distances -- + for mask in subnet_masks: + name = f"_ip_dist_{dist_idx}" + result = result.with_columns( + pl.struct([col_a, col_b]) + .map_elements( + lambda s, _mask=mask, _ca=col_a, _cb=col_b: ( + _subnet_mask_distance(s[_ca], s[_cb], _mask) + ), + return_dtype=pl.Float64, + ) + .alias(name) + ) + dist_idx += 1 + + # -- geographic distance -- + name = f"_ip_dist_{dist_idx}" + result = result.with_columns( + pl.struct([col_a, col_b]) + .map_elements( + lambda s, _ca=col_a, _cb=col_b: _geo_distance(s[_ca], s[_cb]), + return_dtype=pl.Float64, + ) + .alias(name) + ) + dist_idx += 1 + + return result diff --git a/analysis/distance/_normalize.py b/analysis/distance/_normalize.py new file mode 100644 index 0000000..3ca400c --- /dev/null +++ b/analysis/distance/_normalize.py @@ -0,0 +1,52 @@ +"""StandardScaler normalisation (step 2 of preprocessing pipeline).""" + +import numpy as np +import polars as pl +from sklearn.preprocessing import StandardScaler + + +def normalize_features(lf: pl.LazyFrame, feature_cols: list[str]) -> pl.LazyFrame: + """Apply StandardScaler normalisation to numeric feature columns. + + Collects *feature_cols* into a numpy array, replaces ``NaN`` with + zero, centres to μ=0 σ=1 via ``sklearn.preprocessing.StandardScaler``, + then returns the result as *lf* extended with ``_norm_*`` prefixed + columns. + + Parameters + ---------- + lf: + LazyFrame containing *feature_cols*. + feature_cols: + Numeric column names to normalise. Columns not present in + *lf* are silently skipped; non-numeric columns raise an error. + + Returns + ------- + pl.LazyFrame + Input frame with additional ``_norm_{col}`` columns. + + Memory + ------ + Peak memory ≈ collected numpy matrix + scaler workspace + (well under the 800 MB budget for typical inputs). + """ + available = [c for c in feature_cols if c in lf.collect_schema().names()] + if not available: + return lf + + df = lf.select(available).collect(streaming=True) + mat = df.to_numpy().astype(np.float64) + mat = np.nan_to_num(mat, nan=0.0) + + scaler = StandardScaler() + scaled = scaler.fit_transform(mat) + + # Build new columns with _norm_ prefix + df_result = lf.collect(streaming=True) + for i, col in enumerate(available): + df_result = df_result.with_columns( + pl.Series(f'_norm_{col}', scaled[:, i]) + ) + + return df_result.lazy() diff --git a/analysis/distance/_numeric.py b/analysis/distance/_numeric.py new file mode 100644 index 0000000..5a98037 --- /dev/null +++ b/analysis/distance/_numeric.py @@ -0,0 +1,278 @@ +"""Bytes/hex distance (Hamming/popcount) and timestamp FFT phase distance.""" + +import numpy as np +import polars as pl + + +# --------------------------------------------------------------------------- +# Bytes Hamming distance +# --------------------------------------------------------------------------- + +_HEX_POPCOUNT: dict[str, int] = { + '0': 0, '1': 1, '2': 1, '3': 2, + '4': 1, '5': 2, '6': 2, '7': 3, + '8': 1, '9': 2, + 'a': 2, 'b': 3, 'c': 2, 'd': 3, 'e': 3, 'f': 4, + 'A': 2, 'B': 3, 'C': 2, 'D': 3, 'E': 3, 'F': 4, +} + + +def _hex_popcount_normalized(val: str | None) -> float: + """Return the proportion of 1-bits in a hex byte string (0.0–1.0). + + This is the Hamming distance from the all-zero bit string, + normalised by total bits. Returns ``NaN`` for null or + non-hex input. + """ + if val is None: + return float("nan") + s = str(val).replace(' ', '').replace('\t', '').replace('\n', '') + if not s: + return 0.0 + + total_ones = 0 + for ch in s: + ones = _HEX_POPCOUNT.get(ch) + if ones is None: + return float("nan") # non-hex character + total_ones += ones + + total_bits = len(s) * 4 + return total_ones / total_bits + + +def bytes_distance( + lf: pl.LazyFrame, + byte_columns: list[str], +) -> pl.LazyFrame: + """Append normalised Hamming-distance columns for hex byte columns. + + Each byte column is converted from hex to binary, and the + proportion of 1-bits (Hamming distance from the all-zero + bit vector) is computed as a ``_bytes_dist_N`` Float64 column + normalised to [0, 1]. Non-hex / null values produce ``NaN``. + + Parameters + ---------- + lf: + LazyFrame containing the byte columns. + byte_columns: + Column names whose values are hex byte strings + (with optional whitespace separators, e.g. ``"A3 F2"``). + + Returns + ------- + pl.LazyFrame + Input frame with ``_bytes_dist_0``, ``_bytes_dist_1``, … + appended. Columns not present in *lf* are silently skipped. + + Examples + -------- + >>> df = pl.DataFrame({"payload": ["A3F2", "FFFF", None]}) + >>> result = bytes_distance(df.lazy(), ["payload"]) + >>> result.collect() + shape: (3, 2) + ┌─────────┬────────────────┐ + │ payload ┆ _bytes_dist_0 │ + │ --- ┆ --- │ + │ str ┆ f64 │ + ╞═════════╪════════════════╡ + │ A3F2 ┆ 0.625 │ + │ FFFF ┆ 1.0 │ + │ null ┆ NaN │ + └─────────┴────────────────┘ + """ + existing = set(lf.collect_schema().names()) + result = lf + + for i, col_name in enumerate(byte_columns): + if col_name not in existing: + continue + result = result.with_columns( + pl.col(col_name) + .cast(pl.Utf8) + .map_elements(_hex_popcount_normalized, return_dtype=pl.Float64) + .alias(f'_bytes_dist_{i}') + ) + + return result + + +# --------------------------------------------------------------------------- +# Timestamp FFT phase distance +# --------------------------------------------------------------------------- + +# ── threshold: minimum fraction of repeated intervals to use FFT path ── +_FFT_PERIODICITY_THRESHOLD = 0.05 +# ── minimum samples before FFT is attempted ── +_FFT_MIN_SAMPLES = 4 + + +def _parse_timestamp_to_epoch(series: pl.Series) -> pl.Series: + """Convert a timestamp series to Unix epoch seconds (float64). + + Handles int/float (assumed already epoch), pl.Datetime/pl.Date + (ns→seconds), and str/pl.Utf8 (strptime→ns→seconds, with direct + float-cast fallback on parse failure). + """ + dtype = series.dtype + if dtype in (pl.Utf8, pl.String): + try: + parsed = series.str.strptime(pl.Datetime, format=None, strict=False) + except Exception: + return series.cast(pl.Float64) + return parsed.cast(pl.Int64).cast(pl.Float64) / 1_000_000_000.0 + elif dtype in (pl.Datetime, pl.Date): + return series.cast(pl.Int64).cast(pl.Float64) / 1_000_000_000.0 + else: + return series.cast(pl.Float64) + + +def _compute_fft_distances(values: np.ndarray) -> np.ndarray: + """Compute per-row phase distances for a 1-D array. + + Sorts *values*, checks for periodicity, then either: + * FFT path: ``np.fft.fft()`` → ``np.angle()``, mapped back to + original (unsorted) row order via rank lookup. + * Statistical fallback: interval z-scores mapped to row order. + + Returns an ``(n,)`` float64 array, normalised to zero-mean + unit-variance. + """ + n = len(values) + if n < _FFT_MIN_SAMPLES: + return np.full(n, np.nan, dtype=np.float64) + + # Sort once; record rank→position mapping for unsorting later + sorted_order = np.argsort(values) + ranks = np.empty(n, dtype=np.int64) + ranks[sorted_order] = np.arange(n) + sorted_vals = values[sorted_order] + + # ── periodicity check ────────────────────────────────────────── + intervals = np.diff(sorted_vals) + unique_int, counts = np.unique( + np.round(intervals, 6), return_counts=True + ) + repeat_ratio = float(np.sum(counts[counts > 1])) / max(len(intervals), 1) + + if repeat_ratio >= _FFT_PERIODICITY_THRESHOLD: + # ── FFT phase path ───────────────────────────────────────── + centered = sorted_vals - np.mean(sorted_vals) + fft_result = np.fft.fft(centered) + phases = np.angle(fft_result) + dist_values = phases[ranks] + else: + # ── statistical fallback ─────────────────────────────────── + mean_int = float(np.mean(intervals)) + std_int = float(np.std(intervals)) + if std_int > 1e-10: + z_scores = (intervals - mean_int) / std_int + else: + z_scores = np.zeros_like(intervals) + + # Row at sorted position r gets the z-score of interval r-1 + # (preceding interval); row at position 0 reuses last interval. + dist_values = np.empty(n, dtype=np.float64) + for r in range(n): + if r > 0: + dist_values[ranks[r]] = z_scores[r - 1] + else: + dist_values[ranks[0]] = ( + z_scores[-1] if n > 1 else 0.0 + ) + + # Normalise to zero-mean unit-variance + d_std = float(np.std(dist_values)) + d_mean = float(np.mean(dist_values)) + if d_std > 1e-10: + dist_values = (dist_values - d_mean) / d_std + else: + dist_values = dist_values - d_mean + + return dist_values + + +def timestamp_fft_distance( + lf: pl.LazyFrame, ts_columns: list[str] +) -> pl.LazyFrame: + """Append FFT-based phase-distance columns for timestamp columns. + + For each column in *ts_columns* that exists in the schema: + + 1. Parse values to Unix epoch float seconds (handles int, float, + datetime strings, and ``pl.Datetime``). + 2. Collect the column to numpy, sort, then run ``np.fft.fft`` + followed by ``np.angle`` to extract per-row phase values. + 3. If fewer than 5 % of timestamp intervals repeat (weak + periodicity), fall back to statistical features derived + from interval mean, standard deviation, and z-scores. + + Columns are appended as ``_ts_dist_0``, ``_ts_dist_1``, … in the + order *ts_columns* are supplied. Rows whose timestamp is null + receive ``NaN``. + + Parameters + ---------- + lf : pl.LazyFrame + The input LazyFrame with at least one timestamp column. + ts_columns : list[str] + Column names to process. + + Returns + ------- + pl.LazyFrame + A new LazyFrame with ``_ts_dist_N`` columns appended. + + Examples + -------- + >>> import polars as pl + >>> from analysis.distance import timestamp_fft_distance + >>> df = pl.DataFrame({ + ... 'ts': [0, 60, 120, 180, 240, 300], + ... }) + >>> result = timestamp_fft_distance(df.lazy(), ['ts']) + >>> result.collect().columns + ['ts', '_ts_dist_0'] + """ + schema = lf.collect_schema() + available = [c for c in ts_columns if c in schema.names()] + if not available: + return lf + + # Collect all needed columns once + df = lf.select(available).collect(streaming=True) + n_rows = len(df) + + dist_arrays: list[np.ndarray] = [] + + for col in available: + series = df[col] + ts_numeric = _parse_timestamp_to_epoch(series) + + # Identify valid (non-null) rows + valid = ~ts_numeric.is_null() + valid_indices = np.where(valid.to_numpy())[0] + values = ts_numeric.filter(valid).to_numpy().astype(np.float64) + + if len(values) < _FFT_MIN_SAMPLES: + continue + + phase_values = _compute_fft_distances(values) + + # Build full-length array; invalid rows stay NaN + full_dist = np.full(n_rows, np.nan, dtype=np.float64) + full_dist[valid_indices] = phase_values + dist_arrays.append(full_dist) + + if not dist_arrays: + return lf + + # Re-materialise the original frame and bolt on distance columns + result_df = lf.collect(streaming=True) + for i, arr in enumerate(dist_arrays): + result_df = result_df.with_columns( + pl.Series(f'_ts_dist_{i}', arr) + ) + + return result_df.lazy() diff --git a/analysis/distance/_svd.py b/analysis/distance/_svd.py new file mode 100644 index 0000000..2f378c0 --- /dev/null +++ b/analysis/distance/_svd.py @@ -0,0 +1,116 @@ +"""SVD noise removal (step 3 of preprocessing pipeline).""" + +import numpy as np +import polars as pl + + +def svd_denoise( + lf: pl.LazyFrame, + feature_cols: list[str], + variance_threshold: float = 0.95, +) -> tuple[pl.LazyFrame, dict]: + """Remove noise from numeric features via TruncatedSVD. + + 1. Fits TruncatedSVD on *feature_cols*. + 2. Keeps *k* components where cumulative explained variance + ``trace(S[:k]) / trace(S) ≥ variance_threshold`` (default 0.95). + 3. Inverse-transforms to reconstruct the denoised matrix. + 4. **Replaces** the original *feature_cols* with denoised values. + + The noise components (removed by SVD) are returned as a dict for + display, NOT as new feature columns. + + Parameters + ---------- + lf: + LazyFrame containing *feature_cols*. + feature_cols: + Numeric column names to denoise. These columns are replaced + in-place by their denoised reconstructions. + variance_threshold: + Fraction of total variance to preserve (0.0–1.0). + Default 0.95 keeps 95 % of signal and treats 5 % as noise. + + Returns + ------- + tuple[pl.LazyFrame, dict] + ``(denoised_lf, noise_profile)`` where *noise_profile* contains: + - ``kept_components`` : int — number of components kept + - ``total_components`` : int — max possible components + - ``kept_variance`` : float — fraction of variance kept + - ``noise_components`` : list[dict] — per-component noise info + (component index, explained variance, top-5 contributing features) + """ + noise_profile = { + 'kept_components': 0, + 'total_components': 0, + 'kept_variance': 0.0, + 'noise_components': [], + } + + available = [c for c in feature_cols if c in lf.collect_schema().names()] + if not available: + return lf, noise_profile + + df = lf.select(available).collect(streaming=True) + mat = df.to_numpy().astype(np.float64) + mat = np.nan_to_num(mat, nan=0.0) + + n, d = mat.shape + max_components = min(n, d - 1) + if max_components < 1: + return lf, noise_profile + + from sklearn.decomposition import TruncatedSVD + svd = TruncatedSVD(n_components=max_components, random_state=42) + _ = svd.fit_transform(mat) + + total_variance = float(np.sum(svd.explained_variance_)) + if total_variance <= 1e-10: + return lf, noise_profile + + cumsum = np.cumsum(svd.explained_variance_) + for k in range(1, max_components + 1): + if cumsum[k - 1] / total_variance >= variance_threshold: + break + else: + k = max_components + + # Re-fit with chosen k + svd_k = TruncatedSVD(n_components=k, random_state=42) + reduced = svd_k.fit_transform(mat) + reconstructed = svd_k.inverse_transform(reduced) + + # --- Build noise profile (removed components) --- + # The "noise" is the variance in components k..max_components + kept_var_ratio = cumsum[k - 1] / total_variance if k > 0 else 0.0 + noise_profile = { + 'kept_components': k, + 'total_components': max_components, + 'kept_variance': round(float(kept_var_ratio), 4), + } + + # Per-component noise info for removed components (k and beyond) + noise_info = [] + for comp_idx in range(k, min(k + 10, max_components)): + loading = np.abs(svd.components_[comp_idx]) + top5_idx = np.argsort(loading)[-5:][::-1] + top5 = [ + {'feature': available[i], 'strength': round(float(loading[i]), 4)} + for i in top5_idx + ] + noise_info.append({ + 'component': int(comp_idx + 1), + 'explained_var_ratio': round(float(svd.explained_variance_ratio_[comp_idx]) * 100, 3), + 'top_features': top5, + }) + noise_profile['noise_components'] = noise_info + + # --- Replace original columns with denoised values --- + df_result = lf.collect(streaming=True) + for i, col in enumerate(available): + df_result = df_result.with_columns( + pl.Series(col, reconstructed[:, i]) + ) + + return df_result.lazy(), noise_profile diff --git a/analysis/distance/_text.py b/analysis/distance/_text.py new file mode 100644 index 0000000..5b64db0 --- /dev/null +++ b/analysis/distance/_text.py @@ -0,0 +1,164 @@ +"""String distance, enum detection, and boolean distance functions.""" + +import polars as pl + + +# --------------------------------------------------------------------------- +# Enum-like detection helper +# --------------------------------------------------------------------------- + +def _is_enum_like(s: pl.Series) -> bool: + """Check if ≥50% of rows have values appearing ≥2 times.""" + total = len(s) + if total == 0: + return False + non_null = s.drop_nulls() + if len(non_null) == 0: + return False + vc = non_null.value_counts() + if 'count' not in vc.columns: + return False + duplicated = vc.filter(pl.col('count') >= 2) + if len(duplicated) == 0: + return False + duplicated_count: int = duplicated['count'].sum() # type: ignore[union-attr] + return duplicated_count >= 0.5 * total + + +def _levenshtein_normalized(s1: str | None, s2: str) -> float: + """Normalised Levenshtein distance between *s1* and *s2*. + + Empty or ``None`` values return 1.0 (maximum distance). + """ + if s1 is None or s1 == '': + return 1.0 + if s2 is None or s2 == '': + return 1.0 + s1_str = str(s1) + s2_str = str(s2) + max_len = max(len(s1_str), len(s2_str)) + if max_len == 0: + return 0.0 + import Levenshtein + return Levenshtein.distance(s1_str, s2_str) / max_len + + +# --------------------------------------------------------------------------- +# String distance +# --------------------------------------------------------------------------- + +def string_distance(lf: pl.LazyFrame, str_columns: list[str]) -> pl.LazyFrame: + """Compute string distance columns. + + For each column in *str_columns*: + + * If **≥50%** of rows have values appearing ≥2 times, the column is + treated as an enum: 0 when the value equals the mode, 1 otherwise. + * Otherwise, the normalised Levenshtein distance to the mode is + computed for each row. + + Empty or ``None`` values always produce the maximum distance (1.0). + + Parameters + ---------- + lf: + Input LazyFrame. + str_columns: + List of string column names to compute distances for. + + Returns + ------- + pl.LazyFrame + Input LazyFrame extended with ``_str_dist_0``, ``_str_dist_1``, + ... columns (one per entry in *str_columns*). + """ + df = lf.collect(streaming=True) + dist_exprs: list[pl.Expr] = [] + + for idx, col in enumerate(str_columns): + if col not in df.columns: + continue + + series = df[col].cast(pl.Utf8) + enum_like = _is_enum_like(series) + + non_null = series.drop_nulls() + if len(non_null) > 0: + mode_val = non_null.value_counts().sort('count', descending=True).row(0)[0] + mode_str = str(mode_val) + else: + mode_str = '' + + col_name = f'_str_dist_{idx}' + + if enum_like: + dist_exprs.append( + pl.when(pl.col(col).is_null()) + .then(pl.lit(1.0)) + .when(pl.col(col).cast(pl.Utf8) == pl.lit(mode_str)) + .then(pl.lit(0.0)) + .otherwise(pl.lit(1.0)) + .alias(col_name) + ) + else: + # Capture in closure for map_elements + _mode = mode_str + dist_exprs.append( + pl.col(col).map_elements( + lambda v: _levenshtein_normalized(v, _mode), + return_dtype=pl.Float64, + ).alias(col_name) + ) + + return lf.with_columns(dist_exprs) + + +# --------------------------------------------------------------------------- +# Boolean distance +# --------------------------------------------------------------------------- + +def bool_distance(lf: pl.LazyFrame, bool_columns: list[str]) -> pl.LazyFrame: + """Compute boolean distance columns. + + For each column in *bool_columns*, each row gets: + 0 if its value matches the column mode, 1 otherwise. + ``None`` values always produce 1.0 (maximum distance). + + Parameters + ---------- + lf: + Input LazyFrame. + bool_columns: + List of boolean column names to compute distances for. + + Returns + ------- + pl.LazyFrame + Input LazyFrame extended with ``_bool_dist_0``, ``_bool_dist_1``, + ... columns (one per entry in *bool_columns*). + """ + df = lf.collect(streaming=True) + dist_exprs: list[pl.Expr] = [] + + for idx, col in enumerate(bool_columns): + if col not in df.columns: + continue + + series = df[col] + non_null = series.drop_nulls() + col_name = f'_bool_dist_{idx}' + + if len(non_null) == 0: + dist_exprs.append(pl.lit(0.0).alias(col_name)) + else: + mode_val = non_null.value_counts().sort('count', descending=True).row(0)[0] + dist_exprs.append( + pl.when(pl.col(col).is_null()) + .then(pl.lit(1.0)) + .when(pl.col(col) == pl.lit(mode_val)) + .then(pl.lit(0.0)) + .otherwise(pl.lit(1.0)) + .alias(col_name) + ) + + return lf.with_columns(dist_exprs) diff --git a/analysis/distance/_umap.py b/analysis/distance/_umap.py new file mode 100644 index 0000000..f5a3c5d --- /dev/null +++ b/analysis/distance/_umap.py @@ -0,0 +1,122 @@ +"""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