activitymodel 0.1.0__py3-none-any.whl

activitymodel/__init__.py

@@ -0,0 +1,45 @@
+"""activitymodel: occupant activity calendars from the French time-use survey.
+
+Three layers, three consumers:
+
+* ``survey`` + ``clustering`` + ``pipeline`` build a :class:`DiaryLibrary`
+  from the Enquete Emploi du temps (run inside the buildingdata pipeline).
+* ``harmonise`` puts census individuals in the vocabulary the library was
+  built with (used by the census-occupant dataset and by building_eload).
+* ``matching`` + ``calendar`` turn occupants and a library into annual
+  activity calendars (used by building_eload at simulation time).
+
+No module resolves a data path at import time; data comes in as frames.
+"""
+
+from importlib.metadata import PackageNotFoundError, version
+
+try:
+    __version__ = version("activitymodel")
+except PackageNotFoundError:  # editable checkout without metadata
+    __version__ = "0.0.0"
+
+from .calendar import AnnualCalendar, day_types_for_dates, generate_calendar
+from .harmonise import CensusColumns, harmonise_census_individuals, harmonise_survey_individuals
+from .library import DiaryLibrary
+from .matching import assign_clusters
+from .taxonomy import DETAILED_ACTIVITIES, LEVELS, get_level
+from .vocabulary import DEFAULT_MATCHING_KEYS, DEFAULT_MATCHING_LADDER, AgeClasses
+
+__all__ = [
+    "AgeClasses",
+    "AnnualCalendar",
+    "CensusColumns",
+    "DEFAULT_MATCHING_KEYS",
+    "DEFAULT_MATCHING_LADDER",
+    "DETAILED_ACTIVITIES",
+    "DiaryLibrary",
+    "LEVELS",
+    "__version__",
+    "assign_clusters",
+    "day_types_for_dates",
+    "generate_calendar",
+    "get_level",
+    "harmonise_census_individuals",
+    "harmonise_survey_individuals",
+]

activitymodel/calendar.py

