geometry.py

"""Three-axis spatial-geometric profile of a single text.

Implements a simplified, single-text reading of Vallego (2026),
"A Spatial-Geometric Framework for Discourse Comparison".

Method
------
1. Each axis is defined by two poles, each a set of exemplar sentences.
   Its raw direction is the difference between the two pole centroids
   (high-pole centroid minus low-pole centroid) in embedding space.
2. The three raw directions overlap (they are not orthogonal). They are
   orthogonalised with symmetric (Loewdin) orthogonalisation, which finds
   the orthonormal set closest to the raw directions and, unlike
   Gram-Schmidt, treats all three axes equally (no order dependence).
3. A text is split into sentences; each sentence and the text centroid
   (mean of sentence embeddings) are projected onto the three orthonormal
   directions. These three numbers are the 3D coordinates.
4. Each coordinate is rescaled so that 0% = the low-pole centroid and
   100% = the high-pole centroid on that axis. Values below 0 or above 100
   lie further out than the typical exemplar of that pole.
"""

from __future__ import annotations

import re
from dataclasses import dataclass

import numpy as np
import pandas as pd

AXES = [
    {"id": "eco", "name": "Ecocentrism", "low": "Anthropocentric", "high": "Ecocentric"},
    {"id": "change", "name": "Societal change", "low": "Conservative", "high": "Transformative"},
    {"id": "outlook", "name": "Outlook", "low": "Optimistic", "high": "Pessimistic"},
]


def split_sentences(text: str, min_words: int = 3) -> list[str]:
    parts = re.split(r"(?<=[.!?;])\s+|\n+", text.strip())
    sents = [p.strip() for p in parts if len(p.strip().split()) >= min_words]
    return sents or ([text.strip()] if text.strip() else [])


def _normalise(x: np.ndarray) -> np.ndarray:
    return x / np.linalg.norm(x, axis=-1, keepdims=True).clip(min=1e-12)


def loewdin(D: np.ndarray) -> np.ndarray:
    """Symmetric orthogonalisation of the rows of D: Q = (D D^T)^(-1/2) D."""
    S = D @ D.T
    vals, vecs = np.linalg.eigh(S)
    S_inv_sqrt = vecs @ np.diag(1.0 / np.sqrt(vals)) @ vecs.T
    return S_inv_sqrt @ D


@dataclass
class Space:
    axes: list[dict]
    Q: np.ndarray            # 3 x d, orthonormal axis directions
    raw_dirs: np.ndarray     # 3 x d, unit raw directions (before orthogonalisation)
    lo: np.ndarray           # 3, projection of each low-pole centroid on its own axis
    hi: np.ndarray           # 3, projection of each high-pole centroid on its own axis
    raw_lo: np.ndarray
    raw_hi: np.ndarray
    pole_points: dict        # (axis_id, pole) -> (sentences, coords n x 3 in %)
    pole_centroids: dict     # (axis_id, pole) -> coords (3,) in %
    overlap: pd.DataFrame    # cosine between raw directions
    topic: np.ndarray = None  # unit vector: centre of all pole exemplars ("the topic")

    def relevance(self, E: np.ndarray) -> np.ndarray:
        """Cosine similarity of each (unit-length) embedding to the topic centre."""
        return E @ self.topic

    def to_percent(self, proj: np.ndarray) -> np.ndarray:
        return 100.0 * (proj - self.lo) / (self.hi - self.lo)

    def coords(self, E: np.ndarray) -> np.ndarray:
        return self.to_percent(E @ self.Q.T)

    def raw_percent(self, E: np.ndarray) -> np.ndarray:
        return 100.0 * (E @ self.raw_dirs.T - self.raw_lo) / (self.raw_hi - self.raw_lo)


