<commit_before>#!/usr/bin/env python

from __future__ import print_function, division

import numpy as np

from legacypipe.internal import sharedmem
from legacypipe.internal.io import iter_tractor, parse_filename, get_units, git_version

import argparse
import os, sys
from time import time
from scipy.spatial import cKDTree as KDTree

import fitsio
import platform

print('Running from %s' % platform.node())

def main():
    ns = parse_args()

    if ns.ignore_errors:
        print("Warning: *** Will ignore broken tractor catalog files ***")
        print("         *** Disable -I for final data product.         ***")

    bricks = list_bricks(ns)

    # ADM grab a {FIELD: unit} dict from the first Tractor file.
    unitdict = get_units(bricks[0][1])

    tree, nobj, morecols = read_external(ns.external, ns)

    # get the data type of the match
    brickname, path = bricks[0]
    peek = fitsio.read(path, 1, upper=True)
    matched_catalog = sharedmem.empty(nobj, dtype=peek.dtype)
    matched_catalog['OBJID'] = -1

    matched_distance = sharedmem.empty(nobj, dtype='f4')

    # convert to radian
    tol = ns.tolerance / (60. * 60.)  * (np.pi / 180)

    matched_distance[:] = tol
    nprocessed = np.zeros((), dtype='i8')
    nmatched = np.zeros((), dtype='i8')
    ntotal = np.zeros((), dtype='i8')
    t0 = time()

    with sharedmem.MapReduce(np=ns.numproc) as pool:
        def work(brickname, path):
            try:
                objects = fitsio.read(path, 1, upper=True)
            except:
                if ns.ignore_errors:
                    print ("IO Error on %s" %path)
                    return None, None, None
                else:
                    raise

            pos = radec2pos(objects['RA'], objects['DEC'])
            d, i = tree.query(pos, 1)
            assert (objects['OBJID'] != -1).all()
            with pool.critical:
                mask = d < matched_distance[i]
                mask &= objects['BRICK_PRIMARY']
                i = i[mask]
                matched_catalog[i] = objects[mask][list(matched_catalog.dtype.names)]
                matched_distance[i] = d[mask]
            matched = mask.sum()

            return brickname, matched, len(objects)

        def reduce(brickname, matched, total):
            if brickname is None:
                return
            nprocessed[...] += 1
            nmatched[...] += matched
            ntotal[...] += total
            if ns.verbose:
                if nprocessed % 1000 == 0:
                    print("Processed %d files, %g / second, matched %d / %d objects."
                        % (nprocessed, nprocessed / (time() - t0), nmatched, ntotal)
                        )

        pool.map(work, bricks, star=True, reduce=reduce)

        nrealmatched = (matched_catalog['OBJID'] != -1).sum()
        if ns.verbose:
            print("Processed %d files, %g / second, matched %d / %d objects into %d slots."
                % (nprocessed, nprocessed / (time() - t0),
                    nmatched, ntotal,
                    nrealmatched)
                )

        try:
            os.makedirs(os.path.dirname(ns.dest))
        except OSError:
            pass

        hdr = fitsio.FITSHDR()
        hdr.add_record(dict(name='NMATCHED', value=nrealmatched,
                            comment='Number of unique matches.'))
        hdr.add_record(dict(name='NCOLL', value=nmatched - nrealmatched,
                            comment='Total number of matches.'))
        hdr.add_record(dict(name='NCOLL', value=nrealmatched,
                            comment='Total number of matches.'))
        hdr.add_record(dict(name='RADIUS', value=ns.tolerance,
                            comment='Search radius (arcsec).'))
        value = ns.external
        if len(value) > 68:
            hdr.add_record(dict(name='EXTERNAL', value=value[:67]+'&'))
            while len(value):
                value = value[67:]
                if len(value) == 0:
                    break
                hdr.add_record(dict(name='CONTINUE', value="  '%s%s'" % (
                    value[:67], '&' if len(value) > 67 else '')))
            added_long = True
        else:
            added_long = False

        if added_long:
            hdr.add_record(dict(name='LONGSTRN', value='OGIP 1.0',
                                comment='CONTINUE cards are used'))

        # Optionally add the new columns
        if len(morecols) > 0:
            newdtype = matched_catalog.dtype.descr

            for coldata, col in zip( morecols, ns.copycols ):
                newdtype = newdtype + [(col, coldata.dtype)]
            newdtype = np.dtype(newdtype)

            _matched_catalog = np.empty(matched_catalog.shape, dtype=newdtype)
            for field in matched_catalog.dtype.fields:
                _matched_catalog[field] = matched_catalog[field]
            for coldata, col in zip( morecols, ns.copycols ):
                _matched_catalog[col] = coldata

            matched_catalog = _matched_catalog.copy()
            del _matched_catalog

        for format in ns.format:
            save_file(ns.dest, matched_catalog, hdr, format, unitdict=unitdict)

def save_file(filename, data, header, format, unitdict=None):
    basename = os.path.splitext(filename)[0]
    if format == 'fits':
        units = None
	# ADM derive the units from the data columns if possible.
        if unitdict is not None:
            # ADM some columns from external-match files might not have
            # ADM units, so pass an empty string for external columns.
            units = [unitdict[col] if col in unitdict.keys() else ""
                     for col in data.dtype.names]

        # ADM add the external match code version header dependency.
        dep = [int(key.split("DEPNAM")[-1]) for key in header.keys()
               if 'DEPNAM' in key]
        if len(dep) == 0:
            nextdep = 0
        else:
            nextdep = np.max(dep) + 1
        header["DEPNAM{:02d}".format(nextdep)] = 'match_external'
        header["DEPVER{:02d}".format(nextdep)] = git_version()

        filename = basename + '.fits'
        fitsio.write(filename, data, extname='MATCHED', header=header,
                     clobber=True, units=units)
    elif format == 'hdf5':
        filename = basename + '.hdf5'
        import h5py
        with h5py.File(filename, 'w') as ff:
            dset = ff.create_dataset('MATCHED', data=data)
            for key in header:
                dset.attrs[key] = header[key]
    else:
        raise ValueError("Unknown format")

def radec2pos(ra, dec):
    pos = np.empty(len(ra), ('f4', 3))
    pos[:, 2] = np.sin(dec / 180. * np.pi)
    pos[:, 1] = np.cos(dec / 180. * np.pi)
    pos[:, 0] = pos[:, 1]
    pos[:, 0] *= np.sin(ra / 180. * np.pi)
    pos[:, 1] *= np.cos(ra / 180. * np.pi)
    return pos

def read_external(filename, ns):
    t0 = time()
    cat = fitsio.FITS(filename, upper=True)[1][:]

    if ns.verbose:
        print("reading external catalog took %g seconds." % (time() - t0))
        print("%d objects." % len(cat))

    t0 = time()
    for raname, decname in [
            ('RA', 'DEC'),
            ('PLUG_RA', 'PLUG_DEC')
            ]:
        if raname in cat.dtype.names \
        and decname in cat.dtype.names:
            ra = cat[raname]
            dec = cat[decname]
            if ns.verbose:
                print('using %s/%s for positions.' % (raname, decname))
            break
    else:
        raise KeyError("No RA/DEC or PLUG_RA/PLUG_DEC in the external catalog")

    pos = radec2pos(ra, dec)
    # work around NERSC overcommit issue.
    pos = sharedmem.copy(pos)

    tree = KDTree(pos)

    if ns.verbose:
        print("Building KD-Tree took %g seconds." % (time() - t0))

    morecols = []
    if ns.copycols is not None:
        for col in np.atleast_1d(ns.copycols):
            if col not in cat.dtype.names:
                print('Column {} does not exist in external catalog!'.format(col))
                raise IOError
            morecols.append(cat[col])

    return tree, len(cat), morecols

def list_bricks(ns):
    t0 = time()

    if ns.filelist is not None:
        d = dict([(parse_filename(fn.strip()), fn.strip())
            for fn in open(ns.filelist, 'r').readlines()])
    else:
        d = dict(iter_tractor(ns.src))

    if ns.verbose:
        print('enumerated %d bricks in %g seconds' % (
            len(d), time() - t0))

    #- Load list of bricknames to use
    if ns.bricklist is not None:
        bricklist = np.loadtxt(ns.bricklist, dtype='S8')
        # TODO: skip unknown bricks?
        d = dict([(brickname.decode(), d[brickname])
                             for brickname in bricklist])

    t0 = time()

    bricks = sorted(d.items())

    return bricks

def parse_args():
    ap = argparse.ArgumentParser(
    description="""Match to an external catalogs.
        """
        )

    ap.add_argument("external", help="External catalog. e.g. /global/project/projectdirs/cosmo/work/sdss/cats/specObj-dr12.fits")
    ap.add_argument("src", help="Path to the root directory of all tractor files")
    ap.add_argument("dest", help="Path to the output file")

    ap.add_argument('-f', "--format", choices=['fits', 'hdf5'], nargs='+', default=["fits"],
        help="Format of the output file")

    ap.add_argument('-t', "--tolerance", default=1.5, type=float,
        help="Tolerance of the angular distance for a match, in arcseconds")

    ap.add_argument('-F', "--filelist", default=None,
        help="list of tractor brickfiles to use; this will avoid expensive walking of the path.")

    ap.add_argument('-b', "--bricklist",
        help="""Filename with list of bricknames to include.
                If not set, all bricks in src are included, sorted by brickname.
            """)

    ap.add_argument('-v', "--verbose", action='store_true')

    ap.add_argument('-I', "--ignore-errors", action='store_true')

    ap.add_argument("--numproc", type=int, default=None,
        help="""Number of concurrent processes to use. 0 for sequential execution.
            Default is to use OMP_NUM_THREADS, or the number of cores on the node.""")

    ap.add_argument("--copycols", nargs='*', help="List of columns to copy from external to matched output catalog (e.g., MJD, FIBER, PLATE)", default=None)

    return ap.parse_args()

if __name__ == '__main__':
    main()
<commit_msg>begin to flip the sense of the match of the imaging spectro trees<commit_after>#!/usr/bin/env python

from __future__ import print_function, division

import numpy as np

from legacypipe.internal import sharedmem
from legacypipe.internal.io import iter_tractor, parse_filename, get_units, git_version

import argparse
import os, sys
from time import time
from scipy.spatial import cKDTree as KDTree

import fitsio
import platform

print('Running from %s' % platform.node())

