feat: adds psd calculation code and utilities to handle averaging of periodograms
All checks were successful
gitea-physics/tantri/pipeline/head This commit looks good
All checks were successful
gitea-physics/tantri/pipeline/head This commit looks good
This commit is contained in:
parent
586485e24b
commit
22bb3d876f
GPG Key ID:
BEBAEBF28083E022
@ -4,6 +4,56 @@ import numpy.random
|
|||||||
import typing
|
import typing
|
||||||
from tantri.dipoles.types import DipoleTO, DotPosition, DipoleMeasurementType
|
from tantri.dipoles.types import DipoleTO, DotPosition, DipoleMeasurementType
|
||||||
import tantri.dipoles.supersample
|
import tantri.dipoles.supersample
|
||||||
|
import tantri.util
|
||||||
|
import scipy.stats
|
||||||
|
import scipy.fft
|
||||||
|
import logging
|
||||||
|
|
||||||
|
_logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass
|
||||||
|
class APSDResult:
|
||||||
|
psd_dict: typing.Dict[str, numpy.ndarray]
|
||||||
|
freqs: numpy.ndarray
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass
|
||||||
|
class TimeSeriesResult:
|
||||||
|
series_dict: typing.Dict[str, numpy.ndarray]
|
||||||
|
num_points: int
|
||||||
|
delta_t: float
|
||||||
|
|
||||||
|
def get_time_points(self):
|
||||||
|
return [t * self.delta_t for t in range(self.num_points)]
|
||||||
|
|
||||||
|
def get_apsds(self) -> APSDResult:
|
||||||
|
def sq(a):
|
||||||
|
return numpy.real(a * numpy.conjugate(a))
|
||||||
|
|
||||||
|
def psd(v):
|
||||||
|
return sq(scipy.fft.rfft(v)[1:]) * self.delta_t / self.num_points
|
||||||
|
|
||||||
|
fft_dict = {k: psd(v) for k, v in self.series_dict.items()}
|
||||||
|
freqs = scipy.fft.rfftfreq(self.num_points, self.delta_t)[1:]
|
||||||
|
|
||||||
|
return APSDResult(fft_dict, freqs)
|
||||||
|
|
||||||
|
|
||||||
|
def average_apsds(apsds: typing.Sequence[APSDResult]) -> APSDResult:
|
||||||
|
def mean(list_of_arrays: typing.Sequence[numpy.ndarray]) -> numpy.ndarray:
|
||||||
|
return numpy.mean(numpy.array(list_of_arrays), axis=0)
|
||||||
|
|
||||||
|
if len(apsds) >= 1:
|
||||||
|
for subsequent in apsds[1:]:
|
||||||
|
if not numpy.array_equal(subsequent.freqs, apsds[0].freqs):
|
||||||
|
raise ValueError(
|
||||||
|
f"Could not average apsds, as {subsequent} does not match the frequencies in {apsds[0]}"
|
||||||
|
)
|
||||||
|
freqs = apsds[0].freqs
|
||||||
|
|
||||||
|
average_dict = tantri.util.dict_reduce([apsd.psd_dict for apsd in apsds], mean)
|
||||||
|
return APSDResult(average_dict, freqs)
|
||||||
|
|
||||||
|
|
||||||
@dataclass
|
@dataclass
|
||||||
@ -74,6 +124,27 @@ class WrappedDipole:
|
|||||||
self.flip_state()
|
self.flip_state()
|
||||||
return {k: self.state * v for k, v in self.cache.items()}
|
return {k: self.state * v for k, v in self.cache.items()}
|
||||||
|
|
||||||
|
def time_series(
|
||||||
|
self,
|
||||||
|
dt: float,
|
||||||
|
num_points: int,
|
||||||
|
rng_to_use: typing.Optional[numpy.random.Generator] = None,
|
||||||
|
) -> typing.Dict[str, numpy.ndarray]:
|
||||||
|
|
||||||
|
# don't forget to set rng
|
||||||
|
if rng_to_use is None:
|
||||||
|
rng = numpy.random.default_rng()
|
||||||
|
else:
|
||||||
|
rng = rng_to_use
|
||||||
|
|
||||||
|
# scale effective mu by the sample rate.
|
||||||
|
# mu (or w) has units of events/time, so effective rate is events/time * dt, giving mu per dt
|
||||||
|
eff_mu = dt * self.w
|
||||||
|
events = scipy.stats.poisson.rvs(eff_mu, size=num_points, random_state=rng)
|
||||||
|
telegraph_sequence = numpy.cumprod((-1) ** events)
|
||||||
|
|
||||||
|
return {k: telegraph_sequence * v for k, v in self.cache.items()}
|
||||||
|
|
||||||
def flip_state(self):
|
def flip_state(self):
|
||||||
self.state *= -1
|
self.state *= -1
|
||||||
|
|
||||||
@ -113,6 +184,8 @@ class DipoleTimeSeries:
|
|||||||
self.dipoles = get_wrapped_dipoles(dipoles, dots, measurement_type)
|
self.dipoles = get_wrapped_dipoles(dipoles, dots, measurement_type)
|
||||||
self.state = 0
|
self.state = 0
|
||||||
|
|
||||||
|
self.real_delta_t = dt
|
||||||
|
|
||||||
# we may need to supersample, because of how dumb this process is.
|
# we may need to supersample, because of how dumb this process is.
|
||||||
# let's find our highest frequency
|
# let's find our highest frequency
|
||||||
max_frequency = max(d.w for d in self.dipoles)
|
max_frequency = max(d.w for d in self.dipoles)
|
||||||
@ -135,3 +208,53 @@ class DipoleTimeSeries:
|
|||||||
|
|
||||||
def transition(self) -> typing.Dict[str, float]:
|
def transition(self) -> typing.Dict[str, float]:
|
||||||
return [self._sub_transition() for i in range(self.super_sample_ratio)][-1]
|
return [self._sub_transition() for i in range(self.super_sample_ratio)][-1]
|
||||||
|
|
||||||
|
def _generate_series(
|
||||||
|
self, num_points: int, delta_t: float
|
||||||
|
) -> typing.Dict[str, numpy.ndarray]:
|
||||||
|
|
||||||
|
serieses = [
|
||||||
|
dipole.time_series(delta_t, num_points, self.rng) for dipole in self.dipoles
|
||||||
|
]
|
||||||
|
|
||||||
|
result = {}
|
||||||
|
for series in serieses:
|
||||||
|
for k, v in series.items():
|
||||||
|
if k not in result:
|
||||||
|
result[k] = v
|
||||||
|
else:
|
||||||
|
result[k] += v
|
||||||
|
|
||||||
|
return result
|
||||||
|
|
||||||
|
def generate_series(
|
||||||
|
self, num_points: int, override_delta_t: typing.Optional[float] = None
|
||||||
|
) -> TimeSeriesResult:
|
||||||
|
|
||||||
|
delta_t_to_use: float
|
||||||
|
if override_delta_t is not None:
|
||||||
|
delta_t_to_use = override_delta_t
|
||||||
|
else:
|
||||||
|
delta_t_to_use = self.real_delta_t
|
||||||
|
|
||||||
|
series = self._generate_series(num_points, delta_t_to_use)
|
||||||
|
|
||||||
|
return TimeSeriesResult(
|
||||||
|
series_dict=series, num_points=num_points, delta_t=delta_t_to_use
|
||||||
|
)
|
||||||
|
|
||||||
|
def generate_average_apsd(
|
||||||
|
self,
|
||||||
|
num_series: int,
|
||||||
|
num_time_series_points: int,
|
||||||
|
override_delta_t: typing.Optional[float] = None,
|
||||||
|
) -> APSDResult:
|
||||||
|
|
||||||
|
apsds = [
|
||||||
|
self.generate_series(num_time_series_points, override_delta_t).get_apsds()
|
||||||
|
for _ in range(num_series)
|
||||||
|
]
|
||||||
|
|
||||||
|
_logger.debug(f"Averaging {num_series} series")
|
||||||
|
|
||||||
|
return average_apsds(apsds)
|
||||||
|
|||||||
21
tantri/util.py
Normal file
21
tantri/util.py
Normal file
@ -0,0 +1,21 @@
|
|||||||
|
import typing
|
||||||
|
|
||||||
|
A = typing.TypeVar("A")
|
||||||
|
B = typing.TypeVar("B")
|
||||||
|
|
||||||
|
|
||||||
|
def dict_reduce(
|
||||||
|
list_of_dicts: typing.Sequence[typing.Dict[str, A]],
|
||||||
|
func: typing.Callable[[typing.Sequence[A]], B],
|
||||||
|
) -> typing.Dict[str, B]:
|
||||||
|
"""
|
||||||
|
Reduce over list of dicts with function that can coalesce list of dict values.
|
||||||
|
Assumes the keys in the first dictionary are the same as the keys for every other passed in dictionary.
|
||||||
|
"""
|
||||||
|
keys = list_of_dicts[0].keys()
|
||||||
|
|
||||||
|
collated = {}
|
||||||
|
for key in keys:
|
||||||
|
collated[key] = [dct[key] for dct in list_of_dicts]
|
||||||
|
|
||||||
|
return {k: func(v) for k, v in collated.items()}
|
||||||
@ -0,0 +1,173 @@
|
|||||||
|
# serializer version: 1
|
||||||
|
# name: test_multiple_apsd
|
||||||
|
APSDResult(psd_dict={'dot1': array([4.86766031e-04, 4.19651623e-04, 4.80791366e-04, 4.20369610e-04,
|
||||||
|
4.68533341e-04, 5.16068351e-04, 4.78067393e-04, 4.96604551e-04,
|
||||||
|
5.43620111e-04, 6.05242857e-04, 4.69307958e-04, 4.26121386e-04,
|
||||||
|
4.21269821e-04, 5.09777399e-04, 4.48934052e-04, 5.02187345e-04,
|
||||||
|
3.53749504e-04, 4.09220073e-04, 4.62005256e-04, 4.12107505e-04,
|
||||||
|
3.56366504e-04, 4.26203963e-04, 3.67431248e-04, 3.81536291e-04,
|
||||||
|
3.36701409e-04, 3.92802665e-04, 3.15367650e-04, 3.42416435e-04,
|
||||||
|
3.27988067e-04, 4.26302183e-04, 3.28999025e-04, 4.29334822e-04,
|
||||||
|
3.15934312e-04, 3.32819336e-04, 3.15795846e-04, 2.87406555e-04,
|
||||||
|
2.84604109e-04, 3.09919565e-04, 3.32866212e-04, 2.81185901e-04,
|
||||||
|
2.88400397e-04, 2.82958462e-04, 3.28466378e-04, 3.07422736e-04,
|
||||||
|
2.49590746e-04, 2.94908312e-04, 2.82621957e-04, 2.33740364e-04,
|
||||||
|
2.46403654e-04, 1.97037760e-04, 2.08836420e-04, 2.24201149e-04,
|
||||||
|
1.82098361e-04, 2.18155711e-04, 1.74755943e-04, 2.09890378e-04,
|
||||||
|
1.88711703e-04, 2.07412229e-04, 1.98671240e-04, 1.84363253e-04,
|
||||||
|
2.05256522e-04, 2.06315266e-04, 1.64532047e-04, 1.73349546e-04,
|
||||||
|
1.72297423e-04, 1.75252719e-04, 1.88901743e-04, 1.43942892e-04,
|
||||||
|
1.46827918e-04, 1.88735091e-04, 1.61654033e-04, 1.60315092e-04,
|
||||||
|
1.61148083e-04, 1.72781033e-04, 1.29621040e-04, 1.62732256e-04,
|
||||||
|
1.34561124e-04, 1.37301084e-04, 1.31576063e-04, 1.23245524e-04,
|
||||||
|
1.12170523e-04, 1.34295550e-04, 1.32552039e-04, 1.30137162e-04,
|
||||||
|
1.16949156e-04, 1.21855826e-04, 1.20613195e-04, 1.08900703e-04,
|
||||||
|
1.08944336e-04, 1.00637489e-04, 9.69385008e-05, 9.81953856e-05,
|
||||||
|
9.37454981e-05, 1.02083189e-04, 1.04357064e-04, 1.14643829e-04,
|
||||||
|
8.47293232e-05, 9.76277680e-05, 8.70553104e-05, 9.81617075e-05,
|
||||||
|
9.42707241e-05, 9.09620939e-05, 8.54830181e-05, 7.88240418e-05,
|
||||||
|
8.02623727e-05, 9.62302595e-05, 7.67164243e-05, 7.77811015e-05,
|
||||||
|
8.09418323e-05, 8.97414027e-05, 8.28856237e-05, 9.98695313e-05,
|
||||||
|
7.92885646e-05, 7.36902018e-05, 6.69234483e-05, 8.34417266e-05,
|
||||||
|
7.57212816e-05, 8.36407870e-05, 6.85043662e-05, 7.76366871e-05,
|
||||||
|
6.65382834e-05, 7.83262386e-05, 7.77777187e-05, 6.38865728e-05,
|
||||||
|
7.01979215e-05, 5.84083856e-05, 5.84499229e-05, 6.93109353e-05,
|
||||||
|
6.23115544e-05, 6.21688095e-05, 5.88077275e-05, 5.64378115e-05,
|
||||||
|
5.38906182e-05, 5.06949007e-05, 5.77858340e-05, 5.06382029e-05,
|
||||||
|
6.12581644e-05, 5.85201545e-05, 6.70153542e-05, 5.43651586e-05,
|
||||||
|
5.93095077e-05, 4.97793629e-05, 5.04532539e-05, 4.64558263e-05,
|
||||||
|
5.00685914e-05, 4.64340493e-05, 4.00722922e-05, 5.29851315e-05,
|
||||||
|
4.72715285e-05, 4.00251792e-05, 5.05830989e-05, 4.66462949e-05,
|
||||||
|
4.54261653e-05, 4.15406970e-05, 4.28880715e-05, 4.25299775e-05,
|
||||||
|
4.71275321e-05, 3.53267154e-05, 4.69054813e-05, 3.98743647e-05,
|
||||||
|
3.93733595e-05, 4.30266071e-05, 4.00563784e-05, 3.27610427e-05,
|
||||||
|
5.52096846e-05, 3.81643705e-05, 4.81130598e-05, 3.90264084e-05,
|
||||||
|
3.56034599e-05, 4.28634083e-05, 3.44038682e-05, 3.08594418e-05,
|
||||||
|
3.58784095e-05, 3.94996071e-05, 3.85706443e-05, 3.84364529e-05,
|
||||||
|
3.87594258e-05, 3.78118168e-05, 3.42259058e-05, 2.89087952e-05,
|
||||||
|
3.86139742e-05, 3.22664575e-05, 3.83010064e-05, 3.35970363e-05,
|
||||||
|
3.86956149e-05, 3.18553739e-05, 3.22961324e-05, 3.40172167e-05,
|
||||||
|
2.74114986e-05, 2.80315264e-05, 3.26053291e-05, 3.23915065e-05,
|
||||||
|
3.28025471e-05, 3.56961879e-05, 3.45107814e-05, 2.82432276e-05,
|
||||||
|
3.43371896e-05, 3.45438511e-05, 3.26722196e-05, 2.82155606e-05,
|
||||||
|
2.73606156e-05, 2.93707709e-05, 2.75505951e-05, 2.98210317e-05,
|
||||||
|
2.54917374e-05, 2.84107273e-05, 3.18231714e-05, 2.95071039e-05,
|
||||||
|
3.18535647e-05, 3.15739596e-05, 2.76053090e-05, 2.91494162e-05,
|
||||||
|
2.56099179e-05, 2.65930935e-05, 2.38810622e-05, 2.91172657e-05,
|
||||||
|
2.16992155e-05, 2.65981317e-05, 3.24365767e-05, 2.99653858e-05,
|
||||||
|
2.87122629e-05, 1.94206470e-05, 2.31467057e-05, 2.40234823e-05,
|
||||||
|
2.51974170e-05, 2.71108626e-05, 2.42519040e-05, 2.45864200e-05,
|
||||||
|
2.45549894e-05, 2.23195114e-05, 2.73386853e-05, 2.29818631e-05,
|
||||||
|
2.46604987e-05, 1.78922896e-05, 2.23938891e-05, 2.20241111e-05,
|
||||||
|
2.42116112e-05, 1.98445402e-05, 2.55313734e-05, 2.26527774e-05,
|
||||||
|
2.15151192e-05, 2.18127800e-05, 2.27182456e-05, 2.27288252e-05,
|
||||||
|
2.05411209e-05, 2.21132029e-05, 2.31176274e-05, 2.40069380e-05,
|
||||||
|
2.00570407e-05, 2.06764835e-05, 2.37085332e-05, 2.38822380e-05,
|
||||||
|
2.00727524e-05, 1.96596342e-05, 1.80829028e-05, 2.13122335e-05,
|
||||||
|
2.60279773e-05, 1.64626717e-05, 1.95886373e-05, 1.96501023e-05,
|
||||||
|
1.91121814e-05, 1.98475805e-05, 2.06863308e-05, 2.10617681e-05,
|
||||||
|
1.60439016e-05, 1.76395788e-05, 1.99323695e-05, 1.83692094e-05,
|
||||||
|
1.97611430e-05, 1.98457141e-05, 1.66799987e-05, 1.56871247e-05,
|
||||||
|
1.88746203e-05, 1.94343725e-05, 2.09251343e-05, 1.37803586e-05,
|
||||||
|
1.74899716e-05, 2.14970801e-05, 1.70437092e-05, 1.65712128e-05,
|
||||||
|
1.61940567e-05, 1.55139733e-05, 1.48800183e-05, 1.69475723e-05,
|
||||||
|
1.84826509e-05, 1.72154277e-05, 1.67403929e-05, 1.66184048e-05,
|
||||||
|
1.62022596e-05, 1.71137293e-05, 1.73613401e-05, 1.66715457e-05,
|
||||||
|
1.39654603e-05, 1.42945028e-05, 1.79297700e-05, 1.78270846e-05,
|
||||||
|
1.48656092e-05, 1.49235828e-05, 1.42301779e-05, 1.69017128e-05,
|
||||||
|
1.57454438e-05, 1.44974699e-05, 1.79678483e-05, 1.48166741e-05,
|
||||||
|
1.45941033e-05, 1.58602922e-05, 1.74440242e-05, 1.30825272e-05,
|
||||||
|
1.57815641e-05, 1.50168930e-05, 1.32247492e-05, 1.55218213e-05,
|
||||||
|
1.78815665e-05, 1.58621193e-05, 1.65197160e-05, 1.49212113e-05,
|
||||||
|
1.55013627e-05, 1.30893348e-05, 1.79898659e-05, 1.36082372e-05,
|
||||||
|
1.60777842e-05, 1.56820070e-05, 1.32689759e-05, 1.58210523e-05,
|
||||||
|
1.16323267e-05, 1.41930602e-05, 1.71346939e-05, 1.53278626e-05,
|
||||||
|
1.40518275e-05, 1.44641006e-05, 1.58381902e-05, 1.30738062e-05,
|
||||||
|
1.42820023e-05, 1.47887897e-05, 1.33593630e-05, 1.46463583e-05,
|
||||||
|
1.63399307e-05, 1.31079056e-05, 1.58986831e-05, 1.47558230e-05,
|
||||||
|
1.34851861e-05, 1.49156413e-05, 1.55965672e-05, 1.29409930e-05,
|
||||||
|
1.58323920e-05, 1.43507448e-05, 1.30621022e-05, 1.28539727e-05,
|
||||||
|
1.29045750e-05, 1.25434580e-05, 1.27138037e-05, 1.22219609e-05,
|
||||||
|
1.26286359e-05, 1.49072315e-05, 1.33138715e-05, 1.51940285e-05,
|
||||||
|
1.39394318e-05, 1.26485831e-05, 1.49228303e-05, 1.37758807e-05,
|
||||||
|
1.23917368e-05, 1.49912685e-05, 1.25699601e-05, 1.25546148e-05,
|
||||||
|
1.33653887e-05, 1.17058408e-05, 1.13810167e-05, 1.33450603e-05,
|
||||||
|
1.02977623e-05, 1.32989389e-05, 1.22207557e-05, 1.22465087e-05,
|
||||||
|
1.25895540e-05, 1.12599253e-05, 1.20706307e-05, 1.50940507e-05,
|
||||||
|
1.12761645e-05, 1.53244907e-05, 1.24880988e-05, 1.14930089e-05,
|
||||||
|
1.33779611e-05, 1.31382086e-05, 1.29751531e-05, 1.24094261e-05,
|
||||||
|
1.38971827e-05, 1.35603590e-05, 1.54391840e-05, 1.28131324e-05,
|
||||||
|
1.42400197e-05, 1.11555210e-05, 1.32260372e-05, 1.22531409e-05,
|
||||||
|
1.14425592e-05, 1.19601092e-05, 1.25277452e-05, 1.23125375e-05,
|
||||||
|
1.22767606e-05, 1.04635218e-05, 1.06253061e-05, 8.93562763e-06,
|
||||||
|
1.00333160e-05, 1.29232136e-05, 1.17508220e-05, 1.16497936e-05,
|
||||||
|
1.23274931e-05, 1.33937448e-05, 1.04916813e-05, 1.09252020e-05,
|
||||||
|
1.25324965e-05, 1.14311896e-05, 1.16534623e-05, 1.21081248e-05,
|
||||||
|
1.43694806e-05, 1.08832957e-05, 1.21792077e-05, 1.34939104e-05,
|
||||||
|
1.09226237e-05, 1.27637764e-05, 1.19455633e-05, 1.08491595e-05,
|
||||||
|
1.18889125e-05, 1.22260635e-05, 1.28966891e-05, 1.03095455e-05,
|
||||||
|
1.21839688e-05, 1.20874536e-05, 1.07811837e-05, 1.19102602e-05,
|
||||||
|
1.06934673e-05, 1.19047703e-05, 1.12772337e-05, 1.07964023e-05,
|
||||||
|
1.27185235e-05, 1.01101305e-05, 1.11800717e-05, 1.29810055e-05,
|
||||||
|
1.25800228e-05, 9.16328966e-06, 1.11220365e-05, 1.11230675e-05,
|
||||||
|
9.90734807e-06, 1.06842896e-05, 1.22867699e-05, 1.18574811e-05,
|
||||||
|
1.04973417e-05, 1.16523798e-05, 1.02118388e-05, 1.14134887e-05,
|
||||||
|
1.24914260e-05, 1.01224262e-05, 9.32676591e-06, 1.22892870e-05,
|
||||||
|
1.20458389e-05, 1.09317898e-05, 1.11314350e-05, 1.02885804e-05,
|
||||||
|
1.21274068e-05, 1.10258616e-05, 1.18182068e-05, 9.98152358e-06,
|
||||||
|
1.30981584e-05, 9.76611068e-06, 9.92430559e-06, 1.05367166e-05,
|
||||||
|
1.05989768e-05, 1.02407643e-05, 1.28001975e-05, 1.20461317e-05,
|
||||||
|
1.17645226e-05, 1.36953981e-05, 1.11163052e-05, 1.17238577e-05,
|
||||||
|
8.81321352e-06, 1.00073061e-05, 1.00813125e-05, 1.17426962e-05,
|
||||||
|
1.04744073e-05, 1.07882076e-05, 1.15055060e-05, 1.04735993e-05,
|
||||||
|
1.20913342e-05, 1.07555318e-05, 1.05519506e-05, 9.90229131e-06,
|
||||||
|
1.30087569e-05, 1.02021731e-05, 1.12835177e-05, 1.09575431e-05,
|
||||||
|
1.06145280e-05, 1.14745656e-05, 1.04125376e-05, 9.73386195e-06,
|
||||||
|
9.98096535e-06, 9.70268684e-06, 1.01081360e-05, 1.01637012e-05,
|
||||||
|
9.12613069e-06, 1.03875577e-05, 1.16939553e-05, 1.23720659e-05])}, freqs=array([ 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1. , 1.1,
|
||||||
|
1.2, 1.3, 1.4, 1.5, 1.6, 1.7, 1.8, 1.9, 2. , 2.1, 2.2,
|
||||||
|
2.3, 2.4, 2.5, 2.6, 2.7, 2.8, 2.9, 3. , 3.1, 3.2, 3.3,
|
||||||
|
3.4, 3.5, 3.6, 3.7, 3.8, 3.9, 4. , 4.1, 4.2, 4.3, 4.4,
|
||||||
|
4.5, 4.6, 4.7, 4.8, 4.9, 5. , 5.1, 5.2, 5.3, 5.4, 5.5,
|
||||||
|
5.6, 5.7, 5.8, 5.9, 6. , 6.1, 6.2, 6.3, 6.4, 6.5, 6.6,
|
||||||
|
6.7, 6.8, 6.9, 7. , 7.1, 7.2, 7.3, 7.4, 7.5, 7.6, 7.7,
|
||||||
|
7.8, 7.9, 8. , 8.1, 8.2, 8.3, 8.4, 8.5, 8.6, 8.7, 8.8,
|
||||||
|
8.9, 9. , 9.1, 9.2, 9.3, 9.4, 9.5, 9.6, 9.7, 9.8, 9.9,
|
||||||
|
10. , 10.1, 10.2, 10.3, 10.4, 10.5, 10.6, 10.7, 10.8, 10.9, 11. ,
|
||||||
|
11.1, 11.2, 11.3, 11.4, 11.5, 11.6, 11.7, 11.8, 11.9, 12. , 12.1,
|
||||||
|
12.2, 12.3, 12.4, 12.5, 12.6, 12.7, 12.8, 12.9, 13. , 13.1, 13.2,
|
||||||
|
13.3, 13.4, 13.5, 13.6, 13.7, 13.8, 13.9, 14. , 14.1, 14.2, 14.3,
|
||||||
|
14.4, 14.5, 14.6, 14.7, 14.8, 14.9, 15. , 15.1, 15.2, 15.3, 15.4,
|
||||||
|
15.5, 15.6, 15.7, 15.8, 15.9, 16. , 16.1, 16.2, 16.3, 16.4, 16.5,
|
||||||
|
16.6, 16.7, 16.8, 16.9, 17. , 17.1, 17.2, 17.3, 17.4, 17.5, 17.6,
|
||||||
|
17.7, 17.8, 17.9, 18. , 18.1, 18.2, 18.3, 18.4, 18.5, 18.6, 18.7,
|
||||||
|
18.8, 18.9, 19. , 19.1, 19.2, 19.3, 19.4, 19.5, 19.6, 19.7, 19.8,
|
||||||
|
19.9, 20. , 20.1, 20.2, 20.3, 20.4, 20.5, 20.6, 20.7, 20.8, 20.9,
|
||||||
|
21. , 21.1, 21.2, 21.3, 21.4, 21.5, 21.6, 21.7, 21.8, 21.9, 22. ,
|
||||||
|
22.1, 22.2, 22.3, 22.4, 22.5, 22.6, 22.7, 22.8, 22.9, 23. , 23.1,
|
||||||
|
23.2, 23.3, 23.4, 23.5, 23.6, 23.7, 23.8, 23.9, 24. , 24.1, 24.2,
|
||||||
|
24.3, 24.4, 24.5, 24.6, 24.7, 24.8, 24.9, 25. , 25.1, 25.2, 25.3,
|
||||||
|
25.4, 25.5, 25.6, 25.7, 25.8, 25.9, 26. , 26.1, 26.2, 26.3, 26.4,
|
||||||
|
26.5, 26.6, 26.7, 26.8, 26.9, 27. , 27.1, 27.2, 27.3, 27.4, 27.5,
|
||||||
|
27.6, 27.7, 27.8, 27.9, 28. , 28.1, 28.2, 28.3, 28.4, 28.5, 28.6,
|
||||||
|
28.7, 28.8, 28.9, 29. , 29.1, 29.2, 29.3, 29.4, 29.5, 29.6, 29.7,
|
||||||
|
29.8, 29.9, 30. , 30.1, 30.2, 30.3, 30.4, 30.5, 30.6, 30.7, 30.8,
|
||||||
|
30.9, 31. , 31.1, 31.2, 31.3, 31.4, 31.5, 31.6, 31.7, 31.8, 31.9,
|
||||||
|
32. , 32.1, 32.2, 32.3, 32.4, 32.5, 32.6, 32.7, 32.8, 32.9, 33. ,
|
||||||
|
33.1, 33.2, 33.3, 33.4, 33.5, 33.6, 33.7, 33.8, 33.9, 34. , 34.1,
|
||||||
|
34.2, 34.3, 34.4, 34.5, 34.6, 34.7, 34.8, 34.9, 35. , 35.1, 35.2,
|
||||||
|
35.3, 35.4, 35.5, 35.6, 35.7, 35.8, 35.9, 36. , 36.1, 36.2, 36.3,
|
||||||
|
36.4, 36.5, 36.6, 36.7, 36.8, 36.9, 37. , 37.1, 37.2, 37.3, 37.4,
|
||||||
|
37.5, 37.6, 37.7, 37.8, 37.9, 38. , 38.1, 38.2, 38.3, 38.4, 38.5,
|
||||||
|
38.6, 38.7, 38.8, 38.9, 39. , 39.1, 39.2, 39.3, 39.4, 39.5, 39.6,
|
||||||
|
39.7, 39.8, 39.9, 40. , 40.1, 40.2, 40.3, 40.4, 40.5, 40.6, 40.7,
|
||||||
|
40.8, 40.9, 41. , 41.1, 41.2, 41.3, 41.4, 41.5, 41.6, 41.7, 41.8,
|
||||||
|
41.9, 42. , 42.1, 42.2, 42.3, 42.4, 42.5, 42.6, 42.7, 42.8, 42.9,
|
||||||
|
43. , 43.1, 43.2, 43.3, 43.4, 43.5, 43.6, 43.7, 43.8, 43.9, 44. ,
|
||||||
|
44.1, 44.2, 44.3, 44.4, 44.5, 44.6, 44.7, 44.8, 44.9, 45. , 45.1,
|
||||||
|
45.2, 45.3, 45.4, 45.5, 45.6, 45.7, 45.8, 45.9, 46. , 46.1, 46.2,
|
||||||
|
46.3, 46.4, 46.5, 46.6, 46.7, 46.8, 46.9, 47. , 47.1, 47.2, 47.3,
|
||||||
|
47.4, 47.5, 47.6, 47.7, 47.8, 47.9, 48. , 48.1, 48.2, 48.3, 48.4,
|
||||||
|
48.5, 48.6, 48.7, 48.8, 48.9, 49. , 49.1, 49.2, 49.3, 49.4, 49.5,
|
||||||
|
49.6, 49.7, 49.8, 49.9, 50. ]))
|
||||||
|
# ---
|
||||||
@ -0,0 +1,141 @@
|
|||||||
|
# serializer version: 1
|
||||||
|
# name: test_single_dipole_time_series_apsd_new_series
|
||||||
|
APSDResult(psd_dict={'dot1': array([7.15705084e-05, 4.78051535e-04, 8.55068290e-05, 7.55543647e-04,
|
||||||
|
5.38624754e-04, 3.49785878e-04, 3.15360383e-03, 2.25042289e-04,
|
||||||
|
2.25514864e-05, 2.77820792e-04, 1.57875639e-03, 1.44404306e-03])}, freqs=array([0.4, 0.8, 1.2, 1.6, 2. , 2.4, 2.8, 3.2, 3.6, 4. , 4.4, 4.8]))
|
||||||
|
# ---
|
||||||
|
# name: test_single_dipole_time_series_psd
|
||||||
|
list([
|
||||||
|
list([
|
||||||
|
0j,
|
||||||
|
(1.0245563439837635+0j),
|
||||||
|
]),
|
||||||
|
list([
|
||||||
|
(0.4+0j),
|
||||||
|
(-0.06997893548071246-0.04298727993617441j),
|
||||||
|
]),
|
||||||
|
list([
|
||||||
|
(0.8+0j),
|
||||||
|
(-0.039450404619039356-0.2278123801105552j),
|
||||||
|
]),
|
||||||
|
list([
|
||||||
|
(1.2000000000000002+0j),
|
||||||
|
(0.06879271540452367+0.18112796265494657j),
|
||||||
|
]),
|
||||||
|
list([
|
||||||
|
(1.6+0j),
|
||||||
|
(0.08886792286401579-0.13758873132600968j),
|
||||||
|
]),
|
||||||
|
list([
|
||||||
|
(2+0j),
|
||||||
|
(0.11176347216411753-0.08120091560477474j),
|
||||||
|
]),
|
||||||
|
list([
|
||||||
|
(2.4000000000000004+0j),
|
||||||
|
(0.2691233202491688-0.010346527463152275j),
|
||||||
|
]),
|
||||||
|
list([
|
||||||
|
(2.8000000000000003+0j),
|
||||||
|
(0.2576145129611034-0.3358195628624988j),
|
||||||
|
]),
|
||||||
|
list([
|
||||||
|
(3.2+0j),
|
||||||
|
(-0.0710407429754301+0.08850346984926659j),
|
||||||
|
]),
|
||||||
|
list([
|
||||||
|
(3.6+0j),
|
||||||
|
(-0.22596960267390404+0.010400520410955974j),
|
||||||
|
]),
|
||||||
|
list([
|
||||||
|
(4+0j),
|
||||||
|
(0.016306070833852868-0.05018492576136252j),
|
||||||
|
]),
|
||||||
|
list([
|
||||||
|
(4.4+0j),
|
||||||
|
(0.1514065333713827-0.22893435273278817j),
|
||||||
|
]),
|
||||||
|
list([
|
||||||
|
(4.800000000000001+0j),
|
||||||
|
(-0.0024668424412067763+0.17633416864658075j),
|
||||||
|
]),
|
||||||
|
])
|
||||||
|
# ---
|
||||||
|
# name: test_single_dipole_time_series_psd_new_series
|
||||||
|
TimeSeriesResult(series_dict={'dot1': array([-0.0853797, 0.0853797, -0.0853797, -0.0853797, -0.0853797,
|
||||||
|
0.0853797, 0.0853797, -0.0853797, 0.0853797, -0.0853797,
|
||||||
|
0.0853797, -0.0853797, 0.0853797, -0.0853797, -0.0853797,
|
||||||
|
-0.0853797, 0.0853797, -0.0853797, -0.0853797, 0.0853797,
|
||||||
|
0.0853797, -0.0853797, -0.0853797, 0.0853797, 0.0853797])}, num_points=25, delta_t=0.1)
|
||||||
|
# ---
|
||||||
|
# name: test_two_dipole_time_series_apsd_new_series
|
||||||
|
APSDResult(psd_dict={'dot1': array([2.01355098e-02, 2.28728894e-02, 2.68887738e-02, 8.57694875e-02,
|
||||||
|
3.32025308e-02, 1.09906295e-02, 6.20193451e-03, 9.25526310e-03,
|
||||||
|
2.54435863e-02, 2.06494403e-02, 2.30801063e-02, 1.46212821e-04,
|
||||||
|
4.51507207e-03, 9.79219359e-06, 1.78387754e-03, 1.87645428e-05,
|
||||||
|
8.70961095e-03, 2.10590946e-03, 7.94522558e-05, 3.28144577e-03,
|
||||||
|
3.75758394e-03, 3.23552287e-04, 1.16987461e-03, 3.83447107e-04,
|
||||||
|
1.28792075e-03, 4.36287386e-04, 2.63344922e-03, 2.10987153e-03,
|
||||||
|
4.14966325e-04, 3.49214983e-03, 1.34022421e-03, 1.44743780e-03,
|
||||||
|
9.21715166e-04, 1.44488261e-03, 6.16100652e-04, 1.25148127e-04,
|
||||||
|
6.18299895e-04, 7.29981733e-04, 6.64796172e-04, 5.59949875e-04,
|
||||||
|
1.96952019e-03, 4.03586437e-03, 4.91238661e-04, 2.88674291e-04,
|
||||||
|
4.27864253e-04, 6.85945083e-04, 4.09499937e-04, 2.16045557e-04,
|
||||||
|
1.81996683e-03, 5.83175390e-05]), 'dot2': array([7.84223244e-02, 7.80630078e-02, 9.28654994e-02, 2.86646937e-01,
|
||||||
|
1.10740152e-01, 3.57914979e-02, 1.86349599e-02, 3.21123045e-02,
|
||||||
|
8.11466737e-02, 7.00066392e-02, 7.70905805e-02, 2.97613466e-04,
|
||||||
|
1.54886069e-02, 9.46435334e-05, 6.57687748e-03, 8.40879567e-05,
|
||||||
|
2.82600269e-02, 6.98437415e-03, 2.33068759e-04, 1.00329309e-02,
|
||||||
|
1.23323571e-02, 8.53996730e-04, 3.79517235e-03, 1.67958030e-03,
|
||||||
|
3.77088192e-03, 9.61087507e-04, 7.96408605e-03, 6.92973251e-03,
|
||||||
|
2.10466442e-03, 1.08662404e-02, 4.45462316e-03, 4.88388517e-03,
|
||||||
|
3.31628201e-03, 4.28733726e-03, 1.54410149e-03, 3.58990199e-04,
|
||||||
|
1.37659695e-03, 2.36733015e-03, 1.82352870e-03, 1.41013414e-03,
|
||||||
|
5.34997843e-03, 1.16237095e-02, 1.55421093e-03, 1.02900200e-03,
|
||||||
|
1.77722903e-03, 1.64197100e-03, 1.64957728e-03, 9.09951112e-04,
|
||||||
|
6.46482017e-03, 1.31206692e-04])}, freqs=array([ 0.2, 0.4, 0.6, 0.8, 1. , 1.2, 1.4, 1.6, 1.8, 2. , 2.2,
|
||||||
|
2.4, 2.6, 2.8, 3. , 3.2, 3.4, 3.6, 3.8, 4. , 4.2, 4.4,
|
||||||
|
4.6, 4.8, 5. , 5.2, 5.4, 5.6, 5.8, 6. , 6.2, 6.4, 6.6,
|
||||||
|
6.8, 7. , 7.2, 7.4, 7.6, 7.8, 8. , 8.2, 8.4, 8.6, 8.8,
|
||||||
|
9. , 9.2, 9.4, 9.6, 9.8, 10. ]))
|
||||||
|
# ---
|
||||||
|
# name: test_two_dipole_time_series_psd_new_series
|
||||||
|
TimeSeriesResult(series_dict={'dot1': array([ 0.28499068, 0.45575007, 0.45575007, 0.45575007, 0.45575007,
|
||||||
|
-0.28499068, -0.45575007, -0.45575007, -0.45575007, -0.28499068,
|
||||||
|
-0.45575007, -0.45575007, -0.28499068, -0.45575007, -0.28499068,
|
||||||
|
-0.45575007, -0.28499068, -0.28499068, -0.45575007, 0.28499068,
|
||||||
|
0.45575007, 0.28499068, 0.28499068, 0.45575007, 0.28499068,
|
||||||
|
0.28499068, 0.45575007, 0.45575007, 0.28499068, 0.28499068,
|
||||||
|
0.28499068, 0.28499068, 0.28499068, 0.28499068, -0.28499068,
|
||||||
|
-0.45575007, -0.45575007, 0.28499068, 0.28499068, 0.28499068,
|
||||||
|
0.28499068, 0.28499068, 0.28499068, 0.28499068, 0.28499068,
|
||||||
|
0.28499068, 0.28499068, 0.45575007, 0.28499068, 0.45575007,
|
||||||
|
-0.28499068, -0.28499068, -0.45575007, -0.45575007, -0.45575007,
|
||||||
|
-0.45575007, -0.28499068, -0.28499068, -0.28499068, -0.45575007,
|
||||||
|
-0.45575007, -0.28499068, -0.45575007, -0.28499068, -0.28499068,
|
||||||
|
-0.45575007, 0.45575007, 0.28499068, 0.45575007, 0.45575007,
|
||||||
|
0.28499068, 0.45575007, 0.45575007, 0.45575007, -0.28499068,
|
||||||
|
-0.28499068, -0.28499068, -0.28499068, -0.28499068, -0.45575007,
|
||||||
|
-0.28499068, -0.28499068, -0.28499068, -0.45575007, -0.45575007,
|
||||||
|
-0.28499068, -0.28499068, -0.28499068, -0.28499068, -0.28499068,
|
||||||
|
0.45575007, 0.45575007, 0.45575007, 0.45575007, 0.28499068,
|
||||||
|
-0.45575007, -0.45575007, -0.28499068, -0.45575007, -0.28499068]), 'dot2': array([ 0.55234807, 0.80847957, 0.80847957, 0.80847957, 0.80847957,
|
||||||
|
-0.55234807, -0.80847957, -0.80847957, -0.80847957, -0.55234807,
|
||||||
|
-0.80847957, -0.80847957, -0.55234807, -0.80847957, -0.55234807,
|
||||||
|
-0.80847957, -0.55234807, -0.55234807, -0.80847957, 0.55234807,
|
||||||
|
0.80847957, 0.55234807, 0.55234807, 0.80847957, 0.55234807,
|
||||||
|
0.55234807, 0.80847957, 0.80847957, 0.55234807, 0.55234807,
|
||||||
|
0.55234807, 0.55234807, 0.55234807, 0.55234807, -0.55234807,
|
||||||
|
-0.80847957, -0.80847957, 0.55234807, 0.55234807, 0.55234807,
|
||||||
|
0.55234807, 0.55234807, 0.55234807, 0.55234807, 0.55234807,
|
||||||
|
0.55234807, 0.55234807, 0.80847957, 0.55234807, 0.80847957,
|
||||||
|
-0.55234807, -0.55234807, -0.80847957, -0.80847957, -0.80847957,
|
||||||
|
-0.80847957, -0.55234807, -0.55234807, -0.55234807, -0.80847957,
|
||||||
|
-0.80847957, -0.55234807, -0.80847957, -0.55234807, -0.55234807,
|
||||||
|
-0.80847957, 0.80847957, 0.55234807, 0.80847957, 0.80847957,
|
||||||
|
0.55234807, 0.80847957, 0.80847957, 0.80847957, -0.55234807,
|
||||||
|
-0.55234807, -0.55234807, -0.55234807, -0.55234807, -0.80847957,
|
||||||
|
-0.55234807, -0.55234807, -0.55234807, -0.80847957, -0.80847957,
|
||||||
|
-0.55234807, -0.55234807, -0.55234807, -0.55234807, -0.55234807,
|
||||||
|
0.80847957, 0.80847957, 0.80847957, 0.80847957, 0.55234807,
|
||||||
|
-0.80847957, -0.80847957, -0.55234807, -0.80847957, -0.55234807])}, num_points=100, delta_t=0.05)
|
||||||
|
# ---
|
||||||
36
tests/dipoles/time_series_psd/test_apsd_methods.py
Normal file
36
tests/dipoles/time_series_psd/test_apsd_methods.py
Normal file
@ -0,0 +1,36 @@
|
|||||||
|
import pytest
|
||||||
|
import numpy
|
||||||
|
import tantri.dipoles.time_series
|
||||||
|
|
||||||
|
|
||||||
|
def test_apsd_merge():
|
||||||
|
freqs = numpy.array([0.01, 0.1, 1, 10, 100])
|
||||||
|
dict1 = {"t1": numpy.array([1, 2, 3, 4, 5])}
|
||||||
|
a1 = tantri.dipoles.time_series.APSDResult(dict1, freqs)
|
||||||
|
|
||||||
|
dict2 = {"t1": numpy.array([3, 4, 5, 6, 7])}
|
||||||
|
a2 = tantri.dipoles.time_series.APSDResult(dict2, freqs)
|
||||||
|
|
||||||
|
merged = tantri.dipoles.time_series.average_apsds([a1, a2])
|
||||||
|
|
||||||
|
expected = tantri.dipoles.time_series.APSDResult(
|
||||||
|
psd_dict={
|
||||||
|
"t1": numpy.array([2, 3, 4, 5, 6]),
|
||||||
|
},
|
||||||
|
freqs=freqs,
|
||||||
|
)
|
||||||
|
|
||||||
|
numpy.testing.assert_equal(merged.freqs, expected.freqs)
|
||||||
|
numpy.testing.assert_equal(merged.psd_dict, expected.psd_dict)
|
||||||
|
|
||||||
|
|
||||||
|
def test_apsd_merge_mismatch_freqs():
|
||||||
|
dict = {"t1": numpy.array([1, 2, 3, 4, 5])}
|
||||||
|
freqs1 = numpy.array([0.01, 0.1, 1, 10, 100])
|
||||||
|
a1 = tantri.dipoles.time_series.APSDResult(dict, freqs1)
|
||||||
|
|
||||||
|
freqs2 = numpy.array([1, 3, 5, 7, 9])
|
||||||
|
a2 = tantri.dipoles.time_series.APSDResult(dict, freqs2)
|
||||||
|
|
||||||
|
with pytest.raises(ValueError):
|
||||||
|
tantri.dipoles.time_series.average_apsds([a1, a2])
|
||||||
83
tests/dipoles/time_series_psd/test_psd_analytic_compare.py
Normal file
83
tests/dipoles/time_series_psd/test_psd_analytic_compare.py
Normal file
@ -0,0 +1,83 @@
|
|||||||
|
import numpy
|
||||||
|
|
||||||
|
from tantri.dipoles import (
|
||||||
|
DipoleTO,
|
||||||
|
DipoleTimeSeries,
|
||||||
|
DotPosition,
|
||||||
|
DipoleMeasurementType,
|
||||||
|
)
|
||||||
|
|
||||||
|
import logging
|
||||||
|
|
||||||
|
_logger = logging.getLogger(__name__)
|
||||||
|
|
||||||
|
|
||||||
|
def test_multiple_apsd(snapshot):
|
||||||
|
|
||||||
|
dot_name = "dot1"
|
||||||
|
num_series_to_create = 100
|
||||||
|
num_ts_points = 1000
|
||||||
|
delta_t = 0.01
|
||||||
|
|
||||||
|
rng = numpy.random.default_rng(1234)
|
||||||
|
dots = [DotPosition(numpy.array([0, 0, 0]), dot_name)]
|
||||||
|
|
||||||
|
d1 = DipoleTO(
|
||||||
|
numpy.array([0, 0, 10]),
|
||||||
|
numpy.array([5, 3, 2]),
|
||||||
|
15,
|
||||||
|
)
|
||||||
|
|
||||||
|
ts_gen = DipoleTimeSeries(
|
||||||
|
[d1],
|
||||||
|
dots,
|
||||||
|
DipoleMeasurementType.ELECTRIC_POTENTIAL,
|
||||||
|
delta_t,
|
||||||
|
rng_to_use=rng,
|
||||||
|
)
|
||||||
|
|
||||||
|
estimated_psd = ts_gen.generate_average_apsd(num_series_to_create, num_ts_points)
|
||||||
|
|
||||||
|
assert estimated_psd == snapshot
|
||||||
|
|
||||||
|
|
||||||
|
def test_multiple_apsd_compare_analytic():
|
||||||
|
|
||||||
|
dot_name = "dot1"
|
||||||
|
num_series_to_create = 500
|
||||||
|
num_ts_points = 10000
|
||||||
|
delta_t = 0.001
|
||||||
|
|
||||||
|
rng = numpy.random.default_rng(1234)
|
||||||
|
dots = [DotPosition(numpy.array([0, 0, 0]), dot_name)]
|
||||||
|
|
||||||
|
d1 = DipoleTO(
|
||||||
|
numpy.array([0, 0, 10]),
|
||||||
|
numpy.array([5, 3, 2]),
|
||||||
|
15,
|
||||||
|
)
|
||||||
|
|
||||||
|
ts_gen = DipoleTimeSeries(
|
||||||
|
[d1],
|
||||||
|
dots,
|
||||||
|
DipoleMeasurementType.ELECTRIC_POTENTIAL,
|
||||||
|
delta_t,
|
||||||
|
rng_to_use=rng,
|
||||||
|
)
|
||||||
|
|
||||||
|
def s_analytic_potential(f: float, dot: DotPosition, dipole: DipoleTO):
|
||||||
|
g = dipole.w / ((numpy.pi * f) ** 2 + dipole.w**2)
|
||||||
|
rdiff = dot.r - dipole.s
|
||||||
|
|
||||||
|
return g * ((rdiff.dot(dipole.p) / (numpy.linalg.norm(rdiff) ** 3)) ** 2)
|
||||||
|
|
||||||
|
estimated_psd = ts_gen.generate_average_apsd(num_series_to_create, num_ts_points)
|
||||||
|
|
||||||
|
_logger.warning(estimated_psd)
|
||||||
|
|
||||||
|
analytic = numpy.array(
|
||||||
|
[s_analytic_potential(f, dots[0], d1) for f in estimated_psd.freqs]
|
||||||
|
)
|
||||||
|
numpy.testing.assert_allclose(
|
||||||
|
estimated_psd.psd_dict[dot_name], analytic, rtol=1.5, atol=1e-7
|
||||||
|
)
|
||||||
176
tests/dipoles/time_series_psd/test_time_series.py
Normal file
176
tests/dipoles/time_series_psd/test_time_series.py
Normal file
@ -0,0 +1,176 @@
|
|||||||
|
"""
|
||||||
|
Testing whatever related to APSDs who gives a shit.
|
||||||
|
"""
|
||||||
|
|
||||||
|
from tantri.dipoles import (
|
||||||
|
DipoleTO,
|
||||||
|
DipoleTimeSeries,
|
||||||
|
DotPosition,
|
||||||
|
DipoleMeasurementType,
|
||||||
|
)
|
||||||
|
import numpy
|
||||||
|
|
||||||
|
import scipy.fft
|
||||||
|
|
||||||
|
|
||||||
|
def test_single_dipole_time_series_psd(snapshot):
|
||||||
|
|
||||||
|
dot_name = "dot1"
|
||||||
|
num_points = 25
|
||||||
|
delta_t = 0.1
|
||||||
|
|
||||||
|
rng = numpy.random.default_rng(1234)
|
||||||
|
dots = [DotPosition(numpy.array([0, 0, 0]), dot_name)]
|
||||||
|
|
||||||
|
d1 = DipoleTO(
|
||||||
|
numpy.array([0, 0, 10]),
|
||||||
|
numpy.array([5, 3, 2]),
|
||||||
|
15,
|
||||||
|
)
|
||||||
|
|
||||||
|
ts_gen = DipoleTimeSeries(
|
||||||
|
[d1],
|
||||||
|
dots,
|
||||||
|
DipoleMeasurementType.ELECTRIC_POTENTIAL,
|
||||||
|
delta_t,
|
||||||
|
rng_to_use=rng,
|
||||||
|
)
|
||||||
|
# time_series = [ts_gen.transition() for i in range(25)]
|
||||||
|
|
||||||
|
# time_series_to_dict_list = {k: [dic[k] for dic in time_series] for k in time_series[0]}
|
||||||
|
# fft_dict = {k: scipy.fft.rfft(v) for k, v in time_series_to_dict_list.items()}
|
||||||
|
time_series = [ts_gen.transition()[dot_name] for i in range(num_points)]
|
||||||
|
fft = scipy.fft.rfft(time_series)
|
||||||
|
freqs = scipy.fft.rfftfreq(num_points, delta_t)
|
||||||
|
|
||||||
|
result = numpy.array([freqs, fft]).transpose().tolist()
|
||||||
|
|
||||||
|
assert result == snapshot
|
||||||
|
|
||||||
|
|
||||||
|
def test_single_dipole_time_series_psd_new_series(snapshot):
|
||||||
|
|
||||||
|
dot_name = "dot1"
|
||||||
|
num_points = 25
|
||||||
|
delta_t = 0.1
|
||||||
|
|
||||||
|
rng = numpy.random.default_rng(1234)
|
||||||
|
dots = [DotPosition(numpy.array([0, 0, 0]), dot_name)]
|
||||||
|
|
||||||
|
d1 = DipoleTO(
|
||||||
|
numpy.array([0, 0, 10]),
|
||||||
|
numpy.array([5, 3, 2]),
|
||||||
|
15,
|
||||||
|
)
|
||||||
|
|
||||||
|
ts_gen = DipoleTimeSeries(
|
||||||
|
[d1],
|
||||||
|
dots,
|
||||||
|
DipoleMeasurementType.ELECTRIC_POTENTIAL,
|
||||||
|
delta_t,
|
||||||
|
rng_to_use=rng,
|
||||||
|
)
|
||||||
|
|
||||||
|
result = ts_gen.generate_series(num_points)
|
||||||
|
|
||||||
|
assert result == snapshot
|
||||||
|
|
||||||
|
|
||||||
|
def test_two_dipole_time_series_psd_new_series(snapshot):
|
||||||
|
|
||||||
|
num_points = 100
|
||||||
|
delta_t = 0.05
|
||||||
|
|
||||||
|
rng = numpy.random.default_rng(1234)
|
||||||
|
dots = [
|
||||||
|
DotPosition(numpy.array([0, 0, 0]), "dot1"),
|
||||||
|
DotPosition(numpy.array([1, 0, 0]), "dot2"),
|
||||||
|
]
|
||||||
|
|
||||||
|
d1 = DipoleTO(
|
||||||
|
numpy.array([0, 0, 10]),
|
||||||
|
numpy.array([5, 3, 2]),
|
||||||
|
15,
|
||||||
|
)
|
||||||
|
|
||||||
|
d2 = DipoleTO(
|
||||||
|
numpy.array([0, 0, 10]),
|
||||||
|
numpy.array([2, 2, 1]),
|
||||||
|
2,
|
||||||
|
)
|
||||||
|
|
||||||
|
ts_gen = DipoleTimeSeries(
|
||||||
|
[d1, d2],
|
||||||
|
dots,
|
||||||
|
DipoleMeasurementType.ELECTRIC_POTENTIAL,
|
||||||
|
delta_t,
|
||||||
|
rng_to_use=rng,
|
||||||
|
)
|
||||||
|
|
||||||
|
result = ts_gen.generate_series(num_points)
|
||||||
|
|
||||||
|
assert result == snapshot
|
||||||
|
|
||||||
|
|
||||||
|
def test_single_dipole_time_series_apsd_new_series(snapshot):
|
||||||
|
|
||||||
|
dot_name = "dot1"
|
||||||
|
num_points = 25
|
||||||
|
delta_t = 0.1
|
||||||
|
|
||||||
|
rng = numpy.random.default_rng(1234)
|
||||||
|
dots = [DotPosition(numpy.array([0, 0, 0]), dot_name)]
|
||||||
|
|
||||||
|
d1 = DipoleTO(
|
||||||
|
numpy.array([0, 0, 10]),
|
||||||
|
numpy.array([5, 3, 2]),
|
||||||
|
15,
|
||||||
|
)
|
||||||
|
|
||||||
|
ts_gen = DipoleTimeSeries(
|
||||||
|
[d1],
|
||||||
|
dots,
|
||||||
|
DipoleMeasurementType.ELECTRIC_POTENTIAL,
|
||||||
|
delta_t,
|
||||||
|
rng_to_use=rng,
|
||||||
|
)
|
||||||
|
|
||||||
|
result = ts_gen.generate_series(num_points).get_apsds()
|
||||||
|
|
||||||
|
assert result == snapshot
|
||||||
|
|
||||||
|
|
||||||
|
def test_two_dipole_time_series_apsd_new_series(snapshot):
|
||||||
|
|
||||||
|
num_points = 100
|
||||||
|
delta_t = 0.05
|
||||||
|
|
||||||
|
rng = numpy.random.default_rng(1234)
|
||||||
|
dots = [
|
||||||
|
DotPosition(numpy.array([0, 0, 0]), "dot1"),
|
||||||
|
DotPosition(numpy.array([1, 0, 0]), "dot2"),
|
||||||
|
]
|
||||||
|
|
||||||
|
d1 = DipoleTO(
|
||||||
|
numpy.array([0, 0, 10]),
|
||||||
|
numpy.array([5, 3, 2]),
|
||||||
|
15,
|
||||||
|
)
|
||||||
|
|
||||||
|
d2 = DipoleTO(
|
||||||
|
numpy.array([0, 0, 10]),
|
||||||
|
numpy.array([2, 2, 1]),
|
||||||
|
2,
|
||||||
|
)
|
||||||
|
|
||||||
|
ts_gen = DipoleTimeSeries(
|
||||||
|
[d1, d2],
|
||||||
|
dots,
|
||||||
|
DipoleMeasurementType.ELECTRIC_POTENTIAL,
|
||||||
|
delta_t,
|
||||||
|
rng_to_use=rng,
|
||||||
|
)
|
||||||
|
|
||||||
|
result = ts_gen.generate_series(num_points).get_apsds()
|
||||||
|
|
||||||
|
assert result == snapshot
|
||||||
28
tests/test_util.py
Normal file
28
tests/test_util.py
Normal file
@ -0,0 +1,28 @@
|
|||||||
|
import typing
|
||||||
|
import tantri.util
|
||||||
|
import numpy
|
||||||
|
|
||||||
|
|
||||||
|
def test_mean_dict():
|
||||||
|
dict1 = {
|
||||||
|
"squares": numpy.array([1, 4, 9, 16]),
|
||||||
|
"linear": numpy.array([1, 2, 3, 4, 5]),
|
||||||
|
}
|
||||||
|
dict2 = {
|
||||||
|
"squares": numpy.array([2, 8, 18, 32]),
|
||||||
|
"linear": numpy.array([2, 4, 6, 8, 10]),
|
||||||
|
}
|
||||||
|
|
||||||
|
def mean(list_of_arrays: typing.Sequence[numpy.ndarray]) -> numpy.ndarray:
|
||||||
|
return numpy.mean(numpy.array(list_of_arrays), axis=0)
|
||||||
|
|
||||||
|
result = tantri.util.dict_reduce([dict1, dict2], mean)
|
||||||
|
|
||||||
|
expected = {
|
||||||
|
"squares": 1.5 * numpy.array([1, 4, 9, 16]),
|
||||||
|
"linear": 1.5 * numpy.array([1, 2, 3, 4, 5]),
|
||||||
|
}
|
||||||
|
|
||||||
|
numpy.testing.assert_equal(
|
||||||
|
result, expected, "The reduced dictionary should have matched the expected"
|
||||||
|
)
|
||||||
Loading…
x
Reference in New Issue
Block a user