Skip to content

qlroms.metric

The inner-product weight Mw in one place: every norm, distance and nearest-centroid rule on PHYSICAL states goes through here.

Mw is None (identity), a (N,) diagonal, or a dense / sparse (N, N) matrix -- typically an assembled FEM mass matrix. Weights that are all equal (uniform grids: ks1d, ks2d) change nothing but a scale, so they take the plain Euclidean path and reproduce the unweighted results bit for bit.

apply_weight      M @ x for any Mw            (torch)
weight_kind       none | uniform | diag | matrix
wsq, wsq_rows     x^T M x per column / per row (torch or numpy)
col_norm, fro_norm   column / Frobenius norms in the M metric (plain for none / uniform)
total_weight      1^T M 1, the normaliser of a weighted RMS
metric_features   F with F^T F = Xc^T M Xc, so plain k-means on F clusters in the M metric
CentroidIndex     nearest-centroid rule in the M metric, KDTree-compatible query()
save_weight / load_weight   sidecar file (CSR + sha1) for a matrix Mw
pack_weight / unpack_weight   what a saved model stores about its Mw

CentroidIndex

Nearest-centroid rule in the M metric, with the scipy KDTree query() surface (query(x, k=1) -> (distance, index); x is (N,) or (m, N)).

none / uniform weight: the plain KDTree (identical results to the unweighted rule). diagonal weight: KDTree on sqrt(w)-scaled centroids. matrix weight: exact brute force, score_k = c_k^T M c_k - 2 (M c_k)^T x with M C precomputed once and everything centred on the mean centroid for precision (O(K N) per state). labels(X) skips the distances, which need one extra M product.

Source code in qlroms/metric.py 
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
class CentroidIndex:
    """Nearest-centroid rule in the M metric, with the scipy KDTree query() surface
    (`query(x, k=1) -> (distance, index)`; x is (N,) or (m, N)).

    none / uniform weight: the plain KDTree (identical results to the unweighted rule).
    diagonal weight: KDTree on sqrt(w)-scaled centroids.
    matrix weight: exact brute force, score_k = c_k^T M c_k - 2 (M c_k)^T x with M C
    precomputed once and everything centred on the mean centroid for precision (O(K N)
    per state). `labels(X)` skips the distances, which need one extra M product.
    """

    def __init__(self, centroids, Mw=None):
        C = centroids.detach().cpu().numpy() if torch.is_tensor(centroids) else np.asarray(centroids)
        C = np.asarray(C, dtype=float)
        self.K, self.N = C.shape
        self.kind = weight_kind(Mw)
        self._Mw = Mw
        self._M_np = None
        self._tree = self._sw = None
        self._scale = 1.0        # uniform weight: labels are Euclidean, reported distances are M distances
        if self.kind in ("none", "uniform"):
            self._tree = KDTree(C)
            if self.kind == "uniform":
                self._scale = float(np.sqrt(as_diagonal(Mw)[0]))
        elif self.kind == "diag":
            self._sw = np.sqrt(as_diagonal(Mw))
            self._tree = KDTree(C * self._sw)
        else:
            self._cbar = C.mean(axis=0)
            Cc = C - self._cbar
            self._MC = np.asarray(self._M() @ Cc.T)                        # (N, K)
            self._cc = np.einsum("kn,nk->k", Cc, self._MC)                  # c_k^T M c_k

    def __getstate__(self):
        state = self.__dict__.copy()
        state["_M_np"] = None            # rebuilt on demand from the (shared) weight object
        return state

    def _M(self):
        if self._M_np is None:
            self._M_np = _np_weight(self._Mw)
        return self._M_np

    def _score(self, Q):
        return self._cc[None, :] - 2.0 * ((Q - self._cbar) @ self._MC)

    def labels(self, X) -> np.ndarray:
        """Nearest-centroid index of each row of X (m, N) -> (m,) int."""
        X = np.asarray(X, dtype=float)
        if self.kind == "matrix":
            return self._score(X).argmin(axis=1)
        return np.asarray(self._tree.query(X if self._sw is None else X * self._sw, k=1)[1])

    def distances(self, X) -> np.ndarray:
        """(m, K) M-distances from each row of X (m, N) to every centroid."""
        X = np.asarray(X, dtype=float)
        if self.kind == "matrix":
            Qc = X - self._cbar
            qmq = np.einsum("mn,nm->m", Qc, np.asarray(self._M() @ Qc.T))
            return np.sqrt(np.maximum(self._score(X) + qmq[:, None], 0.0))
        C = self._tree.data
        return self._scale * cdist(X if self._sw is None else X * self._sw, C)

    def query(self, x, k: int = 1):
        x = np.asarray(x, dtype=float)
        if self.kind != "matrix":
            d, i = self._tree.query(x if self._sw is None else x * self._sw, k=k)
            return (d * self._scale if self._scale != 1.0 else d), i
        if k != 1:
            raise NotImplementedError("CentroidIndex with a matrix weight supports k=1 only.")
        single = x.ndim == 1
        X = x[None, :] if single else x
        d = self.distances(X)
        idx = d.argmin(axis=1)
        dist = d[np.arange(len(idx)), idx]
        return (dist[0], idx[0]) if single else (dist, idx)

