fastlisaresponse 1.1.14__cp311-cp311-manylinux_2_17_x86_64.manylinux2014_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.
@@ -0,0 +1,813 @@
1
+ from multiprocessing.sharedctypes import Value
2
+ import numpy as np
3
+ from typing import Optional, List
4
+ import warnings
5
+ from typing import Tuple
6
+ from copy import deepcopy
7
+
8
+ import time
9
+ import h5py
10
+
11
+ from scipy.interpolate import CubicSpline
12
+
13
+ from lisatools.detector import EqualArmlengthOrbits, Orbits
14
+ from lisatools.utils.utility import AET
15
+
16
+ from .utils.parallelbase import FastLISAResponseParallelModule
17
+
18
+
19
+ # TODO: need to update constants setup
20
+ YRSID_SI = 31558149.763545603
21
+
22
+
23
+ def get_factorial(n):
24
+ fact = 1
25
+
26
+ for i in range(1, n + 1):
27
+ fact = fact * i
28
+
29
+ return fact
30
+
31
+
32
+ from math import factorial
33
+
34
+ factorials = np.array([factorial(i) for i in range(30)])
35
+
36
+ C_inv = 3.3356409519815204e-09
37
+
38
+
39
+ class pyResponseTDI(FastLISAResponseParallelModule):
40
+ """Class container for fast LISA response function generation.
41
+
42
+ The class computes the generic time-domain response function for LISA.
43
+ It takes LISA constellation orbital information as input and properly determines
44
+ the response for these orbits numerically. This includes both the projection
45
+ of the gravitational waves onto the LISA constellation arms and combinations \
46
+ of projections into TDI observables. The methods and maths used can be found
47
+ [here](https://arxiv.org/abs/2204.06633).
48
+
49
+ This class is also GPU-accelerated, which is very helpful for Bayesian inference
50
+ methods.
51
+
52
+ Args:
53
+ sampling_frequency (double): The sampling rate in Hz.
54
+ num_pts (int): Number of points to produce for the final output template.
55
+ order (int, optional): Order of Lagrangian interpolation technique. Lower orders
56
+ will be faster. The user must make sure the order is sufficient for the
57
+ waveform being used. (default: 25)
58
+ tdi (str or list, optional): TDI setup. Currently, the stock options are
59
+ :code:`'1st generation'` and :code:`'2nd generation'`. Or the user can provide
60
+ a list of tdi_combinations of the form
61
+ :code:`{"link": 12, "links_for_delay": [21, 13, 31], "sign": 1, "type": "delay"}`.
62
+ :code:`'link'` (`int`) the link index (12, 21, 13, 31, 23, 32) for the projection (:math:`y_{ij}`).
63
+ :code:`'links_for_delay'` (`list`) are the link indexes as a list used for delays
64
+ applied to the link projections.
65
+ ``'sign'`` is the sign in front of the contribution to the TDI observable. It takes the value of `+1` or `-1`.
66
+ ``type`` is either ``"delay"`` or ``"advance"``. It is optional and defaults to ``"delay"``.
67
+ (default: ``"1st generation"``)
68
+ orbits (:class:`Orbits`, optional): Orbits class from LISA Analysis Tools. Works with LISA Orbits
69
+ outputs: ``lisa-simulation.pages.in2p3.fr/orbits/``.
70
+ (default: :class:`EqualArmlengthOrbits`)
71
+ tdi_chan (str, optional): Which TDI channel combination to return. Choices are :code:`'XYZ'`,
72
+ :code:`AET`, or :code:`AE`. (default: :code:`'XYZ'`)
73
+ tdi_orbits (:class:`Orbits`, optional): Set if different orbits from projection.
74
+ Orbits class from LISA Analysis Tools. Works with LISA Orbits
75
+ outputs: ``lisa-simulation.pages.in2p3.fr/orbits/``.
76
+ (default: :class:`EqualArmlengthOrbits`)
77
+ force_backend (str, optional): If given, run this class on the requested backend.
78
+ Options are ``"cpu"``, ``"cuda11x"``, ``"cuda12x"``. (default: ``None``)
79
+
80
+ Attributes:
81
+ A_in (xp.ndarray): Array containing y values for linear spline of A
82
+ during Lagrangian interpolation.
83
+ buffer_integer (int): Self-determined buffer necesary for the given
84
+ value for :code:`order`.
85
+ channels_no_delays (2D np.ndarray): Carrier of link index and sign information
86
+ for arms that do not get delayed during TDI computation.
87
+ deps (double): The spacing between Epsilon values in the interpolant
88
+ for the A quantity in Lagrangian interpolation. Hard coded to
89
+ 1/(:code:`num_A` - 1).
90
+ dt (double): Inverse of the sampling_frequency.
91
+ E_in (xp.ndarray): Array containing y values for linear spline of E
92
+ during Lagrangian interpolation.
93
+ half_order (int): Half of :code:`order` adjusted to be :code:`int`.
94
+ link_inds (xp.ndarray): Link indexes for delays in TDI.
95
+ link_space_craft_0_in (xp.ndarray): Link indexes for receiver on each
96
+ arm of the LISA constellation.
97
+ link_space_craft_1_in (xp.ndarray): Link indexes for emitter on each
98
+ arm of the LISA constellation.
99
+ nlinks (int): The number of links in the constellation. Typically 6.
100
+ num_A (int): Number of points to use for A spline values used in the Lagrangian
101
+ interpolation. This is hard coded to 1001.
102
+ num_channels (int): 3.
103
+ num_pts (int): Number of points to produce for the final output template.
104
+ order (int): Order of Lagrangian interpolation technique.
105
+ sampling_frequency (double): The sampling rate in Hz.
106
+ tdi (str or list): TDI setup.
107
+ tdi_buffer (int): The buffer necessary for all information needed at early times
108
+ for the TDI computation. This is set to 200.
109
+ xp (obj): Either Numpy or Cupy.
110
+
111
+ """
112
+
113
+ def __init__(
114
+ self,
115
+ sampling_frequency,
116
+ num_pts,
117
+ order=25,
118
+ tdi="1st generation",
119
+ orbits: Optional[Orbits] = EqualArmlengthOrbits,
120
+ tdi_orbits: Optional[Orbits] = None,
121
+ tdi_chan="XYZ",
122
+ force_backend=None,
123
+ ):
124
+
125
+ # setup all quantities
126
+ self.sampling_frequency = sampling_frequency
127
+ self.dt = 1 / sampling_frequency
128
+ self.tdi_buffer = 200
129
+
130
+ self.num_pts = num_pts
131
+
132
+ # Lagrangian interpolation setup
133
+ self.order = order
134
+ self.buffer_integer = self.order * 2 + 1
135
+ self.half_order = int((order + 1) / 2)
136
+
137
+ # setup TDI information
138
+ self.tdi = tdi
139
+ self.tdi_chan = tdi_chan
140
+
141
+ super().__init__(force_backend=force_backend)
142
+
143
+ # prepare the interpolation of A and E in the Lagrangian interpolation
144
+ self._fill_A_E()
145
+
146
+ # setup orbits
147
+ self.response_orbits = orbits
148
+
149
+ if tdi_orbits is None:
150
+ tdi_orbits = self.response_orbits
151
+
152
+ self.tdi_orbits = tdi_orbits
153
+
154
+ if self.num_pts * self.dt > self.response_orbits.t_base.max():
155
+ warnings.warn(
156
+ "Input number of points is longer in time than available orbital information. Trimming to fit orbital information."
157
+ )
158
+ self.num_pts = int(self.response_orbits.t_base.max() / self.dt)
159
+
160
+ # setup spacecraft links indexes
161
+
162
+ # setup TDI info
163
+ self._init_TDI_delays()
164
+
165
+ @property
166
+ def response_gen(self) -> callable:
167
+ """CPU/GPU function for generating the projections."""
168
+ return self.backend.get_response_wrap
169
+
170
+ @property
171
+ def tdi_gen(self) -> callable:
172
+ """CPU/GPU function for generating tdi."""
173
+ return self.backend.get_tdi_delays_wrap
174
+
175
+ @property
176
+ def pycppDetector_fastlisa(self):
177
+ return self.backend.pycppDetector_fastlisa
178
+
179
+ @property
180
+ def xp(self) -> object:
181
+ return self.backend.xp
182
+
183
+ @property
184
+ def response_orbits(self) -> Orbits:
185
+ """Response function orbits."""
186
+ return self._response_orbits
187
+
188
+ @response_orbits.setter
189
+ def response_orbits(self, orbits: Orbits) -> None:
190
+ """Set response orbits."""
191
+
192
+ if orbits is None:
193
+ orbits = EqualArmlengthOrbits()
194
+
195
+ assert isinstance(orbits, Orbits)
196
+
197
+ self._response_orbits = deepcopy(orbits)
198
+
199
+ if not self._response_orbits.configured:
200
+ self._response_orbits.configure(linear_interp_setup=True)
201
+
202
+ @property
203
+ def tdi_orbits(self) -> Orbits:
204
+ """TDI function orbits."""
205
+ return self._tdi_orbits
206
+
207
+ @tdi_orbits.setter
208
+ def tdi_orbits(self, orbits: Orbits) -> None:
209
+ """Set TDI orbits."""
210
+
211
+ if orbits is None:
212
+ orbits = EqualArmlengthOrbits()
213
+
214
+ assert isinstance(orbits, Orbits)
215
+ assert orbits.backend.name.split("_")[-1] == self.backend.name.split("_")[-1]
216
+
217
+ self._tdi_orbits = deepcopy(orbits)
218
+
219
+ if not self._tdi_orbits.configured:
220
+ self._tdi_orbits.configure(linear_interp_setup=True)
221
+
222
+ @property
223
+ def citation(self):
224
+ """Get citations for use of this code"""
225
+
226
+ return """
227
+ # TODO add
228
+ """
229
+
230
+ @classmethod
231
+ def supported_backends(cls):
232
+ return ["fastlisaresponse_" + _tmp for _tmp in cls.GPU_RECOMMENDED()]
233
+
234
+ def _fill_A_E(self):
235
+ """Set up A and E terms inside the Lagrangian interpolant"""
236
+
237
+ factorials = np.asarray([float(get_factorial(n)) for n in range(40)])
238
+
239
+ # base quantities for linear interpolant over A
240
+ self.num_A = 1001
241
+ self.deps = 1.0 / (self.num_A - 1)
242
+
243
+ eps = np.arange(self.num_A) * self.deps
244
+
245
+ h = self.half_order
246
+
247
+ denominator = factorials[h - 1] * factorials[h]
248
+
249
+ # prepare A
250
+ A_in = np.zeros_like(eps)
251
+ for j, eps_i in enumerate(eps):
252
+ A = 1.0
253
+ for i in range(1, h):
254
+ A *= (i + eps_i) * (i + 1 - eps_i)
255
+
256
+ A /= denominator
257
+ A_in[j] = A
258
+
259
+ self.A_in = self.xp.asarray(A_in)
260
+
261
+ # prepare E
262
+ E_in = self.xp.zeros((self.half_order,))
263
+
264
+ for j in range(1, self.half_order):
265
+ first_term = factorials[h - 1] / factorials[h - 1 - j]
266
+ second_term = factorials[h] / factorials[h + j]
267
+ value = first_term * second_term
268
+ value = value * (-1.0) ** j
269
+ E_in[j - 1] = value
270
+
271
+ self.E_in = self.xp.asarray(E_in)
272
+
273
+ def _init_TDI_delays(self):
274
+ """Initialize TDI specific information"""
275
+
276
+ # setup the actual TDI combination
277
+ if self.tdi in ["1st generation", "2nd generation"]:
278
+ # tdi 1.0
279
+ tdi_combinations = [
280
+ {"link": 13, "links_for_delay": [], "sign": +1},
281
+ {"link": 31, "links_for_delay": [13], "sign": +1},
282
+ {"link": 12, "links_for_delay": [13, 31], "sign": +1},
283
+ {"link": 21, "links_for_delay": [13, 31, 12], "sign": +1},
284
+ {"link": 12, "links_for_delay": [], "sign": -1},
285
+ {"link": 21, "links_for_delay": [12], "sign": -1},
286
+ {"link": 13, "links_for_delay": [12, 21], "sign": -1},
287
+ {"link": 31, "links_for_delay": [12, 21, 13], "sign": -1},
288
+ ]
289
+
290
+ if self.tdi == "2nd generation":
291
+ # tdi 2.0 is tdi 1.0 + additional terms
292
+ tdi_combinations += [
293
+ {"link": 12, "links_for_delay": [13, 31, 12, 21], "sign": +1},
294
+ {"link": 21, "links_for_delay": [13, 31, 12, 21, 12], "sign": +1},
295
+ {
296
+ "link": 13,
297
+ "links_for_delay": [13, 31, 12, 21, 12, 21],
298
+ "sign": +1,
299
+ },
300
+ {
301
+ "link": 31,
302
+ "links_for_delay": [13, 31, 12, 21, 12, 21, 13],
303
+ "sign": +1,
304
+ },
305
+ {"link": 13, "links_for_delay": [12, 21, 13, 31], "sign": -1},
306
+ {"link": 31, "links_for_delay": [12, 21, 13, 31, 13], "sign": -1},
307
+ {
308
+ "link": 12,
309
+ "links_for_delay": [12, 21, 13, 31, 13, 31],
310
+ "sign": -1,
311
+ },
312
+ {
313
+ "link": 21,
314
+ "links_for_delay": [12, 21, 13, 31, 13, 31, 12],
315
+ "sign": -1,
316
+ },
317
+ ]
318
+
319
+ elif isinstance(self.tdi, list):
320
+ tdi_combinations = self.tdi
321
+
322
+ else:
323
+ raise ValueError(
324
+ "tdi kwarg should be '1st generation', '2nd generation', or a list with a specific tdi combination."
325
+ )
326
+ self.tdi_combinations = tdi_combinations
327
+
328
+ @property
329
+ def tdi_combinations(self) -> List:
330
+ """TDI Combination setup"""
331
+ return self._tdi_combinations
332
+
333
+ @tdi_combinations.setter
334
+ def tdi_combinations(self, tdi_combinations: List) -> None:
335
+ """Set TDI combinations and fill out setup."""
336
+ tdi_base_links = []
337
+ tdi_link_combinations = []
338
+ tdi_signs = []
339
+ tdi_operation_index = []
340
+ channels = []
341
+
342
+ tdi_index = 0
343
+ for permutation_number in range(3):
344
+ for tmp in tdi_combinations:
345
+ tdi_base_links.append(
346
+ self._cyclic_permutation(tmp["link"], permutation_number)
347
+ )
348
+ tdi_signs.append(float(tmp["sign"]))
349
+ channels.append(permutation_number)
350
+ if len(tmp["links_for_delay"]) == 0:
351
+ tdi_link_combinations.append(-11)
352
+ tdi_operation_index.append(tdi_index)
353
+
354
+ else:
355
+ for link_delay in tmp["links_for_delay"]:
356
+
357
+ tdi_link_combinations.append(
358
+ self._cyclic_permutation(link_delay, permutation_number)
359
+ )
360
+ tdi_operation_index.append(tdi_index)
361
+
362
+ tdi_index += 1
363
+
364
+ self.tdi_operation_index = self.xp.asarray(tdi_operation_index).astype(
365
+ self.xp.int32
366
+ )
367
+ self.tdi_base_links = self.xp.asarray(tdi_base_links).astype(self.xp.int32)
368
+ self.tdi_link_combinations = self.xp.asarray(tdi_link_combinations).astype(
369
+ self.xp.int32
370
+ )
371
+ self.tdi_signs = self.xp.asarray(tdi_signs).astype(self.xp.float64)
372
+ self.channels = self.xp.asarray(channels).astype(self.xp.int32)
373
+ assert len(self.tdi_link_combinations) == len(self.tdi_operation_index)
374
+
375
+ assert (
376
+ len(self.tdi_base_links)
377
+ == len(np.unique(self.tdi_operation_index))
378
+ == len(self.tdi_signs)
379
+ == len(self.channels)
380
+ )
381
+
382
+ def _cyclic_permutation(self, link, permutation):
383
+ """permute indexes by cyclic permutation"""
384
+ link_str = str(link)
385
+
386
+ out = ""
387
+ for i in range(2):
388
+ sc = int(link_str[i])
389
+ temp = sc + permutation
390
+ if temp > 3:
391
+ temp = temp % 3
392
+ out += str(temp)
393
+
394
+ return int(out)
395
+
396
+ @property
397
+ def y_gw(self):
398
+ """Projections along the arms"""
399
+ return self.y_gw_flat.reshape(self.nlinks, -1)
400
+
401
+ def _data_time_check(
402
+ self, t_data: np.ndarray, input_in: np.ndarray
403
+ ) -> Tuple[np.ndarray, np.ndarray]:
404
+ # remove input data that goes beyond orbital information
405
+ if t_data.max() > self.response_orbits.t.max():
406
+ warnings.warn(
407
+ "Input waveform is longer than available orbital information. Trimming to fit orbital information."
408
+ )
409
+
410
+ max_ind = np.where(t_data <= self.response_orbits.t.max())[0][-1]
411
+
412
+ t_data = t_data[:max_ind]
413
+ input_in = input_in[:max_ind]
414
+ return (t_data, input_in)
415
+
416
+ def get_projections(self, input_in, lam, beta, t0=10000.0):
417
+ """Compute projections of GW signal on to LISA constellation
418
+
419
+ Args:
420
+ input_in (xp.ndarray): Input complex time-domain signal. It should be of the form:
421
+ :math:`h_+ + ih_x`. If using the GPU for the response, this should be a CuPy array.
422
+ lam (double): Ecliptic Longitude in radians.
423
+ beta (double): Ecliptic Latitude in radians.
424
+ t0 (double, optional): Time at which to the waveform. Because of the delays
425
+ and interpolation towards earlier times, the beginning of the waveform
426
+ is garbage. ``t0`` tells the waveform generator where to start the waveform
427
+ compraed to ``t=0``.
428
+
429
+ Raises:
430
+ ValueError: If ``t0`` is not large enough.
431
+
432
+
433
+ """
434
+ self.tdi_start_ind = int(t0 / self.dt)
435
+ # get necessary buffer for TDI
436
+ self.check_tdi_buffer = int(100.0 * self.sampling_frequency) + 4 * self.order
437
+
438
+ from copy import deepcopy
439
+
440
+ tmp_orbits = deepcopy(self.response_orbits.x_base)
441
+ self.projection_buffer = (
442
+ int(
443
+ (
444
+ np.sum(
445
+ tmp_orbits.copy() * tmp_orbits.copy(),
446
+ axis=-1,
447
+ )
448
+ ** (1 / 2)
449
+ ).max()
450
+ * C_inv
451
+ )
452
+ + 4 * self.order
453
+ )
454
+ self.projections_start_ind = self.tdi_start_ind - 2 * self.check_tdi_buffer
455
+
456
+ if self.projections_start_ind < self.projection_buffer:
457
+ raise ValueError(
458
+ "Need to increase t0. The initial buffer is not large enough."
459
+ )
460
+
461
+ # determine sky vectors
462
+ k = np.zeros(3, dtype=np.float64)
463
+ u = np.zeros(3, dtype=np.float64)
464
+ v = np.zeros(3, dtype=np.float64)
465
+
466
+ self.num_total_points = len(input_in)
467
+
468
+ cosbeta = np.cos(beta)
469
+ sinbeta = np.sin(beta)
470
+
471
+ coslam = np.cos(lam)
472
+ sinlam = np.sin(lam)
473
+
474
+ v[0] = -sinbeta * coslam
475
+ v[1] = -sinbeta * sinlam
476
+ v[2] = cosbeta
477
+ u[0] = sinlam
478
+ u[1] = -coslam
479
+ u[2] = 0.0
480
+ k[0] = -cosbeta * coslam
481
+ k[1] = -cosbeta * sinlam
482
+ k[2] = -sinbeta
483
+
484
+ self.nlinks = 6
485
+ k_in = self.xp.asarray(k)
486
+ u_in = self.xp.asarray(u)
487
+ v_in = self.xp.asarray(v)
488
+
489
+ input_in = self.xp.asarray(input_in)
490
+
491
+ t_data = self.xp.arange(len(input_in)) * self.dt
492
+
493
+ t_data, input_in = self._data_time_check(t_data, input_in)
494
+
495
+ assert len(input_in) >= self.num_pts
496
+ y_gw = self.xp.zeros((self.nlinks * self.num_pts,), dtype=self.xp.float64)
497
+
498
+ orbits_in = self.pycppDetector_fastlisa(*self.response_orbits.pycppdetector_args)
499
+
500
+ self.response_gen(
501
+ y_gw,
502
+ t_data,
503
+ k_in,
504
+ u_in,
505
+ v_in,
506
+ self.dt,
507
+ len(input_in),
508
+ input_in,
509
+ len(input_in),
510
+ self.order,
511
+ self.sampling_frequency,
512
+ self.buffer_integer,
513
+ self.A_in,
514
+ self.deps,
515
+ len(self.A_in),
516
+ self.E_in,
517
+ self.projections_start_ind,
518
+ orbits_in,
519
+ )
520
+
521
+ self.y_gw_flat = y_gw
522
+ self.y_gw_length = self.num_pts
523
+
524
+ @property
525
+ def XYZ(self):
526
+ """Return links as an array"""
527
+ return self.delayed_links_flat.reshape(3, -1)
528
+
529
+ def get_tdi_delays(self, y_gw=None):
530
+ """Get TDI combinations from projections.
531
+
532
+ This functions generates the TDI combinations from the projections
533
+ computed with ``get_projections``. It can return XYZ, AET, or AE depending
534
+ on what was input for ``tdi_chan`` into ``__init__``.
535
+
536
+ Args:
537
+ y_gw (xp.ndarray, optional): Projections along each link. Must be
538
+ a 2D ``numpy`` or ``cupy`` array with shape: ``(nlinks, num_pts)``.
539
+ The links must be entered in the proper order in the code.
540
+ The link order is given in the orbits class: ``orbits.LINKS``.
541
+ (Default: None)
542
+
543
+ Returns:
544
+ tuple: (X,Y,Z) or (A,E,T) or (A,E)
545
+
546
+ Raises:
547
+ ValueError: If ``tdi_chan`` is not one of the options.
548
+
549
+
550
+ """
551
+ self.delayed_links_flat = self.xp.zeros(
552
+ (3, self.num_pts), dtype=self.xp.float64
553
+ )
554
+
555
+ # y_gw entered directly
556
+ if y_gw is not None:
557
+ assert y_gw.shape == (len(self.link_space_craft_0_in), self.num_pts)
558
+ self.y_gw_flat = y_gw.flatten().copy()
559
+ self.y_gw_length = self.num_pts
560
+
561
+ elif self.y_gw_flat is None:
562
+ raise ValueError(
563
+ "Need to either enter projection array or have this code determine projections."
564
+ )
565
+
566
+ self.delayed_links_flat = self.delayed_links_flat.flatten()
567
+
568
+ t_data = self.xp.arange(self.y_gw_length) * self.dt
569
+
570
+ num_units = int(self.tdi_operation_index.max() + 1)
571
+
572
+ assert np.all(
573
+ (np.diff(self.tdi_operation_index) == 0)
574
+ | (np.diff(self.tdi_operation_index) == 1)
575
+ )
576
+
577
+ _, unit_starts, unit_lengths = np.unique(
578
+ self.tdi_operation_index,
579
+ return_index=True,
580
+ return_counts=True,
581
+ )
582
+
583
+ unit_starts = unit_starts.astype(np.int32)
584
+ unit_lengths = unit_lengths.astype(np.int32)
585
+
586
+ orbits_in = self.pycppDetector_fastlisa(*self.tdi_orbits.pycppdetector_args)
587
+
588
+ self.tdi_gen(
589
+ self.delayed_links_flat,
590
+ self.y_gw_flat,
591
+ self.y_gw_length,
592
+ self.num_pts,
593
+ t_data,
594
+ unit_starts,
595
+ unit_lengths,
596
+ self.tdi_base_links,
597
+ self.tdi_link_combinations,
598
+ self.tdi_signs,
599
+ self.channels,
600
+ num_units,
601
+ 3, # num channels
602
+ self.order,
603
+ self.sampling_frequency,
604
+ self.buffer_integer,
605
+ self.A_in,
606
+ self.deps,
607
+ len(self.A_in),
608
+ self.E_in,
609
+ self.tdi_start_ind,
610
+ orbits_in,
611
+ )
612
+
613
+ if self.tdi_chan == "XYZ":
614
+ X, Y, Z = self.XYZ
615
+ return X, Y, Z
616
+
617
+ elif self.tdi_chan == "AET" or self.tdi_chan == "AE":
618
+ X, Y, Z = self.XYZ
619
+ A, E, T = AET(X, Y, Z)
620
+ if self.tdi_chan == "AET":
621
+ return A, E, T
622
+
623
+ else:
624
+ return A, E
625
+
626
+ else:
627
+ raise ValueError("tdi_chan must be 'XYZ', 'AET' or 'AE'.")
628
+
629
+
630
+ class ResponseWrapper(FastLISAResponseParallelModule):
631
+ """Wrapper to produce LISA TDI from TD waveforms
632
+
633
+ This class takes a waveform generator that produces :math:`h_+ \pm ih_x`.
634
+ (:code:`flip_hx` is used if the waveform produces :math:`h_+ - ih_x`).
635
+ It takes the complex waveform in the SSB frame and produces the TDI channels
636
+ according to settings chosen for :class:`pyResponseTDI`.
637
+
638
+ The waveform generator must have :code:`kwargs` with :code:`T` for the observation
639
+ time in years and :code:`dt` for the time step in seconds.
640
+
641
+ Args:
642
+ waveform_gen (obj): Function or class (with a :code:`__call__` function) that takes parameters and produces
643
+ :math:`h_+ \pm h_x`.
644
+ Tobs (double): Observation time in years.
645
+ dt (double): Time between time samples in seconds. The inverse of the sampling frequency.
646
+ index_lambda (int): The user will input parameters. The code will read these in
647
+ with the :code:`*args` formalism producing a list. :code:`index_lambda`
648
+ tells the class the index of the ecliptic longitude within this list of
649
+ parameters.
650
+ index_beta (int): The user will input parameters. The code will read these in
651
+ with the :code:`*args` formalism producing a list. :code:`index_beta`
652
+ tells the class the index of the ecliptic latitude (or ecliptic polar angle)
653
+ within this list of parameters.
654
+ t0 (double, optional): Start of returned waveform in seconds leaving ample time for garbage at
655
+ the beginning of the waveform. It also removed the same amount from the end. (Default: 10000.0)
656
+ flip_hx (bool, optional): If True, :code:`waveform_gen` produces :math:`h_+ - ih_x`.
657
+ :class:`pyResponseTDI` takes :math:`h_+ + ih_x`, so this setting will
658
+ multiply the cross polarization term out of the waveform generator by -1.
659
+ (Default: :code:`False`)
660
+ remove_sky_coords (bool, optional): If True, remove the sky coordinates from
661
+ the :code:`*args` list. This should be set to True if the waveform
662
+ generator does not take in the sky information. (Default: :code:`False`)
663
+ is_ecliptic_latitude (bool, optional): If True, the latitudinal sky
664
+ coordinate is the ecliptic latitude. If False, thes latitudinal sky
665
+ coordinate is the polar angle. In this case, the code will
666
+ convert it with :math:`\beta=\pi / 2 - \Theta`. (Default: :code:`True`)
667
+ force_backend (str, optional): If given, run this class on the requested backend.
668
+ Options are ``"cpu"``, ``"cuda11x"``, ``"cuda12x"``. (default: ``None``)
669
+ remove_garbage (bool or str, optional): If True, it removes everything before ``t0``
670
+ and after the end time - ``t0``. If ``str``, it must be ``"zero"``. If ``"zero"``,
671
+ it will not remove the points, but set them to zero. This is ideal for PE. (Default: ``True``)
672
+ n_overide (int, optional): If not ``None``, this will override the determination of
673
+ the number of points, ``n``, from ``int(T/dt)`` to the ``n_overide``. This is used
674
+ if there is an issue matching points between the waveform generator and the response
675
+ model.
676
+ orbits (:class:`Orbits`, optional): Orbits class from LISA Analysis Tools. Works with LISA Orbits
677
+ outputs: ``lisa-simulation.pages.in2p3.fr/orbits/``.
678
+ (default: :class:`EqualArmlengthOrbits`)
679
+ **kwargs (dict, optional): Keyword arguments passed to :class:`pyResponseTDI`.
680
+
681
+ """
682
+
683
+ def __init__(
684
+ self,
685
+ waveform_gen,
686
+ Tobs,
687
+ dt,
688
+ index_lambda,
689
+ index_beta,
690
+ t0=10000.0,
691
+ flip_hx=False,
692
+ remove_sky_coords=False,
693
+ is_ecliptic_latitude=True,
694
+ force_backend=None,
695
+ remove_garbage=True,
696
+ n_overide=None,
697
+ orbits: Optional[Orbits] = EqualArmlengthOrbits,
698
+ **kwargs,
699
+ ):
700
+
701
+ # store all necessary information
702
+ self.waveform_gen = waveform_gen
703
+ self.index_lambda = index_lambda
704
+ self.index_beta = index_beta
705
+ self.dt = dt
706
+ self.t0 = t0
707
+ self.sampling_frequency = 1.0 / dt
708
+ super().__init__(force_backend=force_backend)
709
+
710
+ if orbits is None:
711
+ orbits = EqualArmlengthOrbits()
712
+
713
+ assert isinstance(orbits, Orbits)
714
+
715
+ if Tobs * YRSID_SI > orbits.t_base.max():
716
+ warnings.warn(
717
+ f"Tobs is larger than available orbital information time array. Reducing Tobs to {orbits.t_base.max()}"
718
+ )
719
+ Tobs = orbits.t_base.max() / YRSID_SI
720
+
721
+ if n_overide is not None:
722
+ if not isinstance(n_overide, int):
723
+ raise ValueError("n_overide must be an integer if not None.")
724
+ self.n = n_overide
725
+
726
+ else:
727
+ self.n = int(Tobs * YRSID_SI / dt)
728
+
729
+ self.Tobs = self.n * dt
730
+ self.is_ecliptic_latitude = is_ecliptic_latitude
731
+ self.remove_sky_coords = remove_sky_coords
732
+ self.flip_hx = flip_hx
733
+ self.remove_garbage = remove_garbage
734
+
735
+ # initialize response function class
736
+ self.response_model = pyResponseTDI(
737
+ self.sampling_frequency, self.n, orbits=orbits, force_backend=force_backend, **kwargs
738
+ )
739
+
740
+ self.Tobs = (self.n * self.response_model.dt) / YRSID_SI
741
+
742
+ @property
743
+ def xp(self) -> object:
744
+ return self.backend.xp
745
+
746
+ @property
747
+ def citation(self):
748
+ """Get citations for use of this code"""
749
+
750
+ return """
751
+ # TODO add
752
+ """
753
+
754
+ @classmethod
755
+ def supported_backends(cls):
756
+ return ["fastlisaresponse_" + _tmp for _tmp in cls.GPU_RECOMMENDED()]
757
+
758
+ def __call__(self, *args, **kwargs):
759
+ """Run the waveform and response generation
760
+
761
+ Args:
762
+ *args (list): Arguments to the waveform generator. This must include
763
+ the sky coordinates.
764
+ **kwargs (dict): kwargs necessary for the waveform generator.
765
+
766
+ Return:
767
+ list: TDI Channels.
768
+
769
+ """
770
+
771
+ args = list(args)
772
+
773
+ # get sky coords
774
+ beta = args[self.index_beta]
775
+ lam = args[self.index_lambda]
776
+
777
+ # remove them from the list if waveform generator does not take them
778
+ if self.remove_sky_coords:
779
+ args.pop(self.index_beta)
780
+ args.pop(self.index_lambda)
781
+
782
+ # transform polar angle
783
+ if not self.is_ecliptic_latitude:
784
+ beta = np.pi / 2.0 - beta
785
+
786
+ # add the new Tobs and dt info to the waveform generator kwargs
787
+ kwargs["T"] = self.Tobs
788
+ kwargs["dt"] = self.dt
789
+
790
+ # get the waveform
791
+ h = self.waveform_gen(*args, **kwargs)
792
+
793
+ if self.flip_hx:
794
+ h = h.real - 1j * h.imag
795
+
796
+ self.response_model.get_projections(h, lam, beta, t0=self.t0)
797
+ tdi_out = self.response_model.get_tdi_delays()
798
+
799
+ out = list(tdi_out)
800
+ if self.remove_garbage is True: # bool
801
+ for i in range(len(out)):
802
+ out[i] = out[i][
803
+ self.response_model.tdi_start_ind : -self.response_model.tdi_start_ind
804
+ ]
805
+
806
+ elif isinstance(self.remove_garbage, str): # bool
807
+ if self.remove_garbage != "zero":
808
+ raise ValueError("remove_garbage must be True, False, or 'zero'.")
809
+ for i in range(len(out)):
810
+ out[i][: self.response_model.tdi_start_ind] = 0.0
811
+ out[i][-self.response_model.tdi_start_ind :] = 0.0
812
+
813
+ return out