<commit_before>"""Model a process lifecycle as a list of functions.
"""
from __future__ import absolute_import, division, print_function, unicode_literals

import re
import sys
import types
import traceback

from dependency_injection import resolve_dependencies

class FunctionNotFound(Exception):
    def __str__(self):
        return "The function '{0}' isn't in this lifecycle.".format(*self.args)

class Lifecycle(object):
    """
    """

    short_circuit = False

    def __init__(self, dotted_name):
        self.module = self._load_module_from_dotted_name(dotted_name)
        self.functions = self._load_functions_from_module(self.module)

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

    def get_names(self):
        return [f.func_name for f in self]

    def insert_after(self, name, newfunc):
        """Insert newfunc in the list right after the function named name.
        """
        self.insert_relative_to(name, newfunc, relative_position=1)

    def insert_before(self, name, newfunc):
        """Insert newfunc in the list right before the function named name.
        """
        self.insert_relative_to(name, newfunc, relative_position=-1)

    def insert_relative_to(self, name, newfunc, relative_position):
        func = self.resolve_name_to_function(name)
        index = self.functions.index(func) + relative_position
        self.functions.insert(index, newfunc)

    def remove(self, name):
        """Remove the function named name from the list.
        """
        func = self.resolve_name_to_function(name)
        self.functions.remove(func)

    def resolve_name_to_function(self, name):
        """Given the name of a function in the list, return the function itself.
        """
        func = None
        for func in self.functions:
            if func.func_name == name:
                break
        if func is None:
            raise FunctionNotFound(name)
        return func

    def run(self, state, through=None):
        """Given a state dictionary, run through the functions in the list.
        """
        if through is not None:
            if through not in self.get_names():
                raise FunctionNotFound(through)
        # XXX bring these back when we've sorted out logging
        #print()

        if 'lifecycle' not in state:    state['lifecycle'] = self
        if 'state' not in state:        state['state'] = state
        if 'exc_info' not in state:     state['exc_info'] = None

        for function in self.functions:
            function_name = function.func_name
            try:
                deps = resolve_dependencies(function, state)
                if 'exc_info' in deps.signature.required and state['exc_info'] is None:
                    pass    # Hook needs an exc_info but we don't have it.
                    #print("{0:>48}  \x1b[33;1mskipped\x1b[0m".format(function_name))
                elif 'exc_info' not in deps.signature.parameters and state['exc_info'] is not None:
                    pass    # Hook doesn't want an exc_info but we have it.
                    #print("{0:>48}  \x1b[33;1mskipped\x1b[0m".format(function_name))
                else:
                    new_state = function(**deps.as_kwargs)
                    #print("{0:>48}  \x1b[32;1mdone\x1b[0m".format(function_name))
                    if new_state is not None:
                        state.update(new_state)
            except:
                #print("{0:>48}  \x1b[31;1mfailed\x1b[0m".format(function_name))
                state['exc_info'] = sys.exc_info()[:2] + (traceback.format_exc().strip(),)
                if self.short_circuit:
                    raise

            if through is not None and function_name == through:
                break

        if state['exc_info'] is not None:
            raise

        return state

    # Helpers for loading from a file.
    # ================================

    def _load_module_from_dotted_name(self, dotted_name):
        class Module(object): pass
        module = Module()  # let's us use getattr to traverse down
        exec 'import {0}'.format(dotted_name) in module.__dict__
        for name in dotted_name.split('.'):
            module = getattr(module, name)
        return module

    def _load_functions_from_module(self, module):
        """Given a module object, return a list of functions from the module, sorted by lineno.
        """
        functions_with_lineno = []
        for name in dir(module):
            if name.startswith('_'):
                continue
            obj = getattr(module, name)
            if type(obj) != types.FunctionType:
                continue
            func = obj
            lineno = func.func_code.co_firstlineno
            functions_with_lineno.append((lineno, func))
        functions_with_lineno.sort()
        return [func for lineno, func in functions_with_lineno]

# Filters
# =======

def by_lambda(filter_lambda):
    def wrap(flow_step):
        def wrapped_flow_step_by_lambda(*args,**kwargs):
            if filter_lambda():
                return flow_step(*args,**kwargs)
        wrapped_flow_step_by_lambda.func_name = flow_step.func_name
        return wrapped_flow_step_by_lambda
    return wrap

def by_regex(regex_tuples, default=True):
    """A filter for flow steps.

    regex_tuples is a list of (regex, filter?) where if the regex matches the
    requested URI, then the flow is applied or not based on if filter? is True
    or False.

    For example:

        from aspen.flows.filter import by_regex

        @by_regex( ( ("/secret/agenda", True), ( "/secret.*", False ) ) )
        def use_public_formatting(request):
            ...

    would call the 'use_public_formatting' flow step only on /secret/agenda
    and any other URLs not starting with /secret.

    """
    regex_res = [ (re.compile(regex), disposition) \
                           for regex, disposition in regex_tuples.iteritems() ]
    def filter_flow_step(flow_step):
        def flow_step_filter(request, *args):
            for regex, disposition in regex_res:
                if regex.matches(request.line.uri):
                    if disposition:
                        return flow_step(*args)
            if default:
                return flow_step(*args)
        flow_step_filter.func_name = flow_step.func_name
        return flow_step_filter
    return filter_flow_step

def by_dict(truthdict, default=True):
    """A filter for hooks.

    truthdict is a mapping of URI -> filter? where if the requested URI is a
    key in the dict, then the hook is applied based on the filter? value.

    """
    def filter_flow_step(flow_step):
        def flow_step_filter(request, *args):
            do_hook = truthdict.get(request.line.uri, default)
            if do_hook:
                return flow_step(*args)
        flow_step_filter.func_name = flow_step.func_name
        return flow_step_filter
    return filter_flow_step
<commit_msg>Add installation docs<commit_after>"""Model a process lifecycle as a list of functions.

Installation
------------

:py:mod:`lifecycle` is available on `GitHub`_ and on `PyPI`_::

    $ pip install lifecycle

We `test <https://travis-ci.org/gittip/lifecycle.py>`_ against
Python 2.6, 2.7, 3.2, and 3.3.

:py:mod:`lifecycle` is in the `public domain`_.

.. _GitHub: https://github.com/gittip/lifecycle.py
.. _PyPI: https://pypi.python.org/pypi/lifecycle
.. _public domain: http://creativecommons.org/publicdomain/zero/1.0/

API
---

"""
from __future__ import absolute_import, division, print_function, unicode_literals

import re
import sys
import types
import traceback

from dependency_injection import resolve_dependencies

class FunctionNotFound(Exception):
    def __str__(self):
        return "The function '{0}' isn't in this lifecycle.".format(*self.args)

class Lifecycle(object):
    """
    """

    short_circuit = False

    def __init__(self, dotted_name):
        self.module = self._load_module_from_dotted_name(dotted_name)
        self.functions = self._load_functions_from_module(self.module)

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

    def get_names(self):
        return [f.func_name for f in self]

    def insert_after(self, name, newfunc):
        """Insert newfunc in the list right after the function named name.
        """
        self.insert_relative_to(name, newfunc, relative_position=1)

    def insert_before(self, name, newfunc):
        """Insert newfunc in the list right before the function named name.
        """
        self.insert_relative_to(name, newfunc, relative_position=-1)

    def insert_relative_to(self, name, newfunc, relative_position):
        func = self.resolve_name_to_function(name)
        index = self.functions.index(func) + relative_position
        self.functions.insert(index, newfunc)

    def remove(self, name):
        """Remove the function named name from the list.
        """
        func = self.resolve_name_to_function(name)
        self.functions.remove(func)

    def resolve_name_to_function(self, name):
        """Given the name of a function in the list, return the function itself.
        """
        func = None
        for func in self.functions:
            if func.func_name == name:
                break
        if func is None:
            raise FunctionNotFound(name)
        return func

    def run(self, state, through=None):
        """Given a state dictionary, run through the functions in the list.
        """
        if through is not None:
            if through not in self.get_names():
                raise FunctionNotFound(through)
        # XXX bring these back when we've sorted out logging
        #print()

        if 'lifecycle' not in state:    state['lifecycle'] = self
        if 'state' not in state:        state['state'] = state
        if 'exc_info' not in state:     state['exc_info'] = None

        for function in self.functions:
            function_name = function.func_name
            try:
                deps = resolve_dependencies(function, state)
                if 'exc_info' in deps.signature.required and state['exc_info'] is None:
                    pass    # Hook needs an exc_info but we don't have it.
                    #print("{0:>48}  \x1b[33;1mskipped\x1b[0m".format(function_name))
                elif 'exc_info' not in deps.signature.parameters and state['exc_info'] is not None:
                    pass    # Hook doesn't want an exc_info but we have it.
                    #print("{0:>48}  \x1b[33;1mskipped\x1b[0m".format(function_name))
                else:
                    new_state = function(**deps.as_kwargs)
                    #print("{0:>48}  \x1b[32;1mdone\x1b[0m".format(function_name))
                    if new_state is not None:
                        state.update(new_state)
            except:
                #print("{0:>48}  \x1b[31;1mfailed\x1b[0m".format(function_name))
                state['exc_info'] = sys.exc_info()[:2] + (traceback.format_exc().strip(),)
                if self.short_circuit:
                    raise

            if through is not None and function_name == through:
                break

        if state['exc_info'] is not None:
            raise

        return state

    # Helpers for loading from a file.
    # ================================

    def _load_module_from_dotted_name(self, dotted_name):
        class Module(object): pass
        module = Module()  # let's us use getattr to traverse down
        exec 'import {0}'.format(dotted_name) in module.__dict__
        for name in dotted_name.split('.'):
            module = getattr(module, name)
        return module

    def _load_functions_from_module(self, module):
        """Given a module object, return a list of functions from the module, sorted by lineno.
        """
        functions_with_lineno = []
        for name in dir(module):
            if name.startswith('_'):
                continue
            obj = getattr(module, name)
            if type(obj) != types.FunctionType:
                continue
            func = obj
            lineno = func.func_code.co_firstlineno
            functions_with_lineno.append((lineno, func))
        functions_with_lineno.sort()
        return [func for lineno, func in functions_with_lineno]

# Filters
# =======

def by_lambda(filter_lambda):
    def wrap(flow_step):
        def wrapped_flow_step_by_lambda(*args,**kwargs):
            if filter_lambda():
                return flow_step(*args,**kwargs)
        wrapped_flow_step_by_lambda.func_name = flow_step.func_name
        return wrapped_flow_step_by_lambda
    return wrap

def by_regex(regex_tuples, default=True):
    """A filter for flow steps.

    regex_tuples is a list of (regex, filter?) where if the regex matches the
    requested URI, then the flow is applied or not based on if filter? is True
    or False.

    For example:

        from aspen.flows.filter import by_regex

        @by_regex( ( ("/secret/agenda", True), ( "/secret.*", False ) ) )
        def use_public_formatting(request):
            ...

    would call the 'use_public_formatting' flow step only on /secret/agenda
    and any other URLs not starting with /secret.

    """
    regex_res = [ (re.compile(regex), disposition) \
                           for regex, disposition in regex_tuples.iteritems() ]
    def filter_flow_step(flow_step):
        def flow_step_filter(request, *args):
            for regex, disposition in regex_res:
                if regex.matches(request.line.uri):
                    if disposition:
                        return flow_step(*args)
            if default:
                return flow_step(*args)
        flow_step_filter.func_name = flow_step.func_name
        return flow_step_filter
    return filter_flow_step

def by_dict(truthdict, default=True):
    """A filter for hooks.

    truthdict is a mapping of URI -> filter? where if the requested URI is a
    key in the dict, then the hook is applied based on the filter? value.

    """
    def filter_flow_step(flow_step):
        def flow_step_filter(request, *args):
            do_hook = truthdict.get(request.line.uri, default)
            if do_hook:
                return flow_step(*args)
        flow_step_filter.func_name = flow_step.func_name
        return flow_step_filter
    return filter_flow_step
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
"""
Created on Sun Jan  3 00:30:36 2016

@author: p
"""
import urllib2
import xmltodict
import pandas as pd
from datetime import datetime
from httplib import BadStatusLine
import matplotlib.pyplot as plt
import numpy as np
import avgMeths

