"""Structure learning over embedded heterogeneous fields.
:func:`mixle.inference.learn_bayesian_network` can model flat records containing
categorical, count, real, and fixed-length vector fields. This module adapts
free-text fields into that workflow by embedding each text value with
:func:`mixle.represent.fit_embedder` and representing the field as a
multivariate vector node.
The learned graph relates the embedding vector to other fields through
multivariate conditional-linear-Gaussian factors or by using vector components
as parent design features. Representative examples per embedding cluster remain
available for interpretation, but the model itself uses the full vector.
"""
from __future__ import annotations
from collections.abc import Sequence
from dataclasses import dataclass
from typing import Any
import numpy as np
from mixle.inference.structure import _is_discrete
def _lloyd(vectors: np.ndarray, k: int, seed: int, iters: int = 25) -> np.ndarray:
"""Seeded k-means centroids on unit vectors (for readable per-cluster representatives only)."""
rng = np.random.RandomState(seed)
k = min(int(k), len(vectors))
centroids = vectors[rng.choice(len(vectors), size=k, replace=False)].copy()
for _ in range(iters):
assign = np.argmax(vectors @ centroids.T, axis=1)
for j in range(k):
members = vectors[assign == j]
if len(members):
c = members.mean(axis=0)
centroids[j] = c / max(float(np.linalg.norm(c)), 1e-12)
else:
centroids[j] = vectors[int(np.argmin(np.max(vectors @ centroids.T, axis=1)))]
return centroids
[docs]
@dataclass
class EmbeddedFieldCodec:
"""One text field's bridge into the graph: embed to a vector; centroids are for interpretability."""
field: int
embedder: Any
centroids: np.ndarray
representatives: list[str] # one training example per cluster, nearest its centroid (for describe())
[docs]
def vector(self, value: Any) -> np.ndarray:
"""Embed one field value as a single ``(dim,)`` vector.
Uses :meth:`~mixle.represent.api.Embedder.transform_one` rather than ``transform``.
``transform`` decides one-vs-batch by guessing from the outer container, and its own
docstring says so; handed a *string* it can read the string as the container and answer
with a block of rows, one per character. That block then went into a single field slot,
which is how two unrelated texts ended up scoring identically -- nothing downstream was
looking at the field's own content by then.
"""
one = getattr(self.embedder, "transform_one", None)
raw = one(str(value)) if callable(one) else self.embedder.transform(str(value))
vec = np.asarray(raw, dtype=np.float64)
if vec.ndim != 1:
raise ValueError(
f"embedder returned shape {vec.shape} for one field value; a single value must embed "
"to exactly one (dim,) vector"
)
return vec
[docs]
def vectors(self, values: Sequence[Any]) -> np.ndarray:
"""Embed a sequence of field values as an ``(n, dim)`` matrix, one row per value."""
values = [str(v) for v in values]
batch = getattr(self.embedder, "transform_batch", None)
raw = batch(values) if callable(batch) else self.embedder.transform(values)
mat = np.asarray(raw, dtype=np.float64)
if mat.ndim != 2 or mat.shape[0] != len(values):
raise ValueError(
f"embedder returned shape {mat.shape} for {len(values)} field values; a batch must "
"embed to exactly one (dim,) row per value"
)
return mat
[docs]
class EmbeddedStructureModel:
"""A discovered dependency graph over records whose text fields ride in as embedding VECTORS."""
def __init__(self, net: Any, codecs: dict[int, EmbeddedFieldCodec]) -> None:
self.net = net
self.codecs = codecs
def __str__(self) -> str:
e = ", ".join(f"{p}->{c}" for p, c in self.net.edges())
return f"EmbeddedStructureModel(text_fields={sorted(self.codecs)}, edges=[{e or 'none'}])"
[docs]
def encode_record(self, x: tuple) -> tuple:
"""The record with each text field replaced by its embedding vector."""
vals = list(x)
for i, codec in self.codecs.items():
vals[i] = codec.vector(vals[i])
return tuple(vals)
[docs]
def encode_records(self, rows: Sequence[tuple]) -> list[tuple]:
"""Replace configured text fields with embedding vectors for each record.
:meth:`EmbeddedFieldCodec.vectors` guarantees one row per value, so the pairing below is
total. It used to be a bare ``zip`` over an unchecked embedder result: an embedder that
returned fewer vectors than it was given left the surplus records holding raw *text* in a
slot the graph reads as numeric, and the batch silently shrank downstream instead of
reporting that the embedder had underproduced.
"""
rows = [list(r) for r in rows]
for i, codec in self.codecs.items():
vecs = codec.vectors([r[i] for r in rows])
for r, v in zip(rows, vecs, strict=True):
r[i] = v
return [tuple(r) for r in rows]
[docs]
def edges(self) -> list[tuple[int, int]]:
"""Return discovered graph edges."""
return self.net.edges()
[docs]
def density_semantics(self):
"""Always :attr:`~mixle.stats.compute.pdist.DensitySemantics.LIKELIHOOD_FACTOR`.
The score is exact for what it computes, but it is not a normalized law over records: text
fields are replaced by embedding vectors before scoring, embedding is not a
measure-preserving change of variables, and no Jacobian relates the two measures. Saying so
in the docstring was not enough -- the machine-readable contract (the ``ExactDensity``
capability, ``describe()``, and every combinator that folds child semantics through
:func:`~mixle.stats.compute.pdist.join_density_semantics`) still read this as an exact
density, so a model composed over one of these silently inherited an exactness claim.
"""
from mixle.stats.compute.pdist import DensitySemantics
return DensitySemantics.LIKELIHOOD_FACTOR
[docs]
def log_density(self, x: tuple) -> float:
"""Log density of ``x`` **in the embedded space**, not of the original record.
Worth being exact about, because the name invites the stronger reading. Each configured text
field is replaced by its embedding vector before scoring, and the graph then evaluates a
density over *that* representation. Embedding is not a measure-preserving change of
variables and no Jacobian relates the two, so this number is not the log density of the text
record and does not integrate to one over any space of text.
What it is good for is comparison: between records whose non-text fields differ, and between
records whose text embeds to genuinely different vectors. What it cannot support is an
absolute probability claim about a record, or a comparison against a density from a model
that scores text some other way. Two different texts that embed close together will score
close together -- that is the representation working as designed, not the model failing to
tell them apart.
"""
return float(self.net.log_density(self.encode_record(x)))
[docs]
def seq_log_density(self, rows: Sequence[tuple]) -> np.ndarray:
"""Per-record log densities in the embedded space; see :meth:`log_density` for what that means."""
embedded = self.encode_records(list(rows))
return np.asarray(self.net.seq_log_density(self.net.dist_to_encoder().seq_encode(embedded)))
[docs]
def describe(self) -> dict[str, Any]:
"""The discovered structure with per-cluster representative examples for each text field."""
return {
"edges": self.net.edges(),
"text_fields": {
i: {f"c{j}": rep for j, rep in enumerate(codec.representatives)} for i, codec in self.codecs.items()
},
}
[docs]
def learn_structure_embedded(
data: Sequence[tuple],
*,
text_fields: Sequence[int] | str = "auto",
n_clusters: int = 8,
embed_dim: int = 16,
seed: int = 0,
max_parents: int = 2,
**embed_kw: Any,
) -> EmbeddedStructureModel:
"""Discover cross-field structure where some fields are free text (see module docstring).
Args:
data: flat tuple records; text fields may hold arbitrary strings.
text_fields: which field indices to embed, or ``"auto"`` -- every string-valued field with too
many distinct values to be a categorical.
n_clusters: number of representative clusters surfaced by ``describe()`` (interpretability only;
the model uses the full embedding vector, not a cluster code).
embed_dim: embedding dimension for :func:`mixle.represent.fit_embedder` -- the vector node's dim.
**embed_kw: forwarded to ``fit_embedder`` (``epochs``, ``hidden``, ``feature_dim``, ...).
"""
from mixle.inference.bayesian_network import learn_bayesian_network
from mixle.represent import fit_embedder
if n_clusters < 1:
raise ValueError(f"n_clusters must be >= 1, got {n_clusters}")
rows = [tuple(r) for r in data]
if len(rows) < 40:
raise ValueError("learn_structure_embedded needs at least 40 records")
n_fields = len(rows[0])
if text_fields == "auto":
text_fields = [
i
for i in range(n_fields)
if all(isinstance(r[i], str) for r in rows) and not _is_discrete([r[i] for r in rows])
]
fields = sorted(int(i) for i in text_fields)
if not fields:
raise ValueError(
"no text fields to embed: pass text_fields= explicitly, or use learn_bayesian_network "
"directly for records without free-text fields"
)
codecs: dict[int, EmbeddedFieldCodec] = {}
for i in fields:
values = [str(r[i]) for r in rows]
emb = fit_embedder(values, dim=int(embed_dim), kind="text", seed=seed, **embed_kw)
vecs = emb.corpus_vectors
centroids = _lloyd(vecs, int(n_clusters), seed)
sims = vecs @ centroids.T
reps = [values[int(np.argmax(sims[:, j]))] for j in range(len(centroids))]
codecs[i] = EmbeddedFieldCodec(field=i, embedder=emb, centroids=centroids, representatives=reps)
model = EmbeddedStructureModel(net=None, codecs=codecs) # encode_records needs only the codecs
embedded = model.encode_records(rows)
model.net = learn_bayesian_network(embedded, max_parents=max_parents)
return model