qlroms.atlas

Atlas: the shared global coordinate system over a collection of local charts.

Implements sections 1/4 of docs/streamlined_vision_building_blocks.md: - build_global_assimilation_basis: the atlas basis Phi_g -- the Gram-Schmidt (eigendecomposition) of the union of all local chart bases and centroid offsets, so every local affine space is contained in it and the local<->atlas maps (Tgk/dgk/Tkg/dkg) are exact inverses (docs/shared_global.txt). - Atlas: the dynamics-free collection of charts, with pairwise transition operators and (optionally) the exact global atlas. Mw-aware: pass the chart inner-product weight (None = identity, (N,) diagonal, or a dense/sparse (N, N) mass matrix) and every projection/transition/atlas build uses it. - fit_charts: raw snapshots -> Atlas, the one-call offline geometry build.

Atlas

Dynamics-free chart provider exposing the surface qlESN/qlOpinf consumes: r, K, dt, device, rdtype, tree, tmap, tshift, _get_cluster_rom, recover_state. Tgk/dgk/Tkg/dkg (exact atlas switching, see build_global_atlas) are built by default from the chart geometry unless method='pairwise' is requested; g_mean_all/mu_bar_all (a separate, dynamics-related zero-mean refinement) are always None here -- qlESN treats all of these as optional.

Source code in qlroms/atlas.py

