Skip to content

qlroms.transitions

TransitionMaps dataclass

Transition operators with selectable mapping strategy.

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.

- "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 metrics approximate (O(N) per step; needs centroids/Phi_all, see from_model). Mass-weighted when wt is set: sqrt(sum(wt * (c - x)^2)), the metric a weighted (Mw-diagonal) atlas reproduces exactly; plain Euclidean when wt is None. The map itself falls back to atlas/pairwise (exact by construction, so the metric choice is the only thing "physical" changes).

Source code in qlroms/transitions.py
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 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
@dataclass
class TransitionMaps:
    """Transition operators with selectable mapping strategy.

    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.
        - "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 metrics approximate (O(N) per step; needs
            centroids/Phi_all, see from_model). Mass-weighted when wt is set:
            sqrt(sum(wt * (c - x)^2)), the metric a weighted (Mw-diagonal) atlas
            reproduces exactly; plain Euclidean when wt is None. The map itself
            falls back to atlas/pairwise (exact by construction, so the metric
            choice is the only thing "physical" changes).
    """

    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
    wt: Any | None = None          # (N,)-reshapeable spatial weight, "physical" distances only
    method: TRANSITION_METHODS = "auto"

    @classmethod
    def from_model(cls, model: Any, method: TRANSITION_METHODS = "auto") -> TransitionMaps:
        # Spatial weight for "physical" distances: ks2d-style models expose `wt` (the
        # quadrature-weight grid); Atlas/qlROM carry only the chart weight Mw, whose
        # 1-D (diagonal) form is the same vector. A full (N, N) mass matrix is not a
        # spatial weight vector -> None (plain Euclidean, previous behavior).
        wt = getattr(model, "wt", None)
        if wt is None:
            Mw = getattr(model, "Mw", None)
            if Mw is not None and getattr(Mw, "ndim", None) == 1:
                wt = Mw
        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),
            wt=wt,
            method=method,
        )

    @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)
            if self.wt is not None:
                wt = torch.as_tensor(self.wt, dtype=diff.dtype, device=diff.device).reshape(-1)
                d = torch.sqrt((wt * diff**2).sum(dim=1))
            else:
                d = diff.norm(dim=1)
            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
        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
 96
 97
 98
 99
100
101
102
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
104
105
106
107
108
109
110
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
112
113
114
115
116
117
118
119
120
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
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
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)
        if self.wt is not None:
            wt = torch.as_tensor(self.wt, dtype=diff.dtype, device=diff.device).reshape(-1)
            d = torch.sqrt((wt * diff**2).sum(dim=1))
        else:
            d = diff.norm(dim=1)
        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
    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)