import numpy as np
import types
from collections import OrderedDict
from radvel import kepler
from radvel.basis import Basis
texdict = {
'per': 'P',
'logper': '\\ln{P}',
'logk': '\\ln{K}',
'tc': 'T\\rm{conj}',
'secosw': '\\sqrt{e}\\cos{\\omega}',
'sesinw': '\\sqrt{e}\\sin{\\omega}',
'ecosw': 'e\\cos{\\omega}',
'esinw': 'e\\sin{\\omega}',
'e': 'e',
'w': '\\omega',
'tp': 'T\\rm{peri}',
'k': 'K',
'gamma_': '\\gamma_{\\rm ',
'logjit_': '\\ln{\\sigma_{\\rm jit}}_{\\rm ',
'jit_': '\\sigma_{\\rm ',
'dvdt': '\\dot{\\gamma}',
'curv': '\\ddot{\\gamma}'
}
[docs]class Parameters(OrderedDict):
"""Object to store the orbital parameters.
Parameters to describe a radial velocity orbit
stored as an OrderedDict
Args:
num_planets (int): Number of planets in model
basis (string): parameterization of orbital parameters. See
``radvel.basis.Basis`` for a list of valid basis strings.
planet_letters (Dictionary[optional): custom map to match the planet
numbers in the Parameter object to planet letters.
Default {1: 'b', 2: 'c', etc.}. The keys of this dictionary must
all be integers.
Attributes:
basis (radvel.Basis): Basis object
planet_parameters (list): orbital parameters contained within the
specified basis
num_planets (int): number of planets in the model
Examples:
>>> import radvel
# create a Parameters object for a 2-planet system with
# custom planet number to letter mapping
>>> params = radvel.Parameters(2, planet_letters={1:'d', 2:'e'})
"""
def __init__(self, num_planets, basis='per tc secosw sesinw logk',
planet_letters=None):
super(Parameters, self).__init__()
basis = Basis(basis,num_planets)
self.planet_parameters = basis.name.split()
for num_planet in range(1,1+num_planets):
for parameter in self.planet_parameters:
new_name = self._sparameter(parameter, num_planet)
self.__setitem__(new_name, Parameter())
if planet_letters is not None:
for k in planet_letters.keys():
assert isinstance(k, int), """\
Parameters: ERROR: The planet_letters dictionary \
should have only integers as keys."""
self.basis = basis
self.num_planets = num_planets
self.planet_letters = planet_letters
def __reduce__(self):
red = (self.__class__, (self.num_planets,
self.basis.name,
self.planet_letters),
None,None,iter(self.items()))
return red
[docs] def tex_labels(self, param_list=None):
"""Map Parameters keys to pretty TeX code representations.
Args:
param_list (list): (optional) Manually pass a list of parameter labels
Returns:
dict: dictionary mapping Parameters keys to TeX code
"""
if param_list is None:
param_list = self.keys()
tex_labels = {}
for k in param_list:
n = k[-1]
p = k[:-1]
if n.isdigit(): tex_labels[k] = self._planet_texlabel(p, n)
elif k in texdict.keys(): tex_labels[k] = "$%s$" % texdict[k]
elif p not in self.planet_parameters:
for tex in texdict.keys():
if tex in k and len(tex) > 1:
tex_labels[k] = "$%s}$" % k.replace(tex, texdict[tex])
if k not in tex_labels.keys():
tex_labels[k] = k
return tex_labels
def _sparameter(self, parameter, num_planet):
return '{0}{1}'.format(parameter, num_planet)
def _planet_texlabel(self, parameter, num_planet):
pname = texdict.get(parameter, parameter)
if self.planet_letters is not None:
lett_planet = self.planet_letters[int(num_planet)]
else:
lett_planet = chr(int(num_planet)+97)
return '$%s_{%s}$' % (pname, lett_planet)
[docs]class Parameter(object):
"""Object to store attributes of each orbital parameter
Attributes:
value (float): value of parameter.
vary (Bool): True if parameter is allowed to vary in
MCMC fits, false if fixed.
mcmcscale (float): step size to be used for MCMC fitting
"""
def __init__(self, value=None, vary=True, mcmcscale=None):
self.value = value
self.vary = vary
self.mcmcscale = mcmcscale
def _equals(self, other):
"""method to assess the equivalence of two Parameter objects"""
if isinstance(other,self.__class__):
return (self.value == other.value) and (self.vary == other.vary) \
and (self.mcmcscale == other.mcmcscale)
def __repr__(self):
s = "Parameter object: value = {}, vary = {}, mcmc scale = {}".format(self.value,
self.vary, self.mcmcscale)
return s
if __name__ == "__main__":
a = Parameter(value=1.3)
a.mcmcscale = 100.
print(a)
[docs]class RVModel(object):
"""
Generic RV Model
This class defines the methods common to all RV modeling
classes. The different RV models, having different
parameterizations inherit from this class.
"""
def __init__(self, params, time_base=0):
self.num_planets = params.num_planets
self.params = params
self.time_base = time_base
if 'dvdt' not in params.keys():
self.params['dvdt']=Parameter(value=0.)
if 'curv' not in params.keys():
self.params['curv']=Parameter(value=0.)
def __call__(self, t, planet_num=None):
"""Compute the radial velocity
Includes all Keplerians and additional trends.
Args:
t (array of floats): Timestamps to calculate the RV model
planet_num (Optional[int]): calculate the RV model for a single
planet within a multi-planet system
Returns:
vel (array of floats): Radial velocity at each time in `t`
"""
vel = np.zeros( len(t) )
params_synth = self.params.basis.to_synth(self.params)
if planet_num == None:
planets = range(1, self.num_planets+1)
else:
planets = [planet_num]
for num_planet in planets:
per = params_synth['per{}'.format(num_planet)].value
tp = params_synth['tp{}'.format(num_planet)].value
e = params_synth['e{}'.format(num_planet)].value
w = params_synth['w{}'.format(num_planet)].value
k = params_synth['k{}'.format(num_planet)].value
orbel_synth = np.array([per, tp, e, w, k])
vel+=kepler.rv_drive(t, orbel_synth)
vel+=self.params['dvdt'].value * ( t - self.time_base )
vel+=self.params['curv'].value * ( t - self.time_base )**2
return vel