<commit_before>"""
These "scans" bundle a Message generator with an instance of the RunEngine,
combining two separate concepts -- instructions and execution -- into one
object. This makes the interface less flexible and somewhat less "Pythonic"
but more convenient to users who are not acclimated to using tab-completion and
"up-arrow" for quick iteration during experiments.

This module is meant to be run in a namespace where several global
variables have been defined. If some variables are left undefined, the
associated scans will be not usable.

    dets
    h_motor
    k_motor
    l_motor
    th_motor
    tth_motor
    temp_controller

Page numbers in the code comments refer to the SPEC manual at
http://www.certif.com/downloads/css_docs/spec_man.pdf
"""
from inspect import signature
from bluesky import scans
from boltons.iterutils import chunked

class _PrimitiveScan:

    def __init__(self, RE):
        self.RE = RE
        self.params = list(signature(self._scan_class).parameters.keys())
        RE_params = list(signature(RE.__call__).parameters.keys())
        if set(RE_params) & set(self.params):
            raise AssertionError("The names of the scan's arguments clash "
                                 "the RunEngine arguments. Use different "
                                 "names. Avoid: {0}".format(RE_params))

    def __call__(self, *args, subs=None, **kwargs):
        global dets
        scan_kwargs = dict()
        # Any kwargs valid for the scan go to the scan, not the RE.
        for k, v in kwargs.items():
            if k in self.params:
                scan_kwargs[k] = kwargs.pop(k)
        scan = self._scan_class(dets, *args, **scan_kwargs)
        self.RE(scan, **kwargs)  # to pass to RE, you must use kwargs

### Mid-level base classes ###

# These are responsible for popping off the time arg and adjusting the
# interval. SPEC counts "bonds;" idiomatic Python counts "sites."

class _OuterProductScan(_PrimitiveScan):
    def __call__(self, *args, **kwargs):
        time = args.pop(-1)
        _set_acquire_time(time)
        for i, _ in enumerate(chunked(args, 4)):
            # intervals -> intervals + 1
            args[4*i + 3] += 1
        super.__call__(*args, **kwargs)

class _InnerProductScan(_PrimitiveScan):
    def __call__(self, *args, **kwargs):
        time = args.pop(-1)
        intervals = args.pop(-1) + 1
        _set_acquire_time(time)
        super.__call__(intervals, *args, **kwargs)

class _StepScan(_PrimitiveScan):

    def __call__(self, motor, start, finish, intervals, time, *, **kwargs):
        _set_acquire_time(time)
        super.__call__(motor, start, finish, intervals + 1, **kwargs)

class _HardcodedMotorStepScan(_PrimitiveScan):
    # Subclasses must define self._motor as a property.

    def __call__(self, start, finish, intervals, time, *, **kwargs):
        _set_acquire_time(time)
        super.__call__(self._motor, start, finish, intervals + 1, **kwargs)

### Counts (p. 140) ###

class Count(_PrimitiveScan):
    "ct"
    _scan_class = scans.Count

    def __call__(self, time=1, *, **kwargs):
        _set_acquire_time(time)
        super.__call__(**kwargs)

### Motor Scans (p. 146) ###

class AbsoluteScan(_StepScan):
    "ascan"
    _scan_class = scans.LinAscan

class OuterProductAbsoluteScan(_OuterProductScan):
    "mesh"
    _scan_class = scans.OuterProductAbsoluteScan

class InnerProductAbsoluteScan(_InnerProductScan):
    "a2scan, a3scan, etc."
    _scan_class = scans.InnerProductAbsoluteScan

class DeltaScan(_StepScan):
    "dscan (also known as lup)"
    _scan_class = scans.LinDscan

class InnerProductDeltaScan(_InnerProductScan):
    "d2scan, d3scan, etc."
    _scan_class = scans.InnerProductDeltaScan

class ThetaTwoThetaScan(_InnerProductScan):
    "th2th"
    _scan_class = scans.InnerProductDeltaScan

    def __call__(self, start, finish, intervals, time, *, **kwargs):
        global th_motor
        global tth_motor
        super.__call__(tth_motor, start, finish,
                       th_motor, start/2, finish/2, intervals, time, **kwargs)

### Temperature Scans (p. 148) ###

class _TemperatureScan(_HardcodedMotorStepScan):
    @property
    def _motor(self):
        global temp_controller
        return temp_controller

class AbsoluteTemperatureScan(_TemperatureScan):
    "tscan"
    _scan_class = scans.LinAscan

class DeltaTemperatureScan(_TemperatureScan):
    "dtscan"
    _scan_class = scans.LinDscan

### Basic Reciprocal Space Scans (p. 147) ###

class HScan(_HardcodedMotorStepScan):
    "hscan"
    _scan_class = scans.LinAscan

    @property
    def _motor(self):
        global h_motor
        return h_motor

class KScan(_HardcodedMotorStepScan):
    "kscan"
    _scan_class = scans.LinAscan

    @property
    def _motor(self):
        global l_motor
        return l_motor

class LScan(_HardcodedMotorStepScan):
    "lscan"
    _scan_class = scans.LinAscan

    @property
    def _motor(self):
        global l_motor
        return l_motor

class OuterProductHKLScan(_PrimitiveScan):
    "hklmesh"
    _scan_class = scans.OuterProductAbsoluteScan

    def __call__(self, Q1, start1, finish1, intervals1, Q2, start2, finish2,
                 intervals2, time, *, **kwargs):
        # To be clear, like all other functions in this module, this
        # eye-gouging API is for compatbility with SPEC, not the author's
        # idea of good Python code.
        global h_motor, k_motor, l_motor
        _set_acquire_time(time)
        _motor_mapping = {'H': h_motor, 'K': k_motor, 'L': l_motor}
        motor1 = _motor_mapping[Q1]
        motor2 = _motor_mapping[Q2]
        # Note that intervals + 1 is handled in the base class.
        super.__call__(motor1, start1, finish1, intervals1,
                       motor2, start2, finish2, intervals2, **kwargs)

class InnerProductHKLScan(_PrimitiveScan):
    "hklscan"
    _scan_class = scans.InnerProductAbsoluteScan

    def __call__(self, start_h, finish_h, start_k, finish_k, start_l, finish_l,
                 intervals, time, *, **kwargs):
        global h_motor, k_motor, l_motor
        _set_acquire_time(time)
        super.__call__(self, intervals, start_h, finish_h, start_k, finish_k,
                       start_l, finish_l, **kwargs)

### Special Reciprocal Space Scans ###

# TODO:
# klradial
# hlradial
# hkradial
# klcircle
# hlcircle
# hkcircle

def _set_acquire_time(time):
    global dets
    for det in dets:
        det.acquire_time = time
<commit_msg>DOC: Comment on global variables.<commit_after>"""
These "scans" bundle a Message generator with an instance of the RunEngine,
combining two separate concepts -- instructions and execution -- into one
object. This makes the interface less flexible and somewhat less "Pythonic"
but more convenient to users who are not acclimated to using tab-completion and
"up-arrow" for quick iteration during experiments.

This module is meant to be run in a namespace where several global
variables have been defined. If some variables are left undefined, the
associated scans will be not usable.

    dets  # list of detectors
    master_det  # detector to use for tw
    master_det_field  # detector field to use for tw
    h_motor
    k_motor
    l_motor
    th_motor
    tth_motor
    temp_controller

Page numbers in the code comments refer to the SPEC manual at
http://www.certif.com/downloads/css_docs/spec_man.pdf
"""
from inspect import signature
from bluesky import scans
from boltons.iterutils import chunked

class _PrimitiveScan:

    def __init__(self, RE):
        self.RE = RE
        self.params = list(signature(self._scan_class).parameters.keys())
        RE_params = list(signature(RE.__call__).parameters.keys())
        if set(RE_params) & set(self.params):
            raise AssertionError("The names of the scan's arguments clash "
                                 "the RunEngine arguments. Use different "
                                 "names. Avoid: {0}".format(RE_params))

    def __call__(self, *args, subs=None, **kwargs):
        global dets
        scan_kwargs = dict()
        # Any kwargs valid for the scan go to the scan, not the RE.
        for k, v in kwargs.items():
            if k in self.params:
                scan_kwargs[k] = kwargs.pop(k)
        scan = self._scan_class(dets, *args, **scan_kwargs)
        self.RE(scan, **kwargs)  # to pass to RE, you must use kwargs

### Mid-level base classes ###

# These are responsible for popping off the time arg and adjusting the
# interval. SPEC counts "bonds;" idiomatic Python counts "sites."

class _OuterProductScan(_PrimitiveScan):
    def __call__(self, *args, **kwargs):
        time = args.pop(-1)
        _set_acquire_time(time)
        for i, _ in enumerate(chunked(args, 4)):
            # intervals -> intervals + 1
            args[4*i + 3] += 1
        super.__call__(*args, **kwargs)

class _InnerProductScan(_PrimitiveScan):
    def __call__(self, *args, **kwargs):
        time = args.pop(-1)
        intervals = args.pop(-1) + 1
        _set_acquire_time(time)
        super.__call__(intervals, *args, **kwargs)

class _StepScan(_PrimitiveScan):

    def __call__(self, motor, start, finish, intervals, time, *, **kwargs):
        _set_acquire_time(time)
        super.__call__(motor, start, finish, intervals + 1, **kwargs)

class _HardcodedMotorStepScan(_PrimitiveScan):
    # Subclasses must define self._motor as a property.

    def __call__(self, start, finish, intervals, time, *, **kwargs):
        _set_acquire_time(time)
        super.__call__(self._motor, start, finish, intervals + 1, **kwargs)

### Counts (p. 140) ###

class Count(_PrimitiveScan):
    "ct"
    _scan_class = scans.Count

    def __call__(self, time=1, *, **kwargs):
        _set_acquire_time(time)
        super.__call__(**kwargs)

### Motor Scans (p. 146) ###

class AbsoluteScan(_StepScan):
    "ascan"
    _scan_class = scans.LinAscan

class OuterProductAbsoluteScan(_OuterProductScan):
    "mesh"
    _scan_class = scans.OuterProductAbsoluteScan

class InnerProductAbsoluteScan(_InnerProductScan):
    "a2scan, a3scan, etc."
    _scan_class = scans.InnerProductAbsoluteScan

class DeltaScan(_StepScan):
    "dscan (also known as lup)"
    _scan_class = scans.LinDscan

class InnerProductDeltaScan(_InnerProductScan):
    "d2scan, d3scan, etc."
    _scan_class = scans.InnerProductDeltaScan

class ThetaTwoThetaScan(_InnerProductScan):
    "th2th"
    _scan_class = scans.InnerProductDeltaScan

    def __call__(self, start, finish, intervals, time, *, **kwargs):
        global th_motor
        global tth_motor
        super.__call__(tth_motor, start, finish,
                       th_motor, start/2, finish/2, intervals, time, **kwargs)

