Add stage-5 unknown-format discovery: header clustering + calibration
Implements phase A of the DESIGN_clustering.md design: a Dirichlet-process mixture of per-position categoricals over the first 32 header bytes (257-symbol alphabet) that clusters the binary/ pile by file format, plus signature extraction and promotion nomination. All base-Julia (a Lanczos loggamma keeps the Dirichlet-multinomial marginal dependency-free). - src/cluster.jl: header_symbols feature extraction, collapsed Gibbs sampler (phase A), sequential CRP-predictive assignment (phase B core), signatures/ promotion, and ARI/V-measure calibration metrics. - bin/cluster_calibrate.jl: grid-tunes hyperparameters against magic-collapsed ground truth and cross-checks a model-free NCD (gzip) baseline. - FS_CLUSTER_*/FS_PROMOTE_* config knobs; wire cluster.jl into the module. - Tests for the three DESIGN §10 assertions plus the model primitives. Calibrated defaults (n=32, alpha=1.0, beta=0.1) recover known formats at ARI 0.77 (0.885 excl. tar); docx+zip and the ELF family merge correctly and the NCD baseline agrees. DESIGN §11 records the results and three assumptions the data corrected (tar/ELF header-zero merge, the cold-start seeding deadlock, and the Bernoulli signature / Occam-penalized restart scoring).
This commit is contained in:
518
src/cluster.jl
Normal file
518
src/cluster.jl
Normal file
@@ -0,0 +1,518 @@
|
||||
# Stage-5: unknown-format discovery by Bayesian header clustering.
|
||||
#
|
||||
# See model/DESIGN_clustering.md for the full rationale. In brief: files that
|
||||
# stage-3 sorted into `binary/` are the `:unknown` sink — genuinely
|
||||
# unrecognized bytes. This stage clusters them by *file format* (not producer)
|
||||
# using the first `HEADER_N` header bytes, modeled as a Dirichlet-process
|
||||
# mixture of per-position categoricals over a 257-symbol alphabet
|
||||
# (byte 0–255, plus symbol 257 = "past EOF"). Each cluster's signature is a
|
||||
# magic-number template that can be promoted into the classifier's fast path.
|
||||
#
|
||||
# This file is deliberately dependency-light: everything below is base Julia
|
||||
# (only `log`, no `SpecialFunctions`), so it drops into the existing module and
|
||||
# the offline calibration script alike without touching the Manifest. The model
|
||||
# is categorical on purpose — do NOT reuse model.jl's [0,1] byte scaling here
|
||||
# (that metric is meaningful for the Lux net and meaningless for header bytes,
|
||||
# where 0x89 and 0x88 are not "close"; see DESIGN §4).
|
||||
|
||||
"Number of leading header bytes modeled per file (the feature window). DESIGN §4b."
|
||||
const HEADER_N = 32
|
||||
|
||||
"Alphabet size: byte values 0–255 plus one extra symbol for 'past end of file'."
|
||||
const ALPHABET = 257
|
||||
|
||||
"The 'past EOF' symbol (1-based index `ALPHABET`). A file shorter than a given
|
||||
position emits this here — real, discriminative signal for fixed-length formats,
|
||||
and it avoids colliding zero-padding with genuine 0x00 header bytes (DESIGN §4)."
|
||||
const PAST_EOF = ALPHABET
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Feature extraction
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
"""
|
||||
header_symbols(path; n=HEADER_N) -> Vector{Int}
|
||||
|
||||
Read the first `n` bytes of the file at `path` and map them to a length-`n`
|
||||
vector of 1-based categorical symbols: byte value `b` → `b + 1` (so `1..256`),
|
||||
and every position at or beyond end-of-file → `PAST_EOF` (`257`). Never reads
|
||||
more than `n` bytes, so memory stays flat regardless of file size.
|
||||
"""
|
||||
function header_symbols(path::AbstractString; n::Integer=HEADER_N)
|
||||
syms = fill(PAST_EOF, n)
|
||||
open(path, "r") do io
|
||||
bytes = read(io, n)
|
||||
@inbounds for i in eachindex(bytes)
|
||||
syms[i] = Int(bytes[i]) + 1
|
||||
end
|
||||
end
|
||||
return syms
|
||||
end
|
||||
|
||||
"""
|
||||
header_matrix(paths; n=HEADER_N) -> Matrix{Int}
|
||||
|
||||
Stack `header_symbols` for every path into an `n × length(paths)` matrix (one
|
||||
column per file), the input layout the Gibbs sampler and predictive scorer both
|
||||
consume.
|
||||
"""
|
||||
function header_matrix(paths::AbstractVector{<:AbstractString}; n::Integer=HEADER_N)
|
||||
X = Matrix{Int}(undef, n, length(paths))
|
||||
for (j, p) in enumerate(paths)
|
||||
X[:, j] = header_symbols(p; n=n)
|
||||
end
|
||||
return X
|
||||
end
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Model: DP mixture of per-position categoricals (Dirichlet-Categorical)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
"""
|
||||
ClusterStats
|
||||
|
||||
Sufficient statistics for one cluster: a per-position count table `counts`
|
||||
(`n × ALPHABET`; `counts[i, v]` = how many member files show symbol `v` at
|
||||
position `i`) and the member count `members`. These are exactly what phase-B
|
||||
persists per catalog entry, and everything the collapsed predictive needs.
|
||||
A slot with `members == 0` is inactive (reusable) — the Gibbs sweep prunes
|
||||
emptied clusters without renumbering, so surviving cluster ids stay stable.
|
||||
"""
|
||||
mutable struct ClusterStats
|
||||
counts::Matrix{Int} # n × ALPHABET
|
||||
members::Int
|
||||
end
|
||||
|
||||
ClusterStats(n::Integer) = ClusterStats(zeros(Int, n, ALPHABET), 0)
|
||||
|
||||
"Add file `x` (a length-n symbol vector) into cluster `c`'s sufficient stats."
|
||||
function add!(c::ClusterStats, x::AbstractVector{<:Integer})
|
||||
@inbounds for i in eachindex(x)
|
||||
c.counts[i, x[i]] += 1
|
||||
end
|
||||
c.members += 1
|
||||
return c
|
||||
end
|
||||
|
||||
"Remove file `x` from cluster `c`'s sufficient stats (inverse of `add!`)."
|
||||
function remove!(c::ClusterStats, x::AbstractVector{<:Integer})
|
||||
@inbounds for i in eachindex(x)
|
||||
c.counts[i, x[i]] -= 1
|
||||
end
|
||||
c.members -= 1
|
||||
return c
|
||||
end
|
||||
|
||||
"""
|
||||
log_predictive(c, x, β) -> Float64
|
||||
|
||||
Log probability that file `x` was produced by cluster `c` under the collapsed
|
||||
Dirichlet-Categorical predictive, given `c`'s current counts: at each position
|
||||
`i`, `p(x_i | c) = (counts[i, x_i] + β) / (members + ALPHABET·β)`, summed in log
|
||||
space over positions. Call with `c` NOT containing `x` (Gibbs excludes the point
|
||||
being resampled), so an emptied cluster reduces to the uniform prior `1/ALPHABET`
|
||||
per position — identical to a brand-new cluster, as it should be.
|
||||
"""
|
||||
function log_predictive(c::ClusterStats, x::AbstractVector{<:Integer}, β::Float64)
|
||||
denom = log(c.members + ALPHABET * β)
|
||||
s = 0.0
|
||||
@inbounds for i in eachindex(x)
|
||||
s += log(c.counts[i, x[i]] + β) - denom
|
||||
end
|
||||
return s
|
||||
end
|
||||
|
||||
"Log likelihood of `x` under the fixed uniform component (each position uniform
|
||||
over the 257 symbols): `n · log(1/ALPHABET)`. Used for both the never-adaptive
|
||||
background 'junk drawer' and the prior predictive of a fresh cluster (DESIGN §4a)."
|
||||
log_uniform(n::Integer) = -n * log(ALPHABET)
|
||||
|
||||
# Lanczos approximation to log Γ(x) for x > 0, so partition scoring (below) needs
|
||||
# the Dirichlet-multinomial marginal's gamma terms without pulling in
|
||||
# SpecialFunctions — keeping this stage dependency-flat (no Manifest churn).
|
||||
# g = 7, standard coefficients; accurate to ~1e-14 over the range we use.
|
||||
const _LANCZOS_G = 7
|
||||
const _LANCZOS_C = (0.99999999999980993, 676.5203681218851, -1259.1392167224028,
|
||||
771.32342877765313, -176.61502916214059, 12.507343278686905,
|
||||
-0.13857109526572012, 9.9843695780195716e-6, 1.5056327351493116e-7)
|
||||
|
||||
function loggamma(x::Float64)
|
||||
x < 0.5 && return log(π / sin(π * x)) - loggamma(1.0 - x) # reflection
|
||||
x -= 1.0
|
||||
a = _LANCZOS_C[1]
|
||||
t = x + _LANCZOS_G + 0.5
|
||||
@inbounds for i in 1:_LANCZOS_G + 1
|
||||
a += _LANCZOS_C[i + 1] / (x + i)
|
||||
end
|
||||
return 0.5 * log(2π) + (x + 0.5) * log(t) - t + log(a)
|
||||
end
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Phase A: collapsed Gibbs sampler (offline — the science)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
"""
|
||||
GibbsResult
|
||||
|
||||
Output of `gibbs_cluster`: `assignments` (one per input file; `0` = absorbed by
|
||||
the background junk drawer, positive ints = cluster id), the surviving
|
||||
`clusters` keyed by id, and `score` (the partition's collapsed pseudo-likelihood,
|
||||
used to rank restarts).
|
||||
"""
|
||||
struct GibbsResult
|
||||
assignments::Vector{Int}
|
||||
clusters::Dict{Int,ClusterStats}
|
||||
score::Float64
|
||||
end
|
||||
|
||||
"""
|
||||
gibbs_cluster(X; α, β, bg_mass, sweeps, restarts, rng) -> GibbsResult
|
||||
|
||||
Cluster the columns of `X` (an `n × N` header-symbol matrix) with a collapsed
|
||||
Gibbs sampler over a CRP/Dirichlet-Categorical mixture plus a fixed uniform
|
||||
background (DESIGN §5). Unknown *k* falls out of the CRP natively.
|
||||
|
||||
Per point, per sweep, the point is removed from its cluster and reassigned by
|
||||
sampling from the CRP-predictive weights:
|
||||
|
||||
* existing cluster `k`: `members_k · exp(log_predictive)`
|
||||
* background: `bg_mass · (1/ALPHABET)^n` (never adapts)
|
||||
* a fresh cluster: `α · (1/ALPHABET)^n`
|
||||
|
||||
A uniform/high-entropy blob matches no structured cluster, and background vs.
|
||||
fresh is then decided by `bg_mass` vs. `α`; with `bg_mass ≥ α` such blobs are
|
||||
absorbed rather than minting singletons. `restarts` independent runs are made
|
||||
from different seeds and the highest-scoring partition is returned (a cheap,
|
||||
base-only stand-in for the offline Binder/VI point-summary the design defers).
|
||||
"""
|
||||
function gibbs_cluster(X::AbstractMatrix{<:Integer};
|
||||
α::Float64=1.0, β::Float64=0.5, bg_mass::Float64=5.0,
|
||||
sweeps::Integer=80, restarts::Integer=4,
|
||||
rng::AbstractRNG=Random.default_rng())
|
||||
best = nothing
|
||||
for _ in 1:restarts
|
||||
r = _gibbs_once(X; α=α, β=β, bg_mass=bg_mass, sweeps=sweeps, rng=rng)
|
||||
if best === nothing || r.score > best.score
|
||||
best = r
|
||||
end
|
||||
end
|
||||
return best
|
||||
end
|
||||
|
||||
function _gibbs_once(X::AbstractMatrix{<:Integer};
|
||||
α::Float64, β::Float64, bg_mass::Float64,
|
||||
sweeps::Integer, rng::AbstractRNG)
|
||||
n, N = size(X)
|
||||
# Seed every file in its own singleton (NOT the background). Cold-starting
|
||||
# from the background deadlocks: at a format's first file, a fresh cluster
|
||||
# and the background are equally uniform, so with bg_mass ≥ α the background
|
||||
# always wins and no real cluster is ever seeded. Singleton init sidesteps
|
||||
# this — same-format singletons merge and snowball, while a lone
|
||||
# random-blob singleton dissolves on resample and is reclaimed by the
|
||||
# (stickier) background. See DESIGN §4a.
|
||||
z = collect(1:N)
|
||||
clusters = Dict{Int,ClusterStats}()
|
||||
for j in 1:N
|
||||
c = ClusterStats(n)
|
||||
add!(c, view(X, :, j))
|
||||
clusters[j] = c
|
||||
end
|
||||
next_id = N + 1
|
||||
log_u = log_uniform(n)
|
||||
log_bg = log(bg_mass) + log_u
|
||||
log_new = log(α) + log_u
|
||||
|
||||
# Reused across every point-visit so the sampler's hot loop allocates nothing
|
||||
# per step (2M+ visits per run): `idbuf[t]` is the cluster id whose weight is
|
||||
# `logw[t+1]` (logw[1] = background, logw[end] = fresh). Rebuilding these with
|
||||
# fresh `Vector`/`collect(keys(...))` each step was both slow and enough GC
|
||||
# churn to trip a Julia GC segfault on long grid runs.
|
||||
idbuf = Int[]
|
||||
logw = Float64[]
|
||||
for _ in 1:sweeps
|
||||
for j in 1:N
|
||||
x = view(X, :, j)
|
||||
|
||||
# Remove point j from its current component.
|
||||
zj = z[j]
|
||||
if zj > 0
|
||||
c = clusters[zj]
|
||||
remove!(c, x)
|
||||
if c.members == 0
|
||||
delete!(clusters, zj) # prune emptied cluster; id retired
|
||||
end
|
||||
end
|
||||
|
||||
# Candidate log-weights: background, each live cluster, fresh.
|
||||
empty!(idbuf); empty!(logw)
|
||||
push!(logw, log_bg)
|
||||
for (k, c) in clusters
|
||||
push!(idbuf, k)
|
||||
push!(logw, log(c.members) + log_predictive(c, x, β))
|
||||
end
|
||||
push!(logw, log_new)
|
||||
|
||||
# Gumbel-max sample from the categorical over components.
|
||||
pick = _gumbel_argmax(logw, rng)
|
||||
|
||||
if pick == 1
|
||||
z[j] = 0 # background
|
||||
elseif pick == length(logw)
|
||||
id = next_id; next_id += 1 # fresh cluster
|
||||
c = ClusterStats(n)
|
||||
add!(c, x)
|
||||
clusters[id] = c
|
||||
z[j] = id
|
||||
else
|
||||
id = idbuf[pick - 1]
|
||||
add!(clusters[id], x)
|
||||
z[j] = id
|
||||
end
|
||||
end
|
||||
end
|
||||
|
||||
return GibbsResult(z, clusters, partition_logmarginal(X, z, clusters, α, β))
|
||||
end
|
||||
|
||||
"Argmax of `logw .+ Gumbel noise` — an exact draw from softmax(logw) without
|
||||
normalizing (numerically safe for the tiny header-likelihood magnitudes)."
|
||||
function _gumbel_argmax(logw::AbstractVector{Float64}, rng::AbstractRNG)
|
||||
best_i = 1
|
||||
best_v = -Inf
|
||||
@inbounds for i in eachindex(logw)
|
||||
g = logw[i] - log(-log(rand(rng)))
|
||||
if g > best_v
|
||||
best_v = g
|
||||
best_i = i
|
||||
end
|
||||
end
|
||||
return best_i
|
||||
end
|
||||
|
||||
"""
|
||||
partition_logmarginal(X, z, clusters, α, β) -> Float64
|
||||
|
||||
The joint log-evidence `log p(z, X)` of a partition under the CRP prior and the
|
||||
collapsed Dirichlet-Categorical likelihood — the principled score for ranking
|
||||
Gibbs restarts (higher = better). It is the sum of:
|
||||
|
||||
* the Dirichlet-multinomial **marginal** of each cluster's per-position counts,
|
||||
`lΓ(Aβ) − lΓ(mₖ+Aβ) + Σ_v [lΓ(c_v+β) − lΓ(β)]`, whose normalizer supplies the
|
||||
Occam penalty that a plain product-of-predictives lacks — it is what makes a
|
||||
*merged, heterogeneous* cluster score **worse** than two clean ones (an
|
||||
earlier pseudo-likelihood scorer omitted this and wrongly rewarded merging);
|
||||
* the CRP prior over the clustered points, `K·log α + Σₖ lΓ(mₖ) + lΓ(α) −
|
||||
lΓ(α+N_clustered)`, penalizing gratuitous extra clusters; and
|
||||
* the fixed uniform term for background-assigned files.
|
||||
"""
|
||||
function partition_logmarginal(X::AbstractMatrix{<:Integer}, z::AbstractVector{<:Integer},
|
||||
clusters::Dict{Int,ClusterStats}, α::Float64, β::Float64)
|
||||
n, N = size(X)
|
||||
Aβ = ALPHABET * β
|
||||
lg_Aβ = loggamma(Aβ)
|
||||
lg_β = loggamma(β)
|
||||
s = 0.0
|
||||
# Dirichlet-Categorical marginal likelihood, per cluster × position.
|
||||
for (_, c) in clusters
|
||||
lg_denom = loggamma(c.members + Aβ)
|
||||
@inbounds for i in 1:n, v in 1:ALPHABET
|
||||
cv = c.counts[i, v]
|
||||
cv > 0 && (s += loggamma(cv + β) - lg_β)
|
||||
end
|
||||
s += n * (lg_Aβ - lg_denom)
|
||||
end
|
||||
# CRP prior over the partition of the clustered points.
|
||||
n_bg = count(==(0), z)
|
||||
n_clustered = N - n_bg
|
||||
K = length(clusters)
|
||||
s += K * log(α) + loggamma(α) - loggamma(α + n_clustered)
|
||||
for (_, c) in clusters
|
||||
s += loggamma(float(c.members))
|
||||
end
|
||||
# Background files: independent, uniform.
|
||||
s += n_bg * log_uniform(n)
|
||||
return s
|
||||
end
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Phase B: sequential CRP-predictive assignment (online — the catalog)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
"""
|
||||
assign_file(x, clusters, ids; α, β, bg_mass) -> Int
|
||||
|
||||
Deterministically assign a single file `x` against an existing catalog: the
|
||||
same CRP-predictive rule as Gibbs but at the **argmax** (no sampling) with the
|
||||
current assignments held fixed (DESIGN §5B). Returns the id of the chosen
|
||||
cluster, `0` for the background, or `-1` to signal "mint a new cluster". `ids`
|
||||
is the caller's stable ordering of `keys(clusters)`.
|
||||
|
||||
A new cluster is minted (`-1`) only when the fresh-cluster weight strictly wins.
|
||||
Fresh and background share the same `(1/ALPHABET)^n` likelihood (one file, however
|
||||
structured, is indistinguishable from a uniform blob until a *second* like file
|
||||
appears), so this reduces to `α > bg_mass`. Under the calibrated `bg_mass > α`,
|
||||
minting is therefore effectively off on the live path **by design**: a novel file
|
||||
that matches nothing parks in the background, and genuinely new formats are
|
||||
discovered by the periodic **offline Gibbs compaction** re-clustering that
|
||||
residue (DESIGN §5), not by single-file minting. Because ids are frozen at birth
|
||||
by the caller, there is no label switching.
|
||||
"""
|
||||
function assign_file(x::AbstractVector{<:Integer}, clusters::Dict{Int,ClusterStats},
|
||||
ids::AbstractVector{<:Integer};
|
||||
α::Float64=1.0, β::Float64=0.5, bg_mass::Float64=5.0)
|
||||
n = length(x)
|
||||
log_u = log_uniform(n)
|
||||
best_kind = :bg # :bg, :existing, :new
|
||||
best_id = 0
|
||||
best = log(bg_mass) + log_u
|
||||
new_w = log(α) + log_u
|
||||
if new_w > best
|
||||
best = new_w; best_kind = :new
|
||||
end
|
||||
for k in ids
|
||||
c = clusters[k]
|
||||
w = log(c.members) + log_predictive(c, x, β)
|
||||
if w > best
|
||||
best = w; best_kind = :existing; best_id = k
|
||||
end
|
||||
end
|
||||
return best_kind === :bg ? 0 : best_kind === :new ? -1 : best_id
|
||||
end
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Signatures and promotion (closing the loop to the classifier — DESIGN §6)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
"""
|
||||
signature(c; peak_threshold=0.9, β=0.5) -> Vector{Union{Int,Nothing}}
|
||||
|
||||
Turn a cluster's counts into a libmagic-style template: at each position, if the
|
||||
modal symbol's posterior probability exceeds `peak_threshold`, that byte is
|
||||
*required* (returned as the 0–255 byte value, or `PAST_EOF`); otherwise the
|
||||
position is a wildcard (`nothing`). The vector of required bytes IS the
|
||||
magic-number template — the whole point of the categorical model (DESIGN §4).
|
||||
|
||||
Peakedness is a **Bernoulli** question ("is this position fixed to byte `v`, or
|
||||
not?"), so it uses a 2-way posterior mean `(count + β)/(members + 2β)` — NOT the
|
||||
257-way mixture predictive. The alphabet-wide version would smear the estimate
|
||||
across 257 symbols (`members + 257β` in the denominator), pulling even a
|
||||
unanimous position below any sane threshold once β is small — which would make
|
||||
promotion impossible. This decouples signature detection from the clustering
|
||||
pseudocount and the alphabet size.
|
||||
"""
|
||||
function signature(c::ClusterStats; peak_threshold::Float64=0.9, β::Float64=0.5)
|
||||
n = size(c.counts, 1)
|
||||
sig = Vector{Union{Int,Nothing}}(nothing, n)
|
||||
c.members == 0 && return sig
|
||||
denom = c.members + 2β
|
||||
@inbounds for i in 1:n
|
||||
v = argmax(view(c.counts, i, :))
|
||||
p = (c.counts[i, v] + β) / denom
|
||||
if p > peak_threshold
|
||||
sig[i] = v == PAST_EOF ? PAST_EOF : v - 1 # back to raw byte value
|
||||
end
|
||||
end
|
||||
return sig
|
||||
end
|
||||
|
||||
"Number of fixed (non-wildcard) positions in a signature — its 'magic length'."
|
||||
magic_positions(sig::AbstractVector) = count(!isnothing, sig)
|
||||
|
||||
"""
|
||||
is_promotable(c, sig; min_members=20, min_magic=3) -> Bool
|
||||
|
||||
A cluster qualifies for *nomination* (still human-gated, DESIGN §6) when it has
|
||||
at least `min_members` files AND at least `min_magic` fixed signature positions.
|
||||
The background (id 0) is never passed here — it is never promotable by design.
|
||||
"""
|
||||
function is_promotable(c::ClusterStats, sig::AbstractVector;
|
||||
min_members::Integer=20, min_magic::Integer=3)
|
||||
return c.members >= min_members && magic_positions(sig) >= min_magic
|
||||
end
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Calibration metrics (DESIGN §7): agreement of recovered clusters vs. truth
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
"Map a label vector to consecutive integer ids and a group→indices table."
|
||||
function _groups(labels::AbstractVector)
|
||||
g = Dict{Any,Vector{Int}}()
|
||||
for (i, l) in enumerate(labels)
|
||||
push!(get!(g, l, Int[]), i)
|
||||
end
|
||||
return g
|
||||
end
|
||||
|
||||
"""
|
||||
adjusted_rand_index(a, b) -> Float64
|
||||
|
||||
Adjusted Rand Index between two labelings of the same items: 1.0 = identical
|
||||
partitions (up to relabeling), ~0.0 = chance agreement, can go negative. This is
|
||||
the §7 calibration objective — grid-tuning maximizes ARI of recovered-vs-truth
|
||||
(magic-collapsed) labels. Hand-rolled to keep the dependency footprint flat;
|
||||
matches `Clustering.randindex`.
|
||||
"""
|
||||
function adjusted_rand_index(a::AbstractVector, b::AbstractVector)
|
||||
length(a) == length(b) || throw(DimensionMismatch("label vectors differ in length"))
|
||||
n = length(a)
|
||||
n < 2 && return 1.0
|
||||
ga = collect(values(_groups(a)))
|
||||
gb = collect(values(_groups(b)))
|
||||
# Contingency-table sum of C(n_ij, 2).
|
||||
comb2(x) = x * (x - 1) / 2
|
||||
sa = Set.(ga)
|
||||
index = 0.0
|
||||
for A in sa, B in gb
|
||||
nij = count(in(A), B)
|
||||
index += comb2(nij)
|
||||
end
|
||||
sum_a = sum(comb2(length(g)) for g in ga)
|
||||
sum_b = sum(comb2(length(g)) for g in gb)
|
||||
total = comb2(n)
|
||||
expected = sum_a * sum_b / total
|
||||
maxi = (sum_a + sum_b) / 2
|
||||
denom = maxi - expected
|
||||
return denom == 0 ? 1.0 : (index - expected) / denom
|
||||
end
|
||||
|
||||
"""
|
||||
v_measure(truth, pred; β=1.0) -> (v, homogeneity, completeness)
|
||||
|
||||
Entropy-based cluster agreement (Rosenberg & Hirschberg): homogeneity (each
|
||||
predicted cluster holds one true class), completeness (each true class stays in
|
||||
one predicted cluster), and their weighted harmonic mean `v`. Reported alongside
|
||||
ARI in §7 calibration as a second, differently-biased view.
|
||||
"""
|
||||
function v_measure(truth::AbstractVector, pred::AbstractVector; β::Float64=1.0)
|
||||
n = length(truth)
|
||||
n == 0 && return (1.0, 1.0, 1.0)
|
||||
gt = _groups(truth)
|
||||
gp = _groups(pred)
|
||||
entropy(g) = -sum((length(v) / n) * log(length(v) / n) for v in values(g))
|
||||
H_C = entropy(gt)
|
||||
H_K = entropy(gp)
|
||||
# Conditional entropies via the contingency table.
|
||||
H_CK = 0.0 # H(truth | pred)
|
||||
H_KC = 0.0 # H(pred | truth)
|
||||
for (_, P) in gp
|
||||
Ps = Set(P)
|
||||
for (_, C) in gt
|
||||
nij = count(in(Ps), C)
|
||||
nij == 0 && continue
|
||||
H_CK -= (nij / n) * log(nij / length(P))
|
||||
end
|
||||
end
|
||||
for (_, C) in gt
|
||||
Cs = Set(C)
|
||||
for (_, P) in gp
|
||||
nij = count(in(Cs), P)
|
||||
nij == 0 && continue
|
||||
H_KC -= (nij / n) * log(nij / length(C))
|
||||
end
|
||||
end
|
||||
homogeneity = H_C == 0 ? 1.0 : 1 - H_CK / H_C
|
||||
completeness = H_K == 0 ? 1.0 : 1 - H_KC / H_K
|
||||
v = (homogeneity + completeness == 0) ? 0.0 :
|
||||
(1 + β) * homogeneity * completeness / (β * homogeneity + completeness)
|
||||
return (v, homogeneity, completeness)
|
||||
end
|
||||
Reference in New Issue
Block a user