class Atlas:
    """Dynamics-free chart provider exposing the surface qlESN/qlOpinf consumes:
    r, K, dt, device, rdtype, tree, tmap, tshift, _get_cluster_rom, recover_state.
    Tgk/dgk/Tkg/dkg (exact atlas switching, see build_global_atlas) are built by
    default from the chart geometry unless method='pairwise' is requested;
    g_mean_all/mu_bar_all (a separate, dynamics-related zero-mean refinement) are always
    None here -- qlESN treats all of these as optional."""

    def __init__(self, centroids, Phi_all, dt: float = 1.0, Mw=None,
                 tmap=None, tshift=None, tree=None,
                 qbar_g=None, method: str = "auto"):
        """
        Args:
            centroids: (K, N) cluster centroids.
            Phi_all:   (N, r, K) local POD bases, Mw-orthonormal.
            dt:        snapshot spacing (metadata; used by fit timing only).
            Mw:        chart inner-product weight (see qlroms.charts.Chart): None for
                       M = I, a (N,) diagonal, or a dense/sparse (N, N) mass matrix.
                       Used by the per-chart projections, the pairwise transition
                       operators, and build_global_atlas.
            tmap/tshift/tree: prebuilt pairwise operators (from_rom); built here if None.
            qbar_g: optional atlas reference state. When omitted, the atlas is built
                about the mean of the stored chart centroids, which is sufficient for
                the exact affine-space inclusion/mapping identities. fit_charts later
                overwrites this with the training-snapshot mean when available.
            method: transition strategy. Default 'auto' builds the exact shared atlas
                immediately and uses it when available. Use 'pairwise' to keep only
                tmap/tshift switching.
        """
        self.Phi_all = Phi_all if torch.is_tensor(Phi_all) else torch.as_tensor(Phi_all)
        self.centroids = torch.as_tensor(centroids, dtype=self.Phi_all.dtype, device=self.Phi_all.device)
        self.device, self.rdtype = self.Phi_all.device, self.Phi_all.dtype
        self.dt = float(dt)
        self.Mw = Mw
        self.K = int(self.centroids.shape[0])
        self.r = int(self.Phi_all.shape[1])

        self.Phi_g = self.Tgk = self.dgk = self.Tkg = self.dkg = None
        self.qbar_g = None if qbar_g is None else torch.as_tensor(
            qbar_g, dtype=self.rdtype, device=self.device).reshape(-1)
        self.g_mean_all = self.mu_bar_all = None   # separate zero-mean refinement; not built here

        self.tree = tree if tree is not None else (
            KDTree(self.centroids.detach().cpu().numpy()) if self.K > 1 else None)
        if tmap is None or tshift is None:
            tmap, tshift = self._build_transition_operators()
        self.tmap, self.tshift = tmap, tshift
        self._charts: list[Chart | None] = [None] * self.K
        self.transitions = TransitionMaps.from_model(self, method="pairwise")
        self.method = method

    @property
    def method(self) -> str:
        return self.transitions.method

    @method.setter
    def method(self, value: str) -> None:
        method = _normalize_transition_method(value)
        if method != "pairwise" and self.K > 1 and not self.transitions.has_atlas:
            qbar_g = self.qbar_g if self.qbar_g is not None else self.centroids.mean(dim=0)
            self._attach_global_atlas(qbar_g)
        self.transitions = TransitionMaps.from_model(self, method=method)

    @property
    def N(self) -> int:
        """Physical state dimension (rows of Phi_all), as on the case qlROMs."""
        return int(self.Phi_all.shape[0])

    def _build_transition_operators(self):
        """Pairwise chart-i -> chart-j coordinate maps, for all K*K pairs (including i==j,
        the identity: tmap[i,i]=I, tshift[i,i]=0), under the inner-product weight Mw
        (identity when None):

            a_j = Phi_j^T Mw [ (Phi_i a_i + c_i) - c_j ]               (physical round trip)
                = (Phi_j^T Mw Phi_i) a_i + Phi_j^T Mw (c_i - c_j)
                = tmap[i,j] @ a_i       + tshift[i,j]                  (r x r, r -- small)

        The N-dim Phi_i/Phi_j/centroids are only touched HERE, once, offline, to fold the
        physical round trip into two small cached tensors. At runtime applying
        tmap/tshift is a pure (r,r)@(r,) op -- no physical-space vector is ever
        materialized during a run.

        Exactness: a_j is the EXACT chart-j projection of chart i's OWN reconstruction
        (Phi_i a_i + c_i), so it equals the true chart-j coordinate of the underlying
        physical state only up to chart i's reconstruction (truncation) error -- verified
        empirically in qlesn_LL_washout_checks.ipynb (invariant 1a vs 1b).

        Returns:
            tmap: (K, K, r, r), tshift: (K, K, r).
        """
        K, r = self.K, self.r
        tmap = torch.zeros(K, K, r, r, dtype=self.rdtype, device=self.device)
        tshift = torch.zeros(K, K, r, dtype=self.rdtype, device=self.device)
        MPhi = [apply_weight(self.Mw, self.Phi_all[:, :, j]) for j in range(K)]   # Mw Phi_j
        for i in range(K):
            Phi_i = self.Phi_all[:, :, i]
            for j in range(K):
                tmap[i, j] = MPhi[j].T @ Phi_i
                tshift[i, j] = MPhi[j].T @ (self.centroids[i] - self.centroids[j])
        return tmap, tshift

    def _attach_global_atlas(
        self,
        qbar_g,
        tol: float = 1e-11,
        assume_tgk_transpose: bool = True,
    ) -> "Atlas":
        qbar_t = torch.as_tensor(qbar_g, dtype=self.rdtype, device=self.device).reshape(-1)
        self.qbar_g = qbar_t
        self.Phi_g, self.Tgk, self.dgk, self.Tkg, self.dkg = build_global_assimilation_basis(
            self.Phi_all,
            self.centroids,
            self.qbar_g,
            Mw=self.Mw,
            tol=tol,
            assume_tgk_transpose=assume_tgk_transpose,
        )
        return self

    def build_global_atlas(
        self,
        X,
        tol: float = 1e-11,
        assume_tgk_transpose: bool = True,
    ) -> "Atlas":
        """Attach the exact global assimilation basis (Phi_g/Tgk/dgk), so
        qlroms/esn_base.rom_distance switches charts on EXACT physical distances instead
        of falling back to the distorted in-chart tshift rule (on a real run: 100% vs
        83% correct distance ordering). Uses the atlas' inner-product weight Mw.

        Args:
            X: either a (N, Nt) snapshot matrix whose mean defines qbar_g -- typically
                the same matrix fit_charts was called on -- or an explicit (N,) atlas
                reference state.
            tol: relative eigenvalue cutoff dropping redundant directions (see
                build_global_assimilation_basis).
            assume_tgk_transpose: pass through to build_global_assimilation_basis;
                defaults to True so Tgk is enforced as Tkg^T at build time.
        Returns:
            self, for chaining (fit_charts(...).build_global_atlas(X) reads naturally,
            though fit_charts already calls this by default).
        """
        if self.K <= 1:
            return self          # no other chart to be exact ABOUT; nothing to build
        Xt = torch.as_tensor(X, dtype=self.rdtype, device=self.device)
        qbar_g = Xt if Xt.ndim == 1 else Xt.mean(dim=1)
        self._attach_global_atlas(
            qbar_g,
            tol=tol,
            assume_tgk_transpose=assume_tgk_transpose,
        )
        self.transitions = TransitionMaps.from_model(self, method=_normalize_transition_method(self.method))
        return self

    def _get_cluster_rom(self, k: int) -> Chart:
        chart = self._charts[k]
        if chart is None:
            chart = Chart(self.Phi_all[:, :, k], self.centroids[k], self.r, Mw=self.Mw)
            self._charts[k] = chart
        return chart

    def project_state(self, state, cluster_id=None) -> torch.Tensor:
        """(N,) or (N, m) physical -> (r+1, m) augmented coords, nearest-chart per
        column with the chart id in the last row -- the inverse of recover_state and
        the same convention as the case qlROMs, so ensemble/workflow code can
        run on a bare Atlas."""
        s = torch.as_tensor(state, dtype=self.rdtype, device=self.device)
        if s.ndim == 1:
            s = s[:, None]     # (r+1, 1) out, matching the case qlROMs (no squeeze back)
        if cluster_id is not None:
            cids = torch.full((s.shape[1],), int(cluster_id), dtype=torch.long, device=self.device)
        elif self.K == 1:
            cids = torch.zeros(s.shape[1], dtype=torch.long, device=self.device)
        else:
            tree = self.tree
            assert tree is not None
            _, labels = tree.query(s.detach().cpu().numpy().T, k=1)
            cids = torch.as_tensor(labels, device=self.device).reshape(-1).long()
        a = torch.empty(self.r, s.shape[1], dtype=self.rdtype, device=self.device)
        for k in cids.unique():
            mask = cids == k
            a[:, mask] = self._get_cluster_rom(int(k)).project_state(s[:, mask])
        return torch.cat([a, cids[None, :].to(self.rdtype)], dim=0)

    def recover_state(self, apod) -> torch.Tensor:
        """Augmented (r+1, m) coords (cluster id in the last row) -> (N, m) physical,
        matching qlROM.recover_state's per-column id dispatch (what qlESN.recover_state
        delegates to)."""
        a = torch.as_tensor(apod, dtype=self.rdtype, device=self.device)
        squeezed = a.ndim == 1
        if squeezed:
            a = a[:, None]
        if a.shape[0] == self.r + 1:                 # augmented: ids in last row
            cids, a_pure = a[-1].long(), a[:-1]
        else:                                        # plain (r, m): single chart 0
            cids, a_pure = torch.zeros(a.shape[1], dtype=torch.long, device=self.device), a
        N = self.Phi_all.shape[0]
        out = torch.empty(N, a_pure.shape[1], dtype=self.rdtype, device=self.device)
        for k in cids.unique():
            mask = cids == k
            out[:, mask] = self._get_cluster_rom(int(k)).recover_state(a_pure[:, mask])
        return out.squeeze(-1) if squeezed else out

    @classmethod
    def from_rom(cls, rom) -> "Atlas":
        """Adapt an existing qlROM's charts (exact copy of its centroids/Phi/tmap/tshift,
        preserving whatever inner-product weight it was built with)."""
        method = getattr(getattr(rom, "transitions", None), "method", "auto")
        has_atlas = all(
            getattr(rom, name, None) is not None
            for name in ("Phi_g", "qbar_g", "Tgk", "dgk", "Tkg", "dkg")
        )
        atlas = cls(rom.centroids, rom.Phi_all, dt=getattr(rom, "dt", 1.0),
                    Mw=getattr(rom, "Mw", None),
                    tmap=rom.tmap, tshift=rom.tshift, tree=getattr(rom, "tree", None),
                    qbar_g=getattr(rom, "qbar_g", None),
                    method="pairwise" if has_atlas else method)
        if has_atlas:
            atlas.Phi_g = getattr(rom, "Phi_g")
            atlas.qbar_g = getattr(rom, "qbar_g")
            atlas.Tgk = getattr(rom, "Tgk")
            atlas.dgk = getattr(rom, "dgk")
            atlas.Tkg = getattr(rom, "Tkg")
            atlas.dkg = getattr(rom, "dkg")
            atlas.transitions = TransitionMaps.from_model(atlas, method=_normalize_transition_method(method))
        atlas.g_mean_all = getattr(rom, "g_mean_all", None)
        atlas.mu_bar_all = getattr(rom, "mu_bar_all", None)
        return atlas

    @classmethod
    def from_charts(cls, charts, dt: float = 1.0, *, qbar_g=None, method: TRANSITION_METHODS = "auto") -> "Atlas":
        """
        Assemble an Atlas from K already-built Chart objects (e.g. Chart.from_fenics
        conversions of a dolfinx POD stack): stacks their Phi/centroids, takes the
        inner-product weight from the first chart (all charts of one atlas share it),
        and by default also builds the exact shared atlas from those chart objects.
        Pass method='pairwise' to keep only pairwise switching.
        """
        Phi_all = torch.stack([c.Phi for c in charts], dim=2)          # (N, r, K)
        centroids = torch.stack([c.centroid.reshape(-1) for c in charts])  # (K, N)
        return cls(centroids,
                   Phi_all,
                   dt=dt, Mw=charts[0].Mw, qbar_g=qbar_g, method=method)

    def save(self, path, **extra):
        """One compressed npz holding the atlas core (centroids, Phi_all, dt, qbar_g)
        plus any caller metadata arrays. POD bases are nested (the leading r columns
        of a rank-R basis are the rank-r basis for the same partition), so one save
        at large R serves every smaller r through ``Atlas.load(path, r=r)``."""
        import numpy as np
        qbar_g = self.qbar_g if self.qbar_g is not None else self.centroids.mean(dim=0)
        np.savez_compressed(
            path, centroids=self.centroids.cpu().numpy(),
            Phi_all=self.Phi_all.cpu().numpy(), dt=self.dt,
            qbar_g=(qbar_g.cpu().numpy() if torch.is_tensor(qbar_g) else qbar_g), **extra)

    @classmethod
    def load(cls, source, r: int | None = None, **kwargs) -> "Atlas":
        """Rebuild an Atlas from ``save()``'s npz -- a path, or an already-open
        mapping of its arrays (e.g. a dict of memory-mapped .npy files). ``r``
        truncates every chart basis to its leading r modes (nested POD); transition
        operators and the shared atlas are rebuilt by the constructor."""
        import os

        import numpy as np
        z = (source if hasattr(source, "__getitem__")
             and not isinstance(source, (str, bytes, os.PathLike))
             else np.load(source, allow_pickle=False))
        Phi = np.asarray(z["Phi_all"])
        if r is not None:
            if r > Phi.shape[1]:
                raise ValueError(f"r={r} exceeds the saved basis rank {Phi.shape[1]}")
            Phi = Phi[:, :r, :]
        return cls(np.asarray(z["centroids"]), Phi, dt=float(z["dt"]),
                   qbar_g=np.asarray(z["qbar_g"]), **kwargs)