### Temperature Scans (p. 148) ###

class _TemperatureScan(_HardcodedMotorStepScan):
    @property
    def _motor(self):
        global temp_controller
        return temp_controller

class AbsoluteTemperatureScan(_TemperatureScan):
    "tscan"
    _scan_class = scans.LinAscan

class DeltaTemperatureScan(_TemperatureScan):
    "dtscan"
    _scan_class = scans.LinDscan

### Basic Reciprocal Space Scans (p. 147) ###

class HScan(_HardcodedMotorStepScan):
    "hscan"
    _scan_class = scans.LinAscan

    @property
    def _motor(self):
        global h_motor
        return h_motor

class KScan(_HardcodedMotorStepScan):
    "kscan"
    _scan_class = scans.LinAscan

    @property
    def _motor(self):
        global l_motor
        return l_motor

class LScan(_HardcodedMotorStepScan):
    "lscan"
    _scan_class = scans.LinAscan

    @property
    def _motor(self):
        global l_motor
        return l_motor

class OuterProductHKLScan(_PrimitiveScan):
    "hklmesh"
    _scan_class = scans.OuterProductAbsoluteScan

    def __call__(self, Q1, start1, finish1, intervals1, Q2, start2, finish2,
                 intervals2, time, *, **kwargs):
        # To be clear, like all other functions in this module, this
        # eye-gouging API is for compatbility with SPEC, not the author's
        # idea of good Python code.
        global h_motor, k_motor, l_motor
        _set_acquire_time(time)
        _motor_mapping = {'H': h_motor, 'K': k_motor, 'L': l_motor}
        motor1 = _motor_mapping[Q1]
        motor2 = _motor_mapping[Q2]
        # Note that intervals + 1 is handled in the base class.
        super.__call__(motor1, start1, finish1, intervals1,
                       motor2, start2, finish2, intervals2, **kwargs)

class InnerProductHKLScan(_PrimitiveScan):
    "hklscan"
    _scan_class = scans.InnerProductAbsoluteScan

    def __call__(self, start_h, finish_h, start_k, finish_k, start_l, finish_l,
                 intervals, time, *, **kwargs):
        global h_motor, k_motor, l_motor
        _set_acquire_time(time)
        super.__call__(self, intervals, start_h, finish_h, start_k, finish_k,
                       start_l, finish_l, **kwargs)

### Special Reciprocal Space Scans ###

# TODO:
# klradial
# hlradial
# hkradial
# klcircle
# hlcircle
# hkcircle

def _set_acquire_time(time):
    global dets
    for det in dets:
        det.acquire_time = time
<|endoftext|>


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

import sys
import os
import copy

import numpy as np
import scipy as sp
import scipy.stats

from collections import OrderedDict

import camoco as co
import pandas as pd
import matplotlib.pylab as plt
import statsmodels.api as sm

lowess = sm.nonparametric.lowess

def mean_confidence_interval(data):
    return np.mean(data), confidence_interval(data)

def confidence_interval(data, confidence=0.95):
    a = 1.0*np.array(data)
    n = len(a)
    m, se = np.mean(a), scipy.stats.sem(a)
    h = se * sp.stats.t._ppf((1+confidence)/2., n-1)
    return 1.96*se

class NearestDict(OrderedDict):
    '''
        This extension overrides the get item method
        of dict where if a key does not exist, it returns
        the nearst key which does exist.
    '''
    def __getitem__(self,key):
        'Returns the nearest key which exists'
        return dict.__getitem__(self,min(self.keys(),key=lambda x: abs(x-key)))

def locality(args):
    # Initiate for args
    # Generate output dirs

    ont = co.GWAS(args.GWAS)

    # If all, grab a generater
    if args.term == 'all':
        terms = ont.iter_terms()
    else:
        # Otherwise get the term out of the GWAS
        terms = [ont[args.term]]

    # Grab the COB object
    cob = co.COB(args.cob)

    # Create a plot comparing spread of specified shit
    xaxes_key,yaxes_key = [x.replace('-','_') for x in args.axes]
    xaxes = getattr(args,xaxes_key)
    yaxes = getattr(args,yaxes_key)

    #Open a log file
    gs = plt.GridSpec(len(xaxes),len(yaxes))
    fig = plt.figure(figsize=(4*len(yaxes),4*len(xaxes)))

    # This gets a list of functions so we can map the string passed into
    # args.plot to a plotting function
    possibles = globals().copy()
    possibles.update(locals())

    # Add in text for axes
    for term in terms:
        for plot in args.plot:
            for i,xaxis in enumerate(xaxes):
                for j,yaxis in enumerate(yaxes):
                    ax = fig.add_subplot(gs[i,j])
                    print("Looking at {}:{} x {}:{}".format(
                        xaxes_key,xaxis,
                        yaxes_key,yaxis,
                    ))
                    perm_args = copy.deepcopy(args)
                    # Set the values we aren't permuting to their first value
                    # I hate this
                    # for arg in perm_args.permutable:
                    #    if arg not in [xaxes_key,yaxes_key]:
                    #        setattr(perm_args,arg,getattr(args,arg)[1])
                    # Change the relevent things in the permuted args
                    setattr(perm_args,xaxes_key,xaxis)
                    setattr(perm_args,yaxes_key,yaxis)
                    # Generate data using permuted arguments
                    loc,bsloc,fdr = generate_data(cob,term,perm_args)
                    # Get the plot function based on args
                    plot_function = possibles.get('plot_'+plot)
                    # Plot the data
                    plot_function(perm_args,loc,bsloc,fdr,ax)
                    if i == 0:
                        ax.set_title(yaxis)
                    if j == 0:
                        ax.set_ylabel(xaxis)
                    if i == 0 and j == 0:
                        ax.set_title(str(yaxes_key)+' '+str(yaxis))
                        ax.set_ylabel(str(xaxes_key)+' '+str(xaxis))
            plt.tight_layout()
            plt.savefig("Locality_{}_{}.png".format(
                plot,
                term.id
            ))

def generate_data(cob,term,args):
    '''
        Generates the data according to parameters in args
    '''

    effective_loci = term.effective_loci(
            window_size=args.candidate_window_size
    )

    candidate_genes = cob.refgen.candidate_genes(
        effective_loci,
        gene_limit=args.candidate_gene_limit
    )

    # Find the empirical Locality
    loc = cob.locality(
        candidate_genes,
        include_regression=True
    )

    # Find the Bootstrapped Locality
    bsloc = pd.concat(
            [cob.locality(
                cob.refgen.bootstrap_candidate_genes(
                    effective_loci,
                    gene_limit=args.candidate_gene_limit
                ),
                bootstrap_name=x
            ) for x in range(args.num_bootstraps)]
    )

    '''---------------------------------------------------
        Empirical and SD Calculations
    '''

    # We need to perform regression for the entire bootstrap dataset
    OLS = sm.OLS(bsloc['local'],bsloc['global']).fit()
    bsloc['fitted'] = OLS.fittedvalues
    bsloc['resid'] = OLS.resid
    # Remove global degree larger than empirical
    bslowess = lowess(
            bsloc['resid'],
            bsloc['fitted'],
            frac=0.15,
            it=5,
            delta=0.1*len(bsloc),
            is_sorted=False
    )

    # Windowing
    bsloc = bsloc.sort('fitted')
    # Find out how many tick there are with X items per window
    num_windows = len(bsloc)//args.window_size
    window_ticks = len(bsloc)//num_windows
    bsloc['window'] = [int(i/window_ticks) for i in range(len(bsloc))]
    # If there are not many in the last window, change it second to last
    max_window = max(bsloc['window'])
    if sum(bsloc['window'] == max_window) < args.window_size / 2:
        bsloc.loc[bsloc['window'] == max_window, 'window'] = max_window-1
    # create a dictionary so we can map the empirical data later
    win_map = NearestDict({
        # Good god this is a hack --
        # group the pandas df by window and calculate the mean fitted value,
        # create dict from that and reverser keys and values
        fitted:window for window,fitted in bsloc.groupby('window').apply(
                lambda df: np.max(df['fitted'])
            ).to_dict().items()
    })
    # create a dict of window to std mapping
    win_std = bsloc.groupby('window').apply(lambda df: df['resid'].std())
    # perform lowess on the std_windows
    win_std = NearestDict(
        win_std.to_dict()
    )
    fit_std = {f:win_std[w]for f,w in win_map.items()}

    # Create a dict where keys are fitted valus and
    # values are that fitted values std
    fit_std = NearestDict(
        pd.DataFrame(
            lowess(
                np.array(list(fit_std.values())),
                np.array(list(fit_std.keys())),
                is_sorted=False
            ),columns=['fitted','sd']
        ).sort('fitted').set_index('fitted').to_dict()['sd']
    )

    # Calculate the s.d. of the residuals in each window
    # divide empirical residuals by the s.d. in their respective window
    loc['bs_std'] = [fit_std[x] for x in loc['fitted']]
    loc['zscore'] = [x['resid']/x['bs_std'] for i,x in loc.iterrows()]

    '''---------------------------------------------------
        FDR Calculations
    '''
    # Repeat bootstraps to assess global FDR
    fdr = pd.concat(
            [cob.locality(
                cob.refgen.bootstrap_candidate_genes(
                    term.effective_loci(
                        window_size=args.candidate_window_size
                    ),
                    gene_limit=args.candidate_gene_limit
                ),
                bootstrap_name=x,
                include_regression=True
            ) for x in range(args.num_bootstraps)]
    ).sort('global')
    OLS = sm.OLS(fdr['local'],fdr['global']).fit()

    # Remove global degree larger than empirical
    fdr = fdr[fdr['global'] <= max(loc['global'])]
    fdr['window'] = [int(x/window_ticks) for x in fdr['global']]

    # calculate z-scores for the global
    fdr['bs_std'] = [fit_std[x] for x in fdr['fitted']]
    fdr['zscore'] = [x['resid']/x['bs_std'] for i,x in fdr.iterrows()]

    # Give em the gold
    return loc,bsloc,fdr

def plot_scatter(args,loc,bsloc,fdr,ax):

    ''' ---------------------------------------------------
        Plotting
    '''
    # Y axis is local degree (what we are TRYING to predict)
    ax.set_ylim(0,max(loc['local']))
    ax.set_xlim(0,max(loc['global']))
    ax.set_xlabel('Number Global Interactions')
    ax.set_ylabel('Number Local Interactions')

    # UGH! map lowess
    fdrlowess = lowess(
        fdr['local'],fdr['global'],
        frac=0.15,it=5,delta=0.1*len(fdr),
        is_sorted=False
    )
    # plot the bootstrap points
    ax.plot(fdr['global'],fdr['local'],'ro',alpha=0.05,label='Bootstraps')
    # plot the OLS lowess line
    ci = fdr.groupby('window')['fitted','global'].agg(
        [np.mean,confidence_interval]
    )

    #for win,df in fdr.groupby('bootstrap_name'):
    #    ax.plot(df['global'],df['fitted'],alpha=1)

    ax.errorbar(
        ci['global','mean'],ci['fitted','mean'],
        yerr=ci['fitted','confidence_interval'],
        color='g',label='Bootstrap OLS'
    )

    #plot the empirical data
    ax.plot(loc['global'],loc['local'],'bo',alpha=1,label='Empirical')
    ax.plot(loc['global'],loc['fitted'],'k-',alpha=1,label='Empirical OLS')
    # finish plots
    #legend = ax.legend(loc='best')

    '''---------------------------------------------------
        FDR Plotting
    '''