class usgs:
    @staticmethod
    def getelev(x):
        '''
        Uses USGS elevation service to retrieve elevation

        Input
        -----
        x, y = longitude and latitude of point where elevation is desired

        Output
        ------
        ned float elevation of location in meters
        '''
        elev = "http://ned.usgs.gov/epqs/pqs.php?x="+str(x[0])+"&y="+str(x[1])+"&units=Meters&output=xml"
        try:
            response = urllib2.urlopen(elev)
            html = response.read()
            d = xmltodict.parse(html)
            g = float(d['USGS_Elevation_Point_Query_Service']['Elevation_Query']['Elevation'])
        except(BadStatusLine):
            try:
                response = urllib2.urlopen(elev)
                html = response.read()
                d = xmltodict.parse(html)
                g = float(d['USGS_Elevation_Point_Query_Service']['Elevation_Query']['Elevation'])
            except(BadStatusLine):
                print "could not fetch %s" % html
                g = 0
                pass
        return g

    @staticmethod
    def USGSID(x):
        def dms(dec):
            DD = str(int(abs(dec)))
            MM = str(int((abs(dec) - int(DD))*60)).zfill(2)
            SS = str(int(round((((abs(dec) - int(DD))*60) - int(MM))*60, 0))).zfill(2)
            return DD+MM+SS
        return dms(x[1])+dms(x[0])+'01'

    def getInfo(self, html):
        '''
        Input
        -----
        html = location of data to be queried <http://waterservices.usgs.gov>

        Output
        ------
        df = Pandas Dataframe containing data downloaded from USGS
        '''

        try:
            linefile = urllib2.urlopen(html).readlines()
            numlist=[]
            num=0
            for line in linefile:
                if line.startswith("#"):
                    numlist.append(num)
                num += 1
            numlist.append(numlist[-1]+2)
            df = pd.read_table(html, sep="\t",skiprows=numlist)#, comment='#')
            return df
        except(BadStatusLine):
            try:
                linefile = urllib2.urlopen(html).readlines()
                numlist=[]
                num=0
                for line in linefile:
                    if line.startswith("#"):
                        numlist.append(num)
                    num += 1
                numlist.append(numlist[-1]+2)
                df = pd.read_table(html, sep="\t",skiprows=numlist)#, comment='#')
                return df
            except(BadStatusLine):
                print "could not fetch %s" % html
                pass

    def getQInfo(self,html):
        '''
        Input
        -----
        html = location of data to be queried <http://waterservices.usgs.gov>

        Output
        ------
        df = Pandas Dataframe containing data downloaded from USGS
        '''

        response = urllib2.urlopen(html)
        htmlresp = response.read()

        USGSdict = xmltodict.parse(htmlresp)
        f={}
        for i in range(len(USGSdict['ns1:timeSeriesResponse']['ns1:timeSeries'])):
            f[USGSdict['ns1:timeSeriesResponse']['ns1:timeSeries'][i]['@name'][5:13]] = pd.DataFrame(USGSdict['ns1:timeSeriesResponse']['ns1:timeSeries'][i]['ns1:values']['ns1:value'])
        df = pd.concat(f)
        df.reset_index(inplace=True)
        df['datetime'] = pd.to_datetime(df['@dateTime'])
        df['Q_cfs'] = pd.to_numeric(df['#text'])
        df.set_index('datetime',inplace=True)
        df.drop([u'level_1', '#text', '@dateTime'],axis=1,inplace=True)
        df.columns = ['site_no','qualifiers','value']
        return df

    def parsesitelist(self, ListOfSites):
        '''
        Takes a list of USGS sites and turns it into a string format that can be used in the html REST format
        '''
        siteno = str(ListOfSites).replace(" ","")
        siteno = siteno.replace("]","")
        siteno = siteno.replace("[","")
        siteno = siteno.replace("','",",")
        siteno = siteno.replace("'","")
        siteno = siteno.replace('"',"")
        return siteno

    def getStationInfo(self, sitenos):
        '''
        INPUT
        -----
        sitenos = list of usgs sites to get site info

        RETURNS
        -------
        siteinfo = pandas dataframe with station info
        '''
        siteno = self.parsesitelist(sitenos)
        html = "http://waterservices.usgs.gov/nwis/site/?format=rdb&sites=" + siteno + "&siteOutput=expanded"
        siteinfo = self.getInfo(html)
        return siteinfo

    def getStationInfoFromHUC(self, HUCS, sitetype=['GW'],datatype=['gw']):
        '''
        HUCS = list of HUCS to find sites
        sitetype = list of types of site you are searching for; options include
            ST = Stream,SP = Spring,GW = Groundwater, SB = Other underground
            defaults to GW
        datatype = list of data types you are searching for; options include
            iv = instantaneous ,dv = daily values, sv = site visit, gw = groundwater level, qw = water quality, id =historical instantaneous
            defaults to gw
        RETURNS
        -------
        Pandas dataframe of sites in HUC
        '''
        sitetypes = self.parsesitelist(sitetype)
        datatypes = self.parsesitelist(datatype)
        HUC = self.parsesitelist(HUCS)
        stationhtml = "http://waterservices.usgs.gov/nwis/site/?format=rdb,1.0&huc=" + str(HUC) + "&siteType=" + sitetypes + "&hasDataTypeCd=" + datatypes
        siteinfo = self.getInfo(stationhtml)
        return siteinfo

    def getStationsfromHUC(self, HUCS):
        HUC = self.parsesitelist(HUCS)
        stationhtml = "http://waterservices.usgs.gov/nwis/site/?format=rdb,1.0&huc=" + str(HUC) + "&siteType=GW&hasDataTypeCd=gw"
        sites = self.getInfo(stationhtml)
        stations = list(sites['site_no'].values)
        stations = [str(i) for i in stations]
        return stations

    def getWLfromHUC(self, HUCS):
        HUC = self.parsesitelist(HUCS)
        html = "http://waterservices.usgs.gov/nwis/gwlevels/?format=rdb&huc="+str(HUC)+"&startDT=1800-01-01&endDT="+str(datetime.today().year)+"-"+str(datetime.today().month).zfill(2)+"-"+str(datetime.today().day).zfill(2)
        wls = self.getInfo(html)
        return wls

    def getWLfromSite(self, sitenos):
        siteno = self.parsesitelist(sitenos)
        html = "http://waterservices.usgs.gov/nwis/gwlevels/?format=rdb&sites="+str(siteno)+"&startDT=1800-01-01&endDT="+str(datetime.today().year)+"-"+str(datetime.today().month).zfill(2)+"-"+str(datetime.today().day).zfill(2)
        wls = self.getInfo(html)
        return wls

    def getQfromSites(self, ListOfSites):
        '''
        get discharge data from site list
        '''
        siteno = self.parsesitelist(ListOfSites)
        html = "http://waterservices.usgs.gov/nwis/dv/?format=waterml,1.1&sites="+str(siteno)+"&parameterCd=00060"+"&startDT=1800-01-01&endDT="+str(datetime.today().year)+"-"+str(datetime.today().month).zfill(2)+"-"+str(datetime.today().day).zfill(2)
        wls = self.getQInfo(html)
        #wls.columns = ['agency_cd','site_no','datetime','discharge_cfs','quality_code']
        return wls

    def getQfromHUC(self, HUCS):
        '''
        get discharge data from HUC
        '''
        siteno = self.parsesitelist(HUCS)
        html = "http://waterservices.usgs.gov/nwis/dv/?format=rdb&huc="+str(HUC)+"&parameterCd=00060"+"&startDT=1800-01-01&endDT="+str(datetime.today().year)+"-"+str(datetime.today().month).zfill(2)+"-"+str(datetime.today().day).zfill(2)
        wls = self.getInfo(html)
        return wls

    def cleanGWL(self, data):
        '''
        Drops water level data of suspect quality based on lev_status_cd

        returns Pandas DataFrame
        '''
        CleanData = data[~data['lev_status_cd'].isin(['Z', 'R', 'V', 'P', 'O', 'F', 'W', 'G', 'S', 'C', 'E', 'N'])]
        return CleanData

    def WLStatdf(self, siteinfo, data):
        '''
        generates average water level statistics for a huc or list of hucs
        INPUT
        -----
        siteinfo = pandas dataframe of site information of nwis sites (made using a get station info function)
        data = pandas dataframe of data from nwis sites (made using a get station data function)

        RETURNS
        -------
        wlLongStatsGroups = pandas dataframe of standardized water levels over duration of measurement
        wlLongStatsGroups2 = pandas dataframe of change in average water levels over duration of measurement
        '''

        try:
            data.drop([u'agency_cd', u'site_tp_cd'], inplace=True, axis=1)
        except(ValueError):
            pass
        stationWL = pd.merge(data, siteinfo, on='site_no', how='left')

        stationWL['date'], stationWL['Year'], stationWL['Month'] = zip(*stationWL['lev_dt'].apply(lambda x: avgMeths.getyrmnth(x),1))
        stationWL.reset_index(inplace=True)
        stationWL.set_index('date',inplace=True)
        stationWL = self.cleanGWL(stationWL)
        # get averages by year, month, and site number
        grpstat = stationWL.groupby('site_no')['lev_va'].agg([np.std,np.mean,np.median, np.min,np.max,np.size]).reset_index()
        USGS_Site_Inf = stationWL.groupby('site_no')['lev_dt'].agg([np.min,np.max,np.size]).reset_index()
        USGS_Site_Info = USGS_Site_Inf[USGS_Site_Inf['size']>50]
        wlLong = stationWL[stationWL['site_no'].isin(list(USGS_Site_Info['site_no'].values))]
        wlLongStats = pd.merge(wlLong,grpstat, on='site_no', how='left')
        wlLongStats['stdWL'] = wlLongStats[['lev_va','mean','std']].apply(lambda x: avgMeths.stndrd(x),1 )
        wlLongStats['YRMO'] = wlLongStats[['Year','Month']].apply(lambda x: avgMeths.yrmo(x),1)
        wlLongStats['date'] = wlLongStats[['Year','Month']].apply(lambda x: avgMeths.adddate(x),1)
        self.wlMonthPlot = wlLongStats.groupby(['Month'])['stdWL'].mean().to_frame().plot()
        wlLongStats['levDiff'] = wlLongStats['lev_va'].diff()

        wlLongStatsGroups = wlLongStats.groupby(['date'])['stdWL'].agg({'mean':np.mean,'median':np.median,
                                                                        'standard':np.std, 'cnt':(lambda x: np.count_nonzero(~np.isnan(x))),
                                                                        'err':(lambda x: 1.96*avgMeths.sumstats(x))})
        wlLongStatsGroups2 = wlLongStats.groupby(['date'])['levDiff'].agg({'mean':np.mean,'median':np.median, 'standard':np.std, 'cnt':(lambda x: np.count_nonzero(~np.isnan(x))), 'err':(lambda x: 1.96*avgMeths.sumstats(x))})

        wlLongStatsGroups['meanpluserr'] = wlLongStatsGroups['mean'] + wlLongStatsGroups['err']
        wlLongStatsGroups['meanminuserr'] = wlLongStatsGroups['mean'] - wlLongStatsGroups['err']

        wlLongStatsGroups2['meanpluserr'] = wlLongStatsGroups2['mean'] + wlLongStatsGroups2['err']
        wlLongStatsGroups2['meanminuserr'] = wlLongStatsGroups2['mean'] - wlLongStatsGroups2['err']

        return wlLongStatsGroups, wlLongStatsGroups2

    def HUCplot(self, siteinfo, data):
        '''
        Generates Statistics plots of NWIS WL data

        INPUT
        -----
        siteinfo = pandas dataframe of site information of nwis sites (made using a get station info function)
        data = pandas dataframe of data from nwis sites (made using a get station data function)

        RETURNS
        -------
        self.stand = standardized statistics
        self.diffs = difference statistics
        self.zPlot = seasonal variation plot
        self.wlPlot = plots of stadardized and differenced wl variations of duration of measurement
        '''
        df1,df2 = self.WLStatdf(siteinfo, data)
        wlLongSt = df1[df1['cnt']>2]
        wlLongSt2 = df2[df2['cnt']>2]

        self.stand = wlLongSt
        self.diffs = wlLongSt2

        fig1 = plt.figure()
        x = wlLongSt.index
        y = wlLongSt['mean']
        plt.plot(x,y,label='Average Groundwater Level Variation')
        plt.fill_between(wlLongSt.index, wlLongSt['meanpluserr'], wlLongSt['meanminuserr'],
                         facecolor='blue', alpha=0.4, linewidth=0.5, label= "Std Error")
        plt.grid(which='both')
        plt.ylabel('Depth to Water z-score')
        plt.xticks(rotation=45)
        self.zPlot = fig1

        fig2 = plt.figure()
        x = wlLongSt2.index
        y = wlLongSt2['mean']
        plt.plot(x,y,label='Average Groundwater Level Changes')
        plt.fill_between(wlLongSt.index, wlLongSt2['meanpluserr'], wlLongSt2['meanminuserr'],
                         facecolor='blue', alpha=0.4, linewidth=0.5, label= "Std Error")
        plt.grid(which='both')
        plt.ylabel('Change in Average Depth to Water (ft)')
        plt.xticks(rotation=45)
        self.wlPlot = fig2

        return fig1, fig2, wlLongSt, wlLongSt2

    def get_huc(self, x):
        import json
        """
        Receive the content of ``url``, parse it as JSON and return the object.

        Parameters
        ----------
        x = [longitude, latitude]

        Returns
        -------
        HUC12, HUC12_Name
        """
        huc_url = 'http://services.nationalmap.gov/arcgis/rest/services/nhd/mapserver/6/query?geometry='+ str(x[0])\
    +'%2C'+ str(x[1]) + '&geometryType=esriGeometryPoint&inSR=4326&spatialRel=esriSpatialRelIntersects&relationParam=&outFields=HUC12%2C'+\
    '+Name&returnGeometry=false&returnDistinctValues=true&f=pjson'
        response = urllib2.urlopen(huc_url)
        data = response.read().decode("utf-8")
        return json.loads(data)['features'][0]['attributes']['HUC12'],json.loads(data)['features'][0]['attributes']['NAME']