def build_space(poles: pd.DataFrame, encode, axes=AXES) -> Space:
    """poles: DataFrame with columns axis, pole ('low'/'high'), text.
    encode: function list[str] -> np.ndarray (n x d)."""
    emb, cents = {}, {}
    for ax in axes:
        for pole in ("low", "high"):
            texts = poles[(poles.axis == ax["id"]) & (poles.pole == pole)].text.tolist()
            if len(texts) < 2:
                raise ValueError(f"Axis '{ax['id']}' pole '{pole}' needs at least 2 exemplars.")
            E = _normalise(np.asarray(encode(texts), dtype=float))
            emb[(ax["id"], pole)] = (texts, E)
            cents[(ax["id"], pole)] = E.mean(axis=0)

    raw = np.stack([cents[(a["id"], "high")] - cents[(a["id"], "low")] for a in axes])
    raw_unit = _normalise(raw)
    Q = loewdin(raw)

    lo = np.array([cents[(a["id"], "low")] @ Q[k] for k, a in enumerate(axes)])
    hi = np.array([cents[(a["id"], "high")] @ Q[k] for k, a in enumerate(axes)])
    raw_lo = np.array([cents[(a["id"], "low")] @ raw_unit[k] for k, a in enumerate(axes)])
    raw_hi = np.array([cents[(a["id"], "high")] @ raw_unit[k] for k, a in enumerate(axes)])

    names = [f"{a['low']}–{a['high']}" for a in axes]
    overlap = pd.DataFrame(raw_unit @ raw_unit.T, index=names, columns=names).round(2)

    topic = _normalise(np.vstack([E for _, E in emb.values()]).mean(axis=0))
    space = Space(axes, Q, raw_unit, lo, hi, raw_lo, raw_hi, {}, {}, overlap, topic)
    for key, (texts, E) in emb.items():
        space.pole_points[key] = (texts, space.coords(E))
        space.pole_centroids[key] = space.coords(cents[key][None, :])[0]
    return space


def reading(pct: float, low: str, high: str) -> str:
    if pct < 10:
        return f"Strongly {low.lower()}"
    if pct < 40:
        return f"Leans {low.lower()}"
    if pct <= 60:
        return "Between the poles"
    if pct <= 90:
        return f"Leans {high.lower()}"
    return f"Strongly {high.lower()}"


def analyse(text: str, space: Space, encode, threshold: float | None = None) -> dict:
    """threshold: if given, only sentences whose similarity to the topic centre is at
    least this value are analysed (relevance filter). If none pass, all are used."""
    all_sents = split_sentences(text)
    if not all_sents:
        raise ValueError("Please enter some text.")
    E_all = _normalise(np.asarray(encode(all_sents), dtype=float))
    rel = space.relevance(E_all)
    keep = np.ones(len(all_sents), dtype=bool) if threshold is None else rel >= threshold
    fallback = bool(threshold is not None and not keep.any())
    if fallback:
        keep[:] = True

    sents = [s for s, k in zip(all_sents, keep) if k]
    E = E_all[keep]
    all_xyz = space.coords(E_all)
    centroid = E.mean(axis=0)
    sent_xyz = space.coords(E)
    text_xyz = space.coords(centroid[None, :])[0]
    raw_pct = space.raw_percent(centroid[None, :])[0]

    rows = []
    for k, ax in enumerate(space.axes):
        rows.append({
            "Axis (0% → 100%)": f"{ax['low']} → {ax['high']}",
            "Position": round(float(text_xyz[k]), 1),
            "Reading": reading(float(text_xyz[k]), ax["low"], ax["high"]),
            "Spread across sentences (±)": round(float(sent_xyz[:, k].std()), 1) if len(sents) > 1 else None,
            "Before orthogonalisation": round(float(raw_pct[k]), 1),
        })
    cloud = cloud_measures(E, sent_xyz, space)
    return {
        "sentences": sents,                 # analysed sentences
        "sent_xyz": sent_xyz,
        "sent_rel": rel[keep],
        "off_sentences": [s for s, k in zip(all_sents, keep) if not k],
        "off_xyz": all_xyz[~keep],
        "off_rel": rel[~keep],
        "n_total": len(all_sents),
        "threshold": threshold,
        "fallback": fallback,
        "unfiltered_xyz": space.coords(E_all.mean(axis=0)[None, :])[0],
        "text_xyz": text_xyz,
        "table": pd.DataFrame(rows),
        "cloud": cloud,
        "cloud_table": cloud_table(cloud, space),
    }


