Skip to content

pelinker.analysis

compute_adjusted_rand_index(y_true, y_pred)

Compute clustering quality via adjusted Rand index (ARI).

Parameters:

Name Type Description Default
y_true ndarray

True labels (e.g., property names)

required
y_pred ndarray

Predicted cluster labels

required

Returns:

Type Description
float

ARI score.

Source code in pelinker/analysis.py
def compute_adjusted_rand_index(y_true: np.ndarray, y_pred: np.ndarray) -> float:
    """
    Compute clustering quality via adjusted Rand index (ARI).

    Args:
        y_true: True labels (e.g., property names)
        y_pred: Predicted cluster labels

    Returns:
        ARI score.
    """
    # Filter out noise points (label -1) for accuracy computation
    valid_mask = y_pred != -1
    if not valid_mask.any():
        return 0.0

    y_true_valid = y_true[valid_mask]
    y_pred_valid = y_pred[valid_mask]

    if len(y_true_valid) == 0:
        return 0.0

    ari = adjusted_rand_score(y_true_valid, y_pred_valid)
    return float(ari)

compute_clustering_fit_metrics(clusterer, manifold_df, *, min_cluster_size, cluster_labels)

DBCV, ARI vs entity, and cluster counts for a fitted HDBSCAN model.

Source code in pelinker/analysis.py
def compute_clustering_fit_metrics(
    clusterer: object,
    manifold_df: pd.DataFrame,
    *,
    min_cluster_size: int,
    cluster_labels: np.ndarray,
) -> ClusteringFitMetrics:
    """DBCV, ARI vs ``entity``, and cluster counts for a fitted HDBSCAN model."""
    labels = np.asarray(cluster_labels, dtype=np.int64).ravel()
    n = int(labels.shape[0])
    label_set = set(labels.tolist())
    n_clusters_emergent = len(label_set) - (1 if -1 in label_set else 0)
    noise_count = int(np.sum(labels == -1))
    noise_fraction = float(noise_count) / float(n) if n > 0 else 0.0

    rv = getattr(clusterer, "relative_validity_", None)
    dbcv: float | None
    if rv is None:
        dbcv = None
    else:
        rv_f = float(rv)
        if math.isnan(rv_f) or math.isinf(rv_f):
            dbcv = None
        else:
            dbcv = rv_f

    ari_score: float | None
    if "entity" in manifold_df.columns and len(manifold_df) == n:
        property_labels = manifold_df["entity"].astype("category").cat.codes.values
        ari_score = compute_adjusted_rand_index(property_labels, labels)
    else:
        ari_score = None

    return ClusteringFitMetrics(
        min_cluster_size=min_cluster_size,
        dbcv=dbcv,
        ari=ari_score,
        n_clusters_emergent=n_clusters_emergent,
        noise_fraction=noise_fraction,
        n_samples=n,
    )

compute_kb_generality_scores(embeddings, labels, k_neighbors=10, metric='cosine', word_frequencies=None, density_weight=0.5)

Compute generality scores for entities based on KB statistics.

Combines embedding-space density with label simplicity to identify generic vs specific terms. Generic terms tend to have: - Many similar neighbors (high density) - High average similarity to neighbors - Shorter, simpler labels (fewer words, common words) - Central position in semantic space

Parameters:

Name Type Description Default
embeddings ndarray

Array of shape (n_points, n_features) containing embeddings

required
labels list[str]

List of labels corresponding to embeddings

required
k_neighbors int

Number of nearest neighbors to consider

10
metric str

Distance metric ('cosine' or 'euclidean')

'cosine'
word_frequencies Mapping[str, float] | None

Optional word frequency mapping for simplicity scoring

None
density_weight float

Weight for embedding density vs label simplicity (0.0 = pure simplicity, 1.0 = pure density)

0.5

Returns:

Type Description
ndarray

Array of generality scores (higher = more generic), shape (n_points,)

