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

1from __future__ import annotations 

2 

3import warnings 

4from typing import Callable, Tuple 

5 

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 

12 

13 

14class Wave: 

15 r""" 

16 Signal vector with associated information. 

17 

18 Attributes 

19 ========== 

20 signal : ndarray, optional 

21 Signal vector. Default is an empty array. 

22 

23 ts : float, optional 

24 Sampling time in ps. Default is 1.0. 

25 

26 t0 : float, optional 

27 Absolute time associated with first data point. Default is 0.0. 

28 

29 metadata : dict, optional 

30 Dictionary of metadata associated with the wave object. Default is an 

31 empty dictionary. 

32 """ 

33 

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 

45 

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__()})") 

50 

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 

57 

58 @property 

59 def t(self) -> ArrayLike: 

60 r""" 

61 Generate array of sampled times. 

62 

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 

70 

71 @property 

72 def f(self) -> ArrayLike: 

73 r""" 

74 Generate array of sampled frequencies. 

75 

76 Returns 

77 ------- 

78 ndarray 

79 Array of frequencies associated with the signal. Generate with 

80 numpy.rfftfreq. 

81 

82 """ 

83 return 

84 

85 @property 

86 def spectrum(self) -> ArrayLike: 

87 r""" 

88 Complex spectrum of signal. 

89 

90 Returns 

91 ------- 

92 ndarray 

93 Real Fourier transform of signal. Generate with numpy.rfft 

94 

95 """ 

96 return 

97 

98 @property 

99 def psd(self) -> ArrayLike: 

100 r""" 

101 Power spectral density of signal. 

102 

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'. 

109 

110 """ 

111 return 

112 

113 def load(self, filepath: str) -> None: 

114 r""" 

115 Load ``Wave`` object from a data file. 

116 

117 Parameters 

118 ---------- 

119 filepath : str 

120 File path to read. 

121 

122 Returns 

123 ------- 

124 

125 """ 

126 return 

127 

128 

129def fftfreq(n, ts): 

130 """Computes the positive and negative frequencies sampled in the FFT. 

131 

132 Parameters 

133 ---------- 

134 n : int 

135 Number of time samples 

136 ts: float 

137 Sampling time 

138 

139 Returns 

140 ------- 

141 f : ndarray 

142 Frequency vector (1/``ts``) of length n containing the sample 

143 frequencies. 

144 """ 

145 

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)] 

151 

152 return f 

153 

154 

155def noisevar(sigma: ArrayLike, mu: ArrayLike, ts: float) -> ArrayLike: 

156 r""" 

157 Compute the time-domain noise variance. 

158 

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. 

172 

173 Returns 

174 ------- 

175 ndarray 

176 Time-domain noise variance. 

177 """ 

178 sigma = np.asarray(sigma) 

179 mu = np.asarray(mu) 

180 

181 n = mu.shape[0] 

182 w = 2 * np.pi * rfftfreq(n, ts) 

183 mudot = irfft(1j * w * rfft(mu), n=n) 

184 

185 return sigma[0] ** 2 + (sigma[1] * mu) ** 2 + (sigma[2] * mudot) ** 2 

186 

187 

188def noiseamp(sigma: ArrayLike, mu: ArrayLike, ts: float) -> ArrayLike: 

189 r""" 

190 Compute the time-domain noise amplitude. 

191 

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. 

204 

205 Returns 

206 ------- 

207 ndarray 

208 Time-domain noise amplitude, in signal units. 

209 

210 """ 

211 

212 return np.sqrt(noisevar(sigma, mu, ts)) 

213 

214 

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. 

226 

227 Parameters 

228 ---------- 

229 

230 n : int 

231 Number of samples. 

232 

233 ts : float 

234 Sampling time. 

235 

236 t0 : float 

237 Pulse center. 

238 

239 a : float, optional 

240 Peak amplitude. 

241 

242 taur : float, optional 

243 Current pulse rise time. 

244 

245 tauc : float, optional 

246 Current pulse decay time. 

247 

248 fwhm : float, optional 

249 Laser pulse FWHM. 

250 

251 Returns 

252 ------- 

253 

254 ndarray 

255 Signal array. 

256 

257 ndarray 

258 Array of time samples. 

259 

260 """ 

261 taul = fwhm / np.sqrt(2 * np.log(2)) 

262 

263 f = rfftfreq(n, ts) 

264 

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) 

269 

270 t2 = ts * np.arange(n) 

271 

272 y = irfft(np.conj(s), n=n) 

273 y = a * y / np.max(y) 

274 

275 return y, t2 

276 

277 

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 

313 

314 if filename is not None: 

315 data = pd.read_csv(filename, header=None, delimiter="\t") 

316 

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 

324 