<commit_msg>Update usgsGis.py<commit_after># -*- coding: utf-8 -*-
"""
Created on Sun Jan  3 00:30:36 2016

@author: p
"""
import urllib2
import xmltodict
import pandas as pd
from datetime import datetime
from httplib import BadStatusLine
import matplotlib.pyplot as plt
import numpy as np
import avgMeths

class usgs:
    @staticmethod
    def getelev(x):
        '''
        Uses USGS elevation service to retrieve elevation

        Input
        -----
        x, y = longitude and latitude of point where elevation is desired

        Output
        ------
        ned float elevation of location in meters
        '''
        elev = "http://ned.usgs.gov/epqs/pqs.php?x="+str(x[0])+"&y="+str(x[1])+"&units=Meters&output=xml"
        try:
            response = urllib2.urlopen(elev)
            html = response.read()
            d = xmltodict.parse(html)
            g = float(d['USGS_Elevation_Point_Query_Service']['Elevation_Query']['Elevation'])
        except(BadStatusLine):
            try:
                response = urllib2.urlopen(elev)
                html = response.read()
                d = xmltodict.parse(html)
                g = float(d['USGS_Elevation_Point_Query_Service']['Elevation_Query']['Elevation'])
            except(BadStatusLine):
                print "could not fetch %s" % html
                g = 0
                pass
        return g

    @staticmethod
    def USGSID(x):
        def dms(dec):
            DD = str(int(abs(dec)))
            MM = str(int((abs(dec) - int(DD))*60)).zfill(2)
            SS = str(int(round((((abs(dec) - int(DD))*60) - int(MM))*60, 0))).zfill(2)
            return DD+MM+SS
        return dms(x[1])+dms(x[0])+'01'

    def getInfo(self, html):
        '''
        Input
        -----
        html = location of data to be queried <http://waterservices.usgs.gov>

        Output
        ------
        df = Pandas Dataframe containing data downloaded from USGS
        '''

        try:
            linefile = urllib2.urlopen(html).readlines()
            numlist=[]
            num=0
            for line in linefile:
                if line.startswith("#"):
                    numlist.append(num)
                num += 1
            numlist.append(numlist[-1]+2)
            df = pd.read_table(html, sep="\t",skiprows=numlist)#, comment='#')
            return df
        except(BadStatusLine):
            try:
                linefile = urllib2.urlopen(html).readlines()
                numlist=[]
                num=0
                for line in linefile:
                    if line.startswith("#"):
                        numlist.append(num)
                    num += 1
                numlist.append(numlist[-1]+2)
                df = pd.read_table(html, sep="\t",skiprows=numlist)#, comment='#')
                return df
            except(BadStatusLine):
                print "could not fetch %s" % html
                pass

    def getQInfo(self,html):
        '''
        Input
        -----
        html = location of data to be queried <http://waterservices.usgs.gov>

        Output
        ------
        df = Pandas Dataframe containing data downloaded from USGS
        '''
        linefile = urllib2.urlopen(html).readlines()
        numlist=[]
        num=0
        for line in linefile:
            if line.startswith("#"):
                numlist.append(num)
            num += 1
        numlist.append(numlist[-1]+2)
        df = pd.read_table(html, sep="\t",skiprows=numlist, index_col='datetime',parse_dates=True,na)
        df.columns = ['agency','site','discharge','qual']
        df['discharge'] = pd.to_numeric(df['discharge'])
        return df

    def parsesitelist(self, ListOfSites):
        '''
        Takes a list of USGS sites and turns it into a string format that can be used in the html REST format
        '''
        siteno = str(ListOfSites).replace(" ","")
        siteno = siteno.replace("]","")
        siteno = siteno.replace("[","")
        siteno = siteno.replace("','",",")
        siteno = siteno.replace("'","")
        siteno = siteno.replace('"',"")
        return siteno

    def getStationInfo(self, sitenos):
        '''
        INPUT
        -----
        sitenos = list of usgs sites to get site info

        RETURNS
        -------
        siteinfo = pandas dataframe with station info
        '''
        siteno = self.parsesitelist(sitenos)
        html = "http://waterservices.usgs.gov/nwis/site/?format=rdb&sites=" + siteno + "&siteOutput=expanded"
        siteinfo = self.getInfo(html)
        return siteinfo

    def getStationInfoFromHUC(self, HUCS, sitetype=['GW'],datatype=['gw']):
        '''
        HUCS = list of HUCS to find sites
        sitetype = list of types of site you are searching for; options include
            ST = Stream,SP = Spring,GW = Groundwater, SB = Other underground
            defaults to GW
        datatype = list of data types you are searching for; options include
            iv = instantaneous ,dv = daily values, sv = site visit, gw = groundwater level, qw = water quality, id =historical instantaneous
            defaults to gw
        RETURNS
        -------
        Pandas dataframe of sites in HUC
        '''
        sitetypes = self.parsesitelist(sitetype)
        datatypes = self.parsesitelist(datatype)
        HUC = self.parsesitelist(HUCS)
        stationhtml = "http://waterservices.usgs.gov/nwis/site/?format=rdb,1.0&huc=" + str(HUC) + "&siteType=" + sitetypes + "&hasDataTypeCd=" + datatypes
        siteinfo = self.getInfo(stationhtml)
        return siteinfo

    def getStationsfromHUC(self, HUCS):
        HUC = self.parsesitelist(HUCS)
        stationhtml = "http://waterservices.usgs.gov/nwis/site/?format=rdb,1.0&huc=" + str(HUC) + "&siteType=GW&hasDataTypeCd=gw"
        sites = self.getInfo(stationhtml)
        stations = list(sites['site_no'].values)
        stations = [str(i) for i in stations]
        return stations

    def getWLfromHUC(self, HUCS):
        HUC = self.parsesitelist(HUCS)
        html = "http://waterservices.usgs.gov/nwis/gwlevels/?format=rdb&huc="+str(HUC)+"&startDT=1800-01-01&endDT="+str(datetime.today().year)+"-"+str(datetime.today().month).zfill(2)+"-"+str(datetime.today().day).zfill(2)
        wls = self.getInfo(html)
        return wls

    def getWLfromSite(self, sitenos):
        siteno = self.parsesitelist(sitenos)
        html = "http://waterservices.usgs.gov/nwis/gwlevels/?format=rdb&sites="+str(siteno)+"&startDT=1800-01-01&endDT="+str(datetime.today().year)+"-"+str(datetime.today().month).zfill(2)+"-"+str(datetime.today().day).zfill(2)
        wls = self.getInfo(html)
        return wls

    def getQfromSites(self, ListOfSites):
        '''
        get discharge data from site list
        '''
        siteno = self.parsesitelist(ListOfSites)
        html = "http://waterservices.usgs.gov/nwis/dv/?format=waterml,1.1&sites="+str(siteno)+"&parameterCd=00060"+"&startDT=1800-01-01&endDT="+str(datetime.today().year)+"-"+str(datetime.today().month).zfill(2)+"-"+str(datetime.today().day).zfill(2)
        wls = self.getQInfo(html)
        #wls.columns = ['agency_cd','site_no','datetime','discharge_cfs','quality_code']
        return wls

    def getQfromHUC(self, HUCS):
        '''
        get discharge data from HUC
        '''
        siteno = self.parsesitelist(HUCS)
        html = "http://waterservices.usgs.gov/nwis/dv/?format=rdb&huc="+str(HUC)+"&parameterCd=00060"+"&startDT=1800-01-01&endDT="+str(datetime.today().year)+"-"+str(datetime.today().month).zfill(2)+"-"+str(datetime.today().day).zfill(2)
        wls = self.getInfo(html)
        return wls

    def cleanGWL(self, data):
        '''
        Drops water level data of suspect quality based on lev_status_cd

        returns Pandas DataFrame
        '''
        CleanData = data[~data['lev_status_cd'].isin(['Z', 'R', 'V', 'P', 'O', 'F', 'W', 'G', 'S', 'C', 'E', 'N'])]
        return CleanData

    def WLStatdf(self, siteinfo, data):
        '''
        generates average water level statistics for a huc or list of hucs
        INPUT
        -----
        siteinfo = pandas dataframe of site information of nwis sites (made using a get station info function)
        data = pandas dataframe of data from nwis sites (made using a get station data function)

        RETURNS
        -------
        wlLongStatsGroups = pandas dataframe of standardized water levels over duration of measurement
        wlLongStatsGroups2 = pandas dataframe of change in average water levels over duration of measurement
        '''

        try:
            data.drop([u'agency_cd', u'site_tp_cd'], inplace=True, axis=1)
        except(ValueError):
            pass
        stationWL = pd.merge(data, siteinfo, on='site_no', how='left')

        stationWL['date'], stationWL['Year'], stationWL['Month'] = zip(*stationWL['lev_dt'].apply(lambda x: avgMeths.getyrmnth(x),1))
        stationWL.reset_index(inplace=True)
        stationWL.set_index('date',inplace=True)
        stationWL = self.cleanGWL(stationWL)
        # get averages by year, month, and site number
        grpstat = stationWL.groupby('site_no')['lev_va'].agg([np.std,np.mean,np.median, np.min,np.max,np.size]).reset_index()
        USGS_Site_Inf = stationWL.groupby('site_no')['lev_dt'].agg([np.min,np.max,np.size]).reset_index()
        USGS_Site_Info = USGS_Site_Inf[USGS_Site_Inf['size']>50]
        wlLong = stationWL[stationWL['site_no'].isin(list(USGS_Site_Info['site_no'].values))]
        wlLongStats = pd.merge(wlLong,grpstat, on='site_no', how='left')
        wlLongStats['stdWL'] = wlLongStats[['lev_va','mean','std']].apply(lambda x: avgMeths.stndrd(x),1 )
        wlLongStats['YRMO'] = wlLongStats[['Year','Month']].apply(lambda x: avgMeths.yrmo(x),1)
        wlLongStats['date'] = wlLongStats[['Year','Month']].apply(lambda x: avgMeths.adddate(x),1)
        self.wlMonthPlot = wlLongStats.groupby(['Month'])['stdWL'].mean().to_frame().plot()
        wlLongStats['levDiff'] = wlLongStats['lev_va'].diff()

        wlLongStatsGroups = wlLongStats.groupby(['date'])['stdWL'].agg({'mean':np.mean,'median':np.median,
                                                                        'standard':np.std, 'cnt':(lambda x: np.count_nonzero(~np.isnan(x))),
                                                                        'err':(lambda x: 1.96*avgMeths.sumstats(x))})
        wlLongStatsGroups2 = wlLongStats.groupby(['date'])['levDiff'].agg({'mean':np.mean,'median':np.median, 'standard':np.std, 'cnt':(lambda x: np.count_nonzero(~np.isnan(x))), 'err':(lambda x: 1.96*avgMeths.sumstats(x))})

        wlLongStatsGroups['meanpluserr'] = wlLongStatsGroups['mean'] + wlLongStatsGroups['err']
        wlLongStatsGroups['meanminuserr'] = wlLongStatsGroups['mean'] - wlLongStatsGroups['err']

        wlLongStatsGroups2['meanpluserr'] = wlLongStatsGroups2['mean'] + wlLongStatsGroups2['err']
        wlLongStatsGroups2['meanminuserr'] = wlLongStatsGroups2['mean'] - wlLongStatsGroups2['err']

        return wlLongStatsGroups, wlLongStatsGroups2

    def HUCplot(self, siteinfo, data):
        '''
        Generates Statistics plots of NWIS WL data

        INPUT
        -----
        siteinfo = pandas dataframe of site information of nwis sites (made using a get station info function)
        data = pandas dataframe of data from nwis sites (made using a get station data function)

        RETURNS
        -------
        self.stand = standardized statistics
        self.diffs = difference statistics
        self.zPlot = seasonal variation plot
        self.wlPlot = plots of stadardized and differenced wl variations of duration of measurement
        '''
        df1,df2 = self.WLStatdf(siteinfo, data)
        wlLongSt = df1[df1['cnt']>2]
        wlLongSt2 = df2[df2['cnt']>2]

        self.stand = wlLongSt
        self.diffs = wlLongSt2

        fig1 = plt.figure()
        x = wlLongSt.index
        y = wlLongSt['mean']
        plt.plot(x,y,label='Average Groundwater Level Variation')
        plt.fill_between(wlLongSt.index, wlLongSt['meanpluserr'], wlLongSt['meanminuserr'],
                         facecolor='blue', alpha=0.4, linewidth=0.5, label= "Std Error")
        plt.grid(which='both')
        plt.ylabel('Depth to Water z-score')
        plt.xticks(rotation=45)
        self.zPlot = fig1

        fig2 = plt.figure()
        x = wlLongSt2.index
        y = wlLongSt2['mean']
        plt.plot(x,y,label='Average Groundwater Level Changes')
        plt.fill_between(wlLongSt.index, wlLongSt2['meanpluserr'], wlLongSt2['meanminuserr'],
                         facecolor='blue', alpha=0.4, linewidth=0.5, label= "Std Error")
        plt.grid(which='both')
        plt.ylabel('Change in Average Depth to Water (ft)')
        plt.xticks(rotation=45)
        self.wlPlot = fig2

        return fig1, fig2, wlLongSt, wlLongSt2

    def get_huc(self, x):
        import json
        """
        Receive the content of ``url``, parse it as JSON and return the object.

        Parameters
        ----------
        x = [longitude, latitude]

        Returns
        -------
        HUC12, HUC12_Name
        """
        huc_url = 'http://services.nationalmap.gov/arcgis/rest/services/nhd/mapserver/6/query?geometry='+ str(x[0])\
    +'%2C'+ str(x[1]) + '&geometryType=esriGeometryPoint&inSR=4326&spatialRel=esriSpatialRelIntersects&relationParam=&outFields=HUC12%2C'+\
    '+Name&returnGeometry=false&returnDistinctValues=true&f=pjson'
        response = urllib2.urlopen(huc_url)
        data = response.read().decode("utf-8")
        return json.loads(data)['features'][0]['attributes']['HUC12'],json.loads(data)['features'][0]['attributes']['NAME']