def plot_fdr(args,loc,bsloc,fdr,ax):
    # Plot the empirical Z-score distributions
    zscores = list(range(1,8))
    zloc = [
        sum(np.logical_and(
            loc['zscore'] >= x ,
            loc['local'] >= args.min_fdr_degree
        ))
        for x in zscores
    ]
    ax.plot(zscores,zloc,'bo',label='Empirical Zscore > x')
    # plot the fdr scores spread
    zcdf = pd.DataFrame(
        [ mean_confidence_interval(
            fdr.groupby('bootstrap_name').apply(
                lambda df: sum(np.logical_and(
                    df['zscore'] >= x,
                    df['local'] >= args.min_fdr_degree
                ))
            )
        ) for x in zscores ],
        columns=['mean','ci']
    )
    #zcdf['emp'] = zloc
    #zcdf['el'] = args.term
    #zcdf['min_fdr_degree'] = args.min_fdr_degree
    #zcdf['can_gene_limit'] = args.candidate_gene_limit
    ax.errorbar(
        zscores,
        zcdf['mean'],
        yerr=zcdf['ci'],
        label='Bootstrap Z-scores',
        color='red'
    )
    ax.set_xlabel('Z-Score')
    ax.set_ylabel('Number of Genes > Z')
    ax.set_title('Z Score FDR')

def plot_resid(args,loc,bsloc,fdr,ax):
    resids = list(range(0,int(max(loc['resid'])+2)))
    rcdf = [sum(loc['resid'] >= x) for x in resids]
    ax.plot(resids,rcdf,'bo',label='Empirical Residuals')
    rcdf = pd.DataFrame(
         [ mean_confidence_interval(
            fdr.groupby('bootstrap_name').apply(
                lambda df: sum(df['resid'] >= x )
            )
        ) for x in resids ],
        columns=['mean','ci']
    )
    ax.errorbar(
        resids,
        rcdf['mean'],
        yerr=rcdf['ci'],
        label='Bootstrap Residuals',
        color='red'
    )
    ax.set_xlabel('Residual')
    ax.set_ylabel('Number of Genes > X')
    ax.set_title('Residual FDR')
    legend = ax.legend(loc='best')

    # Plot the predicted vs residuals

    '''---------------------------------------------------
        Residual Plotting
    '''

    ax.set_xlabel('Fitted')
    ax.set_ylabel('Residual')
    ax.set_title('Model Fitting')
    ax.plot(loc['fitted'],loc['resid'],'bo')
    # plot the Z-scores for the std
    ax.plot(list(fit_std.keys()),list(fit_std.values()),'g.',marker='.')
    # plot where windows are

def plot_data(args,loc,bsloc,fdr,ax):
    ax.xaxis.set_visible(False)
    ax.yaxis.set_visible(False)
    ax.text(.1,.1,'''
        COB: {}
        Ontology: {}
        Term: {}
        Num Raw SNPS: {}
        Num Effective SNPS: {}
        Num Genes: {}
        CI Window Size: {}
        Num. Bootstraps: {}
        Min FDR Degree: {}
        Candidate Gene Window Size: {}
        Candidate Gene Limit: {}
        Num SNPs: {}
    '''.format(
        args.cob,
        args.ontology,
        args.term,
        len(term.locus_list),
        len(effective_loci),
        len(candidate_genes),
        args.window_size,
        args.num_bootstraps,
        args.min_fdr_degree,
        args.candidate_window_size,
        args.candidate_gene_limit,
        args.num_snps)
    )

    out_string = '_'.join(map(str,[
        args.term,
        args.cob,
        args.window_size,
        args.min_fdr_degree,
        args.candidate_gene_limit,
        args.candidate_window_size
    ]))
<commit_msg>Locality working.<commit_after>#!/usr/bin/env python3

import sys
import os
import copy

import numpy as np
import scipy as sp
import scipy.stats

from collections import OrderedDict

import camoco as co
import pandas as pd
import matplotlib.pylab as plt
import statsmodels.api as sm

lowess = sm.nonparametric.lowess
co.cf.logging.log_level = 'quiet'

def mean_confidence_interval(data):
    return np.mean(data), confidence_interval(data)

def confidence_interval(data, confidence=0.95):
    a = 1.0*np.array(data)
    n = len(a)
    m, se = np.mean(a), scipy.stats.sem(a)
    h = se * sp.stats.t._ppf((1+confidence)/2., n-1)
    return 1.96*se

class NearestDict(OrderedDict):
    '''
        This extension overrides the get item method
        of dict where if a key does not exist, it returns
        the nearst key which does exist.
    '''
    def __getitem__(self,key):
        'Returns the nearest key which exists'
        return dict.__getitem__(self,min(self.keys(),key=lambda x: abs(x-key)))

def locality(args):
    # Initiate for args
    # Generate output dirs

    # Grab the COB object
    cob = co.COB(args.cob)
    gwas = co.GWAS(args.gwas)

    # If all, grab a generater
    if 'all' in args.terms:
        terms = gwas.iter_terms()
    else:
        # Otherwise get the term out of the GWAS
        terms = [gwas[x] for x in args.terms]

    # Add in text for axes
    for term in terms:
        fig,(ax1,ax2,ax3) = plt.subplots(1,3)
        # Change the relevent things in the permuted args
        # Generate data using permuted arguments
        loc,bsloc,fdr = generate_data(cob,term,args)
        # Plot the data
        plot_data(args,loc,bsloc,fdr,ax1)
        plot_scatter(args,loc,bsloc,fdr,ax2)
        plot_fdr(args,loc,bsloc,fdr,ax3)
        plt.tight_layout()
        plt.savefig("{}_{}.png".format(
            args.out,
            term.id
        ))
        plt.close()

def generate_data(cob,term,args):
    '''
        Generates the data according to parameters in args
    '''
    if args.snp2gene == 'effective':
        loci = sorted(term.effective_loci(
            window_size=args.candidate_window_size
        ))
    elif args.snp2gene == 'strongest':
        loci = term.strongest_loci(
            window_size=args.candidate_window_size,
            attr=args.strongest_attr,
            lowest=args.strongest_higher
        )
    else:
        raise ValueError('{} not valid snp2gene mapping'.format(args.snp2gene))

    candidate_genes = cob.refgen.candidate_genes(
        loci,
        flank_limit=args.candidate_flank_limit
    )

    # Find the empirical Locality
    loc = cob.locality(
        candidate_genes,
        include_regression=True
    )

    # Find the Bootstrapped Locality
    bsloc = pd.concat(
            [cob.locality(
                cob.refgen.bootstrap_candidate_genes(
                    loci,
                    flank_limit=args.candidate_flank_limit
                ),
                bootstrap_name=x
            ) for x in range(args.num_bootstraps)]
    )

    '''---------------------------------------------------
        Empirical and SD Calculations
    '''

    # We need to perform regression for the entire bootstrap dataset
    OLS = sm.OLS(bsloc['local'],bsloc['global']).fit()
    bsloc['fitted'] = OLS.fittedvalues
    bsloc['resid'] = OLS.resid
    # Remove global degree larger than empirical
    bslowess = lowess(
            bsloc['resid'],
            bsloc['fitted'],
            frac=0.15,
            it=5,
            delta=0.1*len(bsloc),
            is_sorted=False
    )

    # Windowing
    bsloc = bsloc.sort('fitted')
    # Find out how many tick there are with X items per window
    num_windows = len(bsloc)//args.regression_window_size
    window_ticks = len(bsloc)//num_windows
    bsloc['window'] = [int(i/window_ticks) for i in range(len(bsloc))]
    # If there are not many in the last window, change it second to last
    max_window = max(bsloc['window'])
    if sum(bsloc['window'] == max_window) < args.regression_window_size / 2:
        bsloc.loc[bsloc['window'] == max_window, 'window'] = max_window-1
    # create a dictionary so we can map the empirical data later
    win_map = NearestDict({
        # Good god this is a hack --
        # group the pandas df by window and calculate the mean fitted value,
        # create dict from that and reverser keys and values
        fitted:window for window,fitted in bsloc.groupby('window').apply(
                lambda df: np.max(df['fitted'])
            ).to_dict().items()
    })
    # create a dict of window to std mapping
    win_std = bsloc.groupby('window').apply(lambda df: df['resid'].std())
    # perform lowess on the std_windows
    win_std = NearestDict(
        win_std.to_dict()
    )
    fit_std = {f:win_std[w]for f,w in win_map.items()}

    # Create a dict where keys are fitted valus and
    # values are that fitted values std
    fit_std = NearestDict(
        pd.DataFrame(
            lowess(
                np.array(list(fit_std.values())),
                np.array(list(fit_std.keys())),
                is_sorted=False
            ),columns=['fitted','sd']
        ).sort('fitted').set_index('fitted').to_dict()['sd']
    )

    # Calculate the s.d. of the residuals in each window
    # divide empirical residuals by the s.d. in their respective window
    loc['bs_std'] = [fit_std[x] for x in loc['fitted']]
    loc['zscore'] = [x['resid']/x['bs_std'] for i,x in loc.iterrows()]

    '''---------------------------------------------------
        FDR Calculations
    '''
    # Repeat bootstraps to assess global FDR
    fdr = pd.concat(
            [cob.locality(
                cob.refgen.bootstrap_candidate_genes(
                    loci,
                    flank_limit=args.candidate_flank_limit
                ),
                bootstrap_name=x,
                include_regression=True
            ) for x in range(args.num_bootstraps)]
    ).sort('global')
    OLS = sm.OLS(fdr['local'],fdr['global']).fit()

    # Remove global degree larger than empirical
    fdr = fdr[fdr['global'] <= max(loc['global'])]
    fdr['window'] = [int(x/window_ticks) for x in fdr['global']]

    # calculate z-scores for the global
    fdr['bs_std'] = [fit_std[x] for x in fdr['fitted']]
    fdr['zscore'] = [x['resid']/x['bs_std'] for i,x in fdr.iterrows()]

    # Give em the gold
    return loc,bsloc,fdr

