示例#1
0
def main():
    # TODO: bad, hard-coded...
    # base_path = '/Volumes/ProjectData/gaia-comoving-followup/'
    base_path = '../data/'
    db_path = path.join(base_path, 'db.sqlite')
    engine = db_connect(db_path)
    session = Session()

    credentials = dict(user='******', password='******')
    Gaia.login(**credentials)

    for obs in session.query(Observation).all():
        q = session.query(Photometry).join(Observation).filter(Observation.id == obs.id).count()
        if q > 0:
            logger.debug('Photometry already exists')
            continue

        if obs.tgas_source is None:
            continue

        tgas_source_id = obs.tgas_source.source_id
        res = get_photometry(tgas_source_id)

        phot_kw = dict()
        for col in result_columns:
            phot_kw[col] = res[col]

        phot = Photometry(**phot_kw)
        phot.observation = obs
        session.add(phot)
        session.commit()
示例#2
0
def main(db_file, pool, overwrite=False):

    # HACK:
    base_path = '../data/'
    db_path = path.join(base_path, 'db.sqlite')
    engine = db_connect(db_path)
    session = Session()

    # HACK:
    from astropy.table import Table
    tbl = Table.read('../paper/figures/group_llr_dv_tbl.ecsv',
                     format='ascii.ecsv')

    worker = Worker(db_path=db_path,
                    samples_path=path.abspath('../data/isochrone_samples'),
                    overwrite=overwrite)

    # A little bit of a hack
    comoving = tbl['R_RV'] > tbl['R_mu']
    tasks = session.query(Observation.id)\
                   .filter(Observation.group_id.in_(tbl['group_id'][comoving]))\
                   .group_by(Observation.group_id).all()
    tasks = tasks[:1]

    session.close()

    for r in pool.map(worker, tasks, callback=worker.callback):
        pass

    pool.close()
    sys.exit(0)
示例#3
0
def main(db_path, run_name, overwrite=False, pool=None):

    if pool is None:
        pool = schwimmbad.SerialPool()

    # connect to the database
    engine = db_connect(db_path)
    # engine.echo = True
    logger.debug("Connected to database at '{}'".format(db_path))

    # create a new session for interacting with the database
    session = Session()

    root_path, _ = path.split(db_path)
    plot_path = path.join(root_path, 'plots', run_name)
    if not path.exists(plot_path):
        os.makedirs(plot_path, exist_ok=True)

    # get object to correct the observed RV's
    rv_corr = RVCorrector(session, run_name)

    observations = session.query(Observation).join(Run)\
                          .filter(Run.name == run_name).all()

    for obs in observations:
        q = session.query(RVMeasurement).join(Observation)\
                   .filter(Observation.id == obs.id)

        if q.count() > 0 and not overwrite:
            logger.debug('RV measurement already complete for object '
                         '{0} in file {1}'.format(obs.object,
                                                  obs.filename_raw))
            continue

        elif q.count() > 1:
            raise RuntimeError(
                'Multiple RV measurements found for object {0}'.format(obs))

        elif len(obs.measurements) == 0:
            logger.debug(
                'Observation {0} has no line measurements.'.format(obs))
            continue

        corrected_rv, err, flag = rv_corr.get_corrected_rv(obs)

        # remove previous RV measurements
        if q.count() > 0:
            session.delete(q.one())
            session.commit()

        rv_meas = RVMeasurement(rv=corrected_rv, err=err, flag=flag)
        rv_meas.observation = obs
        session.add(rv_meas)
        session.commit()

    pool.close()
示例#4
0
    def work(self, id):
        engine = db_connect(self.db_path)
        session = Session()
        obs = session.query(Observation).filter(Observation.id == id).one()
        model = obs_to_starmodel(obs)

        # initial conditions for emcee walkers
        p0 = []
        m0, age0, feh0 = model.ic.random_points(self.nwalkers,
                                                minmass=0.01,
                                                maxmass=10.,
                                                minfeh=-1,
                                                maxfeh=1)
        _, max_distance = model.bounds('distance')
        _, max_AV = model.bounds('AV')
        d0 = 10**(np.random.uniform(0,
                                    np.log10(max_distance),
                                    size=self.nwalkers))
        AV0 = np.random.uniform(0, max_AV, size=self.nwalkers)
        p0 += [m0]
        p0 += [age0, feh0, d0, AV0]

        p0 = np.array(p0).T
        npars = p0.shape[1]

        logger.debug('Running emcee - initial sampling...')
        sampler = emcee.EnsembleSampler(self.nwalkers, npars, model.lnpost)

        pos, prob, _ = sampler.run_mcmc(p0, self.ninit)

        # cull the weak walkers
        best_ix = sampler.flatlnprobability.argmax()
        best_p0 = (sampler.flatchain[best_ix][None] +
                   np.random.normal(0, 1E-5, size=(self.nwalkers, npars)))

        sampler.reset()
        logger.debug('burn-in...')
        pos, prob, _ = sampler.run_mcmc(best_p0, self.nburn)

        sampler.reset()
        logger.debug('sampling...')
        _ = sampler.run_mcmc(pos, self.niter)

        model._sampler = sampler
        model._make_samples(0.01)

        return id, model