N property

Physical state dimension (rows of Phi_all), as on the case qlROMs.

__init__(centroids, Phi_all, dt=1.0, Mw=None, tmap=None, tshift=None, tree=None, qbar_g=None, method='auto')

Args: centroids: (K, N) cluster centroids. Phi_all: (N, r, K) local POD bases, Mw-orthonormal. dt: snapshot spacing (metadata; used by fit timing only). Mw: chart inner-product weight (see qlroms.charts.Chart): None for M = I, a (N,) diagonal, or a dense/sparse (N, N) mass matrix. Used by the per-chart projections, the pairwise transition operators, and build_global_atlas. tmap/tshift/tree: prebuilt pairwise operators (from_rom); built here if None. qbar_g: optional atlas reference state. When omitted, the atlas is built about the mean of the stored chart centroids, which is sufficient for the exact affine-space inclusion/mapping identities. fit_charts later overwrites this with the training-snapshot mean when available. method: transition strategy. Default 'auto' builds the exact shared atlas immediately and uses it when available. Use 'pairwise' to keep only tmap/tshift switching.

Source code in qlroms/atlas.py

def __init__(self, centroids, Phi_all, dt: float = 1.0, Mw=None,
             tmap=None, tshift=None, tree=None,
             qbar_g=None, method: str = "auto"):
    """
    Args:
        centroids: (K, N) cluster centroids.
        Phi_all:   (N, r, K) local POD bases, Mw-orthonormal.
        dt:        snapshot spacing (metadata; used by fit timing only).
        Mw:        chart inner-product weight (see qlroms.charts.Chart): None for
                   M = I, a (N,) diagonal, or a dense/sparse (N, N) mass matrix.
                   Used by the per-chart projections, the pairwise transition
                   operators, and build_global_atlas.
        tmap/tshift/tree: prebuilt pairwise operators (from_rom); built here if None.
        qbar_g: optional atlas reference state. When omitted, the atlas is built
            about the mean of the stored chart centroids, which is sufficient for
            the exact affine-space inclusion/mapping identities. fit_charts later
            overwrites this with the training-snapshot mean when available.
        method: transition strategy. Default 'auto' builds the exact shared atlas
            immediately and uses it when available. Use 'pairwise' to keep only
            tmap/tshift switching.
    """
    self.Phi_all = Phi_all if torch.is_tensor(Phi_all) else torch.as_tensor(Phi_all)
    self.centroids = torch.as_tensor(centroids, dtype=self.Phi_all.dtype, device=self.Phi_all.device)
    self.device, self.rdtype = self.Phi_all.device, self.Phi_all.dtype
    self.dt = float(dt)
    self.Mw = Mw
    self.K = int(self.centroids.shape[0])
    self.r = int(self.Phi_all.shape[1])

    self.Phi_g = self.Tgk = self.dgk = self.Tkg = self.dkg = None
    self.qbar_g = None if qbar_g is None else torch.as_tensor(
        qbar_g, dtype=self.rdtype, device=self.device).reshape(-1)
    self.g_mean_all = self.mu_bar_all = None   # separate zero-mean refinement; not built here

    self.tree = tree if tree is not None else (
        KDTree(self.centroids.detach().cpu().numpy()) if self.K > 1 else None)
    if tmap is None or tshift is None:
        tmap, tshift = self._build_transition_operators()
    self.tmap, self.tshift = tmap, tshift
    self._charts: list[Chart | None] = [None] * self.K
    self.transitions = TransitionMaps.from_model(self, method="pairwise")
    self.method = method

