diff --git a/tantri/dipoles/time_series.py b/tantri/dipoles/time_series.py index 32fa994..79dcfbb 100644 --- a/tantri/dipoles/time_series.py +++ b/tantri/dipoles/time_series.py @@ -4,6 +4,56 @@ import numpy.random import typing from tantri.dipoles.types import DipoleTO, DotPosition, DipoleMeasurementType 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 @@ -74,6 +124,27 @@ class WrappedDipole: self.flip_state() 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): self.state *= -1 @@ -113,6 +184,8 @@ class DipoleTimeSeries: self.dipoles = get_wrapped_dipoles(dipoles, dots, measurement_type) self.state = 0 + self.real_delta_t = dt + # we may need to supersample, because of how dumb this process is. # let's find our highest frequency max_frequency = max(d.w for d in self.dipoles) @@ -135,3 +208,53 @@ class DipoleTimeSeries: def transition(self) -> typing.Dict[str, float]: 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) diff --git a/tantri/util.py b/tantri/util.py new file mode 100644 index 0000000..01c62ef --- /dev/null +++ b/tantri/util.py @@ -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()} diff --git a/tests/dipoles/time_series_psd/__snapshots__/test_psd_analytic_compare.ambr b/tests/dipoles/time_series_psd/__snapshots__/test_psd_analytic_compare.ambr new file mode 100644 index 0000000..dcbc1b4 --- /dev/null +++ b/tests/dipoles/time_series_psd/__snapshots__/test_psd_analytic_compare.ambr @@ -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. ])) +# --- diff --git a/tests/dipoles/time_series_psd/__snapshots__/test_time_series.ambr b/tests/dipoles/time_series_psd/__snapshots__/test_time_series.ambr new file mode 100644 index 0000000..6b238f8 --- /dev/null +++ b/tests/dipoles/time_series_psd/__snapshots__/test_time_series.ambr @@ -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) +# --- diff --git a/tests/dipoles/time_series_psd/test_apsd_methods.py b/tests/dipoles/time_series_psd/test_apsd_methods.py new file mode 100644 index 0000000..5482de3 --- /dev/null +++ b/tests/dipoles/time_series_psd/test_apsd_methods.py @@ -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]) diff --git a/tests/dipoles/time_series_psd/test_psd_analytic_compare.py b/tests/dipoles/time_series_psd/test_psd_analytic_compare.py new file mode 100644 index 0000000..3498a5c --- /dev/null +++ b/tests/dipoles/time_series_psd/test_psd_analytic_compare.py @@ -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 + ) diff --git a/tests/dipoles/time_series_psd/test_time_series.py b/tests/dipoles/time_series_psd/test_time_series.py new file mode 100644 index 0000000..eda75e9 --- /dev/null +++ b/tests/dipoles/time_series_psd/test_time_series.py @@ -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 diff --git a/tests/test_util.py b/tests/test_util.py new file mode 100644 index 0000000..744659f --- /dev/null +++ b/tests/test_util.py @@ -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" + )