labels(X)

Nearest-centroid index of each row of X (m, N) -> (m,) int.

Source code in qlroms/metric.py 
243
244
245
246
247
248
def labels(self, X) -> np.ndarray:
    """Nearest-centroid index of each row of X (m, N) -> (m,) int."""
    X = np.asarray(X, dtype=float)
    if self.kind == "matrix":
        return self._score(X).argmin(axis=1)
    return np.asarray(self._tree.query(X if self._sw is None else X * self._sw, k=1)[1])

distances(X)

(m, K) M-distances from each row of X (m, N) to every centroid.

Source code in qlroms/metric.py 
250
251
252
253
254
255
256
257
258
def distances(self, X) -> np.ndarray:
    """(m, K) M-distances from each row of X (m, N) to every centroid."""
    X = np.asarray(X, dtype=float)
    if self.kind == "matrix":
        Qc = X - self._cbar
        qmq = np.einsum("mn,nm->m", Qc, np.asarray(self._M() @ Qc.T))
        return np.sqrt(np.maximum(self._score(X) + qmq[:, None], 0.0))
    C = self._tree.data
    return self._scale * cdist(X if self._sw is None else X * self._sw, C)

apply_weight(Mw, x)

M @ x for the chart inner-product weight: identity when Mw is None, elementwise for a (N,) diagonal, matmul for a dense/sparse (N, N) matrix.

Source code in qlroms/metric.py 
33
34
35
36
37
38
39
40
41
42
def apply_weight(Mw, x: torch.Tensor) -> torch.Tensor:
    """M @ x for the chart inner-product weight: identity when Mw is None,
    elementwise for a (N,) diagonal, matmul for a dense/sparse (N, N) matrix."""
    if Mw is None:
        return x
    if not torch.is_tensor(Mw):
        Mw = as_torch_weight(Mw, x)              # ndarray / scipy sparse / PETSc
    if Mw.ndim == 1:
        return Mw[:, None] * x
    return torch.sparse.mm(Mw, x) if Mw.is_sparse else Mw @ x

as_torch_weight(Mw, like)

Mw as a torch weight on the dtype/device of like (sparse matrices stay sparse COO).