<|endoftext|>


<commit_before>"""Test suite for abdi_processrepo."""
import unittest

import phldef_conduit
import phlmail_mocksender

import abdi_processrepo
import abdmail_mailer
import abdt_conduitmock
import abdt_naming
import abdtst_devphabgit

#import phlsys_conduit

#import abdt_conduit

# factors affecting a review:
#  age of the revisions
#  editing the review page
#  new revisions on the review branch
#  rewriting history on the review branch
#  author names
#  author accounts
#  base branch
#  availability of the git repo
#  availability of the phabricator instance

class Test(unittest.TestCase):

    def __init__(self, data):
        super(Test, self).__init__(data)
        self.reviewer = phldef_conduit.ALICE.user
        self.author_account = phldef_conduit.BOB
        self.conduit = None
        self.clone = None
        self.dev_phab = None
        self.mailer = None
        self.mock_sender = None

    def setUp(self):
        self.dev_phab = abdtst_devphabgit.Collaboration(
            self.author_account, phldef_conduit.PHAB)

        self.clone = self.dev_phab.clone

        self.dev_phab.dev_commit_new_file("README")
        self.dev_phab.dev_push_branch("master")
        self.dev_phab.phab_fetch()

        # sys_conduit = phlsys_conduit.Conduit(
        #     phldef_conduit.TEST_URI,
        #     phldef_conduit.PHAB.user,
        #     phldef_conduit.PHAB.certificate)

        #self.conduit = abdt_conduit.Conduit(sys_conduit)
        self.conduit = abdt_conduitmock.ConduitMock()

        self.mock_sender = phlmail_mocksender.MailSender()
        self.mailer = abdmail_mailer.Mailer(
            self.mock_sender,
            ["admin@server.test"],
            "http://server.fake/testrepo.git",
            "http://phabricator.server.fake/")

    def _phabUpdate(self):
        self.dev_phab.phab_fetch()
        abdi_processrepo.processUpdatedRepo(
            self.conduit, self.clone, "origin", self.mailer)

    def _phabUpdateWithExpectationsHelper(
            self, total=None, bad=None, emails=None):
        abdi_processrepo.processUpdatedRepo(
            self.conduit, self.clone, "origin", self.mailer)
        if total is not None:
            self.assertEqual(
                self.dev_phab.count_phab_working_branches(), total)
        if bad is not None:
            self.assertEqual(
                self.dev_phab.count_phab_bad_working_branches(), bad)
        if emails is not None:
            self.assertEqual(len(self.mock_sender.mails), emails)

    def _phabUpdateWithExpectations(self, total=None, bad=None, emails=None):
        self.dev_phab.phab_fetch()

        # multiple updates should have the same result if we are
        # not fetching and assuming the data in Phabricator
        # doesn't change.
        self._phabUpdateWithExpectationsHelper(total, bad, emails)
        self._phabUpdateWithExpectationsHelper(total, bad, emails)
        self._phabUpdateWithExpectationsHelper(total, bad, emails)

    def _devSetAuthorAccount(self, account):
        self.dev_phab.dev_clone.set_name_email(account.user, account.email)

    def _phabSetAuthorAccount(self, account):
        self.clone.set_name_email(account.user, account.email)

    def _devPushNewFile(
            self, filename, has_reviewer=True, has_plan=True, contents=""):
        reviewer = self.reviewer if has_reviewer else None
        plan = "testplan" if has_plan else None
        self.dev_phab.dev_commit_new_file(filename, plan, reviewer, contents)
        self.dev_phab.dev_push()

    def _getTheOnlyReviewId(self):
        branches = self.clone.get_remote_branches("origin")
        wbList = abdt_naming.getWorkingBranches(branches)
        self.assertEqual(len(wbList), 1)
        wb = wbList[0]
        return wb.id

    def _acceptTheOnlyReview(self):
        reviewid = self._getTheOnlyReviewId()
        self.conduit.accept_revision_as_user(reviewid, self.reviewer)

    def _abandonTheOnlyReview(self):
        reviewid = self._getTheOnlyReviewId()
        self.conduit.abandon_revision(reviewid)

    def test_nothingToDo(self):
        # nothing to process
        abdi_processrepo.processUpdatedRepo(
            self.conduit, self.clone, "origin", self.mailer)

    def test_simpleWorkflow(self):
        self.dev_phab.dev_checkout_push_new_branch(
            "ph-review/simpleWorkflow/master")
        self._devPushNewFile("NEWFILE")
        self._phabUpdateWithExpectations(total=1, bad=0)
        self._devPushNewFile("NEWFILE2")
        self._phabUpdateWithExpectations(total=1, bad=0)
        self._acceptTheOnlyReview()
        self._phabUpdateWithExpectations(total=0, bad=0, emails=0)

    def test_badMsgWorkflow(self):
        self.dev_phab.dev_checkout_push_new_branch(
            "ph-review/badMsgWorkflow/master")
        self._devPushNewFile("NEWFILE", has_plan=False)
        self._phabUpdateWithExpectations(total=1, bad=0)
        self._acceptTheOnlyReview()
        self._phabUpdateWithExpectations(total=0, bad=0, emails=0)

    def test_noReviewerWorkflow(self):
        self.dev_phab.dev_checkout_push_new_branch(
            "ph-review/noReviewerWorkflow/master")
        self._devPushNewFile("NEWFILE", has_reviewer=False)
        self._phabUpdateWithExpectations(total=1, bad=0)
        self._acceptTheOnlyReview()
        self._phabUpdateWithExpectations(total=0, bad=0, emails=0)

    def test_badBaseWorkflow(self):
        self.dev_phab.dev_checkout_push_new_branch(
            "ph-review/badBaseWorkflow/blaster")
        self._devPushNewFile("NEWFILE", has_plan=False)
        self._phabUpdateWithExpectations(total=1, bad=1)

        # delete the bad branch
        self.dev_phab.dev_push_delete_branch(
            "ph-review/badBaseWorkflow/blaster")

        self._phabUpdateWithExpectations(total=0, bad=0, emails=0)

    def test_withReservedBranch(self):
        self.dev_phab.dev_checkout_push_new_branch("dev/phab/reserve")
        self.dev_phab.dev_checkout_push_new_branch("ph-review/reserve/master")
        self._devPushNewFile("reserve")
        self._phabUpdateWithExpectations(total=1, bad=0)
        self._acceptTheOnlyReview()
        self._phabUpdateWithExpectations(total=0, bad=0, emails=0)

    def test_noBaseWorkflow(self):
        self.dev_phab.dev_checkout_push_new_branch("ph-review/noBaseWorkflow")
        self._devPushNewFile("NEWFILE", has_plan=False)

        # TODO: handle no base properly
        # self._phabUpdateWithExpectations(total=1, bad=1)

        # delete the bad branch
        self.dev_phab.dev_push_delete_branch("ph-review/noBaseWorkflow")

        self._phabUpdateWithExpectations(total=0, bad=0, emails=0)

    # TODO: test_notBasedWorkflow
    # TODO: test_noCommitWorkflow

    def test_badAuthorWorkflow(self):
        self._devSetAuthorAccount(phldef_conduit.NOTAUSER)
        self.dev_phab.dev_checkout_push_new_branch(
            "ph-review/badAuthorWorkflow/master")
        self._devPushNewFile("NEWFILE")
        self._phabUpdateWithExpectations(total=1, bad=1, emails=1)
        self._devPushNewFile("NEWFILE2")
        self._phabUpdateWithExpectations(total=1, bad=1, emails=2)
        self.dev_phab.dev_reset_branch_to_master(
            "ph-review/badAuthorWorkflow/master")
        self._devSetAuthorAccount(self.author_account)
        self._devPushNewFile("NEWFILE")
        self._phabUpdateWithExpectations(total=1, bad=0, emails=2)
        self._acceptTheOnlyReview()
        self._phabUpdateWithExpectations(total=0, bad=0, emails=2)

    def test_abandonedWorkflow(self):
        self.dev_phab.dev_checkout_push_new_branch(
            "ph-review/abandonedWorkflow/master")
        self._devPushNewFile("NEWFILE")
        self._phabUpdateWithExpectations(total=1, bad=0)
        self._abandonTheOnlyReview()
        self._phabUpdateWithExpectations(total=1, bad=0)
        self._devPushNewFile("NEWFILE2")
        self._phabUpdateWithExpectations(total=1, bad=0)
        self._acceptTheOnlyReview()
        self._phabUpdateWithExpectations(total=0, bad=0, emails=0)

    def test_emptyMergeWorkflow(self):
        self.dev_phab.dev_checkout_push_new_branch("temp/emptyMerge/master")
        self._devPushNewFile("NEWFILE")
        self._phabUpdateWithExpectations(total=0, bad=0)

        # move back to master and land a conflicting change
        self.dev_phab.dev_checkout("master")
        self.dev_phab.dev_checkout_push_new_branch(
            "ph-review/emptyMerge/master")
        self._devPushNewFile("NEWFILE")
        self._phabUpdateWithExpectations(total=1, bad=0)
        self._acceptTheOnlyReview()
        self._phabUpdateWithExpectations(total=0, bad=0)

        # move back to original and try to push and land
        self.dev_phab.dev_checkout("temp/emptyMerge/master")
        self.dev_phab.dev_checkout_push_new_branch(
            "ph-review/emptyMerge2/master")
        self._phabUpdateWithExpectations(total=1, bad=0)
        self._acceptTheOnlyReview()
        self._phabUpdateWithExpectations(total=1, bad=1)

        # 'resolve' by abandoning our change
        self.dev_phab.dev_push_delete_branch("ph-review/emptyMerge2/master")
        self._phabUpdateWithExpectations(total=0, bad=0, emails=0)

    def test_mergeConflictWorkflow(self):
        self.dev_phab.dev_checkout_push_new_branch("temp/mergeConflict/master")
        self._devPushNewFile("NEWFILE", contents="hello")
        self._phabUpdateWithExpectations(total=0, bad=0)

        # move back to master and land a conflicting change
        self.dev_phab.dev_checkout("master")
        self.dev_phab.dev_checkout_push_new_branch(
            "ph-review/mergeConflict/master")
        self._devPushNewFile("NEWFILE", contents="goodbye")
        self._phabUpdateWithExpectations(total=1, bad=0)
        self._acceptTheOnlyReview()
        self._phabUpdateWithExpectations(total=0, bad=0)

        # move back to original and try to push and land
        self.dev_phab.dev_checkout_push_new_branch(
            "ph-review/mergeConflict2/master",
            base="temp/mergeConflict/master")
        self._phabUpdateWithExpectations(total=1, bad=0)
        self._acceptTheOnlyReview()
        self._phabUpdateWithExpectations(total=1, bad=1)

        # 'resolve' by forcing our change through
        print "force our change"
        self.dev_phab.dev_fetch()
        self.dev_phab.dev_merge_keep_ours("origin/master")
        self.dev_phab.dev_push_branch("ph-review/mergeConflict2/master")

        print "update again"
        self._phabUpdateWithExpectations(total=1, bad=0)
        self._acceptTheOnlyReview()
        self._phabUpdateWithExpectations(total=0, bad=0, emails=0)

    def test_changeAlreadyMergedOnBase(self):
        self.dev_phab.dev_checkout_push_new_branch("landing_branch")
        self._devPushNewFile("NEWFILE")
        self.dev_phab.dev_checkout_push_new_branch(
            "ph-review/alreadyMerged/landing_branch")
        self._phabUpdateWithExpectations(total=1, bad=1)

        # reset the landing branch back to master to resolve
        self.dev_phab.dev_reset_branch_to_master("landing_branch")

        self._phabUpdateWithExpectations(total=1, bad=0)
        self._acceptTheOnlyReview()
        self._phabUpdateWithExpectations(total=0, bad=0, emails=0)

    def test_commandeeredUpdate(self):
        self.dev_phab.dev_checkout_push_new_branch(
            "ph-review/commandeeredUpdate/master")
        self._devPushNewFile("NEWFILE")
        self._phabUpdateWithExpectations(total=1, bad=0)

        reviewid = self._getTheOnlyReviewId()
        self.conduit.commandeer_revision_as_user(
            reviewid, phldef_conduit.ALICE.user)

        self._devPushNewFile("NEWFILE2")
        self._phabUpdateWithExpectations(total=1, bad=0)

    def test_commandeeredLand(self):
        self.dev_phab.dev_checkout_push_new_branch(
            "ph-review/commandeeredLand/master")
        self._devPushNewFile("NEWFILE")
        self._phabUpdateWithExpectations(total=1, bad=0)
        self._devPushNewFile("NEWFILE2")
        self._phabUpdateWithExpectations(total=1, bad=0)

        reviewid = self._getTheOnlyReviewId()
        self.conduit.commandeer_revision_as_user(
            reviewid, phldef_conduit.PHAB.user)

        self._acceptTheOnlyReview()
        self._phabUpdateWithExpectations(total=0, bad=0, emails=0)

    def test_createHugeReview(self):
        self.dev_phab.dev_checkout_push_new_branch(
            "ph-review/createHugeReview/master")
        lots = "h\n" * 1 * 1024 * 1024
        self._devPushNewFile("NEWFILE", contents=lots)
        self._phabUpdateWithExpectations(total=1, bad=1)

    def test_hugeUpdateToReview(self):
        self.dev_phab.dev_checkout_push_new_branch(
            "ph-review/hugeUpdateReview/master")
        self._devPushNewFile("NEWFILE")
        self._phabUpdateWithExpectations(total=1, bad=0)
        lots = "h\n" * 1 * 1024 * 1024
        self._devPushNewFile("NEWFILE2", contents=lots)
        self._phabUpdateWithExpectations(total=1, bad=1)

    # TODO: test landing when origin has been updated underneath us
    # TODO: test landing when dependent review hasn't been landed

    def tearDown(self):
        self.dev_phab.close()

