Coverage for starry/kepler.py : 97%
Hot-keys on this page
r m x p toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1# -*- coding: utf-8 -*-
2from . import config
3from ._constants import *
4from .maps import MapBase, RVBase, ReflectedBase
5from ._core import OpsSystem, math, linalg
6import numpy as np
7from astropy import units
8from inspect import getmro
9import matplotlib.pyplot as plt
10from matplotlib.animation import FuncAnimation
11from IPython.display import HTML
12import os
13import logging
15logger = logging.getLogger("starry.maps")
18__all__ = ["Primary", "Secondary", "System"]
21class Body(object):
22 """A generic body. Must be subclassed."""
24 def __init__(
25 self,
26 map,
27 r=1.0,
28 m=1.0,
29 prot=1.0,
30 t0=0.0,
31 theta0=0.0,
32 length_unit=units.Rsun,
33 mass_unit=units.Msun,
34 time_unit=units.day,
35 angle_unit=units.degree,
36 **kwargs,
37 ):
38 # Surface map
39 self.map = map
41 # Units
42 self.length_unit = length_unit
43 self.mass_unit = mass_unit
44 self.time_unit = time_unit
45 self.angle_unit = angle_unit
47 # Attributes
48 self.r = r
49 self.m = m
50 self.prot = prot
51 self.t0 = t0
52 self.theta0 = theta0
54 @property
55 def length_unit(self):
56 """An ``astropy.units`` unit defining the length metric for this body."""
57 return self._length_unit
59 @length_unit.setter
60 def length_unit(self, value):
61 assert value.physical_type == "length"
62 self._length_unit = value
63 self._length_factor = value.in_units(units.Rsun)
65 @property
66 def mass_unit(self):
67 """An ``astropy.units`` unit defining the mass metric for this body."""
68 return self._mass_unit
70 @mass_unit.setter
71 def mass_unit(self, value):
72 assert value.physical_type == "mass"
73 self._mass_unit = value
74 self._mass_factor = value.in_units(units.Msun)
76 @property
77 def time_unit(self):
78 """An ``astropy.units`` unit defining the time metric for this body."""
79 return self._time_unit
81 @time_unit.setter
82 def time_unit(self, value):
83 assert value.physical_type == "time"
84 self._time_unit = value
85 self._time_factor = value.in_units(units.day)
87 @property
88 def angle_unit(self):
89 """An ``astropy.units`` unit defining the angle metric for this body."""
90 return self._angle_unit
92 @angle_unit.setter
93 def angle_unit(self, value):
94 assert value.physical_type == "angle"
95 self._angle_unit = value
96 self._angle_factor = value.in_units(units.radian)
98 @property
99 def _angle_unit(self):
100 return self._map._angle_unit
102 @_angle_unit.setter
103 def _angle_unit(self, value):
104 self._map._angle_unit = value
106 @property
107 def _angle_factor(self):
108 return self._map._angle_factor
110 @_angle_factor.setter
111 def _angle_factor(self, value):
112 self._map._angle_factor = value
114 @property
115 def map(self):
116 """The surface map for this body."""
117 return self._map
119 @map.setter
120 def map(self, value):
121 assert MapBase in getmro(
122 type(value)
123 ), "The `map` attribute must be a `starry` map instance."
124 self._map = value
126 @property
127 def r(self):
128 """The radius in units of :py:attr:`length_unit`."""
129 return self._r / self._length_factor
131 @r.setter
132 def r(self, value):
133 self._r = math.cast(value * self._length_factor)
135 @property
136 def m(self):
137 """The mass in units of :py:attr:`mass_unit`."""
138 return self._m / self._mass_factor
140 @m.setter
141 def m(self, value):
142 self._m = math.cast(value * self._mass_factor)
144 @property
145 def prot(self):
146 """The rotation period in units of :py:attr:`time_unit`."""
147 return self._prot / self._time_factor
149 @prot.setter
150 def prot(self, value):
151 self._prot = math.cast(value * self._time_factor)
153 @property
154 def t0(self):
155 """A reference time in units of :py:attr:`time_unit`."""
156 return self._t0 / self._time_factor
158 @t0.setter
159 def t0(self, value):
160 self._t0 = math.cast(value * self._time_factor)
162 @property
163 def theta0(self):
164 """The map rotational phase at time :py:attr:`t0`."""
165 return self._theta0 / self._angle_factor
167 @theta0.setter
168 def theta0(self, value):
169 self._theta0 = math.cast(value * self._angle_factor)
171 def _check_kwargs(self, method, kwargs):
172 if not config.quiet:
173 for key in kwargs.keys():
174 message = "Invalid keyword `{0}` in call to `{1}()`. Ignoring."
175 message = message.format(key, method)
176 logger.warning(message)
179class Primary(Body):
180 """A primary (central) body.
182 Args:
183 map: The surface map of this body. This should be an instance
184 returned by :py:func:`starry.Map`.
185 r (scalar, optional): The radius of the body in units of
186 :py:attr:`length_unit`. Defaults to 1.0.
187 m (scalar, optional): The mass of the body in units of
188 :py:attr:`mass_unit`. Defaults to 1.0.
189 prot (scalar, optional): The rotation period of the body in units of
190 :py:attr:`time_unit`. Defaults to 1.0.
191 t0 (scalar, optional): A reference time in units of
192 :py:attr:`time_unit`. Defaults to 0.0.
193 theta0 (scalar, optional): The rotational phase of the map at time
194 :py:attr:`t0` in units of :py:attr:`angle_unit`. Defaults to 0.0.
195 length_unit (optional): An ``astropy.units`` unit defining the
196 distance metric for this object. Defaults to
197 :py:attr:`astropy.units.Rsun.`
198 mass_unit (optional): An ``astropy.units`` unit defining the
199 mass metric for this object. Defaults to
200 :py:attr:`astropy.units.Msun.`
201 time_unit (optional): An ``astropy.units`` unit defining the
202 time metric for this object. Defaults to
203 :py:attr:`astropy.units.day.`
204 angle_unit (optional): An ``astropy.units`` unit defining the
205 angular metric for this object. Defaults to
206 :py:attr:`astropy.units.degree.`
207 """
209 def __init__(self, map, **kwargs):
210 # Initialize `Body`
211 super(Primary, self).__init__(map, **kwargs)
212 for kw in [
213 "r",
214 "m",
215 "prot",
216 "t0",
217 "theta0",
218 "length_unit",
219 "mass_unit",
220 "time_unit",
221 "angle_unit",
222 ]:
223 kwargs.pop(kw, None)
224 self._check_kwargs("Primary", kwargs)
227class Secondary(Body):
228 """A secondary (orbiting) body.
230 Args:
231 map: The surface map of this body. This should be an instance
232 returned by :py:func:`starry.Map`.
233 r (scalar, optional): The radius of the body in units of
234 :py:attr:`length_unit`. Defaults to 1.0.
235 m (scalar, optional): The mass of the body in units of
236 :py:attr:`mass_unit`. Defaults to 1.0.
237 a (scalar, optional): The semi-major axis of the body in units of
238 :py:attr:`time_unit`. Defaults to 1.0. If :py:attr:`porb` is
239 also provided, this value is ignored.
240 porb (scalar, optional): The orbital period of the body in units of
241 :py:attr:`time_unit`. Defaults to 1.0. Setting this value
242 overrides :py:attr:`a`.
243 prot (scalar, optional): The rotation period of the body in units of
244 :py:attr:`time_unit`. Defaults to 1.0.
245 t0 (scalar, optional): A reference time in units of
246 :py:attr:`time_unit`. This is taken to be the time of a reference
247 transit. Defaults to 0.0.
248 ecc (scalar, optional): The orbital eccentricity of the body.
249 Defaults to 0.
250 w, omega (scalar, optional): The argument of pericenter of the body
251 in units of :py:attr:`angle_unit`. Defaults to 90 degrees.
252 Omega (scalar, optional): The longitude of ascending node of the
253 body in units of :py:attr:`angle_unit`. Defaults to 0 degrees.
254 inc (scalar, optional): The orbital inclination of the body in
255 units of :py:attr:`angle_unit`. Defaults to 90 degrees.
256 theta0 (scalar, optional): The rotational phase of the map at time
257 :py:attr:`t0` in units of :py:attr:`angle_unit`. Defaults to
258 0.0.
259 length_unit (optional): An ``astropy.units`` unit defining the
260 distance metric for this object. Defaults to
261 :py:attr:`astropy.units.Rsun.`
262 mass_unit (optional): An ``astropy.units`` unit defining the
263 mass metric for this object. Defaults to
264 :py:attr:`astropy.units.Msun.`
265 time_unit (optional): An ``astropy.units`` unit defining the
266 time metric for this object. Defaults to
267 :py:attr:`astropy.units.day.`
268 angle_unit (optional): An ``astropy.units`` unit defining the
269 angular metric for this object. Defaults to
270 :py:attr:`astropy.units.degree.`
271 """
273 def __init__(self, map, **kwargs):
274 # Initialize `Body`
275 super(Secondary, self).__init__(map, **kwargs)
276 for kw in [
277 "r",
278 "m",
279 "prot",
280 "t0",
281 "theta0",
282 "length_unit",
283 "mass_unit",
284 "time_unit",
285 "angle_unit",
286 ]:
287 kwargs.pop(kw, None)
289 # Attributes
290 if kwargs.get("porb", None) is not None:
291 self.porb = kwargs.pop("porb", None)
292 elif kwargs.get("a", None) is not None:
293 self.a = kwargs.pop("a", None)
294 else:
295 raise ValueError("Must provide a value for either `porb` or `a`.")
296 self.ecc = kwargs.pop("ecc", 0.0)
297 self.w = kwargs.pop(
298 "w", kwargs.pop("omega", 0.5 * np.pi / self._angle_factor)
299 )
300 self.Omega = kwargs.pop("Omega", 0.0)
301 self.inc = kwargs.pop("inc", 0.5 * np.pi / self._angle_factor)
302 self._check_kwargs("Secondary", kwargs)
304 @property
305 def porb(self):
306 """The orbital period in units of :py:attr:`time_unit`.
308 .. note::
309 Setting this value overrides the value of :py:attr:`a`.
310 """
311 if self._porb == 0.0:
312 return None
313 else:
314 return self._porb / self._time_factor
316 @porb.setter
317 def porb(self, value):
318 self._porb = math.cast(value * self._time_factor)
319 self._a = 0.0
321 @property
322 def a(self):
323 """The semi-major axis in units of :py:attr:`length_unit`.
325 .. note::
326 Setting this value overrides the value of :py:attr:`porb`.
327 """
328 if self._a == 0.0:
329 return None
330 else:
331 return self._a / self._length_factor
333 @a.setter
334 def a(self, value):
335 self._a = math.cast(value * self._length_factor)
336 self._porb = 0.0
338 @property
339 def ecc(self):
340 """The orbital eccentricity."""
341 return self._ecc
343 @ecc.setter
344 def ecc(self, value):
345 self._ecc = value
347 @property
348 def w(self):
349 """The longitude of pericenter in units of :py:attr:`angle_unit`."""
350 return self._w / self._angle_factor
352 @w.setter
353 def w(self, value):
354 self._w = math.cast(value * self._angle_factor)
356 @property
357 def omega(self):
358 """Alias for the longitude of pericenter :py:attr:`w`."""
359 return self.w
361 @omega.setter
362 def omega(self, value):
363 self.w = value
365 @property
366 def Omega(self):
367 """The longitude of ascending node in units of :py:attr:`angle_unit`."""
368 return self._Omega / self._angle_factor
370 @Omega.setter
371 def Omega(self, value):
372 self._Omega = math.cast(value * self._angle_factor)
374 @property
375 def inc(self):
376 """The orbital inclination in units of :py:attr:`angle_unit`."""
377 return self._inc / self._angle_factor
379 @inc.setter
380 def inc(self, value):
381 self._inc = math.cast(value * self._angle_factor)
384class System(object):
385 """A system of bodies in Keplerian orbits about a central primary body.
387 Args:
388 primary (:py:class:`Primary`): The central body.
389 secondaries (:py:class:`Secondary`): One or more secondary bodies
390 in orbit about the primary.
391 time_unit (optional): An ``astropy.units`` unit defining the
392 time metric for this object. Defaults to
393 :py:attr:`astropy.units.day.`
394 light_delay (bool, optional): Account for the light travel time
395 delay to the barycenter of the system? Default is False.
396 texp (scalar): The exposure time of each observation. This can be a
397 scalar or a tensor with the same shape as ``t``. If ``texp`` is
398 provided, ``t`` is assumed to indicate the timestamp at the middle
399 of an exposure of length ``texp``.
400 oversample (int): The number of function evaluations to use when
401 numerically integrating the exposure time.
402 order (int): The order of the numerical integration scheme. This must
403 be one of the following: ``0`` for a centered Riemann sum
404 (equivalent to the "resampling" procedure suggested by Kipping 2010),
405 ``1`` for the trapezoid rule, or ``2`` for Simpson’s rule.
406 """
408 def _no_spectral(self):
409 if self._primary._map.nw is not None: # pragma: no cover
410 raise NotImplementedError(
411 "Method not yet implemented for spectral maps."
412 )
414 def __init__(
415 self,
416 primary,
417 *secondaries,
418 time_unit=units.day,
419 light_delay=False,
420 texp=None,
421 oversample=7,
422 order=0,
423 ):
424 # Units
425 self.time_unit = time_unit
426 self._light_delay = bool(light_delay)
427 if texp is None:
428 self._texp = 0.0
429 else:
430 self._texp = texp
431 assert self._texp >= 0.0, "Parameter `texp` must be >= 0."
432 self._oversample = int(oversample)
433 assert self._oversample > 0, "Parameter `oversample` must be > 0."
434 self._order = int(order)
435 assert self._order in [0, 1, 2], "Invalid value for parameter `order`."
437 # Primary body
438 assert (
439 type(primary) is Primary
440 ), "Argument `primary` must be an instance of `Primary`."
441 assert (
442 primary._map.__props__["reflected"] == False
443 ), "Reflected light map not allowed for the primary body."
444 self._primary = primary
445 self._rv = primary._map.__props__["rv"]
447 # Secondary bodies
448 assert len(secondaries) > 0, "There must be at least one secondary."
449 for sec in secondaries:
450 assert type(sec) is Secondary, (
451 "Argument `*secondaries` must be a sequence of "
452 "`Secondary` instances."
453 )
454 assert (
455 sec._map.nw == self._primary._map.nw
456 ), "All bodies must have the same number of wavelength bins `nw`."
457 assert sec._map.__props__["rv"] == self._rv, (
458 "Radial velocity must be enabled "
459 "for either all or none of the bodies."
460 )
462 reflected = [sec._map.__props__["reflected"] for sec in secondaries]
463 if np.all(reflected):
464 self._reflected = True
465 elif np.any(reflected):
466 raise ValueError(
467 "Reflected light must be enabled "
468 "for either all or none of the secondaries."
469 )
470 else:
471 self._reflected = False
472 self._secondaries = secondaries
474 # All bodies
475 self._bodies = [self._primary] + list(self._secondaries)
477 # Indices of each of the bodies in the design matrix
478 Ny = [self._primary._map.Ny] + [
479 sec._map.Ny for sec in self._secondaries
480 ]
481 self._inds = []
482 cur = 0
483 for N in Ny:
484 self._inds.append(cur + np.arange(N))
485 cur += N
487 # Theano ops class
488 self.ops = OpsSystem(
489 self._primary,
490 self._secondaries,
491 reflected=self._reflected,
492 rv=self._rv,
493 light_delay=self._light_delay,
494 texp=self._texp,
495 oversample=self._oversample,
496 order=self._order,
497 )
499 # Solve stuff
500 self._flux = None
501 self._C = None
502 self._solution = None
503 self._solved_bodies = []
505 @property
506 def light_delay(self):
507 """Account for the light travel time delay? *Read-only*"""
508 return self._light_delay
510 @property
511 def texp(self):
512 """The exposure time in units of :py:attr:`time_unit`. *Read-only*"""
514 @property
515 def oversample(self):
516 """Oversample factor when integrating over exposure time. *Read-only*"""
517 return self._oversample
519 @property
520 def order(self):
521 """The order of the numerical integration scheme. *Read-only*
523 - ``0``: a centered Riemann sum
524 - ``1``: trapezoid rule
525 - ``2``: Simpson’s rule
526 """
527 return self._order
529 @property
530 def time_unit(self):
531 """An ``astropy.units`` unit defining the time metric for the system."""
532 return self._time_unit
534 @time_unit.setter
535 def time_unit(self, value):
536 assert value.physical_type == "time"
537 self._time_unit = value
538 self._time_factor = value.in_units(units.day)
540 @property
541 def primary(self):
542 """The primary (central) object in the Keplerian system."""
543 return self._primary
545 @property
546 def secondaries(self):
547 """A list of the secondary (orbiting) object(s) in the Keplerian system."""
548 return self._secondaries
550 @property
551 def bodies(self):
552 """A list of all objects in the Keplerian system."""
553 return self._bodies
555 @property
556 def map_indices(self):
557 """A list of the indices corresponding to each body in the design matrix."""
558 return self._inds
560 def show(
561 self,
562 t,
563 cmap="plasma",
564 res=300,
565 interval=75,
566 file=None,
567 figsize=(3, 3),
568 html5_video=True,
569 window_pad=1.0,
570 ):
571 """Visualize the Keplerian system.
573 Args:
574 t (scalar or vector): The time(s) at which to evaluate the orbit and
575 the map in units of :py:attr:`time_unit`.
576 cmap (string or colormap instance, optional): The matplotlib colormap
577 to use. Defaults to ``plasma``.
578 res (int, optional): The resolution of the map in pixels on a
579 side. Defaults to 300.
580 figsize (tuple, optional): Figure size in inches. Default is
581 (3, 3) for orthographic maps and (7, 3.5) for rectangular
582 maps.
583 interval (int, optional): Interval between frames in milliseconds
584 (animated maps only). Defaults to 75.
585 file (string, optional): The file name (including the extension)
586 to save the animation to (animated maps only). Defaults to None.
587 html5_video (bool, optional): If rendering in a Jupyter notebook,
588 display as an HTML5 video? Default is True. If False, displays
589 the animation using Javascript (file size will be larger.)
590 window_pad (float, optional): Padding around the primary in units
591 of the primary radius. Bodies outside of this window will be
592 cropped. Default is 1.0.
593 """
594 # Not yet implemented
595 if self._primary._map.nw is not None: # pragma: no cover
596 raise NotImplementedError(
597 "Method not implemented for spectral maps."
598 )
600 # Render the maps & get the orbital positions
601 if self._rv:
602 self._primary.map._set_RV_filter()
603 for sec in self._secondaries:
604 sec.map._set_RV_filter()
605 img_pri, img_sec, x, y, z = self.ops.render(
606 math.reshape(math.to_array_or_tensor(t), [-1]) * self._time_factor,
607 res,
608 self._primary._r,
609 self._primary._m,
610 self._primary._prot,
611 self._primary._t0,
612 self._primary._theta0,
613 self._primary._map._amp,
614 self._primary._map._inc,
615 self._primary._map._obl,
616 self._primary._map._y,
617 self._primary._map._u,
618 self._primary._map._f,
619 self._primary._map._alpha,
620 math.to_array_or_tensor([sec._r for sec in self._secondaries]),
621 math.to_array_or_tensor([sec._m for sec in self._secondaries]),
622 math.to_array_or_tensor([sec._prot for sec in self._secondaries]),
623 math.to_array_or_tensor([sec._t0 for sec in self._secondaries]),
624 math.to_array_or_tensor(
625 [sec._theta0 for sec in self._secondaries]
626 ),
627 self._get_periods(),
628 math.to_array_or_tensor([sec._ecc for sec in self._secondaries]),
629 math.to_array_or_tensor([sec._w for sec in self._secondaries]),
630 math.to_array_or_tensor([sec._Omega for sec in self._secondaries]),
631 math.to_array_or_tensor([sec._inc for sec in self._secondaries]),
632 math.to_array_or_tensor(
633 [sec._map._amp for sec in self._secondaries]
634 ),
635 math.to_array_or_tensor(
636 [sec._map._inc for sec in self._secondaries]
637 ),
638 math.to_array_or_tensor(
639 [sec._map._obl for sec in self._secondaries]
640 ),
641 math.to_array_or_tensor(
642 [sec._map._y for sec in self._secondaries]
643 ),
644 math.to_array_or_tensor(
645 [sec._map._u for sec in self._secondaries]
646 ),
647 math.to_array_or_tensor(
648 [sec._map._f for sec in self._secondaries]
649 ),
650 math.to_array_or_tensor(
651 [sec._map._alpha for sec in self._secondaries]
652 ),
653 )
655 # Convert to units of the primary radius
656 fac = np.reshape(
657 [sec._length_factor for sec in self._secondaries], [-1, 1]
658 )
659 fac = fac * self._primary._r
660 x, y, z = x / fac, y / fac, z / fac
661 r = math.to_array_or_tensor([sec._r for sec in self._secondaries])
662 r = r / self._primary._r
664 # Evaluate if needed
665 if config.lazy:
666 img_pri = img_pri.eval()
667 img_sec = img_sec.eval()
668 x = x.eval()
669 y = y.eval()
670 z = z.eval()
671 r = r.eval()
673 # We need this to be of shape (nplanet, nframe)
674 x = x.T
675 y = y.T
676 z = z.T
678 # Ensure we have an array of frames
679 if len(img_pri.shape) == 3:
680 nframes = img_pri.shape[0]
681 else: # pragma: no cover
682 nframes = 1
683 img_pri = np.reshape(img_pri, (1,) + img_pri.shape)
684 img_sec = np.reshape(img_sec, (1,) + img_sec.shape)
685 animated = nframes > 1
687 # Set up the plot
688 fig, ax = plt.subplots(1, figsize=figsize)
689 ax.axis("off")
690 ax.set_xlim(-1.0 - window_pad, 1.0 + window_pad)
691 ax.set_ylim(-1.0 - window_pad, 1.0 + window_pad)
693 # Render the first frame
694 img = [None for n in range(1 + len(self._secondaries))]
695 circ = [None for n in range(1 + len(self._secondaries))]
696 extent = np.array([-1.0, 1.0, -1.0, 1.0])
697 img[0] = ax.imshow(
698 img_pri[0],
699 origin="lower",
700 extent=extent,
701 cmap=cmap,
702 interpolation="none",
703 vmin=np.nanmin(img_pri),
704 vmax=np.nanmax(img_pri),
705 animated=animated,
706 zorder=0.0,
707 )
708 circ[0] = plt.Circle(
709 (0, 0), 1, color="k", fill=False, zorder=1e-3, lw=2
710 )
711 ax.add_artist(circ[0])
712 for i, _ in enumerate(self._secondaries):
713 extent = np.array([x[i, 0], x[i, 0], y[i, 0], y[i, 0]]) + (
714 r[i] * np.array([-1.0, 1.0, -1.0, 1.0])
715 )
716 img[i + 1] = ax.imshow(
717 img_sec[i, 0],
718 origin="lower",
719 extent=extent,
720 cmap=cmap,
721 interpolation="none",
722 vmin=np.nanmin(img_sec),
723 vmax=np.nanmax(img_sec),
724 animated=animated,
725 zorder=z[i, 0],
726 )
727 circ[i] = plt.Circle(
728 (x[i, 0], y[i, 0]),
729 r[i],
730 color="k",
731 fill=False,
732 zorder=z[i, 0] + 1e-3,
733 lw=2,
734 )
735 ax.add_artist(circ[i])
737 # Animation
738 if animated:
740 def updatefig(k):
742 # Update Primary map
743 img[0].set_array(img_pri[k])
745 # Update Secondary maps & positions
746 for i, _ in enumerate(self._secondaries):
747 extent = np.array([x[i, k], x[i, k], y[i, k], y[i, k]]) + (
748 r[i] * np.array([-1.0, 1.0, -1.0, 1.0])
749 )
750 if np.any(np.abs(extent) < 1.0 + window_pad):
751 img[i + 1].set_array(img_sec[i, k])
752 img[i + 1].set_extent(extent)
753 img[i + 1].set_zorder(z[i, k])
754 circ[i].center = (x[i, k], y[i, k])
755 circ[i].set_zorder(z[i, k] + 1e-3)
757 return img + circ
759 ani = FuncAnimation(
760 fig, updatefig, interval=interval, blit=False, frames=nframes
761 )
763 # Business as usual
764 if (file is not None) and (file != ""):
765 if file.endswith(".mp4"):
766 ani.save(file, writer="ffmpeg")
767 elif file.endswith(".gif"):
768 ani.save(file, writer="imagemagick")
769 else: # pragma: no cover
770 # Try and see what happens!
771 ani.save(file)
772 plt.close()
773 else: # pragma: no cover
774 try:
775 if "zmqshell" in str(type(get_ipython())):
776 plt.close()
777 if html5_video:
778 display(HTML(ani.to_html5_video()))
779 else:
780 display(HTML(ani.to_jshtml()))
781 else:
782 raise NameError("")
783 except NameError:
784 plt.show()
785 plt.close()
787 # Matplotlib generates an annoying empty
788 # file when producing an animation. Delete it.
789 try:
790 os.remove("None0000000.png")
791 except FileNotFoundError:
792 pass
794 else:
796 if (file is not None) and (file != ""):
797 fig.savefig(file)
798 plt.close()
799 else: # pragma: no cover
800 plt.show()
802 if self._rv:
803 self._primary.map._unset_RV_filter()
804 for sec in self._secondaries:
805 sec.map._unset_RV_filter()
807 def design_matrix(self, t):
808 """Compute the system flux design matrix at times ``t``.
810 .. note::
812 This is the *unweighted* design matrix, i.e., it does not
813 include the scaling by the amplitude of each body's map.
814 To perform this weighting, do
816 .. code-block:: python
818 X = sys.design_matrix(**kwargs)
819 for i, body in zip(sys.map_indices, sys.bodies):
820 X[:, i] *= body.map.amp
822 Args:
823 t (scalar or vector): An array of times at which to evaluate
824 the design matrix in units of :py:attr:`time_unit`.
825 """
826 return self.ops.X(
827 math.reshape(math.to_array_or_tensor(t), [-1]) * self._time_factor,
828 self._primary._r,
829 self._primary._m,
830 self._primary._prot,
831 self._primary._t0,
832 self._primary._theta0,
833 math.ones_like(self._primary._map._amp),
834 self._primary._map._inc,
835 self._primary._map._obl,
836 self._primary._map._u,
837 self._primary._map._f,
838 self._primary._map._alpha,
839 math.to_array_or_tensor([sec._r for sec in self._secondaries]),
840 math.to_array_or_tensor([sec._m for sec in self._secondaries]),
841 math.to_array_or_tensor([sec._prot for sec in self._secondaries]),
842 math.to_array_or_tensor([sec._t0 for sec in self._secondaries]),
843 math.to_array_or_tensor(
844 [sec._theta0 for sec in self._secondaries]
845 ),
846 self._get_periods(),
847 math.to_array_or_tensor([sec._ecc for sec in self._secondaries]),
848 math.to_array_or_tensor([sec._w for sec in self._secondaries]),
849 math.to_array_or_tensor([sec._Omega for sec in self._secondaries]),
850 math.to_array_or_tensor([sec._inc for sec in self._secondaries]),
851 math.to_array_or_tensor(
852 [math.ones_like(sec._map._amp) for sec in self._secondaries]
853 ),
854 math.to_array_or_tensor(
855 [sec._map._inc for sec in self._secondaries]
856 ),
857 math.to_array_or_tensor(
858 [sec._map._obl for sec in self._secondaries]
859 ),
860 math.to_array_or_tensor(
861 [sec._map._u for sec in self._secondaries]
862 ),
863 math.to_array_or_tensor(
864 [sec._map._f for sec in self._secondaries]
865 ),
866 math.to_array_or_tensor(
867 [sec._map._alpha for sec in self._secondaries]
868 ),
869 )
871 def flux(self, t, total=True):
872 """Compute the system flux at times ``t``.
874 Args:
875 t (scalar or vector): An array of times at which to evaluate
876 the flux in units of :py:attr:`time_unit`.
877 total (bool, optional): Return the total system flux? Defaults to
878 True. If False, returns arrays corresponding to the flux
879 from each body.
880 """
881 X = self.design_matrix(t)
883 # Weight the ylms by amplitude
884 if self._reflected:
885 # If we're doing reflected light, scale the amplitude of
886 # each of the secondaries by the amplitude of the primary
887 # (the illumination source).
888 ay = [self._primary.map.amp * self._primary._map._y] + [
889 self._primary.map.amp * body.map.amp * body._map._y
890 for body in self._secondaries
891 ]
892 else:
893 ay = [body.map.amp * body._map._y for body in self._bodies]
895 if total:
896 return math.dot(X, math.concatenate(ay))
897 else:
898 return [
899 math.dot(X[:, idx], ay[i]) for i, idx in enumerate(self._inds)
900 ]
902 def rv(self, t, keplerian=True, total=True):
903 """Compute the observed radial velocity of the system at times ``t``.
905 Args:
906 t (scalar or vector): An array of times at which to evaluate
907 the radial velocity in units of :py:attr:`time_unit`.
908 keplerian (bool): Include the Keplerian component of the radial
909 velocity of the primary? Default is True. If False, this
910 method returns a model for only the radial velocity anomaly
911 due to transits (the Rossiter-McLaughlin effect) and
912 time-variable surface features (Doppler tomography) for all
913 bodies in the system.
914 total (bool, optional): Return the total system RV? Defaults to
915 True. If False, returns arrays corresponding to the RV
916 contribution from each body.
918 """
919 rv = self.ops.rv(
920 math.reshape(math.to_array_or_tensor(t), [-1]) * self._time_factor,
921 self._primary._r,
922 self._primary._m,
923 self._primary._prot,
924 self._primary._t0,
925 self._primary._theta0,
926 self._primary._map._amp,
927 self._primary._map._inc,
928 self._primary._map._obl,
929 self._primary._map._y,
930 self._primary._map._u,
931 self._primary._map._alpha,
932 self._primary._map._veq,
933 math.to_array_or_tensor([sec._r for sec in self._secondaries]),
934 math.to_array_or_tensor([sec._m for sec in self._secondaries]),
935 math.to_array_or_tensor([sec._prot for sec in self._secondaries]),
936 math.to_array_or_tensor([sec._t0 for sec in self._secondaries]),
937 math.to_array_or_tensor(
938 [sec._theta0 for sec in self._secondaries]
939 ),
940 self._get_periods(),
941 math.to_array_or_tensor([sec._ecc for sec in self._secondaries]),
942 math.to_array_or_tensor([sec._w for sec in self._secondaries]),
943 math.to_array_or_tensor([sec._Omega for sec in self._secondaries]),
944 math.to_array_or_tensor([sec._inc for sec in self._secondaries]),
945 math.to_array_or_tensor(
946 [sec._map._amp for sec in self._secondaries]
947 ),
948 math.to_array_or_tensor(
949 [sec._map._inc for sec in self._secondaries]
950 ),
951 math.to_array_or_tensor(
952 [sec._map._obl for sec in self._secondaries]
953 ),
954 math.to_array_or_tensor(
955 [sec._map._y for sec in self._secondaries]
956 ),
957 math.to_array_or_tensor(
958 [sec._map._u for sec in self._secondaries]
959 ),
960 math.to_array_or_tensor(
961 [sec._map._alpha for sec in self._secondaries]
962 ),
963 math.to_array_or_tensor(
964 [sec._map._veq for sec in self._secondaries]
965 ),
966 np.array(keplerian),
967 )
968 if total:
969 return math.sum(rv, axis=0)
970 else:
971 return rv
973 def position(self, t):
974 """Compute the Cartesian positions of all bodies at times ``t``.
976 Args:
977 t (scalar or vector): An array of times at which to evaluate
978 the position in units of :py:attr:`time_unit`.
979 """
980 x, y, z = self.ops.position(
981 math.reshape(math.to_array_or_tensor(t), [-1]) * self._time_factor,
982 self._primary._m,
983 self._primary._t0,
984 math.to_array_or_tensor([sec._m for sec in self._secondaries]),
985 math.to_array_or_tensor([sec._t0 for sec in self._secondaries]),
986 self._get_periods(),
987 math.to_array_or_tensor([sec._ecc for sec in self._secondaries]),
988 math.to_array_or_tensor([sec._w for sec in self._secondaries]),
989 math.to_array_or_tensor([sec._Omega for sec in self._secondaries]),
990 math.to_array_or_tensor([sec._inc for sec in self._secondaries]),
991 )
992 fac = np.reshape(
993 [self._primary._length_factor]
994 + [sec._length_factor for sec in self._secondaries],
995 [-1, 1],
996 )
997 return (x / fac, y / fac, z / fac)
999 def _get_periods(self):
1000 periods = [None for sec in self._secondaries]
1001 for i, sec in enumerate(self._secondaries):
1002 if sec.porb:
1003 periods[i] = sec.porb
1004 else:
1005 periods[i] = (
1006 (2 * np.pi)
1007 * sec._a ** (3 / 2)
1008 / (math.sqrt(G_grav * (self._primary._m + sec._m)))
1009 )
1010 return math.to_array_or_tensor(periods)
1012 def set_data(self, flux, C=None, cho_C=None):
1013 """Set the data vector and covariance matrix.
1015 This method is required by the :py:meth:`solve` method, which
1016 analytically computes the posterior over surface maps for all bodies
1017 in the system given a dataset and a prior, provided both are described
1018 as multivariate Gaussians.
1020 Args:
1021 flux (vector): The observed system light curve.
1022 C (scalar, vector, or matrix): The data covariance. This may be
1023 a scalar, in which case the noise is assumed to be
1024 homoscedastic, a vector, in which case the covariance
1025 is assumed to be diagonal, or a matrix specifying the full
1026 covariance of the dataset. Default is None. Either `C` or
1027 `cho_C` must be provided.
1028 cho_C (matrix): The lower Cholesky factorization of the data
1029 covariance matrix. Defaults to None. Either `C` or
1030 `cho_C` must be provided.
1031 """
1032 self._flux = math.cast(flux)
1033 self._C = linalg.Covariance(C=C, cho_C=cho_C, N=self._flux.shape[0])
1035 def solve(self, *, design_matrix=None, t=None):
1036 """Solve the least-squares problem for the posterior over maps for all bodies.
1038 This method solves the generalized least squares problem given a system
1039 light curve and its covariance (set via the :py:meth:`set_data` method)
1040 and a Gaussian prior on the spherical harmonic coefficients
1041 (set via the :py:meth:`set_prior` method). The map amplitudes and
1042 coefficients of each of the bodies in the system are then set to the
1043 maximum a posteriori (MAP) solution.
1045 Args:
1046 design_matrix (matrix, optional): The flux design matrix, the
1047 quantity returned by :py:meth:`design_matrix`. Default is
1048 None, in which case this is computed based on ``kwargs``.
1049 t (vector, optional): The vector of times at which to evaluate
1050 :py:meth:`design_matrix`, if a design matrix is not provided.
1051 Default is None.
1053 Returns:
1054 The posterior mean for the spherical harmonic \
1055 coefficients `l > 0` and the Cholesky factorization of the \
1056 posterior covariance of all of the bodies in the system, \
1057 stacked in order (primary, followed by each of the secondaries \
1058 in the order they were provided.)
1060 .. note::
1061 Users may call the :py:meth:`draw` method of this class to draw
1062 from the posterior after calling :py:meth:`solve`.
1063 """
1064 # TODO: Implement for spectral maps?
1065 self._no_spectral()
1067 # Check that the data is set
1068 if self._flux is None or self._C is None:
1069 raise ValueError("Please provide a dataset with `set_data()`.")
1071 # Get the full design matrix
1072 if design_matrix is None:
1073 assert t is not None, "Please provide a time vector `t`."
1074 design_matrix = self.design_matrix(t)
1075 X = math.cast(design_matrix)
1077 # Get the data vector
1078 f = math.cast(self._flux)
1080 # Check for bodies whose priors are set
1081 self._solved_bodies = []
1082 inds = []
1083 dense_L = False
1084 for k, body in enumerate(self._bodies):
1086 if body.map._mu is None or body.map._L is None:
1088 # Subtract out this term from the data vector,
1089 # since it is fixed
1090 f -= body.map.amp * math.dot(X[:, self._inds[k]], body.map.y)
1092 else:
1094 # Add to our list of indices/bodies to solve for
1095 inds.extend(self._inds[k])
1096 self._solved_bodies.append(body)
1097 if body.map._L.kind in ["matrix", "cholesky"]:
1098 dense_L = True
1100 # Do we have at least one body?
1101 if len(self._solved_bodies) == 0:
1102 raise ValueError("Please provide a prior for at least one body.")
1104 # Keep only the terms we'll solve for
1105 X = X[:, inds]
1107 # Stack our priors
1108 mu = math.concatenate([body.map._mu for body in self._solved_bodies])
1110 if not dense_L:
1111 # We can just concatenate vectors
1112 LInv = math.concatenate(
1113 [
1114 body.map._L.inverse * math.ones(body.map.Ny)
1115 for body in self._solved_bodies
1116 ]
1117 )
1118 else:
1119 # FACT: The inverse of a block diagonal matrix
1120 # is the block diagonal matrix of the inverses.
1121 LInv = math.block_diag(
1122 *[
1123 body.map._L.inverse * math.eye(body.map.Ny)
1124 for body in self._solved_bodies
1125 ]
1126 )
1128 # Compute the MAP solution
1129 self._solution = linalg.solve(X, f, self._C.cholesky, mu, LInv)
1131 # Set all the map vectors
1132 x, _ = self._solution
1133 n = 0
1134 for body in self._solved_bodies:
1135 inds = slice(n, n + body.map.Ny)
1136 body.map.amp = x[inds][0]
1137 body.map[1:, :] = x[inds][1:] / body.map.amp
1138 n += body.map.Ny
1140 # Return the mean and covariance
1141 return self._solution
1143 @property
1144 def solution(self):
1145 r"""The posterior probability distribution for the maps in the system.
1147 This is a tuple containing the mean and lower Cholesky factorization of the
1148 covariance of the amplitude-weighted spherical harmonic coefficient vectors,
1149 obtained by solving the regularized least-squares problem
1150 via the :py:meth:`solve` method.
1152 Note that to obtain the actual covariance matrix from the lower Cholesky
1153 factorization :math:`L`, simply compute :math:`L L^\top`.
1155 Note also that this is the posterior for the **amplitude-weighted**
1156 map vectors. Under this convention, the map amplitude is equal to the
1157 first term of the vector of each body and the spherical harmonic coefficients are
1158 equal to the vector normalized by the first term.
1159 """
1160 if self._solution is None:
1161 raise ValueError("Please call `solve()` first.")
1162 return self._solution
1164 def draw(self):
1165 """
1166 Draw a map from the posterior distribution and set
1167 the :py:attr:`y` map vector of each body.
1169 Users should call :py:meth:`solve` to enable this attribute.
1170 """
1171 if self._solution is None:
1172 raise ValueError("Please call `solve()` first.")
1174 # Number of coefficients
1175 N = np.sum([body.map.Ny for body in self._solved_bodies])
1177 # Fast multivariate sampling using the Cholesky factorization
1178 yhat, cho_ycov = self._solution
1179 u = math.cast(np.random.randn(N))
1180 x = yhat + math.dot(cho_ycov, u)
1182 # Set all the map vectors
1183 n = 0
1184 for body in self._solved_bodies:
1185 inds = slice(n, n + body.map.Ny)
1186 body.map.amp = x[inds][0]
1187 body.map[1:, :] = x[inds][1:] / body.map.amp
1188 n += body.map.Ny
1190 def lnlike(self, *, design_matrix=None, t=None, woodbury=True):
1191 """Returns the log marginal likelihood of the data given a design matrix.
1193 This method computes the marginal likelihood (marginalized over the
1194 spherical harmonic coefficients of all bodies) given a system
1195 light curve and its covariance (set via the :py:meth:`set_data` method)
1196 and a Gaussian prior on the spherical harmonic coefficients
1197 (set via the :py:meth:`set_prior` method).
1199 Args:
1200 design_matrix (matrix, optional): The flux design matrix, the
1201 quantity returned by :py:meth:`design_matrix`. Default is
1202 None, in which case this is computed based on ``kwargs``.
1203 t (vector, optional): The vector of times at which to evaluate
1204 :py:meth:`design_matrix`, if a design matrix is not provided.
1205 Default is None.
1206 woodbury (bool, optional): Solve the linear problem using the
1207 Woodbury identity? Default is True. The
1208 `Woodbury identity <https://en.wikipedia.org/wiki/Woodbury_matrix_identity>`_
1209 is used to speed up matrix operations in the case that the
1210 number of data points is much larger than the number of
1211 spherical harmonic coefficients. In this limit, it can
1212 speed up the code by more than an order of magnitude. Keep
1213 in mind that the numerical stability of the Woodbury identity
1214 is not great, so if you're getting strange results try
1215 disabling this. It's also a good idea to disable this in the
1216 limit of few data points and large spherical harmonic degree.
1218 Returns:
1219 lnlike: The log marginal likelihood.
1220 """
1221 # TODO: Implement for spectral maps?
1222 self._no_spectral()
1224 # Check that the data is set
1225 if self._flux is None or self._C is None:
1226 raise ValueError("Please provide a dataset with `set_data()`.")
1228 # Get the full design matrix
1229 if design_matrix is None:
1230 assert t is not None, "Please provide a time vector `t`."
1231 design_matrix = self.design_matrix(t)
1232 X = math.cast(design_matrix)
1234 # Get the data vector
1235 f = math.cast(self._flux)
1237 # Check for bodies whose priors are set
1238 self._solved_bodies = []
1239 inds = []
1240 dense_L = False
1241 for k, body in enumerate(self._bodies):
1243 if body.map._mu is None or body.map._L is None:
1245 # Subtract out this term from the data vector,
1246 # since it is fixed
1247 f -= body.map.amp * math.dot(X[:, self._inds[k]], body.map.y)
1249 else:
1251 # Add to our list of indices/bodies to solve for
1252 inds.extend(self._inds[k])
1253 self._solved_bodies.append(body)
1254 if body.map._L.kind in ["matrix", "cholesky"]:
1255 dense_L = True
1257 # Do we have at least one body?
1258 if len(self._solved_bodies) == 0:
1259 raise ValueError("Please provide a prior for at least one body.")
1261 # Keep only the terms we'll solve for
1262 X = X[:, inds]
1264 # Stack our priors
1265 mu = math.concatenate([body.map._mu for body in self._solved_bodies])
1267 # Compute the likelihood
1268 if woodbury:
1269 if not dense_L:
1270 # We can just concatenate vectors
1271 LInv = math.concatenate(
1272 [
1273 body.map._L.inverse * math.ones(body.map.Ny)
1274 for body in self._solved_bodies
1275 ]
1276 )
1277 else:
1278 LInv = math.block_diag(
1279 *[
1280 body.map._L.inverse * math.eye(body.map.Ny)
1281 for body in self._solved_bodies
1282 ]
1283 )
1284 lndetL = math.cast(
1285 [body.map._L.lndet for body in self._solved_bodies]
1286 )
1287 return linalg.lnlike_woodbury(
1288 X, f, self._C.inverse, mu, LInv, self._C.lndet, lndetL
1289 )
1290 else:
1291 if not dense_L:
1292 # We can just concatenate vectors
1293 L = math.concatenate(
1294 [
1295 body.map._L.value * math.ones(body.map.Ny)
1296 for body in self._solved_bodies
1297 ]
1298 )
1299 else:
1300 L = math.block_diag(
1301 *[
1302 body.map._L.value * math.eye(body.map.Ny)
1303 for body in self._solved_bodies
1304 ]
1305 )
1306 return linalg.lnlike(X, f, self._C.value, mu, L)