def main():
    ns = parse_args()

    if ns.ignore_errors:
        print("Warning: *** Will ignore broken tractor catalog files ***")
        print("         *** Disable -I for final data product.         ***")

    bricks = list_bricks(ns)

    # ADM grab a {FIELD: unit} dict from the first Tractor file.
    unitdict = get_units(bricks[0][1])

    tree, nobj, morecols = read_external(ns.external, ns)

    # get the data type of the match
    brickname, path = bricks[0]
    peek = fitsio.read(path, 1, upper=True)
    matched_catalog = sharedmem.empty(nobj, dtype=peek.dtype)
    matched_catalog['OBJID'] = -1

    matched_distance = sharedmem.empty(nobj, dtype='f4')

    # convert to radian
    tol = ns.tolerance / (60. * 60.)  * (np.pi / 180)

    matched_distance[:] = tol
    nprocessed = np.zeros((), dtype='i8')
    nmatched = np.zeros((), dtype='i8')
    ntotal = np.zeros((), dtype='i8')
    t0 = time()

    with sharedmem.MapReduce(np=ns.numproc) as pool:
        def work(brickname, path):
            try:
                objects = fitsio.read(path, 1, upper=True)
            except:
                if ns.ignore_errors:
                    print ("IO Error on %s" %path)
                    return None, None, None
                else:
                    raise

            # ADM limit to just PRIMARY objects from imaging.
            ii = objects["BRICK_PRIMARY"]
            objects = objects[ii]

            # ADM build a tree for the imaging objects.
            pos = radec2pos(objects['RA'], objects['DEC'])
            imtree = KDTree(pos)

            # ADM match objects in imaging to objects in spectroscopy so
            # ADM every spectroscopic object finds its nearest match
            d, i = imtree.query(tree.data, 1)
            assert (objects['OBJID'] != -1).all()
            with pool.critical:
                mask = d < matched_distance[i]
                mask &= objects['BRICK_PRIMARY']
                i = i[mask]
                matched_catalog[i] = objects[mask][list(matched_catalog.dtype.names)]
                matched_distance[i] = d[mask]
            matched = mask.sum()

            return brickname, matched, len(objects)

        def reduce(brickname, matched, total):
            if brickname is None:
                return
            nprocessed[...] += 1
            nmatched[...] += matched
            ntotal[...] += total
            if ns.verbose:
                if nprocessed % 1000 == 0:
                    print("Processed %d files, %g / second, matched %d / %d objects."
                        % (nprocessed, nprocessed / (time() - t0), nmatched, ntotal)
                        )

        pool.map(work, bricks, star=True, reduce=reduce)

        nrealmatched = (matched_catalog['OBJID'] != -1).sum()
        if ns.verbose:
            print("Processed %d files, %g / second, matched %d / %d objects into %d slots."
                % (nprocessed, nprocessed / (time() - t0),
                    nmatched, ntotal,
                    nrealmatched)
                )

        try:
            os.makedirs(os.path.dirname(ns.dest))
        except OSError:
            pass

        hdr = fitsio.FITSHDR()
        hdr.add_record(dict(name='NMATCHED', value=nrealmatched,
                            comment='Number of unique matches.'))
        hdr.add_record(dict(name='NCOLL', value=nmatched - nrealmatched,
                            comment='Total number of matches.'))
        hdr.add_record(dict(name='NCOLL', value=nrealmatched,
                            comment='Total number of matches.'))
        hdr.add_record(dict(name='RADIUS', value=ns.tolerance,
                            comment='Search radius (arcsec).'))
        value = ns.external
        if len(value) > 68:
            hdr.add_record(dict(name='EXTERNAL', value=value[:67]+'&'))
            while len(value):
                value = value[67:]
                if len(value) == 0:
                    break
                hdr.add_record(dict(name='CONTINUE', value="  '%s%s'" % (
                    value[:67], '&' if len(value) > 67 else '')))
            added_long = True
        else:
            added_long = False

        if added_long:
            hdr.add_record(dict(name='LONGSTRN', value='OGIP 1.0',
                                comment='CONTINUE cards are used'))

        # Optionally add the new columns
        if len(morecols) > 0:
            newdtype = matched_catalog.dtype.descr

            for coldata, col in zip( morecols, ns.copycols ):
                newdtype = newdtype + [(col, coldata.dtype)]
            newdtype = np.dtype(newdtype)

            _matched_catalog = np.empty(matched_catalog.shape, dtype=newdtype)
            for field in matched_catalog.dtype.fields:
                _matched_catalog[field] = matched_catalog[field]
            for coldata, col in zip( morecols, ns.copycols ):
                _matched_catalog[col] = coldata

            matched_catalog = _matched_catalog.copy()
            del _matched_catalog

        for format in ns.format:
            save_file(ns.dest, matched_catalog, hdr, format, unitdict=unitdict)

def save_file(filename, data, header, format, unitdict=None):
    basename = os.path.splitext(filename)[0]
    if format == 'fits':
        units = None
	# ADM derive the units from the data columns if possible.
        if unitdict is not None:
            # ADM some columns from external-match files might not have
            # ADM units, so pass an empty string for external columns.
            units = [unitdict[col] if col in unitdict.keys() else ""
                     for col in data.dtype.names]

        # ADM add the external match code version header dependency.
        dep = [int(key.split("DEPNAM")[-1]) for key in header.keys()
               if 'DEPNAM' in key]
        if len(dep) == 0:
            nextdep = 0
        else:
            nextdep = np.max(dep) + 1
        header["DEPNAM{:02d}".format(nextdep)] = 'match_external'
        header["DEPVER{:02d}".format(nextdep)] = git_version()

        filename = basename + '.fits'
        fitsio.write(filename, data, extname='MATCHED', header=header,
                     clobber=True, units=units)
    elif format == 'hdf5':
        filename = basename + '.hdf5'
        import h5py
        with h5py.File(filename, 'w') as ff:
            dset = ff.create_dataset('MATCHED', data=data)
            for key in header:
                dset.attrs[key] = header[key]
    else:
        raise ValueError("Unknown format")

def radec2pos(ra, dec):
    pos = np.empty(len(ra), ('f4', 3))
    pos[:, 2] = np.sin(dec / 180. * np.pi)
    pos[:, 1] = np.cos(dec / 180. * np.pi)
    pos[:, 0] = pos[:, 1]
    pos[:, 0] *= np.sin(ra / 180. * np.pi)
    pos[:, 1] *= np.cos(ra / 180. * np.pi)
    return pos

def read_external(filename, ns):
    t0 = time()
    cat = fitsio.FITS(filename, upper=True)[1][:]

    if ns.verbose:
        print("reading external catalog took %g seconds." % (time() - t0))
        print("%d objects." % len(cat))

    t0 = time()
    for raname, decname in [
            ('RA', 'DEC'),
            ('PLUG_RA', 'PLUG_DEC')
            ]:
        if raname in cat.dtype.names \
        and decname in cat.dtype.names:
            ra = cat[raname]
            dec = cat[decname]
            if ns.verbose:
                print('using %s/%s for positions.' % (raname, decname))
            break
    else:
        raise KeyError("No RA/DEC or PLUG_RA/PLUG_DEC in the external catalog")

    pos = radec2pos(ra, dec)
    # work around NERSC overcommit issue.
    pos = sharedmem.copy(pos)

    tree = KDTree(pos)

    if ns.verbose:
        print("Building KD-Tree took %g seconds." % (time() - t0))

    morecols = []
    if ns.copycols is not None:
        for col in np.atleast_1d(ns.copycols):
            if col not in cat.dtype.names:
                print('Column {} does not exist in external catalog!'.format(col))
                raise IOError
            morecols.append(cat[col])

    return tree, len(cat), morecols

def list_bricks(ns):
    t0 = time()

    if ns.filelist is not None:
        d = dict([(parse_filename(fn.strip()), fn.strip())
            for fn in open(ns.filelist, 'r').readlines()])
    else:
        d = dict(iter_tractor(ns.src))

    if ns.verbose:
        print('enumerated %d bricks in %g seconds' % (
            len(d), time() - t0))

    #- Load list of bricknames to use
    if ns.bricklist is not None:
        bricklist = np.loadtxt(ns.bricklist, dtype='S8')
        # TODO: skip unknown bricks?
        d = dict([(brickname.decode(), d[brickname])
                             for brickname in bricklist])

    t0 = time()

    bricks = sorted(d.items())

    return bricks

def parse_args():
    ap = argparse.ArgumentParser(
    description="""Match to an external catalogs.
        """
        )

    ap.add_argument("external", help="External catalog. e.g. /global/project/projectdirs/cosmo/work/sdss/cats/specObj-dr12.fits")
    ap.add_argument("src", help="Path to the root directory of all tractor files")
    ap.add_argument("dest", help="Path to the output file")

    ap.add_argument('-f', "--format", choices=['fits', 'hdf5'], nargs='+', default=["fits"],
        help="Format of the output file")

    ap.add_argument('-t', "--tolerance", default=1.5, type=float,
        help="Tolerance of the angular distance for a match, in arcseconds")

    ap.add_argument('-F', "--filelist", default=None,
        help="list of tractor brickfiles to use; this will avoid expensive walking of the path.")

    ap.add_argument('-b', "--bricklist",
        help="""Filename with list of bricknames to include.
                If not set, all bricks in src are included, sorted by brickname.
            """)

    ap.add_argument('-v', "--verbose", action='store_true')

    ap.add_argument('-I', "--ignore-errors", action='store_true')

    ap.add_argument("--numproc", type=int, default=None,
        help="""Number of concurrent processes to use. 0 for sequential execution.
            Default is to use OMP_NUM_THREADS, or the number of cores on the node.""")

    ap.add_argument("--copycols", nargs='*', help="List of columns to copy from external to matched output catalog (e.g., MJD, FIBER, PLATE)", default=None)

    return ap.parse_args()

if __name__ == '__main__':
    main()
<|endoftext|>


<commit_before># coding: utf8
from __future__ import unicode_literals
import os
import io
import sys

from nose.tools import *
from mock import patch, Mock
import newick
from clldutils.path import Path

from beastling.configuration import (
    Configuration, get_glottolog_data, _BEAST_MAX_LENGTH,
)
from beastling.beastxml import BeastXml
from .util import WithConfigAndTempDir, config_path

def check_lat_lon(provided, target_lat, target_lon):
    prov_lat, prov_lon = provided
    return round(prov_lat, 2) == target_lat and round(prov_lon, 2) == target_lon