def plot_scatter(args,loc,bsloc,fdr,ax):

    ''' ---------------------------------------------------
        Plotting
    '''
    # Y axis is local degree (what we are TRYING to predict)
    ax.set_ylim(0,max(loc['local']))
    ax.set_xlim(0,max(loc['global']))
    ax.set_xlabel('Number Global Interactions')
    ax.set_ylabel('Number Local Interactions')

    # UGH! map lowess
    fdrlowess = lowess(
        fdr['local'],fdr['global'],
        frac=0.15,it=5,delta=0.1*len(fdr),
        is_sorted=False
    )
    # plot the bootstrap points
    ax.plot(fdr['global'],fdr['local'],'ro',alpha=0.05,label='Bootstraps')
    # plot the OLS lowess line
    ci = fdr.groupby('window')['fitted','global'].agg(
        [np.mean,confidence_interval]
    )

    #for win,df in fdr.groupby('bootstrap_name'):
    #    ax.plot(df['global'],df['fitted'],alpha=1)

    ax.errorbar(
        ci['global','mean'],ci['fitted','mean'],
        yerr=ci['fitted','confidence_interval'],
        color='g',label='Bootstrap OLS'
    )

    #plot the empirical data
    ax.plot(loc['global'],loc['local'],'bo',alpha=1,label='Empirical')
    ax.plot(loc['global'],loc['fitted'],'k-',alpha=1,label='Empirical OLS')
    # finish plots
    #legend = ax.legend(loc='best')

    '''---------------------------------------------------
        FDR Plotting
    '''

def plot_fdr(args,loc,bsloc,fdr,ax):
    # Plot the empirical Z-score distributions
    zscores = list(range(1,8))
    zloc = [
        sum(np.logical_and(
            loc['zscore'] >= x ,
            loc['local'] >= args.min_fdr_degree
        ))
        for x in zscores
    ]
    ax.plot(zscores,zloc,'bo',label='Empirical Zscore > x')
    # plot the fdr scores spread
    zcdf = pd.DataFrame(
        [ mean_confidence_interval(
            fdr.groupby('bootstrap_name').apply(
                lambda df: sum(np.logical_and(
                    df['zscore'] >= x,
                    df['local'] >= args.min_fdr_degree
                ))
            )
        ) for x in zscores ],
        columns=['mean','ci']
    )
    ax.errorbar(
        zscores,
        zcdf['mean'],
        yerr=zcdf['ci'],
        label='Bootstrap Z-scores',
        color='red'
    )
    ax.set_xlabel('Z-Score')
    ax.set_ylabel('Number of Genes > Z')
    ax.set_title('Z Score FDR')

def plot_data(args,loc,bsloc,fdr,ax):
    ax.xaxis.set_visible(True)
    ax.yaxis.set_visible(True)
    ax.text(0,0,'''
        COB: {}
        Ontology: {}
        Term: {}
    '''.format(
        args.cob,
        args.gwas,
        args.terms
    ))
<|endoftext|>


<commit_before>import string

import requests
from bs4 import BeautifulSoup
import tinycss
from lxml import etree

# This is how many pixels occupy each square,
# and tell us which position each crossword starts in
square_size = 29

def get_crossword(id, type='cryptic'):
    url = "http://www.theguardian.com/crosswords/" + type + "/" + id

    # Get the page
    r = requests.get(url)
    soup = BeautifulSoup(r.text)

    # Let's grab the Title, Author and Date
    title = soup.find("div", {'id': 'main-article-info'}).h1.text

    article_attributes = soup.find("ul", {'class': 'article-attributes'})

    author = article_attributes.find("li", {'class': 'byline'}).a.text

    # Some more excellent string splitting here ..
    # Reliant on the time always being 00.00
    date = article_attributes.find("li", {'class': 'publication'}).text
    date = string.strip((date.split(','))[1].split('00')[0])

    # We get the clues- and the layout of the clues
    raw_clues_across = soup.find("div", {'class': 'clues-col'})
    raw_clues_down = soup.find("div", {'class': 'clues-col last'})
    raw_layout = soup.find("div", {'id': 'grid'})

    # We're grabbing the 'crossword' div and will grab the solutions
    raw_answers = soup.find("div", {'id': 'wrapper'})
    raw_answers = raw_answers.find("div", {'class': 'crossword'})

    # These dicts are where all clues, locations amd solutions live
    across = {}
    down = {}

    # Populate the clues
    for li in raw_clues_across.findAll('li'):
        # We strip off the start of the string since it has rubbish
        clue = li.text[5:]

        # Empty dict for each clue ID
        across[li.span.text] = {}

        # Populate the clue with the stripped (cleaned) string
        across[li.span.text]['clue'] = string.strip(clue)

    # Same for down
    for li in raw_clues_down.findAll('li'):
        clue = li.text[5:]
        clue_id = li.span.text.split(",")[0]
        down[clue_id] = {}
        down[clue_id]['clue'] = string.strip(clue)

    # We look for each word which is part of the 'across' class
    for div in raw_layout.findAll('div', {'class': 'across'}):

        # We parse the inline styles associated with it
        parser = tinycss.make_parser()
        stylesheet = parser.parse_style_attr(div['style'])

        # Then we go through and pick each useful position out of the styles
        for dec in stylesheet[0]:

            if dec.name == "left":
                left = dec.value[0].value / square_size
            elif dec.name == "top":
                top = dec.value[0].value / square_size
            else:
                pass

        # With those, we can get the position the start of the clue is in
        position = left + top*15

        # Then save to the correct clue ID
        across[div['id'].split('-')[0]]['launch_spot'] = position

    # Same for down
    for div in raw_layout.findAll('div', {'class': 'down'}):

        parser = tinycss.make_parser()
        stylesheet = parser.parse_style_attr(div['style'])

        for dec in stylesheet[0]:

            if dec.name == "left":
                left = dec.value[0].value / square_size
            elif dec.name == "top":
                top = dec.value[0].value / square_size
            else:
                pass

        position = left + top*15
        # print down
        down[div['id'].split('-')[0]]['launch_spot'] = position

    # Now we collect answers
    for line in str(raw_answers.script).split("\n"):

        # Only interested in lines with the 'solutions' array
        if "solutions[" in line:
            separated_answer = line.split('"')

            # We can get certain values by splitting at strategic locations
            key = separated_answer[1].split('-')
            res = separated_answer[3]

            # This is reliant on the order of the answers going down the page
            # If not, we have problems: More code would be needed

            # If it's the first letter then the value won't exist (key error)
            # If that happens, we start it off with the first letter.

            try:
                if key[1] == 'across':
                    across[key[0]]['answer'] += res
                elif key[1] == 'down':
                    down[key[0]]['answer'] += res

            except KeyError:
                if key[1] == 'across':
                    across[key[0]]['answer'] = res
                elif key[1] == 'down':
                    down[key[0]]['answer'] = res

    # Then we start creating the XML
    root = etree.Element('crossword')

    child = etree.Element('Title')
    child.attrib['v'] = title + ', ' + date
    root.append(child)

    child = etree.Element('Author')
    child.attrib['v'] = author
    root.append(child)

    child = etree.Element('Category')
    child.attrib['v'] = "The Guardian Crosswords"
    root.append(child)

    child = etree.Element('Copyright')
    child.attrib['v'] = "The Guardian"
    root.append(child)

    child = etree.Element('Editor')
    child.attrib['v'] = "The Editor"
    root.append(child)

    child = etree.Element('Width')
    child.attrib['v'] = "15"
    root.append(child)

    child = etree.Element('Height')
    child.attrib['v'] = "15"
    root.append(child)

    # Here we actually add each clue and solution in to the XML
    across_child = etree.Element('across')

    i = 1
    for key, value in across.iteritems():
        child = etree.Element('a' + str(i))
        child.attrib['a'] = value['answer']
        child.attrib['c'] = value['clue']
        child.attrib['n'] = str(value['launch_spot'])
        child.attrib['cn'] = key
        across_child.append(child)

        i += 1

    root.append(across_child)

    down_child = etree.Element('down')

    i = 1
    for key, value in down.iteritems():
        child = etree.Element('d' + str(i))
        child.attrib['a'] = value['answer']
        child.attrib['c'] = value['clue']
        child.attrib['n'] = str(value['launch_spot'])
        child.attrib['cn'] = key
        down_child.append(child)

        i += 1

    root.append(down_child)

    s = etree.tostring(root)

    return s

if __name__ == "__main__":
    print get_crossword('26285')
<commit_msg>Return 'None' when The Guardian gives a 404<commit_after>import string

import requests
from bs4 import BeautifulSoup
import tinycss
from lxml import etree

# This is how many pixels occupy each square,
# and tell us which position each crossword starts in
square_size = 29