#------------------------------------------------------------------------------
# Copyright (C) 2012 Bloomberg L.P.
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to
# deal in the Software without restriction, including without limitation the
# rights to use, copy, modify, merge, publish, distribute, sublicense, and/or
# sell copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
# IN THE SOFTWARE.
#------------------------------- END-OF-FILE ----------------------------------
<commit_msg>Minor, re-order abdi_processrepo__t imports<commit_after>"""Test suite for abdi_processrepo."""
import unittest

import phldef_conduit
import phlmail_mocksender
#import phlsys_conduit

import abdi_processrepo
import abdmail_mailer
#import abdt_conduit
import abdt_conduitmock
import abdt_naming
import abdtst_devphabgit

# factors affecting a review:
#  age of the revisions
#  editing the review page
#  new revisions on the review branch
#  rewriting history on the review branch
#  author names
#  author accounts
#  base branch
#  availability of the git repo
#  availability of the phabricator instance

class Test(unittest.TestCase):

    def __init__(self, data):
        super(Test, self).__init__(data)
        self.reviewer = phldef_conduit.ALICE.user
        self.author_account = phldef_conduit.BOB
        self.conduit = None
        self.clone = None
        self.dev_phab = None
        self.mailer = None
        self.mock_sender = None

    def setUp(self):
        self.dev_phab = abdtst_devphabgit.Collaboration(
            self.author_account, phldef_conduit.PHAB)

        self.clone = self.dev_phab.clone

        self.dev_phab.dev_commit_new_file("README")
        self.dev_phab.dev_push_branch("master")
        self.dev_phab.phab_fetch()

        # sys_conduit = phlsys_conduit.Conduit(
        #     phldef_conduit.TEST_URI,
        #     phldef_conduit.PHAB.user,
        #     phldef_conduit.PHAB.certificate)

        #self.conduit = abdt_conduit.Conduit(sys_conduit)
        self.conduit = abdt_conduitmock.ConduitMock()

        self.mock_sender = phlmail_mocksender.MailSender()
        self.mailer = abdmail_mailer.Mailer(
            self.mock_sender,
            ["admin@server.test"],
            "http://server.fake/testrepo.git",
            "http://phabricator.server.fake/")

    def _phabUpdate(self):
        self.dev_phab.phab_fetch()
        abdi_processrepo.processUpdatedRepo(
            self.conduit, self.clone, "origin", self.mailer)

    def _phabUpdateWithExpectationsHelper(
            self, total=None, bad=None, emails=None):
        abdi_processrepo.processUpdatedRepo(
            self.conduit, self.clone, "origin", self.mailer)
        if total is not None:
            self.assertEqual(
                self.dev_phab.count_phab_working_branches(), total)
        if bad is not None:
            self.assertEqual(
                self.dev_phab.count_phab_bad_working_branches(), bad)
        if emails is not None:
            self.assertEqual(len(self.mock_sender.mails), emails)

    def _phabUpdateWithExpectations(self, total=None, bad=None, emails=None):
        self.dev_phab.phab_fetch()

        # multiple updates should have the same result if we are
        # not fetching and assuming the data in Phabricator
        # doesn't change.
        self._phabUpdateWithExpectationsHelper(total, bad, emails)
        self._phabUpdateWithExpectationsHelper(total, bad, emails)
        self._phabUpdateWithExpectationsHelper(total, bad, emails)

    def _devSetAuthorAccount(self, account):
        self.dev_phab.dev_clone.set_name_email(account.user, account.email)

    def _phabSetAuthorAccount(self, account):
        self.clone.set_name_email(account.user, account.email)

    def _devPushNewFile(
            self, filename, has_reviewer=True, has_plan=True, contents=""):
        reviewer = self.reviewer if has_reviewer else None
        plan = "testplan" if has_plan else None
        self.dev_phab.dev_commit_new_file(filename, plan, reviewer, contents)
        self.dev_phab.dev_push()

    def _getTheOnlyReviewId(self):
        branches = self.clone.get_remote_branches("origin")
        wbList = abdt_naming.getWorkingBranches(branches)
        self.assertEqual(len(wbList), 1)
        wb = wbList[0]
        return wb.id

    def _acceptTheOnlyReview(self):
        reviewid = self._getTheOnlyReviewId()
        self.conduit.accept_revision_as_user(reviewid, self.reviewer)

    def _abandonTheOnlyReview(self):
        reviewid = self._getTheOnlyReviewId()
        self.conduit.abandon_revision(reviewid)

    def test_nothingToDo(self):
        # nothing to process
        abdi_processrepo.processUpdatedRepo(
            self.conduit, self.clone, "origin", self.mailer)

    def test_simpleWorkflow(self):
        self.dev_phab.dev_checkout_push_new_branch(
            "ph-review/simpleWorkflow/master")
        self._devPushNewFile("NEWFILE")
        self._phabUpdateWithExpectations(total=1, bad=0)
        self._devPushNewFile("NEWFILE2")
        self._phabUpdateWithExpectations(total=1, bad=0)
        self._acceptTheOnlyReview()
        self._phabUpdateWithExpectations(total=0, bad=0, emails=0)

    def test_badMsgWorkflow(self):
        self.dev_phab.dev_checkout_push_new_branch(
            "ph-review/badMsgWorkflow/master")
        self._devPushNewFile("NEWFILE", has_plan=False)
        self._phabUpdateWithExpectations(total=1, bad=0)
        self._acceptTheOnlyReview()
        self._phabUpdateWithExpectations(total=0, bad=0, emails=0)

    def test_noReviewerWorkflow(self):
        self.dev_phab.dev_checkout_push_new_branch(
            "ph-review/noReviewerWorkflow/master")
        self._devPushNewFile("NEWFILE", has_reviewer=False)
        self._phabUpdateWithExpectations(total=1, bad=0)
        self._acceptTheOnlyReview()
        self._phabUpdateWithExpectations(total=0, bad=0, emails=0)

    def test_badBaseWorkflow(self):
        self.dev_phab.dev_checkout_push_new_branch(
            "ph-review/badBaseWorkflow/blaster")
        self._devPushNewFile("NEWFILE", has_plan=False)
        self._phabUpdateWithExpectations(total=1, bad=1)

        # delete the bad branch
        self.dev_phab.dev_push_delete_branch(
            "ph-review/badBaseWorkflow/blaster")

        self._phabUpdateWithExpectations(total=0, bad=0, emails=0)

    def test_withReservedBranch(self):
        self.dev_phab.dev_checkout_push_new_branch("dev/phab/reserve")
        self.dev_phab.dev_checkout_push_new_branch("ph-review/reserve/master")
        self._devPushNewFile("reserve")
        self._phabUpdateWithExpectations(total=1, bad=0)
        self._acceptTheOnlyReview()
        self._phabUpdateWithExpectations(total=0, bad=0, emails=0)

    def test_noBaseWorkflow(self):
        self.dev_phab.dev_checkout_push_new_branch("ph-review/noBaseWorkflow")
        self._devPushNewFile("NEWFILE", has_plan=False)

        # TODO: handle no base properly
        # self._phabUpdateWithExpectations(total=1, bad=1)

        # delete the bad branch
        self.dev_phab.dev_push_delete_branch("ph-review/noBaseWorkflow")

        self._phabUpdateWithExpectations(total=0, bad=0, emails=0)

    # TODO: test_notBasedWorkflow
    # TODO: test_noCommitWorkflow

    def test_badAuthorWorkflow(self):
        self._devSetAuthorAccount(phldef_conduit.NOTAUSER)
        self.dev_phab.dev_checkout_push_new_branch(
            "ph-review/badAuthorWorkflow/master")
        self._devPushNewFile("NEWFILE")
        self._phabUpdateWithExpectations(total=1, bad=1, emails=1)
        self._devPushNewFile("NEWFILE2")
        self._phabUpdateWithExpectations(total=1, bad=1, emails=2)
        self.dev_phab.dev_reset_branch_to_master(
            "ph-review/badAuthorWorkflow/master")
        self._devSetAuthorAccount(self.author_account)
        self._devPushNewFile("NEWFILE")
        self._phabUpdateWithExpectations(total=1, bad=0, emails=2)
        self._acceptTheOnlyReview()
        self._phabUpdateWithExpectations(total=0, bad=0, emails=2)

    def test_abandonedWorkflow(self):
        self.dev_phab.dev_checkout_push_new_branch(
            "ph-review/abandonedWorkflow/master")
        self._devPushNewFile("NEWFILE")
        self._phabUpdateWithExpectations(total=1, bad=0)
        self._abandonTheOnlyReview()
        self._phabUpdateWithExpectations(total=1, bad=0)
        self._devPushNewFile("NEWFILE2")
        self._phabUpdateWithExpectations(total=1, bad=0)
        self._acceptTheOnlyReview()
        self._phabUpdateWithExpectations(total=0, bad=0, emails=0)

    def test_emptyMergeWorkflow(self):
        self.dev_phab.dev_checkout_push_new_branch("temp/emptyMerge/master")
        self._devPushNewFile("NEWFILE")
        self._phabUpdateWithExpectations(total=0, bad=0)

        # move back to master and land a conflicting change
        self.dev_phab.dev_checkout("master")
        self.dev_phab.dev_checkout_push_new_branch(
            "ph-review/emptyMerge/master")
        self._devPushNewFile("NEWFILE")
        self._phabUpdateWithExpectations(total=1, bad=0)
        self._acceptTheOnlyReview()
        self._phabUpdateWithExpectations(total=0, bad=0)

        # move back to original and try to push and land
        self.dev_phab.dev_checkout("temp/emptyMerge/master")
        self.dev_phab.dev_checkout_push_new_branch(
            "ph-review/emptyMerge2/master")
        self._phabUpdateWithExpectations(total=1, bad=0)
        self._acceptTheOnlyReview()
        self._phabUpdateWithExpectations(total=1, bad=1)

        # 'resolve' by abandoning our change
        self.dev_phab.dev_push_delete_branch("ph-review/emptyMerge2/master")
        self._phabUpdateWithExpectations(total=0, bad=0, emails=0)

    def test_mergeConflictWorkflow(self):
        self.dev_phab.dev_checkout_push_new_branch("temp/mergeConflict/master")
        self._devPushNewFile("NEWFILE", contents="hello")
        self._phabUpdateWithExpectations(total=0, bad=0)

        # move back to master and land a conflicting change
        self.dev_phab.dev_checkout("master")
        self.dev_phab.dev_checkout_push_new_branch(
            "ph-review/mergeConflict/master")
        self._devPushNewFile("NEWFILE", contents="goodbye")
        self._phabUpdateWithExpectations(total=1, bad=0)
        self._acceptTheOnlyReview()
        self._phabUpdateWithExpectations(total=0, bad=0)

        # move back to original and try to push and land
        self.dev_phab.dev_checkout_push_new_branch(
            "ph-review/mergeConflict2/master",
            base="temp/mergeConflict/master")
        self._phabUpdateWithExpectations(total=1, bad=0)
        self._acceptTheOnlyReview()
        self._phabUpdateWithExpectations(total=1, bad=1)

        # 'resolve' by forcing our change through
        print "force our change"
        self.dev_phab.dev_fetch()
        self.dev_phab.dev_merge_keep_ours("origin/master")
        self.dev_phab.dev_push_branch("ph-review/mergeConflict2/master")

        print "update again"
        self._phabUpdateWithExpectations(total=1, bad=0)
        self._acceptTheOnlyReview()
        self._phabUpdateWithExpectations(total=0, bad=0, emails=0)

    def test_changeAlreadyMergedOnBase(self):
        self.dev_phab.dev_checkout_push_new_branch("landing_branch")
        self._devPushNewFile("NEWFILE")
        self.dev_phab.dev_checkout_push_new_branch(
            "ph-review/alreadyMerged/landing_branch")
        self._phabUpdateWithExpectations(total=1, bad=1)

        # reset the landing branch back to master to resolve
        self.dev_phab.dev_reset_branch_to_master("landing_branch")

        self._phabUpdateWithExpectations(total=1, bad=0)
        self._acceptTheOnlyReview()
        self._phabUpdateWithExpectations(total=0, bad=0, emails=0)

    def test_commandeeredUpdate(self):
        self.dev_phab.dev_checkout_push_new_branch(
            "ph-review/commandeeredUpdate/master")
        self._devPushNewFile("NEWFILE")
        self._phabUpdateWithExpectations(total=1, bad=0)

        reviewid = self._getTheOnlyReviewId()
        self.conduit.commandeer_revision_as_user(
            reviewid, phldef_conduit.ALICE.user)

        self._devPushNewFile("NEWFILE2")
        self._phabUpdateWithExpectations(total=1, bad=0)

    def test_commandeeredLand(self):
        self.dev_phab.dev_checkout_push_new_branch(
            "ph-review/commandeeredLand/master")
        self._devPushNewFile("NEWFILE")
        self._phabUpdateWithExpectations(total=1, bad=0)
        self._devPushNewFile("NEWFILE2")
        self._phabUpdateWithExpectations(total=1, bad=0)

        reviewid = self._getTheOnlyReviewId()
        self.conduit.commandeer_revision_as_user(
            reviewid, phldef_conduit.PHAB.user)

        self._acceptTheOnlyReview()
        self._phabUpdateWithExpectations(total=0, bad=0, emails=0)

    def test_createHugeReview(self):
        self.dev_phab.dev_checkout_push_new_branch(
            "ph-review/createHugeReview/master")
        lots = "h\n" * 1 * 1024 * 1024
        self._devPushNewFile("NEWFILE", contents=lots)
        self._phabUpdateWithExpectations(total=1, bad=1)

    def test_hugeUpdateToReview(self):
        self.dev_phab.dev_checkout_push_new_branch(
            "ph-review/hugeUpdateReview/master")
        self._devPushNewFile("NEWFILE")
        self._phabUpdateWithExpectations(total=1, bad=0)
        lots = "h\n" * 1 * 1024 * 1024
        self._devPushNewFile("NEWFILE2", contents=lots)
        self._phabUpdateWithExpectations(total=1, bad=1)

    # TODO: test landing when origin has been updated underneath us
    # TODO: test landing when dependent review hasn't been landed

    def tearDown(self):
        self.dev_phab.close()