Source code in qlroms/metric.py 
71
72
73
74
75
76
77
78
79
80
81
82
83
def as_torch_weight(Mw, like: torch.Tensor):
    """Mw as a torch weight on the dtype/device of `like` (sparse matrices stay sparse COO)."""
    if Mw is None:
        return None
    if torch.is_tensor(Mw):
        return Mw.to(dtype=like.dtype, device=like.device)
    M = _np_weight(Mw)
    if sp.issparse(M):
        M = M.tocoo()
        idx = torch.from_numpy(np.vstack([M.row, M.col]).astype(np.int64))
        return torch.sparse_coo_tensor(idx, torch.from_numpy(M.data), size=M.shape,
                                       dtype=like.dtype, device=like.device)
    return torch.as_tensor(M, dtype=like.dtype, device=like.device)

weight_kind(Mw)

'none' (identity), 'uniform' (all-equal diagonal: a pure scale, Euclidean), 'diag' (non-uniform diagonal) or 'matrix'.

Source code in qlroms/metric.py 
86
87
88
89
90
91
92
93
94
95
96
97
def weight_kind(Mw) -> str:
    """'none' (identity), 'uniform' (all-equal diagonal: a pure scale, Euclidean),
    'diag' (non-uniform diagonal) or 'matrix'."""
    if Mw is None:
        return "none"
    M = _np_weight(Mw)
    if M.ndim == 1:
        return "uniform" if bool(np.all(M == M[0])) else "diag"
    if sp.issparse(M) and M.nnz == M.shape[0] and np.all(M.indices == np.arange(M.shape[0])):
        w = np.asarray(M.data)
        return "uniform" if bool(np.all(w == w[0])) else "diag"      # a sparse matrix that is diagonal
    return "matrix"

as_diagonal(Mw)

The (N,) diagonal of a 'diag'/'uniform' weight (1-D or diagonal sparse matrix) as ndarray.

Source code in qlroms/metric.py 
100
101
102
103
def as_diagonal(Mw):
    """The (N,) diagonal of a 'diag'/'uniform' weight (1-D or diagonal sparse matrix) as ndarray."""
    M = _np_weight(Mw)
    return M if M.ndim == 1 else np.asarray(M.data)

wsq(Mw, X)

x^T M x for X of shape (N,) or (N, m) -> scalar or (m,). torch in, torch out; numpy in, numpy out.

Source code in qlroms/metric.py 
116
117
118
119
120
121
122
123
124
125
126
127
128
def wsq(Mw, X):
    """x^T M x for X of shape (N,) or (N, m) -> scalar or (m,). torch in, torch out;
    numpy in, numpy out."""
    if torch.is_tensor(X):
        squeeze = X.ndim == 1
        Xm = X[:, None] if squeeze else X
        out = (Xm * apply_weight(as_torch_weight(Mw, Xm), Xm)).sum(dim=0)
        return out[0] if squeeze else out
    X = np.asarray(X)
    squeeze = X.ndim == 1
    Xm = X[:, None] if squeeze else X
    out = (Xm * _wapply_np(_np_weight(Mw), Xm)).sum(axis=0)
    return out[0] if squeeze else out

wsq_rows(Mw, D)

x^T M x for the ROWS of D (m, N), torch or numpy -> (m,). Same as wsq(Mw, D.T).

Source code in qlroms/metric.py 
131
132
133
def wsq_rows(Mw, D):
    """x^T M x for the ROWS of D (m, N), torch or numpy -> (m,). Same as wsq(Mw, D.T)."""
    return wsq(Mw, D.T)

wnorm(Mw, X)

sqrt of wsq (clamped at zero against round-off).

Source code in qlroms/metric.py 
136
137
138
139
def wnorm(Mw, X):
    """sqrt of wsq (clamped at zero against round-off)."""
    s = wsq(Mw, X)
    return torch.sqrt(s.clamp_min(0)) if torch.is_tensor(s) else np.sqrt(np.maximum(s, 0.0))

col_norm(Mw, X, kind=None)

Norm of each column of X (N, m) in the M metric -> (m,). None and uniform weights give the plain 2-norm (a uniform scale cancels in every relative error). Pass the weight's kind when calling in a loop, to avoid re-classifying a large matrix.

