def arc_distance_theano_alloc_prepare(dtype='float64'):
"""
Calculates the pairwise arc distance between all points in vector a and b.
"""
a = tensor.matrix(dtype=str(dtype))
b = tensor.matrix(dtype=str(dtype))
# Theano don't implement all case of tile, so we do the equivalent with alloc.
#theta_1 = tensor.tile(a[:, 0], (b.shape[0], 1)).T
theta_1 = tensor.alloc(a[:, 0], b.shape[0], b.shape[0]).T
phi_1 = tensor.alloc(a[:, 1], b.shape[0], b.shape[0]).T
theta_2 = tensor.alloc(b[:, 0], a.shape[0], a.shape[0])
phi_2 = tensor.alloc(b[:, 1], a.shape[0], a.shape[0])
temp = (tensor.sin((theta_2 - theta_1) / 2)**2
+
tensor.cos(theta_1) * tensor.cos(theta_2)
* tensor.sin((phi_2 - phi_1) / 2)**2)
distance_matrix = 2 * (tensor.arctan2(tensor.sqrt(temp),
tensor.sqrt(1 - temp)))
name = "arc_distance_theano_alloc"
rval = theano.function([a, b],
distance_matrix,
name=name)
rval.__name__ = name
return rval
def arc_distance_theano_broadcast_prepare(dtype='float64'):
"""
Calculates the pairwise arc distance between all points in vector a and b.
"""
a = tensor.matrix(dtype=str(dtype))
b = tensor.matrix(dtype=str(dtype))
theta_1 = a[:, 0][None, :]
theta_2 = b[:, 0][None, :]
phi_1 = a[:, 1][:, None]
phi_2 = b[:, 1][None, :]
temp = (tensor.sin((theta_2 - theta_1) / 2)**2
+
tensor.cos(theta_1) * tensor.cos(theta_2)
* tensor.sin((phi_2 - phi_1) / 2)**2)
distance_matrix = 2 * (tensor.arctan2(tensor.sqrt(temp),
tensor.sqrt(1 - temp)))
name = "arc_distance_theano_broadcast"
rval = theano.function([a, b],
distance_matrix,
name=name)
rval.__name__ = name
return rval
def get_phase(states):
v, w, r = states
### convert to centered complex coordinates
Vcenter = -0.22
Wcenter = 0.6
x = v - Vcenter
y = w - Wcenter
angle = T.arctan2(y, x)
### take the mean of unit vectors
mag = T.sqrt(x**2 + y**2)
x = x / mag
y = y / mag
mean = T.arctan2(y.mean(-1), x.mean(-1))
### calculate angles around the mean
angle = T.mod(angle - mean[:,None] + np.pi, 2*np.pi) - np.pi
std = T.sqrt((angle**2).mean(-1))
return std
def get_phase(states):
v, w = states
angle = T.switch(w > 0,
np.pi * v.clip(0, 1),
w * (np.pi / T.abs_(T.min(w))))
mean = T.arctan2(T.sin(angle).mean(axis=-1),
T.cos(angle).mean(axis=-1))
### calculate angles around the mean
angle = T.mod(angle + (np.pi - mean[:,None]), 2*np.pi) - np.pi
std = T.sqrt((angle**2).mean(-1))
return std
def hdist(a, b):
lat1 = a[:, 0] * deg2rad
lon1 = a[:, 1] * deg2rad
lat2 = b[:, 0] * deg2rad
lon2 = b[:, 1] * deg2rad
dlat = abs(lat1-lat2)
dlon = abs(lon1-lon2)
al = tensor.sin(dlat/2)**2 + tensor.cos(lat1) * tensor.cos(lat2) * (tensor.sin(dlon/2)**2)
d = tensor.arctan2(tensor.sqrt(al), tensor.sqrt(const(1)-al))
hd = const(2) * rearth * d
return tensor.switch(tensor.eq(hd, float('nan')), (a-b).norm(2, axis=1), hd)
def dist(self, y_pred, y):
"""A helper method that computes distance between two points
on the surface of earth according to their coordinates.
Inputs are tensors.
"""
y_pred_ra = T.deg2rad(y_pred)
y_ra = T.deg2rad(y)
lat1 = y_pred_ra[:, 0]
lat2 = y_ra[:, 0]
dlon = (y_pred_ra - y_ra)[:, 1]
EARTH_R = 6372.8
y = T.sqrt(
(T.cos(lat2) * T.sin(dlon)) ** 2
+ (T.cos(lat1) * T.sin(lat2) - T.sin(lat1) * T.cos(lat2) * T.cos(dlon)) ** 2
)
x = T.sin(lat1) * T.sin(lat2) + T.cos(lat1) * T.cos(lat2) * T.cos(dlon)
c = T.arctan2(y, x)
return EARTH_R * c
def backward(self, y):
return tt.arctan2(tt.sin(y), tt.cos(y))
def clog(re, im):
log_re = 0.5 * T.log(re**2 + im**2)
log_im = T.arctan2(im, re)
return log_re, log_im
def backward(self, y):
return tt.arctan2(y[0], y[1])
def __init__(self,
period=None, a=None, t0=0.0,
incl=None, b=None, duration=None,
ecc=None, omega=None, m_planet=0.0,
m_star=None, r_star=None, rho_star=None,
m_planet_units=None, rho_star_units=None,
model=None,
contact_points_kwargs=None,
**kwargs):
add_citations_to_model(self.__citations__, model=model)
self.gcc_to_sun = (
(constants.M_sun / constants.R_sun**3).to(u.g / u.cm**3).value)
self.G_grav = constants.G.to(u.R_sun**3 / u.M_sun / u.day**2).value
self.kepler_op = KeplerOp(**kwargs)
# Parameters
self.period = tt.as_tensor_variable(period)
self.t0 = tt.as_tensor_variable(t0)
self.m_planet = tt.as_tensor_variable(m_planet)
if m_planet_units is not None:
self.m_planet *= (1 * m_planet_units).to(u.M_sun).value
self.a, self.period, self.rho_star, self.r_star, self.m_star = \
self._get_consistent_inputs(a, period, rho_star, r_star, m_star,
rho_star_units)
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
self.K0 = self.n * self.a / self.m_total
# Set up the contact points calculation
if contact_points_kwargs is None:
contact_points_kwargs = dict()
# Eccentricity
self.contact_points_op = ContactPointsOp(**contact_points_kwargs)
if ecc is None:
self.ecc = None
self.M0 = 0.5 * np.pi + tt.zeros_like(self.n)
self.tref = self.t0 - self.M0 / self.n
incl_factor = 1
else:
self.ecc = tt.as_tensor_variable(ecc)
if omega is None:
raise ValueError("both e and omega must be provided")
self.omega = tt.as_tensor_variable(omega)
self.cos_omega = tt.cos(self.omega)
self.sin_omega = tt.sin(self.omega)
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)
self.tref = self.t0 - self.M0 / self.n
ome2 = 1 - self.ecc**2
self.K0 /= tt.sqrt(ome2)
incl_factor = (1 + self.ecc * self.sin_omega) / ome2
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 = tt.as_tensor_variable(b)
self.cos_incl = incl_factor*self.b*self.r_star/self.a_planet
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 = tt.as_tensor_variable(incl)
self.cos_incl = tt.cos(self.incl)
self.b = self.a_planet*self.cos_incl/(incl_factor*self.r_star)
elif duration is not None:
if self.ecc is None:
raise ValueError("fitting with duration only works for "
"eccentric orbits")
self.duration = tt.as_tensor_variable(duration)
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 = incl_factor*self.b*self.r_star/self.a_planet
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
self.sin_incl = tt.sin(self.incl)
def backward(self, y):
return T.arctan2(T.sin(y), T.cos(y))
def __init__(self,
period=None, a=None, t0=None, t_periastron=None,
incl=None, b=None,
duration=None, ecc=None, omega=None,
Omega=None, m_planet=0.0, m_star=None,
r_star=None, rho_star=None,
m_planet_units=None, rho_star_units=None,
model=None,
contact_points_kwargs=None,
**kwargs):
add_citations_to_model(self.__citations__, model=model)
self.gcc_to_sun = (
(constants.M_sun / constants.R_sun**3).to(u.g / u.cm**3).value)
self.au_to_R_sun = (constants.au / constants.R_sun).value
self.G_grav = constants.G.to(u.R_sun**3 / u.M_sun / u.day**2).value
self.kepler_op = KeplerOp(**kwargs)
# Parameters
self.period = tt.as_tensor_variable(period)
self.m_planet = tt.as_tensor_variable(m_planet)
if m_planet_units is not None:
self.m_planet *= (1 * m_planet_units).to(u.M_sun).value
self.a, self.period, self.rho_star, self.r_star, self.m_star = \
self._get_consistent_inputs(a, period, rho_star, r_star, m_star,
rho_star_units)
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
self.K0 = self.n * self.a / self.m_total
# Set up the contact points calculation
if contact_points_kwargs is None:
contact_points_kwargs = dict()
if Omega is None:
self.Omega = None
else:
self.Omega = tt.as_tensor_variable(Omega)
self.cos_Omega = tt.cos(self.Omega)
self.sin_Omega = tt.sin(self.Omega)
# Eccentricity
self.contact_points_op = ContactPointsOp(**contact_points_kwargs)
if ecc is None:
self.ecc = None
self.M0 = 0.5 * np.pi + tt.zeros_like(self.n)
incl_factor = 1
else:
self.ecc = tt.as_tensor_variable(ecc)
if omega is None:
raise ValueError("both e and omega must be provided")
self.omega = tt.as_tensor_variable(omega)
self.cos_omega = tt.cos(self.omega)
self.sin_omega = tt.sin(self.omega)
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
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 = tt.as_tensor_variable(b)
self.cos_incl = incl_factor*self.b*self.r_star/self.a_planet
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 = tt.as_tensor_variable(incl)
self.cos_incl = tt.cos(self.incl)
self.b = self.a_planet*self.cos_incl/(incl_factor*self.r_star)
elif duration is not None:
if self.ecc is None:
raise ValueError("fitting with duration only works for "
"eccentric orbits")
self.duration = tt.as_tensor_variable(duration)
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 = incl_factor*self.b*self.r_star/self.a_planet
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 = 0.0
if t0 is None:
self.t_periastron = tt.as_tensor_variable(t_periastron)
self.t0 = self.t_periastron + self.M0 / self.n
else:
self.t0 = tt.as_tensor_variable(t0)
self.t_periastron = self.t0 - self.M0 / self.n
self.tref = self.t_periastron
self.sin_incl = tt.sin(self.incl)