class Tests(WithConfigAndTempDir):
    def _make_cfg(self, *names):
        return self.make_cfg([config_path(name).as_posix() for name in names])

    def _make_bad_cfg(self, name):
        return self.make_cfg(config_path(name, bad=True).as_posix())

    def test_get_glottolog_geo(self):
        geodata = self.tmp.joinpath('glottolog-2.5-geo.csv')
        with geodata.open('w', encoding='utf8') as fp:
            fp.write('x')

        with patch(
                'beastling.configuration.user_data_dir',
                new=Mock(return_value=self.tmp.as_posix())):
            self.assertEqual(Path(get_glottolog_data('geo', '2.5')), geodata)

    def test_get_glottolog_newick(self):
        with self.tmp.joinpath('glottolog-2.5.newick').open('w', encoding='utf8') as fp:
            fp.write('(B [abcd1234],C [abcd1234])A [abcd1234];')

        with patch(
                'beastling.configuration.user_data_dir',
                new=Mock(return_value=self.tmp.as_posix())):
            trees = newick.read(get_glottolog_data('newick', '2.5'))
            self.assertEqual(trees[0].name, 'A [abcd1234]')

    def test_get_glottolog_data_download(self):
        data_dir = os.path.join(self.tmp.as_posix(), 'data')

        class URLopener(object):
            def retrieve(self, url, fname):
                with io.open(fname, 'w', encoding='utf8') as fp:
                    fp.write('(B [abcd1234],C [abcd1234])A [abcd1234];')

        class URLopenerError(object):
            def retrieve(self, url, fname):
                raise IOError()

        with patch.multiple(
            'beastling.configuration',
            user_data_dir=Mock(return_value=data_dir),
            URLopener=URLopenerError,
        ):
            with self.assertRaises(ValueError):
                get_glottolog_data('newick', '2.5')

        with patch.multiple(
            'beastling.configuration',
            user_data_dir=Mock(return_value=data_dir),
            URLopener=URLopener,
        ):
            assert get_glottolog_data('newick', '2.5')

    def test_families(self):
        cfg1 = self._make_cfg('glottolog_families')
        cfg2 = self._make_cfg('glottolog_families_from_file')
        cfg1.process()
        cfg2.process()
        self.assertEqual(cfg1.languages, cfg2.languages)

    def test_config(self):
        cfg = Configuration(configfile={
            'admin': {

            },
            'languages': {
                'monophyly': True,
                'starting_tree': 'T',
                'sample_topology': False,
                'sample_branch_lengths': False,
            },
            'calibration': {
                'abcd1234, efgh5678': '10-20',
            },
            'model': {
                'binarised': True,
                'minimum_data': '4.5'
            },
            'model2': {
                'binarized': True,
            },
        })
        self.assertAlmostEqual(cfg.calibration_configs['abcd1234, efgh5678'], "10-20")
        self.assertTrue(cfg.model_configs[1]['binarised'])

        with self.assertRaisesRegexp(ValueError, 'Value for overlap') as e:
            Configuration(configfile={'languages': {'overlap': 'invalid'}, 'models': {}})

        with self.assertRaisesRegexp(ValueError, 'Config file') as e:
            Configuration(configfile={'languages': {}})

    @raises(ValueError)
    def test_no_data(self):
        cfg = self._make_bad_cfg("no_data")
        cfg.process()

    @raises(ValueError)
    def test_no_langs(self):
        cfg = self._make_bad_cfg("no_langs")
        cfg.process()

    @raises(ValueError)
    def test_no_langs(self):
        cfg = self._make_bad_cfg("no_model_sec")
        cfg.process()

    @raises(ValueError)
    def test_no_langs(self):
        cfg = self._make_bad_cfg("no_model")
        cfg.process()

    @raises(ValueError)
    def test_no_langs(self):
        cfg = self._make_bad_cfg("unknown_model")
        cfg.process()

    @raises(ValueError)
    def test_bad_overlap(self):
        cfg = self._make_bad_cfg("bad_overlap")
        cfg.process()

    @raises(ValueError)
    def test_bad_frequencies(self):
        cfg = self._make_bad_cfg("bad_frequencies")
        cfg.process()
        # This is an error in a model section, which is only raised
        # when the model is constructed as XML.
        BeastXml(cfg)

    def test_calibration(self):
        config = self._make_cfg('basic', 'calibration')
        config.process()
        self.assertIn('Austronesian', config.calibrations)
        v = config.calibrations['Austronesian']
        xml1 = BeastXml(config).tostring().decode('utf8')

        # Now remove one calibration point ...
        del config.calibrations['Austronesian']
        xml2 = BeastXml(config).tostring().decode('utf8')
        self.assertNotEqual(
            len(xml1.split('CalibrationDistribution.')), len(xml2.split('CalibrationDistribution.')))

        # ... and add it back in with using the glottocode:
        config.calibrations['aust1307'] = v
        xml2 = BeastXml(config).tostring().decode('utf8')
        self.assertEqual(
            len(xml1.split('CalibrationDistribution.')), len(xml2.split('CalibrationDistribution.')))

    def test_calibration_string_formats(self):
        # Test lower bound format
        config = self._make_cfg('basic', 'calibration_lower_bound')
        config.process()
        self.assertEqual(list(config.calibrations.values())[0].dist, "uniform")
        self.assertEqual(list(config.calibrations.values())[0].param2, str(sys.maxsize))
        # Test upper bound format
        config = self._make_cfg('basic', 'calibration_upper_bound')
        config.process()
        self.assertEqual(list(config.calibrations.values())[0].dist, "uniform")
        self.assertEqual(list(config.calibrations.values())[0].param1, 0.0)

        # Test range and param formats for all three distributions
        for dist in ('normal', 'lognormal', 'uniform'):
            for style in ('range', 'params'):
                config = self._make_cfg('basic', 'calibration_%s_%s' % (dist, style))
                config.process()
            self.assertEqual(list(config.calibrations.values())[0].dist, dist)

    @raises(ValueError)
    def test_calibration_bad_bounds(self):
        config = Configuration(configfile=[
            config_path("basic").as_posix(),
            config_path("bad_cal_endpoints", bad=True).as_posix(),
            ])
        config.process()

    def test_overlong_chain(self):
        config = self._make_cfg('basic')
        config.chainlength = 9e999
        config.process()
        self.assertEqual(config.chainlength, _BEAST_MAX_LENGTH)

    def test_file_embedding(self):
        config = self._make_cfg('glottolog_families_from_file','embed_data')
        xml = BeastXml(config).tostring().decode('utf8')
        # Check for evidence of data
        self.assertTrue("aari1239,1,1,1,1,1,1,?,1,?,1" in xml)
        # Check for evidence of families
        self.assertTrue("Malayo-Polynesian" in xml)

    def test_minimum_data(self):
        # f8 has 60% missing data.  By default it should be included...
        config = self._make_cfg('basic')
        config.process()
        self.assertTrue("f8" in config.models[0].features)
        # ...but if we insist on 75% data or more it should be excluded...
        config = self._make_cfg('basic', 'minimum_data')
        config.process()
        self.assertEqual(config.models[0].minimum_data, 75.0)
        self.assertTrue("f8" not in config.models[0].features)

    def test_pruned_rlc(self):
        # Make sure pruned trees are disabled if used in conjunction with RLC
        config = self._make_cfg('basic', 'pruned', 'random')
        self.assertTrue(config.model_configs[0]["pruned"])
        config.process()
        self.assertFalse(config.models[0].pruned)

    def test_no_monophyly_geo(self):
        # Make sure that geographic sampling without monophyly constraints emits a warning
        config = self._make_cfg('basic', 'geo', 'geo_sampled')
        config.process()
        self.assertTrue(any(["[WARNING] Geographic sampling" in m for m in config.messages]))

    def test_ascertainment_auto_setting(self):
        # Without calibration, there should be no ascertainment...
        config = self._make_cfg('basic')
        config.process()
        self.assertFalse(config.models[0].ascertained)
        # But with it there should...
        config = self._make_cfg('basic', 'calibration')
        config.process()
        self.assertTrue(config.models[0].ascertained)
        # Unless, of course, we have constant data...
        config = self._make_cfg('covarion_multistate', 'calibration')
        config.model_configs[0]["remove_constant_features"] = False
        config.process()
        self.assertFalse(config.models[0].ascertained)

    def test_ascertainment_override(self):
        # Make sure we can override the above automagic
        config = self._make_cfg('basic', 'ascertainment_true')
        config.process()
        self.assertTrue(config.models[0].ascertained)
        # And with calibrations...
        config = self._make_cfg('basic', 'calibration', 'ascertainment_false')
        config.process()
        self.assertFalse(config.models[0].ascertained)

    def test_bad_ascertainment(self):
        # Make sure we refuse to produce a misspecified model
        config = self._make_cfg('covarion_multistate','ascertainment_true')
        config.model_configs[0]["remove_constant_features"] = False
        with self.assertRaises(ValueError):
            config.process()

    def test_binarisation_ascertainment(self):
        # Even with ascertainment = False, ascertainment should still
        # be done for recoded data
        config = self._make_cfg('covarion_binarised','ascertainment_false')
        config.process()
        xml = BeastXml(config).tostring().decode('utf8')
        self.assertTrue("ascertained=\"true\"" in xml)
        self.assertTrue("excludeto=\"1\"" in xml)

    def test_user_locations(self):
        # First check that we correctly load Glottolog's locations for aiw and abp
        config = self._make_cfg('basic', 'geo')
        config.process()
        print(config.locations["aiw"])
        self.assertTrue(check_lat_lon(config.locations["aiw"], 5.95, 36.57))
        self.assertTrue(check_lat_lon(config.locations["abp"], 15.41, 120.20))
        # Now check that we can overwrite just one of these...
        config = self._make_cfg('basic', 'geo', 'geo_user_loc')
        config.process()
        self.assertTrue(check_lat_lon(config.locations["aiw"], 4.20, 4.20))
        self.assertTrue(check_lat_lon(config.locations["abp"], 15.41, 120.20))
        # Now check that we can overwrite them both using multiple files
        config = self._make_cfg('basic', 'geo', 'geo_user_loc_multifile')
        config.process()
        self.assertTrue(check_lat_lon(config.locations["aiw"], 4.20, 4.20))
        self.assertTrue(check_lat_lon(config.locations["abp"], 6.66, 6.66))
<commit_msg>Kill debugging print.<commit_after># coding: utf8
from __future__ import unicode_literals
import os
import io
import sys

from nose.tools import *
from mock import patch, Mock
import newick
from clldutils.path import Path

from beastling.configuration import (
    Configuration, get_glottolog_data, _BEAST_MAX_LENGTH,
)
from beastling.beastxml import BeastXml
from .util import WithConfigAndTempDir, config_path

def check_lat_lon(provided, target_lat, target_lon):
    prov_lat, prov_lon = provided
    return round(prov_lat, 2) == target_lat and round(prov_lon, 2) == target_lon