Source code in qlroms/metric.py 
142
143
144
145
146
147
148
149
def col_norm(Mw, X, kind: str | None = None):
    """Norm of each column of X (N, m) in the M metric -> (m,). None and uniform weights
    give the plain 2-norm (a uniform scale cancels in every relative error). Pass the
    weight's `kind` when calling in a loop, to avoid re-classifying a large matrix."""
    kind = weight_kind(Mw) if kind is None else kind
    if kind in ("none", "uniform"):
        return torch.linalg.norm(X, dim=0) if torch.is_tensor(X) else np.linalg.norm(X, axis=0)
    return wnorm(Mw, X)

fro_norm(Mw, X, kind=None)

sqrt(sum_j x_j^T M x_j) over all columns of X (N, m): the Frobenius norm in the M metric (plain Frobenius for None / uniform weights).

Source code in qlroms/metric.py 
152
153
154
155
156
157
158
159
def fro_norm(Mw, X, kind: str | None = None):
    """sqrt(sum_j x_j^T M x_j) over all columns of X (N, m): the Frobenius norm in the M
    metric (plain Frobenius for None / uniform weights)."""
    kind = weight_kind(Mw) if kind is None else kind
    if kind in ("none", "uniform"):
        return torch.linalg.norm(X) if torch.is_tensor(X) else np.linalg.norm(X)
    s = wsq(Mw, X).sum()
    return torch.sqrt(s.clamp_min(0)) if torch.is_tensor(s) else np.sqrt(max(float(s), 0.0))

total_weight(Mw, N)

1^T M 1: the total measure. sqrt(r^T M r / total_weight) is the weighted RMS of r, which equals the plain RMS for a uniform weight (and for Mw=None).

Source code in qlroms/metric.py 
162
163
164
165
166
167
168
def total_weight(Mw, N: int) -> float:
    """1^T M 1: the total measure. sqrt(r^T M r / total_weight) is the weighted RMS of r,
    which equals the plain RMS for a uniform weight (and for Mw=None)."""
    M = _np_weight(Mw)
    if M is None:
        return float(N)
    return float(M.sum()) if M.ndim == 1 else float(np.ones(N) @ (M @ np.ones(N)))

metric_features(X, Mw, *, tol=1e-12)

Feature matrix F (rank, n) whose columns are the snapshots of X (N, n) in a coordinate system where plain Euclidean distance IS the M distance: F^T F = Xc^T M Xc, so |F_i - F_j|^2 = (x_i - x_j)^T M (x_i - x_j).

Diagonal weights: sqrt(w) * X. Matrix weights: sqrt(Lambda) V^T from the snapshot Gram Xc^T M Xc = V Lambda V^T (n x n; directions below tol * lambda_max are dropped) -- exact, and cheaper than an SVD of the (N, n) snapshot matrix when N >> n.

Source code in qlroms/metric.py 
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
def metric_features(X: torch.Tensor, Mw, *, tol: float = 1e-12) -> torch.Tensor:
    """Feature matrix F (rank, n) whose columns are the snapshots of X (N, n) in a
    coordinate system where plain Euclidean distance IS the M distance:
    F^T F = Xc^T M Xc, so |F_i - F_j|^2 = (x_i - x_j)^T M (x_i - x_j).

    Diagonal weights: sqrt(w) * X. Matrix weights: sqrt(Lambda) V^T from the snapshot Gram
    Xc^T M Xc = V Lambda V^T (n x n; directions below tol * lambda_max are dropped) -- exact,
    and cheaper than an SVD of the (N, n) snapshot matrix when N >> n."""
    kind = weight_kind(Mw)
    if kind in ("none", "uniform"):
        return X
    if kind == "diag":
        w = torch.as_tensor(as_diagonal(Mw), dtype=X.dtype, device=X.device)
        return torch.sqrt(w)[:, None] * X
    Mt = as_torch_weight(Mw, X)
    Xc = X - X.mean(dim=1, keepdim=True)
    G = Xc.T @ apply_weight(Mt, Xc)
    evals, evecs = torch.linalg.eigh(0.5 * (G + G.T))
    keep = evals > tol * evals[-1].clamp_min(0)
    return torch.sqrt(evals[keep])[:, None] * evecs[:, keep].T

