pycorpdiff 0.1.0a0__py3-none-any.whl

pycorpdiff/tokenize.py

@@ -0,0 +1,110 @@
+"""Tokenizer protocol and the default regex tokenizer.
+
+The :class:`Tokenizer` protocol is the package's only extension point for
+language-specific preprocessing. Adapters around spaCy, Stanza, jieba,
+fugashi, etc. need to satisfy a single ``__call__(text: str) -> list[str]``
+contract — no inheritance, no registration.
+"""
+
+from __future__ import annotations
+
+import re
+from dataclasses import dataclass, field
+from typing import Protocol, runtime_checkable
+
+
+@runtime_checkable
+class Tokenizer(Protocol):
+    """Anything callable that maps a string to a list of token strings."""
+
+    def __call__(self, text: str) -> list[str]: ...
+
+
+@dataclass(frozen=True)
+class RegexTokenizer:
+    """A minimal Unicode-aware regex tokenizer used as the default.
+
+    The default pattern matches sequences of word characters (``\\w+``),
+    which under Python's default regex engine is Unicode-aware and
+    therefore safe for non-Latin scripts at the granularity of "word-like
+    runs of letters / digits / underscores". Researchers needing
+    language-specific behaviour (lemmatisation, segmentation of CJK
+    scripts, MWE handling, etc.) should plug in a spaCy/Stanza/jieba
+    adapter instead.
+    """
+
+    pattern: str = r"\w+"
+    lowercase: bool = True
+    _compiled: re.Pattern[str] = field(init=False, repr=False, compare=False)
+
+    def __post_init__(self) -> None:
+        object.__setattr__(self, "_compiled", re.compile(self.pattern, re.UNICODE))
+
+    def __call__(self, text: str) -> list[str]:
+        if self.lowercase:
+            text = text.lower()
+        return self._compiled.findall(text)
+
+
+@dataclass(frozen=True)
+class NgramTokenizer:
+    """A tokenizer wrapper that emits joined n-grams from a base tokenizer.
+
+    Wraps any :class:`Tokenizer` and yields ``n``-token sequences joined
+    by ``sep``. The output is a flat ``list[str]`` of joined n-grams,
+    which means every downstream analytical surface — keyness,
+    dispersion, collocation, semantic shift — treats them as ordinary
+    terms with no special-casing needed.
+
+    Parameters
+    ----------
+    base
+        The underlying tokenizer producing unigrams. Defaults to
+        :class:`RegexTokenizer`.
+    n
+        N-gram order. ``2`` = bigrams, ``3`` = trigrams. Must be ``>=1``.
+    sep
+        Joiner string. Default ``"_"`` matches gensim's convention and
+        sidesteps the usual ambiguity of whitespace-joined n-grams when
+        the base tokenizer itself strips whitespace.
+    include_lower
+        If ``True``, the output also includes every ``k``-gram for
+        ``k < n`` (so ``n=3`` emits unigrams + bigrams + trigrams).
+        Useful when you want a single Comparison to keyness-rank
+        single words *and* their multi-word collocations side-by-side.
+
+    Examples
+    --------
+    >>> from pycorpdiff.tokenize import NgramTokenizer
+    >>> tok = NgramTokenizer(n=2)
+    >>> tok("the cat sat on the mat")
+    ['the_cat', 'cat_sat', 'sat_on', 'on_the', 'the_mat']
+    >>> NgramTokenizer(n=2, include_lower=True)("a b c")
+    ['a', 'b', 'c', 'a_b', 'b_c']
+    """
+
+    base: Tokenizer = field(default_factory=RegexTokenizer)
+    n: int = 2
+    sep: str = "_"
+    include_lower: bool = False
+
+    def __post_init__(self) -> None:
+        if self.n < 1:
+            raise ValueError(f"n must be >= 1; got {self.n}")
+
+    def __call__(self, text: str) -> list[str]:
+        unigrams = self.base(text)
+        if self.n == 1:
+            return unigrams
+        out: list[str] = []
+        start = 1 if self.include_lower else self.n
+        for k in range(start, self.n + 1):
+            if k == 1:
+                out.extend(unigrams)
+                continue
+            sep = self.sep
+            out.extend(
+                sep.join(unigrams[i : i + k])
+                for i in range(len(unigrams) - k + 1)
+            )
+        return out