class Tests(WithConfigAndTempDir):
    def _make_cfg(self, *names):
        return self.make_cfg([config_path(name).as_posix() for name in names])

    def _make_bad_cfg(self, name):
        return self.make_cfg(config_path(name, bad=True).as_posix())

    def test_get_glottolog_geo(self):
        geodata = self.tmp.joinpath('glottolog-2.5-geo.csv')
        with geodata.open('w', encoding='utf8') as fp:
            fp.write('x')

        with patch(
                'beastling.configuration.user_data_dir',
                new=Mock(return_value=self.tmp.as_posix())):
            self.assertEqual(Path(get_glottolog_data('geo', '2.5')), geodata)

    def test_get_glottolog_newick(self):
        with self.tmp.joinpath('glottolog-2.5.newick').open('w', encoding='utf8') as fp:
            fp.write('(B [abcd1234],C [abcd1234])A [abcd1234];')

        with patch(
                'beastling.configuration.user_data_dir',
                new=Mock(return_value=self.tmp.as_posix())):
            trees = newick.read(get_glottolog_data('newick', '2.5'))
            self.assertEqual(trees[0].name, 'A [abcd1234]')

    def test_get_glottolog_data_download(self):
        data_dir = os.path.join(self.tmp.as_posix(), 'data')

        class URLopener(object):
            def retrieve(self, url, fname):
                with io.open(fname, 'w', encoding='utf8') as fp:
                    fp.write('(B [abcd1234],C [abcd1234])A [abcd1234];')

        class URLopenerError(object):
            def retrieve(self, url, fname):
                raise IOError()

        with patch.multiple(
            'beastling.configuration',
            user_data_dir=Mock(return_value=data_dir),
            URLopener=URLopenerError,
        ):
            with self.assertRaises(ValueError):
                get_glottolog_data('newick', '2.5')

        with patch.multiple(
            'beastling.configuration',
            user_data_dir=Mock(return_value=data_dir),
            URLopener=URLopener,
        ):
            assert get_glottolog_data('newick', '2.5')

    def test_families(self):
        cfg1 = self._make_cfg('glottolog_families')
        cfg2 = self._make_cfg('glottolog_families_from_file')
        cfg1.process()
        cfg2.process()
        self.assertEqual(cfg1.languages, cfg2.languages)

    def test_config(self):
        cfg = Configuration(configfile={
            'admin': {

            },
            'languages': {
                'monophyly': True,
                'starting_tree': 'T',
                'sample_topology': False,
                'sample_branch_lengths': False,
            },
            'calibration': {
                'abcd1234, efgh5678': '10-20',
            },
            'model': {
                'binarised': True,
                'minimum_data': '4.5'
            },
            'model2': {
                'binarized': True,
            },
        })
        self.assertAlmostEqual(cfg.calibration_configs['abcd1234, efgh5678'], "10-20")
        self.assertTrue(cfg.model_configs[1]['binarised'])

        with self.assertRaisesRegexp(ValueError, 'Value for overlap') as e:
            Configuration(configfile={'languages': {'overlap': 'invalid'}, 'models': {}})

        with self.assertRaisesRegexp(ValueError, 'Config file') as e:
            Configuration(configfile={'languages': {}})

    @raises(ValueError)
    def test_no_data(self):
        cfg = self._make_bad_cfg("no_data")
        cfg.process()

    @raises(ValueError)
    def test_no_langs(self):
        cfg = self._make_bad_cfg("no_langs")
        cfg.process()

    @raises(ValueError)
    def test_no_langs(self):
        cfg = self._make_bad_cfg("no_model_sec")
        cfg.process()

    @raises(ValueError)
    def test_no_langs(self):
        cfg = self._make_bad_cfg("no_model")
        cfg.process()

    @raises(ValueError)
    def test_no_langs(self):
        cfg = self._make_bad_cfg("unknown_model")
        cfg.process()

    @raises(ValueError)
    def test_bad_overlap(self):
        cfg = self._make_bad_cfg("bad_overlap")
        cfg.process()

    @raises(ValueError)
    def test_bad_frequencies(self):
        cfg = self._make_bad_cfg("bad_frequencies")
        cfg.process()
        # This is an error in a model section, which is only raised
        # when the model is constructed as XML.
        BeastXml(cfg)

    def test_calibration(self):
        config = self._make_cfg('basic', 'calibration')
        config.process()
        self.assertIn('Austronesian', config.calibrations)
        v = config.calibrations['Austronesian']
        xml1 = BeastXml(config).tostring().decode('utf8')

        # Now remove one calibration point ...
        del config.calibrations['Austronesian']
        xml2 = BeastXml(config).tostring().decode('utf8')
        self.assertNotEqual(
            len(xml1.split('CalibrationDistribution.')), len(xml2.split('CalibrationDistribution.')))

        # ... and add it back in with using the glottocode:
        config.calibrations['aust1307'] = v
        xml2 = BeastXml(config).tostring().decode('utf8')
        self.assertEqual(
            len(xml1.split('CalibrationDistribution.')), len(xml2.split('CalibrationDistribution.')))

    def test_calibration_string_formats(self):
        # Test lower bound format
        config = self._make_cfg('basic', 'calibration_lower_bound')
        config.process()
        self.assertEqual(list(config.calibrations.values())[0].dist, "uniform")
        self.assertEqual(list(config.calibrations.values())[0].param2, str(sys.maxsize))
        # Test upper bound format
        config = self._make_cfg('basic', 'calibration_upper_bound')
        config.process()
        self.assertEqual(list(config.calibrations.values())[0].dist, "uniform")
        self.assertEqual(list(config.calibrations.values())[0].param1, 0.0)

        # Test range and param formats for all three distributions
        for dist in ('normal', 'lognormal', 'uniform'):
            for style in ('range', 'params'):
                config = self._make_cfg('basic', 'calibration_%s_%s' % (dist, style))
                config.process()
            self.assertEqual(list(config.calibrations.values())[0].dist, dist)

    @raises(ValueError)
    def test_calibration_bad_bounds(self):
        config = Configuration(configfile=[
            config_path("basic").as_posix(),
            config_path("bad_cal_endpoints", bad=True).as_posix(),
            ])
        config.process()

    def test_overlong_chain(self):
        config = self._make_cfg('basic')
        config.chainlength = 9e999
        config.process()
        self.assertEqual(config.chainlength, _BEAST_MAX_LENGTH)

    def test_file_embedding(self):
        config = self._make_cfg('glottolog_families_from_file','embed_data')
        xml = BeastXml(config).tostring().decode('utf8')
        # Check for evidence of data
        self.assertTrue("aari1239,1,1,1,1,1,1,?,1,?,1" in xml)
        # Check for evidence of families
        self.assertTrue("Malayo-Polynesian" in xml)

    def test_minimum_data(self):
        # f8 has 60% missing data.  By default it should be included...
        config = self._make_cfg('basic')
        config.process()
        self.assertTrue("f8" in config.models[0].features)
        # ...but if we insist on 75% data or more it should be excluded...
        config = self._make_cfg('basic', 'minimum_data')
        config.process()
        self.assertEqual(config.models[0].minimum_data, 75.0)
        self.assertTrue("f8" not in config.models[0].features)

    def test_pruned_rlc(self):
        # Make sure pruned trees are disabled if used in conjunction with RLC
        config = self._make_cfg('basic', 'pruned', 'random')
        self.assertTrue(config.model_configs[0]["pruned"])
        config.process()
        self.assertFalse(config.models[0].pruned)

    def test_no_monophyly_geo(self):
        # Make sure that geographic sampling without monophyly constraints emits a warning
        config = self._make_cfg('basic', 'geo', 'geo_sampled')
        config.process()
        self.assertTrue(any(["[WARNING] Geographic sampling" in m for m in config.messages]))

    def test_ascertainment_auto_setting(self):
        # Without calibration, there should be no ascertainment...
        config = self._make_cfg('basic')
        config.process()
        self.assertFalse(config.models[0].ascertained)
        # But with it there should...
        config = self._make_cfg('basic', 'calibration')
        config.process()
        self.assertTrue(config.models[0].ascertained)
        # Unless, of course, we have constant data...
        config = self._make_cfg('covarion_multistate', 'calibration')
        config.model_configs[0]["remove_constant_features"] = False
        config.process()
        self.assertFalse(config.models[0].ascertained)

    def test_ascertainment_override(self):
        # Make sure we can override the above automagic
        config = self._make_cfg('basic', 'ascertainment_true')
        config.process()
        self.assertTrue(config.models[0].ascertained)
        # And with calibrations...
        config = self._make_cfg('basic', 'calibration', 'ascertainment_false')
        config.process()
        self.assertFalse(config.models[0].ascertained)

    def test_bad_ascertainment(self):
        # Make sure we refuse to produce a misspecified model
        config = self._make_cfg('covarion_multistate','ascertainment_true')
        config.model_configs[0]["remove_constant_features"] = False
        with self.assertRaises(ValueError):
            config.process()

    def test_binarisation_ascertainment(self):
        # Even with ascertainment = False, ascertainment should still
        # be done for recoded data
        config = self._make_cfg('covarion_binarised','ascertainment_false')
        config.process()
        xml = BeastXml(config).tostring().decode('utf8')
        self.assertTrue("ascertained=\"true\"" in xml)
        self.assertTrue("excludeto=\"1\"" in xml)

    def test_user_locations(self):
        # First check that we correctly load Glottolog's locations for aiw and abp
        config = self._make_cfg('basic', 'geo')
        config.process()
        self.assertTrue(check_lat_lon(config.locations["aiw"], 5.95, 36.57))
        self.assertTrue(check_lat_lon(config.locations["abp"], 15.41, 120.20))
        # Now check that we can overwrite just one of these...
        config = self._make_cfg('basic', 'geo', 'geo_user_loc')
        config.process()
        self.assertTrue(check_lat_lon(config.locations["aiw"], 4.20, 4.20))
        self.assertTrue(check_lat_lon(config.locations["abp"], 15.41, 120.20))
        # Now check that we can overwrite them both using multiple files
        config = self._make_cfg('basic', 'geo', 'geo_user_loc_multifile')
        config.process()
        self.assertTrue(check_lat_lon(config.locations["aiw"], 4.20, 4.20))
        self.assertTrue(check_lat_lon(config.locations["abp"], 6.66, 6.66))
<|endoftext|>


<commit_before># Lint as: python3
# Copyright 2020 Google LLC. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""E2E test using kubeflow orchestrator for taxi template."""

import datetime
import os
import subprocess
import tarfile
import urllib.request

from absl import logging
import docker
import kfp
import tensorflow as tf
from tfx.dsl.io import fileio
from tfx.experimental.templates import test_utils
from tfx.orchestration import test_utils as orchestration_test_utils
from tfx.orchestration.kubeflow import test_utils as kubeflow_test_utils
from tfx.utils import retry
from tfx.utils import telemetry_utils
import yaml
from google.cloud import storage