@@ -0,0 +1,286 @@
+"""Expand cluster assignments into annual activity calendars.
+
+An :class:`AnnualCalendar` is stored as an ``int32`` matrix ``day_index`` of
+shape ``(n_occupants, n_dates)`` pointing into the library's day pool: about
+1.5 KiB per occupant-year instead of 51 KiB for the expanded ``uint8`` codes
+and several hundred KiB for one polars column per occupant. Codes are
+materialised per occupant slice with :meth:`AnnualCalendar.codes`, which is the
+shape building_eload's chunked dwelling model wants (occupant-major, one shared
+time axis, 144 slots per day).
+
+The time axis is the survey's clock (``Europe/Paris`` wall-clock, no DST
+handling here); building_eload re-places it on its UTC axis with
+``place_local_clock_curve_on_utc_axis``.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from datetime import date, datetime, timedelta
+
+import numpy as np
+import polars as pl
+
+from .library import CLOCK, SLOT_MINUTES, SLOTS_PER_DAY, DiaryLibrary, cluster_column
+from .matching import MATCH_LEVEL_COLUMN, assign_clusters
+from .taxonomy import DETAILED_ACTIVITIES, get_level
+from .vocabulary import DEFAULT_MATCHING_LADDER
+
+# --------------------------------------------------------------------------- #
+# Calendar of day types
+# --------------------------------------------------------------------------- #
+
+
+def easter_sunday(year: int) -> date:
+    """Gregorian Easter (Meeus/Jones/Butcher)."""
+    a = year % 19
+    b, c = divmod(year, 100)
+    d, e = divmod(b, 4)
+    f = (b + 8) // 25
+    g = (b - f + 1) // 3
+    h = (19 * a + b - d - g + 15) % 30
+    i, k = divmod(c, 4)
+    m = (32 + 2 * e + 2 * i - h - k) % 7
+    n = (a + 11 * h + 22 * m) // 451
+    month, day = divmod(h + m - 7 * n + 114, 31)
+    return date(year, month, day + 1)
+
+
+def french_public_holidays(year: int) -> list[date]:
+    """The eleven national public holidays of metropolitan France."""
+    easter = easter_sunday(year)
+    return sorted(
+        [
+            date(year, 1, 1),
+            easter + timedelta(days=1),
+            date(year, 5, 1),
+            date(year, 5, 8),
+            easter + timedelta(days=39),
+            easter + timedelta(days=50),
+            date(year, 7, 14),
+            date(year, 8, 15),
+            date(year, 11, 1),
+            date(year, 11, 11),
+            date(year, 12, 25),
+        ]
+    )
+
+
+def year_dates(year: int) -> np.ndarray:
+    """All dates of ``year`` as ``datetime64[D]``."""
+    return np.arange(np.datetime64(f"{year}-01-01"), np.datetime64(f"{year + 1}-01-01"))
+
+
+def day_types_for_dates(
+    dates: np.ndarray,
+    holiday_day_type: str | None = "sunday",
+    holidays: list[date] | None = None,
+) -> np.ndarray:
+    """``weekday`` / ``saturday`` / ``sunday`` per date, holidays optionally remapped.
+
+    Args:
+        dates: ``datetime64[D]`` array.
+        holiday_day_type: day type given to public holidays (``None`` = keep
+            the calendar day type).
+        holidays: explicit holiday dates; default is the French national list
+            of every year present in ``dates``.
+    """
+    dates = np.asarray(dates, dtype="datetime64[D]")
+    weekday = (dates.astype("datetime64[D]").astype(np.int64) + 3) % 7  # 1970-01-01 was a Thursday
+    types = np.where(weekday == 5, "saturday", np.where(weekday == 6, "sunday", "weekday")).astype(
+        object
+    )
+    if holiday_day_type is not None:
+        if holidays is None:
+            years = np.unique(dates.astype("datetime64[Y]").astype(int) + 1970)
+            holidays = [d for y in years for d in french_public_holidays(int(y))]
+        holiday_set = np.array([np.datetime64(d) for d in holidays], dtype="datetime64[D]")
+        types[np.isin(dates, holiday_set)] = holiday_day_type
+    return types
+
+
+# --------------------------------------------------------------------------- #
+# Day sampling
+# --------------------------------------------------------------------------- #
+
+
+def sample_days(
+    library: DiaryLibrary,
+    assignments: pl.DataFrame,
+    date_day_types: np.ndarray,
+    rng: np.random.Generator,
+) -> np.ndarray:
+    """Draw one library day per (occupant, date) from the assigned cluster.
+
+    Args:
+        library: the diary library.
+        assignments: output of :func:`~activitymodel.matching.assign_clusters`,
+            in occupant order.
+        date_day_types: day type of each date (see :func:`day_types_for_dates`).
+        rng: generator driving the draw.
+
+    Returns:
+        ``int32`` matrix ``(n_occupants, n_dates)`` of row positions in
+        ``library.codes``. Within a cluster, days are drawn with replacement,
+        proportionally to their survey weight.
+    """
+    n_occ, n_dates = assignments.height, len(date_day_types)
+    day_index = np.full((n_occ, n_dates), -1, dtype=np.int32)
+    day_type_col = library.days.get_column("day_type").to_numpy().astype(object)
+    cluster_col = library.days.get_column("cluster_id").to_numpy()
+    weight_col = library.days.get_column("weight").to_numpy()
+    for day_type in library.day_types:
+        date_mask = date_day_types == day_type
+        if not date_mask.any():
+            continue
+        occupant_clusters = assignments.get_column(cluster_column(day_type)).to_numpy()
+        for cluster_id in np.unique(occupant_clusters):
+            pool = np.flatnonzero((cluster_col == cluster_id) & (day_type_col == day_type))
+            if pool.size == 0:
+                raise ValueError(f"cluster {cluster_id} has no {day_type} day in the library")
+            occupant_mask = occupant_clusters == cluster_id
+            count = int(occupant_mask.sum()) * int(date_mask.sum())
+            p = weight_col[pool] / weight_col[pool].sum()
+            draws = rng.choice(pool, size=count, p=p).astype(np.int32)
+            day_index[np.ix_(occupant_mask, date_mask)] = draws.reshape(
+                int(occupant_mask.sum()), -1
+            )
+    unmapped = date_day_types[(day_index < 0).any(axis=0)]
+    if unmapped.size:
+        raise ValueError(f"dates with day types outside the library: {sorted(set(unmapped))}")
+    return day_index
+
+
+# --------------------------------------------------------------------------- #
+# Annual calendar object
+# --------------------------------------------------------------------------- #
+
+
+@dataclass
+class AnnualCalendar:
+    """Compact annual calendars for a set of occupants.
+
+    Attributes:
+        library: the diary library the days come from.
+        occupant_ids: occupant identifiers, in row order.
+        dates: ``datetime64[D]`` dates covered (a whole year by default).
+        day_index: ``int32`` ``(n_occupants, n_dates)`` positions in ``library.codes``.
+        assignments: cluster per day type and match level per occupant.
+        time_zone: clock the slots are expressed in.
+    """
+
+    library: DiaryLibrary
+    occupant_ids: np.ndarray
+    dates: np.ndarray
+    day_index: np.ndarray
+    assignments: pl.DataFrame
+    time_zone: str = CLOCK
+
+    @property
+    def n_occupants(self) -> int:
+        return self.day_index.shape[0]
+
+    @property
+    def n_steps(self) -> int:
+        return self.day_index.shape[1] * SLOTS_PER_DAY
+
+    def timeline(self) -> pl.Series:
+        """Naive ``Datetime`` axis, one stamp per slot, local clock."""
+        start = datetime.combine(self.dates[0].astype(date), datetime.min.time())
+        end = start + timedelta(days=len(self.dates))
+        return pl.datetime_range(
+            start, end, interval=f"{SLOT_MINUTES}m", closed="left", eager=True, time_unit="us"
+        ).alias("datetime")
+
+    def codes(self, start: int = 0, stop: int | None = None, level: str = "detailed") -> np.ndarray:
+        """``uint8`` codes ``(n, n_steps)`` for occupants ``start:stop``, at ``level``."""
+        block = self.library.codes[self.day_index[start:stop]]
+        block = block.reshape(block.shape[0], -1)
+        if level != "detailed":
+            block = get_level(level).project(block)
+        return block
+
+    def to_long(self, start: int = 0, stop: int | None = None) -> pl.DataFrame:
+        """Occupant-major long frame: ``occupant_id``, ``datetime``, ``activity`` (Enum).
+
+        This is the layout building_eload's dwelling model consumes: every
+        occupant's whole year contiguous, time ascending.
+        """
+        codes = self.codes(start, stop)
+        n = codes.shape[0]
+        timeline = self.timeline()
+        activity = (
+            pl.Series("activity", codes.ravel(), dtype=pl.UInt8)
+            .to_frame()
+            .select(
+                pl.col("activity").replace_strict(
+                    dict(enumerate(DETAILED_ACTIVITIES)),
+                    return_dtype=pl.Enum(list(DETAILED_ACTIVITIES)),
+                )
+            )
+            .to_series()
+        )
+        ids = pl.Series("occupant_id", np.repeat(self.occupant_ids[start:stop], timeline.len()))
+        return pl.DataFrame([ids, pl.concat([timeline] * n) if n else timeline.head(0), activity])
+
+    def to_wide_book(self, start: int = 0, stop: int | None = None) -> pl.DataFrame:
+        """``datetime`` + one ``String`` column per occupant (``occupant_<i>``).
+
+        The legacy building_eload / buildingdata "occupant book" layout. Kept
+        for comparisons and for the current ``activity_file`` parameter; it is
+        the memory-hungry representation, not the one to build on.
+        """
+        codes = self.codes(start, stop)
+        names = np.array(DETAILED_ACTIVITIES, dtype=object)
+        columns = [self.timeline()]
+        for i, row in enumerate(codes, start=start):
+            columns.append(pl.Series(f"occupant_{i}", names[row], dtype=pl.String))
+        return pl.DataFrame(columns)
+
+    def match_level_counts(self) -> pl.DataFrame:
+        return (
+            self.assignments.group_by(MATCH_LEVEL_COLUMN)
+            .agg(pl.len().alias("n"))
+            .sort(MATCH_LEVEL_COLUMN)
+        )
+
+
+def generate_calendar(
+    library: DiaryLibrary,
+    occupants: pl.DataFrame,
+    year: int,
+    seed: int | None = 0,
+    ladder=DEFAULT_MATCHING_LADDER,
+    holiday_day_type: str | None = "sunday",
+    occupant_id: str = "occupant_id",
+) -> AnnualCalendar:
+    """Match occupants to clusters and draw a whole year of days for each.
+
+    Args:
+        library: the diary library.
+        occupants: one row per occupant with ``occupant_id`` and the harmonised
+            matching columns (see :mod:`activitymodel.harmonise`).
+        year: calendar year to generate (leap years give 366 days).
+        seed: master seed; matching and day sampling use two child streams so
+            changing the year does not change the cluster assignment.
+        ladder: matching fallback ladder.
+        holiday_day_type: day type used for French public holidays.
+        occupant_id: identifier column.
+
+    Returns:
+        An :class:`AnnualCalendar`.
+    """
+    match_rng, day_rng = (np.random.default_rng(s) for s in np.random.SeedSequence(seed).spawn(2))
+    assignments = assign_clusters(
+        library, occupants, seed=match_rng, ladder=ladder, occupant_id=occupant_id
+    )
+    dates = year_dates(year)
+    day_types = day_types_for_dates(dates, holiday_day_type=holiday_day_type)
+    day_index = sample_days(library, assignments, day_types, day_rng)
+    return AnnualCalendar(
+        library=library,
+        occupant_ids=occupants.get_column(occupant_id).to_numpy(),
+        dates=dates,
+        day_index=day_index,
+        assignments=assignments,
+    )

activitymodel/clustering.py

@@ -0,0 +1,220 @@
+"""Cluster coded diaries into daily activity patterns.
+
+The previous implementation ran Ward linkage on integer activity labels with a
+Jaccard metric, which is not a metric on ordinal codes and not the distance
+Ward minimises. Here a day is a one-hot matrix (144 slots x activities); the
+Euclidean distance between two such rows is a scaled Hamming distance between
+the two days, which is exactly what Ward's within-cluster variance measures.
+
+Two methods share the same representation:
+
+* ``ward`` (default): principal components of the one-hot matrix, then Ward
+  linkage (``fastcluster.linkage_vector`` when installed, scipy otherwise), cut
+  at the requested number of clusters. Deterministic.
+* ``kmeans``: weighted k-means++ on the same components. Honours the survey
+  weights, cheaper on very large pools, seeded.
+
+Cluster counts are a *choice* recorded in the library manifest;
+:func:`distance_growth` and :func:`silhouette_by_k` give the curves used to
+choose them.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+
+import numpy as np
+from scipy.cluster.hierarchy import fcluster
+from scipy.cluster.hierarchy import linkage as _scipy_linkage
+
+from .taxonomy import get_level
+
+try:  # optional memory-saving linkage
+    from fastcluster import linkage_vector as _fast_linkage
+except ModuleNotFoundError:  # pragma: no cover - exercised only without the extra
+    _fast_linkage = None
+
+
+@dataclass(frozen=True)
+class ClusteringOptions:
+    """How diaries are clustered, per day type."""
+
+    #: Number of clusters per day type.
+    n_clusters: dict[str, int] = field(
+        default_factory=lambda: {"weekday": 16, "saturday": 10, "sunday": 10}
+    )
+    method: str = "ward"
+    #: Taxonomy level whose activities define the one-hot representation.
+    level: str = "detailed"
+    #: Consecutive slots averaged together (3 = 30-minute resolution).
+    slot_step: int = 1
+    #: Principal components kept before clustering (``None`` = no reduction).
+    n_components: int | None = 50
+    seed: int = 0
+
+
+def one_hot(codes: np.ndarray, level: str = "detailed", slot_step: int = 1) -> np.ndarray:
+    """``uint8`` diaries ``(n, 144)`` -> float32 ``(n, slots * activities)``.
+
+    With ``slot_step > 1`` the one-hot vectors of consecutive slots are
+    averaged, giving the share of each activity within the coarser slot.
+    """
+    lvl = get_level(level)
+    projected = lvl.project(codes)
+    n, slots = projected.shape
+    if slots % slot_step:
+        raise ValueError(f"slot_step {slot_step} does not divide {slots} slots")
+    eye = np.eye(lvl.n_activities, dtype=np.float32)
+    matrix = eye[projected]  # (n, slots, activities)
+    if slot_step > 1:
+        matrix = matrix.reshape(n, slots // slot_step, slot_step, lvl.n_activities).mean(axis=2)
+    return matrix.reshape(n, -1)
+
+
+def reduce_components(matrix: np.ndarray, n_components: int | None) -> np.ndarray:
+    """Project onto the leading principal components (centred SVD)."""
+    if n_components is None or n_components >= min(matrix.shape):
+        return matrix
+    centred = matrix - matrix.mean(axis=0, keepdims=True)
+    # Economy SVD on the smaller Gram side keeps memory at O(d^2).
+    cov = centred.T @ centred
+    values, vectors = np.linalg.eigh(cov)
+    order = np.argsort(values)[::-1][:n_components]
+    return (centred @ vectors[:, order]).astype(np.float32)
+
+
+def ward_linkage(matrix: np.ndarray) -> np.ndarray:
+    """Ward linkage matrix, memory-saving when fastcluster is available."""
+    data = np.ascontiguousarray(matrix, dtype=np.float64)
+    if _fast_linkage is not None:
+        return _fast_linkage(data, method="ward")
+    return _scipy_linkage(data, method="ward")
+
+
+def cut(linkage: np.ndarray, n_clusters: int) -> np.ndarray:
+    """Cut a linkage matrix into ``n_clusters`` labels ``0..n_clusters-1``."""
+    return fcluster(linkage, t=n_clusters, criterion="maxclust").astype(np.int32) - 1
+
+
+def distance_growth(linkage: np.ndarray, max_k: int = 60) -> np.ndarray:
+    """Merge height reached when going from ``k`` to ``k-1`` clusters, for k=2..max_k.
+
+    Returns an array of shape ``(max_k - 1, 2)`` with columns ``k`` and
+    ``height``. A knee in this curve is the usual elbow criterion.
+    """
+    heights = linkage[:, 2][::-1]  # last merge first
+    ks = np.arange(2, min(max_k, len(heights) + 1) + 1)
+    return np.column_stack([ks, heights[ks - 2]])
+
+
+def weighted_kmeans(
+    matrix: np.ndarray,
+    n_clusters: int,
+    weights: np.ndarray | None,
+    rng: np.random.Generator,
+    n_init: int = 4,
+    max_iter: int = 100,
+) -> np.ndarray:
+    """Weighted Lloyd's algorithm with k-means++ seeding. Returns labels."""
+    n = matrix.shape[0]
+    w = np.ones(n) if weights is None else np.asarray(weights, dtype=float)
+    best_labels, best_inertia = None, np.inf
+    for _ in range(n_init):
+        centres = _kmeanspp(matrix, n_clusters, w, rng)
+        labels = None
+        for _ in range(max_iter):
+            distances = _sq_distances(matrix, centres)
+            new_labels = distances.argmin(axis=1)
+            if labels is not None and np.array_equal(new_labels, labels):
+                break
+            labels = new_labels
+            for k in range(n_clusters):
+                mask = labels == k
+                if mask.any():
+                    centres[k] = np.average(matrix[mask], axis=0, weights=w[mask])
+                else:  # re-seed an empty cluster on the farthest point
+                    centres[k] = matrix[distances.min(axis=1).argmax()]
+        inertia = float((w * distances[np.arange(n), labels]).sum())
+        if inertia < best_inertia:
+            best_inertia, best_labels = inertia, labels.copy()
+    return _relabel_by_size(best_labels, n_clusters)
+
+
+def _sq_distances(matrix: np.ndarray, centres: np.ndarray) -> np.ndarray:
+    # Clamped at zero: the expanded form can go slightly negative in float32.
+    return np.maximum(
+        (matrix**2).sum(axis=1)[:, None]
+        - 2 * matrix @ centres.T
+        + (centres**2).sum(axis=1)[None, :],
+        0.0,
+    )
+
+
+def _kmeanspp(matrix, n_clusters, weights, rng):
+    n = matrix.shape[0]
+    centres = np.empty((n_clusters, matrix.shape[1]), dtype=matrix.dtype)
+    centres[0] = matrix[rng.choice(n, p=weights / weights.sum())]
+    closest = _sq_distances(matrix, centres[:1]).ravel()
+    for k in range(1, n_clusters):
+        p = weights * closest
+        p = p / p.sum() if p.sum() > 0 else weights / weights.sum()
+        centres[k] = matrix[rng.choice(n, p=p)]
+        closest = np.minimum(closest, _sq_distances(matrix, centres[k : k + 1]).ravel())
+    return centres
+
+
+def _relabel_by_size(labels: np.ndarray, n_clusters: int) -> np.ndarray:
+    """Renumber clusters by decreasing size so labels are stable to read."""
+    counts = np.bincount(labels, minlength=n_clusters)
+    order = np.argsort(-counts, kind="stable")
+    mapping = np.empty(n_clusters, dtype=np.int32)
+    mapping[order] = np.arange(n_clusters, dtype=np.int32)
+    return mapping[labels]
+
+
+def cluster_codes(
+    codes: np.ndarray,
+    n_clusters: int,
+    options: ClusteringOptions,
+    weights: np.ndarray | None = None,
+    rng: np.random.Generator | None = None,
+) -> np.ndarray:
+    """Cluster one pool of diaries (one day type). Returns labels ``0..n_clusters-1``."""
+    if n_clusters < 1:
+        raise ValueError("n_clusters must be at least 1")
+    if codes.shape[0] <= n_clusters:
+        raise ValueError(f"cannot form {n_clusters} clusters from {codes.shape[0]} diaries")
+    matrix = reduce_components(
+        one_hot(codes, options.level, options.slot_step), options.n_components
+    )
+    if options.method == "ward":
+        return _relabel_by_size(cut(ward_linkage(matrix), n_clusters), n_clusters)
+    if options.method == "kmeans":
+        rng = np.random.default_rng(options.seed) if rng is None else rng
+        return weighted_kmeans(matrix, n_clusters, weights, rng)
+    raise ValueError(f"unknown clustering method {options.method!r}")
+
+
+def silhouette_by_k(
+    matrix: np.ndarray, linkage: np.ndarray, ks, rng, sample: int = 3000
+) -> np.ndarray:
+    """Mean silhouette on a random sample for each ``k`` (columns ``k``, ``silhouette``)."""
+    n = matrix.shape[0]
+    idx = rng.choice(n, size=min(sample, n), replace=False)
+    sub = matrix[idx]
+    d = np.sqrt(np.maximum(_sq_distances(sub, sub), 0))
+    out = []
+    for k in ks:
+        labels = cut(linkage, k)[idx]
+        s = np.zeros(len(idx))
+        for i in range(len(idx)):
+            own = labels == labels[i]
+            own[i] = False
+            a = d[i, own].mean() if own.any() else 0.0
+            b = min(
+                (d[i, labels == other].mean() for other in np.unique(labels) if other != labels[i]),
+                default=0.0,
+            )
+            s[i] = 0.0 if max(a, b) == 0 else (b - a) / max(a, b)
+        out.append((k, float(s.mean())))
+    return np.array(out)

activitymodel/diagnostics/__init__.py

@@ -0,0 +1,14 @@
+"""Checks that the library and the matching behave: profiles, distances, effect scores."""
+
+from .association import association_scores, weighted_cramers_v
+from .distance import matrix_distance
+from .profiles import activity_probability_matrix, calendar_profile, profile_frame
+
+__all__ = [
+    "activity_probability_matrix",
+    "association_scores",
+    "calendar_profile",
+    "matrix_distance",
+    "profile_frame",
+    "weighted_cramers_v",
+]

activitymodel/diagnostics/association.py

@@ -0,0 +1,93 @@
+"""How much each socio-demographic variable explains the cluster a diary falls in.
+
+Used to order the matching ladder: the variables that carry the least signal
+are the first to be dropped when a cell is empty.
+"""
+
+from __future__ import annotations
+
+import numpy as np
+import polars as pl
+
+
+def weighted_cramers_v(x: np.ndarray, y: np.ndarray, weights: np.ndarray) -> float:
+    """Weighted Cramer's V between two integer-coded categorical arrays."""
+    n_x, n_y = int(x.max()) + 1, int(y.max()) + 1
+    table = np.zeros((n_x, n_y))
+    np.add.at(table, (x, y), weights)
+    total = table.sum()
+    if total <= 0 or n_x < 2 or n_y < 2:
+        return 0.0
+    expected = np.outer(table.sum(axis=1), table.sum(axis=0)) / total
+    valid = expected > 0
+    chi2 = (((table - expected) ** 2)[valid] / expected[valid]).sum()
+    return float(np.sqrt(max(chi2 / (total * min(n_x - 1, n_y - 1)), 0.0)))
+
+
+def association_scores(
+    frame: pl.DataFrame,
+    target: str,
+    variables: list[str],
+    weight: str | None = None,
+    n_permutations: int = 199,
+    seed: int = 0,
+    group: str | None = None,
+) -> pl.DataFrame:
+    """Cramer's V of each variable with ``target`` plus a permutation p-value.
+
+    Args:
+        frame: one row per diary with the target (cluster label), the
+            variables and optionally a weight and a grouping column (the
+            person, so permutations respect the two-diaries-per-person design).
+        target: cluster label column.
+        variables: categorical columns to score.
+        weight: weight column (``None`` = unweighted).
+        n_permutations: permutations for the p-value.
+        seed: RNG seed.
+        group: column defining the permutation unit.
+
+    Returns:
+        Frame ``variable``, ``n_categories``, ``cramers_v``, ``p_value``, sorted
+        by decreasing association.
+    """
+    rng = np.random.default_rng(seed)
+    y = frame.get_column(target).cast(pl.String).rank("dense").cast(pl.Int64).to_numpy() - 1
+    w = (
+        np.ones(frame.height)
+        if weight is None
+        else frame.get_column(weight).cast(pl.Float64).fill_null(0).to_numpy()
+    )
+    if group is not None:
+        groups = frame.get_column(group).cast(pl.String).rank("dense").cast(pl.Int64).to_numpy() - 1
+    rows = []
+    for variable in variables:
+        x = (
+            frame.get_column(variable)
+            .cast(pl.String)
+            .fill_null("<null>")
+            .rank("dense")
+            .cast(pl.Int64)
+            .to_numpy()
+            - 1
+        )
+        observed = weighted_cramers_v(x, y, w)
+        at_least = 0
+        for _ in range(n_permutations):
+            if group is None:
+                shuffled = rng.permutation(x)
+            else:
+                # permute the variable at the group level: every diary of a
+                # person receives the value of another person.
+                first = np.unique(groups, return_index=True)[1]
+                values = x[first]
+                shuffled = rng.permutation(values)[np.searchsorted(groups[first], groups)]
+            at_least += weighted_cramers_v(shuffled, y, w) >= observed
+        rows.append(
+            {
+                "variable": variable,
+                "n_categories": int(x.max()) + 1,
+                "cramers_v": observed,
+                "p_value": (at_least + 1) / (n_permutations + 1),
+            }
+        )
+    return pl.DataFrame(rows).sort("cramers_v", descending=True)

activitymodel/diagnostics/distance.py

@@ -0,0 +1,33 @@
+"""Distance between two probability matrices with the same row/column layout."""
+
+from __future__ import annotations
+
+import numpy as np
+
+
+def matrix_distance(a, b, method: str = "tv", normalize: bool = True) -> float:
+    """Mean per-row distance in ``[0, 1]`` between two probability matrices.
+
+    ``tv`` is the mean total-variation distance, ``hellinger`` the mean
+    Hellinger distance. 0 means identical rows, 1 means disjoint support on
+    every row.
+    """
+    a = np.asarray(a, dtype=float)
+    b = np.asarray(b, dtype=float)
+    if a.shape != b.shape:
+        raise ValueError(f"shape mismatch: {a.shape} vs {b.shape}")
+    if normalize:
+        a, b = _normalise_rows(a), _normalise_rows(b)
+    if method == "tv":
+        per_row = 0.5 * np.abs(a - b).sum(axis=1)
+    elif method == "hellinger":
+        per_row = np.sqrt(0.5 * ((np.sqrt(a) - np.sqrt(b)) ** 2).sum(axis=1))
+    else:
+        raise ValueError("method must be 'tv' or 'hellinger'")
+    return float(per_row.mean())
+
+
+def _normalise_rows(m: np.ndarray) -> np.ndarray:
+    totals = m.sum(axis=1, keepdims=True)
+    totals[totals == 0] = 1.0
+    return m / totals