refactor: split distance.py (1114 lines) into analysis/distance/ package
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.
This commit is contained in:
PM-pinou
File diff suppressed because it is too large
Load Diff
@@ -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",
|
||||
]
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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()
|
||||
@@ -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()
|
||||
@@ -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
|
||||
@@ -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)
|
||||
@@ -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
|
||||
Reference in New Issue
Block a user