#------------------------------------------------------------------------------
# Copyright (C) 2012 Bloomberg L.P.
#
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to
# deal in the Software without restriction, including without limitation the
# rights to use, copy, modify, merge, publish, distribute, sublicense, and/or
# sell copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
#
# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.  IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS
# IN THE SOFTWARE.
#------------------------------- END-OF-FILE ----------------------------------
<|endoftext|>


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

# INPUT:
#   volume_src
#   volume_dest
#   paramreg (class inherited in sct_register_multimodal)
#   (mask)
#   fname_warp
# OUTPUT:
#   none
#   (write warping field)

import os, sys, commands

# Get path of the toolbox
status, path_sct = commands.getstatusoutput('echo $SCT_DIR')
# Append path that contains scripts, to be able to load modules
sys.path.append(path_sct + '/scripts')

import sct_utils as sct
import nibabel
from sct_orientation import get_orientation, set_orientation
from scipy import ndimage
from scipy.io import loadmat
from numpy import array, asarray, zeros, int8, mean, std, sqrt, convolve, hanning
from copy import copy
from msct_image import Image
from numpy.linalg import inv
import matplotlib.pyplot as plt
from scipy.interpolate import splrep,splev
from sct_register_multimodal import Paramreg
import time
import csv
from math import acos

# FUNCTION
# register_seg
# First step to register a segmentation onto another by calculating precisely the displacement.
# (The segmentations can be of different size but the output segmentation must be smaller thant the input segmentation)?????? necessary or done before????

def register_seg(seg_input, seg_dest):
    seg_input_img = nibabel.load(seg_input)
    seg_dest_img = nibabel.load(seg_dest)
    seg_input_data = seg_input_img.get_data()
    seg_dest_data = seg_dest_img.get_data()

    x_center_of_mass_input = [0 for i in range(seg_input_data.shape[2])]
    y_center_of_mass_input = [0 for i in range(seg_input_data.shape[2])]
    print '\nGet center of mass of the input segmentation for each slice (corresponding to a slice int the output segmentation)...' #different if size of the two seg are different
    #TO DO: select only the slices corresponding to the output segmentation
    for iz in xrange(seg_dest_data.shape[2]):
        x_center_of_mass_input[iz], y_center_of_mass_input[iz] = ndimage.measurements.center_of_mass(array(seg_input_data[:,:,iz]))

    x_center_of_mass_output = [0 for i in range(seg_input_data.shape[2])]
    y_center_of_mass_output = [0 for i in range(seg_input_data.shape[2])]
    print '\nGet center of mass of the output segmentation for each slice ...'
    for iz in xrange(seg_dest_data.shape[2]):
        x_center_of_mass_output[iz], y_center_of_mass_output[iz] = ndimage.measurements.center_of_mass(array(seg_dest_data[:,:,iz]))

    x_displacement = [0 for i in range(seg_input_data.shape[2])]
    y_displacement = [0 for i in range(seg_input_data.shape[2])]
    print '\nGet displacement by voxel...'
    for iz in xrange(seg_dest_data.shape[2]):
        x_displacement[iz] = -(x_center_of_mass_output[iz] - x_center_of_mass_input[iz])    #strangely, this is the inverse of x_displacement when the same equation defines y_displacement
        y_displacement[iz] = y_center_of_mass_output[iz] - y_center_of_mass_input[iz]

    return x_displacement, y_displacement