build_global_atlas(X, tol=1e-11, assume_tgk_transpose=True)

Attach the exact global assimilation basis (Phi_g/Tgk/dgk), so qlroms/esn_base.rom_distance switches charts on EXACT physical distances instead of falling back to the distorted in-chart tshift rule (on a real run: 100% vs 83% correct distance ordering). Uses the atlas' inner-product weight Mw.

Parameters:

Name	Type	Description	Default
X		either a (N, Nt) snapshot matrix whose mean defines qbar_g -- typically the same matrix fit_charts was called on -- or an explicit (N,) atlas reference state.	required
tol	float	relative eigenvalue cutoff dropping redundant directions (see build_global_assimilation_basis).	1e-11
assume_tgk_transpose	bool	pass through to build_global_assimilation_basis; defaults to True so Tgk is enforced as Tkg^T at build time.	True

Returns: self, for chaining (fit_charts(...).build_global_atlas(X) reads naturally, though fit_charts already calls this by default).

Source code in qlroms/atlas.py

def build_global_atlas(
    self,
    X,
    tol: float = 1e-11,
    assume_tgk_transpose: bool = True,
) -> "Atlas":
    """Attach the exact global assimilation basis (Phi_g/Tgk/dgk), so
    qlroms/esn_base.rom_distance switches charts on EXACT physical distances instead
    of falling back to the distorted in-chart tshift rule (on a real run: 100% vs
    83% correct distance ordering). Uses the atlas' inner-product weight Mw.

    Args:
        X: either a (N, Nt) snapshot matrix whose mean defines qbar_g -- typically
            the same matrix fit_charts was called on -- or an explicit (N,) atlas
            reference state.
        tol: relative eigenvalue cutoff dropping redundant directions (see
            build_global_assimilation_basis).
        assume_tgk_transpose: pass through to build_global_assimilation_basis;
            defaults to True so Tgk is enforced as Tkg^T at build time.
    Returns:
        self, for chaining (fit_charts(...).build_global_atlas(X) reads naturally,
        though fit_charts already calls this by default).
    """
    if self.K <= 1:
        return self          # no other chart to be exact ABOUT; nothing to build
    Xt = torch.as_tensor(X, dtype=self.rdtype, device=self.device)
    qbar_g = Xt if Xt.ndim == 1 else Xt.mean(dim=1)
    self._attach_global_atlas(
        qbar_g,
        tol=tol,
        assume_tgk_transpose=assume_tgk_transpose,
    )
    self.transitions = TransitionMaps.from_model(self, method=_normalize_transition_method(self.method))
    return self

project_state(state, cluster_id=None)

(N,) or (N, m) physical -> (r+1, m) augmented coords, nearest-chart per column with the chart id in the last row -- the inverse of recover_state and the same convention as the case qlROMs, so ensemble/workflow code can run on a bare Atlas.

Source code in qlroms/atlas.py

def project_state(self, state, cluster_id=None) -> torch.Tensor:
    """(N,) or (N, m) physical -> (r+1, m) augmented coords, nearest-chart per
    column with the chart id in the last row -- the inverse of recover_state and
    the same convention as the case qlROMs, so ensemble/workflow code can
    run on a bare Atlas."""
    s = torch.as_tensor(state, dtype=self.rdtype, device=self.device)
    if s.ndim == 1:
        s = s[:, None]     # (r+1, 1) out, matching the case qlROMs (no squeeze back)
    if cluster_id is not None:
        cids = torch.full((s.shape[1],), int(cluster_id), dtype=torch.long, device=self.device)
    elif self.K == 1:
        cids = torch.zeros(s.shape[1], dtype=torch.long, device=self.device)
    else:
        tree = self.tree
        assert tree is not None
        _, labels = tree.query(s.detach().cpu().numpy().T, k=1)
        cids = torch.as_tensor(labels, device=self.device).reshape(-1).long()
    a = torch.empty(self.r, s.shape[1], dtype=self.rdtype, device=self.device)
    for k in cids.unique():
        mask = cids == k
        a[:, mask] = self._get_cluster_rom(int(k)).project_state(s[:, mask])
    return torch.cat([a, cids[None, :].to(self.rdtype)], dim=0)