示例#5
0
def main(db_path, run_root_path, drop_all=False, overwrite=False, **kwargs):

    # Make sure the specified paths actually exist
    db_path = path.abspath(db_path)
    run_root_path = path.abspath(run_root_path)
    for path_ in [path.dirname(db_path), run_root_path]:
        if not path.exists(path_):
            raise ValueError("Path '{0}' doesn't exist!".format(path_))

    # --------------------------------------------------------------------------
    # These are relative paths, so the script needs to be run from the
    #   scripts path...

    # ID table for mapping group index to TGAS row
    ID_tbl = Table.read('../data/star_identifier.csv')

    # TGAS table
    logger.debug("Loading TGAS data...")
    tgas = Table.read('../../gaia-comoving-stars/data/stacked_tgas.fits')

    # Catalog of velocities for Bensby's HIP stars:
    bensby = Table.read('../data/bensbyrv_bestunique.csv')

    # --------------------------------------------------------------------------

    # connect to the database
    engine = db_connect(db_path, ensure_db_exists=True)
    # engine.echo = True
    logger.debug("Connected to database at '{}'".format(db_path))

    if drop_all: # remove all tables and replace
        Base.metadata.drop_all()
        Base.metadata.create_all()

    # create a new session for interacting with the database
    session = Session()

    logger.debug("Loading SpectralLineInfo table")

    line_info = OrderedDict()
    # air wavelength of Halpha -- wavelength calibration from comp lamp is done
    #   at air wavelengths, so this is where Halpha should be, right?
    line_info['Halpha'] = 6562.8*u.angstrom

    # [OI] emission lines -- wavelengths from:
    #   http://physics.nist.gov/PhysRefData/ASD/lines_form.html
    line_info['[OI] 5577'] = 5577.3387*u.angstrom
    line_info['[OI] 6300'] = 6300.304*u.angstrom
    line_info['[OI] 6364'] = 6363.776*u.angstrom

    for name, wvln in line_info.items():
        n = session.query(SpectralLineInfo).filter(SpectralLineInfo.name == name).count()
        if n == 0:
            logger.debug('Loading line {0} at {1}'.format(name, wvln))
            line = SpectralLineInfo(name=name, wavelength=wvln)
            session.add(line)
            session.commit()
        else:
            logger.debug('Line {0} already loaded'.format(name))

    # Create an entry for this observing run
    data_path, run_name = path.split(run_root_path)
    logger.info("Path to night paths: {0}".format(data_path))
    n = session.query(Run).filter(Run.name == run_name).count()
    if n == 0:
        logger.debug('Adding run {0} to database'.format(run_name))
        run = Run(name=run_name)
        session.add(run)
        session.commit()

    elif n == 1:
        logger.debug('Loading run from database'.format(run_name))
        run = session.query(Run).filter(Run.name == run_name).limit(1).one()

    else:
        raise RuntimeError("F**k.")

    # Now we need to go through each processed night of data and load all of the
    # relevant observations of sources.

    # First we get the column names for the Observation and TGASSource tables
    obs_columns = [str(c).split('.')[1] for c in Observation.__table__.columns]
    tgassource_columns = [str(c).split('.')[1]
                          for c in TGASSource.__table__.columns]

    # Here's where there's a bit of hard-coded bewitchery - the nights (within
    # each run) have to be labeled 'n1', 'n2', and etc. Sorry.
    glob_pattr_proc = path.join(data_path, 'processed', run_name, 'n?')
    for proc_night_path in glob.glob(glob_pattr_proc):
        night = path.basename(proc_night_path)
        night_id = int(night[1])
        logger.debug('Loading night {0}...'.format(night_id))

        observations = []
        tgas_sources = []
        prior_rvs = []

        glob_pattr_1d = path.join(proc_night_path, '1d_*.fit')
        for path_1d in ProgressBar(glob.glob(glob_pattr_1d)):
            hdr = fits.getheader(path_1d)

            # skip all except OBJECT observations
            if hdr['IMAGETYP'] != 'OBJECT':
                continue

            basename = path.basename(path_1d)[3:]
            logger.log(1, 'loading row for {0}'.format(basename))

            kw = dict()

            # construct filenames using hard-coded bullshit
            kw['filename_raw'] = basename
            kw['filename_p'] = 'p_' + basename
            kw['filename_1d'] = '1d_' + basename

            # check if this filename is already in the database, if so, drop it
            base_query = session.query(Observation)\
                                .filter(Observation.filename_raw == kw['filename_raw'])
            already_loaded = base_query.count() > 0

            if already_loaded and overwrite:
                base_query.delete()
                session.commit()

            elif already_loaded:
                logger.debug('Object {0} [{1}] already loaded'
                             .format(hdr['OBJECT'],
                                     path.basename(kw['filename_raw'])))
                continue

            # read in header of 1d file and store keywords that exist as columns
            kw.update(fits_header_to_cols(hdr, obs_columns))

            # HACK: skip empty object name
            if len(str(hdr['OBJECT'])) == 0:
                logger.warning('SKIPPING - empty OBJECT key')
                continue

            # get group id from object name
            if '-' in str(hdr['OBJECT']):
                # Per APW and SMOH's convention

                split_name = hdr['OBJECT'].split('-')
                kw['group_id'] = int(split_name[0])

                # because: reasons
                if kw['group_id'] == 10:
                    tgas_row_idx = int(split_name[1])
                else:
                    smoh_idx = int(split_name[1])
                    tgas_row_idx = ID_tbl[smoh_idx]['tgas_row']
                tgas_row = tgas[tgas_row_idx]

                # query Simbad to get all possible names for this target
                if tgas_row['hip'] > 0:
                    object_name = 'HIP{0}'.format(tgas_row['hip'])
                else:
                    object_name = 'TYC {0}'.format(tgas_row['tycho2_id'])
                logger.log(1, 'common name: {0}'.format(object_name))

                try:
                    all_ids = Simbad.query_objectids(object_name)['ID'].astype(str)
                except Exception as e:
                    logger.warning('Simbad query_objectids failed for "{0}" '
                                   'with error: {1}'
                                   .format(object_name, str(e)))
                    all_ids = []

                logger.log(1, 'this is a group object')

                if len(all_ids) > 0:
                    logger.log(1, 'other names for this object: {0}'
                               .format(', '.join(all_ids)))
                else:
                    logger.log(1, 'simbad names for this object could not be '
                               'retrieved')

            elif (isinstance(hdr['OBJECT'], int) or
                    str(hdr['OBJECT']).startswith('k') or
                    hdr['OBJECT'][0].isdigit()):
                # Assume it's a KIC number - per Ruth and Dan's convention

                if isinstance(hdr['OBJECT'], int):
                    object_name = 'KIC {0:d}'.format(hdr['OBJECT'])

                elif hdr['OBJECT'].startswith('k'):
                    object_name = 'KIC {0}'.format(hdr['OBJECT'][1:])

                else:
                    object_name = 'KIC {0}'.format(hdr['OBJECT'])

                # query Simbad to get all possible names for this target
                logger.log(1, 'common name: {0}'.format(object_name))

                try:
                    all_ids = Simbad.query_objectids(object_name)['ID'].astype(str)
                except Exception as e:
                    logger.warning('Simbad query_objectids failed for "{0}" '
                                   'with error: {1}'
                                   .format(object_name, str(e)))
                    all_ids = []

                logger.log(1, 'this is a KIC object')

                if len(all_ids) > 0:
                    logger.log(1, 'other names for this object: {0}'
                               .format(', '.join(all_ids)))
                else:
                    logger.log(1, 'simbad names for this object could not be '
                               'retrieved')

                # get the Tycho 2 ID, if it has one
                hip_id = [id_ for id_ in all_ids if 'HIP' in id_]
                tyc_id = [id_ for id_ in all_ids if 'TYC' in id_]
                if hip_id:
                    hip_id = int(hip_id[0].replace('HIP', '').strip())
                    logger.log(1, 'source has HIP id: {0}'.format(hip_id))
                    tgas_row_idx = np.where(tgas['hip'] == hip_id)[0]

                    if len(tgas_row_idx) == 0:
                        tgas_row_idx = None
                    else:
                        tgas_row = tgas[tgas_row_idx]

                elif tyc_id:
                    tyc_id = tyc_id[0].replace('TYC', '').strip()
                    logger.log(1, 'source has tycho 2 id: {0}'.format(tyc_id))
                    tgas_row_idx = np.where(tgas['tycho2_id'] == tyc_id)[0]

                    if len(tgas_row_idx) == 0:
                        tgas_row_idx = None
                    else:
                        tgas_row = tgas[tgas_row_idx]

                else:
                    logger.log(1, 'source has no HIP or TYC id.')
                    tgas_row_idx = None

                # result_table = Simbad.query_object(object_name)

            else:
                object_name = hdr['OBJECT']
                logger.log(1, 'common name: {0}'.format(object_name))
                logger.log(1, 'this is not a group object')

                # query Simbad to get all possible names for this target
                try:
                    all_ids = Simbad.query_objectids(object_name)['ID'].astype(str)
                except Exception as e:
                    logger.warning('SKIPPING: Simbad query_objectids failed for '
                                   '"{0}" with error: {1}'
                                   .format(object_name, str(e)))
                    continue

                # get the Tycho 2 ID, if it has one
                hip_id = [id_ for id_ in all_ids if 'HIP' in id_]
                tyc_id = [id_ for id_ in all_ids if 'TYC' in id_]
                if hip_id:
                    hip_id = int(hip_id[0].replace('HIP', '').strip())
                    logger.log(1, 'source has HIP id: {0}'.format(hip_id))
                    tgas_row_idx = np.where(tgas['hip'] == hip_id)[0]

                    if len(tgas_row_idx) == 0:
                        tgas_row_idx = None
                    else:
                        tgas_row = tgas[tgas_row_idx]

                elif tyc_id:
                    tyc_id = tyc_id[0].replace('TYC', '').strip()
                    logger.log(1, 'source has tycho 2 id: {0}'.format(tyc_id))
                    tgas_row_idx = np.where(tgas['tycho2_id'] == tyc_id)[0]

                    if len(tgas_row_idx) == 0:
                        tgas_row_idx = None
                    else:
                        tgas_row = tgas[tgas_row_idx]

                else:
                    logger.log(1, 'source has no tycho 2 id.')
                    tgas_row_idx = None

            # store relevant names / IDs
            simbad_info_kw = dict()
            for id_ in all_ids:
                if id_.lower().startswith('hd'):
                    simbad_info_kw['hd_id'] = id_[2:]

                elif id_.lower().startswith('hip'):
                    simbad_info_kw['hip_id'] = id_[3:]

                elif id_.lower().startswith('tyc'):
                    simbad_info_kw['tyc_id'] = id_[3:]

                elif id_.lower().startswith('2mass'):
                    simbad_info_kw['twomass_id'] = id_[5:]

            for k,v in simbad_info_kw.items():
                simbad_info_kw[k] = v.strip()

            simbad_info = SimbadInfo(**simbad_info_kw)

            # Compute barycenter velocity given coordinates of where the
            # telescope was pointing and observation time
            t = Time(hdr['JD'], format='jd', scale='utc')
            sc = coord.SkyCoord(ra=hdr['RA'], dec=hdr['DEC'],
                                unit=(u.hourangle, u.degree))
            kw['v_bary'] = bary_vel_corr(t, sc, location=kitt_peak)

            obs = Observation(night=night_id, **kw)
            obs.run = run

            # Get the TGAS data if the source is in TGAS
            if tgas_row_idx is not None:
                logger.log(1, 'TGAS row: {0}'.format(tgas_row_idx))

                tgas_kw = dict()
                tgas_kw['row_index'] = tgas_row_idx
                for name in tgas.colnames:
                    if name in tgassource_columns:
                        tgas_kw[name] = tgas_row[name]

                job = Gaia.launch_job(gaia_query.format(tgas_kw['source_id'][0]),
                                      dump_to_file=False)
                res = job.get_results()

                if len(res) == 0:
                    logger.warning("No 2MASS data found for: {0}"
                                   .format(tgas_kw['source_id']))

                elif len(res) == 1:
                    tgas_kw['J'] = res['j_m'][0]
                    tgas_kw['J_err'] = res['j_msigcom'][0]
                    tgas_kw['H'] = res['h_m'][0]
                    tgas_kw['H_err'] = res['h_msigcom'][0]
                    tgas_kw['Ks'] = res['ks_m'][0]
                    tgas_kw['Ks_err'] = res['ks_msigcom'][0]

                tgas_source = TGASSource(**tgas_kw)
                tgas_sources.append(tgas_source)

                obs.tgas_source = tgas_source

            else:
                logger.log(1, 'TGAS row could not be found.')

            obs.simbad_info = simbad_info
            observations.append(obs)

            # retrieve a previous measurement from the literature
            result = get_best_rv(obs)
            if result is not None:
                rv, rv_err, rv_qual, rv_bibcode, rv_source = result

                prv = PriorRV(rv=rv*u.km/u.s, err=rv_err*u.km/u.s,
                              qual=rv_qual, bibcode=rv_bibcode,
                              source=rv_source)
                obs.prior_rv = prv
                prior_rvs.append(prv)

            logger.log(1, '-'*68)

        session.add_all(observations)
        session.add_all(tgas_sources)
        session.add_all(prior_rvs)
        session.commit()

    # Last thing to do is cross-match with the Bensby catalog to
    #   replace velocities when they are better
    for sim_info in session.query(SimbadInfo)\
                           .filter(SimbadInfo.hip_id != None).all():
        hip_id = 'HIP' + str(sim_info.hip_id)
        row = bensby[bensby['OBJECT'] == hip_id]
        if len(row) > 0:
            sim_info.rv = row['velValue']
            sim_info.rv_qual = row['quality']
            sim_info.rv_bibcode = row['bibcode']
            session.flush()

    session.close()