class TaxiTemplateKubeflowE2ETest(test_utils.BaseEndToEndTest):

  _POLLING_INTERVAL_IN_SECONDS = 10
  _MAX_POLLING_COUNT = 20 * 6  # 20 min.

  _DATA_DIRECTORY_NAME = 'template_data'

  # The following environment variables need to be set prior to calling the test
  # in this file. All variables are required and do not have a default.

  # The base container image name to use when building the image used in tests.
  _BASE_CONTAINER_IMAGE = os.environ['KFP_E2E_BASE_CONTAINER_IMAGE']

  # The src path to use to build docker image
  _REPO_BASE = os.environ['KFP_E2E_SRC']

  # The project id to use to run tests.
  _GCP_PROJECT_ID = os.environ['KFP_E2E_GCP_PROJECT_ID']

  # The GCP region in which the end-to-end test is run.
  _GCP_REGION = os.environ['KFP_E2E_GCP_REGION']

  # The GCP bucket to use to write output artifacts.
  # This default bucket name is valid for KFP marketplace deployment since KFP
  # version 0.5.0.
  _BUCKET_NAME = _GCP_PROJECT_ID + '-kubeflowpipelines-default'

  def setUp(self):
    super().setUp()
    random_id = orchestration_test_utils.random_id()
    self._pipeline_name = 'taxi-template-kubeflow-e2e-test-' + random_id
    logging.info('Pipeline: %s', self._pipeline_name)
    self._namespace = 'kubeflow'
    self._endpoint = self._get_endpoint(self._namespace)
    self._kfp_client = kfp.Client(host=self._endpoint)
    logging.info('ENDPOINT: %s', self._endpoint)

    if ':' not in self._BASE_CONTAINER_IMAGE:
      self._base_container_image = '{}:{}'.format(self._BASE_CONTAINER_IMAGE,
                                                  random_id)
      self._prepare_base_container_image()
    else:
      self._base_container_image = self._BASE_CONTAINER_IMAGE
    self._target_container_image = 'gcr.io/{}/{}:{}'.format(
        self._GCP_PROJECT_ID, 'taxi-template-kubeflow-e2e-test', random_id)

    self._prepare_skaffold()

  def tearDown(self):
    super(TaxiTemplateKubeflowE2ETest, self).tearDown()
    self._cleanup_kfp()

  def _cleanup_kfp(self):
    self._delete_base_container_image()
    self._delete_target_container_image()
    self._delete_caip_model()
    self._delete_runs()
    self._delete_pipeline()
    self._delete_pipeline_data()

  def _get_kfp_runs(self):
    # CLI uses experiment_name which is the same as pipeline_name.
    experiment_id = self._kfp_client.get_experiment(
        experiment_name=self._pipeline_name).id
    response = self._kfp_client.list_runs(experiment_id=experiment_id)
    return response.runs

  @retry.retry(ignore_eventual_failure=True)
  def _delete_caip_model(self):
    model_name = self._pipeline_name.replace('-', '_')
    kubeflow_test_utils.delete_ai_platform_model(model_name)

  @retry.retry(ignore_eventual_failure=True)
  def _delete_runs(self):
    for run in self._get_kfp_runs():
      self._kfp_client._run_api.delete_run(id=run.id)

  @retry.retry(ignore_eventual_failure=True)
  def _delete_pipeline(self):
    self._runCli([
        'pipeline', 'delete', '--engine', 'kubeflow', '--pipeline_name',
        self._pipeline_name
    ])

  @retry.retry(ignore_eventual_failure=True)
  def _delete_pipeline_data(self):
    path = 'tfx_pipeline_output/{}'.format(self._pipeline_name)
    orchestration_test_utils.delete_gcs_files(self._GCP_PROJECT_ID,
                                              self._BUCKET_NAME, path)
    path = '{}/{}'.format(self._DATA_DIRECTORY_NAME, self._pipeline_name)
    orchestration_test_utils.delete_gcs_files(self._GCP_PROJECT_ID,
                                              self._BUCKET_NAME, path)

  @retry.retry(ignore_eventual_failure=True)
  def _delete_docker_image(self, image):
    subprocess.check_output(['gcloud', 'container', 'images', 'delete', image])
    client = docker.from_env()
    client.images.remove(image=image)

  def _delete_base_container_image(self):
    if self._base_container_image == self._BASE_CONTAINER_IMAGE:
      return  # Didn't generate a base image for the test.
    self._delete_docker_image(self._base_container_image)

  def _delete_target_container_image(self):
    self._delete_docker_image(self._target_container_image)

  def _get_endpoint(self, namespace):
    cmd = 'kubectl describe configmap inverse-proxy-config -n {}'.format(
        namespace)
    output = subprocess.check_output(cmd.split())
    for line in output.decode('utf-8').split('\n'):
      if line.endswith('googleusercontent.com'):
        return line

  def _prepare_data(self):
    client = storage.Client(project=self._GCP_PROJECT_ID)
    bucket = client.bucket(self._BUCKET_NAME)
    blob = bucket.blob('{}/{}/data.csv'.format(self._DATA_DIRECTORY_NAME,
                                               self._pipeline_name))
    blob.upload_from_filename('data/data.csv')

  def _prepare_base_container_image(self):
    orchestration_test_utils.build_docker_image(self._base_container_image,
                                                self._REPO_BASE)

  def _prepare_skaffold(self):
    self._skaffold = os.path.join(self._temp_dir, 'skaffold')
    urllib.request.urlretrieve(
        'https://storage.googleapis.com/skaffold/releases/latest/skaffold-linux-amd64',
        self._skaffold)
    os.chmod(self._skaffold, 0o775)

  def _create_pipeline(self):
    result = self._runCli([
        'pipeline',
        'create',
        '--engine',
        'kubeflow',
        '--pipeline_path',
        'kubeflow_dag_runner.py',
        '--endpoint',
        self._endpoint,
        '--build-target-image',
        self._target_container_image,
        '--skaffold-cmd',
        self._skaffold,
        '--build-base-image',
        self._base_container_image,
    ])
    self.assertEqual(0, result.exit_code)

  def _update_pipeline(self):
    result = self._runCli([
        'pipeline',
        'update',
        '--engine',
        'kubeflow',
        '--pipeline_path',
        'kubeflow_dag_runner.py',
        '--endpoint',
        self._endpoint,
        '--skaffold-cmd',
        self._skaffold,
    ])
    self.assertEqual(0, result.exit_code)

  def _run_pipeline(self):
    result = self._runCli([
        'run',
        'create',
        '--engine',
        'kubeflow',
        '--pipeline_name',
        self._pipeline_name,
        '--endpoint',
        self._endpoint,
    ])
    self.assertEqual(0, result.exit_code)
    run_id = self._parse_run_id(result.output)
    self._wait_until_completed(run_id)
    kubeflow_test_utils.print_failure_log_for_run(self._endpoint, run_id,
                                                  self._namespace)

  def _parse_run_id(self, output: str):
    run_id_lines = [
        line for line in output.split('\n')
        if '| {} |'.format(self._pipeline_name) in line
    ]
    self.assertLen(run_id_lines, 1)
    return run_id_lines[0].split('|')[2].strip()

  def _wait_until_completed(self, run_id: str):
    # This timeout will never expire. polling_count * interval == 20min.
    timeout = datetime.timedelta(hours=1)
    end_state = kubeflow_test_utils.poll_kfp_with_retry(
        self._endpoint, run_id, self._MAX_POLLING_COUNT, timeout,
        self._POLLING_INTERVAL_IN_SECONDS)
    self.assertEqual(end_state.lower(), kubeflow_test_utils.KFP_SUCCESS_STATUS)

  def _check_telemetry_label(self):
    file_path = os.path.join(self._project_dir,
                             '{}.tar.gz'.format(self._pipeline_name))
    self.assertTrue(fileio.exists(file_path))

    with tarfile.TarFile.open(file_path).extractfile(
        'pipeline.yaml') as pipeline_file:
      self.assertIsNotNone(pipeline_file)
      pipeline = yaml.safe_load(pipeline_file)
      metadata = [
          c['metadata'] for c in pipeline['spec']['templates'] if 'dag' not in c
      ]
      for m in metadata:
        self.assertEqual('tfx-template',
                         m['labels'][telemetry_utils.LABEL_KFP_SDK_ENV])

  def testPipeline(self):
    self._copyTemplate('taxi')
    os.environ['KUBEFLOW_HOME'] = os.path.join(self._temp_dir, 'kubeflow')

    # Uncomment all variables in config.
    self._uncommentMultiLineVariables(
        os.path.join('pipeline', 'configs.py'), [
            'GOOGLE_CLOUD_REGION',
            'BIG_QUERY_WITH_DIRECT_RUNNER_BEAM_PIPELINE_ARGS',
            'BIG_QUERY_QUERY', 'DATAFLOW_BEAM_PIPELINE_ARGS',
            'GCP_AI_PLATFORM_TRAINING_ARGS', 'GCP_AI_PLATFORM_SERVING_ARGS'
        ])
    self._replaceFileContent(
        os.path.join('pipeline', 'configs.py'), [
            ('GOOGLE_CLOUD_REGION = \'\'',
             'GOOGLE_CLOUD_REGION = \'{}\''.format(self._GCP_REGION)),
            ('\'imageUri\': \'gcr.io/\' + GOOGLE_CLOUD_PROJECT + \'/tfx-pipeline\'',
             '\'imageUri\': \'{}\''.format(self._target_container_image)),
        ])

    # Prepare data
    self._prepare_data()
    self._replaceFileContent('kubeflow_dag_runner.py', [
        ('DATA_PATH = \'gs://{}/tfx-template/data/taxi/\'.format(configs.GCS_BUCKET_NAME)',
         'DATA_PATH = \'gs://{{}}/{}/{}\'.format(configs.GCS_BUCKET_NAME)'
         .format(self._DATA_DIRECTORY_NAME, self._pipeline_name)),
    ])

    # Create a pipeline with only one component.
    self._create_pipeline()
    self._run_pipeline()

    self._check_telemetry_label()

    # Update the pipeline to include all components.
    updated_pipeline_file = self._addAllComponents()
    logging.info('Updated %s to add all components to the pipeline.',
                 updated_pipeline_file)
    self._update_pipeline()
    self._run_pipeline()

    # Enable BigQuery
    self._uncomment(
        os.path.join('pipeline', 'pipeline.py'),
        ['query: Text,', 'example_gen = BigQueryExampleGen('])
    self._uncomment('kubeflow_dag_runner.py', [
        'query=configs.BIG_QUERY_QUERY',
        'beam_pipeline_args=configs\n',
        '.BIG_QUERY_WITH_DIRECT_RUNNER_BEAM_PIPELINE_ARGS,',
    ])
    logging.info('Added BigQueryExampleGen to pipeline.')
    self._update_pipeline()
    self._run_pipeline()

    # TODO(b/173065862) Re-enable Dataflow tests after timeout is resolved.
#     # Enable Dataflow
#     self._comment('kubeflow_dag_runner.py', [
#         'beam_pipeline_args=configs\n',
#         '.BIG_QUERY_WITH_DIRECT_RUNNER_BEAM_PIPELINE_ARGS',
#     ])
#     self._uncomment('kubeflow_dag_runner.py', [
#         'beam_pipeline_args=configs.DATAFLOW_BEAM_PIPELINE_ARGS',
#     ])
#     logging.info('Added Dataflow to pipeline.')
#     self._update_pipeline()
#     self._run_pipeline()

#     # Enable CAIP extension.
#     self._comment('kubeflow_dag_runner.py', [
#         'beam_pipeline_args=configs.DATAFLOW_BEAM_PIPELINE_ARGS',
#     ])
    self._uncomment('kubeflow_dag_runner.py', [
        'ai_platform_training_args=configs.GCP_AI_PLATFORM_TRAINING_ARGS,',
        'ai_platform_serving_args=configs.GCP_AI_PLATFORM_SERVING_ARGS,',
    ])
    logging.info('Using CAIP trainer and pusher.')
    self._update_pipeline()
    self._run_pipeline()

if __name__ == '__main__':
  logging.set_verbosity(logging.INFO)
  tf.test.main()
<commit_msg>Fixes wrong Python pipeline definition file name in the Kubeflow E2E test.<commit_after># Lint as: python3
# Copyright 2020 Google LLC. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""E2E test using kubeflow orchestrator for taxi template."""

import datetime
import os
import subprocess
import tarfile
import urllib.request

from absl import logging
import docker
import kfp
import tensorflow as tf
from tfx.dsl.io import fileio
from tfx.experimental.templates import test_utils
from tfx.orchestration import test_utils as orchestration_test_utils
from tfx.orchestration.kubeflow import test_utils as kubeflow_test_utils
from tfx.utils import retry
from tfx.utils import telemetry_utils
import yaml
from google.cloud import storage