recover_state(apod)

Augmented (r+1, m) coords (cluster id in the last row) -> (N, m) physical, matching qlROM.recover_state's per-column id dispatch (what qlESN.recover_state delegates to).

Source code in qlroms/atlas.py

def recover_state(self, apod) -> torch.Tensor:
    """Augmented (r+1, m) coords (cluster id in the last row) -> (N, m) physical,
    matching qlROM.recover_state's per-column id dispatch (what qlESN.recover_state
    delegates to)."""
    a = torch.as_tensor(apod, dtype=self.rdtype, device=self.device)
    squeezed = a.ndim == 1
    if squeezed:
        a = a[:, None]
    if a.shape[0] == self.r + 1:                 # augmented: ids in last row
        cids, a_pure = a[-1].long(), a[:-1]
    else:                                        # plain (r, m): single chart 0
        cids, a_pure = torch.zeros(a.shape[1], dtype=torch.long, device=self.device), a
    N = self.Phi_all.shape[0]
    out = torch.empty(N, a_pure.shape[1], dtype=self.rdtype, device=self.device)
    for k in cids.unique():
        mask = cids == k
        out[:, mask] = self._get_cluster_rom(int(k)).recover_state(a_pure[:, mask])
    return out.squeeze(-1) if squeezed else out

from_rom(rom) classmethod

Adapt an existing qlROM's charts (exact copy of its centroids/Phi/tmap/tshift, preserving whatever inner-product weight it was built with).

Source code in qlroms/atlas.py

@classmethod
def from_rom(cls, rom) -> "Atlas":
    """Adapt an existing qlROM's charts (exact copy of its centroids/Phi/tmap/tshift,
    preserving whatever inner-product weight it was built with)."""
    method = getattr(getattr(rom, "transitions", None), "method", "auto")
    has_atlas = all(
        getattr(rom, name, None) is not None
        for name in ("Phi_g", "qbar_g", "Tgk", "dgk", "Tkg", "dkg")
    )
    atlas = cls(rom.centroids, rom.Phi_all, dt=getattr(rom, "dt", 1.0),
                Mw=getattr(rom, "Mw", None),
                tmap=rom.tmap, tshift=rom.tshift, tree=getattr(rom, "tree", None),
                qbar_g=getattr(rom, "qbar_g", None),
                method="pairwise" if has_atlas else method)
    if has_atlas:
        atlas.Phi_g = getattr(rom, "Phi_g")
        atlas.qbar_g = getattr(rom, "qbar_g")
        atlas.Tgk = getattr(rom, "Tgk")
        atlas.dgk = getattr(rom, "dgk")
        atlas.Tkg = getattr(rom, "Tkg")
        atlas.dkg = getattr(rom, "dkg")
        atlas.transitions = TransitionMaps.from_model(atlas, method=_normalize_transition_method(method))
    atlas.g_mean_all = getattr(rom, "g_mean_all", None)
    atlas.mu_bar_all = getattr(rom, "mu_bar_all", None)
    return atlas

from_charts(charts, dt=1.0, *, qbar_g=None, method='auto') classmethod

Assemble an Atlas from K already-built Chart objects (e.g. Chart.from_fenics conversions of a dolfinx POD stack): stacks their Phi/centroids, takes the inner-product weight from the first chart (all charts of one atlas share it), and by default also builds the exact shared atlas from those chart objects. Pass method='pairwise' to keep only pairwise switching.

Source code in qlroms/atlas.py

@classmethod
def from_charts(cls, charts, dt: float = 1.0, *, qbar_g=None, method: TRANSITION_METHODS = "auto") -> "Atlas":
    """
    Assemble an Atlas from K already-built Chart objects (e.g. Chart.from_fenics
    conversions of a dolfinx POD stack): stacks their Phi/centroids, takes the
    inner-product weight from the first chart (all charts of one atlas share it),
    and by default also builds the exact shared atlas from those chart objects.
    Pass method='pairwise' to keep only pairwise switching.
    """
    Phi_all = torch.stack([c.Phi for c in charts], dim=2)          # (N, r, K)
    centroids = torch.stack([c.centroid.reshape(-1) for c in charts])  # (K, N)
    return cls(centroids,
               Phi_all,
               dt=dt, Mw=charts[0].Mw, qbar_g=qbar_g, method=method)

save(path, **extra)

One compressed npz holding the atlas core (centroids, Phi_all, dt, qbar_g) plus any caller metadata arrays. POD bases are nested (the leading r columns of a rank-R basis are the rank-r basis for the same partition), so one save at large R serves every smaller r through Atlas.load(path, r=r).

Source code in qlroms/atlas.py

def save(self, path, **extra):
    """One compressed npz holding the atlas core (centroids, Phi_all, dt, qbar_g)
    plus any caller metadata arrays. POD bases are nested (the leading r columns
    of a rank-R basis are the rank-r basis for the same partition), so one save
    at large R serves every smaller r through ``Atlas.load(path, r=r)``."""
    import numpy as np
    qbar_g = self.qbar_g if self.qbar_g is not None else self.centroids.mean(dim=0)
    np.savez_compressed(
        path, centroids=self.centroids.cpu().numpy(),
        Phi_all=self.Phi_all.cpu().numpy(), dt=self.dt,
        qbar_g=(qbar_g.cpu().numpy() if torch.is_tensor(qbar_g) else qbar_g), **extra)

load(source, r=None, **kwargs) classmethod

Rebuild an Atlas from save()'s npz -- a path, or an already-open mapping of its arrays (e.g. a dict of memory-mapped .npy files). r truncates every chart basis to its leading r modes (nested POD); transition operators and the shared atlas are rebuilt by the constructor.

Source code in qlroms/atlas.py