pycorpdiff/viz/__init__.py

@@ -0,0 +1,37 @@
+"""Visualisation helpers — altair-first, matplotlib for paper-grade figures.
+
+Every Result type's ``.plot()`` method delegates here. Plot functions
+also accept a bare DataFrame so users can call
+``pcd.viz.keyness_volcano(df)`` directly without going through a Result.
+
+altair is an optional dependency declared in the ``viz`` extra. Each
+plot function lazily imports altair on first call; the friendly
+ImportError lives at that boundary.
+"""
+
+from __future__ import annotations
+
+from .bocpd import bocpd_plot
+from .causal_impact import causal_impact_plot
+from .collocation import collocation_diverging_bar
+from .dispersion import dispersion_plot
+from .forecast import forecast_plot
+from .keyness import keyness_top_n_bar, keyness_volcano
+from .network import network_plot
+from .scattertext import scattertext_plot
+from .semantic_forecast import semantic_forecast_plot
+from .trajectory import trajectory_with_ci
+
+__all__ = [
+    "bocpd_plot",
+    "causal_impact_plot",
+    "collocation_diverging_bar",
+    "dispersion_plot",
+    "forecast_plot",
+    "keyness_top_n_bar",
+    "keyness_volcano",
+    "network_plot",
+    "scattertext_plot",
+    "semantic_forecast_plot",
+    "trajectory_with_ci",
+]

pycorpdiff/viz/bocpd.py