325 keys = data[0][0:ind].to_list() 

326 vals = data[1][0:ind].to_list() 

327 

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]) 

335 

336 else: 

337 msg = "The data key is not defied" 

338 raise Warning(msg) 

339 

340 self.time = data[0][ind:].to_numpy(dtype=float) 

341 self.amplitude = data[1][ind:].to_numpy(dtype=float) 

342 

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]) 

348 

349 # Calculate fft 

350 famp = np.fft.fft(self.amplitude) 

351 self.famp = famp[0: int(np.floor(len(famp) / 2))] 

352 

353 

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. 

357 

358 Parameters 

359 ----------- 

360 

361 tau : float 

362 Delay. 

363 

364 n : int 

365 Number of samples. 

366 

367 ts: float, optional 

368 Sampling time. 

369 

370 Returns 

371 ------- 

372 h: ndarray 

373 Transfer matrix with shape (n, n). 

374 

375 """ 

376 

377 # Fourier method 

378 f = rfftfreq(n, ts) 

379 w = 2 * np.pi * f 

380 

381 imp = irfft(np.exp(-1j * w * tau), n=n) 

382 

383 # computes the n by n transformation matrix 

384 h = scipy.linalg.toeplitz(imp, np.roll(np.flipud(imp), 1)) 

385 

386 return h 

387 

388 

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. 

397 

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. 

410 

411 Returns 

412 ------- 

413 xadj : ndarray 

414 Adjusted data array. 

415 

416 """ 

417 x = np.asarray(x) 

418 

419 [n, m] = x.shape 

420 

421 if a is None: 

422 a = np.ones((m,)) 

423 else: 

424 a = np.asarray(a) 

425 

426 if eta is None: 

427 eta = np.zeros((m,)) 

428 else: 

429 eta = np.asarray(eta) 

430 

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]) 

436 

437 return xadj 

438 

439 

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. 

450 

451 Parameters 

452 ---------- 

453 fun : callable 

454 Transfer function, in the form fun(theta,w), -iwt convention. 

455 

456 theta : array_like 

457 Input parameters for the function. 

458 

459 xx : array_like 

460 Measured input signal. 

461 

462 yy : array_like 

463 Measured output signal. 

464 

465 sigmax : array_like 

466 Noise covariance matrix of the input signal. 

467 

468 sigmay : array_like 

469 Noise covariance matrix of the output signal. 

470 

471 ts : float 

472 Sampling time. 

473 

474 Returns 

475 ------- 

476 res : array_like 

477 

478 

479 """ 

480 n = xx.shape[0] 

481 wfft = 2 * np.pi * rfftfreq(n, ts) 

482 h = np.conj(fun(theta, wfft)) 

483 

484 ry = yy - irfft(rfft(xx) * h, n=n) 

485 vy = np.diag(sigmay ** 2) 

486 

487 htilde = irfft(h, n=n) 

488 

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 

496 

497 w = np.dot(np.eye(n), scipy.linalg.inv(scipy.linalg.sqrtm(uy + vy))) 

498 res = np.dot(w, ry) 

499 

500 return res 

501 

502 

503def tdtf(fun: Callable, theta: ArrayLike, n: int, ts: float) -> ArrayLike: 

504 """ 

505 Computes the time-domain transfer matrix for a frequency response function. 

506 

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. 

513 

514 theta : array_like 

515 Input parameters for the function. 

516 

517 n : int 

518 Number of time samples. 

519 

520 ts : array_like 

521 Sampling time. 

522 

523 Returns 

524 ------- 

525 h : array_like 

526 Transfer matrix with size (n,n). 

527 

