def get_celerite_matrices(self, x, diag, **kwargs): x = tt.as_tensor_variable(x) diag = tt.as_tensor_variable(diag) ar, cr, ac, bc, cc, dc = self.coefficients a = diag + tt.sum(ar) + tt.sum(ac) arg = dc[None, :] * x[:, None] cos = tt.cos(arg) sin = tt.sin(arg) z = tt.zeros_like(x) U = tt.concatenate( ( ar[None, :] + z[:, None], ac[None, :] * cos + bc[None, :] * sin, ac[None, :] * sin - bc[None, :] * cos, ), axis=1, ) V = tt.concatenate( (tt.ones_like(ar)[None, :] + z[:, None], cos, sin), axis=1, ) c = tt.concatenate((cr, cc, cc)) return c, a, U, V
def grad(self, inputs, gradients): outputs = self(*inputs) grads = ( tt.zeros_like(outputs[n]) if isinstance(b.type, theano.gradient.DisconnectedType) else b for n, b in enumerate(gradients[: len(self.spec["outputs"])]) ) return self.rev_op(*chain(inputs, outputs, grads))
def __init__(self, gp, *args, **kwargs): if not HAS_PYMC3: raise ImportError( "pymc3 is required to use the CeleriteNormal distribution") super().__init__(*args, **kwargs) self.gp = gp self.mean = ( self.median) = self.mode = self.gp.mean_value + tt.zeros_like( self.gp._t)
def grad(self, inputs, gradients): M, e = inputs sinf, cosf = self(M, e) ecosf = e * cosf ome2 = 1 - e**2 dfdM = (1 + ecosf)**2 / ome2**1.5 dfde = (2 + ecosf) * sinf / ome2 bM = tt.zeros_like(M) be = tt.zeros_like(M) if not isinstance(gradients[0].type, aesara.gradient.DisconnectedType): bM += gradients[0] * cosf * dfdM be += gradients[0] * cosf * dfde if not isinstance(gradients[1].type, aesara.gradient.DisconnectedType): bM -= gradients[1] * sinf * dfdM be -= gradients[1] * sinf * dfde return [bM, be]
def _compute_light_curve(self, b, r, los=None): """Compute the light curve for a set of impact parameters and radii .. note:: The stellar radius is *not* included in this method so the coordinates should be in units of the star's radius. Args: b (tensor): A tensor of impact parameter values. r (tensor): A tensor of radius ratios with the same shape as ``b``. los (Optional[tensor]): The coordinates of the body along the line-of-sight. If ``los > 0`` the body is between the observer and the source. """ b = as_tensor_variable(b) if los is None: los = tt.ones_like(b) lc = quad_limbdark_light_curve(self.c, b, r) return tt.switch(tt.gt(los, 0), lc, tt.zeros_like(lc))
def get_relative_position(self, t, light_delay=False): """The planets' positions relative to the star Args: t: The times where the position should be evaluated. Returns: The components of the position vector at ``t`` in units of ``R_sun``. """ if light_delay: raise NotImplementedError( "Light travel time delay is not implemented for simple orbits") dt = tt.mod(tt.shape_padright(t) - self._ref_time, self.period) dt -= self._half_period x = tt.squeeze(self.speed * dt) y = tt.squeeze(self._b_norm + tt.zeros_like(dt)) m = tt.abs_(dt) < 0.5 * self.duration z = tt.squeeze(m * 1.0 - (~m) * 1.0) return x, y, z
def _get_consistent_inputs(a, period, rho_star, r_star, m_star, m_planet): if a is None and period is None: raise ValueError("values must be provided for at least one of a " "and period") if m_planet is not None: m_planet = as_tensor_variable(to_unit(m_planet, u.M_sun)) if a is not None: a = as_tensor_variable(to_unit(a, u.R_sun)) if m_planet is None: m_planet = tt.zeros_like(a) if period is not None: period = as_tensor_variable(to_unit(period, u.day)) if m_planet is None: m_planet = tt.zeros_like(period) # Compute the implied density if a and period are given implied_rho_star = False if a is not None and period is not None: if rho_star is not None or m_star is not None: raise ValueError("if both a and period are given, you can't " "also define rho_star or m_star") # Default to a stellar radius of 1 if not provided if r_star is None: r_star = as_tensor_variable(1.0) else: r_star = as_tensor_variable(to_unit(r_star, u.R_sun)) # Compute the implied mass via Kepler's 3rd law m_tot = 4 * np.pi * np.pi * a**3 / (G_grav * period**2) # Compute the implied density m_star = m_tot - m_planet vol_star = 4 * np.pi * r_star**3 / 3.0 rho_star = m_star / vol_star implied_rho_star = True # Make sure that the right combination of stellar parameters are given if r_star is None and m_star is None: r_star = 1.0 if rho_star is None: m_star = 1.0 if (not implied_rho_star) and sum( arg is None for arg in (rho_star, r_star, m_star)) != 1: raise ValueError("values must be provided for exactly two of " "rho_star, m_star, and r_star") if rho_star is not None and not implied_rho_star: if has_unit(rho_star): rho_star = as_tensor_variable( to_unit(rho_star, u.M_sun / u.R_sun**3)) else: rho_star = as_tensor_variable(rho_star) / gcc_per_sun if r_star is not None: r_star = as_tensor_variable(to_unit(r_star, u.R_sun)) if m_star is not None: m_star = as_tensor_variable(to_unit(m_star, u.M_sun)) # Work out the stellar parameters if rho_star is None: rho_star = 3 * m_star / (4 * np.pi * r_star**3) elif r_star is None: r_star = (3 * m_star / (4 * np.pi * rho_star))**(1 / 3) elif m_star is None: m_star = 4 * np.pi * r_star**3 * rho_star / 3.0 # Work out the planet parameters if a is None: a = (G_grav * (m_star + m_planet) * period**2 / (4 * np.pi**2))**(1.0 / 3) elif period is None: period = (2 * np.pi * a**(3 / 2) / (tt.sqrt(G_grav * (m_star + m_planet)))) return a, period, rho_star * gcc_per_sun, r_star, m_star, m_planet
def __init__(self, period=None, a=None, t0=None, t_periastron=None, incl=None, b=None, duration=None, ecc=None, omega=None, sin_omega=None, cos_omega=None, Omega=None, m_planet=0.0, m_star=None, r_star=None, rho_star=None, ror=None, model=None, **kwargs): add_citations_to_model(self.__citations__, model=model) if "m_planet_units" in kwargs: deprecation_warning( "'m_planet_units' is deprecated; Use `with_unit` instead") m_planet = with_unit(m_planet, kwargs.pop("m_planet_units")) if "rho_star_units" in kwargs: deprecation_warning( "'rho_star_units' is deprecated; Use `with_unit` instead") rho_star = with_unit(rho_star, kwargs.pop("rho_star_units")) self.jacobians = defaultdict(lambda: defaultdict(None)) daordtau = None if ecc is None and duration is not None: if r_star is None: r_star = as_tensor_variable(1.0) if b is None: raise ValueError( "'b' must be provided for a circular orbit with a " "'duration'") if ror is None: warnings.warn( "When using the 'duration' parameter in KeplerianOrbit, " "the 'ror' parameter should also be provided.", UserWarning, ) aor, daordtau = get_aor_from_transit_duration(duration, period, b, ror=ror) a = r_star * aor duration = None inputs = _get_consistent_inputs(a, period, rho_star, r_star, m_star, m_planet) ( self.a, self.period, self.rho_star, self.r_star, self.m_star, self.m_planet, ) = inputs self.m_total = self.m_star + self.m_planet self.n = 2 * np.pi / self.period self.a_star = self.a * self.m_planet / self.m_total self.a_planet = -self.a * self.m_star / self.m_total # Track the Jacobian between the duration and a if daordtau is not None: dadtau = self.r_star * daordtau self.jacobians["duration"]["a"] = dadtau self.jacobians["duration"]["a_star"] = (dadtau * self.m_planet / self.m_total) self.jacobians["duration"]["a_planet"] = (-dadtau * self.m_star / self.m_total) # rho = 3 * pi * (a/R)**3 / (G * P**2) # -> drho / d(a/R) = 9 * pi * (a/R)**2 / (G * P**2) self.jacobians["duration"]["rho_star"] = ( 9 * np.pi * (self.a / self.r_star)**2 * daordtau * gcc_per_sun / (G_grav * self.period**2)) self.K0 = self.n * self.a / self.m_total if Omega is None: self.Omega = None else: self.Omega = as_tensor_variable(Omega) self.cos_Omega = tt.cos(self.Omega) self.sin_Omega = tt.sin(self.Omega) # Eccentricity if ecc is None: self.ecc = None self.M0 = 0.5 * np.pi + tt.zeros_like(self.n) incl_factor = 1 else: self.ecc = as_tensor_variable(ecc) if omega is not None: if sin_omega is not None and cos_omega is not None: raise ValueError( "either 'omega' or 'sin_omega' and 'cos_omega' can be " "provided") self.omega = as_tensor_variable(omega) self.cos_omega = tt.cos(self.omega) self.sin_omega = tt.sin(self.omega) elif sin_omega is not None and cos_omega is not None: self.cos_omega = as_tensor_variable(cos_omega) self.sin_omega = as_tensor_variable(sin_omega) self.omega = tt.arctan2(self.sin_omega, self.cos_omega) else: raise ValueError("both e and omega must be provided") opsw = 1 + self.sin_omega E0 = 2 * tt.arctan2( tt.sqrt(1 - self.ecc) * self.cos_omega, tt.sqrt(1 + self.ecc) * opsw, ) self.M0 = E0 - self.ecc * tt.sin(E0) ome2 = 1 - self.ecc**2 self.K0 /= tt.sqrt(ome2) incl_factor = (1 + self.ecc * self.sin_omega) / ome2 # The Jacobian for the transform cos(i) -> b self.dcosidb = self.jacobians["b"]["cos_incl"] = (incl_factor * self.r_star / self.a) if b is not None: if incl is not None or duration is not None: raise ValueError( "only one of 'incl', 'b', and 'duration' can be given") self.b = as_tensor_variable(b) self.cos_incl = self.dcosidb * self.b self.incl = tt.arccos(self.cos_incl) elif incl is not None: if duration is not None: raise ValueError( "only one of 'incl', 'b', and 'duration' can be given") self.incl = as_tensor_variable(incl) self.cos_incl = tt.cos(self.incl) self.b = self.cos_incl / self.dcosidb elif duration is not None: # This assertion should never be hit because of the first # conditional in this method, but let's keep it here anyways assert self.ecc is not None self.duration = as_tensor_variable(to_unit(duration, u.day)) c = tt.sin(np.pi * self.duration * incl_factor / self.period) c2 = c * c aor = self.a_planet / self.r_star esinw = self.ecc * self.sin_omega self.b = tt.sqrt( (aor**2 * c2 - 1) / (c2 * esinw**2 + 2 * c2 * esinw + c2 - self.ecc**4 + 2 * self.ecc**2 - 1)) self.b *= 1 - self.ecc**2 self.cos_incl = self.dcosidb * self.b self.incl = tt.arccos(self.cos_incl) else: zla = tt.zeros_like(self.a) self.incl = 0.5 * np.pi + zla self.cos_incl = zla self.b = zla if t0 is not None and t_periastron is not None: raise ValueError("you can't define both t0 and t_periastron") if t0 is None and t_periastron is None: t0 = tt.zeros_like(self.period) if t0 is None: self.t_periastron = as_tensor_variable(t_periastron) self.t0 = self.t_periastron + self.M0 / self.n else: self.t0 = as_tensor_variable(t0) self.t_periastron = self.t0 - self.M0 / self.n self.tref = self.t_periastron - self.t0 self.sin_incl = tt.sin(self.incl)
def in_transit(self, t, r=0.0, texp=None, light_delay=False): """Get a list of timestamps that are in transit Args: t (vector): A vector of timestamps to be evaluated. r (Optional): The radii of the planets. texp (Optional[float]): The exposure time. Returns: The indices of the timestamps that are in transit. """ if light_delay: raise NotImplementedError( "Light travel time delay not yet implemented for `in_transit`") z = tt.zeros_like(self.a) r = as_tensor_variable(r) + z R = self.r_star + z # Wrap the times into time since transit hp = 0.5 * self.period dt = tt.mod(self._warp_times(t) + hp, self.period) - hp if self.ecc is None: # Equation 14 from Winn (2010) k = r / R arg = tt.square(1 + k) - tt.square(self.b) factor = R / (self.a * self.sin_incl) hdur = hp * tt.arcsin(factor * tt.sqrt(arg)) / np.pi t_start = -hdur t_end = hdur flag = z else: M_contact = ops.contact_points( self.a, self.ecc, self.cos_omega, self.sin_omega, self.cos_incl + z, self.sin_incl + z, R + r, ) flag = M_contact[2] t_start = (M_contact[0] - self.M0) / self.n t_start = tt.mod(t_start + hp, self.period) - hp t_end = (M_contact[1] - self.M0) / self.n t_end = tt.mod(t_end + hp, self.period) - hp t_start = tt.switch(tt.gt(t_start, 0.0), t_start - self.period, t_start) t_end = tt.switch(tt.lt(t_end, 0.0), t_end + self.period, t_end) if texp is not None: t_start -= 0.5 * texp t_end += 0.5 * texp mask = tt.any(tt.and_(dt >= t_start, dt <= t_end), axis=-1) result = ifelse( tt.all(tt.eq(flag, 0)), tt.arange(t.shape[0])[mask], tt.arange(t.shape[0]), ) return result
def _get_true_anomaly(self, t, _pad=True): M = (self._warp_times(t, _pad=_pad) - self.tref) * self.n if self.ecc is None: return tt.sin(M), tt.cos(M) sinf, cosf = ops.kepler(M, self.ecc + tt.zeros_like(M)) return sinf, cosf
def get_light_curve( self, orbit=None, r=None, t=None, texp=None, oversample=7, order=0, use_in_transit=None, light_delay=False, ): """Get the light curve for an orbit at a set of times Args: orbit: An object with a ``get_relative_position`` method that takes a tensor of times and returns a list of Cartesian coordinates of a set of bodies relative to the central source. This method should return three tensors (one for each coordinate dimension) and each tensor should have the shape ``append(t.shape, r.shape)`` or ``append(t.shape, oversample, r.shape)`` when ``texp`` is given. The first two coordinate dimensions are treated as being in the plane of the sky and the third coordinate is the line of sight with positive values pointing *away* from the observer. For an example, take a look at :class:`orbits.KeplerianOrbit`. r (tensor): The radius of the transiting body in the same units as ``r_star``. This should have a shape that is consistent with the coordinates returned by ``orbit``. In general, this means that it should probably be a scalar or a vector with one entry for each body in ``orbit``. Note that this is a different quantity than the planet-to-star radius ratio; do not confuse the two! t (tensor): The times where the light curve should be evaluated. texp (Optional[tensor]): The exposure time of each observation. This can be a scalar or a tensor with the same shape as ``t``. If ``texp`` is provided, ``t`` is assumed to indicate the timestamp at the *middle* of an exposure of length ``texp``. oversample (Optional[int]): The number of function evaluations to use when numerically integrating the exposure time. order (Optional[int]): The order of the numerical integration scheme. This must be one of the following: ``0`` for a centered Riemann sum (equivalent to the "resampling" procedure suggested by Kipping 2010), ``1`` for the trapezoid rule, or ``2`` for Simpson's rule. use_in_transit (Optional[bool]): If ``True``, the model will only be evaluated for the data points expected to be in transit as computed using the ``in_transit`` method on ``orbit``. """ if orbit is None: raise ValueError("missing required argument 'orbit'") if r is None: raise ValueError("missing required argument 'r'") if t is None: raise ValueError("missing required argument 't'") use_in_transit = (not light_delay if use_in_transit is None else use_in_transit) r = as_tensor_variable(r) r = tt.reshape(r, (r.size, )) t = as_tensor_variable(t) # If use_in_transit, we should only evaluate the model at times where # at least one planet is transiting if use_in_transit: transit_model = tt.shape_padleft(tt.zeros_like(r), t.ndim) + tt.shape_padright( tt.zeros_like(t), r.ndim) inds = orbit.in_transit(t, r=r, texp=texp, light_delay=light_delay) t = t[inds] # Handle exposure time integration if texp is None: tgrid = t rgrid = tt.shape_padleft(r, tgrid.ndim) + tt.shape_padright( tt.zeros_like(tgrid), r.ndim) else: texp = as_tensor_variable(texp) oversample = int(oversample) oversample += 1 - oversample % 2 stencil = np.ones(oversample) # Construct the exposure time integration stencil if order == 0: dt = np.linspace(-0.5, 0.5, 2 * oversample + 1)[1:-1:2] elif order == 1: dt = np.linspace(-0.5, 0.5, oversample) stencil[1:-1] = 2 elif order == 2: dt = np.linspace(-0.5, 0.5, oversample) stencil[1:-1:2] = 4 stencil[2:-1:2] = 2 else: raise ValueError("order must be <= 2") stencil /= np.sum(stencil) if texp.ndim == 0: dt = texp * dt else: if use_in_transit: dt = tt.shape_padright(texp[inds]) * dt else: dt = tt.shape_padright(texp) * dt tgrid = tt.shape_padright(t) + dt # Madness to get the shapes to work out... rgrid = tt.shape_padleft(r, tgrid.ndim) + tt.shape_padright( tt.zeros_like(tgrid), 1) # Evalute the coordinates of the transiting body in the plane of the # sky coords = orbit.get_relative_position(tgrid, light_delay=light_delay) b = tt.sqrt(coords[0]**2 + coords[1]**2) b = tt.reshape(b, rgrid.shape) los = tt.reshape(coords[2], rgrid.shape) lc = self._compute_light_curve(b / orbit.r_star, rgrid / orbit.r_star, los / orbit.r_star) if texp is not None: stencil = tt.shape_padright(tt.shape_padleft(stencil, t.ndim), 1) lc = tt.sum(stencil * lc, axis=t.ndim) if use_in_transit: transit_model = tt.set_subtensor(transit_model[inds], lc) return transit_model else: return lc
def logp(self, value): return tt.zeros_like(tt.as_tensor_variable(value))
def get_star_position(self, t, light_delay=False): nothing = tt.zeros_like(as_tensor_variable(t)) return nothing, nothing, nothing
def _zeros_like(self, tensor): return tt.zeros_like(tensor)
def duration_to_eccentricity(func, duration, ror, **kwargs): # pragma: no cover num_planets = kwargs.pop("num_planets", 1) orbit_type = kwargs.pop("orbit_type", KeplerianOrbit) name = kwargs.get("name", "dur_ecc") inputs = _get_consistent_inputs( kwargs.get("a", None), kwargs.get("period", None), kwargs.get("rho_star", None), kwargs.get("r_star", None), kwargs.get("m_star", None), kwargs.get("rho_star_units", None), kwargs.get("m_planet", 0.0), kwargs.get("m_planet_units", None), ) a, period, rho_star, r_star, m_star, m_planet = inputs b = kwargs.get("b", 0.0) s = tt.sin(kwargs["omega"]) umax_inv = tt.switch(tt.lt(s, 0), tt.sqrt(1 - s**2), 1.0) const = (period * tt.shape_padright(r_star) * tt.sqrt((1 + ror)**2 - b**2)) const /= np.pi * a u = duration / const e1 = -s * u**2 / ((s * u)**2 + 1) e2 = tt.sqrt((s**2 - 1) * u**2 + 1) / ((s * u)**2 + 1) models = [] logjacs = [] logprobs = [] for args in product(*(zip("np", (-1, 1)) for _ in range(num_planets))): labels, signs = zip(*args) # Compute the eccentricity branch ecc = tt.stack([e1[i] + signs[i] * e2[i] for i in range(num_planets)]) # Work out the Jacobian valid_ecc = tt.and_(tt.lt(ecc, 1.0), tt.ge(ecc, 0.0)) logjac = tt.switch( tt.all(valid_ecc), tt.sum(0.5 * tt.log(1 - ecc**2) + 2 * tt.log(s * ecc + 1) - tt.log(tt.abs_(s + ecc)) - tt.log(const)), -np.inf, ) ecc = tt.switch(valid_ecc, ecc, tt.zeros_like(ecc)) # Create a sub-model to capture this component with pm.Model(name="dur_ecc_" + "_".join(labels)) as model: pm.Deterministic("ecc", ecc) orbit = orbit_type(ecc=ecc, **kwargs) logprob = tt.sum(func(orbit)) models.append(model) logjacs.append(logjac) logprobs.append(logprob) # Compute the marginalized likelihood logjacs = tt.stack(logjacs) logprobs = tt.stack(logprobs) logprob = tt.switch( tt.gt(1.0 / u, umax_inv), tt.sum(pm.logsumexp(logprobs + logjacs)), -np.inf, ) pm.Potential(name + "_logp", logprob) pm.Deterministic(name + "_logjacs", logjacs) pm.Deterministic(name + "_logprobs", logprobs) # Loop over models and compute the marginalized values for all the # parameters and deterministics norm = tt.sum(pm.logsumexp(logjacs)) logw = tt.switch( tt.gt(1.0 / u, umax_inv), logjacs - norm, -np.inf + tt.zeros_like(logjacs), ) pm.Deterministic(name + "_logw", logw) for k in models[0].named_vars.keys(): name = k[len(models[0].name) + 1:] pm.Deterministic( name, sum( tt.exp(logw[i]) * model.named_vars[model.name + "_" + name] for i, model in enumerate(models)), )
def get_coefficients(self): c1 = self.term1.coefficients c2 = self.term2.coefficients # First compute real terms ar = [] cr = [] ar.append(tt.flatten(c1[0][:, None] * c2[0][None, :])) cr.append(tt.flatten(c1[1][:, None] + c2[1][None, :])) # Then the complex terms ac = [] bc = [] cc = [] dc = [] # real * complex ac.append(tt.flatten(c1[0][:, None] * c2[2][None, :])) bc.append(tt.flatten(c1[0][:, None] * c2[3][None, :])) cc.append(tt.flatten(c1[1][:, None] + c2[4][None, :])) dc.append(tt.flatten(tt.zeros_like(c1[1])[:, None] + c2[5][None, :])) ac.append(tt.flatten(c2[0][:, None] * c1[2][None, :])) bc.append(tt.flatten(c2[0][:, None] * c1[3][None, :])) cc.append(tt.flatten(c2[1][:, None] + c1[4][None, :])) dc.append(tt.flatten(tt.zeros_like(c2[1])[:, None] + c1[5][None, :])) # complex * complex aj, bj, cj, dj = c1[2:] ak, bk, ck, dk = c2[2:] ac.append( tt.flatten( 0.5 * (aj[:, None] * ak[None, :] + bj[:, None] * bk[None, :]) ) ) bc.append( tt.flatten( 0.5 * (bj[:, None] * ak[None, :] - aj[:, None] * bk[None, :]) ) ) cc.append(tt.flatten(cj[:, None] + ck[None, :])) dc.append(tt.flatten(dj[:, None] - dk[None, :])) ac.append( tt.flatten( 0.5 * (aj[:, None] * ak[None, :] - bj[:, None] * bk[None, :]) ) ) bc.append( tt.flatten( 0.5 * (bj[:, None] * ak[None, :] + aj[:, None] * bk[None, :]) ) ) cc.append(tt.flatten(cj[:, None] + ck[None, :])) dc.append(tt.flatten(dj[:, None] + dk[None, :])) return [ tt.concatenate(vals, axis=0) if len(vals) else tt.zeros(0, dtype=self.dtype) for vals in (ar, cr, ac, bc, cc, dc) ]