def get_crossword(id, type='cryptic'):
    url = "http://www.theguardian.com/crosswords/" + type + "/" + id

    # Get the page
    r = requests.get(url)
    if r.status_code == 404:
        return None

    soup = BeautifulSoup(r.text)

    # Let's grab the Title, Author and Date
    title = soup.find("div", {'id': 'main-article-info'}).h1.text

    article_attributes = soup.find("ul", {'class': 'article-attributes'})

    author = article_attributes.find("li", {'class': 'byline'}).a.text

    # Some more excellent string splitting here ..
    # Reliant on the time always being 00.00
    date = article_attributes.find("li", {'class': 'publication'}).text
    date = string.strip((date.split(','))[1].split('00')[0])

    # We get the clues- and the layout of the clues
    raw_clues_across = soup.find("div", {'class': 'clues-col'})
    raw_clues_down = soup.find("div", {'class': 'clues-col last'})
    raw_layout = soup.find("div", {'id': 'grid'})

    # We're grabbing the 'crossword' div and will grab the solutions
    raw_answers = soup.find("div", {'id': 'wrapper'})
    raw_answers = raw_answers.find("div", {'class': 'crossword'})

    # These dicts are where all clues, locations amd solutions live
    across = {}
    down = {}

    # Populate the clues
    for li in raw_clues_across.findAll('li'):
        # We strip off the start of the string since it has rubbish
        clue = li.text[5:]

        # Empty dict for each clue ID
        across[li.span.text] = {}

        # Populate the clue with the stripped (cleaned) string
        across[li.span.text]['clue'] = string.strip(clue)

    # Same for down
    for li in raw_clues_down.findAll('li'):
        clue = li.text[5:]
        clue_id = li.span.text.split(",")[0]
        down[clue_id] = {}
        down[clue_id]['clue'] = string.strip(clue)

    # We look for each word which is part of the 'across' class
    for div in raw_layout.findAll('div', {'class': 'across'}):

        # We parse the inline styles associated with it
        parser = tinycss.make_parser()
        stylesheet = parser.parse_style_attr(div['style'])

        # Then we go through and pick each useful position out of the styles
        for dec in stylesheet[0]:

            if dec.name == "left":
                left = dec.value[0].value / square_size
            elif dec.name == "top":
                top = dec.value[0].value / square_size
            else:
                pass

        # With those, we can get the position the start of the clue is in
        position = left + top*15

        # Then save to the correct clue ID
        across[div['id'].split('-')[0]]['launch_spot'] = position

    # Same for down
    for div in raw_layout.findAll('div', {'class': 'down'}):

        parser = tinycss.make_parser()
        stylesheet = parser.parse_style_attr(div['style'])

        for dec in stylesheet[0]:

            if dec.name == "left":
                left = dec.value[0].value / square_size
            elif dec.name == "top":
                top = dec.value[0].value / square_size
            else:
                pass

        position = left + top*15
        # print down
        down[div['id'].split('-')[0]]['launch_spot'] = position

    # Now we collect answers
    for line in str(raw_answers.script).split("\n"):

        # Only interested in lines with the 'solutions' array
        if "solutions[" in line:
            separated_answer = line.split('"')

            # We can get certain values by splitting at strategic locations
            key = separated_answer[1].split('-')
            res = separated_answer[3]

            # This is reliant on the order of the answers going down the page
            # If not, we have problems: More code would be needed

            # If it's the first letter then the value won't exist (key error)
            # If that happens, we start it off with the first letter.

            try:
                if key[1] == 'across':
                    across[key[0]]['answer'] += res
                elif key[1] == 'down':
                    down[key[0]]['answer'] += res

            except KeyError:
                if key[1] == 'across':
                    across[key[0]]['answer'] = res
                elif key[1] == 'down':
                    down[key[0]]['answer'] = res

    # Then we start creating the XML
    root = etree.Element('crossword')

    child = etree.Element('Title')
    child.attrib['v'] = title + ', ' + date
    root.append(child)

    child = etree.Element('Author')
    child.attrib['v'] = author
    root.append(child)

    child = etree.Element('Category')
    child.attrib['v'] = "The Guardian Crosswords"
    root.append(child)

    child = etree.Element('Copyright')
    child.attrib['v'] = "The Guardian"
    root.append(child)

    child = etree.Element('Editor')
    child.attrib['v'] = "The Editor"
    root.append(child)

    child = etree.Element('Width')
    child.attrib['v'] = "15"
    root.append(child)

    child = etree.Element('Height')
    child.attrib['v'] = "15"
    root.append(child)

    # Here we actually add each clue and solution in to the XML
    across_child = etree.Element('across')

    i = 1
    for key, value in across.iteritems():
        child = etree.Element('a' + str(i))
        child.attrib['a'] = value['answer']
        child.attrib['c'] = value['clue']
        child.attrib['n'] = str(value['launch_spot'])
        child.attrib['cn'] = key
        across_child.append(child)

        i += 1

    root.append(across_child)

    down_child = etree.Element('down')

    i = 1
    for key, value in down.iteritems():
        child = etree.Element('d' + str(i))
        child.attrib['a'] = value['answer']
        child.attrib['c'] = value['clue']
        child.attrib['n'] = str(value['launch_spot'])
        child.attrib['cn'] = key
        down_child.append(child)

        i += 1

    root.append(down_child)

    s = etree.tostring(root)

    return s

if __name__ == "__main__":
    print get_crossword('26285')
<|endoftext|>


<commit_before># Copyright (c) 2010-2013, GEM Foundation.
#
# OpenQuake is free software: you can redistribute it and/or modify it
# under the terms of the GNU Affero General Public License as published
# by the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# OpenQuake is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with OpenQuake.  If not, see <http://www.gnu.org/licenses/>.

"""
Core functionality for the classical PSHA hazard calculator.
"""
import collections

import numpy

from openquake.hazardlib.imt import from_string

from openquake.engine import logs, writer
from openquake.engine.calculators.hazard import general
from openquake.engine.calculators.hazard.classical import (
    post_processing as post_proc)
from openquake.engine.db import models
from openquake.engine.utils import tasks
from openquake.engine.performance import EnginePerformanceMonitor

MAX_RUPTURES = 1000  # if there are more ruptures, spawn a separate task

@tasks.oqtask
def compute_hazard_curves(job_id, sources, tom, gsims_by_rlz):
    """
    This task computes R2 * I hazard curves (each one is a
    numpy array of S * L floats) from the given source_ruptures
    pairs.

    :param job_id:
    :param sources:
        a list of pairs (source, ruptures) where ruptures is a
        a (sub)set of ruptures generated by the source
    :param gsim_dicts:
        a list of gsim dictionaries, one for each GMPE realization
    """
    results = []
    hc = models.HazardCalculation.objects.get(oqjob=job_id)
    total_sites = len(hc.site_collection)
    imts = general.im_dict_to_hazardlib(
        hc.intensity_measure_types_and_levels)
    curves = dict((rlz, dict((imt, numpy.ones([total_sites, len(imts[imt])]))
                             for imt in imts))
                  for rlz in gsims_by_rlz)
    num_sources = len(sources)
    for source in sources:
        s_sites = source.filter_sites_by_distance_to_source(
            hc.maximum_distance, hc.site_collection
        ) if hc.maximum_distance else hc.site_collection
        if s_sites is None:
            continue
        ruptures = list(source.iter_ruptures(tom))
        if num_sources > 1 and len(ruptures) > MAX_RUPTURES:  # spawn a task
            results.append(
                compute_hazard_curves.delay(
                    job_id, [source], tom, gsims_by_rlz))
            continue
        for rupture in ruptures:
            r_sites = rupture.source_typology.\
                filter_sites_by_distance_to_rupture(
                    rupture, hc.maximum_distance, s_sites
                ) if hc.maximum_distance else s_sites
            if r_sites is None:
                continue
            prob = rupture.get_probability_one_or_more_occurrences()
            for rlz, curv in curves.iteritems():
                gsim = gsims_by_rlz[rlz][rupture.tectonic_region_type]
                sctx, rctx, dctx = gsim.make_contexts(r_sites, rupture)
                for imt in imts:
                    poes = gsim.get_poes(sctx, rctx, dctx, imt, imts[imt],
                                         hc.truncation_level)
                    curv[imt] *= r_sites.expand(
                        (1. - prob) ** poes, total_sites, placeholder=1)
    # the 0 here is a shortcut for filtered sources giving no contribution;
    # this is essential for performance, we want to avoid returning
    # big arrays of zeros (MS)
    dic = dict((rlz, [0 if (curv[imt] == 1.0).all() else 1. - curv[imt]
                      for imt in sorted(imts)])
               for rlz, curv in curves.iteritems())
    results.append(dic)
    return results

class ClassicalHazardCalculator(general.BaseHazardCalculator):
    """
    Classical PSHA hazard calculator. Computes hazard curves for a given set of
    points.

    For each realization of the calculation, we randomly sample source models
    and GMPEs (Ground Motion Prediction Equations) from logic trees.
    """

    core_calc_task = compute_hazard_curves

    def pre_execute(self):
        """
        Do pre-execution work. At the moment, this work entails:
        parsing and initializing sources, parsing and initializing the
        site model (if there is one), parsing vulnerability and
        exposure files and generating logic tree realizations. (The
        latter piece basically defines the work to be done in the
        `execute` phase.).
        """
        super(ClassicalHazardCalculator, self).pre_execute()
        self.imtls = self.hc.intensity_measure_types_and_levels
        realizations = self._get_realizations()
        n_rlz = len(realizations)
        n_levels = sum(len(lvls) for lvls in self.imtls.itervalues()
                       ) / float(len(self.imtls))
        n_sites = len(self.hc.site_collection)
        total = n_rlz * len(self.imtls) * n_levels * n_sites
        logs.LOG.info('Considering %d realization(s), %d IMT(s), %d level(s) '
                      'and %d sites, total %d', n_rlz, len(self.imtls),
                      n_levels, n_sites, total)
        self.curves_by_rlz = collections.OrderedDict(
            (rlz, [numpy.zeros((n_sites, len(self.imtls[imt])))
                   for imt in self.imtls])
            for rlz in realizations)

    @EnginePerformanceMonitor.monitor
    def task_completed(self, task_results):
        """
        This is used to incrementally update hazard curve results by combining
        an initial value with some new results. (Each set of new results is
        computed over only a subset of seismic sources defined in the
        calculation model.)

        :param task_result:
            XXX
            A pair (curves_by_imt, ordinal) where curves_by_imt is a
            list of 2-D numpy arrays representing the new results which need
            to be combined with the current value. These should be the same
            shape as self.curves_by_rlz[i][j] where i is the realization
            ordinal and j the IMT ordinal.
        """
        for res in task_results:
            for dic in (res.get() if hasattr(res, 'result') else [res]):
                for rlz, curves_by_imt in dic.iteritems():
                    for j, curves in enumerate(curves_by_imt):
                        # j is the IMT index
                        self.curves_by_rlz[rlz][j] = 1. - (
                            1. - self.curves_by_rlz[rlz][j]) * (1. - curves)
            self.log_percent()

    # this could be parallelized in the future, however in all the cases
    # I have seen until now, the serialized approach is fast enough (MS)
    @EnginePerformanceMonitor.monitor
    def save_hazard_curves(self):
        """
        Post-execution actions. At the moment, all we do is finalize the hazard
        curve results.
        """
        imtls = self.hc.intensity_measure_types_and_levels
        for rlz, curves_by_imt in self.curves_by_rlz.iteritems():
            # create a new `HazardCurve` 'container' record for each
            # realization (virtual container for multiple imts)
            models.HazardCurve.objects.create(
                output=models.Output.objects.create_output(
                    self.job, "hc-multi-imt-rlz-%s" % rlz.id,
                    "hazard_curve_multi"),
                lt_realization=rlz,
                imt=None,
                investigation_time=self.hc.investigation_time)

            # create a new `HazardCurve` 'container' record for each
            # realization for each intensity measure type
            for imt, curves in zip(sorted(imtls), curves_by_imt):
                hc_im_type, sa_period, sa_damping = from_string(imt)

                # save output
                hco = models.Output.objects.create(
                    oq_job=self.job,
                    display_name="Hazard Curve rlz-%s" % rlz.id,
                    output_type='hazard_curve',
                )

                # save hazard_curve
                haz_curve = models.HazardCurve.objects.create(
                    output=hco,
                    lt_realization=rlz,
                    investigation_time=self.hc.investigation_time,
                    imt=hc_im_type,
                    imls=imtls[imt],
                    sa_period=sa_period,
                    sa_damping=sa_damping,
                )

                # save hazard_curve_data
                points = self.hc.points_to_compute()
                logs.LOG.info('saving %d hazard curves for %s, imt=%s',
                              len(points), hco, imt)
                writer.CacheInserter.saveall([
                    models.HazardCurveData(
                        hazard_curve=haz_curve,
                        poes=list(poes),
                        location='POINT(%s %s)' % (p.longitude, p.latitude),
                        weight=rlz.weight)
                    for p, poes in zip(points, curves)])
        del self.curves_by_rlz  # save memory for the post_processing phase

    post_execute = save_hazard_curves

    def post_process(self):
        """
        Optionally generates aggregate curves, hazard maps and
        uniform_hazard_spectra.
        """
        logs.LOG.debug('> starting post processing')

        # means/quantiles:
        if self.hc.mean_hazard_curves or self.hc.quantile_hazard_curves:
            self.do_aggregate_post_proc()

        # hazard maps:
        # required for computing UHS
        # if `hazard_maps` is false but `uniform_hazard_spectra` is true,
        # just don't export the maps
        if self.hc.hazard_maps or self.hc.uniform_hazard_spectra:
            self.parallelize(
                post_proc.hazard_curves_to_hazard_map_task,
                post_proc.hazard_curves_to_hazard_map_task_arg_gen(self.job),
                self.log_percent)

        if self.hc.uniform_hazard_spectra:
            post_proc.do_uhs_post_proc(self.job)

        logs.LOG.debug('< done with post processing')
