-
Notifications
You must be signed in to change notification settings - Fork 1
/
ppo.py
1633 lines (1369 loc) · 64.5 KB
/
ppo.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936
937
938
939
940
941
942
943
944
945
946
947
948
949
950
951
952
953
954
955
956
957
958
959
960
961
962
963
964
965
966
967
968
969
970
971
972
973
974
975
976
977
978
979
980
981
982
983
984
985
986
987
988
989
990
991
992
993
994
995
996
997
998
999
1000
#!/usr/bin/env python
# -*- coding: utf-8 -*-
'''
ppo.py |github|
---------------
The main Python interface to the C photodynamics code. Allows users to compute
and plot planet-planet occultation light curves, as well as transits, secondary
eclipses, phase curves, mutual transits, planet-moon events, and more.
.. role:: raw-html(raw)
:format: html
.. |github| replace:: :raw-html:`<a href = "https://github.com/rodluger/planetplanet/blob/master/planetplanet/photo/ppo.py"><i class="fa fa-github" aria-hidden="true"></i></a>`
'''
from __future__ import division, print_function, absolute_import, \
unicode_literals
from ..constants import *
from ..detect import jwst
from .eyeball import DrawEyeball, GetAngles
from .structs import *
import ctypes
import numpy as np
np.seterr(invalid = 'ignore')
import os, shutil
import sysconfig
from numpy.ctypeslib import ndpointer, as_ctypes
import matplotlib
import matplotlib.pyplot as pl
import matplotlib.animation as animation
from matplotlib.ticker import MaxNLocator
from matplotlib.widgets import Button
from matplotlib.collections import LineCollection
from matplotlib.colors import LinearSegmentedColormap
from scipy.integrate import quad
from tqdm import tqdm
from six import string_types
import numba
import rebound
__all__ = ['System']
# Find system suffix and import the shared library
suffix = sysconfig.get_config_var('EXT_SUFFIX')
if suffix is None:
suffix = ".so"
dn = os.path.dirname
libppo = ctypes.cdll.LoadLibrary(os.path.join(dn(dn(dn(
os.path.abspath(__file__)))),
"libppo" + suffix))
def _colorline(ax, x, y, color = (0, 0, 0), **kwargs):
'''
Plots the curve `y(x)` with linearly increasing alpha.
Adapted from `http://nbviewer.jupyter.org/github/dpsanders/matplotlib-examples/blob/master/colorline.ipynb`_.
'''
# A bit hacky... But there doesn't seem to be
# an easy way to get the hex code for a named color...
if isinstance(color, string_types):
if color.startswith("#"):
hex = color[1:]
else:
if len(color) == 1:
if color == 'k':
color = 'black'
elif color == 'r':
color = 'red'
elif color == 'b':
color = 'blue'
elif color == 'g':
color = 'green'
elif color == 'y':
color = 'yellow'
elif color == 'w':
color = 'white'
else:
# ?!
color = 'black'
hex = matplotlib.colors.cnames[color.lower()][1:]
r, g, b = tuple(int(hex[i:i+2], 16) / 255. for i in (0, 2, 4))
else:
r, g, b = color
colors = [(r, g, b, i) for i in np.linspace(0, 1, 3)]
cmap = LinearSegmentedColormap.from_list('alphacmap', colors, N = 1000)
points = np.array([x, y]).T.reshape(-1, 1, 2)
segments = np.concatenate([points[:-1], points[1:]], axis=1)
lc = LineCollection(segments, array = np.linspace(0.0, 1.0, len(x)),
cmap = cmap, norm = pl.Normalize(0.0, 1.0), **kwargs)
ax.add_collection(lc)
return lc
class _Animation(object):
'''
An animation class for occultation movies.
'''
def __init__(self, t, fig, axim, tracker, occ, pt_xz, pt_zy, body, bodies,
occultors, interval = 50, gifname = None, xy = None,
quiet = False):
'''
'''
self.t = t
self.fig = fig
self.axim = axim
self.tracker = tracker
self.occ = occ
self.pt_xz = pt_xz
self.pt_zy = pt_zy
self.body = body
self.bodies = bodies
self.occultors = occultors
self.xy = xy
self.pause = True
self.animation = animation.FuncAnimation(self.fig, self.animate,
frames = 100,
interval = interval,
repeat = True)
# Save?
if gifname is not None:
self.pause = False
if not gifname.endswith('.gif'):
gifname += '.gif'
if not quiet:
print("Saving %s..." % gifname)
self.animation.save(gifname, writer = 'imagemagick',
fps = 20, dpi = 150)
self.pause = True
# Toggle button
self.axbutton = pl.axes([0.185, 0.12, 0.1, 0.03])
self.button = Button(self.axbutton, 'Play', color = 'lightgray')
self.button.on_clicked(self.toggle)
def toggle(self, event):
'''
Pause/play the animation. Unfortunately I haven't been able to
figure out how to freeze the animation so that it resumes at the
same frame it was paused on...
'''
if self.pause:
self.button.label.set_text('Pause')
else:
self.button.label.set_text('Play')
self.pause ^= True
def animate(self, j):
'''
Play frame `j` of the animation.
'''
if not self.pause:
# Normalize the time index
j = int(j * len(self.t) / 100.)
# Time tracker
self.tracker.set_xdata(self.bodies[0].time_hr[self.t[j]])
# Occultor images
x0 = self.body.x_hr[self.t[j]]
y0 = self.body.y_hr[self.t[j]]
for k, occultor in enumerate(self.occultors):
if self.xy is None:
xo, yo = occultor.x_hr[self.t[j]] - x0, \
occultor.y_hr[self.t[j]] - y0
else:
xo, yo = self.xy(occultor.x_hr[self.t[j]] - x0,
occultor.y_hr[self.t[j]] - y0)
self.occ[k].center = (xo / self.body._r, yo / self.body._r)
# BODY orbits
for k, b in enumerate(self.bodies):
self.pt_xz[k].set_xdata(b.x_hr[self.t[j]])
self.pt_xz[k].set_ydata(b.z_hr[self.t[j]])
self.pt_zy[k].set_xdata(b.z_hr[self.t[j]])
self.pt_zy[k].set_ydata(b.y_hr[self.t[j]])
class System(object):
'''
The planetary system class. This is the main interface to the
photodynamical core. Instantiate with all bodies in the system and the
desired settings, passed as :py:obj:`kwargs`.
:param bodies: Any number of :py:func:`Planet`, :py:func:`Moon`, or \
:py:func:`Star` instances comprising all the bodies in the system. \
The first body is assumed to be the primary.
:param bool nbody: Uses the :py:obj:`REBOUND` N-body code to compute \
orbits. Default :py:obj:`True`
:param float keptol: Kepler solver tolerance. Default `1.e-15`
:param int maxkepiter: Maximum number of Kepler solver iterations. \
Default `100`
:param str kepsolver: Kepler solver (`newton` | `mdfast`). Default `newton`
:param float timestep: Timestep in days for the N-body solver. \
Default `0.01`
:param bool adaptive: Adaptive grid for limb-darkened bodies? \
Default :py:obj:`True`
:param bool quiet: Suppress output? Default :py:obj:`False`
:param float mintheta: Absolute value of the minimum phase angle in \
degrees. Below this angle, elliptical boundaries of constant \
surface brightness on the planet surface are treated as vertical \
lines. Default `1.`
:param int maxvertices: Maximum number of vertices allowed in the area \
computation. Default `999`
:param int maxfunctions: Maximum number of functions allowed in the area \
computation. Default `999`
:param int oversample: Oversampling factor for each exposure. Default `1`
:param float distance: Distance to the system in parsecs. Default `10.`
:param bool circleopt: Solve the simpler quadratic problem for \
circle-ellipse intersections when the axes of the ellipse are \
equal to within :math:`10^{-10}`? Default :py:obj:`True`
:param bool batmanopt: Use the :py:mod:`batman` algorithm to compute \
light curves of radially symmetric bodies? This can significantly \
speed up the code. Default :py:obj:`True`
:param str integrator: The N-body integrator \
(:py:obj:`whfast` | :py:obj:`ias15`) to use. \
Default :py:obj:`ias15`
'''
def __init__(self, *bodies, **kwargs):
'''
Initialize the class.
'''
# Initialize
self.bodies = bodies
self.settings = SETTINGS(**kwargs)
self._reset()
def _reset(self):
'''
Applies all settings, makes bodies accessible as properties, resets
flags, and computes some preliminary orbital information for the
system.
'''
# Move params set by the user over to the settings class
for param in self.settings.params:
if hasattr(self, param):
setattr(self.settings, param, getattr(self, param))
delattr(self, param)
# Make planets accessible as properties
self.primary = self.bodies[0]
for body in self.bodies:
setattr(self, body.name, body)
# Force an update of `tperi0`
body.ecc = body._ecc
self._names = np.array([p.name for p in self.bodies])
self.colors = [b.color for b in self.bodies]
# Evaluate the body host names
for i, body in enumerate(self.bodies):
if body.host is None:
# Host is CM of all interior bodies
body.host = CM(*self.bodies[:i])
body._host = -1
elif type(body.host) is int:
# Host is an index
body._host = body.host
body.host = self.bodies[body.host]
elif isinstance(body.host, string_types):
# Host is another body in the system
body.host = getattr(self, body.host)
body._host = np.argmax(self._names == body.host.name)
elif body.host in self.bodies:
# Host is another body in the system
body._host = np.argmax(self._names == body.host.name)
else:
try:
assert body.host.body_type == 'cm', "Error!"
except:
# Something went wrong
raise ValueError("Invalid `host` setting for body `%s`."
% body.name)
# Compute the semi-major axis for each body (in Earth radii)
for body in self.bodies:
body._computed = False
body.a = ((body.per * DAYSEC) ** 2 * G
* (body.host._m + body._m) * MEARTH
/ (4 * np.pi ** 2)) ** (1. / 3.) / REARTH
# Reset continuum
self._continuum = np.empty((0,), dtype = 'float64')
# Reset animations
self._animations = []
# Reset flag
self._computed = False
def _malloc(self, nt, nw, noccultors = 0, noccultations = 0):
'''
Allocate memory for all the C arrays.
'''
# Check that the first body is a star
assert self.bodies[0].body_type == 'star', \
'The first body must be a :py:class:`Star`.'
# Force cartesian coords for the star.
self.bodies[0].cartesian = True
# Check that if there are N stars, the first N bodies are stars
nstars = np.sum([int(b.body_type == 'star') for b in self.bodies])
assert np.all([b.body_type == 'star' for b in self.bodies[:nstars]]), \
'Stars must be passed as the first arguments to `System`,' \
'before any planets or moons.'
self.settings.nstars = nstars
self.nstars = nstars
# Check that if there are multiple stars, Keplerian solver is off
if nstars > 1:
assert self.settings.nbody, "N-body integrator must be selected" \
"for systems with multiple stars."
# If there's more than one star, disable the RadiativeEquilibriumMap
if 'star' in [b.body_type for b in self.bodies[1:]]:
for b in self.bodies[1:]:
assert b.radiancemap.maptype not in [MAP_ELLIPTICAL_DEFAULT,
MAP_ELLIPTICAL_CUSTOM], "Elliptically-symmetric radiance "\
"maps not implemented for multiple-star systems."
# Initialize the C interface
self._Orbits = libppo.Orbits
self._Orbits.restype = ctypes.c_int
self._Orbits.argtypes = [ctypes.c_int,
ctypes.ARRAY(ctypes.c_double, nt),
ctypes.c_int,
ctypes.POINTER(ctypes.POINTER(BODY)),
SETTINGS]
self._Flux = libppo.Flux
self._Flux.restype = ctypes.c_int
self._Flux.argtypes = [ctypes.c_int,
ctypes.ARRAY(ctypes.c_double, nt),
ctypes.c_int,
ctypes.ARRAY(ctypes.c_double, nw),
ctypes.ARRAY(ctypes.c_double, nw * nt),
ctypes.c_int,
ctypes.POINTER(ctypes.POINTER(BODY)),
SETTINGS]
# Are we just checking for occultations?
if (noccultors > 0) and (noccultations > 0):
self._NextOccultation = libppo.NextOccultation
self._NextOccultation.restype = ctypes.c_int
self._NextOccultation.argtypes = \
[ctypes.c_int,
ctypes.ARRAY(ctypes.c_double, nt),
ctypes.c_int,
ctypes.POINTER(ctypes.POINTER(BODY)),
SETTINGS,
ctypes.c_int,
ctypes.c_int,
ctypes.ARRAY(ctypes.c_int, noccultors),
ctypes.c_int,
ctypes.ARRAY(ctypes.c_double, noccultations),
ctypes.ARRAY(ctypes.c_int, noccultations),
ctypes.ARRAY(ctypes.c_double, noccultations)]
# Allocate memory for all the arrays
for body in self.bodies:
body.time = np.zeros(nt)
body._time = np.ctypeslib.as_ctypes(body.time)
body.wavelength = np.zeros(nw)
body._wavelength = np.ctypeslib.as_ctypes(body.wavelength)
body.x = np.zeros(nt)
body.x[0] = body.x0
body._x = np.ctypeslib.as_ctypes(body.x)
body.y = np.zeros(nt)
body.y[0] = body.y0
body._y = np.ctypeslib.as_ctypes(body.y)
body.z = np.zeros(nt)
body.z[0] = body.z0
body._z = np.ctypeslib.as_ctypes(body.z)
body.vx = np.zeros(nt)
body.vx[0] = body.vx0
body._vx = np.ctypeslib.as_ctypes(body.vx)
body.vy = np.zeros(nt)
body.vy[0] = body.vy0
body._vy = np.ctypeslib.as_ctypes(body.vy)
body.vz = np.zeros(nt)
body.vz[0] = body.vz0
body._vz = np.ctypeslib.as_ctypes(body.vz)
body.occultor = np.zeros(nt, dtype = 'int32')
body._occultor = np.ctypeslib.as_ctypes(body.occultor)
body.total_flux = np.zeros(nw)
body._total_flux = np.ctypeslib.as_ctypes(body.total_flux)
# HACK: The fact that flux is 2d is a nightmare for ctypes. We will
# treat it as a 1d array within C and keep track of the row/column
# indices by hand...
body.flux = np.zeros((nt, nw))
body._flux1d = body.flux.reshape(-1)
body._flux = np.ctypeslib.as_ctypes(body._flux1d)
# HACK: Same for the limb darkening coefficients
body._u1d = body.u.reshape(-1)
body._u = np.ctypeslib.as_ctypes(body._u1d)
# Dimensions
body.nu = len(body.u)
body.nt = nt
body.nw = nw
# A pointer to a pointer to `BODY`. This is an array of `n`
# `BODY` instances, passed by reference. The contents can all be
# accessed through `bodies`
# NOTE: Before I subclassed BODY, this used to be
# >>> self._ptr_bodies = (ctypes.POINTER(BODY) * \
# >>> len(self.bodies))(*[ctypes.pointer(p) for p in self.bodies])
# I now cast the `Planet`, `Star`, and `Moon` instances as `BODY`
# pointers, as per https://stackoverflow.com/a/37827528
self._ptr_bodies = (ctypes.POINTER(BODY) * len(self.bodies)) \
(*[ctypes.cast(ctypes.byref(p),
ctypes.POINTER(BODY)) for p in self.bodies])
def compute(self, time, lambda1 = 5, lambda2 = 15, R = 100):
'''
Compute the full system light curve over a given `time` array between
wavelengths `lambda1` and `lambda2` at resolution `R`. This method
runs the photodynamical core, populates all bodies with their
individual light curves, and flags all occultation events.
:param array_like time: The times at which to evaluate the light \
curve (BJD − 2,450,000)
:param float lambda1: The start point of the wavelength grid in \
microns. Default `5`
:param float lambda2: The end point of the wavelength grid in \
microns. Default `15`
:param float R: The spectrum resolution, \
:math:`R = \\frac{\lambda}{\Delta\lambda}`. Default `100`
'''
# Reset
self._reset()
# Compute the wavelength grid
wav = [lambda1]
while(wav[-1] < lambda2):
wav.append(wav[-1] + wav[-1] / R)
wavelength = np.array(wav)
# Compute all limb darkening coefficients
for body in self.bodies:
body.u = [None for ld in body.limbdark]
for n, ld in enumerate(body.limbdark):
if callable(ld):
body.u[n] = ld(wavelength)
elif not hasattr(ld, '__len__'):
body.u[n] = ld * np.ones_like(wavelength)
else:
raise Exception("Limb darkening coefficients must be "
+ "provided as a list of scalars or "
+ "as a list of functions.")
body.u = np.array(body.u)
# Convert from microns to meters
wavelength *= 1e-6
# Oversample the time array to generate light curves
# observed w/ finite exposure time.
# Ensure `oversample` is odd.
if self.settings.oversample % 2 == 0:
oversample = self.settings.oversample + 1
else:
oversample = self.settings.oversample
if oversample == 1:
time_hr = np.array(time)
else:
tmid = (time[:-1] + time[1:]) / 2.
tmid = np.concatenate(([tmid[0] - (tmid[1] - tmid[0])], tmid,
[tmid[-1] + (tmid[-1] - tmid[-2])]))
time_hr = [np.linspace(tmid[i], tmid[i + 1], oversample)
for i in range(len(tmid) - 1)]
time_hr = np.concatenate(time_hr)
# Continuum flux
self._continuum = np.zeros(len(time_hr) * len(wavelength))
# Allocate memory
self._malloc(len(time_hr), len(wavelength))
# Call the light curve routine
err = self._Flux(len(time_hr), np.ctypeslib.as_ctypes(time_hr),
len(wavelength), np.ctypeslib.as_ctypes(wavelength),
np.ctypeslib.as_ctypes(self._continuum),
len(self.bodies), self._ptr_bodies, self.settings)
assert err <= 0, "Error in C routine `Flux` (%d)." % err
self._computed = True
# Downbin the continuum flux to the original time array
self._continuum_hr = np.array(self._continuum)
if self.settings.oversample > 1:
self._continuum = np.mean(self._continuum.reshape(-1, oversample,
len(wavelength)), axis = 1)
# Downbin other arrays to original time array
for body in self.bodies:
# Store the oversampled light curve
body.time_hr = np.array(body.time)
body.flux_hr = np.array(body.flux)
body.occultor_hr = np.array(body.occultor)
body.x_hr = np.array(body.x)
body.y_hr = np.array(body.y)
body.z_hr = np.array(body.z)
body.vx_hr = np.array(body.vx)
body.vy_hr = np.array(body.vy)
body.vz_hr = np.array(body.vz)
# Store the binned light curve
if self.settings.oversample > 1:
# Revert to original time array
body.time = time
# Average the flux over the exposure
body.flux = np.mean(body.flux.reshape(-1, oversample,
len(wavelength)),
axis = 1)
# Take the XYZ position at the bin center
body.x = body.x[oversample // 2::oversample]
body.y = body.y[oversample // 2::oversample]
body.z = body.z[oversample // 2::oversample]
body.vx = body.vx[oversample // 2::oversample]
body.vy = body.vy[oversample // 2::oversample]
body.vz = body.vz[oversample // 2::oversample]
# Get all bodies that occult at some point over the exposure
# The operation `bitwise_or.reduce` is *magical*
body.occultor = np.bitwise_or.reduce(body.occultor.reshape(-1,
oversample), axis = 1)
# Loop over all bodies to check for occultations
for body in self.bodies:
# Set the flag
body._computed = True
# Loop over all possible occultors
for occ in range(len(self.bodies)):
# Indices of occultation
inds = np.where(body.occultor_hr & 2 ** occ)[0]
# Store each one individually
if len(inds):
inds = np.split(inds, 1 + np.where(np.diff(inds) > 1)[0])
# Loop over each occultation event
for i in inds:
# Add a little padding
di = (i[-1] - i[0]) // 2
ia = max(0, i[0] - di)
ib = min(len(body.time_hr), i[-1] + di + 1)
i = np.array(list(range(ia, i[0])) + list(i)
+ list(range(i[-1] + 1, ib)))
# Store it
body._inds.append(i)
def compute_orbits(self, time):
'''
Run the dynamical code to compute the positions of all bodies over a
given `time` array. This method does not compute light curves, but it
does check for occultations.
:param array_like time: The times at which to store the body \
positions (BJD − 2,450,000)
'''
# Reset
self._reset()
for body in self.bodies:
body.u = np.array([], dtype = float)
self._malloc(len(time), 1)
# Call the light curve routine
err = self._Orbits(len(time), np.ctypeslib.as_ctypes(time),
len(self.bodies), self._ptr_bodies, self.settings)
assert err <= 0, "Error in C routine `Orbits` (%d)." % err
def observe(self, save = None, filter = 'f1500w', stack = 1,
instrument = 'jwst', alpha_err = 0.7, figsize = (12,4),
time_unit = 'BJD − 2,450,000'):
'''
Run telescope observability calculations for a system that has had its
lightcurve computed. Calculates a synthetic noised lightcurve in the
user specified filter.
:param bool save: Save a text file and a plot. Default :py:obj:`None`
:param filter: Filter name or :py:func:`Filter` object. \
Default `'f1500w'`
:type filter: str or :py:func:`Filter`
:param int stack: Number of exposures to stack. Default `1`
:param str instrument: Telescope instrument to use. Default `'jwst'`
'''
# Have we computed the light curves?
assert self._computed, "Please run `compute()` first."
# Get the filter "wheel"
if instrument.lower() == 'jwst':
wheel = jwst.get_miri_filter_wheel()
atel = 25.0
thermal = True
elif instrument.lower() == 'ost':
# Using JWST filters for OST for now
wheel = jwst.get_miri_filter_wheel()
# create new filter for 50 microns
filt50 = jwst.create_tophat_filter(47.5, 52.5, dlam=0.1,
Tput=0.3, name="OST50")
# Will have mirror between 8-15m, let's say: 13.5m
# This gives a telescope area of 144 m^2
atel = np.pi * (13.5 / 2.) ** 2
# No thermal noise
thermal = False
elif instrument.lower() == 'spitzer':
wheel = jwst.get_spitzer_filter_wheel()
atel = np.pi * (0.85 / 2.) ** 2
thermal = False
else:
raise ValueError("Invalid instrument.")
# Get the filter
if type(filter).__name__ == "Filter":
self.filter = filter
elif filter.lower() in [f.name.lower() for f in wheel]:
self.filter = wheel[np.argmax([f.name.lower() == filter.lower()
for f in wheel])]
else:
raise ValueError("Invalid filter.")
# Check that filter wavelength width is fully contained within
# the photodynamical wavelength grid
assert (self.filter.eff_wl - self.filter.eff_dwl/2.) > \
np.min(self.wavelength), "Wavelength grid does not extend short"\
" enough for filter."
assert (self.filter.eff_wl + self.filter.eff_dwl/2.) < \
np.max(self.wavelength), "Wavelength grid does not extend long"\
" enough for filter."
# Compute lightcurve in this filter
self.filter.compute_lightcurve(self.time, self.flux, self.continuum,
self.wavelength, stack = stack,
time_hr = self.time_hr,
flux_hr = self.flux_hr,
atel = atel, thermal = thermal,
quiet = self.settings.quiet)
# Setup plot
fig, ax = pl.subplots(figsize = figsize)
ax.set_title(r"%s" % self.filter.name)
ax.set_ylabel("Relative Flux", fontweight = 'bold', fontsize = 10)
ax.set_xlabel("Time [%s]" % time_unit, fontweight = 'bold',
fontsize = 10)
# Plot lightcurve
self.filter.lightcurve.plot(ax0 = ax, alpha_err = alpha_err)
# Determine average precision attained in lightcurve
ppm = np.mean(self.filter.lightcurve.sig
/ self.filter.lightcurve.obs) * 1e6
if not self.settings.quiet:
print("Average lightcurve precision: %.3f ppm" %ppm)
"""
Determine SNR on each event:
"""
# Mask in and out of occultation times
inmask = np.logical_or.reduce([planet.occultor > 0
for planet in self.bodies])
outmask = ~inmask
arr = np.arange(len(inmask))
# Break lightcurve mask into event chunks
n = 0
prev = False
events = []
for i in range(len(inmask)):
if inmask[i]:
if not prev:
n += 1
events.append([i])
else:
events[n-1].append(i)
prev = inmask[i]
# Determine SNR on event detections
SNRs = []
signal = []
noise = []
for i in range(n):
# In-event indices
mask = events[i]
# System photons over event
Nsys = self.filter.lightcurve.Nsys[mask]
# Background photons over event
Nback = self.filter.lightcurve.Nback[mask]
# Continuum photons over event
Ncont = self.filter.lightcurve.Ncont[mask]
# Compute signal of and noise on the event
S = np.sum(np.fabs(Ncont - Nsys)) / np.sum(Nsys + Nback) * 1e6
N = np.sqrt(np.sum(Nsys + Nback)) / np.sum(Nsys + Nback) * 1e6
signal.append(S)
noise.append(N)
# Compute the actual SNR on event. Note that this is NOT
# the sum of the signals divided by the sum of the noises!
# We need to add the SNR of each *datapoint* individually
# in quadrature.
SNR = np.sqrt(np.sum((Ncont - Nsys) ** 2 / (Nsys + Nback)))
SNRs.append(SNR)
# Annotate event SNRs on each event
for i in range(n):
k = np.argmin(self.filter.lightcurve.obs[np.array(events[i])])
j = events[i][k]
ax.annotate(r"$%.1f \sigma$" % SNRs[i],
xy = (self.filter.lightcurve.time[
int(np.mean(events[i]))],
self.filter.lightcurve.obs[j]
- 1 * self.filter.lightcurve.sig[j]),
ha = 'center',
va = 'center', color = "black",
fontweight = 'bold',
fontsize = 10, xytext = (0, -15),
textcoords = 'offset points')
# Save SNR and ampl as attributes in Lightcurve
self.filter.lightcurve.event_SNRs = SNRs
self.filter.lightcurve.event_signal = signal
self.filter.lightcurve.event_noise = noise
# Save to disk?
if save is not None:
# Compose data array to save
data = np.array([self.filter.lightcurve.time,
self.filter.lightcurve.obs,
self.filter.lightcurve.sig]).T
# Did the user provide a file name?
if save is True:
# No
np.savetxt("jwst_lc_%s.txt" % self.filter.name,
data, fmt = str("%.6e"),
header = "time [days] flux" +
" error",
comments = "")
fig.savefig("jwst_lc_%s.png" % self.filter.name,
bbox_inches="tight")
else:
# Yes
np.savetxt("%s.txt" % save,
data, fmt = str("%.6e"),
header = "time [days] flux"+
" error",
comments = "")
fig.savefig("%s.png" % save, bbox_inches="tight")
# Adjust for on-screen viewing
fig.subplots_adjust(bottom = 0.2)
return fig, ax
def next_occultation(self, occulted, occultors = None,
tstart = OCTOBER_08_2016,
tend = OCTOBER_08_2016 + 100,
dt = 0.001, noccultations = 1):
'''
Computes the time of the next occultation of body `occulted`.
:param occulted: The occulted body, passed as a string corresponding \
to the body name, the body instance, or the index of the body.
:param occultors: The occultor(s), passed as a list of strings, body \
instances, or indices. Default is to consider occultations by \
all bodies in the system.
:param float tstart: Time at which to start searching for \
occultations (BJD − 2,450,000). Default `8000.` \
(12:00:00 UT October 8, 2016)
:param float tend: Time at which to end the search if fewer than \
`noccultations` occultations have been found \
(BJD − 2,450,000). Default `8100.`
:param float dt: The search timestep in days. Occultations shorter \
than this value will be missed. Default `0.001` \
(about 2 minutes)
:param int noccultations: The number of occultations to search for. \
Once this many occultations have been found, halts the N-body \
integration and returns. Default `1`
:returns: Arrays corresponding to the times of the occultations \
(BJD − 2,450,000), the occulting bodies, and the durations \
(in minutes)
'''
assert tend > tstart, "The end time must be larger than the start time!"
# Convert `occulted` to an index
if isinstance(occulted, string_types):
occulted = np.argmax([b.name == occulted for b in self.bodies])
elif occulted in self.bodies:
occulted = np.argmax(np.array(self.bodies) == occulted)
else:
try:
if occulted < len(self.bodies):
pass
else:
raise TypeError("")
except TypeError:
raise ValueError("Invalid value for `occulted`.")
# If not set, default to occultations by all bodies
if occultors is None:
occultors = np.arange(len(self.bodies))
# Convert `occultors` to an array of indices
occultors = np.atleast_1d(occultors)
for i, occultor in enumerate(occultors):
if isinstance(occultor, string_types) or \
(type(occultor) is np.str_):
occultors[i] = np.argmax([b.name == occultor
for b in self.bodies])
elif occultor in self.bodies:
occultors[i] = np.argmax(np.array(self.bodies) == occultor)
else:
try:
if occultor < len(self.bodies):
pass
else:
raise TypeError("")
except TypeError:
raise ValueError("Invalid value for `occultors`.")
occultors = np.array(occultors, dtype = 'int32')
noccultors = len(occultors)
# Get the time array
time = np.arange(tstart, tend, dt)
# Construct empty occultation arrays
occultation_times = np.zeros(noccultations, dtype = 'float64') * np.nan
occultation_inds = np.zeros(noccultations, dtype = 'int32')
occultation_durs = np.zeros(noccultations, dtype = 'float64')
# Reset and allocate memory
self._reset()
for body in self.bodies:
body.u = np.array([], dtype = float)
self._malloc(len(time), 1, noccultors = noccultors,
noccultations = noccultations)
# Final check: currently only supported for the N-body integrator.
# TODO: Make this work for the Keplerian integrator.
assert self.settings.nbody, "Currently, `next_occultation` works " + \
"only with the N-Body solver. Please set `nbody = True`."
# Call the orbit routine
err = self._NextOccultation(len(time), np.ctypeslib.as_ctypes(time),
len(self.bodies),
self._ptr_bodies, self.settings, occulted,
noccultors,
np.ctypeslib.as_ctypes(occultors),
noccultations,
np.ctypeslib.as_ctypes(occultation_times),
np.ctypeslib.as_ctypes(occultation_inds),
np.ctypeslib.as_ctypes(occultation_durs))
times = []
occultors = []
durations = []
for t, i, d in zip(occultation_times, occultation_inds,
occultation_durs):
if not np.isnan(t):
if not self.settings.quiet:
print(self.bodies[occulted].name + " occulted by " +
self.bodies[i].name + " at t = %8.3f days" % t +
" lasting %5.2f minutes." % (d / MINUTE))
times.append(t)
occultors.append(self.bodies[i])
durations.append(d / MINUTE)
return times, occultors, durations
def plot_occultation(self, body, time, wavelength = 15., interval = 50,
gifname = None, spectral = False,
time_unit = 'BJD − 2,450,000',
trail_dur = 1., full_lightcurve = False,
**kwargs):
'''
Plots and animates an occultation event.
:param body: The occulted body
:type body: :py:class:`BODY`
:param float time: The approximate time of the occultation event \
(BJD − 2,450,000)
:param float wavelength: The wavelength in microns at which to plot \
the light curve. Must be within the wavelength grid. \
Default `15`
:param int interval: The interval between frames in the animation in \
ms. Default `50`
:param str gifname: If specified, saves the occultation animation as \
a :py:obj:`gif` in the current directory. Default :py:obj:`None`
:param bool spectral: Plot the light curve at different wavelengths? \
If :py:obj:`True`, plots the first, middle, and last \
wavelength in the wavelength grid. If :py:obj:`False`, plots \
the specified `wavelength`. Default :py:obj:`True`
:param float trail_dur: Duration of the orbital trails in days when \
plotting the orbits. The plotting code actually calls \
:py:mod:`rebound` to compute these. Default `1`
:param bool full_lightcurve: Plot the full system light curve? If \
:py:obj:`True`, the light curve will include the contribution \
of any other bodies being occulted at the given time. Default \
:py:obj:`False`, in which case only the flux from :py:obj:`body`\
is plotted.
:param kwargs: Any additional keyword arguments to be passed to \
:py:func:`plot_image`
:returns: :py:obj:`(fig, ax1, ax2, ax3)`, a figure instance and \
its three axes
.. plot::
:align: center
from scripts import occultation
import matplotlib.pyplot as pl
occultation.plot()
pl.show()
'''
# Have we computed the light curves?
assert self._computed, "Please run `compute()` first."
# Get the wavelength index
assert (wavelength >= self.bodies[0].wavelength[0]) and \
(wavelength >= self.bodies[0].wavelength[-1]), \
"Wavelength value outside of computed grid."
w = np.argmax(1e-6 * wavelength <= self.bodies[0].wavelength)
# Check file name
if gifname is not None:
if gifname.endswith(".gif"):
gifname = gifname[:-4]
# Get the occulted body
if isinstance(body, string_types):
p = np.argmax(self._names == body)