class TaxiTemplateKubeflowE2ETest(test_utils.BaseEndToEndTest):

  _POLLING_INTERVAL_IN_SECONDS = 10
  _MAX_POLLING_COUNT = 20 * 6  # 20 min.

  _DATA_DIRECTORY_NAME = 'template_data'

  # The following environment variables need to be set prior to calling the test
  # in this file. All variables are required and do not have a default.

  # The base container image name to use when building the image used in tests.
  _BASE_CONTAINER_IMAGE = os.environ['KFP_E2E_BASE_CONTAINER_IMAGE']

  # The src path to use to build docker image
  _REPO_BASE = os.environ['KFP_E2E_SRC']

  # The project id to use to run tests.
  _GCP_PROJECT_ID = os.environ['KFP_E2E_GCP_PROJECT_ID']

  # The GCP region in which the end-to-end test is run.
  _GCP_REGION = os.environ['KFP_E2E_GCP_REGION']

  # The GCP bucket to use to write output artifacts.
  # This default bucket name is valid for KFP marketplace deployment since KFP
  # version 0.5.0.
  _BUCKET_NAME = _GCP_PROJECT_ID + '-kubeflowpipelines-default'

  def setUp(self):
    super().setUp()
    random_id = orchestration_test_utils.random_id()
    self._pipeline_name = 'taxi-template-kubeflow-e2e-test-' + random_id
    logging.info('Pipeline: %s', self._pipeline_name)
    self._namespace = 'kubeflow'
    self._endpoint = self._get_endpoint(self._namespace)
    self._kfp_client = kfp.Client(host=self._endpoint)
    logging.info('ENDPOINT: %s', self._endpoint)

    if ':' not in self._BASE_CONTAINER_IMAGE:
      self._base_container_image = '{}:{}'.format(self._BASE_CONTAINER_IMAGE,
                                                  random_id)
      self._prepare_base_container_image()
    else:
      self._base_container_image = self._BASE_CONTAINER_IMAGE
    self._target_container_image = 'gcr.io/{}/{}:{}'.format(
        self._GCP_PROJECT_ID, 'taxi-template-kubeflow-e2e-test', random_id)

    self._prepare_skaffold()

  def tearDown(self):
    super(TaxiTemplateKubeflowE2ETest, self).tearDown()
    self._cleanup_kfp()

  def _cleanup_kfp(self):
    self._delete_base_container_image()
    self._delete_target_container_image()
    self._delete_caip_model()
    self._delete_runs()
    self._delete_pipeline()
    self._delete_pipeline_data()

  def _get_kfp_runs(self):
    # CLI uses experiment_name which is the same as pipeline_name.
    experiment_id = self._kfp_client.get_experiment(
        experiment_name=self._pipeline_name).id
    response = self._kfp_client.list_runs(experiment_id=experiment_id)
    return response.runs

  @retry.retry(ignore_eventual_failure=True)
  def _delete_caip_model(self):
    model_name = self._pipeline_name.replace('-', '_')
    kubeflow_test_utils.delete_ai_platform_model(model_name)

  @retry.retry(ignore_eventual_failure=True)
  def _delete_runs(self):
    for run in self._get_kfp_runs():
      self._kfp_client._run_api.delete_run(id=run.id)

  @retry.retry(ignore_eventual_failure=True)
  def _delete_pipeline(self):
    self._runCli([
        'pipeline', 'delete', '--engine', 'kubeflow', '--pipeline_name',
        self._pipeline_name
    ])

  @retry.retry(ignore_eventual_failure=True)
  def _delete_pipeline_data(self):
    path = 'tfx_pipeline_output/{}'.format(self._pipeline_name)
    orchestration_test_utils.delete_gcs_files(self._GCP_PROJECT_ID,
                                              self._BUCKET_NAME, path)
    path = '{}/{}'.format(self._DATA_DIRECTORY_NAME, self._pipeline_name)
    orchestration_test_utils.delete_gcs_files(self._GCP_PROJECT_ID,
                                              self._BUCKET_NAME, path)

  @retry.retry(ignore_eventual_failure=True)
  def _delete_docker_image(self, image):
    subprocess.check_output(['gcloud', 'container', 'images', 'delete', image])
    client = docker.from_env()
    client.images.remove(image=image)

  def _delete_base_container_image(self):
    if self._base_container_image == self._BASE_CONTAINER_IMAGE:
      return  # Didn't generate a base image for the test.
    self._delete_docker_image(self._base_container_image)

  def _delete_target_container_image(self):
    self._delete_docker_image(self._target_container_image)

  def _get_endpoint(self, namespace):
    cmd = 'kubectl describe configmap inverse-proxy-config -n {}'.format(
        namespace)
    output = subprocess.check_output(cmd.split())
    for line in output.decode('utf-8').split('\n'):
      if line.endswith('googleusercontent.com'):
        return line

  def _prepare_data(self):
    client = storage.Client(project=self._GCP_PROJECT_ID)
    bucket = client.bucket(self._BUCKET_NAME)
    blob = bucket.blob('{}/{}/data.csv'.format(self._DATA_DIRECTORY_NAME,
                                               self._pipeline_name))
    blob.upload_from_filename('data/data.csv')

  def _prepare_base_container_image(self):
    orchestration_test_utils.build_docker_image(self._base_container_image,
                                                self._REPO_BASE)

  def _prepare_skaffold(self):
    self._skaffold = os.path.join(self._temp_dir, 'skaffold')
    urllib.request.urlretrieve(
        'https://storage.googleapis.com/skaffold/releases/latest/skaffold-linux-amd64',
        self._skaffold)
    os.chmod(self._skaffold, 0o775)

  def _create_pipeline(self):
    result = self._runCli([
        'pipeline',
        'create',
        '--engine',
        'kubeflow',
        '--pipeline_path',
        'kubeflow_runner.py',
        '--endpoint',
        self._endpoint,
        '--build-target-image',
        self._target_container_image,
        '--skaffold-cmd',
        self._skaffold,
        '--build-base-image',
        self._base_container_image,
    ])
    self.assertEqual(0, result.exit_code)

  def _update_pipeline(self):
    result = self._runCli([
        'pipeline',
        'update',
        '--engine',
        'kubeflow',
        '--pipeline_path',
        'kubeflow_runner.py',
        '--endpoint',
        self._endpoint,
        '--skaffold-cmd',
        self._skaffold,
    ])
    self.assertEqual(0, result.exit_code)

  def _run_pipeline(self):
    result = self._runCli([
        'run',
        'create',
        '--engine',
        'kubeflow',
        '--pipeline_name',
        self._pipeline_name,
        '--endpoint',
        self._endpoint,
    ])
    self.assertEqual(0, result.exit_code)
    run_id = self._parse_run_id(result.output)
    self._wait_until_completed(run_id)
    kubeflow_test_utils.print_failure_log_for_run(self._endpoint, run_id,
                                                  self._namespace)

  def _parse_run_id(self, output: str):
    run_id_lines = [
        line for line in output.split('\n')
        if '| {} |'.format(self._pipeline_name) in line
    ]
    self.assertLen(run_id_lines, 1)
    return run_id_lines[0].split('|')[2].strip()

  def _wait_until_completed(self, run_id: str):
    # This timeout will never expire. polling_count * interval == 20min.
    timeout = datetime.timedelta(hours=1)
    end_state = kubeflow_test_utils.poll_kfp_with_retry(
        self._endpoint, run_id, self._MAX_POLLING_COUNT, timeout,
        self._POLLING_INTERVAL_IN_SECONDS)
    self.assertEqual(end_state.lower(), kubeflow_test_utils.KFP_SUCCESS_STATUS)

  def _check_telemetry_label(self):
    file_path = os.path.join(self._project_dir,
                             '{}.tar.gz'.format(self._pipeline_name))
    self.assertTrue(fileio.exists(file_path))

    with tarfile.TarFile.open(file_path).extractfile(
        'pipeline.yaml') as pipeline_file:
      self.assertIsNotNone(pipeline_file)
      pipeline = yaml.safe_load(pipeline_file)
      metadata = [
          c['metadata'] for c in pipeline['spec']['templates'] if 'dag' not in c
      ]
      for m in metadata:
        self.assertEqual('tfx-template',
                         m['labels'][telemetry_utils.LABEL_KFP_SDK_ENV])

  def testPipeline(self):
    self._copyTemplate('taxi')
    os.environ['KUBEFLOW_HOME'] = os.path.join(self._temp_dir, 'kubeflow')

    # Uncomment all variables in config.
    self._uncommentMultiLineVariables(
        os.path.join('pipeline', 'configs.py'), [
            'GOOGLE_CLOUD_REGION',
            'BIG_QUERY_WITH_DIRECT_RUNNER_BEAM_PIPELINE_ARGS',
            'BIG_QUERY_QUERY', 'DATAFLOW_BEAM_PIPELINE_ARGS',
            'GCP_AI_PLATFORM_TRAINING_ARGS', 'GCP_AI_PLATFORM_SERVING_ARGS'
        ])
    self._replaceFileContent(
        os.path.join('pipeline', 'configs.py'), [
            ('GOOGLE_CLOUD_REGION = \'\'',
             'GOOGLE_CLOUD_REGION = \'{}\''.format(self._GCP_REGION)),
            ('\'imageUri\': \'gcr.io/\' + GOOGLE_CLOUD_PROJECT + \'/tfx-pipeline\'',
             '\'imageUri\': \'{}\''.format(self._target_container_image)),
        ])

    # Prepare data
    self._prepare_data()
    self._replaceFileContent('kubeflow_runner.py', [
        ('DATA_PATH = \'gs://{}/tfx-template/data/taxi/\'.format(configs.GCS_BUCKET_NAME)',
         'DATA_PATH = \'gs://{{}}/{}/{}\'.format(configs.GCS_BUCKET_NAME)'
         .format(self._DATA_DIRECTORY_NAME, self._pipeline_name)),
    ])

    # Create a pipeline with only one component.
    self._create_pipeline()
    self._run_pipeline()

    self._check_telemetry_label()

    # Update the pipeline to include all components.
    updated_pipeline_file = self._addAllComponents()
    logging.info('Updated %s to add all components to the pipeline.',
                 updated_pipeline_file)
    self._update_pipeline()
    self._run_pipeline()

    # Enable BigQuery
    self._uncomment(
        os.path.join('pipeline', 'pipeline.py'),
        ['query: Text,', 'example_gen = BigQueryExampleGen('])
    self._uncomment('kubeflow_runner.py', [
        'query=configs.BIG_QUERY_QUERY',
        'beam_pipeline_args=configs\n',
        '.BIG_QUERY_WITH_DIRECT_RUNNER_BEAM_PIPELINE_ARGS,',
    ])
    logging.info('Added BigQueryExampleGen to pipeline.')
    self._update_pipeline()
    self._run_pipeline()

    # TODO(b/173065862) Re-enable Dataflow tests after timeout is resolved.
#     # Enable Dataflow
#     self._comment('kubeflow_runner.py', [
#         'beam_pipeline_args=configs\n',
#         '.BIG_QUERY_WITH_DIRECT_RUNNER_BEAM_PIPELINE_ARGS',
#     ])
#     self._uncomment('kubeflow_runner.py', [
#         'beam_pipeline_args=configs.DATAFLOW_BEAM_PIPELINE_ARGS',
#     ])
#     logging.info('Added Dataflow to pipeline.')
#     self._update_pipeline()
#     self._run_pipeline()

#     # Enable CAIP extension.
#     self._comment('kubeflow_runner.py', [
#         'beam_pipeline_args=configs.DATAFLOW_BEAM_PIPELINE_ARGS',
#     ])
    self._uncomment('kubeflow_runner.py', [
        'ai_platform_training_args=configs.GCP_AI_PLATFORM_TRAINING_ARGS,',
        'ai_platform_serving_args=configs.GCP_AI_PLATFORM_SERVING_ARGS,',
    ])
    logging.info('Using CAIP trainer and pusher.')
    self._update_pipeline()
    self._run_pipeline()

if __name__ == '__main__':
  logging.set_verbosity(logging.INFO)
  tf.test.main()
<|endoftext|>


<commit_before>#!/usr/bin/env python3

import argparse
import itertools
from time import gmtime, strftime

from eva_cttv_pipeline import clinvar_xml_utils

class OntologyUri:
    db_to_uri_dict = {
        'orphanet': 'http://www.orpha.net/ORDO/Orphanet_{}',
        'omim':     'http://identifiers.org/omim/{}',
        'efo':      'http://www.ebi.ac.uk/efo/{}',
        'mesh':     'http://identifiers.org/mesh/{}',
        'medgen':   'http://identifiers.org/medgen/{}',
        'mondo':    'http://purl.obolibrary.org/obo/MONDO_{}',
    }

    def __init__(self, id_, db):
        self.id_ = id_
        self.db = db
        self.uri = self.db_to_uri_dict[self.db.lower()].format(self.id_)

    def __str__(self):
        return self.uri

def write_zooma_record(clinvar_acc, variant_id, trait_name, ontology_uri, date, outfile):
    zooma_output_list = [clinvar_acc,
                         variant_id,
                         'disease',
                         trait_name,
                         str(ontology_uri),
                         'clinvar-xrefs',
                         date]
    outfile.write('\t'.join(zooma_output_list) + '\n')