<commit_msg>Fixed an ordering bug<commit_after># Copyright (c) 2010-2013, GEM Foundation.
#
# OpenQuake is free software: you can redistribute it and/or modify it
# under the terms of the GNU Affero General Public License as published
# by the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# OpenQuake is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with OpenQuake.  If not, see <http://www.gnu.org/licenses/>.

"""
Core functionality for the classical PSHA hazard calculator.
"""
import collections

import numpy

from openquake.hazardlib.imt import from_string

from openquake.engine import logs, writer
from openquake.engine.calculators.hazard import general
from openquake.engine.calculators.hazard.classical import (
    post_processing as post_proc)
from openquake.engine.db import models
from openquake.engine.utils import tasks
from openquake.engine.performance import EnginePerformanceMonitor

MAX_RUPTURES = 1000  # if there are more ruptures, spawn a separate task

@tasks.oqtask
def compute_hazard_curves(job_id, sources, tom, gsims_by_rlz):
    """
    This task computes R2 * I hazard curves (each one is a
    numpy array of S * L floats) from the given source_ruptures
    pairs.

    :param job_id:
    :param sources:
        a list of pairs (source, ruptures) where ruptures is a
        a (sub)set of ruptures generated by the source
    :param gsim_dicts:
        a list of gsim dictionaries, one for each GMPE realization
    """
    results = []
    hc = models.HazardCalculation.objects.get(oqjob=job_id)
    total_sites = len(hc.site_collection)
    imts = general.im_dict_to_hazardlib(
        hc.intensity_measure_types_and_levels)
    curves = dict((rlz, dict((imt, numpy.ones([total_sites, len(imts[imt])]))
                             for imt in imts))
                  for rlz in gsims_by_rlz)
    num_sources = len(sources)
    for source in sources:
        s_sites = source.filter_sites_by_distance_to_source(
            hc.maximum_distance, hc.site_collection
        ) if hc.maximum_distance else hc.site_collection
        if s_sites is None:
            continue
        ruptures = list(source.iter_ruptures(tom))
        if num_sources > 1 and len(ruptures) > MAX_RUPTURES:  # spawn a task
            results.append(
                compute_hazard_curves.delay(
                    job_id, [source], tom, gsims_by_rlz))
            continue
        for rupture in ruptures:
            r_sites = rupture.source_typology.\
                filter_sites_by_distance_to_rupture(
                    rupture, hc.maximum_distance, s_sites
                ) if hc.maximum_distance else s_sites
            if r_sites is None:
                continue
            prob = rupture.get_probability_one_or_more_occurrences()
            for rlz, curv in curves.iteritems():
                gsim = gsims_by_rlz[rlz][rupture.tectonic_region_type]
                sctx, rctx, dctx = gsim.make_contexts(r_sites, rupture)
                for imt in imts:
                    poes = gsim.get_poes(sctx, rctx, dctx, imt, imts[imt],
                                         hc.truncation_level)
                    curv[imt] *= r_sites.expand(
                        (1. - prob) ** poes, total_sites, placeholder=1)
    # the 0 here is a shortcut for filtered sources giving no contribution;
    # this is essential for performance, we want to avoid returning
    # big arrays of zeros (MS)
    dic = dict((rlz, [0 if (curv[imt] == 1.0).all() else 1. - curv[imt]
                      for imt in sorted(imts)])
               for rlz, curv in curves.iteritems())
    results.append(dic)
    return results

class ClassicalHazardCalculator(general.BaseHazardCalculator):
    """
    Classical PSHA hazard calculator. Computes hazard curves for a given set of
    points.

    For each realization of the calculation, we randomly sample source models
    and GMPEs (Ground Motion Prediction Equations) from logic trees.
    """

    core_calc_task = compute_hazard_curves

    def pre_execute(self):
        """
        Do pre-execution work. At the moment, this work entails:
        parsing and initializing sources, parsing and initializing the
        site model (if there is one), parsing vulnerability and
        exposure files and generating logic tree realizations. (The
        latter piece basically defines the work to be done in the
        `execute` phase.).
        """
        super(ClassicalHazardCalculator, self).pre_execute()
        self.imtls = self.hc.intensity_measure_types_and_levels
        realizations = self._get_realizations()
        n_rlz = len(realizations)
        n_levels = sum(len(lvls) for lvls in self.imtls.itervalues()
                       ) / float(len(self.imtls))
        n_sites = len(self.hc.site_collection)
        total = n_rlz * len(self.imtls) * n_levels * n_sites
        logs.LOG.info('Considering %d realization(s), %d IMT(s), %d level(s) '
                      'and %d sites, total %d', n_rlz, len(self.imtls),
                      n_levels, n_sites, total)
        self.curves_by_rlz = collections.OrderedDict(
            (rlz, [numpy.zeros((n_sites, len(self.imtls[imt])))
                   for imt in sorted(self.imtls)])
            for rlz in realizations)

    @EnginePerformanceMonitor.monitor
    def task_completed(self, task_results):
        """
        This is used to incrementally update hazard curve results by combining
        an initial value with some new results. (Each set of new results is
        computed over only a subset of seismic sources defined in the
        calculation model.)

        :param task_result:
            XXX
            A pair (curves_by_imt, ordinal) where curves_by_imt is a
            list of 2-D numpy arrays representing the new results which need
            to be combined with the current value. These should be the same
            shape as self.curves_by_rlz[i][j] where i is the realization
            ordinal and j the IMT ordinal.
        """
        for res in task_results:
            for dic in (res.get() if hasattr(res, 'result') else [res]):
                for rlz, curves_by_imt in dic.iteritems():
                    for j, curves in enumerate(curves_by_imt):
                        # j is the IMT index
                        self.curves_by_rlz[rlz][j] = 1. - (
                            1. - self.curves_by_rlz[rlz][j]) * (1. - curves)
            self.log_percent()

    # this could be parallelized in the future, however in all the cases
    # I have seen until now, the serialized approach is fast enough (MS)
    @EnginePerformanceMonitor.monitor
    def save_hazard_curves(self):
        """
        Post-execution actions. At the moment, all we do is finalize the hazard
        curve results.
        """
        imtls = self.hc.intensity_measure_types_and_levels
        for rlz, curves_by_imt in self.curves_by_rlz.iteritems():
            # create a new `HazardCurve` 'container' record for each
            # realization (virtual container for multiple imts)
            models.HazardCurve.objects.create(
                output=models.Output.objects.create_output(
                    self.job, "hc-multi-imt-rlz-%s" % rlz.id,
                    "hazard_curve_multi"),
                lt_realization=rlz,
                imt=None,
                investigation_time=self.hc.investigation_time)

            # create a new `HazardCurve` 'container' record for each
            # realization for each intensity measure type
            for imt, curves in zip(sorted(imtls), curves_by_imt):
                hc_im_type, sa_period, sa_damping = from_string(imt)

                # save output
                hco = models.Output.objects.create(
                    oq_job=self.job,
                    display_name="Hazard Curve rlz-%s" % rlz.id,
                    output_type='hazard_curve',
                )

                # save hazard_curve
                haz_curve = models.HazardCurve.objects.create(
                    output=hco,
                    lt_realization=rlz,
                    investigation_time=self.hc.investigation_time,
                    imt=hc_im_type,
                    imls=imtls[imt],
                    sa_period=sa_period,
                    sa_damping=sa_damping,
                )

                # save hazard_curve_data
                points = self.hc.points_to_compute()
                logs.LOG.info('saving %d hazard curves for %s, imt=%s',
                              len(points), hco, imt)
                writer.CacheInserter.saveall([
                    models.HazardCurveData(
                        hazard_curve=haz_curve,
                        poes=list(poes),
                        location='POINT(%s %s)' % (p.longitude, p.latitude),
                        weight=rlz.weight)
                    for p, poes in zip(points, curves)])
        del self.curves_by_rlz  # save memory for the post_processing phase

    post_execute = save_hazard_curves

    def post_process(self):
        """
        Optionally generates aggregate curves, hazard maps and
        uniform_hazard_spectra.
        """
        logs.LOG.debug('> starting post processing')

        # means/quantiles:
        if self.hc.mean_hazard_curves or self.hc.quantile_hazard_curves:
            self.do_aggregate_post_proc()

        # hazard maps:
        # required for computing UHS
        # if `hazard_maps` is false but `uniform_hazard_spectra` is true,
        # just don't export the maps
        if self.hc.hazard_maps or self.hc.uniform_hazard_spectra:
            self.parallelize(
                post_proc.hazard_curves_to_hazard_map_task,
                post_proc.hazard_curves_to_hazard_map_task_arg_gen(self.job),
                self.log_percent)

        if self.hc.uniform_hazard_spectra:
            post_proc.do_uhs_post_proc(self.job)

        logs.LOG.debug('< done with post processing')
<|endoftext|>


<commit_before># Copyright 2014 Open Source Robotics Foundation, Inc.
#
# 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.

from __future__ import print_function

import argparse
from datetime import date
import os
import pkg_resources
import sys

from catkin_tools.common import is_tty

from catkin_tools.config import get_verb_aliases
from catkin_tools.config import initialize_config

from catkin_tools.terminal_color import fmt
from catkin_tools.terminal_color import set_color
from catkin_tools.terminal_color import test_colors

CATKIN_COMMAND_VERB_GROUP = 'catkin_tools.commands.catkin.verbs'

def list_verbs():
    verbs = []
    for entry_point in pkg_resources.iter_entry_points(group=CATKIN_COMMAND_VERB_GROUP):
        verbs.append(entry_point.name)
    return verbs

def load_verb_description(verb_name):
    for entry_point in pkg_resources.iter_entry_points(group=CATKIN_COMMAND_VERB_GROUP):
        if entry_point.name == verb_name:
            return entry_point.load()

def default_argument_preprocessor(args):
    extras = {}
    return args, extras