@classmethod
def load(cls, source, r: int | None = None, **kwargs) -> "Atlas":
    """Rebuild an Atlas from ``save()``'s npz -- a path, or an already-open
    mapping of its arrays (e.g. a dict of memory-mapped .npy files). ``r``
    truncates every chart basis to its leading r modes (nested POD); transition
    operators and the shared atlas are rebuilt by the constructor."""
    import os

    import numpy as np
    z = (source if hasattr(source, "__getitem__")
         and not isinstance(source, (str, bytes, os.PathLike))
         else np.load(source, allow_pickle=False))
    Phi = np.asarray(z["Phi_all"])
    if r is not None:
        if r > Phi.shape[1]:
            raise ValueError(f"r={r} exceeds the saved basis rank {Phi.shape[1]}")
        Phi = Phi[:, :r, :]
    return cls(np.asarray(z["centroids"]), Phi, dt=float(z["dt"]),
               qbar_g=np.asarray(z["qbar_g"]), **kwargs)

build_global_assimilation_basis(Phi_all, centroids, qbar_g, Mw=None, tol=1e-11, assume_tgk_transpose=True)

Global M-orthonormal assimilation basis and exact local<->global maps.

Implements docs/shared_global.txt: Phi_g = build_atlas(U) U = [Phi_1..Phi_K, d_1..d_K], where d_k=(c_k-qbar_g) so every local affine space is contained in the global one and the maps below are exact inverses (up to the numerical rank of U).

Parameters:

Name	Type	Description	Default
Phi_all		(Nh, r, K) local POD bases.	required
centroids		(K, Nh) state-space centroids.	required
qbar_g		(Nh,) global reference (snapshot mean).	required
Mw		chart inner-product weight: None for M = I, a (Nh,) mass-matrix diagonal, or a dense/sparse (Nh, Nh) matrix (e.g. an assembled FEM mass matrix) -- applied through qlroms.charts.apply_weight.	None
tol		relative eigenvalue cutoff dropping redundant directions of U. Default: 1e-11 NB: At large K*r (e.g. K=40, r=50 -> U has 2040 columns), dividing by their near-zero eigenvalues below (Vr/sqrt(Lr)) pushes the exact-atlas residual (assert below) over 1e-6. Lowering tol keeps more of the real structure the offsets need to reconstruct exactly; verified across seeds that 1e-11 stays well short of where Phi_g's own M-orthonormality (Phi_g.T @ M @ Phi_g == I) starts degrading (1e-12 and below).	1e-11
assume_tgk_transpose	bool	if True (default), enforce Tgk[k] = Tkg[k].T at build time. dgk/dkg are still computed independently and are NOT tied by this flag.	True

Returns: Phi_g (Nh, rg), Tgk (K, rg, r), dgk (K, rg), Tkg (K, r, rg), dkg(K, r), where z = Tgk[k] @ a_k + dgk[k]; #local->atlas a_k = Tkg[k] @ z + dkg[k]; #atlas->local

Source code in qlroms/atlas.py

def build_global_assimilation_basis(
    Phi_all,
    centroids,
    qbar_g,
    Mw=None,
    tol=1e-11,
    assume_tgk_transpose: bool = True,
):
    """Global M-orthonormal assimilation basis and exact local<->global maps.

    Implements docs/shared_global.txt: Phi_g = build_atlas(U)
    U = [Phi_1..Phi_K, d_1..d_K], where d_k=(c_k-qbar_g) so every local affine space is
    contained in the global one and the maps below are exact inverses (up to the
    numerical rank of U).

    Args:
        Phi_all:   (Nh, r, K) local POD bases.
        centroids: (K, Nh) state-space centroids.
        qbar_g:    (Nh,) global reference (snapshot mean).
        Mw:        chart inner-product weight: None for M = I, a (Nh,) mass-matrix
                   diagonal, or a dense/sparse (Nh, Nh) matrix (e.g. an assembled FEM
                   mass matrix) -- applied through qlroms.charts.apply_weight.
        tol:       relative eigenvalue cutoff dropping redundant directions of U.
                   Default: 1e-11
                        NB: At large K*r (e.g. K=40, r=50 -> U has 2040 columns),
                        dividing by their near-zero eigenvalues below (Vr/sqrt(Lr))
                        pushes the exact-atlas residual (assert below) over 1e-6.
                        Lowering tol keeps more of the real structure the offsets
                        need to reconstruct exactly; verified across seeds that
                        1e-11 stays well short of where Phi_g's own M-orthonormality
                        (Phi_g.T @ M @ Phi_g == I) starts degrading (1e-12 and below).
        assume_tgk_transpose: if True (default), enforce Tgk[k] = Tkg[k].T at
               build time. dgk/dkg are still computed independently and are
               NOT tied by this flag.
    Returns:
        Phi_g (Nh, rg),
        Tgk (K, rg, r),
        dgk (K, rg),
        Tkg (K, r, rg),
        dkg(K, r),
            where
                z = Tgk[k] @ a_k + dgk[k];  #local->atlas
                a_k = Tkg[k] @ z + dkg[k];  #atlas->local
    """
    Nh, r, K = Phi_all.shape
    if Mw is not None and not torch.is_tensor(Mw):
        Mw = torch.as_tensor(Mw, dtype=Phi_all.dtype, device=Phi_all.device)
    # snapshots often live on the CPU while the bases sit on the model's device
    centroids = torch.as_tensor(centroids, dtype=Phi_all.dtype, device=Phi_all.device)
    qbar_g = torch.as_tensor(qbar_g, dtype=Phi_all.dtype, device=Phi_all.device)

    offsets = (centroids - qbar_g).T                                  # (Nh, K)
    U = torch.cat([Phi_all[:, :, k] for k in range(K)] +
                  [offsets], dim=1)                                   # (Nh, K*r+K)

    # Build the global assimilation basis Phi_g. This is the Gram-Schmidt
    # (eigendecomposition) of the union of all local chart bases and centroid offsets.
    # Depending on the weight Mw, we can compute the M-orthonormal basis using different methods:
    if Mw is None:
        # unweighted: use standard QR
        Phi_g = _rank_truncated_qr_basis(U, tol)
    elif Mw.ndim == 1 and torch.all(Mw > 0):
        # weighted diagonal: use sqrt(Mw) trick to reduce to unweighted QR
        sqrtw = torch.sqrt(Mw)
        Y = sqrtw[:, None] * U
        Qw = _rank_truncated_qr_basis(Y, tol)
        Phi_g = Qw / sqrtw[:, None]
    else:
        # else: manual QR with M-inner product
        G = U.T @ apply_weight(Mw, U)
        evals, evecs = torch.linalg.eigh(G)
        keep = evals > tol * evals[-1]
        Vr, Lr = evecs[:, keep], evals[keep]
        Phi_g = U @ (Vr / torch.sqrt(Lr))

    # Phi_g shape: (Nh, rg), where rg is the dimension of the shared coordinates used
    # for the assimilation basis. Now compute the exact local<->global maps.
    rg = Phi_g.shape[1]

    PhigW = apply_weight(Mw, Phi_g)                                  # (Nh, rg)
    Tgk = torch.zeros(K, rg, r, dtype=Phi_g.dtype, device=Phi_g.device)
    dgk = torch.zeros(K, rg,   dtype=Phi_g.dtype, device=Phi_g.device)
    Tkg = torch.zeros(K, r, rg, dtype=Phi_g.dtype, device=Phi_g.device)
    dkg = torch.zeros(K, r,     dtype=Phi_g.dtype, device=Phi_g.device)
    for k in range(K):
        Phik = Phi_all[:, :, k]
        PhikW = apply_weight(Mw, Phik)
        dgk[k] = PhigW.T @ (centroids[k] - qbar_g)
        Tkg[k] = Phik.T @ PhigW
        dkg[k] = PhikW.T @ (qbar_g - centroids[k])
        if assume_tgk_transpose:
            Tgk[k] = Tkg[k].T
        else:
            Tgk[k] = PhigW.T @ Phik

    # Verify the centroids are on the atlas: c_k - qbar_g must be reproduced exactly by
    # Phi_g @ dgk[k] (the offsets are columns of U, up to the tol truncation). Transition
    # distances computed in atlas coordinates are the true physical distances only then.
    resid = (centroids - qbar_g) - dgk @ Phi_g.T                      # (K, Nh)
    rel = resid.norm(dim=1) / (centroids - qbar_g).norm(dim=1).clamp_min(1e-30)
    assert float(rel.max()) < 1e-6, \
        f"centroids off the atlas (max rel err {float(rel.max()):.2e}); lower tol."
    return Phi_g, Tgk, dgk, Tkg, dkg

