Coverage for /opt/hostedtoolcache/Python/3.10.20/x64/lib/python3.10/site-packages/equine/utils.py: 100%
168 statements
« prev ^ index » next coverage.py v7.14.3, created at 2026-06-25 02:01 +0000
1# Copyright 2024, MASSACHUSETTS INSTITUTE OF TECHNOLOGY
2# Subject to FAR 52.227-11 – Patent Rights – Ownership by the Contractor (May 2014).
3# SPDX-License-Identifier: MIT
5import sys
6from collections import OrderedDict
7from typing import Any, Union
9import icontract
10import torch
11from beartype import beartype
12from torchmetrics.classification import (
13 MulticlassAccuracy,
14 MulticlassCalibrationError,
15 MulticlassConfusionMatrix,
16 MulticlassF1Score,
17)
19from .equine import Equine
20from .equine_output import EquineOutput
23@icontract.ensure(lambda result, module: result is module)
24@beartype
25def prepare_jit_module(module: torch.nn.Module) -> torch.nn.Module:
26 """
27 Make an ``nn.Module`` safe to pass to ``torch.jit.script`` on Python 3.14+.
29 Starting with Python 3.14 (PEP 649/749) a class's ``__annotations__`` is
30 exposed through a ``type`` getset descriptor instead of being stored in the
31 class ``__dict__``. ``torch.jit``'s scripting checker reads
32 ``__annotations__`` from the module *instance*, where attribute lookup only
33 consults the ``__dict__`` of each class in the MRO. For any ``nn.Module``
34 that declares no class-level annotations this lookup misses and
35 ``nn.Module.__getattr__`` raises ``AttributeError``, breaking
36 ``torch.jit.script``. Materializing each submodule's own class annotations
37 onto its instance ``__dict__`` lets that lookup succeed. This is a no-op on
38 Python < 3.14.
40 Parameters
41 ----------
42 module : torch.nn.Module
43 The module that is about to be scripted with ``torch.jit.script``.
45 Returns
46 -------
47 torch.nn.Module
48 The same module, returned for convenient inline use.
49 """
50 if sys.version_info >= (3, 14):
51 for submodule in module.modules():
52 if "__annotations__" not in submodule.__dict__:
53 submodule.__dict__["__annotations__"] = dict(
54 type(submodule).__annotations__
55 )
56 return module
59@icontract.require(lambda y_hat, y_test: y_hat.size(dim=0) == y_test.size(dim=0))
60@icontract.ensure(lambda result: result >= 0.0)
61@beartype
62def brier_score(y_hat: torch.Tensor, y_test: torch.Tensor) -> float:
63 """
64 Compute the Brier score for a multiclass problem:
65 $$ \\frac{1}{N} \\sum_{i=1}^{N} \\sum_{j=1}^{M} (f_{ij} - o_{ij})^2 , $$
66 where $f_{ij}$ is the predicted probability of class $j$ for inference sample $i$
67 and $o_{ij}$ is the one-hot encoded ground truth label.
69 Parameters
70 ----------
71 y_hat : torch.Tensor
72 Probabilities for each class.
73 y_test : torch.Tensor
74 Integer argument class labels (ground truth).
76 Returns
77 -------
78 float
79 Brier score.
80 """
81 _, num_classes = y_hat.size()
82 one_hot_y_test = torch.nn.functional.one_hot(y_test.long(), num_classes=num_classes)
83 bs = torch.mean(torch.sum((y_hat - one_hot_y_test) ** 2, dim=1)).item()
84 return bs
87@icontract.require(lambda y_hat, y_test: y_hat.size(dim=0) == y_test.size(dim=0))
88@icontract.ensure(lambda result: result <= 1.0)
89@beartype
90def brier_skill_score(y_hat: torch.Tensor, y_test: torch.Tensor) -> float:
91 """
92 Compute the Brier skill score as compared to randomly guessing.
94 Parameters
95 ----------
96 y_hat : torch.Tensor
97 Probabilities for each class.
98 y_test : torch.Tensor
99 Integer argument class labels (ground truth).
101 Returns
102 -------
103 float
104 Brier skill score.
105 """
106 _, num_classes = y_hat.size()
107 random_guess = (1.0 / num_classes) * torch.ones(y_hat.size())
108 bs0 = brier_score(random_guess, y_test)
109 bs1 = brier_score(y_hat, y_test)
110 bss = 1.0 - bs1 / bs0
111 return bss
114@icontract.require(lambda y_hat, y_test: y_hat.size(dim=0) == y_test.size(dim=0))
115@icontract.ensure(lambda result: (0.0 <= result) and (result <= 1.0))
116@beartype
117def expected_calibration_error(y_hat: torch.Tensor, y_test: torch.Tensor) -> float:
118 """
119 Compute the expected calibration error (ECE) for a multiclass problem.
121 Parameters
122 ----------
123 y_hat : torch.Tensor
124 Probabilities for each class.
125 y_test : torch.Tensor
126 Class label indices (ground truth).
128 Returns
129 -------
130 float
131 Expected calibration error.
132 """
133 _, num_classes = y_hat.size()
134 metric = MulticlassCalibrationError(num_classes=num_classes, n_bins=25, norm="l1")
135 ece = metric(y_hat, y_test).item()
136 return ece
139@icontract.require(
140 lambda train_y, selected_labels: len(selected_labels) <= len(train_y)
141)
142@icontract.ensure(
143 lambda result, selected_labels: set(result.keys()).issubset(set(selected_labels))
144)
145@beartype
146def _get_shuffle_idxs_by_class(
147 train_y: torch.Tensor, selected_labels: list
148) -> dict[Any, torch.Tensor]:
149 """
150 Internal helper function to randomly select indices of example classes for a given
151 set of labels.
153 Parameters
154 ----------
155 train_y : torch.Tensor
156 Label data.
157 selected_labels : list
158 list of unique labels found in the label data.
160 Returns
161 -------
162 dict[Any, torch.Tensor]
163 Tensor of indices corresponding to each label.
164 """
165 shuffled_idxs_by_class = OrderedDict()
166 for label in selected_labels:
167 label_idxs = torch.argwhere(train_y == label).squeeze()
168 shuffled_idxs_by_class[label] = label_idxs[torch.randperm(label_idxs.shape[0])]
170 return shuffled_idxs_by_class
173@icontract.require(lambda train_x, train_y: len(train_x) <= len(train_y))
174@icontract.require(
175 lambda selected_labels, train_x: (
176 (0 < len(selected_labels)) & (len(selected_labels) < len(train_x))
177 )
178)
179@icontract.require(
180 lambda support_size, train_x: (0 < support_size) & (support_size < len(train_x))
181)
182@icontract.require(
183 lambda support_size, selected_labels, train_x: (
184 support_size * len(selected_labels) <= len(train_x)
185 )
186)
187@icontract.require(
188 lambda selected_labels, shuffled_indexes: (
189 (len(shuffled_indexes.keys()) == len(selected_labels))
190 if shuffled_indexes is not None
191 else True
192 )
193)
194@icontract.ensure(
195 lambda result, selected_labels: len(result.keys()) == len(selected_labels)
196)
197@beartype
198def generate_support(
199 train_x: torch.Tensor,
200 train_y: torch.Tensor,
201 support_size: int,
202 selected_labels: list[Any],
203 shuffled_indexes: Union[None, dict[Any, torch.Tensor]] = None,
204) -> OrderedDict[int, torch.Tensor]:
205 """
206 Randomly select `support_size` examples of `way` classes from the examples in
207 `train_x` with corresponding labels in `train_y` and return them as a dictionary.
209 Parameters
210 ----------
211 train_x : torch.Tensor
212 Input training data.
213 train_y : torch.Tensor
214 Corresponding classification labels.
215 support_size : int
216 Number of support examples for each class.
217 selected_labels : list
218 Selected class labels to generate examples from.
219 shuffled_indexes: Union[None, dict[Any, torch.Tensor]], optional
220 Simply use the precomputed indexes if they are available
222 Returns
223 -------
224 OrderedDict[int, torch.Tensor]
225 Ordered dictionary of class labels with corresponding support examples.
226 """
227 labels, counts = torch.unique(train_y, return_counts=True)
228 if shuffled_indexes is None:
229 for label, count in list(zip(labels, counts)):
230 if (label in selected_labels) and (count < support_size):
231 raise ValueError(f"Not enough support examples in class {label}")
232 shuffled_idxs = _get_shuffle_idxs_by_class(train_y, selected_labels)
233 else:
234 shuffled_idxs = shuffled_indexes
236 support = OrderedDict[int, torch.Tensor]()
237 for label in selected_labels:
238 shuffled_x = train_x[shuffled_idxs[label]]
240 assert torch.unique(train_y[shuffled_idxs[label]]).tolist() == [label], (
241 "Not enough support for label " + str(label)
242 )
243 selected_support = shuffled_x[:support_size]
244 support[int(label)] = selected_support
246 return support
249@icontract.require(lambda train_x: len(train_x.shape) >= 2)
250@icontract.require(lambda train_y: len(train_y.shape) == 1)
251@icontract.require(lambda support_size: support_size > 1)
252@icontract.require(lambda way: way > 0)
253@icontract.require(lambda episode_size: episode_size > 0)
254@icontract.ensure(lambda result: len(result) == 3)
255@icontract.ensure(lambda result: result[1].shape[0] == result[2].shape[0])
256@icontract.ensure(lambda way, result: len(result[0]) == way)
257@icontract.ensure(
258 lambda support_size, result: all(
259 len(support) == support_size for support in result[0].values()
260 )
261)
262@beartype
263def generate_episode(
264 train_x: torch.Tensor,
265 train_y: torch.Tensor,
266 support_size: int,
267 way: int,
268 episode_size: int,
269) -> tuple[OrderedDict[int, torch.Tensor], torch.Tensor, torch.Tensor]:
270 """
271 Generate a single episode of data for a few-shot learning task.
273 Parameters
274 ----------
275 train_x : torch.Tensor
276 Input training data.
277 train_y : torch.Tensor
278 Corresponding classification labels.
279 support_size : int
280 Number of support examples for each class.
281 way : int
282 Number of classes in the episode.
283 episode_size : int
284 Total number of examples in the episode.
286 Returns
287 -------
288 tuple[dict[Any, torch.Tensor], torch.Tensor, torch.Tensor]
289 tuple of support examples, query examples, and query labels.
290 """
291 labels, counts = torch.unique(train_y, return_counts=True)
292 if way > len(labels):
293 raise ValueError(
294 f"The way (#classes in each episode), {way}, must be <= number of labels, {len(labels)}"
295 )
297 selected_labels = sorted(
298 labels[torch.randperm(labels.shape[0])][:way].tolist()
299 ) # need to be in same order every time
301 for label, count in list(zip(labels, counts)):
302 if (label in selected_labels) and (count < support_size):
303 raise ValueError(f"Not enough support examples in class {label}")
304 shuffled_idxs = _get_shuffle_idxs_by_class(train_y, selected_labels)
306 support = generate_support(
307 train_x, train_y, support_size, selected_labels, shuffled_idxs
308 )
310 examples_per_task = episode_size // way
312 episode_data_list = []
313 episode_label_list = []
314 episode_support = OrderedDict()
315 for episode_label, label in enumerate(selected_labels):
316 shuffled_x = train_x[shuffled_idxs[label]]
317 shuffled_y = torch.Tensor(
318 [episode_label] * len(shuffled_idxs[label])
319 ) # need sequential labels for episode
321 num_remaining_examples = shuffled_x.shape[0] - support_size
322 assert num_remaining_examples > 0, (
323 "Cannot have "
324 + str(num_remaining_examples)
325 + " left with support_size "
326 + str(support_size)
327 + " and shape "
328 + str(shuffled_x.shape)
329 + " from train_x shaped "
330 + str(train_x.shape)
331 )
332 episode_end_idx = support_size + min(num_remaining_examples, examples_per_task)
334 episode_data_list.append(shuffled_x[support_size:episode_end_idx])
335 episode_label_list.append(shuffled_y[support_size:episode_end_idx])
336 episode_support[episode_label] = support[label]
338 episode_x = torch.concat(episode_data_list)
339 episode_y = torch.concat(episode_label_list)
341 return episode_support, episode_x, episode_y.squeeze().to(torch.long)
344@icontract.require(
345 lambda eq_preds, true_y: eq_preds.classes.size(dim=0) == true_y.size(dim=0)
346)
347@beartype
348def generate_model_metrics(
349 eq_preds: EquineOutput, true_y: torch.Tensor
350) -> dict[str, Any]:
351 """
352 Generate various metrics for evaluating a model's performance.
354 Parameters
355 ----------
356 eq_preds : EquineOutput
357 Model predictions.
358 true_y : torch.Tensor
359 True class labels.
361 Returns
362 -------
363 dict[str, Any]
364 Dictionary of model metrics.
365 """
366 pred_y = torch.argmax(eq_preds.classes, dim=1)
367 accuracy = MulticlassAccuracy(num_classes=eq_preds.classes.shape[1])
368 f1_score = MulticlassF1Score(num_classes=eq_preds.classes.shape[1], average="micro")
369 confusion_matrix = MulticlassConfusionMatrix(num_classes=eq_preds.classes.shape[1])
370 metrics = {
371 "accuracy": accuracy(true_y, pred_y),
372 "microF1Score": f1_score(true_y, pred_y),
373 "confusionMatrix": confusion_matrix(true_y, pred_y).tolist(),
374 "brierScore": brier_score(eq_preds.classes, true_y),
375 "brierSkillScore": brier_skill_score(eq_preds.classes, true_y),
376 "expectedCalibrationError": expected_calibration_error(
377 eq_preds.classes, true_y
378 ),
379 }
380 return metrics
383@icontract.require(lambda Y: len(Y.shape) == 1)
384@icontract.ensure(
385 lambda result: all("label" in d and "numExamples" in d for d in result)
386)
387@icontract.ensure(lambda result: all(d["numExamples"] >= 0 for d in result))
388@beartype
389def get_num_examples_per_label(Y: torch.Tensor) -> list[dict[str, Any]]:
390 """
391 Get the number of examples per label in the given tensor.
393 Parameters
394 ----------
395 Y : torch.Tensor
396 Tensor of class labels.
398 Returns
399 -------
400 list[dict[str, Any]]
401 list of dictionaries containing label and number of examples.
402 """
403 tensor_labels, tensor_counts = Y.unique(return_counts=True)
405 examples_per_label = []
406 for i, label in enumerate(tensor_labels):
407 examples_per_label.append(
408 {"label": label.item(), "numExamples": tensor_counts[i].item()}
409 )
411 return examples_per_label
414@icontract.require(lambda train_y: train_y.shape[0] > 0)
415@beartype
416def generate_train_summary(
417 model: Equine, train_y: torch.Tensor, date_trained: str
418) -> dict[str, Any]:
419 """
420 Generate a summary of the training data.
422 Parameters
423 ----------
424 model : Equine
425 Model object.
426 train_y : torch.Tensor
427 Training labels.
428 date_trained : str
429 Date of training.
431 Returns
432 -------
433 dict[str, Any]
434 Dictionary containing training summary.
435 """
436 train_summary = {
437 "numTrainExamples": get_num_examples_per_label(train_y),
438 "dateTrained": date_trained,
439 "modelType": model.__class__.__name__,
440 }
441 return train_summary
444@icontract.require(
445 lambda eq_preds, test_y: test_y.shape[0] == eq_preds.classes.shape[0]
446)
447@beartype
448def generate_model_summary(
449 model: Equine,
450 eq_preds: EquineOutput,
451 test_y: torch.Tensor,
452) -> dict[str, Any]:
453 """
454 Generate a summary of the model's performance.
456 Parameters
457 ----------
458 model : Equine
459 Model object.
460 eq_preds : EquineOutput
461 Model predictions.
462 test_y : torch.Tensor
463 True class labels.
465 Returns
466 -------
467 dict[str, Any]
468 Dictionary containing model summary.
469 """
470 summary = generate_model_metrics(eq_preds, test_y)
471 summary["numTestExamples"] = get_num_examples_per_label(test_y)
472 summary.update(model.train_summary) # union of train_summary and generated metrics
474 return summary
477@icontract.require(lambda cov: cov.shape[-2] == cov.shape[-1])
478def mahalanobis_distance_nosq(x: torch.Tensor, cov: torch.Tensor) -> torch.Tensor:
479 """
480 Compute Mahalanobis distance $x^T C x$ (without square root), assume cov is symmetric positive definite
482 Parameters
483 ----------
484 x : torch.Tensor
485 vectors to compute distances for
486 cov : torch.Tensor
487 covariance matrix, assumes first dimension is number of classes
488 """
489 U, S, _ = torch.linalg.svd(cov)
490 S_inv_sqrt = torch.stack(
491 [torch.diag(torch.sqrt(1.0 / S[i])) for i in range(S.shape[0])], dim=0
492 )
493 prod = torch.matmul(S_inv_sqrt, torch.transpose(U, 1, 2))
494 dist = torch.sum(torch.square(torch.matmul(prod, x)), dim=1)
495 return dist
498@icontract.require(
499 lambda X, Y: X.shape[0] == Y.shape[0],
500 "X and Y must have the same number of samples.",
501)
502@icontract.require(
503 lambda test_size: 0.0 < test_size < 1.0, "test_size must be between 0 and 1."
504)
505@icontract.ensure(
506 lambda result: len(result) == 4, "Function must return four elements."
507)
508@icontract.ensure(
509 lambda X, result: result[0].shape[0] + result[1].shape[0] == X.shape[0],
510 "Total samples must be preserved.",
511)
512@icontract.ensure(
513 lambda Y, result: result[2].shape[0] + result[3].shape[0] == Y.shape[0],
514 "Total labels must be preserved.",
515)
516@icontract.ensure(
517 lambda result: result[0].shape[0] == result[2].shape[0],
518 "Train features and labels must match in size.",
519)
520@icontract.ensure(
521 lambda result: result[1].shape[0] == result[3].shape[0],
522 "Test features and labels must match in size.",
523)
524@beartype
525def stratified_train_test_split(
526 X: torch.Tensor, Y: torch.Tensor, test_size: float
527) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
528 """
529 A pytorch-ified version of sklearn's train_test_split with data stratification
531 Parameters
532 ----------
533 X : torch.Tensor
534 Input features tensor of shape (n_samples, n_features).
535 Y : torch.Tensor
536 Labels tensor of shape (n_samples,).
537 test_size : float
538 Proportion of the dataset to include in the test split (between 0.0 and 1.0).
540 Returns
541 -------
542 train_x : torch.Tensor
543 Training set features.
544 calib_x : torch.Tensor
545 Test set features.
546 train_y : torch.Tensor
547 Training set labels.
548 calib_y : torch.Tensor
549 Test set labels.
550 """
551 unique_classes, class_counts = torch.unique(Y, return_counts=True)
552 test_counts = (class_counts.float() * test_size).round().long()
553 train_indices = []
554 test_indices = []
556 for cls, test_count in zip(unique_classes, test_counts):
557 cls_indices = torch.where(Y == cls)[0]
558 cls_indices = cls_indices[torch.randperm(len(cls_indices))]
559 test_idx = cls_indices[:test_count]
560 train_idx = cls_indices[test_count:]
561 train_indices.append(train_idx)
562 test_indices.append(test_idx)
564 train_indices = torch.cat(train_indices)
565 test_indices = torch.cat(test_indices)
567 train_x = X[train_indices]
568 train_y = Y[train_indices]
569 calib_x = X[test_indices]
570 calib_y = Y[test_indices]
572 return train_x, calib_x, train_y, calib_y