def create_subparsers(parser, verbs):
    verbs = sorted(verbs)
    verb_array_str = '[' + ' | '.join(verbs) + ']'
    verb_list_str = 'Call `catkin VERB -h` for help on each verb listed below:\n'
    for verb in verbs:
        desc = load_verb_description(verb)
        verb_list_str += '\n  %s\t%s' % (desc['verb'], desc['description'])

    subparser = parser.add_subparsers(
        title='catkin command',
        metavar=verb_array_str,
        description=verb_list_str,
        dest='verb'
    )

    argument_preprocessors = {}

    for verb in verbs:
        desc = load_verb_description(verb)
        cmd_parser = subparser.add_parser(desc['verb'], description=desc['description'])
        cmd_parser = desc['prepare_arguments'](cmd_parser)

        cmd_parser.set_defaults(main=desc['main'])

        if 'argument_preprocessor' in desc:
            argument_preprocessors[verb] = desc['argument_preprocessor']
        else:
            argument_preprocessors[verb] = default_argument_preprocessor

    return argument_preprocessors

def expand_one_verb_alias(sysargs, verb_aliases, used_aliases):
    """Iterate through sysargs looking for expandable verb aliases.

    When a verb alias is found, sysargs is modified to effectively expand the alias.
    The alias is removed from verb_aliases and added to used_aliases.
    After finding and expanding an alias, this function returns True.
    If no alias is found to be expanded, this function returns False.
    """
    cmd = os.path.basename(sys.argv[0])
    for index, arg in enumerate(sysargs):
        if arg.startswith('-'):
            # Not a verb, continue through the arguments
            continue
        if arg in used_aliases:
            print(fmt(
                "@!@{gf}==>@| Expanding alias '@!@{yf}" + arg +
                "@|' was previously expanded, ignoring this time to prevent infinite recursion."
            ))
        if arg in verb_aliases:
            before = [] if index == 0 else sysargs[:index - 1]
            after = [] if index == len(sysargs) else sysargs[index + 1:]
            sysargs[:] = before + verb_aliases[arg].split() + after
            print(fmt(
                "@!@{gf}==>@| Expanding alias "
                "'@!@{yf}{alias}@|' "
                "from '@{yf}{before} @!{alias}@{boldoff}{after}@|' "
                "to '@{yf}{before} @!{expansion}@{boldoff}{after}@|'"
            ).format(
                alias=arg,
                expansion=verb_aliases[arg],
                before=' '.join([cmd] + before),
                after=(' '.join([''] + after) if after else '')
            ))
            # Prevent the alias from being used again, to prevent infinite recursion
            used_aliases.append(arg)
            del verb_aliases[arg]
            # Return True since one has been found
            return True
        # Return False since no verb alias was found
        return False

def expand_verb_aliases(sysargs, verb_aliases):
    """Expands aliases in sysargs which are found in verb_aliases until none are found."""
    used_aliases = []
    while expand_one_verb_alias(sysargs, verb_aliases, used_aliases):
        pass
    return sysargs

def catkin_main(sysargs):
    # Initialize config
    try:
        initialize_config()
    except RuntimeError as exc:
        sys.exit("Failed to initialize config: {0}".format(exc))

    # Create a top level parser
    parser = argparse.ArgumentParser(
        description="catkin command", formatter_class=argparse.RawDescriptionHelpFormatter)
    add = parser.add_argument
    add('-a', '--list-aliases', action="store_true", default=False,
        help="Lists the current verb aliases and then quits, all other arguments are ignored")
    add('--test-colors', action='store_true', default=False,
        help="Prints a color test pattern to the screen and then quits, all other arguments are ignored")
    add('--version', action='store_true', default=False,
        help="Prints the catkin_tools version.")
    color_control_group = parser.add_mutually_exclusive_group()
    add = color_control_group.add_argument
    add('--force-color', action='store_true', default=False,
        help='Forces catkin to output in color, even when the terminal does not appear to support it.')
    add('--no-color', action='store_true', default=False,
        help='Forces catkin to not use color in the output, regardless of the detect terminal type.')

    # Generate a list of verbs available
    verbs = list_verbs()

    # Create the subparsers for each verb and collect the argument preprocessors
    argument_preprocessors = create_subparsers(parser, verbs)

    # Get verb aliases
    verb_aliases = get_verb_aliases()

    # Setup sysargs
    sysargs = sys.argv[1:] if sysargs is None else sysargs

    # Get colors config
    no_color = False
    force_color = os.environ.get('CATKIN_TOOLS_FORCE_COLOR', False)
    for arg in sysargs:
        if arg == '--no-color':
            no_color = True
        if arg == '--force-color':
            force_color = True

    if no_color or not force_color and not is_tty(sys.stdout):
        set_color(False)

    # Check for version
    if '--version' in sysargs:
        print('catkin_tools {} (C) 2014-{} Open Source Robotics Foundation'.format(
            pkg_resources.get_distribution('catkin_tools').version,
            date.today().year)
        )
        print('catkin_tools is released under the Apache License,'
              ' Version 2.0 (http://www.apache.org/licenses/LICENSE-2.0)')
        print('---')
        print('Using Python {}'.format(''.join(sys.version.split('\n'))))
        sys.exit(0)

    # Check for --test-colors
    for arg in sysargs:
        if arg == '--test-colors':
            test_colors()
            sys.exit(0)
        if not arg.startswith('-'):
            break

    # Check for --list-aliases
    for arg in sysargs:
        if arg == '--list-aliases' or arg == '-a':
            for alias in sorted(list(verb_aliases.keys())):
                print("{0}: {1}".format(alias, verb_aliases[alias]))
            sys.exit(0)
        if not arg.startswith('-'):
            break

    # Do verb alias expansion
    sysargs = expand_verb_aliases(sysargs, verb_aliases)

    # Determine the verb, splitting arguments into pre and post verb
    verb = None
    pre_verb_args = []
    post_verb_args = []
    for index, arg in enumerate(sysargs):
        # If the arg does not start with a `-` then it is a positional argument
        # The first positional argument must be the verb
        if not arg.startswith('-'):
            verb = arg
            post_verb_args = sysargs[index + 1:]
            break
        # If the `-h` or `--help` option comes before the verb, parse_args
        if arg in ['-h', '--help']:
            parser.parse_args(sysargs)
        # Otherwise it is a pre-verb option
        pre_verb_args.append(arg)

    # Error on no verb provided
    if verb is None:
        print(parser.format_usage())
        sys.exit("Error: No verb provided.")
    # Error on unknown verb provided
    if verb not in verbs:
        print(parser.format_usage())
        sys.exit("Error: Unknown verb '{0}' provided.".format(verb))

    # First allow the verb's argument preprocessor to strip any args
    # and return any "extra" information it wants as a dict
    processed_post_verb_args, extras = argument_preprocessors[verb](post_verb_args)
    # Then allow argparse to process the left over post-verb arguments along
    # with the pre-verb arguments and the verb itself
    args = parser.parse_args(pre_verb_args + [verb] + processed_post_verb_args)
    # Extend the argparse result with the extras from the preprocessor
    for key, value in extras.items():
        setattr(args, key, value)

    # Finally call the subparser's main function with the processed args
    # and the extras which the preprocessor may have returned
    sys.exit(args.main(args) or 0)

def main(sysargs=None):
    try:
        catkin_main(sysargs)
    except KeyboardInterrupt:
        print('Interrupted by user!')