# FUNCTION
# register_image
# First step to register a 3D image onto another by splitting the images into 2D images and calculating the displacement by splice by image correlation algorithms.
# (The images can be of different size but the output image must be smaller thant the input image)?????? necessary or done before????
# If the mask is inputed, it must also be 3D.

def register_images(im_input, im_dest, mask='', paramreg=Paramreg(step='0', type='im', algo='Translation', metric='MI', iter='5', shrink='1', smooth='0', gradStep='0.5'), remove_tmp_folder = 1):

    path_i, root_i, ext_i = sct.extract_fname(im_input)
    path_d, root_d, ext_d = sct.extract_fname(im_dest)
    path_m, root_m, ext_m = sct.extract_fname(mask)

    # set metricSize
    if paramreg.metric == 'MI':
        metricSize = '32'  # corresponds to number of bins
    else:
        metricSize = '4'  # corresponds to radius (for CC, MeanSquares...)

    # initiate default parameters of antsRegistration transformation
    ants_registration_params = {'rigid': '', 'affine': '', 'compositeaffine': '', 'similarity': '', 'translation': '',
                                'bspline': ',10', 'gaussiandisplacementfield': ',3,0',
                                'bsplinedisplacementfield': ',5,10', 'syn': ',3,0', 'bsplinesyn': ',3,32'}

    # Get image dimensions and retrieve nz
    print '\nGet image dimensions of destination image...'
    nx, ny, nz, nt, px, py, pz, pt = sct.get_dimension(im_dest)
    print '.. matrix size: '+str(nx)+' x '+str(ny)+' x '+str(nz)
    print '.. voxel size:  '+str(px)+'mm x '+str(py)+'mm x '+str(pz)+'mm'

    # Define x and y displacement as list
    x_displacement = [0 for i in range(nz)]
    y_displacement = [0 for i in range(nz)]
    theta_rotation = [0 for i in range(nz)]

    # create temporary folder
    print('\nCreate temporary folder...')
    path_tmp = 'tmp.'+time.strftime("%y%m%d%H%M%S")
    sct.create_folder(path_tmp)
    print '\nCopy input data...'
    sct.run('cp '+im_input+ ' ' + path_tmp +'/'+ root_i+ext_i)
    sct.run('cp '+im_dest+ ' ' + path_tmp +'/'+ root_d+ext_d)
    if mask:
        sct.run('cp '+mask+ ' '+path_tmp +'/mask.nii.gz')

    # go to temporary folder
    os.chdir(path_tmp)

    # Split input volume along z
    print '\nSplit input volume...'
    sct.run(sct.fsloutput + 'fslsplit '+im_input+' '+root_i+'_z -z')
    #file_anat_split = ['tmp.anat_orient_z'+str(z).zfill(4) for z in range(0,nz,1)]

    # Split destination volume along z
    print '\nSplit destination volume...'
    sct.run(sct.fsloutput + 'fslsplit '+im_dest+' '+root_d+'_z -z')
    #file_anat_split = ['tmp.anat_orient_z'+str(z).zfill(4) for z in range(0,nz,1)]

    # Split mask volume along z
    if mask:
        print '\nSplit mask volume...'
        sct.run(sct.fsloutput + 'fslsplit mask.nii.gz mask_z -z')
        #file_anat_split = ['tmp.anat_orient_z'+str(z).zfill(4) for z in range(0,nz,1)]

    # loop across slices
    for i in range(nz):
        # set masking
        num = numerotation(i)
        if mask:
            masking = '-x mask_z' +num+ '.nii'
        else:
            masking = ''

        cmd = ('isct_antsRegistration '
               '--dimensionality 2 '
               '--transform '+paramreg.algo+'['+paramreg.gradStep +
               ants_registration_params[paramreg.algo.lower()]+'] '
               '--metric '+paramreg.metric+'['+root_d+'_z'+ num +'.nii' +','+root_i+'_z'+ num +'.nii' +',1,'+metricSize+'] '  #[fixedImage,movingImage,metricWeight +nb_of_bins (MI) or radius (other)
               '--convergence '+paramreg.iter+' '
               '--shrink-factors '+paramreg.shrink+' '
               '--smoothing-sigmas '+paramreg.smooth+'mm '
               #'--restrict-deformation 1x1x0 '    # how to restrict? should not restrict here, if transform is precised...?
               '--output [transform_' + num + '] '    #--> file.txt (contains Tx,Ty)    [outputTransformPrefix,<outputWarpedImage>,<outputInverseWarpedImage>]
               '--interpolation BSpline[3] '
               +masking)

        # sct.run(cmd)

        try:
            sct.run(cmd)

            if paramreg.algo == 'Rigid' or paramreg.algo == 'Translation':
                f = 'transform_' +num+ '0GenericAffine.mat'
                matfile = loadmat(f, struct_as_record=True)
                array_transfo = matfile['AffineTransform_double_2_2']
                x_displacement[i] = -array_transfo[4][0]  #is it? or is it y?
                y_displacement[i] = array_transfo[5][0]
                theta_rotation[i] = acos(array_transfo[0])

        except:
                if paramreg.algo == 'Rigid' or paramreg.algo == 'Translation':
                    x_displacement[i] = x_displacement[i-1]  #is it? or is it y?
                    y_displacement[i] = y_displacement[i-1]
                    theta_rotation[i] = theta_rotation[i-1]

        # # get displacement form this slice and complete x and y displacement lists
        # with open('transform_'+num+'.csv') as f:
        #     reader = csv.reader(f)
        #     count = 0
        #     for line in reader:
        #         count += 1
        #         if count == 2:
        #             x_displacement[i] = line[0]
        #             y_displacement[i] = line[1]
        #             f.close()

        # # get matrix of transfo for a rigid transform   (pb slicereg fait une rotation ie le deplacement n'est pas homogene par slice)
        # # recuperer le deplacement ne donnerait pas une liste mais un warping field: mieux vaut recup la matrice output
        # # pb du smoothing du deplacement par slice !!   on peut smoother les param theta tx ty
        # if paramreg.algo == 'Rigid' or paramreg.algo == 'Translation':
        #     f = 'transform_' +num+ '0GenericAffine.mat'
        #     matfile = loadmat(f, struct_as_record=True)
        #     array_transfo = matfile['AffineTransform_double_2_2']
        #     x_displacement[i] = -array_transfo[4][0]  #is it? or is it y?
        #     y_displacement[i] = array_transfo[5][0]
        #     theta_rotation[i] = acos(array_transfo[0])

        #TO DO: different treatment for other algo

    #Delete tmp folder
    os.chdir('../')
    if remove_tmp_folder:
        print('\nRemove temporary files...')
        sct.run('rm -rf '+path_tmp)
    if paramreg.algo == 'Rigid':
        return x_displacement, y_displacement, theta_rotation   # check if the displacement are not inverted (x_dis = -x_disp...)   theta is in radian
    if paramreg.algo == 'Translation':
        return x_displacement, y_displacement

def numerotation(nb):
    if nb < 0:
        print 'ERROR: the number is negative.'
        sys.exit(status = 2)
    elif -1 < nb < 10:
        nb_output = '000'+str(nb)
    elif 9 < nb < 100:
        nb_output = '00'+str(nb)
    elif 99 < nb < 1000:
        nb_output = '0'+str(nb)
    elif 999 < nb < 10000:
        nb_output = str(nb)
    elif nb > 9999:
        print 'ERROR: the number is superior to 9999.'
        sys.exit(status = 2)
    return nb_output

# FUNCTION REGULARIZE
# input: x_displacement,y_displacement
# output: x_displacement_reg,y_displacement_reg
# use msct_smooth smoothing window

# FUNCTION: generate warping field that is homogeneous by slice   (for rigid transform only for now)
#input:
#   image_destination to recover the header
#   x_trans
#   y_trans
#   theta_rot
#   fname_warp

def generate_warping_field(im_dest, x_trans, y_trans, theta_rot=None, fname = 'warping_field.nii.gz'):
    from nibabel import load
    from math import cos, sin

    #Make sure image is in rpi format

    file_dest = load(im_dest)
    hdr_file_dest = file_dest.get_header()
    hdr_warp = hdr_file_dest.copy()

    # Get image dimensions
    print '\nGet image dimensions of destination image...'
    nx, ny, nz, nt, px, py, pz, pt = sct.get_dimension(im_dest)
    print '.. matrix size: '+str(nx)+' x '+str(ny)+' x '+str(nz)
    print '.. voxel size:  '+str(px)+'mm x '+str(py)+'mm x '+str(pz)+'mm'

    # Calculate displacement for each voxel
    data_warp = zeros(((((nx, ny, nz, 1, 3)))))
    if theta_rot != None:
        for k in range(nz):
            for i in range(nx):
                for j in range(ny):
                    # matrix_rot = [[cos(theta_rot[k]), -sin(theta_rot[k])], [sin(theta_rot[k], cos(theta_rot[k]]]
                    # data_warp[i,j,k,0,:] = matrix_rot
                    data_warp[i, j, k, 0, 0] = (cos(theta_rot[k])-1) * i - sin(theta_rot[k]) * j + x_trans[k]
                    data_warp[i, j, k, 0, 1] = sin(theta_rot[k]) * i + (cos(theta_rot[k])-1) * j + y_trans[k]
                    data_warp[i, j, k, 0, 2] = 0
    if theta_rot == None:
        for i in range(nx):
            for j in range(ny):
                for k in range(nz):
                    data_warp[i, j, k, 0, 0] = x_trans[k]
                    data_warp[i, j, k, 0, 1] = y_trans[k]
                    data_warp[i, j, k, 0, 2] = 0

    hdr_warp.set_intent('vector', (), '')
    hdr_warp.set_data_dtype('float32')
    img = nibabel.Nifti1Image(data_warp, None, hdr_warp)
    nibabel.save(img, fname)<commit_msg>OPT: now allow different size of input and output<commit_after>#!/usr/bin/env python

# INPUT:
#   volume_src
#   volume_dest
#   paramreg (class inherited in sct_register_multimodal)
#   (mask)
#   fname_warp
# OUTPUT:
#   none
#   (write warping field)

import os, sys, commands

# Get path of the toolbox
status, path_sct = commands.getstatusoutput('echo $SCT_DIR')
# Append path that contains scripts, to be able to load modules
sys.path.append(path_sct + '/scripts')

import sct_utils as sct
import nibabel
from sct_orientation import get_orientation, set_orientation
from scipy import ndimage
from scipy.io import loadmat
from numpy import array, asarray, zeros, int8, mean, std, sqrt, convolve, hanning
from copy import copy
from msct_image import Image
from numpy.linalg import inv
import matplotlib.pyplot as plt
from scipy.interpolate import splrep,splev
from sct_register_multimodal import Paramreg
import time
import csv
from math import acos

# FUNCTION
# register_seg
# First step to register a segmentation onto another by calculating precisely the displacement.
# (The segmentations can be of different size but the output segmentation must be smaller thant the input segmentation)?????? necessary or done before????

