PM-pinou
1115 lines
36 KiB
Python
1115 lines
36 KiB
Python
"""
|
||
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
|