lisaanalysistools 1.0.10__cp312-cp312-macosx_10_13_x86_64.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.

Potentially problematic release.


This version of lisaanalysistools might be problematic. Click here for more details.

Files changed (43) hide show
  1. lisaanalysistools-1.0.10.dist-info/LICENSE +201 -0
  2. lisaanalysistools-1.0.10.dist-info/METADATA +101 -0
  3. lisaanalysistools-1.0.10.dist-info/RECORD +43 -0
  4. lisaanalysistools-1.0.10.dist-info/WHEEL +5 -0
  5. lisaanalysistools-1.0.10.dist-info/top_level.txt +1 -0
  6. lisatools/__init__.py +0 -0
  7. lisatools/_version.py +4 -0
  8. lisatools/analysiscontainer.py +451 -0
  9. lisatools/cutils/__init__.py +0 -0
  10. lisatools/cutils/detector_cpu.cpython-312-darwin.so +0 -0
  11. lisatools/cutils/include/Detector.hpp +84 -0
  12. lisatools/cutils/include/__init__.py +0 -0
  13. lisatools/cutils/include/global.hpp +28 -0
  14. lisatools/cutils/src/Detector.cpp +307 -0
  15. lisatools/cutils/src/Detector.cu +307 -0
  16. lisatools/cutils/src/__init__.py +0 -0
  17. lisatools/cutils/src/pycppdetector.pyx +255 -0
  18. lisatools/datacontainer.py +309 -0
  19. lisatools/detector.py +704 -0
  20. lisatools/diagnostic.py +977 -0
  21. lisatools/sampling/__init__.py +0 -0
  22. lisatools/sampling/likelihood.py +882 -0
  23. lisatools/sampling/moves/__init__.py +0 -0
  24. lisatools/sampling/moves/skymodehop.py +110 -0
  25. lisatools/sampling/prior.py +646 -0
  26. lisatools/sampling/stopping.py +320 -0
  27. lisatools/sampling/utility.py +411 -0
  28. lisatools/sensitivity.py +899 -0
  29. lisatools/sources/__init__.py +6 -0
  30. lisatools/sources/bbh/__init__.py +1 -0
  31. lisatools/sources/bbh/waveform.py +91 -0
  32. lisatools/sources/defaultresponse.py +36 -0
  33. lisatools/sources/emri/__init__.py +1 -0
  34. lisatools/sources/emri/waveform.py +79 -0
  35. lisatools/sources/gb/__init__.py +1 -0
  36. lisatools/sources/gb/waveform.py +67 -0
  37. lisatools/sources/utils.py +456 -0
  38. lisatools/sources/waveformbase.py +41 -0
  39. lisatools/stochastic.py +291 -0
  40. lisatools/utils/__init__.py +0 -0
  41. lisatools/utils/constants.py +40 -0
  42. lisatools/utils/pointeradjust.py +106 -0
  43. lisatools/utils/utility.py +245 -0