fit_charts(X, K, r, dt=1.0, *, random_state=1, pod_method='exact', pod_kwargs=None, clustering_kwargs=None, cluster_space='physical', Mw=None, build_atlas=True, assume_tgk_transpose=True, time_index=None, include_transition_snapshots=False, device=None, rdtype=torch.float64)

Build an Atlas from raw snapshots: KMeans clusters + per-cluster POD. No dynamics fitted -- this is all a qlESN needs.

Parameters:

Name	Type	Description	Default
X		(N, Ntrain) snapshot matrix, any source (chronological order not required for charts, but qlESN.fit later needs it for dwell segments).	required
K, r		number of clusters and POD modes per cluster.	required
dt	float	snapshot spacing (metadata).	1.0
clustering_kwargs	dict | None	forwarded to fit_clusters (random_state, kmeans_method, assign_overlapping, overlap_tolerance, ...).	None
pod_kwargs	dict | None	extra compute_pod_basis knobs for the randomized methods (oversampling, n_iter, block_size, random_state).	None
cluster_space	CLUSTER_SPACES	which representation of the snapshots k-means clusters on -- "physical" (default), "pod_lossless", "pod_r", or a callable; see clustering_features. For non-physical spaces the centroids are recomputed as physical-space means of the assigned columns afterwards.	'physical'
Mw		chart inner-product weight (None = identity). When given, the per-cluster POD is Mw-weighted (method of snapshots, Mw-orthonormal modes) and the Atlas carries Mw through projections/transitions/global atlas. Clustering stays Euclidean (labels only).	None
build_atlas	bool	also fit the global assimilation basis (Tgk/dgk) so qlESN's chart switching uses exact physical distances (Atlas.build_global_atlas). Default True; set False to skip the extra eigendecomposition (O((Kr+K)^2) or so) when K*r is large and only approximate switching is needed.	True
assume_tgk_transpose	bool	when building the atlas, enforce Tgk = Tkg^T (default True). dgk/dkg remain independently computed.	True
time_index		optional (Ntrain,) monotone integer index, one entry per X column, marking genuine time adjacency (gap > 1 = not consecutive -- e.g. multi-run seams from split.data.concat_wake_runs). Only used by include_transition_snapshots; None assumes one contiguous run.	None
include_transition_snapshots	bool	if True (needs X in chronological order), each cluster k's POD basis is additionally fit on the immediate POST-SWITCH snapshots (the genuinely-next-in-time snapshot of each member that leaves cluster k, which itself belongs to a different cluster). Default False (previous behaviour).	False

Source code in qlroms/atlas.py