Source code in pelinker/analysis.py
def compute_kb_generality_scores(
    embeddings: np.ndarray,
    labels: list[str],
    k_neighbors: int = 10,
    metric: str = "cosine",
    word_frequencies: Mapping[str, float] | None = None,
    density_weight: float = 0.5,
) -> np.ndarray:
    """
    Compute generality scores for entities based on KB statistics.

    Combines embedding-space density with label simplicity to identify generic vs specific terms.
    Generic terms tend to have:
    - Many similar neighbors (high density)
    - High average similarity to neighbors
    - Shorter, simpler labels (fewer words, common words)
    - Central position in semantic space

    Args:
        embeddings: Array of shape (n_points, n_features) containing embeddings
        labels: List of labels corresponding to embeddings
        k_neighbors: Number of nearest neighbors to consider
        metric: Distance metric ('cosine' or 'euclidean')
        word_frequencies: Optional word frequency mapping for simplicity scoring
        density_weight: Weight for embedding density vs label simplicity (0.0 = pure simplicity, 1.0 = pure density)

    Returns:
        Array of generality scores (higher = more generic), shape (n_points,)
    """
    from sklearn.neighbors import NearestNeighbors

    n_points = embeddings.shape[0]
    k_neighbors = min(k_neighbors, n_points - 1)

    if k_neighbors < 1:
        return np.ones(n_points)

    # Normalize embeddings for cosine distance
    if metric == "cosine":
        embeddings_norm = embeddings / (
            np.linalg.norm(embeddings, axis=1, keepdims=True) + 1e-8
        )
    else:
        embeddings_norm = embeddings

    # Find k nearest neighbors for each point
    nn = NearestNeighbors(n_neighbors=k_neighbors + 1, metric=metric)
    nn.fit(embeddings_norm)
    distances, indices = nn.kneighbors(embeddings_norm)

    # Compute embedding-space density scores
    density_scores = np.zeros(n_points)

    for i in range(n_points):
        # Get neighbors (excluding self)
        neighbor_distances = distances[i, 1:]

        # Convert distances to similarities (for cosine: similarity = 1 - distance)
        if metric == "cosine":
            similarities = 1.0 - neighbor_distances
        else:
            # For euclidean, use inverse distance (with smoothing)
            similarities = 1.0 / (1.0 + neighbor_distances)

        # Density = average similarity to neighbors
        # Higher similarity means the term is in a dense region (more generic)
        density_scores[i] = similarities.mean()

    density_scores = np.log(density_scores)
    # Normalize density scores to [0, 1] range
    if density_scores.max() > density_scores.min():
        density_scores_norm = (density_scores - density_scores.min()) / (
            density_scores.max() - density_scores.min()
        )
    else:
        density_scores_norm = np.ones_like(density_scores)

    # Compute label simplicity scores
    if word_frequencies is None:
        word_frequencies = {}

    simplicity_scores = np.zeros(n_points)
    for i, label in enumerate(labels):
        simplicity_metrics = _measure_label_simplicity(
            str(label), word_frequencies=word_frequencies
        )
        simplicity_scores[i] = simplicity_metrics["simplicity_score"]

    simplicity_scores = np.log(simplicity_scores)

    # Normalize simplicity scores to [0, 1] range
    if simplicity_scores.max() > simplicity_scores.min():
        simplicity_scores_norm = (simplicity_scores - simplicity_scores.min()) / (
            simplicity_scores.max() - simplicity_scores.min()
        )
    else:
        simplicity_scores_norm = np.ones_like(simplicity_scores)

    # Combine density and simplicity scores
    # Shorter, simpler terms should be preferred even if density is similar
    generality_scores = (
        density_weight * density_scores_norm
        + (1 - density_weight) * simplicity_scores_norm
    )

    return generality_scores

drop_entities_with_few_mentions(frame, min_mentions_per_entity, *, negative_label=None)

Drop entities with fewer than min_mentions_per_entity rows (same rule as :func:~pelinker.selection.load_selection_frame / mention-level selection eval).

When negative_label is set, that label is never dropped for low mention count (so thin negative tails remain for screener training).

