1
2
3
4
5
6
7
8
9
10
11
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
|
"""
This module provides the Trace class.
"""
import weakref
from typing import Any, Mapping, Sequence
from copy import copy, deepcopy
from numpy import ndarray
import numpy as np
from public import public
@public
class Trace(object):
"""A trace, which has some samples and metadata."""
meta: Mapping[str, Any]
samples: ndarray
def __init__(self, samples: ndarray, meta: Mapping[str, Any] = None, trace_set: Any = None):
"""
Construct a new trace.
:param samples: The sample array of the trace.
:param meta: Metadata associated with the trace.
:param trace_set: A trace set the trace is contained in.
"""
if meta is None:
meta = {}
self.meta = meta
self.samples = samples
self.trace_set = trace_set
def __len__(self):
"""Length of the trace, in samples."""
return len(self.samples)
def __getitem__(self, index):
"""Get the sample at `index`."""
return self.samples[index]
def __setitem__(self, key, value):
"""Set the sample at `key`."""
self.samples[key] = value
def __iter__(self):
"""Iterate over the samples."""
yield from self.samples
@property
def trace_set(self) -> Any:
"""
The trace set this trace is contained in, if any.
"""
if self._trace_set is None:
return None
return self._trace_set()
@trace_set.setter
def trace_set(self, trace_set: Any):
"""
Set the trace set of this trace.
"""
if trace_set is None:
self._trace_set = None
else:
self._trace_set = weakref.ref(trace_set)
def __getstate__(self):
state = self.__dict__.copy()
state["_trace_set"] = None
return state
def __setstate__(self, state):
self._trace_set = None
self.__dict__.update(state)
def __eq__(self, other):
if not isinstance(other, Trace):
return False
return np.array_equal(self.samples, other.samples) and self.meta == other.meta
def with_samples(self, samples: ndarray) -> "Trace":
"""
Construct a copy of this trace, with the same metadata, but samples replaced by `samples`.
:param samples: The samples of the new trace.
:return: The new trace.
"""
return Trace(samples, deepcopy(self.meta))
def __copy__(self):
return Trace(copy(self.samples), copy(self.meta), copy(self.trace_set))
def __deepcopy__(self, memodict={}):
return Trace(deepcopy(self.samples, memo=memodict), deepcopy(self.meta, memo=memodict),
deepcopy(self.trace_set, memo=memodict))
def __repr__(self):
return f"Trace(samples={self.samples!r}, trace_set={self.trace_set!r})"
@public
class CombinedTrace(Trace):
"""A trace that was combined from other traces, `parents`."""
def __init__(self, samples: ndarray, meta: Mapping[str, Any] = None, trace_set: Any = None,
parents: Sequence[Trace] = None):
super().__init__(samples, meta, trace_set=trace_set)
self.parents = None
if parents is not None:
self.parents = weakref.WeakSet(parents)
def __repr__(self):
return f"CombinedTrace(samples={self.samples!r}, trace_set={self.trace_set!r}, parents={self.parents})"
|
