from __future__ import print_function, division
import os, os.path, sys, re, glob
import itertools
from copy import deepcopy
import logging
import json
from .config import on_rtd
if not on_rtd:
import numpy as np
import pandas as pd
import numpy.random as rand
from scipy.stats import gaussian_kde
import scipy
import emcee
import corner
try:
import pymultinest
except ImportError:
logging.warning('PyMultiNest not imported. MultiNest fits will not work.')
import configobj
from astropy.coordinates import SkyCoord
try:
basestring
except NameError:
basestring = str
from .utils import addmags
from .observation import ObservationTree, Observation, Source
from .priors import age_prior, distance_prior, AV_prior, q_prior
from .priors import salpeter_prior, feh_prior
from .isochrone import get_ichrone, Isochrone
def _parse_config_value(v):
try:
val = float(v)
except:
try:
val = [float(x) for x in v]
except:
val = v
#print('{} becomes {}, type={}'.format(v,val,type(val)))
return val
[docs]class StarModel(object):
"""
:param ic:
:class:`Isochrone` object used to model star.
:param obs: (optional)
:class:`ObservationTree` object containing photometry information.
If not provided, then one will be constructed from the provided
keyword arguments (which must include at least one photometric
bandpass). This should only happen in the simplest case
of a single star system---if multiple stars are detected
in any of the observations being used, an :class:`ObservationTree`
should be passed. If `obs` is a string, then it is assumed
to be a filename of an obs summary DataFrame.
:param N:
Number of model stars to assign to each "leaf node" of the
:class:`ObservationTree`. If you want to model a binary star,
provide ``N=2``.
:param **kwargs:
Keyword arguments must be properties of given isochrone, e.g., logg,
feh, Teff, and/or magnitudes. The values represent measurements of
the star, and must be in (value,error) format. All such keyword
arguments will be held in ``self.properties``. ``parallax`` is
also a valid property, and should be provided in miliarcseconds,
as is ``density`` [g/cc], and ``nu_max`` and ``delta_nu``
(asteroseismic parameters in uHz.)
"""
# These are allowable parameters that are not photometric bands
_not_a_band = ('RA','dec','ra','Dec','maxAV','parallax',
'logg','Teff','feh','density', 'separation',
'PA','resolution','relative','N','index', 'id')
_default_name = 'single'
def __init__(self, ic, obs=None, N=1, index=0,
name='', use_emcee=False,
RA=None, dec=None, coords=None,
**kwargs):
self.name = name if name else self._default_name
if coords is None:
if RA is not None and dec is not None:
try:
coords = SkyCoord(RA, dec)
except:
coords = SkyCoord(float(RA), float(dec), unit='deg')
self.coords = coords
self._ic = ic
self.use_emcee = use_emcee
# If obs is not provided, build it
if obs is None:
self._build_obs(**kwargs)
self.obs.define_models(ic, N=N, index=index)
self._add_properties(**kwargs)
elif isinstance(obs, basestring):
df = pd.read_csv(obs)
obs = ObservationTree.from_df(df)
obs.define_models(ic, N=N, index=index)
self.obs = obs
self._add_properties(**kwargs)
else:
self.obs = obs
if len(self.obs.get_model_nodes())==0:
self.obs.define_models(ic, N=N, index=index)
self._add_properties(**kwargs)
self._priors = {'mass':salpeter_prior,
'feh':feh_prior,
'q':q_prior,
'age':age_prior,
'distance':distance_prior,
'AV':AV_prior}
self._bounds = {'mass':None,
'feh':None,
'age':None,
'q':q_prior.bounds,
'distance':distance_prior.bounds,
'AV':AV_prior.bounds}
if 'maxAV' in kwargs:
self.set_bounds(AV=(0, kwargs['maxAV']))
if 'max_distance' in kwargs:
self.set_bounds(distance=(0, kwargs['max_distance']))
self._directory = '.'
self._samples = None
@property
def directory(self):
return self._directory
@property
def ic(self):
if type(self._ic)==type:
self._ic = self._ic()
return self._ic
@classmethod
def _parse_band(cls, kw):
"""Returns photometric band from inifile keyword
"""
m = re.search('([a-zA-Z0-9]+)(_\d+)?', kw)
if m:
if m.group(1) in cls._not_a_band:
return None
else:
return m.group(1)
@classmethod
def get_bands(cls, inifile):
bands = []
c = configobj.ConfigObj(inifile)
for kw,v in c.items():
if type(v) is configobj.Section:
for kw in v:
b = cls._parse_band(kw)
if b is not None:
bands.append(b)
else:
b = cls._parse_band(kw)
if b is not None:
bands.append(b)
return list(set(bands))
[docs] @classmethod
def from_ini(cls, ic, folder='.', ini_file='star.ini', **kwargs):
"""
Initialize a StarModel from a .ini file
The "classic" format (version <= 0.9) should still work for a single star,
where all properties are just listed in the file; e.g.,
J = 10, 0.05
H = 9.5, 0.05
K = 9.0, 0.05
Teff = 5000, 150
If there are multiple stars observed, you can either define them in
the ini file, or use the `obsfile` keyword, pointing to a file with
the summarized photometric observations. In this case, spectroscopic/parallax
info should still be included in the .ini file; e.g.,
obsfile = obs.csv
Teff = 5000, 150
The obsfile should be a comma-separated table with the following columns:
`[name, band, resolution, mag, e_mag, separation, pa, relative]`.
* `name` is the name of instrument
* `band` is the photometric bandpass
* `resolution` is the approximate spatial resolution of instrument
* `mag`, `e_mag` describe magnitude of source (absolute or relative)
* `separation`, `pa` describe position of source
* `relative`: single-bit flag; if 1 then magnitudes taken with this
instrument are assumed to be relative rather than absolute.
If an obsfile is not provided, you can also define all the same information
in the ini file, following these rules:
* Every instrument/survey gets its own [section]. Sections are only
created for different photometric observations.
* if photometry relates to *all* stars in aperture,
there is no extra info in the section, just the photometry. In this case, it is
also assumed that the photometry is absolute. (`relative=False`)
* If 'resolution' is an attribute under a particular survey section (and
'relative' is not explicitly stated), then the survey is assumed to have relative
photometry, and to be listing
information about companion stars. In this case, there must be "separation"
and "PA" included for each companion. If there is more than one companion star,
they must be identifed by tag, e.g., separation_1, PA_1, Ks_1, J_1, etc. The
tag can be anything alphanumeric, but it must be consistent within a particular
section (instrument). If there
is no tag, there is assumed to be only one companion detected.
* If there are no sections, then bands will be interpreted at face value
and will all be assumed to apply to all stars modeled.
* Default is to model each star in the highest-resolution observation as a
single star, at the same distance/age/feh/AV.
The `N` and `index`
parameters may also be provided, to specify the relations between the
model stars. If these are not provided, then `N` will default to `1`
(one model star per star observed in highest-resolution observation)
and `index` will default to all `0` (all stars physically associated).
"""
if not os.path.isabs(ini_file):
ini_file = os.path.join(folder,ini_file)
bands = cls.get_bands(ini_file)
if not isinstance(ic, Isochrone):
ic = get_ichrone(ic, bands)
logging.debug('Initializing StarModel from {}'.format(ini_file))
c = configobj.ConfigObj(ini_file)
RA = c.get('RA')
dec = c.get('dec')
maxAV = c.get('maxAV')
if len(c.sections) == 0:
for k in c:
kwargs[k] = _parse_config_value(c[k])
obs = None
else:
columns = ['name', 'band', 'resolution', 'relative', 'separation', 'pa', 'mag', 'e_mag']
df = pd.DataFrame(columns=columns)
i = 0
for k in c:
if type(c[k]) != configobj.Section:
kwargs[k] = _parse_config_value(c[k])
else:
instrument = k
# Set values of 'resolution' and 'relative'
if 'resolution' in c[k]:
resolution = float(c[k]['resolution'])
relative = True
else:
resolution = 4.0 #default
relative = False
# Overwrite value of 'relative' if it is explicitly set
if 'relative' in c[k]:
relative = c[k]['relative']=='True'
# Check if there are multiple stars (defined by whether
# any separations are listed).
# While we're at it, keep track of tags if they exist,
# and pull out the names of the bands.
multiple = False
tags = []
bands = []
for label in c[k]:
m = re.search('separation(_\w+)?', label)
if m:
multiple = True
if m.group(1) is not None:
if m.group(1) not in tags:
tags.append(m.group(1))
elif re.search('PA', label) or re.search('id', label) or \
label in ['resolution', 'relative']:
continue
else:
# At this point, this should be a photometric band
m = re.search('([a-zA-Z0-9]+)(_\w+)?', label)
b = m.group(1)
if b not in bands:
bands.append(b)
# If a blank tags needs to be created, do so
if len(bands) > 0 and (len(tags)==0 or bands[0] in c[k]):
tags.append('')
# For each band and each star, create a row
for b in bands:
for tag in tags:
if '{}{}'.format(b, tag) not in c[k]:
continue
row = {}
row['name'] = instrument
row['band'] = b
row['resolution'] = resolution
row['relative'] = relative
if 'separation{}'.format(tag) in c[k]:
row['separation'] = c[k]['separation{}'.format(tag)]
row['pa'] = c[k]['PA{}'.format(tag)]
else:
row['separation'] = 0.
row['pa'] = 0.
mag, e_mag = c[k]['{}{}'.format(b,tag)]
row['mag'] = float(mag)
row['e_mag'] = float(e_mag)
if not np.isnan(row['mag']) and not np.isnan(row['e_mag']):
df = df.append(pd.DataFrame(row, index=[i]))
i += 1
# put the reference star in w/ mag=0
if relative:
row = {}
row['name'] = instrument
row['band'] = b
row['resolution'] = resolution
row['relative'] = relative
row['separation'] = 0.
row['pa'] = 0.
row['mag'] = 0.
row['e_mag'] = 0.01
df = df.append(pd.DataFrame(row, index=[i]))
i += 1
obs = ObservationTree.from_df(df)
if 'obsfile' in c:
obs = c['obsfile']
logging.debug('Obs is {}'.format(obs))
if 'name' not in kwargs:
kwargs['name'] = os.path.basename(folder)
new = StarModel(ic, obs=obs, **kwargs)
new._directory = os.path.abspath(folder)
return new
[docs] def print_ascii(self):
"""Prints an ascii representation of the observation tree structure.
"""
return self.obs.print_ascii()
def bounds(self, prop):
if self._bounds[prop] is not None:
return self._bounds[prop]
elif prop=='mass':
lo, hi = (self.ic.minmass, self.ic.maxmass)
self._bounds['mass'] = (lo, hi)
self._priors['mass'].bounds = (lo, hi)
elif prop=='feh':
lo, hi = (self.ic.minfeh, self.ic.maxfeh)
self._bounds['feh'] = (lo, hi)
self._priors['feh'].bounds = (lo, hi)
elif prop=='age':
lo, hi = (self.ic.minage, self.ic.maxage)
self._bounds['age'] = (lo, hi)
self._priors['age'].bounds = (lo, hi)
self._bounds['age'] = (self.ic.minage,
self.ic.maxage)
else:
raise ValueError('Unknown property {}'.format(prop))
return self._bounds[prop]
def set_bounds(self, **kwargs):
for k,v in kwargs.items():
if len(v) != 2:
raise ValueError('Must provide (min, max)')
self._bounds[k] = v
self._priors[k].bounds = v
def _build_obs(self, **kwargs):
"""
Builds ObservationTree out of keyword arguments
Ignores anything that is not a photometric bandpass.
This should not be used if there are multiple stars observed.
Creates self.obs
"""
logging.debug('Building ObservationTree...')
tree = ObservationTree()
for k,v in kwargs.items():
if k in self.ic.bands:
if np.size(v) != 2:
logging.warning('{}={} ignored.'.format(k,v))
# continue
v = [v, np.nan]
o = Observation('', k, 99) #bogus resolution=99
s = Source(v[0], v[1])
o.add_source(s)
logging.debug('Adding {} ({})'.format(s,o))
tree.add_observation(o)
self.obs = tree
def _add_properties(self, **kwargs):
"""
Adds non-photometry properties to ObservationTree
"""
for k,v in kwargs.items():
if k=='parallax':
self.obs.add_parallax(v)
elif k in ['Teff', 'logg', 'feh', 'density']:
par = {k:v}
self.obs.add_spectroscopy(**par)
elif re.search('_', k):
m = re.search('^(\w+)_(\w+)$', k)
prop = m.group(1)
tag = m.group(2)
self.obs.add_spectroscopy(**{prop:v, 'label':'0_{}'.format(tag)})
@property
def param_description(self):
return self.obs.param_description
@property
def param_names(self):
return self.param_description
@property
def mags(self):
return {n.band : n.value[0] for n in self.obs.get_obs_nodes()}
def lnpost(self, p, **kwargs):
lnpr = self.lnprior(p)
if not np.isfinite(lnpr):
return lnpr
return lnpr + self.lnlike(p, **kwargs)
def lnlike(self, p, **kwargs):
lnl = self.obs.lnlike(p, **kwargs)
return lnl
def lnprior(self, p):
N = self.obs.Nstars
i = 0
lnp = 0
for s in self.obs.systems:
age, feh, dist, AV = p[i+N[s]:i+N[s]+4]
for prop, val in zip(['age','feh','distance','AV'],
[age, feh, dist, AV]):
lo,hi = self.bounds(prop)
if val < lo or val > hi:
return -np.inf
lnp += np.log(self.prior(prop, val))
if not np.isfinite(lnp):
logging.debug('lnp=-inf for {}={} (system {})'.format(prop,val,s))
return -np.inf
# Note: this is just assuming proper order.
# Is this OK? Should keep eye out for bugs here.
masses = p[i:i+N[s]]
# Mass prior for primary
lnp += np.log(self.prior('mass', masses[0]))
if not np.isfinite(lnp):
logging.debug('lnp=-inf for mass={} (system {})'.format(masses[0],s))
# Priors for mass ratios
for j in range(N[s]-1):
q = masses[j+1]/masses[0]
qmin, qmax = self.bounds('q')
## The following would enforce MA > MB > MC, but seems to make things very slow:
#if j+1 > 1:
# qmax = masses[j] / masses[0]
lnp += np.log(self.prior('q', q))
if not np.isfinite(lnp):
logging.debug('lnp=-inf for q={} (system {})'.format(q,s))
return -np.inf
i += N[s] + 4
return lnp
def prior(self, prop, val, **kwargs):
return self._priors[prop](val, **kwargs)
@property
def n_params(self):
tot = 0
for _,n in self.obs.Nstars.items():
tot += 4+n
return tot
def mnest_prior(self, cube, ndim, nparams):
i = 0
for _,n in self.obs.Nstars.items():
minmass, maxmass = self.bounds('mass')
for j in range(n):
cube[i+j] = (maxmass - minmass)*cube[i+j] + minmass
for j, par in enumerate(['age','feh','distance','AV']):
lo, hi = self.bounds(par)
cube[i+n+j] = (hi - lo)*cube[i+n+j] + lo
i += 4 + n
[docs] def mnest_loglike(self, cube, ndim, nparams):
"""loglikelihood function for multinest
"""
return self.lnpost(cube)
@property
def labelstring(self):
return '--'.join(['-'.join([n.label for n in l.children]) for l in self.obs.get_obs_leaves()])
def fit(self, **kwargs):
if self.use_emcee:
return self.fit_mcmc(**kwargs)
else:
return self.fit_multinest(**kwargs)
@property
def mnest_basename(self):
"""Full path to basename
"""
if not hasattr(self, '_mnest_basename'):
s = self.labelstring
if s=='0_0':
s = 'single'
elif s=='0_0-0_1':
s = 'binary'
elif s=='0_0-0_1-0_2':
s = 'triple'
s = '{}-{}'.format(self.ic.name, s)
self._mnest_basename = os.path.join('chains', s+'-')
if os.path.isabs(self._mnest_basename):
return self._mnest_basename
else:
return os.path.join(self.directory, self._mnest_basename)
@mnest_basename.setter
def mnest_basename(self, basename):
if os.path.isabs(basename):
self._mnest_basename = basename
else:
self._mnest_basename = os.path.join('chains', basename)
def lnpost_polychord(self, theta):
phi = [0.0] #nDerived
return self.lnpost(theta), phi
def fit_polychord(self, basename, verbose=False, **kwargs):
from .config import POLYCHORD
sys.path.append(POLYCHORD)
import PyPolyChord.PyPolyChord as PolyChord
return PolyChord.run_nested_sampling(self.lnpost_polychord,
self.n_params, 0, file_root=basename, **kwargs)
[docs] def fit_multinest(self, n_live_points=1000, basename=None,
verbose=True, refit=False, overwrite=False,
test=False,
**kwargs):
"""
Fits model using MultiNest, via pymultinest.
:param n_live_points:
Number of live points to use for MultiNest fit.
:param basename:
Where the MulitNest-generated files will live.
By default this will be in a folder named `chains`
in the current working directory. Calling this
will define a `_mnest_basename` attribute for
this object.
:param verbose:
Whether you want MultiNest to talk to you.
:param refit, overwrite:
Set either of these to true if you want to
delete the MultiNest files associated with the
given basename and start over.
:param **kwargs:
Additional keyword arguments will be passed to
:func:`pymultinest.run`.
"""
if basename is not None: #Should this even be allowed?
self.mnest_basename = basename
basename = self.mnest_basename
if verbose:
logging.info('MultiNest basename: {}'.format(basename))
folder = os.path.abspath(os.path.dirname(basename))
if not os.path.exists(folder):
os.makedirs(folder)
#If previous fit exists, see if it's using the same
# observed properties
prop_nomatch = False
propfile = '{}properties.json'.format(basename)
"""
if os.path.exists(propfile):
with open(propfile) as f:
props = json.load(f)
if set(props.keys()) != set(self.properties.keys()):
prop_nomatch = True
else:
for k,v in props.items():
if np.size(v)==2:
if not self.properties[k][0] == v[0] and \
self.properties[k][1] == v[1]:
props_nomatch = True
else:
if not self.properties[k] == v:
props_nomatch = True
if prop_nomatch and not overwrite:
raise ValueError('Properties not same as saved chains ' +
'(basename {}*). '.format(basename) +
'Use overwrite=True to fit.')
"""
if refit or overwrite:
files = glob.glob('{}*'.format(basename))
[os.remove(f) for f in files]
short_basename = self._mnest_basename
mnest_kwargs = dict(n_live_points=n_live_points, outputfiles_basename=short_basename,
verbose=verbose)
for k,v in kwargs.items():
mnest_kwargs[k] = v
if test:
print('pymultinest.run() with the following kwargs: {}'.format(mnest_kwargs))
else:
wd = os.getcwd()
os.chdir(os.path.join(folder, '..'))
pymultinest.run(self.mnest_loglike, self.mnest_prior, self.n_params,
**mnest_kwargs)
os.chdir(wd)
#with open(propfile, 'w') as f:
# json.dump(self.properties, f, indent=2)
self._make_samples()
@property
def mnest_analyzer(self):
"""
PyMultiNest Analyzer object associated with fit.
See PyMultiNest documentation for more.
"""
return pymultinest.Analyzer(self.n_params, self.mnest_basename)
@property
def evidence(self):
"""
Log(evidence) from multinest fit
"""
s = self.mnest_analyzer.get_stats()
return (s['global evidence'],s['global evidence error'])
[docs] def maxlike(self, p0, **kwargs):
""" Finds (local) optimum in parameter space.
"""
def fn(p):
return -self.lnpost(p)
if 'method' not in kwargs:
kwargs['method'] = 'Nelder-Mead'
p0 = [0.8, 9.5, 0.0, 200, 0.2]
fit = scipy.optimize.minimize(fn, p0, **kwargs)
return fit
def sample_from_prior(self, n):
return self.emcee_p0(n)
def emcee_p0(self, nwalkers):
def sample_row(nstars, n=nwalkers):
p = []
m0 = self._priors['mass'].sample(n)
age0 = self._priors['age'].sample(n)
feh0 = self._priors['feh'].sample(n)
d0 = self._priors['distance'].sample(n)
AV0 = self._priors['AV'].sample(n)
for i in range(nstars):
p += [m0 * 0.95**i]
p += [age0, feh0, d0, AV0]
return p
p0 = []
for _,n in self.obs.Nstars.items():
p0 += sample_row(n)
p0 = np.array(p0).T
nbad = 1
while True:
ibad = []
for i, p in enumerate(p0):
if not np.isfinite(self.lnpost(p)):
ibad.append(i)
nbad = len(ibad)
if nbad == 0:
break
pnew = []
for _, n in self.obs.Nstars.items():
pnew += sample_row(n, n=nbad)
pnew = np.array(pnew).T
p0[ibad, :] = pnew
return p0
[docs] def fit_mcmc(self,nwalkers=300,nburn=200,niter=100,
p0=None,initial_burn=None,
ninitial=50, loglike_kwargs=None,
**kwargs):
"""Fits stellar model using MCMC.
:param nwalkers: (optional)
Number of walkers to pass to :class:`emcee.EnsembleSampler`.
Default is 200.
:param nburn: (optional)
Number of iterations for "burn-in." Default is 100.
:param niter: (optional)
Number of for-keeps iterations for MCMC chain.
Default is 200.
:param p0: (optional)
Initial parameters for emcee. If not provided, then chains
will behave according to whether inital_burn is set.
:param initial_burn: (optional)
If `True`, then initialize walkers first with a random initialization,
then cull the walkers, keeping only those with > 15% acceptance
rate, then reinitialize sampling. If `False`, then just do
normal burn-in. Default is `None`, which will be set to `True` if
fitting for distance (i.e., if there are apparent magnitudes as
properties of the model), and `False` if not.
:param ninitial: (optional)
Number of iterations to test walkers for acceptance rate before
re-initializing.
:param loglike_args:
Any arguments to pass to :func:`StarModel.loglike`, such
as what priors to use.
:param **kwargs:
Additional keyword arguments passed to :class:`emcee.EnsembleSampler`
constructor.
:return:
:class:`emcee.EnsembleSampler` object.
"""
#clear any saved _samples
if self._samples is not None:
self._samples = None
npars = self.n_params
if p0 is None:
p0 = self.emcee_p0(nwalkers)
if initial_burn:
sampler = emcee.EnsembleSampler(nwalkers,npars,self.lnpost,
**kwargs)
#ninitial = 300 #should this be parameter?
pos, prob, state = sampler.run_mcmc(p0, ninitial)
# Choose walker with highest final lnprob to seed new one
i,j = np.unravel_index(sampler.lnprobability.argmax(),
sampler.shape)
p0_best = sampler.chain[i,j,:]
print("After initial burn, p0={}".format(p0_best))
p0 = p0_best * (1 + rand.normal(size=p0.shape)*0.001)
print(p0)
else:
p0 = np.array(p0)
p0 = rand.normal(size=(nwalkers,npars))*0.01 + p0.T[None,:]
sampler = emcee.EnsembleSampler(nwalkers,npars,self.lnpost)
pos, prob, state = sampler.run_mcmc(p0, nburn)
sampler.reset()
sampler.run_mcmc(pos, niter, rstate0=state)
self._sampler = sampler
return sampler
@property
def sampler(self):
"""
Sampler object from MCMC run.
"""
if hasattr(self,'_sampler'):
return self._sampler
else:
raise AttributeError('MCMC must be run to access sampler')
def _make_samples(self):
if not self.use_emcee:
filename = '{}post_equal_weights.dat'.format(self.mnest_basename)
try:
chain = np.loadtxt(filename)
try:
lnprob = chain[:,-1]
chain = chain[:,:-1]
except IndexError:
lnprob = np.array([chain[-1]])
chain = np.array([chain[:-1]])
except:
logging.error('Error loading chains from {}'.format(filename))
raise
else:
#select out only walkers with > 0.15 acceptance fraction
ok = self.sampler.acceptance_fraction > 0.15
chain = self.sampler.chain[ok,:,:]
chain = chain.reshape((chain.shape[0]*chain.shape[1],
chain.shape[2]))
lnprob = self.sampler.lnprobability[ok, :].ravel()
df = pd.DataFrame()
i=0
for s,n in self.obs.Nstars.items():
age = chain[:,i+n]
feh = chain[:,i+n+1]
distance = chain[:,i+n+2]
AV = chain[:,i+n+3]
for j in range(n):
mass = chain[:,i+j]
d = self.ic(mass, age, feh,
distance=distance, AV=AV)
for c in d.columns:
df[c+'_{}_{}'.format(s,j)] = d[c]
df['age_{}'.format(s)] = age
df['feh_{}'.format(s)] = feh
df['distance_{}'.format(s)] = distance
df['AV_{}'.format(s)] = AV
i += 4 + n
for b in self.ic.bands:
tot = np.inf
for s,n in self.obs.Nstars.items():
for j in range(n):
tot = addmags(tot,df[b + '_mag_{}_{}'.format(s,j)])
df[b + '_mag'] = tot
df['lnprob'] = lnprob
self._samples = df.copy()
@property
def samples(self):
"""Dataframe with samples drawn from isochrone according to posterior
Columns include both the sampling parameters from the MCMC
fit (mass, age, Fe/H, [distance, A_V]), and also evaluation
of the :class:`Isochrone` at each of these sample points---this
is how chains of physical/observable parameters get produced.
"""
if not hasattr(self,'sampler') and self._samples is None:
raise AttributeError('Must run MCMC (or load from file) '+
'before accessing samples')
if self._samples is not None:
df = self._samples
else:
self._make_samples()
df = self._samples
return df
[docs] def random_samples(self, n):
"""
Returns a random sampling of given size from the existing samples.
:param n:
Number of samples
:return:
:class:`pandas.DataFrame` of length ``n`` with random samples.
"""
samples = self.samples
inds = rand.randint(len(samples),size=int(n))
newsamples = samples.iloc[inds]
newsamples.reset_index(inplace=True)
return newsamples
def triangle(self, *args, **kwargs):
return self.corner(*args, **kwargs)
def corner(self, params, query=None, **kwargs):
df = self.samples
if query is not None:
df = df.query(query)
priors = []
for p in params:
if re.match('mass', p):
priors.append(lambda x: self.prior('mass', x, bounds=self.bounds('mass')))
elif re.match('age', p):
priors.append(lambda x: self.prior('age', x, bounds=self.bounds('age')))
elif re.match('feh', p):
priors.append(lambda x: self.prior('feh', x, bounds=self.bounds('feh')))
elif re.match('distance', p):
priors.append(lambda x: self.prior('distance', x, bounds=self.bounds('distance')))
elif re.match('AV', p):
priors.append(lambda x: self.prior('AV', x, bounds=self.bounds('AV')))
else:
priors.append(None)
try:
fig = corner.corner(df[params], labels=params, priors=priors, **kwargs)
except:
logging.warning("Use Tim's version of corner to plot priors.")
fig = corner.corner(df[params], labels=params, **kwargs)
fig.suptitle(self.name, fontsize=22)
return fig
def triangle_physical(self, *args, **kwargs):
return self.corner_physical(*args, **kwargs)
def corner_plots(self, basename, **kwargs):
fig1, fig2 = self.corner_physical(**kwargs), self.corner_observed(**kwargs)
fig1.savefig(basename + '_physical.png')
fig2.savefig(basename + '_observed.png')
return fig1, fig2
def triangle_plots(self, *args, **kwargs):
return self.corner_plots(*args, **kwargs)
def corner_physical(self, props=['mass','radius','feh','age','distance','AV'], **kwargs):
collective_props = ['feh','age','distance','AV']
indiv_props = [p for p in props if p not in collective_props]
sys_props = [p for p in props if p in collective_props]
props = ['{}_{}'.format(p,l) for p in indiv_props for l in self.obs.leaf_labels]
props += ['{}_{}'.format(p,s) for p in sys_props for s in self.obs.systems]
if 'range' not in kwargs:
rng = [0.995 for p in props]
return self.corner(props, range=rng, **kwargs)
def mag_plot(self, *args, **kwargs):
pass
[docs] def corner_observed(self, **kwargs):
"""Makes corner plot for each observed node magnitude
"""
tot_mags = []
names = []
truths = []
rng = []
for n in self.obs.get_obs_nodes():
labels = [l.label for l in n.get_model_nodes()]
band = n.band
mags = [self.samples['{}_mag_{}'.format(band, l)] for l in labels]
tot_mag = addmags(*mags)
if n.relative:
name = '{} $\Delta${}'.format(n.instrument, n.band)
ref = n.reference
if ref is None:
continue
ref_labels = [l.label for l in ref.get_model_nodes()]
ref_mags = [self.samples['{}_mag_{}'.format(band, l)] for l in ref_labels]
tot_ref_mag = addmags(*ref_mags)
tot_mags.append(tot_mag - tot_ref_mag)
truths.append(n.value[0] - ref.value[0])
else:
name = '{} {}'.format(n.instrument, n.band)
tot_mags.append(tot_mag)
truths.append(n.value[0])
names.append(name)
rng.append((min(truths[-1], np.percentile(tot_mags[-1],0.5)),
max(truths[-1], np.percentile(tot_mags[-1],99.5))))
tot_mags = np.array(tot_mags).T
return corner.corner(tot_mags, labels=names, truths=truths, range=rng, **kwargs)
[docs] def save_hdf(self, filename, path='', overwrite=False, append=False):
"""Saves object data to HDF file (only works if MCMC is run)
Samples are saved to /samples location under given path,
:class:`ObservationTree` is saved to /obs location under given path.
:param filename:
Name of file to save to. Should be .h5 file.
:param path: (optional)
Path within HDF file structure to save to.
:param overwrite: (optional)
If ``True``, delete any existing file by the same name
before writing.
:param append: (optional)
If ``True``, then if a file exists, then just the path
within the file will be updated.
"""
if os.path.exists(filename):
with pd.HDFStore(filename) as store:
if path in store:
if overwrite:
os.remove(filename)
elif not append:
raise IOError('{} in {} exists. Set either overwrite or append option.'.format(path,filename))
if self.samples is not None:
self.samples.to_hdf(filename, path+'/samples')
else:
pd.DataFrame().to_hdf(filename, path+'/samples')
self.obs.save_hdf(filename, path+'/obs', append=True)
with pd.HDFStore(filename) as store:
# store = pd.HDFStore(filename)
attrs = store.get_storer('{}/samples'.format(path)).attrs
attrs.ic_type = type(self.ic)
attrs.ic_bands = list(self.ic.bands)
attrs.use_emcee = self.use_emcee
if hasattr(self, '_mnest_basename'):
attrs._mnest_basename = self._mnest_basename
attrs._bounds = self._bounds
attrs._priors = self._priors
attrs.name = self.name
store.close()
[docs] @classmethod
def load_hdf(cls, filename, path='', name=None):
"""
A class method to load a saved StarModel from an HDF5 file.
File must have been created by a call to :func:`StarModel.save_hdf`.
:param filename:
H5 file to load.
:param path: (optional)
Path within HDF file.
:return:
:class:`StarModel` object.
"""
if not os.path.exists(filename):
raise IOError('{} does not exist.'.format(filename))
store = pd.HDFStore(filename)
try:
samples = store[path+'/samples']
attrs = store.get_storer(path+'/samples').attrs
except:
store.close()
raise
try:
ic = attrs.ic_type(attrs.ic_bands)
except AttributeError:
ic = attrs.ic_type
use_emcee = attrs.use_emcee
mnest = True
try:
basename = attrs._mnest_basename
except AttributeError:
mnest = False
bounds = attrs._bounds
priors = attrs._priors
if name is None:
try:
name = attrs.name
except:
name = ''
store.close()
obs = ObservationTree.load_hdf(filename, path+'/obs', ic=ic)
mod = cls(ic, obs=obs,
use_emcee=use_emcee, name=name)
mod._samples = samples
if mnest:
mod._mnest_basename = basename
mod._directory = os.path.dirname(filename)
return mod
class BinaryStarModel(StarModel):
_default_name = 'binary'
def __init__(self, *args, **kwargs):
kwargs['N'] = 2
super(BinaryStarModel, self).__init__(*args, **kwargs)
@classmethod
def from_ini(cls, *args, **kwargs):
kwargs['N'] = 2
return super(BinaryStarModel, cls).from_ini(*args, **kwargs)
class TripleStarModel(StarModel):
_default_name = 'triple'
def __init__(self, *args, **kwargs):
kwargs['N'] = 3
super(TripleStarModel, self).__init__(*args, **kwargs)
@classmethod
def from_ini(cls, *args, **kwargs):
kwargs['N'] = 3
return super(TripleStarModel, cls).from_ini(*args, **kwargs)
class StarModelGroup(object):
"""A collection of StarModel objects with different model node specifications
Pass a single StarModel, and model nodes will be cleared and replaced with
different variants.
"""
def __init__(self, base_model, max_multiples=1, max_stars=2):
self.base_model = deepcopy(base_model)
self.base_model.obs.clear_models()
self.max_multiples = max_multiples
self.max_stars = max_stars
self.models = []
for N, index in self.model_options:
mod = deepcopy(self.base_model)
mod.obs.define_models(self.ic, N=N, index=index)
self.models.append(mod)
@property
def ic(self):
return self.base_model.ic
@property
def N_stars(self):
return len(self.base_model.obs.leaves)
@property
def N_options(self):
return N_options(self.N_stars, max_multiples=self.max_multiples,
max_stars=self.max_stars)
@property
def index_options(self):
return index_options(self.N_stars)
@property
def model_options(self):
return [(N, index) for N in self.N_options for index in self.index_options]
########## Utility functions ###############
def N_options(N_stars, max_multiples=1, max_stars=2):
return [N for N in itertools.product(np.arange(max_stars) + 1, repeat=N_stars)
if (np.array(N)>1).sum() <= max_multiples]
def index_options(N_stars):
if N_stars==1:
return [0]
options = []
for ind in itertools.product(range(N_stars), repeat=N_stars):
diffs = np.array(ind[1:]) - np.array(ind[:-1])
if ind[0]==0 and diffs.max()<=1:
options.append(ind)
return options