<commit_msg>shell-verbs: adding instruction to use new command for getting shell verbs (#352)<commit_after># Copyright 2014 Open Source Robotics Foundation, Inc.
#
# 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.

from __future__ import print_function

import argparse
from datetime import date
import os
import pkg_resources
import sys

from catkin_tools.common import is_tty

from catkin_tools.config import get_verb_aliases
from catkin_tools.config import initialize_config

from catkin_tools.terminal_color import fmt
from catkin_tools.terminal_color import set_color
from catkin_tools.terminal_color import test_colors

CATKIN_COMMAND_VERB_GROUP = 'catkin_tools.commands.catkin.verbs'

def list_verbs():
    verbs = []
    for entry_point in pkg_resources.iter_entry_points(group=CATKIN_COMMAND_VERB_GROUP):
        verbs.append(entry_point.name)
    return verbs

def load_verb_description(verb_name):
    for entry_point in pkg_resources.iter_entry_points(group=CATKIN_COMMAND_VERB_GROUP):
        if entry_point.name == verb_name:
            return entry_point.load()

def default_argument_preprocessor(args):
    extras = {}
    return args, extras

def create_subparsers(parser, verbs):
    verbs = sorted(verbs)
    verb_array_str = '[' + ' | '.join(verbs) + ']'
    verb_list_str = 'Call `catkin VERB -h` for help on each verb listed below:\n'
    for verb in verbs:
        desc = load_verb_description(verb)
        verb_list_str += '\n  %s\t%s' % (desc['verb'], desc['description'])

    subparser = parser.add_subparsers(
        title='catkin command',
        metavar=verb_array_str,
        description=verb_list_str,
        dest='verb'
    )

    argument_preprocessors = {}

    for verb in verbs:
        desc = load_verb_description(verb)
        cmd_parser = subparser.add_parser(desc['verb'], description=desc['description'])
        cmd_parser = desc['prepare_arguments'](cmd_parser)

        cmd_parser.set_defaults(main=desc['main'])

        if 'argument_preprocessor' in desc:
            argument_preprocessors[verb] = desc['argument_preprocessor']
        else:
            argument_preprocessors[verb] = default_argument_preprocessor

    return argument_preprocessors

def expand_one_verb_alias(sysargs, verb_aliases, used_aliases):
    """Iterate through sysargs looking for expandable verb aliases.

    When a verb alias is found, sysargs is modified to effectively expand the alias.
    The alias is removed from verb_aliases and added to used_aliases.
    After finding and expanding an alias, this function returns True.
    If no alias is found to be expanded, this function returns False.
    """
    cmd = os.path.basename(sys.argv[0])
    for index, arg in enumerate(sysargs):
        if arg.startswith('-'):
            # Not a verb, continue through the arguments
            continue
        if arg in used_aliases:
            print(fmt(
                "@!@{gf}==>@| Expanding alias '@!@{yf}" + arg +
                "@|' was previously expanded, ignoring this time to prevent infinite recursion."
            ))
        if arg in verb_aliases:
            before = [] if index == 0 else sysargs[:index - 1]
            after = [] if index == len(sysargs) else sysargs[index + 1:]
            sysargs[:] = before + verb_aliases[arg].split() + after
            print(fmt(
                "@!@{gf}==>@| Expanding alias "
                "'@!@{yf}{alias}@|' "
                "from '@{yf}{before} @!{alias}@{boldoff}{after}@|' "
                "to '@{yf}{before} @!{expansion}@{boldoff}{after}@|'"
            ).format(
                alias=arg,
                expansion=verb_aliases[arg],
                before=' '.join([cmd] + before),
                after=(' '.join([''] + after) if after else '')
            ))
            # Prevent the alias from being used again, to prevent infinite recursion
            used_aliases.append(arg)
            del verb_aliases[arg]
            # Return True since one has been found
            return True
        # Return False since no verb alias was found
        return False

def expand_verb_aliases(sysargs, verb_aliases):
    """Expands aliases in sysargs which are found in verb_aliases until none are found."""
    used_aliases = []
    while expand_one_verb_alias(sysargs, verb_aliases, used_aliases):
        pass
    return sysargs

def catkin_main(sysargs):
    # Initialize config
    try:
        initialize_config()
    except RuntimeError as exc:
        sys.exit("Failed to initialize config: {0}".format(exc))

    # Create a top level parser
    parser = argparse.ArgumentParser(
        description="catkin command", formatter_class=argparse.RawDescriptionHelpFormatter)
    add = parser.add_argument
    add('-a', '--list-aliases', action="store_true", default=False,
        help="Lists the current verb aliases and then quits, all other arguments are ignored")
    add('--test-colors', action='store_true', default=False,
        help="Prints a color test pattern to the screen and then quits, all other arguments are ignored")
    add('--version', action='store_true', default=False,
        help="Prints the catkin_tools version.")
    color_control_group = parser.add_mutually_exclusive_group()
    add = color_control_group.add_argument
    add('--force-color', action='store_true', default=False,
        help='Forces catkin to output in color, even when the terminal does not appear to support it.')
    add('--no-color', action='store_true', default=False,
        help='Forces catkin to not use color in the output, regardless of the detect terminal type.')

    # Deprecated, moved to `catkin locate --shell-verbs
    add('--locate-extra-shell-verbs', action='store_true', help=argparse.SUPPRESS)

    # Generate a list of verbs available
    verbs = list_verbs()

    # Create the subparsers for each verb and collect the argument preprocessors
    argument_preprocessors = create_subparsers(parser, verbs)

    # Get verb aliases
    verb_aliases = get_verb_aliases()

    # Setup sysargs
    sysargs = sys.argv[1:] if sysargs is None else sysargs

    # Get colors config
    no_color = False
    force_color = os.environ.get('CATKIN_TOOLS_FORCE_COLOR', False)
    for arg in sysargs:
        if arg == '--no-color':
            no_color = True
        if arg == '--force-color':
            force_color = True

    if no_color or not force_color and not is_tty(sys.stdout):
        set_color(False)

    # Check for version
    if '--version' in sysargs:
        print('catkin_tools {} (C) 2014-{} Open Source Robotics Foundation'.format(
            pkg_resources.get_distribution('catkin_tools').version,
            date.today().year)
        )
        print('catkin_tools is released under the Apache License,'
              ' Version 2.0 (http://www.apache.org/licenses/LICENSE-2.0)')
        print('---')
        print('Using Python {}'.format(''.join(sys.version.split('\n'))))
        sys.exit(0)

    # Deprecated option
    if '--locate-extra-shell-verbs' in sysargs:
        print('Please use `catkin locate --shell-verbs` instead of `catkin --locate-extra-shell-verbs`',
              file=sys.stderr)
        sys.exit(0)

    # Check for --test-colors
    for arg in sysargs:
        if arg == '--test-colors':
            test_colors()
            sys.exit(0)
        if not arg.startswith('-'):
            break

    # Check for --list-aliases
    for arg in sysargs:
        if arg == '--list-aliases' or arg == '-a':
            for alias in sorted(list(verb_aliases.keys())):
                print("{0}: {1}".format(alias, verb_aliases[alias]))
            sys.exit(0)
        if not arg.startswith('-'):
            break

    # Do verb alias expansion
    sysargs = expand_verb_aliases(sysargs, verb_aliases)

    # Determine the verb, splitting arguments into pre and post verb
    verb = None
    pre_verb_args = []
    post_verb_args = []
    for index, arg in enumerate(sysargs):
        # If the arg does not start with a `-` then it is a positional argument
        # The first positional argument must be the verb
        if not arg.startswith('-'):
            verb = arg
            post_verb_args = sysargs[index + 1:]
            break
        # If the `-h` or `--help` option comes before the verb, parse_args
        if arg in ['-h', '--help']:
            parser.parse_args(sysargs)
        # Otherwise it is a pre-verb option
        pre_verb_args.append(arg)

    # Error on no verb provided
    if verb is None:
        print(parser.format_usage())
        sys.exit("Error: No verb provided.")
    # Error on unknown verb provided
    if verb not in verbs:
        print(parser.format_usage())
        sys.exit("Error: Unknown verb '{0}' provided.".format(verb))

    # First allow the verb's argument preprocessor to strip any args
    # and return any "extra" information it wants as a dict
    processed_post_verb_args, extras = argument_preprocessors[verb](post_verb_args)
    # Then allow argparse to process the left over post-verb arguments along
    # with the pre-verb arguments and the verb itself
    args = parser.parse_args(pre_verb_args + [verb] + processed_post_verb_args)
    # Extend the argparse result with the extras from the preprocessor
    for key, value in extras.items():
        setattr(args, key, value)

    # Finally call the subparser's main function with the processed args
    # and the extras which the preprocessor may have returned
    sys.exit(args.main(args) or 0)

def main(sysargs=None):
    try:
        catkin_main(sysargs)
    except KeyboardInterrupt:
        print('Interrupted by user!')
<|endoftext|>


<commit_before>import parsley

class Number (object):
    def __init__(self, value):
        self.value = value

    def __repr__(self):
        return "Number({})".format(self.value)

    def excel(self):
        return self.value

class Cell (object):
    def __init__(self, x, y):
        self.x = x
        self.y = y

    def __repr__(self):
        return "Cell({}, {})".format(repr(self.x), repr(self.y))

    def excel(self):
        return "{}{}".format(self.x, self.y)

class Arithmetic (object):
    def __init__(self, operand1, operator, operand2):
        self.operand1 = operand1
        self.operator = operator
        self.operand2 = operand2

    def __repr__(self):
        return "Arithmetic({}, {}, {})".format(
                repr(self.operand1),
                repr(self.operator),
                repr(self.operand2)
        )

    def excel(self, suppress_parens=False):
        # VisiCalc used left-to-right evaluation of arithmetic operators, unless
        # overriden by parentheses.
        # By only allowing explicitly parenthesised arithmetic operators on the
        # RHS, the grammar ensures that the AST's structure matches VisiCalc's
        # order of evaluation.
        # Here, we parenthesise every operation so that Excel will evaluate them
        # the way we want, rather than applying its natural order of operations.
        if suppress_parens:
            fs = "{}{}{}"
        else:
            fs = "({}{}{})"
        return "({}{}{})".format(
            self.operand1.excel(),
            self.operator,
            self.operand2.excel()
        )

class Value (object):
    def __init__(self, value):
        self.value = value

    def __repr__(self):
        return "Value({})".format(self.value)

    def excel(self):
        if isinstance(self.value, Number):
            return self.value
        elif isinstance(self.value, Arithmetic):
            return "={}".format(self.value.excel(suppress_parens=True))
        else:
            return "={}".format(self.value.excel())

class Label (object):
    def __init__(self, first, rest):
        if first == '"':
            self.value = rest
        else:
            self.value = first + rest

    def __repr__(self):
        return "Label({})".format(self.value)

    def excel(self):
        return self.value

grammar = parsley.makeGrammar(r"""
digit = anything:x ?(x in '0123456789') -> x
letter = anything:x ?(x.upper() in 'ABCDEFGHIJKLMNOPQRSTUVWXYZ') -> x
cell = '+'? letter:x digit:y -> Cell(x, y)
number = '+'? digit+:x -> Number(''.join(x))
arithmetic = (arithmetic | parenthesised_arithmetic | cell | number):o1 ('+' | '-' | '*' | '/'):oper (parenthesised_arithmetic | cell | number):o2 -> Arithmetic(o1, oper, o2)
parenthesised_arithmetic = '(' arithmetic:a ')' -> a
value = (arithmetic | cell | number):x -> Value(x)
label = ('"' | letter):first anything*:rest -> Label(first, ''.join(rest))
cell_content = label | value
""", globals())

while True:
    i = raw_input("> ")
    result = grammar(i).cell_content()
    print(repr(result))
    print(result.excel())
<commit_msg>fix(cell_parser): Omit brackets on outermost arithmetic<commit_after>import parsley

class Number (object):
    def __init__(self, value):
        self.value = value

    def __repr__(self):
        return "Number({})".format(self.value)

    def excel(self):
        return self.value

class Cell (object):
    def __init__(self, x, y):
        self.x = x
        self.y = y

    def __repr__(self):
        return "Cell({}, {})".format(repr(self.x), repr(self.y))

    def excel(self):
        return "{}{}".format(self.x, self.y)

class Arithmetic (object):
    def __init__(self, operand1, operator, operand2):
        self.operand1 = operand1
        self.operator = operator
        self.operand2 = operand2

    def __repr__(self):
        return "Arithmetic({}, {}, {})".format(
                repr(self.operand1),
                repr(self.operator),
                repr(self.operand2)
        )

    def excel(self, suppress_parens=False):
        # VisiCalc used left-to-right evaluation of arithmetic operators, unless
        # overriden by parentheses.
        # By only allowing explicitly parenthesised arithmetic operators on the
        # RHS, the grammar ensures that the AST's structure matches VisiCalc's
        # order of evaluation.
        # Here, we parenthesise every operation so that Excel will evaluate them
        # the way we want, rather than applying its natural order of operations.
        if suppress_parens:
            fs = "{}{}{}"
        else:
            fs = "({}{}{})"
        return fs.format(
            self.operand1.excel(),
            self.operator,
            self.operand2.excel()
        )

class Value (object):
    def __init__(self, value):
        self.value = value

    def __repr__(self):
        return "Value({})".format(self.value)

    def excel(self):
        if isinstance(self.value, Number):
            return self.value
        elif isinstance(self.value, Arithmetic):
            return "={}".format(self.value.excel(suppress_parens=True))
        else:
            return "={}".format(self.value.excel())

class Label (object):
    def __init__(self, first, rest):
        if first == '"':
            self.value = rest
        else:
            self.value = first + rest

    def __repr__(self):
        return "Label({})".format(self.value)

    def excel(self):
        return self.value

grammar = parsley.makeGrammar(r"""
digit = anything:x ?(x in '0123456789') -> x
letter = anything:x ?(x.upper() in 'ABCDEFGHIJKLMNOPQRSTUVWXYZ') -> x
cell = '+'? letter:x digit:y -> Cell(x, y)
number = '+'? digit+:x -> Number(''.join(x))
arithmetic = (arithmetic | parenthesised_arithmetic | cell | number):o1 ('+' | '-' | '*' | '/'):oper (parenthesised_arithmetic | cell | number):o2 -> Arithmetic(o1, oper, o2)
parenthesised_arithmetic = '(' arithmetic:a ')' -> a
value = (arithmetic | cell | number):x -> Value(x)
label = ('"' | letter):first anything*:rest -> Label(first, ''.join(rest))
cell_content = label | value
""", globals())

while True:
    i = raw_input("> ")
    result = grammar(i).cell_content()
    print(repr(result))
    print(result.excel())
<|endoftext|>