Coverage for Users/jsd/Library/CloudStorage/OneDrive-SimonFraserUniversity(1sfu)/projects/thztools/src/thztools/thztools.py: 59%
446 statements
« prev ^ index » next coverage.py v7.2.7, created at 2023-06-12 16:26 -0700
« prev ^ index » next coverage.py v7.2.7, created at 2023-06-12 16:26 -0700
1from __future__ import annotations
3import warnings
4from typing import Callable, Tuple
6import numpy as np
7import pandas as pd
8import scipy.linalg # type: ignore
9from numpy.fft import irfft, rfft, rfftfreq
10from numpy.typing import ArrayLike
11from scipy.optimize import minimize # type: ignore
14class Wave:
15 r"""
16 Signal vector with associated information.
18 Attributes
19 ==========
20 signal : ndarray, optional
21 Signal vector. Default is an empty array.
23 ts : float, optional
24 Sampling time in ps. Default is 1.0.
26 t0 : float, optional
27 Absolute time associated with first data point. Default is 0.0.
29 metadata : dict, optional
30 Dictionary of metadata associated with the wave object. Default is an
31 empty dictionary.
32 """
34 def __init__(
35 self,
36 signal: ArrayLike = None,
37 ts: float = 1.0,
38 t0: float = 0.0,
39 metadata: dict = None,
40 ) -> None:
41 self.signal = signal
42 self.ts = ts
43 self.t0 = t0
44 self.metadata = metadata
46 def __repr__(self):
47 return (f"{self.__class__.__name__}(signal={self.signal.__repr__()}, "
48 f"ts={self.ts}, t0={self.t0}, metadata="
49 f"{self.metadata.__repr__()})")
51 def __array__(self, dtype=None):
52 # See
53 # https://numpy.org/doc/stable/user/basics.dispatch.html#basics
54 # -dispatch
55 # for details on NumPy custom array containers
56 return self.signal
58 @property
59 def t(self) -> ArrayLike:
60 r"""
61 Generate array of sampled times.
63 Returns
64 -------
65 ndarray
66 Array of sampled times associated with the signal, beginning with
67 t0 and separated by ts.
68 """
69 return
71 @property
72 def f(self) -> ArrayLike:
73 r"""
74 Generate array of sampled frequencies.
76 Returns
77 -------
78 ndarray
79 Array of frequencies associated with the signal. Generate with
80 numpy.rfftfreq.
82 """
83 return
85 @property
86 def spectrum(self) -> ArrayLike:
87 r"""
88 Complex spectrum of signal.
90 Returns
91 -------
92 ndarray
93 Real Fourier transform of signal. Generate with numpy.rfft
95 """
96 return
98 @property
99 def psd(self) -> ArrayLike:
100 r"""
101 Power spectral density of signal.
103 Returns
104 -------
105 ndarray
106 Power spectral density of signal. Generate with
107 scipy.signal.periodogram. Some of the optional parameters may be
108 useful to include, such as 'window', 'detrend', and 'scaling'.
110 """
111 return
113 def load(self, filepath: str) -> None:
114 r"""
115 Load ``Wave`` object from a data file.
117 Parameters
118 ----------
119 filepath : str
120 File path to read.
122 Returns
123 -------
125 """
126 return
129def fftfreq(n, ts):
130 """Computes the positive and negative frequencies sampled in the FFT.
132 Parameters
133 ----------
134 n : int
135 Number of time samples
136 ts: float
137 Sampling time
139 Returns
140 -------
141 f : ndarray
142 Frequency vector (1/``ts``) of length n containing the sample
143 frequencies.
144 """
146 if n % 2 == 1:
147 f = np.fft.fftfreq(n, ts)
148 else:
149 f = np.fft.fftfreq(n, ts)
150 f[int(n / 2)] = -f[int(n / 2)]
152 return f
155def noisevar(sigma: ArrayLike, mu: ArrayLike, ts: float) -> ArrayLike:
156 r"""
157 Compute the time-domain noise variance.
159 Parameters
160 ----------
161 sigma : array_like
162 Noise parameter array with shape (3, ). The first element corresponds
163 to the amplitude noise, in signal units (ie, the same units as ``mu``);
164 the second element corresponds to multiplicative noise, which is
165 dimensionless; and the third element corresponds to timebase noise, in
166 units of signal/time, where the units for time are the same as for
167 ``ts``.
168 mu : array_like
169 Time-domain signal.
170 ts : float
171 Sampling time.
173 Returns
174 -------
175 ndarray
176 Time-domain noise variance.
177 """
178 sigma = np.asarray(sigma)
179 mu = np.asarray(mu)
181 n = mu.shape[0]
182 w = 2 * np.pi * rfftfreq(n, ts)
183 mudot = irfft(1j * w * rfft(mu), n=n)
185 return sigma[0] ** 2 + (sigma[1] * mu) ** 2 + (sigma[2] * mudot) ** 2
188def noiseamp(sigma: ArrayLike, mu: ArrayLike, ts: float) -> ArrayLike:
189 r"""
190 Compute the time-domain noise amplitude.
192 Parameters
193 ----------
194 sigma : array_like
195 Noise parameter array with shape (3, ). The first element corresponds
196 to the amplitude noise, in signal units (ie, the same units as mu);
197 the second element corresponds to multiplicative noise, which is
198 dimensionless; and the third element corresponds to timebase noise, in
199 units of signal/time, where the units for time are the same as for t.
200 mu : array_like
201 Time-domain signal.
202 ts : float
203 Sampling time.
205 Returns
206 -------
207 ndarray
208 Time-domain noise amplitude, in signal units.
210 """
212 return np.sqrt(noisevar(sigma, mu, ts))
215def thzgen(
216 n: int,
217 ts: float,
218 t0: float,
219 a: float = 1.0,
220 taur: float = 0.3,
221 tauc: float = 0.1,
222 fwhm: float = 0.05,
223) -> tuple[ArrayLike, ArrayLike]:
224 r"""
225 Simulate a terahertz pulse.
227 Parameters
228 ----------
230 n : int
231 Number of samples.
233 ts : float
234 Sampling time.
236 t0 : float
237 Pulse center.
239 a : float, optional
240 Peak amplitude.
242 taur : float, optional
243 Current pulse rise time.
245 tauc : float, optional
246 Current pulse decay time.
248 fwhm : float, optional
249 Laser pulse FWHM.
251 Returns
252 -------
254 ndarray
255 Signal array.
257 ndarray
258 Array of time samples.
260 """
261 taul = fwhm / np.sqrt(2 * np.log(2))
263 f = rfftfreq(n, ts)
265 w = 2 * np.pi * f
266 ell = np.exp(-((w * taul) ** 2) / 2) / np.sqrt(2 * np.pi * taul ** 2)
267 r = 1 / (1 / taur - 1j * w) - 1 / (1 / taur + 1 / tauc - 1j * w)
268 s = -1j * w * (ell * r) ** 2 * np.exp(1j * w * t0)
270 t2 = ts * np.arange(n)
272 y = irfft(np.conj(s), n=n)
273 y = a * y / np.max(y)
275 return y, t2
278class DataPulse:
279 def __init__(self, filename=""):
280 self.AcquisitionTime = None
281 self.Description = None
282 self.TimeConstant = None
283 self.WaitTime = None
284 self.setPoint = None
285 self.scanOffset = None
286 self.temperature = None
287 self.time = None
288 self.amplitude = None
289 self.ChannelAMaximum = None
290 self.ChannelAMinimum = None
291 self.ChannelAVariance = None
292 self.ChannelASlope = None
293 self.ChannelAOffset = None
294 self.ChannelBMaximum = None
295 self.ChannelBMinimum = None
296 self.ChannelBVariance = None
297 self.ChannelBSlope = None
298 self.ChannelBOffset = None
299 self.ChannelCMaximum = None
300 self.ChannelCMinimum = None
301 self.ChannelCVariance = None
302 self.ChannelCSlope = None
303 self.ChannelCOffset = None
304 self.ChannelDMaximum = None
305 self.ChannelDMinimum = None
306 self.ChannelDVariance = None
307 self.ChannelDSlope = None
308 self.ChannelDOffset = None
309 self.dirname = None
310 self.file = None
311 self.filename = filename
312 self.frequency = None
314 if filename is not None:
315 data = pd.read_csv(filename, header=None, delimiter="\t")
317 ind = 0
318 for i in range(data.shape[0]):
319 try:
320 float(data[0][i])
321 except ValueError:
322 ind = i + 1
323 pass
325 keys = data[0][0:ind].to_list()
326 vals = data[1][0:ind].to_list()
328 for k in range(len(keys)):
329 if keys[k] in list(self.__dict__.keys()):
330 try:
331 float(vals[k])
332 setattr(self, keys[k], float(vals[k]))
333 except ValueError:
334 setattr(self, keys[k], vals[k])
336 else:
337 msg = "The data key is not defied"
338 raise Warning(msg)
340 self.time = data[0][ind:].to_numpy(dtype=float)
341 self.amplitude = data[1][ind:].to_numpy(dtype=float)
343 # Calculate frequency range
344 self.frequency = (
345 np.arange(
346 np.floor(len(self.time))) / 2 - 1
347 ).T / (self.time[-1] - self.time[0])
349 # Calculate fft
350 famp = np.fft.fft(self.amplitude)
351 self.famp = famp[0: int(np.floor(len(famp) / 2))]
354def shiftmtx(tau: float, n: int, ts: float = 1) -> ArrayLike:
355 """
356 Shiftmtx computes the n by n transfer matrix for a continuous time-shift.
358 Parameters
359 -----------
361 tau : float
362 Delay.
364 n : int
365 Number of samples.
367 ts: float, optional
368 Sampling time.
370 Returns
371 -------
372 h: ndarray
373 Transfer matrix with shape (n, n).
375 """
377 # Fourier method
378 f = rfftfreq(n, ts)
379 w = 2 * np.pi * f
381 imp = irfft(np.exp(-1j * w * tau), n=n)
383 # computes the n by n transformation matrix
384 h = scipy.linalg.toeplitz(imp, np.roll(np.flipud(imp), 1))
386 return h
389def airscancorrect(
390 x: ArrayLike,
391 *,
392 a: ArrayLike | None = None,
393 eta: ArrayLike | None = None,
394 ts: float = 1.0,
395) -> ArrayLike:
396 """Rescales and shifts each column of the matrix x.
398 Parameters
399 ----------
400 x : array_like
401 Data array with shape (n, m).
402 a : array_like, optional
403 Amplitude correction. If set, a must have shape (m,). Otherwise, no
404 correction is applied.
405 eta : array_like, optional
406 Delay correction. If set, a must have shape (m,). Otherwise, no
407 correction is applied.
408 ts : float, optional
409 Sampling time. Default is 1.0.
411 Returns
412 -------
413 xadj : ndarray
414 Adjusted data array.
416 """
417 x = np.asarray(x)
419 [n, m] = x.shape
421 if a is None:
422 a = np.ones((m,))
423 else:
424 a = np.asarray(a)
426 if eta is None:
427 eta = np.zeros((m,))
428 else:
429 eta = np.asarray(eta)
431 xadj = np.zeros((n, m))
432 # TODO: refactor with broadcasting
433 for i in np.arange(m):
434 s = shiftmtx(-eta[i], n, ts)
435 xadj[:, i] = s @ (x[:, i] / a[i])
437 return xadj
440def costfunlsq(
441 fun: Callable,
442 theta: ArrayLike,
443 xx: ArrayLike,
444 yy: ArrayLike,
445 sigmax: ArrayLike,
446 sigmay: ArrayLike,
447 ts: float,
448) -> ArrayLike:
449 r"""Computes the maximum likelihood cost function.
451 Parameters
452 ----------
453 fun : callable
454 Transfer function, in the form fun(theta,w), -iwt convention.
456 theta : array_like
457 Input parameters for the function.
459 xx : array_like
460 Measured input signal.
462 yy : array_like
463 Measured output signal.
465 sigmax : array_like
466 Noise covariance matrix of the input signal.
468 sigmay : array_like
469 Noise covariance matrix of the output signal.
471 ts : float
472 Sampling time.
474 Returns
475 -------
476 res : array_like
479 """
480 n = xx.shape[0]
481 wfft = 2 * np.pi * rfftfreq(n, ts)
482 h = np.conj(fun(theta, wfft))
484 ry = yy - irfft(rfft(xx) * h, n=n)
485 vy = np.diag(sigmay ** 2)
487 htilde = irfft(h, n=n)
489 uy = np.zeros((n, n))
490 for k in np.arange(n):
491 a = np.reshape(np.roll(htilde, k), (n, 1))
492 b = np.reshape(np.conj(np.roll(htilde, k)), (1, n))
493 uy = uy + np.real(np.dot(a, b)) * sigmax[k] ** 2
494 # uy = uy + np.real(np.roll(htilde, k-1) @ np.roll(htilde, k-1).T) @
495 # sigmax[k]**2
497 w = np.dot(np.eye(n), scipy.linalg.inv(scipy.linalg.sqrtm(uy + vy)))
498 res = np.dot(w, ry)
500 return res
503def tdtf(fun: Callable, theta: ArrayLike, n: int, ts: float) -> ArrayLike:
504 """
505 Computes the time-domain transfer matrix for a frequency response function.
507 Parameters
508 ----------
509 fun : callable
510 Frequency function, in the form fun(theta, w), where theta
511 is a vector of the function parameters. The function should be
512 expressed in the -iwt convention and must be Hermitian.
514 theta : array_like
515 Input parameters for the function.
517 n : int
518 Number of time samples.
520 ts : array_like
521 Sampling time.
523 Returns
524 -------
525 h : array_like
526 Transfer matrix with size (n,n).
528 """
530 # compute the transfer function over positive frequencies
531 if not isinstance(ts, np.ndarray):
532 ts = np.array([ts])
533 else:
534 ts = ts
536 fs = 1 / (ts * n)
537 fp = fs * np.arange(0, (n - 1) // 2 + 1)
538 wp = 2 * np.pi * fp
539 tfunp = fun(theta, wp)
541 # The transfer function is Hermitian, so we evaluate negative frequencies
542 # by taking the complex conjugate of the corresponding positive frequency.
543 # Include the value of the transfer function at the Nyquist frequency for
544 # even n.
545 if n % 2 != 0:
546 tfun = np.concatenate((tfunp, np.conj(np.flipud(tfunp[1:]))))
548 else:
549 wny = np.pi * n * fs
550 # print('tfunp', tfunp)
551 tfun = np.concatenate(
552 (
553 tfunp,
554 np.conj(
555 np.concatenate((fun(theta, wny), np.flipud(tfunp[1:])))
556 ),
557 )
558 )
560 # Evaluate the impulse response by taking the inverse Fourier transform,
561 # taking the complex conjugate first to convert to ... +iwt convention
563 imp = np.real(np.fft.ifft(np.conj(tfun)))
564 h = scipy.linalg.toeplitz(imp, np.roll(np.flipud(imp), 1))
566 return h
569def tdnll(
570 x: ArrayLike,
571 logv: ArrayLike,
572 mu: ArrayLike,
573 a: ArrayLike | None = None,
574 eta: ArrayLike | None = None,
575 ts: float = 1.0,
576 d: ArrayLike | None = None,
577 fix_logv: bool = False,
578 fix_mu: bool = False,
579 fix_a: bool = False,
580 fix_eta: bool = False
581) -> Tuple[ArrayLike, ArrayLike]:
582 r"""
583 Compute negative log-likelihood for the time-domain noise model.
585 Computes the negative log-likelihood function for obtaining the
586 data matrix ``x``, given the parameter dictionary param.
588 Parameters
589 ----------
590 x : ndarray or matrix
591 Data matrix
592 logv : ndarray
593 Array of size (3, ) containing log of noise parameters.
594 mu : ndarray
595 Signal vector of size (n,).
596 a: ndarray, optional
597 Amplitude vector of size (m,).
598 eta : ndarray, optional
599 Delay vector of size (m,).
600 ts : float, optional
601 Sampling time.
602 d : ndarray, optional
603 Derivative matrix, size (n, n).
604 fix_logv : bool, optional
605 Log of noise parameters.
606 fix_mu : bool, optional
607 Signal vector.
608 fix_a : bool, optional
609 Amplitude vector.
610 fix_eta : bool, optional
611 Delay vector.
613 Returns
614 -------
615 nll : callable
616 Negative log-likelihood function
617 gradnll : ndarray
618 Gradient of the negative log-likelihood function
619 """
620 # Parameters to ignore when computing gradnll
621 ignore_a = False
622 ignore_eta = False
624 # Parse and validate function inputs
625 x = np.asarray(x)
626 if x.ndim != 2:
627 raise ValueError("Data array x must be 2D.")
628 n, m = x.shape
630 logv = np.asarray(logv)
631 if logv.size != 3:
632 raise ValueError("Noise parameter array logv must have 3 elements.")
634 mu = np.asarray(mu)
635 if mu.ndim != 1:
636 raise ValueError("Ideal signal vector mu must be 1D.")
637 if mu.size != n:
638 raise ValueError("Size of mu is incompatible with data array x.")
639 mu = np.reshape(mu, (n, 1))
641 if a is None:
642 a = np.ones((m,))
643 ignore_a = True
644 else:
645 a = np.asarray(a)
646 if a.size != m:
647 raise ValueError("Size of a is incompatible with data array x.")
648 a = np.reshape(a, (m, 1))
650 if eta is None:
651 eta = np.ones((m,))
652 ignore_eta = True
653 else:
654 eta = np.asarray(eta)
655 if eta.size != m:
656 raise ValueError("Size of eta is incompatible with data array x.")
657 eta = np.reshape(eta, (m, 1))
659 if d is None:
660 def fun(_, _w):
661 return -1j * _w
663 d = tdtf(fun, 0, n, ts)
665 # Compute variance
666 v = np.exp(logv)
667 v = np.reshape(v, (len(v), 1))
669 # Compute frequency vector and Fourier coefficients of mu
670 f = fftfreq(n, ts)
671 w = 2 * np.pi * f
672 w = w.reshape(len(w), 1)
673 mu_f = np.fft.fft(mu.flatten()).reshape(len(mu), 1)
675 gradcalc = np.logical_not(
676 [
677 [fix_logv],
678 [fix_mu],
679 [fix_a or ignore_a],
680 [fix_eta or ignore_eta],
681 ]
682 )
684 if ignore_eta:
685 zeta = mu * np.conj(a).T
686 zeta_f = np.fft.fft(zeta, axis=0)
687 else:
688 exp_iweta = np.exp(1j * np.tile(w, m) * np.conj(np.tile(eta, n)).T)
689 zeta_f = (
690 np.conj(np.tile(a, n)).T * np.conj(exp_iweta) * np.tile(mu_f,
691 m)
692 )
693 zeta = np.real(np.fft.ifft(zeta_f, axis=0))
695 # Compute negative - log likelihood and gradient
697 # Compute residuals and their squares for subsequent computations
698 res = x - zeta
699 ressq = res ** 2
701 # Simplest case: just variance and signal parameters, A and eta fixed at
702 # defaults
703 if ignore_a and ignore_eta:
704 dmu = np.real(np.fft.ifft(1j * w * mu_f, axis=0))
705 valpha = v[0]
706 vbeta = v[1] * mu ** 2
707 vtau = v[2] * dmu ** 2
708 vtot = valpha + vbeta + vtau
710 resnormsq = ressq / np.tile(vtot, m)
711 nll = (
712 m * n * np.log(2 * np.pi) / 2
713 + (m / 2) * np.sum(np.log(vtot))
714 + np.sum(resnormsq) / 2
715 )
717 # Compute gradient if requested
718 # if nargout > 1:
719 ngrad = np.sum(gradcalc[0:2] * [[3], [n]])
720 gradnll = np.zeros((ngrad, 1))
721 nstart = 0
722 dvar = (vtot - np.mean(ressq, axis=1).reshape(n, 1)) / vtot ** 2
723 if gradcalc[0]:
724 gradnll[nstart] = (m / 2) * np.sum(dvar) * v[0]
725 gradnll[nstart + 1] = (m / 2) * np.sum(mu ** 2 * dvar) * v[1]
726 gradnll[nstart + 2] = (m / 2) * np.sum(dmu ** 2.0 * dvar) * v[2]
727 nstart = nstart + 3
728 if gradcalc[1]:
729 # print('mu shape : ', mu.shape)
730 # print('dvar shape: ', dvar.shape)
731 # print('d shape: ', d.shape)
732 # print('Dmu shape: ', dmu.shape)
733 gradnll[nstart: nstart + n] = m * (
734 v[1] * mu * dvar
735 + v[2] * np.dot(d.T, (dmu * dvar))
736 - np.mean(res, axis=1).reshape(n, 1) / vtot
737 )
739 # Alternative case: A, eta, or both are not set to defaults
740 else:
741 dzeta = np.real(np.fft.ifft(1j * np.tile(w, m) * zeta_f, axis=0))
743 valpha = v[0]
744 vbeta = v[1] * zeta ** 2
745 vtau = v[2] * dzeta ** 2
746 vtot = valpha + vbeta + vtau
748 resnormsq = ressq / vtot
749 nll = (
750 m * n * np.log(2 * np.pi) / 2
751 + np.sum(np.log(vtot)) / 2
752 + np.sum(resnormsq) / 2
753 )
755 # Compute gradient if requested
756 # if nargout > 1:
757 ngrad = np.sum(gradcalc * [[3], [n], [m], [m]])
758 gradnll = np.zeros((ngrad, 1))
759 nstart = 0
760 reswt = res / vtot
761 dvar = (vtot - ressq) / vtot ** 2
762 if gradcalc[0]:
763 # Gradient wrt logv
764 gradnll[nstart] = 0.5 * np.sum(dvar) * v[0]
765 gradnll[nstart + 1] = (
766 0.5 * np.sum(zeta.flatten() ** 2 * dvar.flatten()) * v[1]
767 )
768 gradnll[nstart + 2] = (
769 0.5 * np.sum(dzeta.flatten() ** 2 * dvar.flatten()) * v[2]
770 )
771 nstart = nstart + 3
772 if gradcalc[1]:
773 # Gradient wrt mu
774 p = np.fft.fft(v[1] * dvar * zeta - reswt, axis=0) - 1j * v[
775 2
776 ] * w * np.fft.fft(dvar * dzeta, axis=0)
777 gradnll[nstart: nstart + n] = np.sum(
778 np.conj(a).T * np.real(np.fft.ifft(exp_iweta * p, axis=0)),
779 axis=1,
780 ).reshape(n, 1)
781 nstart = nstart + n
782 if gradcalc[2]:
783 # Gradient wrt A
784 term = (vtot - valpha) * dvar - reswt * zeta
785 if np.any(np.isclose(a, 0)):
786 msg = (
787 "One or more elements of the amplitude vector are "
788 "close to zero "
789 )
790 raise ValueError(msg)
791 gradnll[nstart: nstart + m] = (
792 np.conj(np.sum(term, axis=0)).reshape(m, 1) / a
793 )
794 if not fix_mu:
795 gradnll = np.delete(gradnll, nstart)
796 nstart = nstart + m - 1
797 else:
798 nstart = nstart + m
799 if gradcalc[3]:
800 # Gradient wrt eta
801 ddzeta = np.real(np.fft.ifft(-np.tile(w, m) ** 2 * zeta_f, axis=0))
802 gradnll = np.squeeze(gradnll)
803 gradnll[nstart: nstart + m] = -np.sum(
804 dvar * (zeta * dzeta * v[1] + dzeta * ddzeta * v[2])
805 - reswt * dzeta,
806 axis=0,
807 ).reshape(
808 m,
809 )
811 if not fix_mu:
812 gradnll = np.delete(gradnll, nstart)
813 gradnll = gradnll.flatten()
815 return nll, gradnll
818def tdnoisefit(
819 x: ArrayLike,
820 v0: ArrayLike | None = None,
821 mu0: ArrayLike | None = None,
822 a0: ArrayLike | None = None,
823 eta0: ArrayLike | None = None,
824 ts: float = 1.0,
825 fix_v: bool = False,
826 fix_mu: bool = False,
827 fix_a: bool = True,
828 fix_eta: bool = True,
829 ignore_a: bool = True,
830 ignore_eta: bool = True
831) -> Tuple[dict, float, dict]:
832 r"""
833 Compute time-domain noise model parameters.
835 Computes the noise parameters sigma and the underlying signal vector ``mu``
836 for the data matrix ``x``, where the columns of ``x`` are each noisy
837 measurements of ``mu``.
839 Parameters
840 ----------
841 x : ndarray
842 Data array.
843 v0 : ndarray, optional
844 Initial guess, noise model parameters with size (3,).
845 mu0 : ndarray, optional
846 Initial guess, signal vector with size (n,).
847 a0 : ndarray, optional
848 Initial guess, amplitude vector with size (m,).
849 eta0 : ndarray, optional
850 Initial guess, delay vector with size (m,).
851 ts : float, optional
852 Sampling time
853 fix_v : bool, optional
854 Noise variance parameters.
855 fix_mu : bool, optional
856 Signal vector.
857 fix_a : bool, optional
858 Amplitude vector.
859 fix_eta : bool, optional
860 Delay vector.
861 ignore_a : bool, optional
862 Amplitude vector.
863 ignore_eta : bool, optional
864 Delay vector.
866 Returns
867 --------
868 p : dict
869 Output parameter dictionary containing:
870 eta : ndarray
871 Delay vector.
872 a : ndarray
873 Amplitude vector.
874 mu : ndarray
875 Signal vector.
876 var : ndarray
877 Log of noise parameters
878 fval : float
879 Value of NLL cost function from FMINUNC
880 Diagnostic : dict
881 Dictionary containing diagnostic information
882 err : dic
883 Dictionary containing error of the parameters.
884 grad : ndarray
885 Negative loglikelihood cost function gradient from
886 scipy.optimize.minimize BFGS method.
887 hessian : ndarray
888 Negative loglikelihood cost function hessian from
889 scipy.optimize.minimize BFGS method.
890 """
891 # Parse and validate function inputs
892 x = np.asarray(x)
893 if x.ndim != 2:
894 raise ValueError("Data array x must be 2D.")
895 n, m = x.shape
897 if v0 is None:
898 v0 = np.mean(np.var(x, 1)) * np.array([1, 1, 1])
899 else:
900 v0 = np.asarray(v0)
901 if v0.size != 3:
902 raise ValueError(
903 "Noise parameter array logv must have 3 elements."
904 )
906 if mu0 is None:
907 mu0 = np.mean(x, 1)
908 else:
909 mu0 = np.asarray(mu0)
910 if mu0.size != n:
911 raise ValueError("Size of mu0 is incompatible with data array x.")
913 if a0 is None:
914 a0 = np.ones(m)
915 else:
916 a0 = np.asarray(a0)
917 if a0.size != m:
918 raise ValueError("Size of a0 is incompatible with data array x.")
920 if eta0 is None:
921 eta0 = np.ones(m)
922 else:
923 eta0 = np.asarray(eta0)
924 if eta0.size != m:
925 raise ValueError("Size of eta0 is incompatible with data array x.")
927 mle = {"x0": np.array([])}
928 idxstart = 0
929 idxrange = {}
931 # If fix['logv'], return log(v0); otherwise return logv parameters
932 if fix_v:
934 def setplogv(_):
935 return np.log(v0)
937 else:
938 mle["x0"] = np.concatenate((mle["x0"], np.log(v0)))
939 idxend = idxstart + 3
940 idxrange["logv"] = np.arange(idxstart, idxend)
942 def setplogv(_p):
943 return _p[idxrange["logv"]]
945 idxstart = idxend
947 # If Fix['mu'], return mu0, otherwise, return mu parameters
948 if fix_mu:
950 def setpmu(_):
951 return mu0
953 else:
954 mle["x0"] = np.concatenate((mle["x0"], mu0))
955 idxend = idxstart + n
956 idxrange["mu"] = np.arange(idxstart, idxend)
958 def setpmu(_p):
959 return _p[idxrange["mu"]]
961 idxstart = idxend
962 pass
964 # If ignore_a, return None; if fix_a, return a0; if (!fix_a & fix_mu),
965 # return all a parameters; if !fix_a & !fix_mu, return all a parameters
966 # but first
968 if ignore_a:
970 def setpa(_):
971 return None
973 elif fix_a:
975 def setpa(_):
976 return a0
978 elif fix_mu:
979 mle["x0"] = np.concatenate((mle["x0"], a0))
980 idxend = idxstart + m
981 idxrange["a"] = np.arange(idxstart, idxend)
983 def setpa(_p):
984 return _p[idxrange["a"]]
986 idxstart = idxend
987 else:
988 mle["x0"] = np.concatenate(
989 (mle["x0"], a0[1:] / a0[0])
990 )
991 idxend = idxstart + m - 1
992 idxrange["a"] = np.arange(idxstart, idxend)
994 def setpa(_p):
995 return np.concatenate(([1], _p[idxrange["a"]]), axis=0)
997 idxstart = idxend
998 pass
1000 # If ignore_eta, return None; if fix_eta, return eta0; if !fix_eta &
1001 # fix_mu,return all eta parameters; if !fix_eta & !fix_mu, return all eta
1002 # parameters but first
1004 if ignore_eta:
1006 def setpeta(_):
1007 return None
1009 elif fix_eta:
1011 def setpeta(_):
1012 return eta0
1014 elif fix_mu:
1015 mle["x0"] = np.concatenate((mle["x0"], eta0))
1016 idxend = idxstart + m
1017 idxrange["eta"] = np.arange(idxstart, idxend)
1019 def setpeta(_p):
1020 return _p[idxrange["eta"]]
1022 else:
1023 mle["x0"] = np.concatenate(
1024 (mle["x0"], eta0[1:] - eta0[0])
1025 )
1026 idxend = idxstart + m - 1
1027 idxrange["eta"] = np.arange(idxstart, idxend)
1029 def setpeta(_p):
1030 return np.concatenate(([0], _p[idxrange["eta"]]), axis=0)
1032 pass
1034 def fun(_, _w):
1035 return -1j * _w
1037 d = tdtf(fun, 0, n, ts)
1039 def parsein(_p):
1040 return {
1041 "logv": setplogv(_p),
1042 "mu": setpmu(_p),
1043 "a": setpa(_p),
1044 "eta": setpeta(_p),
1045 "ts": ts,
1046 "d": d,
1047 }
1049 def objective(_p):
1050 return tdnll(x, *parsein(_p).values(),
1051 fix_v, fix_mu, fix_a, fix_eta)[0]
1053 def jacobian(_p):
1054 return tdnll(x, *parsein(_p).values(),
1055 fix_v, fix_mu, fix_a, fix_eta)[1]
1057 mle["objective"] = objective
1058 out = minimize(mle["objective"], mle["x0"], method="BFGS", jac=jacobian)
1060 # The trust-region algorithm returns the Hessian for the next-to-last
1061 # iterate, which may not be near the final point. To check, test for
1062 # positive definiteness by attempting to Cholesky factorize it. If it
1063 # returns an error, rerun the optimization with the quasi-Newton algorithm
1064 # from the current optimal point.
1066 try:
1067 np.linalg.cholesky(np.linalg.inv(out.hess_inv))
1068 hess = np.linalg.inv(out.hess_inv)
1069 except np.linalg.LinAlgError:
1070 print(
1071 "Hessian returned by FMINUNC is not positive definite;\n"
1072 "recalculating with quasi-Newton algorithm"
1073 )
1075 mle["x0"] = out.x
1076 out2 = minimize(
1077 mle["objective"], mle["x0"], method="BFGS", jac=jacobian
1078 )
1079 hess = np.linalg.inv(out2.hess_inv)
1081 # Parse output
1082 p = {}
1083 idxrange = {}
1084 idxstart = 0
1086 if fix_v:
1087 p["var"] = v0
1088 else:
1089 idxend = idxstart + 3
1090 idxrange["logv"] = np.arange(idxstart, idxend)
1091 idxstart = idxend
1092 p["var"] = np.exp(out.x[idxrange["logv"]])
1093 pass
1095 if fix_mu:
1096 p["mu"] = mu0
1097 else:
1098 idxend = idxstart + n
1099 idxrange["mu"] = np.arange(idxstart, idxend)
1100 idxstart = idxend
1101 p["mu"] = out.x[idxrange["mu"]]
1102 pass
1104 if ignore_a or fix_a:
1105 p["a"] = a0
1106 elif fix_mu:
1107 idxend = idxstart + m
1108 idxrange["a"] = np.arange(idxstart, idxend)
1109 idxstart = idxend
1110 p["a"] = out.x[idxrange["a"]]
1111 else:
1112 idxend = idxstart + m - 1
1113 idxrange["a"] = np.arange(idxstart, idxend)
1114 idxstart = idxend
1115 p["a"] = np.concatenate(([1], out.x[idxrange["a"]]), axis=0)
1116 pass
1118 if ignore_eta or fix_eta:
1119 p["eta"] = eta0
1120 elif fix_mu:
1121 idxend = idxstart + m
1122 idxrange["eta"] = np.arange(idxstart, idxend)
1123 p["eta"] = out.x[idxrange["eta"]]
1124 else:
1125 idxend = idxstart + m - 1
1126 idxrange["eta"] = np.arange(idxstart, idxend)
1127 p["eta"] = np.concatenate(([0], out.x[idxrange["eta"]]), axis=0)
1128 pass
1130 p["ts"] = ts
1132 vary_param = np.logical_not(
1133 [
1134 fix_v,
1135 fix_mu,
1136 fix_a or ignore_a,
1137 fix_eta or ignore_eta,
1138 ]
1139 )
1140 diagnostic = {
1141 "grad": out.jac,
1142 "hessian": hess,
1143 "err": {"var": [], "mu": [], "a": [], "eta": []},
1144 }
1145 v = np.dot(np.eye(hess.shape[0]), scipy.linalg.inv(hess))
1146 err = np.sqrt(np.diag(v))
1147 idxstart = 0
1148 if vary_param[0]:
1149 diagnostic["err"]["var"] = np.sqrt(
1150 np.diag(np.diag(p["var"]) * v[0:3, 0:3]) * np.diag(p["var"])
1151 )
1152 idxstart = idxstart + 3
1153 pass
1155 if vary_param[1]:
1156 diagnostic["err"]["mu"] = err[idxstart: idxstart + n]
1157 idxstart = idxstart + n
1158 pass
1160 if vary_param[2]:
1161 if vary_param[1]:
1162 diagnostic["err"]["a"] = err[idxstart: idxstart + m - 1]
1163 idxstart = idxstart + m - 1
1164 else:
1165 diagnostic["err"]["a"] = err[idxstart: idxstart + m]
1166 idxstart = idxstart + m
1167 pass
1169 if vary_param[3]:
1170 if vary_param[1]:
1171 diagnostic["err"]["eta"] = err[idxstart: idxstart + m - 1]
1172 else:
1173 diagnostic["err"]["eta"] = err[idxstart: idxstart + m]
1174 pass
1176 return [p, out.fun, diagnostic]