def process_clinvar_record(clinvar_record, outfile):
    """Extract the variant, trait and ontology from Clinvar record and write them as a ZOOMA feedback record."""
    if clinvar_record.measure is None:
        return
    variant_ids = [variant_id
                   for variant_id in (clinvar_record.measure.rs_id, clinvar_record.measure.nsv_id)
                   if variant_id is not None]
    traits = clinvar_record.traits
    for variant_id, trait in itertools.product(variant_ids, traits):
        if (trait.name is None) or (trait.name.lower() == 'not provided'):
            continue
        for db, identifier, status in trait.xrefs:
            if status != 'current' or db.lower() not in OntologyUri.db_to_uri_dict:
                continue
            ontology_uri = OntologyUri(identifier, db)
            write_zooma_record(clinvar_record.accession, variant_id, trait.name, ontology_uri,
                               strftime('%d/%m/%y %H:%M', gmtime()), outfile)

def main(clinvar_xml, zooma_feedback):
    with open(zooma_feedback, 'wt') as outfile:
        outfile.write('STUDY\tBIOENTITY\tPROPERTY_TYPE\tPROPERTY_VALUE\tSEMANTIC_TAG\tANNOTATOR\tANNOTATION_DATE\n')
        for clinvar_record in clinvar_xml_utils.ClinVarDataset(clinvar_xml):
            process_clinvar_record(clinvar_record, outfile)

parser = argparse.ArgumentParser('Extracts OMIM and MedGen cross-references from ClinVar for submission to ZOOMA')
parser.add_argument('--clinvar-xml',    help='ClinVar XML release',                 required=True)
parser.add_argument('--zooma-feedback', help='Disease string to ontology mappings', required=True)

if __name__ == '__main__':
    args = parser.parse_args()
    main(args.clinvar_xml, args.zooma_feedback)
<commit_msg>Fixed parsing MONDO URIs from ClinVar data<commit_after>#!/usr/bin/env python3

import argparse
import itertools
from time import gmtime, strftime

from eva_cttv_pipeline import clinvar_xml_utils

class OntologyUri:
    # ClinVar stores cross-references in very different formats. This provides their conversion to full IRIs, along with
    # some examples of how this looks like in ClinVar data.
    db_to_uri_conversion = {
        'orphanet': lambda x: f'http://www.orpha.net/ORDO/Orphanet_{x}',  # <XRef ID="1756" DB="Orphanet"/>
        'omim': lambda x: f'https://www.omim.org/entry/{x}',  # <XRef Type="MIM" ID="612773" DB="OMIM"/>
        'efo': lambda x: f'http://www.ebi.ac.uk/efo/{x}',  # <XRef ID="EFO_0005137" DB="EFO"/>
        'mesh': lambda x: f'http://identifiers.org/mesh/{x}',  # <XRef ID="D065630" DB="MeSH"/>
        'medgen': lambda x: f'http://identifiers.org/medgen/{x}',  # <XRef ID="C0235833" DB="MedGen"/>
        # <XRef ID="MONDO:0013353" DB="MONDO"/>
        'mondo': lambda x: 'http://purl.obolibrary.org/obo/{}'.format(x.replace(':', '_')),
    }

    def __init__(self, id_, db):
        self.id_ = id_
        self.db = db
        self.uri = self.db_to_uri_conversion[self.db.lower()](self.id_)

    def __str__(self):
        return self.uri

def write_zooma_record(clinvar_acc, variant_id, trait_name, ontology_uri, date, outfile):
    zooma_output_list = [clinvar_acc,
                         variant_id,
                         'disease',
                         trait_name,
                         str(ontology_uri),
                         'clinvar-xrefs',
                         date]
    outfile.write('\t'.join(zooma_output_list) + '\n')

def process_clinvar_record(clinvar_record, outfile):
    """Extract the variant, trait and ontology from Clinvar record and write them as a ZOOMA feedback record."""
    if clinvar_record.measure is None:
        return
    variant_ids = [variant_id
                   for variant_id in (clinvar_record.measure.rs_id, clinvar_record.measure.nsv_id)
                   if variant_id is not None]
    traits = clinvar_record.traits
    for variant_id, trait in itertools.product(variant_ids, traits):
        if (trait.name is None) or (trait.name.lower() == 'not provided'):
            continue
        for db, identifier, status in trait.xrefs:
            if status != 'current' or db.lower() not in OntologyUri.db_to_uri_conversion:
                continue
            ontology_uri = OntologyUri(identifier, db)
            write_zooma_record(clinvar_record.accession, variant_id, trait.name, ontology_uri,
                               strftime('%d/%m/%y %H:%M', gmtime()), outfile)

def main(clinvar_xml, zooma_feedback):
    with open(zooma_feedback, 'wt') as outfile:
        outfile.write('STUDY\tBIOENTITY\tPROPERTY_TYPE\tPROPERTY_VALUE\tSEMANTIC_TAG\tANNOTATOR\tANNOTATION_DATE\n')
        for clinvar_record in clinvar_xml_utils.ClinVarDataset(clinvar_xml):
            process_clinvar_record(clinvar_record, outfile)

parser = argparse.ArgumentParser('Extracts OMIM and MedGen cross-references from ClinVar for submission to ZOOMA')
parser.add_argument('--clinvar-xml',    help='ClinVar XML release',                 required=True)
parser.add_argument('--zooma-feedback', help='Disease string to ontology mappings', required=True)

if __name__ == '__main__':
    args = parser.parse_args()
    main(args.clinvar_xml, args.zooma_feedback)
<|endoftext|>


<commit_before>import openpathsampling as paths
import collections
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt

try:
    from collections import abc
except ImportError:
    import abc

# based on http://stackoverflow.com/a/3387975
class TransformedDict(abc.MutableMapping):
    """A dictionary that applies an arbitrary key-altering function before
    accessing the keys

    This implementation involves a particular hashing function. It is
    assumed that any two input objects which give the same hash are
    effectively identical, allowing later rehashing based on the same.
    """

    def __init__(self, hash_function, *args, **kwargs):
        self.store = dict()
        self.hash_representatives = dict()
        self.hash_function = hash_function
        self.update(dict(*args, **kwargs))  # use the free update to set keys

    def __getitem__(self, key):
        return self.store[self.hash_function(key)]

    def __setitem__(self, key, value):
        hashed = self.hash_function(key)
        if hashed not in self.hash_representatives:
            self.hash_representatives[hashed] = key
        self.store[hashed] = value

    def __delitem__(self, key):
        hashed = self.hash_function(key)
        del self.store[hashed]
        del self.hash_representatives[hashed]

    def __iter__(self):
        return iter(self.store)

    def __len__(self):
        return len(self.store)

    def rehash(self, new_hash):
        """Create a new TransformedDict with this data and new hash.

        It is up to the user to ensure that the mapping from the old hash to
        the new is a function (i.e., each entry from the old hash can be
        mapped directly onto the new hash).

        For example, this is used to map from a snapshot's coordinates to
        a collective variable based on the coordinates. However, if the
        orignal hash was based on coordinates, but the new hash included
        velocities, the resulting mapping would be invalid. It is up to the
        user to avoid such invalid remappings.
        """
        return TransformedDict(new_hash,
                               {self.hash_representatives[k]: self.store[k]
                                for k in self.store})

class SnapshotByCoordinateDict(TransformedDict):
    """TransformedDict that uses snapshot coordinates as keys.

    This is primarily used to have a unique key for shooting point analysis
    (e.g., committor analysis).
    """
    def __init__(self, *args, **kwargs):
        hash_fcn = lambda x : x.coordinates.tostring()
        super(SnapshotByCoordinateDict, self).__init__(hash_fcn,
                                                       *args, **kwargs)

class ShootingPointAnalysis(SnapshotByCoordinateDict):
    """
    Container and methods for shooting point analysis.

    This is especially useful for analyzing committors, which is
    automatically done on a per-configuration basis, and can also be done
    as a histogram.

    Parameters
    ----------
    steps : iterable of :class:`.MCStep` or None
        input MC steps to analyze; if None, no analysis performed
    states : list of :class:`.Volume`
        volumes to consider as states for the analysis. For pandas output,
        these volumes must be named.
    """
    def __init__(self, steps, states):
        super(ShootingPointAnalysis, self).__init__()
        self.states = states
        if steps is not None:
            self.analyze(steps)

    def analyze(self, steps):
        """Analyze a list of steps, adding to internal results.

        Parameters
        ----------
        steps : iterable of :class:`.MCStep` or None
            MC steps to analyze
        """
        for step in steps:
            total = self.analyze_single_step(step)

    def analyze_single_step(self, step):
        """
        Analyzes final states from a path sampling step. Adds to internal
        results.

        Parameters
        ----------
        step : :class:`.MCStep`
            the step to analyze and add to this analysis

        Returns
        -------
        list of :class:`.Volume`
            the states which are identified as new final states from this
            move
        """
        key = self.step_key(step)
        if key is not None:
            details = step.change.canonical.details
            trial_traj = step.change.canonical.trials[0].trajectory
            init_traj = details.initial_trajectory
            test_points = [s for s in [trial_traj[0], trial_traj[-1]]
                           if s not in [init_traj[0], init_traj[-1]]]

            total = collections.Counter(
                {state: sum([int(state(pt)) for pt in test_points])
                            for state in self.states}
            )
            total_count = sum(total.values())
            # TODO: clarify assertion (at least one endpoint in state)
            assert total_count == 1 or total_count == 2
            try:
                self[key] += total
            except KeyError:
                self[key] = total
        else:
            total = {}

        return [s for s in total.keys() if total[s] > 0]

    @staticmethod
    def step_key(step):
        """
        Returns the key we use for hashing (the shooting snapshot).

        Parameters
        ----------
        step : :class:`.MCStep`
            the step to extract a shooting point from

        Returns
        -------
        :class:`.Snapshot` or None
            the shooting snapshot, or None if this step is not a shooting
            move.
        """
        key = None
        try:
            details = step.change.canonical.details
            shooting_snap = details.shooting_snapshot
        except AttributeError:
            # wrong kind of move (no shooting_snapshot)
            pass
        except IndexError:
            # very wrong kind of move (no trials!)
            pass
        else:
            # easy to change how we define the key
            key = shooting_snap
        return key

    @classmethod
    def from_individual_runs(cls, run_results, states=None):
        """Build shooting point analysis from pairs of shooting point to
        final state.

        Parameters
        ----------
        run_results : list of 2-tuples (:class:`.Snapshot`, :class:`.Volume`)
            the first element in each pair is the shooting point, the second
            is the final volume
        """
        if states is None:
            states = set(s[1] for s in run_results)
        analyzer = ShootingPointAnalysis(None, states)
        for step in run_results:
            key = step[0]
            total = collections.Counter({step[1] : 1})
            try:
                analyzer[key] += total
            except KeyError:
                analyzer[key] = total

        return analyzer

    def committor(self, state, label_function=None):
        """Calculate the (point-by-point) committor.

        This is for the point-by-point (per-configuration) committor, not
        for histograms. See `committor_histogram` for the histogram version.

        Parameters
        ----------
        state : :class:`.Volume`
            the committor is 1.0 if 100% of shots enter this state
        label_function : callable
            the keys for the dictionary that is returned are
            `label_function(snapshot)`; default `None` gives the snapshot as
            key.

        Returns
        -------
        dict :
            mapping labels given by label_function to the committor value
        """
        if label_function is None:
            label_function = lambda s : s
        results = {}
        for k in self:
            out_key = label_function(self.hash_representatives[k])
            counter_k = self.store[k]
            committor = float(counter_k[state]) / sum(counter_k.values())
            results[out_key] = committor
        return results

    @staticmethod
    def _get_key_dim(key):
        try:
            ndim = len(key)
        except TypeError:
            ndim = 1
        if ndim > 2 or ndim < 1:
            raise RuntimeError("Histogram key dimension {0} > 2 or {0} < 1 "
                               + "(key: {1})".format(ndim, key))
        return ndim

    def committor_histogram(self, new_hash, state, bins=10):
        """Calculate the histogrammed version of the committor.

        Parameters
        ----------
        new_hash : callable
            values are histogrammed in bins based on new_hash(snapshot)
        state : :class:`.Volume`
            the committor is 1.0 if 100% of shots enter this state
        bins : see numpy.histogram
            bins input to numpy.histogram

        Returns
        -------
        tuple :
            hist, bins like numpy.histogram, where hist is the histogram
            count and bins is the bins output from numpy.histogram. 2-tuple
            in the case of 1D histogram, 3-tuple in the case of 2D histogram
        """
        rehashed = self.rehash(new_hash)
        r_store = rehashed.store
        count_all = {k : sum(r_store[k].values()) for k in r_store}
        count_state = {k : r_store[k][state] for k in r_store}
        ndim = self._get_key_dim(list(r_store.keys())[0])
        if ndim == 1:
            (all_hist, b) = np.histogram(list(count_all.keys()),
                                         weights=list(count_all.values()),
                                         bins=bins)
            (state_hist, b) = np.histogram(list(count_state.keys()),
                                           weights=list(count_state.values()),
                                           bins=bins)
            b_list = [b]
        elif ndim == 2:
            (all_hist, b_x, b_y) = np.histogram2d(
                x=[k[0] for k in count_all],
                y=[k[1] for k in count_all],
                weights=list(count_all.values()),
                bins=bins
            )
            (state_hist, b_x, b_y) = np.histogram2d(
                x=[k[0] for k in count_state],
                y=[k[1] for k in count_state],
                weights=list(count_state.values()),
                bins=bins
            )
            b_list = [b_x, b_y]
        # if all_hist is 0, state_hist is NaN: ignore warning, return NaN
        with np.errstate(divide='ignore', invalid='ignore'):
            state_frac = np.true_divide(state_hist, all_hist)
        return tuple([state_frac] + b_list)

    def to_pandas(self, label_function=None):
        """
        Pandas dataframe. Row for each configuration, column for each state.

        Parameters
        ----------
        label_function : callable
            takes snapshot, returns index to use for pandas.DataFrame
        """
        transposed = pd.DataFrame(self.store).transpose().to_dict()
        df = pd.DataFrame(transposed)
        df.columns = [s.name for s in transposed.keys()]
        if label_function is None:
            df.index = range(len(df.index))
        else:
            # TODO: is ordering guaranteed here?
            df.index = [label_function(self.hash_representatives[k])
                        for k in self.store]
        return df