@@ -0,0 +1,977 @@
1
+ from __future__ import annotations
2
+ import warnings
3
+ from typing import Optional, Any, Tuple, List
4
+
5
+ import matplotlib.pyplot as plt
6
+
7
+ from eryn.utils import TransformContainer
8
+
9
+ import numpy as np
10
+
11
+ try:
12
+ import cupy as cp
13
+
14
+ except (ModuleNotFoundError, ImportError):
15
+ import numpy as cp
16
+
17
+ pass
18
+
19
+ from .sensitivity import get_sensitivity, SensitivityMatrix
20
+ from .datacontainer import DataResidualArray
21
+ from .utils.utility import get_array_module
22
+
23
+
24
+ def inner_product(
25
+ sig1: np.ndarray | list | DataResidualArray,
26
+ sig2: np.ndarray | list | DataResidualArray,
27
+ dt: Optional[float] = None,
28
+ df: Optional[float] = None,
29
+ f_arr: Optional[float] = None,
30
+ psd: Optional[str | None | np.ndarray | SensitivityMatrix] = "LISASens",
31
+ psd_args: Optional[tuple] = (),
32
+ psd_kwargs: Optional[dict] = {},
33
+ normalize: Optional[bool | str] = False,
34
+ complex: Optional[bool] = False,
35
+ ) -> float | complex:
36
+ """Compute the inner product between two signals weighted by a psd.
37
+
38
+ The inner product between time series :math:`a(t)` and :math:`b(t)` is
39
+
40
+ .. math::
41
+
42
+ \langle a | b \\rangle = 2\int_{f_\\text{min}}^{f_\\text{max}} \\frac{\\tilde{a}(f)^*\\tilde{b}(f) + \\tilde{a}(f)\\tilde{b}(f)^*}{S_n(f)} df\ \ ,
43
+
44
+ where :math:`\\tilde{a}(f)` is the Fourier transform of :math:`a(t)` and :math:`S_n(f)` is the one-sided Power Spectral Density of the noise.
45
+
46
+ The inner product can be left complex using the ``complex`` kwarg.
47
+
48
+ **GPU Capability**: Pass CuPy arrays rather than NumPy arrays.
49
+
50
+ Args:
51
+ sig1: First signal to use for the inner product.
52
+ Can be time-domain or frequency-domain.
53
+ Must be 1D ``np.ndarray``, list of 1D ``np.ndarray``s, or 2D ``np.ndarray``
54
+ across channels with shape ``(nchannels, data length)``.
55
+ sig2: Second signal to use for the inner product.
56
+ Can be time-domain or frequency-domain.
57
+ Must be 1D ``np.ndarray``, list of 1D ``np.ndarray``s, or 2D ``np.ndarray``
58
+ across channels with shape ``(nchannels, data length)``.
59
+ dt: Time step in seconds. If provided, assumes time-domain signals.
60
+ df: Constant frequency spacing. This will assume a frequency domain signal with constant frequency spacing.
61
+ f_arr: Array of specific frequencies at which the signal is given.
62
+ psd: Indicator of what psd to use. If a ``str``, this will be passed as the ``sens_fn`` kwarg to :func:`get_sensitivity`.
63
+ If ``None``, it will be an array of ones. Or, you can pass a 1D ``np.ndarray`` of psd values that must be the same length
64
+ as the frequency domain signals.
65
+ psd_args: Arguments to pass to the psd function if ``type(psd) == str``.
66
+ psd_kwargs: Keyword arguments to pass to the psd function if ``type(psd) == str``.
67
+ normalize: Normalize the inner product. If ``True``, it will normalize the square root of the product of individual signal inner products.
68
+ You can also pass ``"sig1"`` or ``"sig2"`` to normalize with respect to one signal.
69
+ complex: If ``True``, return the complex value of the inner product rather than just its real-valued part.
70
+
71
+ Returns:
72
+ Inner product value.
73
+
74
+ """
75
+ # initial input checks and setup
76
+ sig1 = DataResidualArray(sig1, dt=dt, f_arr=f_arr, df=df)
77
+ sig2 = DataResidualArray(sig2, dt=dt, f_arr=f_arr, df=df)
78
+
79
+ if sig1.nchannels != sig2.nchannels:
80
+ raise ValueError(
81
+ f"Signal 1 has {sig1.nchannels} channels. Signal 2 has {sig2.nchannels} channels. Must be the same."
82
+ )
83
+
84
+ xp = get_array_module(sig1[0])
85
+
86
+ # checks
87
+ for i in range(sig1.nchannels):
88
+ if not type(sig1[0]) == type(sig1[i]) and type(sig1[0]) == type(sig2[i]):
89
+ raise ValueError(
90
+ "Array in sig1, index 0 sets array module. Not all arrays match that module type (Numpy or Cupy)"
91
+ )
92
+
93
+ if sig1.data_length != sig2.data_length:
94
+ raise ValueError(
95
+ "The two signals are two different lengths. Must be the same length."
96
+ )
97
+
98
+ freqs = sig1.f_arr
99
+
100
+ # get psd weighting
101
+ if not isinstance(psd, SensitivityMatrix):
102
+ psd = SensitivityMatrix(freqs, [psd], *psd_args, **psd_kwargs)
103
+
104
+ operational_sets = []
105
+
106
+ if psd.ndim == 3:
107
+ assert psd.shape[0] == psd.shape[1] == sig1.shape[0] == sig2.shape[0]
108
+
109
+ for i in range(psd.shape[0]):
110
+ for j in range(i, psd.shape[1]):
111
+ factor = 1.0 if i == j else 2.0
112
+ operational_sets.append(
113
+ dict(factor=factor, sig1_ind=i, sig2_ind=j, psd_ind=(i, j))
114
+ )
115
+
116
+ elif psd.ndim == 2 and psd.shape[0] > 1:
117
+ assert psd.shape[0] == sig1.shape[0] == sig2.shape[0]
118
+ for i in range(psd.shape[0]):
119
+ operational_sets.append(dict(factor=1.0, sig1_ind=i, sig2_ind=i, psd_ind=i))
120
+
121
+ elif psd.ndim == 2 and psd.shape[0] == 1:
122
+ for i in range(sig1.shape[0]):
123
+ operational_sets.append(dict(factor=1.0, sig1_ind=i, sig2_ind=i, psd_ind=0))
124
+
125
+ else:
126
+ raise ValueError("# TODO")
127
+
128
+ if complex:
129
+ func = lambda x: x
130
+ else:
131
+ func = xp.real
132
+
133
+ # initialize
134
+ out = 0.0
135
+ x = freqs
136
+
137
+ # account for hp and hx if included in time domain signal
138
+ for op_set in operational_sets:
139
+ factor = op_set["factor"]
140
+
141
+ temp1 = sig1[op_set["sig1_ind"]]
142
+ temp2 = sig2[op_set["sig2_ind"]]
143
+ psd_tmp = psd[op_set["psd_ind"]]
144
+
145
+ ind_start = 1 if np.isnan(psd_tmp[0]) else 0
146
+
147
+ y = (
148
+ func(temp1[ind_start:].conj() * temp2[ind_start:]) / psd_tmp[ind_start:]
149
+ ) # assumes right summation rule
150
+ # df is sunk into trapz
151
+ tmp_out = factor * 4 * xp.trapz(y, x=x[ind_start:])
152
+ out += tmp_out
153
+
154
+ # normalize the inner produce
155
+ normalization_value = 1.0
156
+ if normalize is True:
157
+ norm1 = inner_product(
158
+ sig1,
159
+ sig1,
160
+ psd=psd,
161
+ normalize=False,
162
+ )
163
+ norm2 = inner_product(
164
+ sig2,
165
+ sig2,
166
+ psd=psd,
167
+ normalize=False,
168
+ )
169
+
170
+ normalization_value = np.sqrt(norm1 * norm2)
171
+
172
+ elif isinstance(normalize, str):
173
+ if normalize == "sig1":
174
+ sig_to_normalize = sig1
175
+
176
+ elif normalize == "sig2":
177
+ sig_to_normalize = sig2
178
+
179
+ else:
180
+ raise ValueError(
181
+ "If normalizing with respect to sig1 or sig2, normalize kwarg must either be 'sig1' or 'sig2'."
182
+ )
183
+
184
+ normalization_value = inner_product(
185
+ sig_to_normalize,
186
+ sig_to_normalize,
187
+ psd=psd,
188
+ normalize=False,
189
+ )
190
+
191
+ elif normalize is not False:
192
+ raise ValueError("Normalize must be True, False, 'sig1', or 'sig2'.")
193
+
194
+ out /= normalization_value
195
+ return out
196
+
197
+
198
+ def residual_source_likelihood_term(
199
+ data_res_arr: DataResidualArray, **kwargs: dict
200
+ ) -> float | complex:
201
+ """Calculate the source term in the Likelihood for a data residual (d - h).
202
+
203
+ The source term in the likelihood is given by,
204
+
205
+ .. math::
206
+
207
+ \\log{\\mathcal{L}}_\\text{src} = -\\frac{1}{2}\\langle \\vec{d} - \\vec{h} | \\vec{d} - \\vec{h}\\rangle.
208
+
209
+ Args:
210
+ data_res_arr: Data residual.
211
+ **kwargs: Keyword arguments to pass to :func:`inner_product`.
212
+
213
+ Returns:
214
+ Source term Likelihood value.
215
+
216
+ """
217
+ kwargs["normalize"] = False
218
+ ip_val = inner_product(data_res_arr, data_res_arr, **kwargs)
219
+ return -1 / 2.0 * ip_val
220
+
221
+
222
+ def noise_likelihood_term(psd: SensitivityMatrix) -> float:
223
+ """Calculate the noise term in the Likelihood.
224
+
225
+ The noise term in the likelihood is given by,
226
+
227
+ .. math::
228
+
229
+ \\log{\\mathcal{L}}_n = -\\sum \\log{\\vec{S}_n}.
230
+
231
+ Args:
232
+ psd: Sensitivity information.
233
+
234
+ Returns:
235
+ Noise term Likelihood value.
236
+
237
+ """
238
+ fix = np.isnan(psd[:]) | np.isinf(psd[:])
239
+ assert np.sum(fix) == np.prod(psd.shape[:-1]) or np.sum(fix) == 0
240
+ nl_val = -1.0 / 2.0 * np.sum(np.log(psd[~fix]))
241
+ return nl_val
242
+
243
+
244
+ def residual_full_source_and_noise_likelihood(
245
+ data_res_arr: DataResidualArray,
246
+ psd: str | None | np.ndarray | SensitivityMatrix,
247
+ **kwargs: dict,
248
+ ) -> float | complex:
249
+ """Calculate the full Likelihood including noise and source terms.
250
+
251
+ The noise term is calculated with :func:`noise_likelihood_term`.
252
+
253
+ The source term is calcualted with :func:`residual_source_likelihood_term`.
254
+
255
+ Args:
256
+ data_res_arr: Data residual.
257
+ psd: Sensitivity information.
258
+ **kwargs: Keyword arguments to pass to :func:`inner_product`.
259
+
260
+ Returns:
261
+ Full Likelihood value.
262
+
263
+ """
264
+ if not isinstance(psd, SensitivityMatrix):
265
+ # TODO: maybe adjust so it can take a list just like Sensitivity matrix
266
+ psd = SensitivityMatrix(data_res_arr.f_arr, [psd], **kwargs)
267
+
268
+ # remove key
269
+ for key in "psd", "psd_args", "psd_kwargs":
270
+ if key in kwargs:
271
+ kwargs.pop(key)
272
+
273
+ rslt = residual_source_likelihood_term(data_res_arr, psd=psd, **kwargs)
274
+
275
+ nlt = noise_likelihood_term(psd)
276
+ return nlt + rslt
277
+
278
+
279
+ def data_signal_source_likelihood_term(
280
+ data_arr: DataResidualArray, sig_arr: DataResidualArray, **kwargs: dict
281
+ ) -> float | complex:
282
+ """Calculate the source term in the Likelihood for separate signal and data.
283
+
284
+ The source term in the likelihood is given by,
285
+
286
+ .. math::
287
+
288
+ \\log{\\mathcal{L}}_\\text{src} = -\\frac{1}{2}\\left(\\langle \\vec{d} | \\vec{d}\\rangle + \\langle \\vec{h} | \\vec{h}\\rangle - 2\\langle \\vec{d} | \\vec{h}\\rangle \\right)\ \ .
289
+
290
+ Args:
291
+ data_arr: Data.
292
+ sig_arr: Signal.
293
+ **kwargs: Keyword arguments to pass to :func:`inner_product`.
294
+
295
+ Returns:
296
+ Source term Likelihood value.
297
+
298
+ """
299
+ kwargs["normalize"] = False
300
+ d_h = inner_product(data_arr, sig_arr, **kwargs)
301
+ h_h = inner_product(sig_arr, sig_arr, **kwargs)
302
+ d_d = inner_product(data_arr, data_arr, **kwargs)
303
+ return -1 / 2.0 * (d_d + h_h - 2 * d_h)
304
+
305
+
306
+ def data_signal_full_source_and_noise_likelihood(
307
+ data_arr: DataResidualArray,
308
+ sig_arr: DataResidualArray,
309
+ psd: str | None | np.ndarray | SensitivityMatrix,
310
+ **kwargs: dict,
311
+ ) -> float | complex:
312
+ """Calculate the full Likelihood including noise and source terms.
313
+
314
+ Here, the signal is treated separate from the data.
315
+
316
+ The noise term is calculated with :func:`noise_likelihood_term`.
317
+
318
+ The source term is calcualted with :func:`data_signal_source_likelihood_term`.
319
+
320
+ Args:
321
+ data_arr: Data.
322
+ sig_arr: Signal.
323
+ psd: Sensitivity information.
324
+ **kwargs: Keyword arguments to pass to :func:`inner_product`.
325
+
326
+ Returns:
327
+ Full Likelihood value.
328
+
329
+ """
330
+ if not isinstance(psd, SensitivityMatrix):
331
+ # TODO: maybe adjust so it can take a list just like Sensitivity matrix
332
+ psd = SensitivityMatrix(data_arr.f_arr, [psd], **kwargs)
333
+
334
+ # remove key
335
+ for key in "psd", "psd_args", "psd_kwargs":
336
+ if key in kwargs:
337
+ kwargs.pop(key)
338
+
339
+ rslt = data_signal_source_likelihood_term(data_arr, sig_arr, psd=psd, **kwargs)
340
+
341
+ nlt = noise_likelihood_term(psd)
342
+
343
+ return nlt + rslt
344
+
345
+
346
+ def snr(
347
+ sig1: np.ndarray | list | DataResidualArray,
348
+ *args: Any,
349
+ data: Optional[np.ndarray | list | DataResidualArray] = None,
350
+ **kwargs: Any,
351
+ ) -> float:
352
+ """Compute the snr between two signals weighted by a psd.
353
+
354
+ The signal-to-noise ratio of a signal is :math:`\\sqrt{\\langle a|a\\rangle}`.
355
+
356
+ This will be the optimal SNR if ``data==None``. If a data array is given, it will be the observed
357
+ SNR: :math:`\\langle d|a\\rangle/\\sqrt{\\langle a|a\\rangle}`.
358
+
359
+ **GPU Capability**: Pass CuPy arrays rather than NumPy arrays.
360
+
361
+ Args:
362
+ sig1: Signal to use as the templatefor the SNR.
363
+ Can be time-domain or frequency-domain.
364
+ Must be 1D ``np.ndarray``, list of 1D ``np.ndarray``s, or 2D ``np.ndarray``
365
+ across channels with shape ``(nchannels, data length)``.
366
+ *args: Arguments to pass to :func:`inner_product`.
367
+ data: Data becomes the ``sig2`` argument to :func:`inner_product`.
368
+ **kwargs: Keyword arguments to pass to :func:`inner_product`.
369
+
370
+ Returns:
371
+ Optimal or detected SNR value (depending on ``data`` kwarg).
372
+
373
+ """
374
+
375
+ # get optimal SNR
376
+ opt_snr = np.sqrt(inner_product(sig1, sig1, *args, **kwargs).real)
377
+ if data is None:
378
+ return opt_snr
379
+
380
+ else:
381
+ # if inputed data, calculate detected SNR
382
+ det_snr = inner_product(sig1, data, *args, **kwargs).real / opt_snr
383
+ return det_snr
384
+
385
+
386
+ def h_var_p_eps(
387
+ step: float,
388
+ waveform_model: callable,
389
+ params: np.ndarray | list,
390
+ index: int,
391
+ parameter_transforms: Optional[TransformContainer] = None,
392
+ waveform_args: Optional[tuple] = (),
393
+ waveform_kwargs: Optional[dict] = {},
394
+ ) -> np.ndarray: # TODO: check this
395
+ """Calculate the waveform with a perturbation step of the variable V[i]
396
+
397
+ Args:
398
+ waveform_model: Callable function to the waveform generator with signature ``(*params, **waveform_kwargs)``.
399
+ params: Source parameters that are over derivatives (not in fill dict of parameter transforms)
400
+ step: Absolute step size for variable of interest.
401
+ index: Index to parameter of interest.
402
+ parameter_transforms: `TransformContainer <https://mikekatz04.github.io/Eryn/html/user/utils.html#eryn.utils.TransformContainer>`_ object to transform from the derivative parameter basis
403
+ to the waveform parameter basis. This class can also fill in fixed parameters where the derivatives are not being taken.
404
+ waveform_args: args (beyond parameters) for the waveform generator.
405
+ waveform_kwargs: kwargs for the waveform generation.
406
+
407
+ Returns:
408
+ Perturbation to the waveform in the given parameter. Will always be 2D array with shape ``(num channels, data length)``
409
+
410
+ """
411
+ params_p_eps = params.copy()
412
+ params_p_eps[index] += step
413
+
414
+ if parameter_transforms is not None:
415
+ # transform
416
+ params_p_eps = parameter_transforms.transform_base_parameters(
417
+ params_p_eps[None, :]
418
+ )[0]
419
+
420
+ args_in = tuple(params_p_eps) + tuple(waveform_args)
421
+ dh = waveform_model(*args_in, **waveform_kwargs)
422
+
423
+ # adjust output based on waveform model output
424
+ # needs to be 2D array
425
+ if (isinstance(dh, np.ndarray) or isinstance(dh, cp.ndarray)) and dh.ndim == 1:
426
+ xp = get_array_module(dh)
427
+ dh = xp.atleast_2d(dh)
428
+
429
+ elif isinstance(dh, list):
430
+ xp = get_array_module(dh[0])
431
+ dh = xp.asarray(dh)
432
+
433
+ return dh
434
+
435
+
436
+ def dh_dlambda(
437
+ eps: float,
438
+ *args: tuple,
439
+ more_accurate: Optional[bool] = True,
440
+ **kwargs: dict,
441
+ ) -> np.ndarray:
442
+ """Derivative of the waveform
443
+
444
+ Calculate the derivative of the waveform with precision of order (step^4)
445
+ with respect to the variable V in the i direction.
446
+
447
+ Args:
448
+ eps: Absolute **derivative** step size for variable of interest.
449
+ *args: Arguments passed to :func:`h_var_p_eps`.
450
+ more_accurate: If ``True``, run a more accurate derivate requiring 2x more waveform generations.
451
+ **kwargs: Keyword arguments passed to :func:`h_var_p_eps`.
452
+
453
+ Returns:
454
+ Numerical derivative of the waveform with respect to a varibale of interest. Will be 2D array
455
+ with shape ``(num channels, data length)``.
456
+
457
+ """
458
+ if more_accurate:
459
+ # Derivative of the Waveform
460
+ # up
461
+ h_I_up_2eps = h_var_p_eps(2 * eps, *args, **kwargs)
462
+ h_I_up_eps = h_var_p_eps(eps, *args, **kwargs)
463
+ # down
464
+ h_I_down_2eps = h_var_p_eps(-2 * eps, *args, **kwargs)
465
+ h_I_down_eps = h_var_p_eps(-eps, *args, **kwargs)
466
+
467
+ # make sure they are all the same length
468
+ ind_max = np.min(
469
+ [
470
+ h_I_up_2eps.shape[1],
471
+ h_I_up_eps.shape[1],
472
+ h_I_down_2eps.shape[1],
473
+ h_I_down_eps.shape[1],
474
+ ]
475
+ )
476
+
477
+ # error scales as eps^4
478
+ dh_I = (
479
+ -h_I_up_2eps[:, :ind_max]
480
+ + h_I_down_2eps[:, :ind_max]
481
+ + 8 * (h_I_up_eps[:, :ind_max] - h_I_down_eps[:, :ind_max])
482
+ ) / (12 * eps)
483
+
484
+ else:
485
+ # Derivative of the Waveform
486
+ # up
487
+ h_I_up_eps = h_var_p_eps(eps, *args, **kwargs)
488
+ # down
489
+ h_I_down_eps = h_var_p_eps(-eps, *args, **kwargs)
490
+
491
+ ind_max = np.min([h_I_up_eps.shape[1], h_I_down_eps.shape[1]])
492
+
493
+ # TODO: check what error scales as.
494
+ dh_I = (h_I_up_eps[:, :ind_max] - h_I_down_eps[:, :ind_max]) / (2 * eps)
495
+ # Time thta it takes for one variable: approx 5 minutes
496
+
497
+ return dh_I
498
+
499
+
500
+ def info_matrix(
501
+ eps: float | np.ndarray,
502
+ waveform_model: callable,
503
+ params: np.ndarray | list,
504
+ deriv_inds: Optional[np.ndarray | list] = None,
505
+ inner_product_kwargs: Optional[dict] = {},
506
+ return_derivs: Optional[bool] = False,
507
+ **kwargs: dict,
508
+ ) -> np.ndarray | Tuple[np.ndarray, list]:
509
+ """Calculate Information Matrix.
510
+
511
+ This calculates the information matrix for a given waveform model at a given set of parameters.
512
+ The inverse of the information matrix gives the covariance matrix.
513
+
514
+ This is also referred to as the Fisher information matrix, but @MichaelKatz has chosen to leave out the name because of `this <https://www.newstatesman.com/long-reads/2020/07/ra-fisher-and-science-hatred>`_.
515
+
516
+ The info matrix is given by:
517
+
518
+ .. math::
519
+
520
+ M_{ij} = \\langle h_i | h_j \\rangle \\text{ with } h_i = \\frac{\\partial h}{\\partial \\lambda_i}.
521
+
522
+ Args:
523
+ eps: Absolute **derivative** step size for variable of interest. Can be provided as a ``float`` value that applies
524
+ to all variables or an array, one for each parameter being evaluated in the information matrix.
525
+ waveform_model: Callable function to the waveform generator with signature ``(*params, **waveform_kwargs)``.
526
+ params: Source parameters.
527
+ deriv_inds: Subset of parameters of interest for which to calculate the information matrix, by index.
528
+ If ``None``, it will be ``np.arange(len(params))``.
529
+ inner_product_kwargs: Keyword arguments for the inner product function.
530
+ return_derivs: If ``True``, also returns computed numerical derivatives.
531
+ **kwargs: Keyword arguments passed to :func:`dh_dlambda`.
532
+
533
+ Returns:
534
+ If ``not return_derivs``, this will be the information matrix as a numpy array. If ``return_derivs is True``,
535
+ it will be a tuple with the first entry as the information matrix and the second entry as the partial derivatives.
536
+
537
+ """
538
+
539
+ # setup initial information
540
+ num_params = len(params)
541
+
542
+ if deriv_inds is None:
543
+ deriv_inds = np.arange(num_params)
544
+
545
+ num_info_params = len(deriv_inds)
546
+
547
+ if isinstance(eps, float):
548
+ eps = np.full_like(params, eps)
549
+
550
+ # collect derivatives over the variables of interest
551
+ dh = []
552
+ for i, eps_i in zip(deriv_inds, eps):
553
+ # derivative up
554
+ temp = dh_dlambda(eps_i, waveform_model, params, i, **kwargs)
555
+ dh.append(temp)
556
+
557
+ # calculate the components of the symmetric matrix
558
+ info = np.zeros((num_info_params, num_info_params))
559
+ for i in range(num_info_params):
560
+ for j in range(i, num_info_params):
561
+ info[i][j] = inner_product(
562
+ [dh[i][k] for k in range(len(dh[i]))],
563
+ [dh[j][k] for k in range(len(dh[i]))],
564
+ **inner_product_kwargs,
565
+ )
566
+ info[j][i] = info[i][j]
567
+
568
+ if return_derivs:
569
+ return info, dh
570
+ else:
571
+ return info
572
+
573
+
574
+ def covariance(
575
+ *info_mat_args: tuple,
576
+ info_mat: Optional[np.ndarray] = None,
577
+ diagonalize: Optional[bool] = False,
578
+ return_info_mat: Optional[bool] = False,
579
+ precision: Optional[bool] = False,
580
+ **info_mat_kwargs: dict,
581
+ ) -> np.ndarray | list:
582
+ """Calculate covariance matrix for a set of EMRI parameters, computing the information matrix if not supplied.
583
+
584
+ Args:
585
+ *info_mat_args: Set of arguments to pass to :func:`info_matrix`. Not required if ``info_mat`` is not ``None``.
586
+ info_mat: Pre-computed information matrix. If supplied, this matrix will be inverted.
587
+ diagonalize: If ``True``, diagonalizes the covariance matrix.
588
+ return_info_mat: If ``True``, also returns the computed information matrix.
589
+ precision: If ``True``, uses 500-dps precision to compute the information matrix inverse (requires `mpmath <https://mpmath.org>`_). This is typically a good idea as the information matrix can be highly ill-conditioned.
590
+ **info_mat_kwargs: Keyword arguments to pass to :func:`info_matrix`.
591
+
592
+ Returns:
593
+ Covariance matrix. If ``return_info_mat is True``. A list will be returned with the covariance as the first
594
+ entry and the information matrix as the second entry. If ``return_derivs is True`` (keyword argument to :func:`info_matrix`),
595
+ then another entry will be added to the list for the derivatives.
596
+
597
+ """
598
+
599
+ if info_mat is None:
600
+ info_mat = info_matrix(*info_mat_args, **info_mat_kwargs)
601
+
602
+ # parse output properly
603
+ if "return_derivs" in info_mat_kwargs and info_mat_kwargs["return_derivs"]:
604
+ return_derivs = True
605
+ info_mat, dh = info_mat
606
+
607
+ else:
608
+ return_derivs = False
609
+
610
+ # attempt to import and setup mpmath if precision required
611
+ if precision:
612
+ try:
613
+ import mpmath as mp
614
+
615
+ mp.mp.dps = 500
616
+ except ModuleNotFoundError:
617
+ print("mpmath module not installed. Defaulting to low precision...")
618
+ precision = False
619
+
620
+ if precision:
621
+ hp_info_mat = mp.matrix(info_mat.tolist())
622
+ U, S, V = mp.svd_r(hp_info_mat) # singular value decomposition
623
+ temp = mp.diag([val ** (-1) for val in S]) # get S**-1
624
+ temp2 = V.T * temp * U.T # construct pseudo-inverse
625
+ cov = np.array(temp2.tolist(), dtype=np.float64)
626
+ else:
627
+ cov = np.linalg.pinv(info_mat)
628
+
629
+ if diagonalize:
630
+ # get eigeninformation
631
+ eig_vals, eig_vecs = get_eigeninfo(cov, high_precision=precision)
632
+
633
+ # diagonal cov now
634
+ cov = np.dot(np.dot(eig_vecs.T, cov), eig_vecs)
635
+
636
+ # just requesting covariance
637
+ if True not in [return_info_mat, return_derivs]:
638
+ return cov
639
+
640
+ # if desiring more information, create a list capable of variable size
641
+ returns = [cov]
642
+
643
+ # add information matrix
644
+ if return_info_mat:
645
+ returns.append(info_mat)
646
+
647
+ # add derivatives
648
+ if return_derivs:
649
+ returns.append(dh)
650
+
651
+ return returns
652
+
653
+
654
+ def plot_covariance_corner(
655
+ params: np.ndarray,
656
+ cov: np.ndarray,
657
+ nsamp: Optional[int] = 25000,
658
+ fig: Optional[plt.Figure] = None,
659
+ **kwargs: dict,
660
+ ) -> plt.Figure:
661
+ """Construct a corner plot for a given covariance matrix.
662
+
663
+ The `corner <https://corner.readthedocs.io/en/latest/>`_ module is required for this.
664
+
665
+ Args:
666
+ params: The set of parameters used for the event (the mean vector of the covariance matrix).
667
+ cov: Covariance matrix from which to construct the corner plot.
668
+ nsamp: Number of samples to draw from the multivariate distribution.
669
+ fig: Matplotlib :class:`plt.Figure` object. Use this if passing an existing corner plot figure.
670
+ **kwargs: Keyword arguments for the corner plot - see the module documentation for more info.
671
+
672
+ Returns:
673
+ The corner plot figure.
674
+
675
+ """
676
+
677
+ # TODO: add capability for ChainConsumer?
678
+ try:
679
+ import corner
680
+ except ModuleNotFoundError:
681
+ raise ValueError(
682
+ "Attempting to plot using the corner module, but it is not installed."
683
+ )
684
+
685
+ # generate fake samples from the covariance distribution
686
+ samp = np.random.multivariate_normal(params, cov, size=nsamp)
687
+
688
+ # make corner plot
689
+ fig = corner.corner(samp, **kwargs)
690
+ return fig
691
+
692
+
693
+ def plot_covariance_contour(
694
+ params: np.ndarray,
695
+ cov: np.ndarray,
696
+ horizontal_index: int,
697
+ vertical_index: int,
698
+ nsamp: Optional[int] = 25000,
699
+ ax: Optional[plt.Axes] = None,
700
+ **kwargs: dict,
701
+ ) -> plt.Axes | Tuple[plt.Figure, plt.Axes]:
702
+ """Construct a contour plot for a given covariance matrix on a single axis object.
703
+
704
+ The `corner <https://corner.readthedocs.io/en/latest/>`_ module is required for this.
705
+
706
+ Args:
707
+ params: The set of parameters used for the event (the mean vector of the covariance matrix).
708
+ cov: Covariance matrix from which to construct the corner plot.
709
+ horizontal_index: Parameter index to plot along the horizontal axis of the contour plot.
710
+ vertical_index: Parameter index to plot along the vertical axis of the contour plot.
711
+ nsamp: Number of samples to draw from the multivariate distribution.
712
+ fig: Matplotlib :class:`plt.Figure` object. Use this if passing an existing corner plot figure.
713
+ **kwargs: Keyword arguments for the corner plot - see the module documentation for more info.
714
+
715
+ Returns:
716
+ If ``ax`` is provided, the return will be that ax object. If it is not provided, then a
717
+ Matplotlib Figure and Axes obejct is created and returned as a tuple: ``(plt.Figure, plt.Axes)``.
718
+
719
+ """
720
+
721
+ # TODO: add capability for ChainConsumer?
722
+ try:
723
+ import corner
724
+ except ModuleNotFoundError:
725
+ raise ModuleNotFoundError(
726
+ "Attempting to plot using the corner module, but it is not installed."
727
+ )
728
+
729
+ if ax is None:
730
+ fig, ax = plt.subplots(1, 1)
731
+ else:
732
+ fig = None
733
+
734
+ # generate fake samples from the covariance distribution
735
+ samp = np.random.multivariate_normal(params, cov, size=nsamp)
736
+
737
+ x = samp[:, horizontal_index]
738
+ y = samp[:, vertical_index]
739
+
740
+ # make corner plot
741
+ corner.hist2d(x, y, ax=ax, **kwargs)
742
+
743
+ if fig is None:
744
+ return ax
745
+ else:
746
+ return (fig, ax)
747
+
748
+
749
+ def get_eigeninfo(
750
+ arr: np.ndarray, high_precision: Optional[bool] = False
751
+ ) -> Tuple[np.ndarray, np.ndarray]:
752
+ """Performs eigenvalue decomposition and returns the eigenvalues and right-eigenvectors for the supplied fisher/covariance matrix.
753
+
754
+ Args:
755
+ arr: Input matrix for which to perform eigenvalue decomposition.
756
+ high_precision: If ``True``, use 500-dps precision to ensure accurate eigenvalue decomposition
757
+ (requires `mpmath <https://mpmath.org>`_ to be installed). Defaults to False.
758
+
759
+ Returns:
760
+ Tuple containing Eigenvalues and right-Eigenvectors for the supplied array, constructed such that evects[:,k] corresponds to evals[k].
761
+
762
+
763
+ """
764
+
765
+ if high_precision:
766
+ try:
767
+ import mpmath as mp
768
+
769
+ mp.mp.dps = 500
770
+ except ModuleNotFoundError:
771
+ print("mpmath is not installed - using low-precision eigen decomposition.")
772
+ high_precision = False
773
+
774
+ if high_precision:
775
+ hp_arr = mp.matrix(arr.tolist())
776
+ # get eigenvectors
777
+ E, EL, ER = mp.eig(hp_arr, left=True, right=True)
778
+
779
+ # convert back
780
+ evals = np.array(E, dtype=np.float64)
781
+ evects = np.array(ER.tolist(), dtype=np.float64)
782
+
783
+ return evals, evects
784
+
785
+ else:
786
+ evals, evects = np.linalg.eig(arr)
787
+ return evals, evects
788
+
789
+
790
+ def cutler_vallisneri_bias(
791
+ waveform_model_true: callable,
792
+ waveform_model_approx: callable,
793
+ params: np.ndarray,
794
+ eps: float | np.ndarray,
795
+ input_diagnostics: Optional[dict] = None,
796
+ info_mat: Optional[np.ndarray] = None,
797
+ deriv_inds: Optional[np.ndarray | list] = None,
798
+ return_derivs: Optional[bool] = False,
799
+ return_cov: Optional[bool] = False,
800
+ parameter_transforms: Optional[TransformContainer] = None,
801
+ waveform_true_args: Optional[tuple] = (),
802
+ waveform_true_kwargs: Optional[dict] = {},
803
+ waveform_approx_args: Optional[tuple] = (),
804
+ waveform_approx_kwargs: Optional[dict] = {},
805
+ inner_product_kwargs: Optional[dict] = {},
806
+ ) -> list:
807
+ """Calculate the Cutler-Vallisneri bias.
808
+
809
+ # TODO: add basic math
810
+
811
+ Args:
812
+ waveform_model_true: Callable function to the **true** waveform generator with signature ``(*params, **waveform_kwargs)``.
813
+ waveform_model_approx: Callable function to the **approximate** waveform generator with signature ``(*params, **waveform_kwargs)``.
814
+ params: Source parameters.
815
+ eps: Absolute **derivative** step size. See :func:`info_matrix`.
816
+ input_diagnostics: Dictionary including the diagnostic information if it is precomputed. Dictionary must include
817
+ keys ``"cov"`` (covariance matrix, output of :func:`covariance`), ``"h_true"`` (the **true** waveform),
818
+ and ``"dh"`` (derivatives of the waveforms, list of outputs from :func:`dh_dlambda`).
819
+ info_mat: Pre-computed information matrix. If supplied, this matrix will be inverted to find the covariance.
820
+ deriv_inds: Subset of parameters of interest. See :func:`info_matrix`.
821
+ return_derivs: If ``True``, also returns computed numerical derivatives.
822
+ return_cov: If ``True``, also returns computed covariance matrix.
823
+ parameter_transforms: `TransformContainer <https://mikekatz04.github.io/Eryn/html/user/utils.html#eryn.utils.TransformContainer>`_ object. See :func:`info_matrix`.
824
+ waveform_true_args: Arguments for the **true** waveform generator.
825
+ waveform_true_kwargs: Keyword arguments for the **true** waveform generator.
826
+ waveform_approx_args: Arguments for the **approximate** waveform generator.
827
+ waveform_approx_kwargs: Keyword arguments for the **approximate** waveform generator.
828
+ inner_product_kwargs: Keyword arguments for the inner product function.
829
+
830
+ Returns:
831
+ List of return information. By default, it is ``[systematic error, bias]``.
832
+ If ``return_derivs`` or ``return_cov`` are ``True``, they will be added to the list with derivs added before covs.
833
+
834
+ """
835
+
836
+ if deriv_inds is None:
837
+ deriv_inds = np.arange(len(params))
838
+
839
+ if info_mat is not None and input_diagnostics is not None:
840
+ warnings.warn(
841
+ "Provided info_mat and input_diagnostics kwargs. Ignoring info_mat."
842
+ )
843
+
844
+ # adjust parameters to waveform basis
845
+ params_in = parameter_transforms.transform_base_parameters(params.copy())
846
+
847
+ if input_diagnostics is None:
848
+ # get true waveform
849
+ h_true = waveform_model_true(
850
+ *(tuple(params_in) + tuple(waveform_true_args)), **waveform_true_kwargs
851
+ )
852
+
853
+ # get covariance info and waveform derivatives
854
+ cov, dh = covariance(
855
+ eps,
856
+ waveform_model_true,
857
+ params,
858
+ return_derivs=True,
859
+ deriv_inds=deriv_inds,
860
+ info_mat=info_mat,
861
+ parameter_transforms=parameter_transforms,
862
+ waveform_args=waveform_true_args,
863
+ waveform_kwargs=waveform_true_kwargs,
864
+ inner_product_kwargs=inner_product_kwargs,
865
+ )
866
+
867
+ else:
868
+ # pre-computed info
869
+ cov = input_diagnostics["cov"]
870
+ h_true = input_diagnostics["h_true"]
871
+ dh = input_diagnostics["dh"]
872
+
873
+ # get approximate waveform
874
+ h_approx = waveform_model_approx(
875
+ *(tuple(params_in) + tuple(waveform_approx_args)), **waveform_approx_kwargs
876
+ )
877
+
878
+ # adjust/check waveform outputs
879
+ if isinstance(h_true, np.ndarray) and h_true.ndim == 1:
880
+ h_true = [h_true]
881
+ elif isinstance(h_true, np.ndarray) and h_true.ndim == 2:
882
+ h_true = list(h_true)
883
+
884
+ if isinstance(h_approx, np.ndarray) and h_approx.ndim == 1:
885
+ h_approx = [h_approx]
886
+ elif isinstance(h_approx, np.ndarray) and h_approx.ndim == 2:
887
+ h_approx = list(h_approx)
888
+
889
+ assert len(h_approx) == len(h_true)
890
+ assert np.all(
891
+ np.asarray([len(h_approx[i]) == len(h_true[i]) for i in range(len(h_true))])
892
+ )
893
+
894
+ # difference in the waveforms
895
+ diff = [h_true[i] - h_approx[i] for i in range(len(h_approx))]
896
+
897
+ # systematic err
898
+ syst_vec = np.array(
899
+ [
900
+ inner_product(
901
+ dh[k],
902
+ diff,
903
+ **inner_product_kwargs,
904
+ )
905
+ for k in range(len(deriv_inds))
906
+ ]
907
+ )
908
+
909
+ bias = np.dot(cov, syst_vec)
910
+
911
+ # return list
912
+ returns = [syst_vec, bias]
913
+
914
+ # add anything requested
915
+ if return_cov:
916
+ returns.append(cov)
917
+
918
+ if return_derivs:
919
+ returns.append(dh)
920
+
921
+ return returns
922
+
923
+
924
+ def scale_to_snr(
925
+ target_snr: float,
926
+ sig: np.ndarray | list,
927
+ *snr_args: tuple,
928
+ return_orig_snr: Optional[bool] = False,
929
+ **snr_kwargs: dict,
930
+ ) -> np.ndarray | list | Tuple[np.ndarray | list, float]:
931
+ """Calculate the SNR and scale a signal.
932
+
933
+ Args:
934
+ target_snr: Desired SNR value for the injected signal.
935
+ sig: Signal to adjust. A copy will be made.
936
+ Can be time-domain or frequency-domain.
937
+ Must be 1D ``np.ndarray``, list of 1D ``np.ndarray``s, or 2D ``np.ndarray``
938
+ across channels with shape ``(nchannels, data length)``.
939
+ *snr_args: Arguments to pass to :func:`snr`.
940
+ return_orig_snr: If ``True``, return the original SNR in addition to the adjusted data.
941
+ **snr_kwargs: Keyword arguments to pass to :func:`snr`.
942
+
943
+ Returns:
944
+ Returns the copied input signal adjusted to the target SNR. If ``return_orig_snr is True``, the original
945
+ SNR is added as the second entry of a tuple with the adjusted signal (as the first entry in the tuple).
946
+
947
+ """
948
+ # get the snr and adjustment factor
949
+ snr_out = snr(sig, *snr_args, **snr_kwargs)
950
+ factor = target_snr / snr_out
951
+
952
+ # any changes back to the original signal type
953
+ back_single = False
954
+ back_2d_array = False
955
+
956
+ if isinstance(sig, list) is False and sig.ndim == 1:
957
+ sig = [sig]
958
+ back_single = True
959
+
960
+ elif isinstance(sig, list) is False and sig.ndim == 2:
961
+ sig = list(sig)
962
+ back_2d_array = True
963
+
964
+ # adjust
965
+ out = [sig_i * factor for sig_i in sig]
966
+
967
+ # adjust type back to the input type
968
+ if back_2d_array:
969
+ out = np.asarray(out)
970
+
971
+ if back_single:
972
+ out = out[0]
973
+
974
+ if return_orig_snr:
975
+ return (out, snr_out)
976
+
977
+ return out