nearest_labels(tree, X)

Nearest-centroid labels of the rows of X (m, N) for a CentroidIndex, or any object with the scipy KDTree query() surface (a KDTree handed to Atlas by hand).

Source code in qlroms/metric.py 
275
276
277
278
279
280
def nearest_labels(tree, X) -> np.ndarray:
    """Nearest-centroid labels of the rows of X (m, N) for a CentroidIndex, or any object with
    the scipy KDTree query() surface (a KDTree handed to Atlas by hand)."""
    if isinstance(tree, CentroidIndex):
        return tree.labels(X)
    return np.asarray(tree.query(np.asarray(X), k=1)[1]).reshape(-1)

weight_sha1(Mw)

sha1 of the weight's contents (CSR arrays, or the dense/diagonal values).

Source code in qlroms/metric.py 
293
294
295
296
297
298
299
300
301
302
303
304
def weight_sha1(Mw) -> str:
    """sha1 of the weight's contents (CSR arrays, or the dense/diagonal values)."""
    M = _np_weight(Mw)
    h = hashlib.sha1()
    if sp.issparse(M):
        M.sort_indices()
        for a in (M.indptr, M.indices, M.data):
            h.update(np.ascontiguousarray(a).tobytes())
        h.update(str(M.shape).encode())
    else:
        h.update(np.ascontiguousarray(M).tobytes())
    return h.hexdigest()

save_weight(Mw, path)