def register_seg(seg_input, seg_dest):
    seg_input_img = Image(seg_input)
    seg_dest_img = Image(seg_dest)
    seg_input_data = seg_input_img.data
    seg_dest_data = seg_dest_img.data

    x_center_of_mass_input = [0 for i in range(seg_dest_data.shape[2])]
    y_center_of_mass_input = [0 for i in range(seg_dest_data.shape[2])]
    print '\nGet center of mass of the input segmentation for each slice (corresponding to a slice in the output segmentation)...' #different if size of the two seg are different
    #TO DO: select only the slices corresponding to the output segmentation
    coord_origin_dest = seg_dest_img.transfo_pix2phys([[0,0,0]])
    [[x_o, y_o, z_o]] = seg_input_img.transfo_phys2pix(coord_origin_dest)
    for iz in xrange(seg_dest_data.shape[2]):
        x_center_of_mass_input[iz], y_center_of_mass_input[iz] = ndimage.measurements.center_of_mass(array(seg_input_data[:, :, z_o + iz]))

    x_center_of_mass_output = [0 for i in range(seg_dest_data.shape[2])]
    y_center_of_mass_output = [0 for i in range(seg_dest_data.shape[2])]
    print '\nGet center of mass of the output segmentation for each slice ...'
    for iz in xrange(seg_dest_data.shape[2]):
        x_center_of_mass_output[iz], y_center_of_mass_output[iz] = ndimage.measurements.center_of_mass(array(seg_dest_data[:,:,iz]))

    x_displacement = [0 for i in range(seg_input_data.shape[2])]
    y_displacement = [0 for i in range(seg_input_data.shape[2])]
    print '\nGet displacement by voxel...'
    for iz in xrange(seg_dest_data.shape[2]):
        x_displacement[iz] = -(x_center_of_mass_output[iz] - x_center_of_mass_input[iz])    #strangely, this is the inverse of x_displacement when the same equation defines y_displacement
        y_displacement[iz] = y_center_of_mass_output[iz] - y_center_of_mass_input[iz]

    return x_displacement, y_displacement

# FUNCTION
# register_image
# First step to register a 3D image onto another by splitting the images into 2D images and calculating the displacement by splice by image correlation algorithms.
# (The images can be of different size but the output image must be smaller thant the input image)?????? necessary or done before????
# If the mask is inputed, it must also be 3D and it must be in the same space as the destination image.

def register_images(im_input, im_dest, mask='', paramreg=Paramreg(step='0', type='im', algo='Translation', metric='MI', iter='5', shrink='1', smooth='0', gradStep='0.5'), remove_tmp_folder = 1):

    path_i, root_i, ext_i = sct.extract_fname(im_input)
    path_d, root_d, ext_d = sct.extract_fname(im_dest)
    path_m, root_m, ext_m = sct.extract_fname(mask)

    # set metricSize
    if paramreg.metric == 'MI':
        metricSize = '32'  # corresponds to number of bins
    else:
        metricSize = '4'  # corresponds to radius (for CC, MeanSquares...)

    # initiate default parameters of antsRegistration transformation
    ants_registration_params = {'rigid': '', 'affine': '', 'compositeaffine': '', 'similarity': '', 'translation': '',
                                'bspline': ',10', 'gaussiandisplacementfield': ',3,0',
                                'bsplinedisplacementfield': ',5,10', 'syn': ',3,0', 'bsplinesyn': ',3,32'}

    # Get image dimensions and retrieve nz
    print '\nGet image dimensions of destination image...'
    nx, ny, nz, nt, px, py, pz, pt = sct.get_dimension(im_dest)
    print '.. matrix size: '+str(nx)+' x '+str(ny)+' x '+str(nz)
    print '.. voxel size:  '+str(px)+'mm x '+str(py)+'mm x '+str(pz)+'mm'

    # Define x and y displacement as list
    x_displacement = [0 for i in range(nz)]
    y_displacement = [0 for i in range(nz)]
    theta_rotation = [0 for i in range(nz)]

    # create temporary folder
    print('\nCreate temporary folder...')
    path_tmp = 'tmp.'+time.strftime("%y%m%d%H%M%S")
    sct.create_folder(path_tmp)
    print '\nCopy input data...'
    sct.run('cp '+im_input+ ' ' + path_tmp +'/'+ root_i+ext_i)
    sct.run('cp '+im_dest+ ' ' + path_tmp +'/'+ root_d+ext_d)
    if mask:
        sct.run('cp '+mask+ ' '+path_tmp +'/mask.nii.gz')

    # go to temporary folder
    os.chdir(path_tmp)

    # Split input volume along z
    print '\nSplit input volume...'
    sct.run(sct.fsloutput + 'fslsplit '+im_input+' '+root_i+'_z -z')
    #file_anat_split = ['tmp.anat_orient_z'+str(z).zfill(4) for z in range(0,nz,1)]

    # Split destination volume along z
    print '\nSplit destination volume...'
    sct.run(sct.fsloutput + 'fslsplit '+im_dest+' '+root_d+'_z -z')
    #file_anat_split = ['tmp.anat_orient_z'+str(z).zfill(4) for z in range(0,nz,1)]

    # Split mask volume along z
    if mask:
        print '\nSplit mask volume...'
        sct.run(sct.fsloutput + 'fslsplit mask.nii.gz mask_z -z')
        #file_anat_split = ['tmp.anat_orient_z'+str(z).zfill(4) for z in range(0,nz,1)]

    im_dest_img = Image(im_dest)
    im_input_img = Image(im_input)
    coord_origin_dest = im_dest_img.transfo_pix2phys([[0,0,0]])
    [[x_o, y_o, z_o]] = im_input_img.transfo_phys2pix(coord_origin_dest)

    # loop across slices
    for i in range(nz):
        # set masking
        num = numerotation(i)
        num_2 = numerotation(int(num) + z_o)
        if mask:
            masking = '-x mask_z' +num+ '.nii'
        else:
            masking = ''

        cmd = ('isct_antsRegistration '
               '--dimensionality 2 '
               '--transform '+paramreg.algo+'['+paramreg.gradStep +
               ants_registration_params[paramreg.algo.lower()]+'] '
               '--metric '+paramreg.metric+'['+root_d+'_z'+ num +'.nii' +','+root_i+'_z'+ num_2 +'.nii' +',1,'+metricSize+'] '  #[fixedImage,movingImage,metricWeight +nb_of_bins (MI) or radius (other)
               '--convergence '+paramreg.iter+' '
               '--shrink-factors '+paramreg.shrink+' '
               '--smoothing-sigmas '+paramreg.smooth+'mm '
               #'--restrict-deformation 1x1x0 '    # how to restrict? should not restrict here, if transform is precised...?
               '--output [transform_' + num + '] '    #--> file.txt (contains Tx,Ty)    [outputTransformPrefix,<outputWarpedImage>,<outputInverseWarpedImage>]
               '--interpolation BSpline[3] '
               +masking)

        # sct.run(cmd)

        try:
            sct.run(cmd)

            if paramreg.algo == 'Rigid' or paramreg.algo == 'Translation':
                f = 'transform_' +num+ '0GenericAffine.mat'
                matfile = loadmat(f, struct_as_record=True)
                array_transfo = matfile['AffineTransform_double_2_2']
                x_displacement[i] = -array_transfo[4][0]  #is it? or is it y?
                y_displacement[i] = array_transfo[5][0]
                theta_rotation[i] = acos(array_transfo[0])

        except:
                if paramreg.algo == 'Rigid' or paramreg.algo == 'Translation':
                    x_displacement[i] = x_displacement[i-1]  #is it? or is it y?
                    y_displacement[i] = y_displacement[i-1]
                    theta_rotation[i] = theta_rotation[i-1]

        # # get displacement form this slice and complete x and y displacement lists
        # with open('transform_'+num+'.csv') as f:
        #     reader = csv.reader(f)
        #     count = 0
        #     for line in reader:
        #         count += 1
        #         if count == 2:
        #             x_displacement[i] = line[0]
        #             y_displacement[i] = line[1]
        #             f.close()

        # # get matrix of transfo for a rigid transform   (pb slicereg fait une rotation ie le deplacement n'est pas homogene par slice)
        # # recuperer le deplacement ne donnerait pas une liste mais un warping field: mieux vaut recup la matrice output
        # # pb du smoothing du deplacement par slice !!   on peut smoother les param theta tx ty
        # if paramreg.algo == 'Rigid' or paramreg.algo == 'Translation':
        #     f = 'transform_' +num+ '0GenericAffine.mat'
        #     matfile = loadmat(f, struct_as_record=True)
        #     array_transfo = matfile['AffineTransform_double_2_2']
        #     x_displacement[i] = -array_transfo[4][0]  #is it? or is it y?
        #     y_displacement[i] = array_transfo[5][0]
        #     theta_rotation[i] = acos(array_transfo[0])

        #TO DO: different treatment for other algo

    #Delete tmp folder
    os.chdir('../')
    if remove_tmp_folder:
        print('\nRemove temporary files...')
        sct.run('rm -rf '+path_tmp)
    if paramreg.algo == 'Rigid':
        return x_displacement, y_displacement, theta_rotation   # check if the displacement are not inverted (x_dis = -x_disp...)   theta is in radian
    if paramreg.algo == 'Translation':
        return x_displacement, y_displacement

def numerotation(nb):
    if nb < 0:
        print 'ERROR: the number is negative.'
        sys.exit(status = 2)
    elif -1 < nb < 10:
        nb_output = '000'+str(nb)
    elif 9 < nb < 100:
        nb_output = '00'+str(nb)
    elif 99 < nb < 1000:
        nb_output = '0'+str(nb)
    elif 999 < nb < 10000:
        nb_output = str(nb)
    elif nb > 9999:
        print 'ERROR: the number is superior to 9999.'
        sys.exit(status = 2)
    return nb_output

# FUNCTION REGULARIZE
# input: x_displacement,y_displacement
# output: x_displacement_reg,y_displacement_reg
# use msct_smooth smoothing window

# FUNCTION: generate warping field that is homogeneous by slice   (for rigid transform only for now)
#input:
#   image_destination to recover the header
#   x_trans
#   y_trans
#   theta_rot
#   fname_warp

def generate_warping_field(im_dest, x_trans, y_trans, theta_rot=None, fname = 'warping_field.nii.gz'):
    from nibabel import load
    from math import cos, sin

    #Make sure image is in rpi format

    file_dest = load(im_dest)
    hdr_file_dest = file_dest.get_header()
    hdr_warp = hdr_file_dest.copy()

    # Get image dimensions
    print '\nGet image dimensions of destination image...'
    nx, ny, nz, nt, px, py, pz, pt = sct.get_dimension(im_dest)
    print '.. matrix size: '+str(nx)+' x '+str(ny)+' x '+str(nz)
    print '.. voxel size:  '+str(px)+'mm x '+str(py)+'mm x '+str(pz)+'mm'

    # Calculate displacement for each voxel
    data_warp = zeros(((((nx, ny, nz, 1, 3)))))
    if theta_rot != None:
        for k in range(nz):
            for i in range(nx):
                for j in range(ny):
                    # matrix_rot = [[cos(theta_rot[k]), -sin(theta_rot[k])], [sin(theta_rot[k], cos(theta_rot[k]]]
                    # data_warp[i,j,k,0,:] = matrix_rot
                    data_warp[i, j, k, 0, 0] = (cos(theta_rot[k])-1) * i - sin(theta_rot[k]) * j + x_trans[k]
                    data_warp[i, j, k, 0, 1] = sin(theta_rot[k]) * i + (cos(theta_rot[k])-1) * j + y_trans[k]
                    data_warp[i, j, k, 0, 2] = 0
    if theta_rot == None:
        for i in range(nx):
            for j in range(ny):
                for k in range(nz):
                    data_warp[i, j, k, 0, 0] = x_trans[k]
                    data_warp[i, j, k, 0, 1] = y_trans[k]
                    data_warp[i, j, k, 0, 2] = 0

    hdr_warp.set_intent('vector', (), '')
    hdr_warp.set_data_dtype('float32')
    img = nibabel.Nifti1Image(data_warp, None, hdr_warp)
    nibabel.save(img, fname)
    print '\nDONE ! Warping field generated: '+fname<|endoftext|>