528 """ 

529 

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 

535 

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) 

540 

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:])))) 

547 

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 ) 

559 

560 # Evaluate the impulse response by taking the inverse Fourier transform, 

561 # taking the complex conjugate first to convert to ... +iwt convention 

562 

563 imp = np.real(np.fft.ifft(np.conj(tfun))) 

564 h = scipy.linalg.toeplitz(imp, np.roll(np.flipud(imp), 1)) 

565 

566 return h 

567 

568 

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. 

584 

585 Computes the negative log-likelihood function for obtaining the 

586 data matrix ``x``, given the parameter dictionary param. 

587 

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. 

612 

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 

623 

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 

629 

630 logv = np.asarray(logv) 

631 if logv.size != 3: 

632 raise ValueError("Noise parameter array logv must have 3 elements.") 

633 

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)) 

640 

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)) 

649 

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)) 

658 

659 if d is None: 

660 def fun(_, _w): 

661 return -1j * _w 

662 

663 d = tdtf(fun, 0, n, ts) 

664 

665 # Compute variance 

666 v = np.exp(logv) 

667 v = np.reshape(v, (len(v), 1)) 

668 

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) 

674 

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 ) 

683 

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)) 

694 

695 # Compute negative - log likelihood and gradient 

696 

697 # Compute residuals and their squares for subsequent computations 

698 res = x - zeta 

699 ressq = res ** 2 

700 

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 

709 

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 ) 

716 

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 ) 

738 

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)) 

742 

743 valpha = v[0] 

744 vbeta = v[1] * zeta ** 2 

745 vtau = v[2] * dzeta ** 2 

746 vtot = valpha + vbeta + vtau 

747 

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 ) 

754 

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 ) 

810 

811 if not fix_mu: 

812 gradnll = np.delete(gradnll, nstart) 

813 gradnll = gradnll.flatten() 

814 

815 return nll, gradnll 

816 

817 

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. 

834 

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``. 

838 

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. 

865 

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 

896 

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 ) 

905 

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.") 

912 

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.") 

919 

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.") 

926 

927 mle = {"x0": np.array([])} 

928 idxstart = 0 

929 idxrange = {} 

930 

931 # If fix['logv'], return log(v0); otherwise return logv parameters 

932 if fix_v: 

933 

934 def setplogv(_): 

935 return np.log(v0) 

936 

937 else: 

938 mle["x0"] = np.concatenate((mle["x0"], np.log(v0))) 

939 idxend = idxstart + 3 

940 idxrange["logv"] = np.arange(idxstart, idxend) 

941 

942 def setplogv(_p): 

943 return _p[idxrange["logv"]] 

944 

945 idxstart = idxend 

946 

947 # If Fix['mu'], return mu0, otherwise, return mu parameters 

948 if fix_mu: 

949 

950 def setpmu(_): 

951 return mu0 

952 

953 else: 

954 mle["x0"] = np.concatenate((mle["x0"], mu0)) 

955 idxend = idxstart + n 

956 idxrange["mu"] = np.arange(idxstart, idxend) 

957 

958 def setpmu(_p): 

959 return _p[idxrange["mu"]] 

960 

961 idxstart = idxend 

962 pass 

963 

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 

967 

968 if ignore_a: 

969 

970 def setpa(_): 

971 return None 

972 

973 elif fix_a: 

974 

975 def setpa(_): 

976 return a0 

977 

978 elif fix_mu: 

979 mle["x0"] = np.concatenate((mle["x0"], a0)) 

980 idxend = idxstart + m 

981 idxrange["a"] = np.arange(idxstart, idxend) 

982 

983 def setpa(_p): 

984 return _p[idxrange["a"]] 

985 

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) 

993 

994 def setpa(_p): 

995 return np.concatenate(([1], _p[idxrange["a"]]), axis=0) 

996 

997 idxstart = idxend 

998 pass 

999 

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 

1003 

1004 if ignore_eta: 

1005 

1006 def setpeta(_): 

1007 return None 

1008 

1009 elif fix_eta: 

1010 

1011 def setpeta(_): 

1012 return eta0 

1013 

1014 elif fix_mu: 

1015 mle["x0"] = np.concatenate((mle["x0"], eta0)) 

1016 idxend = idxstart + m 

1017 idxrange["eta"] = np.arange(idxstart, idxend) 

1018 

1019 def setpeta(_p): 

1020 return _p[idxrange["eta"]] 

1021 

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) 

1028 

1029 def setpeta(_p): 

1030 return np.concatenate(([0], _p[idxrange["eta"]]), axis=0) 

1031 

1032 pass 

1033 

1034 def fun(_, _w): 

1035 return -1j * _w 

1036 

1037 d = tdtf(fun, 0, n, ts) 

1038 

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 } 

1048 

1049 def objective(_p): 

1050 return tdnll(x, *parsein(_p).values(), 

1051 fix_v, fix_mu, fix_a, fix_eta)[0] 

1052 

1053 def jacobian(_p): 

1054 return tdnll(x, *parsein(_p).values(), 

1055 fix_v, fix_mu, fix_a, fix_eta)[1] 

1056 

1057 mle["objective"] = objective 

1058 out = minimize(mle["objective"], mle["x0"], method="BFGS", jac=jacobian) 

1059 

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. 

1065 

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 ) 

1074 

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) 

1080 

1081 # Parse output 

1082 p = {} 

1083 idxrange = {} 

1084 idxstart = 0 

1085 

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 

1094 

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 

1103 

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 

1117 

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 

1129 

1130 p["ts"] = ts 

1131 

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 

1154 

1155 if vary_param[1]: 

1156 diagnostic["err"]["mu"] = err[idxstart: idxstart + n] 

1157 idxstart = idxstart + n 

1158 pass 

1159 

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 

1168 

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 

1175 

1176 return [p, out.fun, diagnostic]