Source code in pelinker/analysis.py
def drop_entities_with_few_mentions(
    frame: pd.DataFrame,
    min_mentions_per_entity: int,
    *,
    negative_label: str | None = None,
) -> pd.DataFrame:
    """
    Drop entities with fewer than ``min_mentions_per_entity`` rows (same rule as
    :func:`~pelinker.selection.load_selection_frame` / mention-level selection eval).

    When ``negative_label`` is set, that label is never dropped for low mention count
    (so thin negative tails remain for screener training).
    """
    if "entity" not in frame.columns:
        raise ValueError("frame must contain an 'entity' column")
    mention_count = frame["entity"].value_counts()
    low_count = mention_count[
        ~(mention_count >= min_mentions_per_entity)
    ].index.to_list()
    if negative_label is not None:
        low_count = [e for e in low_count if e != negative_label]
    return frame.loc[~frame["entity"].isin(low_count)].copy()

embeddings_dict_to_dataframe(embeddings_dict)

Convert embeddings dictionary to DataFrame format expected by transform artifacts.

Parameters:

Name Type Description Default
embeddings_dict dict[str, tuple[str, Tensor | ndarray]]

Dictionary mapping id -> (label, embedding)

required

Returns:

Type Description
DataFrame

DataFrame with columns: id, label, embed

Source code in pelinker/analysis.py
def embeddings_dict_to_dataframe(
    embeddings_dict: dict[str, tuple[str, torch.Tensor | np.ndarray]],
) -> pd.DataFrame:
    """
    Convert embeddings dictionary to DataFrame format expected by transform artifacts.

    Args:
        embeddings_dict: Dictionary mapping id -> (label, embedding)

    Returns:
        DataFrame with columns: id, label, embed
    """
    embeddings_list = []
    id_list = []
    label_list = []

    for id_val, (label, emb) in embeddings_dict.items():
        if isinstance(emb, torch.Tensor):
            emb_np = emb.detach().cpu().numpy()
        else:
            emb_np = np.array(emb)
        embeddings_list.append(emb_np)
        id_list.append(id_val)
        label_list.append(label)

    return pd.DataFrame({"id": id_list, "label": label_list, "embed": embeddings_list})

evaluate_all_screeners_cv(X_embed, X_manifold, y, entity, orig_idx, screener_cfg, oov_cfg)

Shared-stratified-fold CV for LDA/SVM negative screener, manifold OOV model, and stacked score.

screener_best scores use the ROC winner (LDA vs SVM) on pooled OOS predictions.

When oov_cfg.enabled is False or X_manifold is None, OOV branch is skipped: combined metrics match screener_best and oov_winner_kind is "disabled".

