def underdamped(self): Q = self.Q f = tt.sqrt(tt.maximum(4.0 * Q ** 2 - 1.0, self.eps)) a = self.S0 * self.w0 * Q c = 0.5 * self.w0 / Q empty = tt.zeros(0, dtype=self.dtype) return ( empty, empty, tt.stack([a]), tt.stack([a / f]), tt.stack([c]), tt.stack([c * f]), )
def overdamped(self): Q = self.Q f = tt.sqrt(tt.maximum(1.0 - 4.0 * Q ** 2, self.eps)) empty = tt.zeros(0, dtype=self.dtype) return ( 0.5 * self.S0 * self.w0 * Q * tt.stack([1.0 + 1.0 / f, 1.0 - 1.0 / f]), 0.5 * self.w0 / Q * tt.stack([1.0 - f, 1.0 + f]), empty, empty, empty, empty, )
def get_cl(u1, u2): u1 = as_tensor_variable(u1) u2 = as_tensor_variable(u2) c0 = 1 - u1 - 1.5 * u2 c1 = u1 + 2 * u2 c2 = -0.25 * u2 norm = np.pi * (c0 + c1 / 1.5) return tt.stack([c0, c1, c2]) / norm
def forward(self, x): p = x[0] b = x[1] pl = self.min_radius r11 = (b / (1 + pl) - 1) * (1 - self.Ar) + 1 r21 = (p - pl) / self.dr arg = p - b - self.dr + 1 q1 = (arg / self.dr)**2 q2 = (pl - b + 1) / arg r12 = q1 * self.Ar r22 = q2 y = tt.switch(b <= 1, tt.stack((r11, r21), axis=0), tt.stack((r12, r22), axis=0)) return tt.log(y) - tt.log(1 - y)
def _get_position_and_velocity(self, t, parallax=None): sinf, cosf = self._get_true_anomaly(t) if self.ecc is None: r = 1.0 vx, vy, vz = self._rotate_vector(-self.K0 * sinf, self.K0 * cosf) else: r = (1.0 - self.ecc**2) / (1 + self.ecc * cosf) vx, vy, vz = self._rotate_vector(-self.K0 * sinf, self.K0 * (cosf + self.ecc)) x, y, z = self._rotate_vector(r * cosf, r * sinf) pos = tt.stack((x, y, z), axis=-1) pos = tt.concatenate( ( tt.sum(tt.shape_padright(self.a_star) * pos, axis=0, keepdims=True), tt.shape_padright(self.a_planet) * pos, ), axis=0, ) vel = tt.stack((vx, vy, vz), axis=-1) vel = tt.concatenate( ( tt.sum( tt.shape_padright(self.m_planet) * vel, axis=0, keepdims=True, ), -tt.shape_padright(self.m_star) * vel, ), axis=0, ) if parallax is not None: # convert r into arcseconds pos = pos * (parallax * au_per_R_sun) vel = vel * (parallax * au_per_R_sun) return pos, vel
def backward(self, y): y = tt.nnet.sigmoid(y) r1 = y[0] r2 = y[1] pl, pu = self.min_radius, self.max_radius b1 = (1 + pl) * (1 + (r1 - 1) / (1 - self.Ar)) p1 = pl + r2 * self.dr q1 = r1 / self.Ar q2 = r2 b2 = (1 + pl) + tt.sqrt(q1) * q2 * self.dr p2 = pu - self.dr * tt.sqrt(q1) * (1 - q2) pb = tt.switch( r1 > self.Ar, tt.stack((p1, b1), axis=0), tt.stack((p2, b2), axis=0), ) return pb
def get_light_curve( self, orbit=None, r=None, t=None, texp=None, oversample=7, order=0, use_in_transit=None, light_delay=False, ): if self.num_planets is None: try: vec = orbit.period.tag.test_value except AttributeError: raise ValueError( "Can't compute num_planets, please provide a value") num_planets = len(np.atleast_1d(vec)) else: num_planets = int(self.num_planets) if num_planets <= 1: func = _wrapper( self.base_light_curve, orbit=orbit, r=r, texp=texp, oversample=oversample, order=order, use_in_transit=use_in_transit, light_delay=light_delay, ) mn = orbit.t0 mx = orbit.t0 + orbit.period return interp( 0, tt.mod(t - orbit.t0, orbit.period) + orbit.t0, mn, mx, (mx - mn) / (self.num_phase + 1), func, )[:, None] ys = [] for n in range(num_planets): func = _wrapper( self.base_light_curve, orbit=orbit, r=r, texp=texp, oversample=oversample, order=order, use_in_transit=use_in_transit, light_delay=light_delay, ) mn = orbit.t0[n] mx = orbit.t0[n] + orbit.period[n] ys.append( interp( n, tt.mod(t - orbit.t0[n], orbit.period[n]) + orbit.t0[n], mn, mx, (mx - mn) / (self.num_phase + 1), func, )) return tt.stack(ys, axis=-1)
def __init__(self, *args, **kwargs): ttvs = kwargs.pop("ttvs", None) transit_times = kwargs.pop("transit_times", None) transit_inds = kwargs.pop("transit_inds", None) if ttvs is None and transit_times is None: raise ValueError("one of 'ttvs' or 'transit_times' must be " "defined") if ttvs is not None: self.ttvs = [as_tensor_variable(ttv, ndim=1) for ttv in ttvs] if transit_inds is None: self.transit_inds = [ tt.arange(ttv.shape[0]) for ttv in self.ttvs ] else: self.transit_inds = [ tt.cast(as_tensor_variable(inds, ndim=1), "int64") for inds in transit_inds ] else: # If transit times are given, compute the least squares period and # t0 based on these times. self.transit_times = [] self.ttvs = [] self.transit_inds = [] period = [] t0 = [] for i, times in enumerate(transit_times): times = as_tensor_variable(times, ndim=1) if transit_inds is None: inds = tt.arange(times.shape[0]) else: inds = tt.cast(as_tensor_variable(transit_inds[i]), "int64") self.transit_inds.append(inds) # A convoluted version of linear regression; don't ask N = times.shape[0] sumx = tt.sum(inds) sumx2 = tt.sum(inds**2) sumy = tt.sum(times) sumxy = tt.sum(inds * times) denom = N * sumx2 - sumx**2 slope = (N * sumxy - sumx * sumy) / denom intercept = (sumx2 * sumy - sumx * sumxy) / denom expect = intercept + inds * slope period.append(slope) t0.append(intercept) self.ttvs.append(times - expect) self.transit_times.append(times) kwargs["t0"] = tt.stack(t0) # We'll have two different periods: one that is the mean difference # between transit times and one that is a parameter that sets the # transit shape. If a "period" parameter is not given, these will # be the same. Users will probably want to put a prior relating the # two periods if they use separate values. self.ttv_period = tt.stack(period) if "period" not in kwargs: if "delta_log_period" in kwargs: kwargs["period"] = tt.exp( tt.log(self.ttv_period) + kwargs.pop("delta_log_period")) else: kwargs["period"] = self.ttv_period super(TTVOrbit, self).__init__(*args, **kwargs) if ttvs is not None: self.ttv_period = self.period self.transit_times = [ self.t0[i] + self.period[i] * self.transit_inds[i] + ttv for i, ttv in enumerate(self.ttvs) ] # Compute the full set of transit times self.all_transit_times = [] for i, inds in enumerate(self.transit_inds): expect = self.t0[i] + self.period[i] * tt.arange(inds.max() + 1) self.all_transit_times.append( tt.set_subtensor(expect[inds], self.transit_times[i])) # Set up a histogram for identifying the transit offsets self._bin_edges = [ tt.concatenate(( [tts[0] - 0.5 * self.ttv_period[i]], 0.5 * (tts[1:] + tts[:-1]), [tts[-1] + 0.5 * self.ttv_period[i]], )) for i, tts in enumerate(self.all_transit_times) ] self._bin_values = [ tt.concatenate(([tts[0]], tts, [tts[-1]])) for i, tts in enumerate(self.all_transit_times) ]
def _get_model_dt(self, t): vals = [] for i in range(len(self.ttvs)): inds = tt.extra_ops.searchsorted(self._bin_edges[i], t) vals.append(self._bin_values[i][inds]) return tt.stack(vals, -1)
def forward(self, x): usum = tt.sum(x, axis=0) q = tt.stack([usum**2, 0.5 * x[0] / usum]) return tt.log(q) - tt.log(1 - q)
def backward(self, y): q = tt.nnet.sigmoid(y) sqrtq1 = tt.sqrt(q[0]) twoq2 = 2 * q[1] u = tt.stack([sqrtq1 * twoq2, sqrtq1 * (1 - twoq2)]) return u
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)), )