def main(db_path,
         run_name,
         data_root_path=None,
         filename=None,
         overwrite=False,
         pool=None):

    if pool is None:
        pool = schwimmbad.SerialPool()

    # connect to the database
    engine = db_connect(db_path)
    # engine.echo = True
    logger.debug("Connected to database at '{}'".format(db_path))

    # create a new session for interacting with the database
    session = Session()

    root_path, _ = path.split(db_path)
    if data_root_path is None:
        data_root_path = root_path

    plot_path = path.join(root_path, 'plots', run_name)
    if not path.exists(plot_path):
        os.makedirs(plot_path, exist_ok=True)

    # TODO: there might be some bugs here...
    n_lines = session.query(SpectralLineInfo).count()
    Halpha = session.query(SpectralLineInfo)\
                    .filter(SpectralLineInfo.name == 'Halpha').one()
    OI_lines = session.query(SpectralLineInfo)\
                      .filter(SpectralLineInfo.name.contains('[OI]')).all()

    if filename is None:  # grab all unfinished sources
        observations = session.query(Observation).join(Run)\
                              .filter(Run.name == run_name).all()

    else:  # only process the observation corresponding to this filename
        observations = session.query(Observation).join(Run)\
                              .filter(Run.name == run_name)\
                              .filter(Observation.filename_raw == filename).all()

    for obs in observations:
        measurements = session.query(SpectralLineMeasurement)\
                              .join(Observation)\
                              .filter(Observation.id == obs.id).all()

        if len(measurements) == n_lines and not overwrite:
            logger.debug('All line measurements already complete for object '
                         '{0} in file {1}'.format(obs.object,
                                                  obs.filename_raw))
            continue

        # Read the spectrum data and get wavelength solution
        filebase, _ = path.splitext(obs.filename_1d)
        filename_1d = obs.path_1d(data_root_path)
        spec = Table.read(filename_1d)
        logger.debug('Loaded 1D spectrum for object {0} from file {1}'.format(
            obs.object, filename_1d))

        # Extract region around Halpha
        x, (flux, ivar) = extract_region(
            spec['wavelength'],
            center=Halpha.wavelength.value,
            width=100,
            arrs=[spec['source_flux'], spec['source_ivar']])

        # We start by doing maximum likelihood estimation to fit the line, then
        # use the best-fit parameters to initialize an MCMC run.
        # TODO: need to figure out if it's emission or absorption...for now just
        #   assume absorption
        absorp_emiss = -1.
        lf = VoigtLineFitter(x, flux, ivar, absorp_emiss=absorp_emiss)
        lf.fit()
        fit_pars = lf.get_gp_mean_pars()

        if (not lf.success
                or abs(fit_pars['x0'] - Halpha.wavelength.value) > 16.
                or  # 16 Å = ~700 km/s
                abs(fit_pars['amp']) < 10):  # minimum amplitude - MAGIC NUMBER
            # TODO: should try again with emission line
            logger.error('absorption line has tiny amplitude! did '
                         'auto-determination of absorption/emission fail?')
            # TODO: what now?
            continue

        fig = lf.plot_fit()
        fig.savefig(path.join(plot_path, '{}_maxlike.png'.format(filebase)),
                    dpi=256)
        plt.close(fig)

        # ----------------------------------------------------------------------

        # Run `emcee` instead to sample over GP model parameters:
        if fit_pars['std_G'] < 1E-2:
            lf.gp.freeze_parameter('mean:ln_std_G')

        initial = np.array(lf.gp.get_parameter_vector())
        if initial[4] < -10:  # TODO: ???
            initial[4] = -8.
        if initial[5] < -10:  # TODO: ???
            initial[5] = -8.
        ndim, nwalkers = len(initial), 64

        sampler = emcee.EnsembleSampler(nwalkers,
                                        ndim,
                                        log_probability,
                                        pool=pool,
                                        args=(lf.gp, flux))

        logger.debug("Running burn-in...")
        p0 = initial + 1e-6 * np.random.randn(nwalkers, ndim)
        p0, lp, _ = sampler.run_mcmc(p0, 128)

        logger.debug("Running 2nd burn-in...")
        sampler.reset()
        p0 = p0[lp.argmax()] + 1e-3 * np.random.randn(nwalkers, ndim)
        p0, lp, _ = sampler.run_mcmc(p0, 512)

        logger.debug("Running production...")
        sampler.reset()
        pos, lp, _ = sampler.run_mcmc(p0, 1024)

        fit_kw = dict()
        for i, par_name in enumerate(lf.gp.get_parameter_names()):
            if 'kernel' in par_name: continue

            # remove 'mean:'
            par_name = par_name[5:]

            # skip bg
            if par_name.startswith('bg'): continue

            samples = sampler.flatchain[:, i]

            if par_name.startswith('ln_'):
                par_name = par_name[3:]
                samples = np.exp(samples)

            MAD = np.median(np.abs(samples - np.median(samples)))
            fit_kw[par_name] = np.median(samples)
            fit_kw[par_name + '_error'] = 1.5 * MAD  # convert to ~stddev

        # remove all previous line measurements
        q = session.query(SpectralLineMeasurement).join(Observation)\
                   .filter(Observation.id == obs.id)
        if q.count() > 0:
            for meas in q.all():
                session.delete(meas)
            session.commit()

        slm = SpectralLineMeasurement(**fit_kw)
        slm.info = Halpha
        slm.observation = obs
        session.add(slm)
        session.commit()

        # --------------------------------------------------------------------
        # plot MCMC traces
        fig, axes = plt.subplots(2, 4, figsize=(18, 6))
        for i in range(sampler.dim):
            for walker in sampler.chain[..., i]:
                axes.flat[i].plot(walker,
                                  marker='',
                                  drawstyle='steps-mid',
                                  alpha=0.2)
            axes.flat[i].set_title(lf.gp.get_parameter_names()[i], fontsize=12)
        fig.tight_layout()
        fig.savefig(path.join(plot_path, '{}_mcmc_trace.png'.format(filebase)),
                    dpi=256)
        plt.close(fig)
        # --------------------------------------------------------------------

        # --------------------------------------------------------------------
        # plot samples
        fig, axes = plt.subplots(3, 1, figsize=(10, 10), sharex=True)

        samples = sampler.flatchain
        for s in samples[np.random.randint(len(samples), size=32)]:
            lf.gp.set_parameter_vector(s)
            lf.plot_fit(axes=axes, fit_alpha=0.2)

        fig.tight_layout()
        fig.savefig(path.join(plot_path, '{}_mcmc_fits.png'.format(filebase)),
                    dpi=256)
        plt.close(fig)
        # --------------------------------------------------------------------

        # --------------------------------------------------------------------
        # corner plot
        fig = corner.corner(
            sampler.flatchain[::10, :],
            labels=[x.split(':')[1] for x in lf.gp.get_parameter_names()])
        fig.savefig(path.join(plot_path, '{}_corner.png'.format(filebase)),
                    dpi=256)
        plt.close(fig)
        # --------------------------------------------------------------------

        # compute centroids for sky lines
        sky_centroids = []
        for j, sky_line in enumerate(OI_lines):
            wvln = sky_line.wavelength.value
            x, (flux, ivar) = extract_region(
                spec['wavelength'],
                center=wvln,
                width=32.,  # angstroms
                arrs=[spec['background_flux'], spec['background_ivar']])

            lf = GaussianLineFitter(x, flux, ivar,
                                    absorp_emiss=1.)  # all emission lines

            try:
                lf.fit()
                fit_pars = lf.get_gp_mean_pars()

            except Exception as e:
                logger.warn("Failed to fit sky line {0}:\n{1}".format(
                    sky_line, e))
                lf.success = False
                fit_pars = lf.get_init()
                # OMG this is the biggest effing hack
                fit_pars['amp'] = 0.
                fit_pars['bg_coef'] = None
                fit_pars['x0'] = 0.

            # HACK: hackish signal-to-noise
            max_ = fit_pars['amp'] / np.sqrt(2 * np.pi * fit_pars['std']**2)
            SNR = max_ / np.median(1 / np.sqrt(ivar))

            if (not lf.success or abs(fit_pars['x0'] - wvln) > 4
                    or fit_pars['amp'] < 10 or fit_pars['std'] > 4
                    or SNR < 2.5):
                # failed
                x0 = np.nan * u.angstrom
                title = 'f****d'
                fit_pars['amp'] = 0.

            else:
                x0 = fit_pars['x0'] * u.angstrom
                title = '{:.2f}'.format(fit_pars['amp'])

            if lf.success:
                fig = lf.plot_fit()
                fig.suptitle(title, y=0.95)
                fig.subplots_adjust(top=0.8)
                fig.savefig(path.join(
                    plot_path,
                    '{}_maxlike_sky_{:.0f}.png'.format(filebase, wvln)),
                            dpi=256)
                plt.close(fig)

            # store the sky line measurements
            fit_pars['std_G'] = fit_pars.pop('std')  # HACK
            fit_pars.pop('bg_coef')  # HACK
            slm = SpectralLineMeasurement(**fit_pars)
            slm.info = sky_line
            slm.observation = obs
            session.add(slm)
            session.commit()

            sky_centroids.append(x0)
        sky_centroids = u.Quantity(sky_centroids)

        logger.info('{} [{}]: x0={x0:.3f} σ={err:.3f}\n--------'.format(
            obs.object, filebase, x0=fit_kw['x0'], err=fit_kw['x0_error']))

        session.commit()

    pool.close()