def cloud_measures(E: np.ndarray, sent_xyz: np.ndarray, space: Space) -> dict:
    """Spread (Frobenius norm) and eigenanalysis of the cloud of sentences.

    - frobenius: Frobenius norm of the centred sentence positions (3D, % units),
      divided by sqrt(n) = root-mean-square distance of a sentence from the text centre.
    - captured: Pythagorean decomposition in the full embedding space. The squared
      Frobenius norm of the centred sentence embeddings splits exactly into the part
      lying along the three orthonormal axes and the residual; this is the axes' share.
    - eigen: eigen-decomposition of the 3x3 covariance of the sentence positions.
    """
    n = len(sent_xyz)
    out = {"n": n}
    if n < 2:
        return out
    Xc = sent_xyz - sent_xyz.mean(axis=0)
    out["frobenius"] = float(np.linalg.norm(Xc) / np.sqrt(n))

    Ec = E - E.mean(axis=0)
    total = float(np.sum(Ec ** 2))
    on_axes = float(np.sum((Ec @ space.Q.T) ** 2))
    out["captured"] = on_axes / total if total > 0 else float("nan")

    if n >= 3:
        cov = Xc.T @ Xc / n
        vals, vecs = np.linalg.eigh(cov)
        order = np.argsort(vals)[::-1]
        vals, vecs = np.clip(vals[order], 0, None), vecs[:, order]
        out["eigvals"] = vals
        out["eigvecs"] = vecs                      # columns = principal directions
        out["share_main"] = float(vals[0] / vals.sum()) if vals.sum() > 0 else float("nan")
        v1 = vecs[:, 0]
        out["angles"] = np.degrees(np.arccos(np.clip(np.abs(v1), 0, 1)))  # to each axis
        out["main_axis"] = int(np.argmax(np.abs(v1)))
    return out


def cloud_table(c: dict, space: Space) -> pd.DataFrame:
    rows = []

    def add(measure, value, meaning):
        rows.append({"Measure": measure, "Value": value, "What it means": meaning})

    if c["n"] < 2:
        add("Spread and main direction", "Needs 2+ sentences",
            "These measures describe how the sentences of a text spread out.")
        return pd.DataFrame(rows)

    add("Spread (Frobenius norm)", f"{c['frobenius']:.1f} points",
        "Average distance of a sentence from the text's centre across the three axes. "
        "Lower means a more consistent text.")
    add("Variation captured by the axes", f"{100 * c['captured']:.0f}%",
        "Share of the sentences' total variation in meaning that lies along these three axes. "
        "The rest concerns other things, such as topic or style.")
    if "eigvals" in c:
        ax = space.axes[c["main_axis"]]
        add("Main direction of variation (eigenanalysis)",
            f"Mostly {ax['name']} ({c['angles'][c['main_axis']]:.0f}° from it)",
            "The direction along which the sentences differ most, and the axis it is closest to.")
        add("Variation along the main direction", f"{100 * c['share_main']:.0f}%",
            "Near 33%: sentences vary evenly in all directions. "
            "Near 100%: they vary along a single line.")
    else:
        add("Main direction of variation (eigenanalysis)", "Needs 3+ sentences", "")
    return pd.DataFrame(rows)


def voronoi_zones(points: np.ndarray, box: tuple) -> list:
    """Zones of a 2D map: for each point, the polygon of the plane that is closer to it
    than to any other point (a Voronoi cell), clipped to box = (xmin, xmax, ymin, ymax).
    Pure numpy (Sutherland-Hodgman clipping by perpendicular bisectors)."""
    xmin, xmax, ymin, ymax = box
    square = np.array([[xmin, ymin], [xmax, ymin], [xmax, ymax], [xmin, ymax]], dtype=float)
    zones = []
    for i, p in enumerate(points):
        poly = square
        for j, q in enumerate(points):
            if i == j or np.allclose(p, q):
                continue
            n, mid = q - p, (p + q) / 2.0          # keep the side of the bisector nearer to p
            inside = lambda v: (v - mid) @ n <= 1e-9
            out = []
            for k in range(len(poly)):
                a, b = poly[k], poly[(k + 1) % len(poly)]
                ina, inb = inside(a), inside(b)
                if ina:
                    out.append(a)
                if ina != inb:
                    t = ((mid - a) @ n) / ((b - a) @ n)
                    out.append(a + t * (b - a))
            poly = np.array(out) if out else np.empty((0, 2))
            if len(poly) == 0:
                break
        zones.append(poly)
    return zones