Skip to content

qlroms.transitions

TransitionMaps dataclass

Transition operators with selectable mapping strategy.

Every distance is measured in the inner-product metric of the charts, ||x||_M^2 = x^T Mw x (Mw None = Euclidean, a diagonal, or a full FEM mass matrix).

Methods:

Name Description
- "atlas"

use Tgk/dgk/Tkg/dkg (local->global->local), exact when the atlas contains each local affine space.
- "pairwise"

use tmap/tshift direct chart-to-chart maps. The distance to centroid i of a state a in chart k is exact, d_i^2 = ||tshift[i,k] - a||^2 + oos2[i,k], the in-chart part plus the part of c_i - c_k that lies OUTSIDE chart k's subspace (oos2, a (K, K) matrix of squared M distances computed once from the geometry). Without oos2 (a hand-built TransitionMaps that never saw the centroids) only the in-chart part is available.
- "auto"

atlas when available, else pairwise.
- "physical"

DISTANCES only -- lift the state to physical space (centroid + Phi a) and compare to every cluster centroid directly, the ground truth the other methods reproduce (O(N) per step; needs centroids/Phi_all, see from_model): sqrt((c - x)^T Mw (c - x)), for any Mw. The map itself falls back to atlas/pairwise.
Source code in qlroms/transitions.py 
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
@dataclass
class TransitionMaps:
    """Transition operators with selectable mapping strategy.

    Every distance is measured in the inner-product metric of the charts, ||x||_M^2 =
    x^T Mw x (Mw None = Euclidean, a diagonal, or a full FEM mass matrix).

    Methods:
        - "atlas": use Tgk/dgk/Tkg/dkg (local->global->local), exact when the
            atlas contains each local affine space.
        - "pairwise": use tmap/tshift direct chart-to-chart maps. The distance to
            centroid i of a state a in chart k is exact,
                d_i^2 = ||tshift[i,k] - a||^2 + oos2[i,k],
            the in-chart part plus the part of c_i - c_k that lies OUTSIDE chart k's
            subspace (oos2, a (K, K) matrix of squared M distances computed once from
            the geometry). Without oos2 (a hand-built TransitionMaps that never saw the
            centroids) only the in-chart part is available.
        - "auto": atlas when available, else pairwise.
        - "physical": DISTANCES only -- lift the state to physical space
            (centroid + Phi a) and compare to every cluster centroid directly, the
            ground truth the other methods reproduce (O(N) per step; needs
            centroids/Phi_all, see from_model): sqrt((c - x)^T Mw (c - x)), for any Mw.
            The map itself falls back to atlas/pairwise.
    """

    tmap: Any | None = None
    tshift: Any | None = None
    Tgk: Any | None = None
    dgk: Any | None = None
    Tkg: Any | None = None
    dkg: Any | None = None
    centroids: Any | None = None   # (K, N), "physical" distances only
    Phi_all: Any | None = None     # (N, r, K), "physical" distances only
    Mw: Any | None = None          # inner-product weight, "physical" distances only
    method: TRANSITION_METHODS = "auto"
    oos2: Any | None = None        # (K, K) squared M distance of centroid i outside chart k's subspace
    wkind: tuple | None = None     # cache (id(Mw), weight_kind(Mw)); recomputed whenever Mw is replaced

    def _kind(self) -> str:
        if self.wkind is None or self.wkind[0] != id(self.Mw):
            self.wkind = (id(self.Mw), weight_kind(self.Mw))
        return self.wkind[1]

    @classmethod
    def from_model(cls, model: Any, method: TRANSITION_METHODS = "auto") -> TransitionMaps:
        # The chart inner-product weight of the model (a ks2d-style model that only exposes
        # its quadrature-weight grid `wt` is read through it).
        Mw = getattr(model, "Mw", None)
        if Mw is None and getattr(model, "wt", None) is not None:
            Mw = torch.as_tensor(model.wt).reshape(-1)
        return cls(
            tmap=getattr(model, "tmap", None),
            tshift=getattr(model, "tshift", None),
            Tgk=getattr(model, "Tgk", None),
            dgk=getattr(model, "dgk", None),
            Tkg=getattr(model, "Tkg", None),
            dkg=getattr(model, "dkg", None),
            centroids=getattr(model, "centroids", None),
            Phi_all=getattr(model, "Phi_all", None),
            Mw=Mw,
            method=method,
            oos2=getattr(model, "oos2", None),
            wkind=((id(Mw), model._wkind) if Mw is getattr(model, "Mw", None)
                   and getattr(model, "_wkind", None) else None),
        )

    @property
    def has_pairwise(self) -> bool:
        return self.tmap is not None and self.tshift is not None

    @property
    def has_atlas(self) -> bool:
        return self.Tgk is not None and self.dgk is not None and self.Tkg is not None and self.dkg is not None

    def _use_atlas(self) -> bool:
        if self.method == "physical":
            # the physical choice only affects distances(); maps stay exact either way
            if self.has_atlas:
                return True
            if self.has_pairwise:
                return False
            raise ValueError("No transition maps available (need atlas or pairwise maps).")
        if self.method == "atlas":
            if not self.has_atlas:
                raise ValueError("Atlas transitions requested but Tgk/dgk/Tkg/dkg are missing.")
            return True
        if self.method == "pairwise":
            if not self.has_pairwise:
                raise ValueError("Pairwise transitions requested but tmap/tshift are missing.")
            return False
        if self.has_atlas:
            return True
        if self.has_pairwise:
            return False
        raise ValueError("No transition maps available (need atlas or pairwise maps).")

    def to_atlas(self, a, k_from: int):
        """Lift coordinate(s) a from chart k_from into the shared atlas frame:
        z = Tgk[k_from] a + dgk[k_from]. a: (r,) or (r, m), torch or numpy."""
        if not self.has_atlas:
            raise ValueError("No atlas maps (Tgk/dgk); build the global atlas first.")
        col = (lambda v: v[:, None]) if a.ndim == 2 else (lambda v: v)
        return self.Tgk[k_from] @ a + col(self.dgk[k_from])

    def from_atlas(self, z, k_to: int):
        """Drop atlas coordinate(s) z into chart k_to: a = Tkg[k_to] z + dkg[k_to]
        (the exact inverse of to_atlas on chart k_to's affine subspace)."""
        if not self.has_atlas:
            raise ValueError("No atlas maps (Tkg/dkg); build the global atlas first.")
        col = (lambda v: v[:, None]) if z.ndim == 2 else (lambda v: v)
        return self.Tkg[k_to] @ z + col(self.dkg[k_to])

    def map(self, a, k_from: int, k_to: int):
        """Map coordinate(s) a from chart k_from to chart k_to.

        a: (r,) or (r, m), torch.Tensor or np.ndarray (the ops are type-generic).
        """
        if self._use_atlas():
            return self.from_atlas(self.to_atlas(a, k_from), k_to)
        col = (lambda v: v[:, None]) if a.ndim == 2 else (lambda v: v)
        return self.tmap[k_from, k_to] @ a + col(self.tshift[k_from, k_to])

    def distances(self, a, k_from: int):
        """Distance from state (a in chart k_from) to all cluster centroids.

        Returns:
            (K,) vector with distances in chosen transition metric.
        """
        if self.method == "physical":
            if self.centroids is None or self.Phi_all is None:
                raise ValueError("method='physical' needs centroids and Phi_all (see from_model).")
            a_t = a if torch.is_tensor(a) else torch.as_tensor(a, dtype=self.centroids.dtype,
                                                               device=self.centroids.device)
            x_hat = self.centroids[k_from] + self.Phi_all[:, :, k_from] @ a_t.reshape(-1)   # (N,)
            diff = self.centroids - x_hat[None, :]                                          # (K, N)
            kind = self._kind()
            if kind == "none":
                d = diff.norm(dim=1)
            elif kind in ("uniform", "diag"):
                wt = torch.as_tensor(as_diagonal(self.Mw), dtype=diff.dtype, device=diff.device)
                d = torch.sqrt((wt * diff**2).sum(dim=1))
            else:
                d = torch.sqrt(wsq_rows(self.Mw, diff).clamp_min(0))
            return d if torch.is_tensor(a) else d.cpu().numpy()
        use_atlas = self._use_atlas()
        if use_atlas:
            assert self.Tgk is not None and self.dgk is not None
            z = self.Tgk[k_from] @ a + self.dgk[k_from]
            return (self.dgk - z).norm(dim=1) if torch.is_tensor(a) else np.linalg.norm(self.dgk - z, axis=1)
        assert self.tshift is not None
        if self.oos2 is not None:
            if torch.is_tensor(a):
                return torch.sqrt(((self.tshift[:, k_from] - a) ** 2).sum(dim=1) + self.oos2[:, k_from])
            return np.sqrt(((self.tshift[:, k_from] - a) ** 2).sum(axis=1) + self.oos2[:, k_from])
        return (self.tshift[:, k_from] - a).norm(dim=1) if torch.is_tensor(a) else np.linalg.norm(self.tshift[:, k_from] - a, axis=1)