示例#7
0
def main():
    # TODO: bad, hard-coded...
    # base_path = '/Volumes/ProjectData/gaia-comoving-followup/'
    base_path = '../../data/'
    db_path = path.join(base_path, 'db.sqlite')
    engine = db_connect(db_path)
    session = Session()

    chain_path = path.abspath('./isochrone_chains')
    os.makedirs(chain_path, exist_ok=True)

    # Check out the bottom of "Color-magnitude diagram.ipynb":
    interesting_group_ids = [1500, 1229, 1515]

    all_photometry = OrderedDict([
        ('1500-8455',
         OrderedDict([('J', (6.8379998, 0.021)),
                      ('H', (6.4640002, 0.017000001)),
                      ('K', (6.3369999, 0.017999999)),
                      ('W1', (6.2950001, 0.093000002)),
                      ('W2', (6.2490001, 0.026000001)),
                      ('W3', (6.3330002, 0.015)), ('B', (9.5950003, 0.022)),
                      ('V', (8.5120001, 0.014))])),
        ('1500-1804',
         OrderedDict([('J', (6.9039998, 0.041000001)),
                      ('H', (6.8559999, 0.027000001)),
                      ('K', (6.7989998, 0.017000001)),
                      ('W1', (6.803, 0.064999998)),
                      ('W2', (6.7600002, 0.018999999)),
                      ('W3', (6.8270001, 0.016000001)),
                      ('B', (7.4980001, 0.015)), ('V', (7.289, 0.011))])),
        ('1229-1366',
         OrderedDict([('J', (6.7290001, 0.024)), ('H', (6.2449999, 0.02)),
                      ('K', (6.1529999, 0.023)),
                      ('W1', (6.1799998, 0.096000001)), ('W2', (6.04, 0.035)),
                      ('W3', (6.132, 0.016000001)), ('B', (9.5539999, 0.021)),
                      ('V', (8.4619999, 0.014))])),
        ('1229-7470',
         OrderedDict([
             ('J', (9.1709995, 0.024)), ('H', (8.7959995, 0.026000001)),
             ('K', (8.7299995, 0.022)), ('W1', (8.6669998, 0.023)),
             ('W2', (8.7189999, 0.02)), ('W3', (8.6680002, 0.025)),
             ('B', (11.428, 0.054000001)), ('V', (10.614, 0.039999999))
         ])),
        ('1515-3584',
         OrderedDict([('J', (5.363999843597412, 0.024000000208616257)),
                      ('H', (4.965000152587891, 0.035999998450279236)),
                      ('K', (4.815999984741211, 0.032999999821186066)),
                      ('W1', (4.758, 0.215)), ('W2', (4.565, 0.115)),
                      ('W3', (4.771, 0.015)),
                      ('B', (8.347999572753906, 0.01600000075995922)),
                      ('V', (7.182000160217285, 0.009999999776482582))])),
        ('1515-1834',
         OrderedDict([('J', (8.855999946594238, 0.024000000208616257)),
                      ('H', (8.29699993133545, 0.020999999716877937)),
                      ('K', (8.178999900817871, 0.017999999225139618)),
                      ('W1', (8.117, 0.022)), ('W2', (8.15, 0.019)),
                      ('W3', (8.065, 0.02)),
                      ('B', (12.309000015258789, 0.11999999731779099)),
                      ('V', (11.069999694824219, 0.054999999701976776))]))
    ])

    for k in all_photometry:
        samples_file = path.join(chain_path, '{0}.hdf5'.format(k))

        if path.exists(samples_file):
            logger.info("skipping {0} - samples exist at {1}".format(
                k, samples_file))
            continue

        phot = all_photometry[k]
        obs = session.query(Observation).filter(Observation.object == k).one()
        plx = (obs.tgas_source.parallax, obs.tgas_source.parallax_error)

        # fit an isochrone
        model = StarModel(iso, use_emcee=True, parallax=plx, **phot)
        model.set_bounds(mass=(0.01, 20),
                         feh=(-1, 1),
                         distance=(0, 300),
                         AV=(0, 1))

        # initial conditions for emcee walkers
        nwalkers = 128

        p0 = []
        m0, age0, feh0 = model.ic.random_points(nwalkers,
                                                minmass=0.01,
                                                maxmass=10.,
                                                minfeh=-1,
                                                maxfeh=1)
        _, max_distance = model.bounds('distance')
        _, max_AV = model.bounds('AV')
        d0 = 10**(np.random.uniform(0, np.log10(max_distance), size=nwalkers))
        AV0 = np.random.uniform(0, max_AV, size=nwalkers)
        p0 += [m0]
        p0 += [age0, feh0, d0, AV0]

        p0 = np.array(p0).T
        npars = p0.shape[1]

        # run emcee
        ninit = 256
        nburn = 1024
        niter = 4096

        logger.debug('Running emcee - initial sampling...')
        sampler = emcee.EnsembleSampler(nwalkers, npars, model.lnpost)
        # pos, prob, state = sampler.run_mcmc(p0, ninit)

        for pos, prob, state in tqdm(sampler.sample(p0, iterations=ninit),
                                     total=ninit):
            pass

        # cull the weak walkers
        best_ix = sampler.flatlnprobability.argmax()
        best_p0 = (sampler.flatchain[best_ix][None] +
                   np.random.normal(0, 1E-5, size=(nwalkers, npars)))

        sampler.reset()
        logger.debug('burn-in...')
        for pos, prob, state in tqdm(sampler.sample(best_p0, iterations=nburn),
                                     total=nburn):
            pass
        # pos,_,_ = sampler.run_mcmc(best_p0, nburn)

        sampler.reset()
        logger.debug('sampling...')
        # _ = sampler.run_mcmc(pos, niter)
        for pos, prob, state in tqdm(sampler.sample(pos, iterations=niter),
                                     total=niter):
            pass

        model._sampler = sampler
        model._make_samples(0.08)

        model.samples.to_hdf(samples_file, key='samples')
        # np.save('isochrone_chains/chain.npy', sampler.chain)
        logger.debug('...done and saved!')