Write a weight next to the data it belongs to (CSR components in a compressed npz for a matrix; the directory's weights.json maps sha1 -> file name) and return its sha1. Models fitted with it store only the path and this hash.

Source code in qlroms/metric.py 
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
def save_weight(Mw, path) -> str:
    """Write a weight next to the data it belongs to (CSR components in a compressed npz
    for a matrix; the directory's weights.json maps sha1 -> file name) and return its sha1.
    Models fitted with it store only the path and this hash."""
    path = Path(path)
    path.parent.mkdir(parents=True, exist_ok=True)
    M = _np_weight(Mw)
    if sp.issparse(M):
        np.savez_compressed(path, kind="csr", data=M.data, indices=M.indices, indptr=M.indptr,
                            shape=np.asarray(M.shape))
    else:
        np.savez_compressed(path, kind="dense", data=np.asarray(M))
    sha = weight_sha1(Mw)
    manifest = path.parent / _MANIFEST
    entries = json.loads(manifest.read_text()) if manifest.is_file() else {}
    entries[sha] = path.name
    manifest.write_text(json.dumps(entries, indent=2))
    _WEIGHT_REFS[sha] = str(path.resolve())
    return sha

load_weight(path, sha1=None, dtype=torch.float64)

Read a weight written by save_weight -> torch (sparse COO for a matrix, dense otherwise). If sha1 is given the contents must match it, or a ValueError names both.

Source code in qlroms/metric.py 
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
def load_weight(path, sha1: str | None = None, dtype=torch.float64):
    """Read a weight written by save_weight -> torch (sparse COO for a matrix, dense
    otherwise). If `sha1` is given the contents must match it, or a ValueError names both."""
    path = Path(path)
    if not path.is_file():
        raise FileNotFoundError(f"weight file {path} not found (expected sha1 {sha1}); pass "
                                "Mw=... explicitly or regenerate the file.")
    z = np.load(path, allow_pickle=False)
    if str(z["kind"]) == "csr":
        M = sp.csr_matrix((z["data"], z["indices"], z["indptr"]), shape=tuple(int(s) for s in z["shape"]))
        got = weight_sha1(M)
        out = as_torch_weight(M, torch.zeros(0, dtype=dtype))
    else:
        arr = np.asarray(z["data"])
        got = weight_sha1(arr)
        out = torch.as_tensor(arr, dtype=dtype)
    if sha1 is not None and got != sha1:
        raise ValueError(f"weight file {path} has sha1 {got}, the model was fitted with {sha1}.")
    _WEIGHT_REFS[got] = str(path.resolve())
    return out

default_weight_dirs()

Where sidecars are looked for when a model file does not point at one that exists: the pinball mesh directory of qlroms.utils.paths.

Source code in qlroms/metric.py 
362
363
364
365
366
def default_weight_dirs() -> list[Path]:
    """Where sidecars are looked for when a model file does not point at one that exists:
    the pinball mesh directory of qlroms.utils.paths."""
    from .utils import paths
    return [Path(paths.PINBALL_MESHES)]

pack_weight(Mw, search_dirs=None)

The entries a saved model file stores for its Mw: nothing for None, atlas_Mw (the values) for a diagonal, and atlas_Mw_ref + atlas_Mw_sha1 for a matrix, which must have been written with save_weight / export_mass_matrices (or read with load_weight) so that a sidecar file exists; otherwise a ValueError tells how to create it.

Source code in qlroms/metric.py 
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
def pack_weight(Mw, search_dirs=None) -> dict:
    """The entries a saved model file stores for its Mw: nothing for None, `atlas_Mw` (the
    values) for a diagonal, and `atlas_Mw_ref` + `atlas_Mw_sha1` for a matrix, which must have
    been written with save_weight / export_mass_matrices (or read with load_weight) so that a
    sidecar file exists; otherwise a ValueError tells how to create it."""
    kind = weight_kind(Mw)
    if kind == "none":
        return {}
    if kind in ("uniform", "diag"):
        return {"atlas_Mw": as_diagonal(Mw)}
    sha = weight_sha1(Mw)
    ref = _find_sidecar(sha, default_weight_dirs() if search_dirs is None else search_dirs)
    if ref is None:
        raise ValueError(f"the matrix Mw (sha1 {sha}) has no sidecar file; write it once with "
                         "qlroms.metric.save_weight(Mw, path) (the pinball build does it next to the "
                         "mesh), then save the model again.")
    return {"atlas_Mw_ref": np.asarray(ref), "atlas_Mw_sha1": np.asarray(sha)}

unpack_weight(z, override=None, search_dirs=None)

The Mw a model file was fitted with: override if given, else the stored diagonal, else the sidecar named by (atlas_Mw_ref, atlas_Mw_sha1) -- looked up by hash in the mesh directory when the stored path no longer exists. None if the file has no Mw.

Source code in qlroms/metric.py 
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
def unpack_weight(z, override=None, search_dirs=None):
    """The Mw a model file was fitted with: `override` if given, else the stored diagonal, else
    the sidecar named by (atlas_Mw_ref, atlas_Mw_sha1) -- looked up by hash in the mesh
    directory when the stored path no longer exists. None if the file has no Mw."""
    if override is not None:
        return override
    if "atlas_Mw" in z:
        return np.asarray(z["atlas_Mw"])
    if "atlas_Mw_sha1" not in z:
        return None
    sha = str(np.asarray(z["atlas_Mw_sha1"]))
    ref = str(np.asarray(z["atlas_Mw_ref"])) if "atlas_Mw_ref" in z else ""
    path = ref if ref and Path(ref).is_file() else _find_sidecar(
        sha, default_weight_dirs() if search_dirs is None else search_dirs)
    if path is None:
        raise FileNotFoundError(f"the model was fitted with a mass matrix (sha1 {sha}, last seen at "
                                f"{ref or 'unknown'}) that cannot be found; pass Mw=... to load().")
    return load_weight(path, sha)