to_atlas(a, k_from)

Lift coordinate(s) a from chart k_from into the shared atlas frame: z = Tgk[k_from] a + dgk[k_from]. a: (r,) or (r, m), torch or numpy.

Source code in qlroms/transitions.py 
124
125
126
127
128
129
130
def to_atlas(self, a, k_from: int):
    """Lift coordinate(s) a from chart k_from into the shared atlas frame:
    z = Tgk[k_from] a + dgk[k_from]. a: (r,) or (r, m), torch or numpy."""
    if not self.has_atlas:
        raise ValueError("No atlas maps (Tgk/dgk); build the global atlas first.")
    col = (lambda v: v[:, None]) if a.ndim == 2 else (lambda v: v)
    return self.Tgk[k_from] @ a + col(self.dgk[k_from])

from_atlas(z, k_to)

Drop atlas coordinate(s) z into chart k_to: a = Tkg[k_to] z + dkg[k_to] (the exact inverse of to_atlas on chart k_to's affine subspace).

Source code in qlroms/transitions.py 
132
133
134
135
136
137
138
def from_atlas(self, z, k_to: int):
    """Drop atlas coordinate(s) z into chart k_to: a = Tkg[k_to] z + dkg[k_to]
    (the exact inverse of to_atlas on chart k_to's affine subspace)."""
    if not self.has_atlas:
        raise ValueError("No atlas maps (Tkg/dkg); build the global atlas first.")
    col = (lambda v: v[:, None]) if z.ndim == 2 else (lambda v: v)
    return self.Tkg[k_to] @ z + col(self.dkg[k_to])

map(a, k_from, k_to)

Map coordinate(s) a from chart k_from to chart k_to.

a: (r,) or (r, m), torch.Tensor or np.ndarray (the ops are type-generic).

Source code in qlroms/transitions.py 
140
141
142
143
144
145
146
147
148
def map(self, a, k_from: int, k_to: int):
    """Map coordinate(s) a from chart k_from to chart k_to.

    a: (r,) or (r, m), torch.Tensor or np.ndarray (the ops are type-generic).
    """
    if self._use_atlas():
        return self.from_atlas(self.to_atlas(a, k_from), k_to)
    col = (lambda v: v[:, None]) if a.ndim == 2 else (lambda v: v)
    return self.tmap[k_from, k_to] @ a + col(self.tshift[k_from, k_to])

distances(a, k_from)

Distance from state (a in chart k_from) to all cluster centroids.

Returns:

Type Description

(K,) vector with distances in chosen transition metric.
Source code in qlroms/transitions.py 
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
def distances(self, a, k_from: int):
    """Distance from state (a in chart k_from) to all cluster centroids.

    Returns:
        (K,) vector with distances in chosen transition metric.
    """
    if self.method == "physical":
        if self.centroids is None or self.Phi_all is None:
            raise ValueError("method='physical' needs centroids and Phi_all (see from_model).")
        a_t = a if torch.is_tensor(a) else torch.as_tensor(a, dtype=self.centroids.dtype,
                                                           device=self.centroids.device)
        x_hat = self.centroids[k_from] + self.Phi_all[:, :, k_from] @ a_t.reshape(-1)   # (N,)
        diff = self.centroids - x_hat[None, :]                                          # (K, N)
        kind = self._kind()
        if kind == "none":
            d = diff.norm(dim=1)
        elif kind in ("uniform", "diag"):
            wt = torch.as_tensor(as_diagonal(self.Mw), dtype=diff.dtype, device=diff.device)
            d = torch.sqrt((wt * diff**2).sum(dim=1))
        else:
            d = torch.sqrt(wsq_rows(self.Mw, diff).clamp_min(0))
        return d if torch.is_tensor(a) else d.cpu().numpy()
    use_atlas = self._use_atlas()
    if use_atlas:
        assert self.Tgk is not None and self.dgk is not None
        z = self.Tgk[k_from] @ a + self.dgk[k_from]
        return (self.dgk - z).norm(dim=1) if torch.is_tensor(a) else np.linalg.norm(self.dgk - z, axis=1)
    assert self.tshift is not None
    if self.oos2 is not None:
        if torch.is_tensor(a):
            return torch.sqrt(((self.tshift[:, k_from] - a) ** 2).sum(dim=1) + self.oos2[:, k_from])
        return np.sqrt(((self.tshift[:, k_from] - a) ** 2).sum(axis=1) + self.oos2[:, k_from])
    return (self.tshift[:, k_from] - a).norm(dim=1) if torch.is_tensor(a) else np.linalg.norm(self.tshift[:, k_from] - a, axis=1)

pairwise_oos2(centroids, tshift, Mw)

(K, K) squared M distance of centroid i to chart j's affine space that the chart-j coordinates cannot see: oos2[i, j] = ||c_i - c_j||_M^2 - ||tshift[i, j]||^2 (>= 0). tshift[i, j] = Phi_j^T Mw (c_i - c_j) is the in-chart part of c_i - c_j, so adding oos2 to ||tshift[i, j] - a||^2 gives the exact squared distance from a chart-j state to c_i. Computed once from the geometry; everything is centred on the mean centroid so the difference of large, nearly equal M products keeps its precision.

Source code in qlroms/transitions.py 
14
15
16
17
18
19
20
21
22
23
24
25
def pairwise_oos2(centroids: torch.Tensor, tshift: torch.Tensor, Mw) -> torch.Tensor:
    """(K, K) squared M distance of centroid i to chart j's affine space that the chart-j
    coordinates cannot see:  oos2[i, j] = ||c_i - c_j||_M^2 - ||tshift[i, j]||^2  (>= 0).
    tshift[i, j] = Phi_j^T Mw (c_i - c_j) is the in-chart part of c_i - c_j, so adding oos2
    to ||tshift[i, j] - a||^2 gives the exact squared distance from a chart-j state to c_i.
    Computed once from the geometry; everything is centred on the mean centroid so the
    difference of large, nearly equal M products keeps its precision."""
    Cc = centroids - centroids.mean(dim=0, keepdim=True)                      # (K, N)
    G = Cc @ apply_weight(as_torch_weight(Mw, Cc), Cc.T.contiguous())          # (K, K): c_i^T M c_j
    g = torch.diagonal(G)
    D2 = g[:, None] + g[None, :] - 2.0 * G
    return (D2 - (tshift ** 2).sum(dim=-1)).clamp_min(0.0)