def fit_charts(X, K: int, r: int, dt: float = 1.0, *, random_state: int = 1,
               pod_method: POD_METHODS = "exact",
               pod_kwargs: dict | None = None,
               clustering_kwargs: dict | None = None,
               cluster_space: CLUSTER_SPACES = "physical",
               Mw=None,
               build_atlas: bool = True,
               assume_tgk_transpose: bool = True,
               time_index=None, include_transition_snapshots: bool = False,
               device=None, rdtype=torch.float64) -> Atlas:
    """Build an Atlas from raw snapshots: KMeans clusters + per-cluster POD.
    No dynamics fitted -- this is all a qlESN needs.

    Args:
        X: (N, Ntrain) snapshot matrix, any source (chronological order not required
            for charts, but qlESN.fit later needs it for dwell segments).
        K, r: number of clusters and POD modes per cluster.
        dt: snapshot spacing (metadata).
        clustering_kwargs: forwarded to fit_clusters (random_state, kmeans_method,
            assign_overlapping, overlap_tolerance, ...).
        pod_kwargs: extra compute_pod_basis knobs for the randomized methods
            (oversampling, n_iter, block_size, random_state).
        cluster_space: which representation of the snapshots k-means clusters on --
            "physical" (default), "pod_lossless", "pod_r", or a callable; see
            clustering_features. For non-physical spaces the centroids are recomputed
            as physical-space means of the assigned columns afterwards.
        Mw: chart inner-product weight (None = identity). When given, the per-cluster
            POD is Mw-weighted (method of snapshots, Mw-orthonormal modes) and the
            Atlas carries Mw through projections/transitions/global atlas. Clustering
            stays Euclidean (labels only).
        build_atlas: also fit the global assimilation basis (Tgk/dgk) so qlESN's chart
            switching uses exact physical distances (Atlas.build_global_atlas).
            Default True; set False to skip
            the extra eigendecomposition (O((Kr+K)^2) or so) when K*r is large and only
            approximate switching is needed.
        assume_tgk_transpose: when building the atlas, enforce Tgk = Tkg^T
            (default True). dgk/dkg remain independently computed.
        time_index: optional (Ntrain,) monotone integer index, one entry per X
            column, marking genuine time adjacency (gap > 1 = not consecutive --
            e.g. multi-run seams from split.data.concat_wake_runs). Only used by
            include_transition_snapshots; None assumes one contiguous run.
        include_transition_snapshots: if True (needs X in chronological order),
            each cluster k's POD basis is additionally fit on the immediate
            POST-SWITCH snapshots (the genuinely-next-in-time snapshot of each
            member that leaves cluster k, which itself belongs to a different
            cluster). Default False (previous behaviour).
    """
    Xt = X if torch.is_tensor(X) else torch.as_tensor(X, dtype=rdtype)
    if device is not None:
        Xt = Xt.to(device)
    Xt = Xt.to(rdtype)
    N = Xt.shape[0]
    ckw = dict(random_state=random_state)
    ckw.update(clustering_kwargs or {})

    pk = pod_kwargs or {}
    if K == 1:
        centroids = Xt.mean(dim=1, keepdim=True).T            # (1, N)
        labels = torch.zeros(Xt.shape[1], dtype=torch.long)
        Xpod = Xt
    else:
        feats = clustering_features(Xt, r=r, cluster_space=cluster_space, pod_method=pod_method,
                                    random_state=random_state, **pk)
        centroids, labels, _sizes, Xaug, aug_idx = fit_clusters(feats, K, **ckw)
        if callable(cluster_space) or cluster_space != "physical":
            # labels (and the overlap band -- distances to centroids) were assigned in
            # feature space; aug_idx maps the augmented columns back to physical
            # snapshots, and the model's centroids must live in physical space
            # (means of the assigned columns), as split/build.py did.
            Xpod = Xt[:, aug_idx]
            centroids = torch.stack([Xpod[:, labels == k].mean(dim=1) for k in range(K)])
        else:
            # overlap-safe: Xaug's columns align with labels (repeated for overlapping
            # assignments), so per-cluster POD sees every column assigned to the cluster.
            Xpod = Xaug
        centroids = centroids.to(dtype=rdtype, device=Xt.device)

    # Cluster-switch bookkeeping for include_transition_snapshots: a switch at
    # column j (labels[j] != labels[j+1]) only counts when j and j+1 are genuinely
    # adjacent in time (time_index gap of exactly 1 -- multi-run seams are not
    # physical transitions).
    switch_next = None
    if include_transition_snapshots and K > 1:
        if Xpod.shape[1] != Xt.shape[1]:
            raise NotImplementedError(
                "include_transition_snapshots with assign_overlapping is not supported.")
        Nt = Xt.shape[1]
        ti = (torch.arange(Nt) if time_index is None
              else torch.as_tensor(time_index).detach().cpu().reshape(-1).long())
        if ti.shape != (Nt,):
            raise ValueError(f"time_index must have shape ({Nt},), got {tuple(ti.shape)}.")
        lab_cpu = torch.as_tensor(labels).detach().cpu().long()
        switch_next = ((ti[1:] - ti[:-1]) == 1) & (lab_cpu[1:] != lab_cpu[:-1])

    Phi_all = torch.zeros(N, r, K, dtype=rdtype, device=Xt.device)
    for k in range(K):
        cols = torch.nonzero(labels == k, as_tuple=False).flatten()
        if cols.numel() == 0:
            raise ValueError(f"cluster {k} is empty; reduce K.")
        Xc = Xpod[:, cols] - centroids[k][:, None]           # center on the centroid
        if switch_next is not None:
            outgoing = torch.nonzero(switch_next & (lab_cpu[:-1] == k), as_tuple=False).flatten() + 1
            if outgoing.numel():
                Xc = torch.cat([Xc, Xt[:, outgoing.to(Xt.device)] - centroids[k][:, None]], dim=1)
        Phi_all[:, :, k] = compute_pod_basis(Xc, r, method=pod_method, Mw=Mw, **pk)
    atlas = Atlas(centroids, Phi_all, dt=dt, Mw=Mw, method="pairwise")
    if build_atlas:
        atlas.build_global_atlas(Xt, assume_tgk_transpose=assume_tgk_transpose)
        atlas.method = "auto"
    return atlas