Source code in pelinker/analysis.py
def evaluate_all_screeners_cv(
    X_embed: np.ndarray,
    X_manifold: np.ndarray | None,
    y: np.ndarray,
    entity: np.ndarray,
    orig_idx: np.ndarray,
    screener_cfg: NegativeScreenerConfig,
    oov_cfg: ManifoldOovScreenerConfig,
) -> tuple[AllScreenerCvResult, PerDatapointScores] | None:
    """
    Shared-stratified-fold CV for LDA/SVM negative screener, manifold OOV model, and stacked score.

    ``screener_best`` scores use the ROC winner (LDA vs SVM) on pooled OOS predictions.

    When ``oov_cfg.enabled`` is False or ``X_manifold`` is None, OOV branch is skipped:
    ``combined`` metrics match ``screener_best`` and ``oov_winner_kind`` is ``"disabled"``.
    """
    y_i = np.asarray(y, dtype=np.int64).ravel()
    n_splits = _unified_cv_fold_count(y_i, screener_cfg.cv_n_splits)
    if n_splits is None:
        return None

    splitter = StratifiedKFold(
        n_splits=n_splits,
        shuffle=True,
        random_state=screener_cfg.cv_random_state,
    )
    Xe = np.asarray(X_embed, dtype=np.float64)
    fold_pairs = list(splitter.split(Xe, y_i))
    rs_emb = screener_cfg.cv_random_state

    oov_run = bool(oov_cfg.enabled) and X_manifold is not None
    Xm: np.ndarray | None = (
        np.asarray(X_manifold, dtype=np.float64) if oov_run else None
    )
    winner_kind_oov: ManifoldOovKind | str = "lda"

    if oov_run and Xm is not None:
        oov_run, winner_kind_oov = _pick_oov_winner_by_cv_folds(
            fold_pairs, Xm, y_i, oov_cfg
        )
        if not oov_run:
            Xm = None

    collected = _collect_embedding_cv_folds(fold_pairs, Xe, y_i, rs_emb)
    if collected is None:
        return None
    embed_fold_rows, screener_best_kind = collected

    oov_disabled = not oov_run or Xm is None
    oov_kind_disp = "disabled" if oov_disabled else str(winner_kind_oov)

    scored = _score_folds_with_oov_and_pool(
        embed_fold_rows,
        screener_best_kind=screener_best_kind,
        oov_disabled=oov_disabled,
        winner_kind_oov=winner_kind_oov,
        Xm=Xm,
        y_i=y_i,
        entity=entity,
        orig_idx=orig_idx,
        oov_cfg=oov_cfg,
    )
    if scored is None:
        return None
    (
        pool_orig,
        pool_ent,
        pool_y,
        pool_lda_s,
        pool_svm_s,
        pool_sb,
        pool_oov,
        pool_comb,
        fold_lists,
    ) = scored

    result = _build_all_screener_cv_result(
        fold_lists,
        screener_best_kind=screener_best_kind,
        oov_kind_disp=oov_kind_disp,
        oov_disabled=oov_disabled,
    )

    datapoints = PerDatapointScores(
        orig_idx=list(pool_orig),
        entity=list(pool_ent),
        y_true=list(pool_y),
        screener_lda_score=list(pool_lda_s),
        screener_svm_score=list(pool_svm_s),
        screener_best_score=list(pool_sb),
        oov_score=list(pool_oov),
        combined_score=list(pool_comb),
    )
    return result, datapoints

fit_ambient_screener_with_metrics(dfr, config)

Fit the persisted screener on dfr and report in-sample PR/F1 for detecting negative_label when both classes are present.

Source code in pelinker/analysis.py
def fit_ambient_screener_with_metrics(
    dfr: pd.DataFrame,
    config: NegativeScreenerConfig,
) -> tuple[NegativeClassScreener, NegativeScreenerInSampleMetrics | None]:
    """
    Fit the persisted screener on ``dfr`` and report in-sample PR/F1 for detecting
    ``negative_label`` when both classes are present.
    """
    screener = NegativeClassScreener.fit_from_frame(dfr, config)
    y_true = (dfr["entity"].astype(str).values == config.negative_label).astype(
        np.int64
    )
    n_kb = int(np.sum(y_true == 0))
    n_neg = int(np.sum(y_true == 1))
    if n_kb == 0 or n_neg == 0:
        return screener, None
    X = np.stack(dfr["embed"].values).astype(np.float32, copy=False)
    y_pred = screener.predict_is_negative(X).astype(np.int64)
    prec = float(precision_score(y_true, y_pred, pos_label=1, zero_division=0))
    rec = float(recall_score(y_true, y_pred, pos_label=1, zero_division=0))
    f1 = float(f1_score(y_true, y_pred, pos_label=1, zero_division=0))
    return screener, NegativeScreenerInSampleMetrics(
        precision=prec,
        recall=rec,
        f1=f1,
        n_kb_mentions=n_kb,
        n_negative_label_mentions=n_neg,
        kind=config.kind,
    )

get_word_frequencies_from_library(language='en', wordlist='best')

Get word frequency lookup object from wordfreq library.

Parameters:

Name Type Description Default
language str

Language code (default: "en" for English)

'en'
wordlist str

Wordlist size - "best", "large", or "small" (default: "best")

'best'

Returns:

Type Description
object | None

WordFrequencyLookup object with .get() method, or None if library not available

