def get_value(self, tau0): dt = self.delta ar, cr, a, b, c, d = self.term.coefficients # Format the lags correctly tau0 = tt.abs_(tau0) tau = tau0[..., None] # Precompute some factors dpt = dt + tau dmt = dt - tau # Real parts: # tau > Delta crd = cr * dt cosh = tt.cosh(crd) norm = 2 * ar / crd ** 2 K_large = tt.sum(norm * (cosh - 1) * tt.exp(-cr * tau), axis=-1) # tau < Delta crdmt = cr * dmt K_small = K_large + tt.sum(norm * (crdmt - tt.sinh(crdmt)), axis=-1) # Complex part cd = c * dt dd = d * dt c2 = c ** 2 d2 = d ** 2 c2pd2 = c2 + d2 C1 = a * (c2 - d2) + 2 * b * c * d C2 = b * (c2 - d2) - 2 * a * c * d norm = 1.0 / (dt * c2pd2) ** 2 k0 = tt.exp(-c * tau) cdt = tt.cos(d * tau) sdt = tt.sin(d * tau) # For tau > Delta cos_term = 2 * (tt.cosh(cd) * tt.cos(dd) - 1) sin_term = 2 * (tt.sinh(cd) * tt.sin(dd)) factor = k0 * norm K_large += tt.sum( (C1 * cos_term - C2 * sin_term) * factor * cdt, axis=-1 ) K_large += tt.sum( (C2 * cos_term + C1 * sin_term) * factor * sdt, axis=-1 ) # tau < Delta edmt = tt.exp(-c * dmt) edpt = tt.exp(-c * dpt) cos_term = ( edmt * tt.cos(d * dmt) + edpt * tt.cos(d * dpt) - 2 * k0 * cdt ) sin_term = ( edmt * tt.sin(d * dmt) + edpt * tt.sin(d * dpt) - 2 * k0 * sdt ) K_small += tt.sum(2 * (a * c + b * d) * c2pd2 * dmt * norm, axis=-1) K_small += tt.sum((C1 * cos_term + C2 * sin_term) * norm, axis=-1) return tt.switch(tt.le(tau0, dt), K_small, K_large)
def get_aor_from_transit_duration(duration, period, b, ror=None): """Get the semimajor axis implied by a circular orbit and duration Args: duration: The transit duration period: The orbital period b: The impact parameter of the transit ror: The radius ratio of the planet to the star Returns: The semimajor axis in units of the stellar radius and the Jacobian ``d a / d duration`` """ if ror is None: ror = as_tensor_variable(0.0) b2 = b**2 opk2 = (1 + ror)**2 phi = np.pi * duration / period sinp = tt.sin(phi) cosp = tt.cos(phi) num = tt.sqrt(opk2 - b2 * cosp**2) aor = num / sinp grad = np.pi * cosp * (b2 - opk2) / (num * period * sinp**2) return aor, grad
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 get_coefficients(self): ar, cr, a, b, c, d = self.term.coefficients # Real componenets crd = cr * self.delta coeffs = [2 * ar * (tt.cosh(crd) - 1) / crd ** 2, cr] # Imaginary coefficients cd = c * self.delta dd = d * self.delta c2 = c ** 2 d2 = d ** 2 factor = 2.0 / (self.delta * (c2 + d2)) ** 2 cos_term = tt.cosh(cd) * tt.cos(dd) - 1 sin_term = tt.sinh(cd) * tt.sin(dd) C1 = a * (c2 - d2) + 2 * b * c * d C2 = b * (c2 - d2) - 2 * a * c * d coeffs += [ factor * (C1 * cos_term - C2 * sin_term), factor * (C2 * cos_term + C1 * sin_term), c, d, ] return coeffs
def get_celerite_matrices(self, x, diag, **kwargs): dt = self.delta ar, cr, a, b, c, d = self.term.coefficients # Real part cd = cr * dt delta_diag = 2 * tt.sum(ar * (cd - tt.sinh(cd)) / cd ** 2) # Complex part cd = c * dt dd = d * dt c2 = c ** 2 d2 = d ** 2 c2pd2 = c2 + d2 C1 = a * (c2 - d2) + 2 * b * c * d C2 = b * (c2 - d2) - 2 * a * c * d norm = (dt * c2pd2) ** 2 sinh = tt.sinh(cd) cosh = tt.cosh(cd) delta_diag += 2 * tt.sum( ( C2 * cosh * tt.sin(dd) - C1 * sinh * tt.cos(dd) + (a * c + b * d) * dt * c2pd2 ) / norm ) new_diag = diag + delta_diag return super().get_celerite_matrices(x, new_diag)
def get_value(self, tau): ar, cr, ac, bc, cc, dc = self.coefficients tau = tt.abs_(tau) tau = tt.shape_padright(tau) K = tt.sum(ar * tt.exp(-cr * tau), axis=-1) factor = tt.exp(-cc * tau) K += tt.sum(ac * factor * tt.cos(dc * tau), axis=-1) K += tt.sum(bc * factor * tt.sin(dc * tau), axis=-1) return K
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 _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 compute_gradient(self, t): t = tt.as_tensor_variable(t).astype("float64") dcos = tt.cos(self.omega * t) dsin = tt.sin(self.omega * t) df = 2 * np.pi * t * (self.sin * dcos - self.cos * dsin) return (df, dcos, dsin)
def get_value(self, t): t = tt.as_tensor_variable(t).astype("float64") return ( self.cos * tt.cos(self.omega * t) + self.sin * tt.sin(self.omega * t) )
def compute_gradient(self, t): t = tt.as_tensor_variable(t).astype("float64") damp = tt.sin(self.omega * t + self.phase) dphi = self.amp * tt.cos(self.omega * t + self.phase) df = 2 * np.pi * t * dphi return (df, damp, dphi)
def get_value(self, t): t = tt.as_tensor_variable(t).astype("float64") return self.amp * tt.sin(self.omega * t + self.phase)
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_psd(self, omega): psd0 = self.term.get_psd(omega) arg = 0.5 * self.delta * omega sinc = tt.switch(tt.neq(arg, 0), tt.sin(arg) / arg, tt.ones_like(arg)) return psd0 * sinc ** 2