@@ -0,0 +1,173 @@
+"""BOCPD diagnostic plot — series + run-length posterior heatmap + MAP line.
+
+Three stacked panels sharing the time axis:
+
+1. The input series with detected changepoints flagged.
+2. Heatmap of the run-length posterior P(r_t | data so far) on a
+   log colour scale — the canonical BOCPD diagnostic figure.
+3. MAP run length over time. Visible drops mark changepoints.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import numpy as np
+import pandas as pd
+
+if TYPE_CHECKING:
+    import altair as alt
+
+    from ..temporal.bocpd import BocpdResult
+
+
+def bocpd_plot(
+    result: BocpdResult,
+    *,
+    width: int = 660,
+    height_per_panel: int = 180,
+    max_run_length_shown: int = 40,
+    threshold: int = 3,
+) -> alt.Chart:
+    """Three-panel BOCPD diagnostic chart.
+
+    Parameters
+    ----------
+    result
+        A :class:`BocpdResult`.
+    width
+        Width of each panel (panels are vertically concat'd).
+    height_per_panel
+        Height of each panel.
+    max_run_length_shown
+        Truncate the heatmap at this run length — most posterior mass
+        lives in the first few dozen run lengths.
+    threshold
+        MAP-run-length threshold below which a step is flagged as a
+        detected changepoint in panel 1.
+    """
+    import altair as alt
+
+    # The heatmap panel has T × (max_run_length_shown + 1) cells which
+    # can exceed altair's default 5000-row inline-data limit for longer
+    # series. Disable the limit on the heatmap dataframe specifically —
+    # the data is inline (no external host), so the rendering cost is
+    # bounded by the local browser.
+    alt.data_transformers.disable_max_rows()
+
+    periods = result.series.index
+    if isinstance(periods, pd.PeriodIndex):
+        period_axis: pd.Index = pd.Index(periods.to_timestamp())
+    else:
+        period_axis = pd.Index(periods)
+
+    series_df = pd.DataFrame(
+        {"period": period_axis, "value": result.series.to_numpy(dtype=float)}
+    )
+    flagged_df = result.detected_changepoints(threshold=threshold).copy()
+    if len(flagged_df) and isinstance(flagged_df["period"].iloc[0], pd.Period):
+        flagged_df["period"] = flagged_df["period"].apply(lambda p: p.to_timestamp())
+
+    # ---------- Panel 1: series + flagged changepoints ----------
+    series_line = (
+        alt.Chart(series_df)
+        .mark_line(strokeWidth=2, color="#0b6e7c")
+        .encode(
+            x=alt.X("period:T", title=None),
+            y=alt.Y("value:Q", title=None),
+            tooltip=["period", alt.Tooltip("value:Q", format=".5f")],
+        )
+    )
+    series_points = (
+        alt.Chart(series_df)
+        .mark_point(filled=True, size=40, color="#0b6e7c")
+        .encode(x="period:T", y="value:Q")
+    )
+    cp_rules = (
+        alt.Chart(flagged_df)
+        .mark_rule(color="#e63946", strokeDash=[4, 3], opacity=0.7)
+        .encode(x="period:T")
+    )
+    panel1 = (series_line + series_points + cp_rules).properties(
+        width=width,
+        height=height_per_panel,
+        title=alt.TitleParams(
+            text=f"Observed series — {len(flagged_df)} flagged changepoint(s)",
+            subtitle=f"red lines: MAP run length ≤ {threshold} (hazard = {result.hazard})",
+        ),
+    )
+
+    # ---------- Panel 2: run-length posterior heatmap ----------
+    R = np.asarray(result.run_length_posterior)
+    R_shown = R[:, : max_run_length_shown + 1]
+    # log10 + clip for colour mapping — most posterior values are
+    # tiny floats, log compresses the dynamic range.
+    with np.errstate(divide="ignore"):
+        log_R = np.log10(np.clip(R_shown, 1e-12, 1.0))
+    n_t, n_r = log_R.shape
+    period_repeat = np.tile(period_axis.to_numpy(), n_r)
+    runs = np.repeat(np.arange(n_r), n_t)
+    heat_df = pd.DataFrame(
+        {
+            "period": period_repeat,
+            "run_length": runs,
+            "log_posterior": log_R.T.ravel(),
+        }
+    )
+    heatmap = (
+        alt.Chart(heat_df)
+        .mark_rect()
+        .encode(
+            x=alt.X("period:T", title=None),
+            y=alt.Y("run_length:O", title="run length r", sort="descending"),
+            color=alt.Color(
+                "log_posterior:Q",
+                scale=alt.Scale(scheme="viridis", domain=[-6, 0]),
+                title="log₁₀ P(r | data)",
+            ),
+            tooltip=[
+                "period",
+                "run_length",
+                alt.Tooltip("log_posterior:Q", format=".2f"),
+            ],
+        )
+        .properties(
+            width=width,
+            height=height_per_panel,
+            title="Run-length posterior P(r_t | data through t) — log₁₀ scale",
+        )
+    )
+
+    # ---------- Panel 3: MAP run length over time ----------
+    map_df = pd.DataFrame(
+        {
+            "period": period_axis,
+            "map_run_length": result.map_run_length.to_numpy(dtype=int),
+        }
+    )
+    map_line = (
+        alt.Chart(map_df)
+        .mark_line(strokeWidth=2, color="#1f7a3e")
+        .encode(
+            x=alt.X("period:T", title=None),
+            y=alt.Y("map_run_length:Q", title="MAP run length"),
+            tooltip=["period", "map_run_length"],
+        )
+    )
+    map_points = (
+        alt.Chart(map_df)
+        .mark_point(filled=True, size=40, color="#1f7a3e")
+        .encode(x="period:T", y="map_run_length:Q")
+    )
+    threshold_rule = (
+        alt.Chart(pd.DataFrame({"y": [threshold]}))
+        .mark_rule(color="#888", strokeDash=[3, 3], opacity=0.6)
+        .encode(y="y:Q")
+    )
+    panel3 = (map_line + map_points + threshold_rule).properties(
+        width=width,
+        height=height_per_panel,
+        title="MAP run length — visible drops mark detected changepoints",
+    )
+
+    return alt.vconcat(panel1, heatmap, panel3).resolve_scale(x="shared")  # type: ignore[no-any-return]

pycorpdiff/viz/causal_impact.py

@@ -0,0 +1,142 @@
+"""Three-panel causal-impact plot (Brodersen et al. 2015 style).
+
+- Panel 1 — observed series vs counterfactual (dashed) with CrI band
+- Panel 2 — pointwise effect (observed − counterfactual) with CrI band,
+  zero reference line
+- Panel 3 — cumulative effect with CrI band, zero reference line
+
+Stacked vertically, sharing the time axis. The visual grammar matches
+the figures from Brodersen, Gallusser, Koehler, Remy & Scott (2015).
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import pandas as pd
+
+if TYPE_CHECKING:
+    import altair as alt
+
+    from ..temporal.causal_impact import CausalImpactResult
+
+
+def causal_impact_plot(
+    result: CausalImpactResult,
+    *,
+    width: int = 640,
+    height_per_panel: int = 200,
+) -> alt.Chart:
+    """Three-panel observed/counterfactual + effect plot.
+
+    Parameters
+    ----------
+    result
+        A :class:`CausalImpactResult` from
+        :meth:`TemporalTrajectory.causal_impact`.
+    width
+        Width of each panel.
+    height_per_panel
+        Height of each panel. Total chart height ≈ 3× this.
+    """
+    import altair as alt
+
+    df = result.table.copy()
+    if len(df) and isinstance(df["period"].iloc[0], pd.Period):
+        df["period"] = df["period"].apply(lambda p: p.to_timestamp())
+
+    # The event marker: a vertical rule at result.event_date.
+    event_df = pd.DataFrame({"event": [pd.Timestamp(result.event_date)]})
+    event_rule = (
+        alt.Chart(event_df)
+        .mark_rule(color="#999", strokeDash=[3, 3], strokeWidth=1.5)
+        .encode(x="event:T")
+    )
+    zero_rule = (
+        alt.Chart(pd.DataFrame({"y": [0.0]}))
+        .mark_rule(color="#666", strokeWidth=1)
+        .encode(y="y:Q")
+    )
+
+    x_axis = alt.X("period:T", title=None)
+
+    # ---------- Panel 1: observed vs counterfactual ----------
+    base1 = alt.Chart(df).encode(x=x_axis)
+    cf_band = base1.mark_area(opacity=0.18, color="#888").encode(
+        y=alt.Y("counterfactual_lower:Q", title=f"{result.target} rate"),
+        y2="counterfactual_upper:Q",
+    )
+    cf_line = base1.mark_line(
+        color="#888", strokeDash=[5, 4], strokeWidth=2
+    ).encode(y="counterfactual:Q")
+    observed_line = base1.mark_line(color="#0b6e7c", strokeWidth=2.5).encode(
+        y="observed:Q",
+        tooltip=[
+            "period",
+            alt.Tooltip("observed:Q", format=".5f"),
+            alt.Tooltip("counterfactual:Q", format=".5f"),
+            alt.Tooltip("counterfactual_lower:Q", format=".5f"),
+            alt.Tooltip("counterfactual_upper:Q", format=".5f"),
+        ],
+    )
+    panel1 = (
+        (cf_band + cf_line + observed_line + event_rule)
+        .properties(
+            width=width,
+            height=height_per_panel,
+            title=alt.TitleParams(
+                text=f"Observed (teal) vs counterfactual (gray) — {result.target!r}",
+                subtitle=(
+                    f"event = {pd.Timestamp(result.event_date).date()} · "
+                    f"BSTS local linear trend on {result.n_pre} pre-event periods"
+                ),
+            ),
+        )
+    )
+
+    # ---------- Panel 2: pointwise effect ----------
+    base2 = alt.Chart(df).encode(x=x_axis)
+    pw_band = base2.mark_area(opacity=0.22, color="#e63946").encode(
+        y=alt.Y("pointwise_lower:Q", title="pointwise effect"),
+        y2="pointwise_upper:Q",
+    )
+    pw_line = base2.mark_line(color="#e63946", strokeWidth=2).encode(
+        y="pointwise_effect:Q",
+        tooltip=[
+            "period",
+            alt.Tooltip("pointwise_effect:Q", format=".5f"),
+            alt.Tooltip("pointwise_lower:Q", format=".5f"),
+            alt.Tooltip("pointwise_upper:Q", format=".5f"),
+        ],
+    )
+    panel2 = (pw_band + pw_line + zero_rule + event_rule).properties(
+        width=width,
+        height=height_per_panel,
+        title=alt.TitleParams(
+            text="Pointwise effect (observed − counterfactual)",
+            subtitle=f"{int(result.level * 100)}% credible interval shaded",
+        ),
+    )
+
+    # ---------- Panel 3: cumulative effect ----------
+    base3 = alt.Chart(df).encode(x=x_axis)
+    cum_band = base3.mark_area(opacity=0.22, color="#1f7a3e").encode(
+        y=alt.Y("cumulative_lower:Q", title="cumulative effect"),
+        y2="cumulative_upper:Q",
+    )
+    cum_line = base3.mark_line(color="#1f7a3e", strokeWidth=2).encode(
+        y="cumulative_effect:Q",
+        tooltip=[
+            "period",
+            alt.Tooltip("cumulative_effect:Q", format=".5f"),
+            alt.Tooltip("cumulative_lower:Q", format=".5f"),
+            alt.Tooltip("cumulative_upper:Q", format=".5f"),
+        ],
+    )
+    panel3 = (cum_band + cum_line + zero_rule + event_rule).properties(
+        width=width,
+        height=height_per_panel,
+        title="Cumulative effect",
+    )
+
+    return alt.vconcat(panel1, panel2, panel3).resolve_scale(x="shared")  # type: ignore[no-any-return]

pycorpdiff/viz/collocation.py

@@ -0,0 +1,48 @@
+"""Collocation-shift visualisation — diverging horizontal bar."""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import pandas as pd
+
+if TYPE_CHECKING:
+    import altair as alt
+
+
+def collocation_diverging_bar(
+    df: pd.DataFrame,
+    n: int = 20,
+    width: int = 500,
+    height: int | None = None,
+) -> alt.Chart:
+    """Diverging bar chart of the top ``n`` collocate shifts.
+
+    Positive bars are A-leaning collocates (gained around the target in
+    A); negative bars are B-leaning (lost from A, gained in B). Sorted
+    by signed shift so the eye reads the divergence directly.
+    """
+    import altair as alt
+
+    subset = (
+        df.assign(_abs=df["shift"].abs()).nlargest(n, "_abs").drop(columns="_abs")
+    ).sort_values("shift", ascending=False)
+    if height is None:
+        height = max(200, 18 * len(subset))
+
+    chart = (
+        alt.Chart(subset)
+        .mark_bar()
+        .encode(
+            x=alt.X("shift:Q", title="Shift (A − B)"),
+            y=alt.Y("collocate:N", sort="-x", title=None),
+            color=alt.condition(
+                alt.datum.shift >= 0,
+                alt.value("#1f77b4"),
+                alt.value("#d62728"),
+            ),
+            tooltip=list(subset.columns),
+        )
+        .properties(width=width, height=height)
+    )
+    return chart  # type: ignore[no-any-return]

pycorpdiff/viz/dispersion.py

@@ -0,0 +1,117 @@
+"""Dispersion plots: *where* in a corpus does a term appear?
+
+A term can have the same overall frequency in two corpora and be
+distributed completely differently — even across one corpus, a high-
+frequency word can be clustered in a few documents or evenly spread.
+The dispersion plot answers "where" by marking each occurrence at the
+relevant document index along a horizontal axis.
+
+This is the classic Mosteller / Stubbs / Brezina visualisation for
+"how representative is a frequency count" — companion to the
+:func:`pycorpdiff.keyness.juilland_d` / ``dispersion_dp`` numerical
+measures pycorpdiff already exposes.
+"""
+
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
+
+import pandas as pd
+
+if TYPE_CHECKING:
+    import altair as alt
+
+    from ..corpus import Corpus, CorpusSlice
+
+
+def dispersion_plot(
+    corpus: Corpus | CorpusSlice,
+    targets: str | list[str],
+    width: int = 600,
+    height: int | None = None,
+) -> alt.Chart:
+    """Visualise *where* each ``target`` appears in the corpus.
+
+    Each occurrence of every target becomes a small vertical tick at
+    its document index. Stacked rows show one target at a time, so
+    "this term is concentrated in the first third of the corpus" or
+    "this term is evenly spread" reads off the plot at a glance.
+
+    Parameters
+    ----------
+    corpus
+        A :class:`Corpus` or :class:`CorpusSlice`.
+    targets
+        Single target or a list. Each gets its own horizontal row.
+    width
+        Chart width in pixels.
+    height
+        Chart height. If ``None``, scales with the number of targets.
+
+    Returns
+    -------
+    altair.Chart
+        Interactive chart with per-occurrence ticks coloured by target.
+        Requires the ``[viz]`` extra.
+
+    Example
+    -------
+    >>> import pycorpdiff as pcd
+    >>> corpus = pcd.load_hansard_sample()
+    >>> chart = pcd.viz.dispersion_plot(corpus, ['criminal', 'family'])  # doctest: +SKIP
+    >>> chart.save('dispersion.svg')                                      # doctest: +SKIP
+    """
+    import altair as alt
+
+    target_list = [targets] if isinstance(targets, str) else list(targets)
+    docs_tokens = corpus.tokens()
+    n_docs = len(docs_tokens)
+
+    # Collect (doc_index, target) for every occurrence.
+    rows: list[dict[str, Any]] = []
+    for doc_idx, tokens in enumerate(docs_tokens):
+        for token in tokens:
+            if token in target_list:
+                rows.append({"doc_index": doc_idx, "target": token})
+
+    if rows:
+        points = pd.DataFrame(rows, columns=["doc_index", "target"])
+    else:
+        # Empty corpus or no matches — emit a typed empty frame so
+        # altair can infer column types (it can't infer from empties).
+        points = pd.DataFrame(
+            {
+                "doc_index": pd.Series([], dtype="int64"),
+                "target": pd.Series([], dtype="object"),
+            }
+        )
+
+    if height is None:
+        height = max(60, 40 * len(target_list))
+
+    chart = (
+        alt.Chart(points)
+        .mark_tick(thickness=1.5)
+        .encode(
+            x=alt.X(
+                "doc_index:Q",
+                title=None,
+                scale=alt.Scale(domain=[0, max(1, n_docs - 1)]),
+            ),
+            y=alt.Y("target:N", title=None, sort=target_list),
+            color=alt.Color("target:N", legend=None, sort=target_list),
+            tooltip=[
+                alt.Tooltip("doc_index:Q"),
+                alt.Tooltip("target:N"),
+            ],
+        )
+        .properties(
+            width=width,
+            height=height,
+            title=alt.TitleParams(
+                text="Dispersion plot — where each term occurs in the corpus",
+                subtitle=f"{n_docs} documents on the x-axis, one tick per occurrence",
+            ),
+        )
+    )
+    return chart  # type: ignore[no-any-return]