Source code in pelinker/analysis.py
def get_word_frequencies_from_library(
    language: str = "en",
    wordlist: str = "best",
) -> object | None:
    """
    Get word frequency lookup object from wordfreq library.

    Args:
        language: Language code (default: "en" for English)
        wordlist: Wordlist size - "best", "large", or "small" (default: "best")

    Returns:
        WordFrequencyLookup object with .get() method, or None if library not available
    """
    try:
        from wordfreq import word_frequency  # type: ignore

        # Return a callable that looks up frequencies
        # We'll use a lazy lookup approach
        class WordFrequencyLookup:
            def __init__(self, lang: str, wlist: str):
                self.lang = lang
                self.wlist = wlist
                self._cache: dict[str, float] = {}

            def get(self, word: str, default: float = 0.0) -> float:
                word_lower = word.lower()
                if word_lower not in self._cache:
                    try:
                        self._cache[word_lower] = word_frequency(
                            word_lower, self.lang, wordlist=self.wlist
                        )
                    except (KeyError, ValueError):
                        self._cache[word_lower] = default
                return self._cache[word_lower]

            def __getitem__(self, word: str) -> float:
                return self.get(word)

        return WordFrequencyLookup(language, wordlist)  # type: ignore
    except ImportError:
        return None

pooled_grid_solve_from_metrics_dfs(metrics_dfs, optimization_config=None)

After all bootstrap samples have run a min_cluster_size grid, aggregate their metrics once and return full smoothed-grid diagnostics (including y_objective curve).

Source code in pelinker/analysis.py
def pooled_grid_solve_from_metrics_dfs(
    metrics_dfs: Sequence[pd.DataFrame],
    optimization_config: ClusteringOptimizationConfig | None = None,
) -> SmoothedGridOptimumResult:
    """
    After all bootstrap samples have run a min_cluster_size grid, aggregate their metrics
    once and return full smoothed-grid diagnostics (including ``y_objective`` curve).
    """
    if not metrics_dfs:
        raise ValueError("metrics_dfs must be non-empty")
    config = optimization_config or ClusteringOptimizationConfig()
    aggregated = aggregate_grid_metrics(list(metrics_dfs))
    return solve_optimal_min_cluster_size_from_aggregated(
        aggregated,
        objective=config.grid_objective,
        method=config.optimization_method,
        smooth_window=config.grid_smooth_window,
        plateau_fraction=config.grid_plateau_fraction,
        derivative_rel_tol=config.grid_derivative_rel_tol,
        cluster_count_reward=config.grid_cluster_count_reward,
        n_entities=config.grid_n_entities,
    )

pooled_min_cluster_size_from_metrics_dfs(metrics_dfs, optimization_config=None)

After all bootstrap samples have run a min_cluster_size grid, aggregate their metrics once and return the smoothed (chosen_min_cluster_size, raw objective mean at that grid point). The objective is set by ClusteringOptimizationConfig.grid_objective (default: pooled min–max normalized DBCV and ARI).

Source code in pelinker/analysis.py
def pooled_min_cluster_size_from_metrics_dfs(
    metrics_dfs: Sequence[pd.DataFrame],
    optimization_config: ClusteringOptimizationConfig | None = None,
) -> tuple[int, float]:
    """
    After all bootstrap samples have run a min_cluster_size grid, aggregate their metrics
    once and return the smoothed ``(chosen_min_cluster_size, raw objective mean at that grid point)``.
    The objective is set by ``ClusteringOptimizationConfig.grid_objective`` (default: pooled
    min–max normalized DBCV and ARI).
    """
    solved = pooled_grid_solve_from_metrics_dfs(metrics_dfs, optimization_config)
    return solved.chosen_min_cluster_size, solved.score_mean_at_chosen

split_by_negative_label(dfr, negative_label)

Split a mention frame into a boolean mask of synthetic-negative rows and the manifold frame (KB / non-negative rows only).

Source code in pelinker/analysis.py
def split_by_negative_label(
    dfr: pd.DataFrame,
    negative_label: str,
) -> tuple[np.ndarray, pd.DataFrame]:
    """
    Split a mention frame into a boolean mask of synthetic-negative rows and the
    manifold frame (KB / non-negative rows only).
    """
    neg_mask = dfr["entity"].astype(str).values == negative_label
    manifold_df = dfr.loc[~neg_mask].copy()
    return neg_mask, manifold_df