<commit_msg>really fix Py2 abc import<commit_after>import openpathsampling as paths
import collections
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt

try:
    from collections import abc
except ImportError:
    import collections as abc

# based on http://stackoverflow.com/a/3387975
class TransformedDict(abc.MutableMapping):
    """A dictionary that applies an arbitrary key-altering function before
    accessing the keys

    This implementation involves a particular hashing function. It is
    assumed that any two input objects which give the same hash are
    effectively identical, allowing later rehashing based on the same.
    """

    def __init__(self, hash_function, *args, **kwargs):
        self.store = dict()
        self.hash_representatives = dict()
        self.hash_function = hash_function
        self.update(dict(*args, **kwargs))  # use the free update to set keys

    def __getitem__(self, key):
        return self.store[self.hash_function(key)]

    def __setitem__(self, key, value):
        hashed = self.hash_function(key)
        if hashed not in self.hash_representatives:
            self.hash_representatives[hashed] = key
        self.store[hashed] = value

    def __delitem__(self, key):
        hashed = self.hash_function(key)
        del self.store[hashed]
        del self.hash_representatives[hashed]

    def __iter__(self):
        return iter(self.store)

    def __len__(self):
        return len(self.store)

    def rehash(self, new_hash):
        """Create a new TransformedDict with this data and new hash.

        It is up to the user to ensure that the mapping from the old hash to
        the new is a function (i.e., each entry from the old hash can be
        mapped directly onto the new hash).

        For example, this is used to map from a snapshot's coordinates to
        a collective variable based on the coordinates. However, if the
        orignal hash was based on coordinates, but the new hash included
        velocities, the resulting mapping would be invalid. It is up to the
        user to avoid such invalid remappings.
        """
        return TransformedDict(new_hash,
                               {self.hash_representatives[k]: self.store[k]
                                for k in self.store})

class SnapshotByCoordinateDict(TransformedDict):
    """TransformedDict that uses snapshot coordinates as keys.

    This is primarily used to have a unique key for shooting point analysis
    (e.g., committor analysis).
    """
    def __init__(self, *args, **kwargs):
        hash_fcn = lambda x : x.coordinates.tostring()
        super(SnapshotByCoordinateDict, self).__init__(hash_fcn,
                                                       *args, **kwargs)

class ShootingPointAnalysis(SnapshotByCoordinateDict):
    """
    Container and methods for shooting point analysis.

    This is especially useful for analyzing committors, which is
    automatically done on a per-configuration basis, and can also be done
    as a histogram.

    Parameters
    ----------
    steps : iterable of :class:`.MCStep` or None
        input MC steps to analyze; if None, no analysis performed
    states : list of :class:`.Volume`
        volumes to consider as states for the analysis. For pandas output,
        these volumes must be named.
    """
    def __init__(self, steps, states):
        super(ShootingPointAnalysis, self).__init__()
        self.states = states
        if steps is not None:
            self.analyze(steps)

    def analyze(self, steps):
        """Analyze a list of steps, adding to internal results.

        Parameters
        ----------
        steps : iterable of :class:`.MCStep` or None
            MC steps to analyze
        """
        for step in steps:
            total = self.analyze_single_step(step)

    def analyze_single_step(self, step):
        """
        Analyzes final states from a path sampling step. Adds to internal
        results.

        Parameters
        ----------
        step : :class:`.MCStep`
            the step to analyze and add to this analysis

        Returns
        -------
        list of :class:`.Volume`
            the states which are identified as new final states from this
            move
        """
        key = self.step_key(step)
        if key is not None:
            details = step.change.canonical.details
            trial_traj = step.change.canonical.trials[0].trajectory
            init_traj = details.initial_trajectory
            test_points = [s for s in [trial_traj[0], trial_traj[-1]]
                           if s not in [init_traj[0], init_traj[-1]]]

            total = collections.Counter(
                {state: sum([int(state(pt)) for pt in test_points])
                            for state in self.states}
            )
            total_count = sum(total.values())
            # TODO: clarify assertion (at least one endpoint in state)
            assert total_count == 1 or total_count == 2
            try:
                self[key] += total
            except KeyError:
                self[key] = total
        else:
            total = {}

        return [s for s in total.keys() if total[s] > 0]

    @staticmethod
    def step_key(step):
        """
        Returns the key we use for hashing (the shooting snapshot).

        Parameters
        ----------
        step : :class:`.MCStep`
            the step to extract a shooting point from

        Returns
        -------
        :class:`.Snapshot` or None
            the shooting snapshot, or None if this step is not a shooting
            move.
        """
        key = None
        try:
            details = step.change.canonical.details
            shooting_snap = details.shooting_snapshot
        except AttributeError:
            # wrong kind of move (no shooting_snapshot)
            pass
        except IndexError:
            # very wrong kind of move (no trials!)
            pass
        else:
            # easy to change how we define the key
            key = shooting_snap
        return key

    @classmethod
    def from_individual_runs(cls, run_results, states=None):
        """Build shooting point analysis from pairs of shooting point to
        final state.

        Parameters
        ----------
        run_results : list of 2-tuples (:class:`.Snapshot`, :class:`.Volume`)
            the first element in each pair is the shooting point, the second
            is the final volume
        """
        if states is None:
            states = set(s[1] for s in run_results)
        analyzer = ShootingPointAnalysis(None, states)
        for step in run_results:
            key = step[0]
            total = collections.Counter({step[1] : 1})
            try:
                analyzer[key] += total
            except KeyError:
                analyzer[key] = total

        return analyzer

    def committor(self, state, label_function=None):
        """Calculate the (point-by-point) committor.

        This is for the point-by-point (per-configuration) committor, not
        for histograms. See `committor_histogram` for the histogram version.

        Parameters
        ----------
        state : :class:`.Volume`
            the committor is 1.0 if 100% of shots enter this state
        label_function : callable
            the keys for the dictionary that is returned are
            `label_function(snapshot)`; default `None` gives the snapshot as
            key.

        Returns
        -------
        dict :
            mapping labels given by label_function to the committor value
        """
        if label_function is None:
            label_function = lambda s : s
        results = {}
        for k in self:
            out_key = label_function(self.hash_representatives[k])
            counter_k = self.store[k]
            committor = float(counter_k[state]) / sum(counter_k.values())
            results[out_key] = committor
        return results

    @staticmethod
    def _get_key_dim(key):
        try:
            ndim = len(key)
        except TypeError:
            ndim = 1
        if ndim > 2 or ndim < 1:
            raise RuntimeError("Histogram key dimension {0} > 2 or {0} < 1 "
                               + "(key: {1})".format(ndim, key))
        return ndim

    def committor_histogram(self, new_hash, state, bins=10):
        """Calculate the histogrammed version of the committor.

        Parameters
        ----------
        new_hash : callable
            values are histogrammed in bins based on new_hash(snapshot)
        state : :class:`.Volume`
            the committor is 1.0 if 100% of shots enter this state
        bins : see numpy.histogram
            bins input to numpy.histogram

        Returns
        -------
        tuple :
            hist, bins like numpy.histogram, where hist is the histogram
            count and bins is the bins output from numpy.histogram. 2-tuple
            in the case of 1D histogram, 3-tuple in the case of 2D histogram
        """
        rehashed = self.rehash(new_hash)
        r_store = rehashed.store
        count_all = {k : sum(r_store[k].values()) for k in r_store}
        count_state = {k : r_store[k][state] for k in r_store}
        ndim = self._get_key_dim(list(r_store.keys())[0])
        if ndim == 1:
            (all_hist, b) = np.histogram(list(count_all.keys()),
                                         weights=list(count_all.values()),
                                         bins=bins)
            (state_hist, b) = np.histogram(list(count_state.keys()),
                                           weights=list(count_state.values()),
                                           bins=bins)
            b_list = [b]
        elif ndim == 2:
            (all_hist, b_x, b_y) = np.histogram2d(
                x=[k[0] for k in count_all],
                y=[k[1] for k in count_all],
                weights=list(count_all.values()),
                bins=bins
            )
            (state_hist, b_x, b_y) = np.histogram2d(
                x=[k[0] for k in count_state],
                y=[k[1] for k in count_state],
                weights=list(count_state.values()),
                bins=bins
            )
            b_list = [b_x, b_y]
        # if all_hist is 0, state_hist is NaN: ignore warning, return NaN
        with np.errstate(divide='ignore', invalid='ignore'):
            state_frac = np.true_divide(state_hist, all_hist)
        return tuple([state_frac] + b_list)

    def to_pandas(self, label_function=None):
        """
        Pandas dataframe. Row for each configuration, column for each state.

        Parameters
        ----------
        label_function : callable
            takes snapshot, returns index to use for pandas.DataFrame
        """
        transposed = pd.DataFrame(self.store).transpose().to_dict()
        df = pd.DataFrame(transposed)
        df.columns = [s.name for s in transposed.keys()]
        if label_function is None:
            df.index = range(len(df.index))
        else:
            # TODO: is ordering guaranteed here?
            df.index = [label_function(self.hash_representatives[k])
                        for k in self.store]
        return df
<|endoftext|>