<commit_before># -*- coding: utf-8 -*-
"""Simulation"""
from __future__ import division

import logging
import datetime
import numpy as np
import pandas as pd

import basal_area_increment as incr
from GYPSYNonSpatial import densities_speciescomp_topheight_to_250
from utils import _log_loop_progress, estimate_species_composition
from pygypsy.basal_area_factor import (
    estimate_basal_area_factor_aw,
    estimate_basal_area_factor_sb,
    estimate_basal_area_factor_sw,
    estimate_basal_area_factor_pl,
)
from pygypsy.basal_area_simulate import (
    sim_basal_area_aw,
    sim_basal_area_sw,
    sim_basal_area_sb,
    sim_basal_area_pl,
)
from pygypsy.volume import(
    merchantable_volume,
    gross_total_volume,
)

logger = logging.getLogger(__name__)

SPECIES = ('Aw', 'Sw', 'Sb', 'Pl')

DEFAULT_UTILIZATIONS = {
    "aw": {
        "topDiamInsideBark": 7,
        "stumpDiamOutsideBark": 13,
        "stumpHeight": 0.3
    },
    "sw": {
        "topDiamInsideBark": 7,
        "stumpDiamOutsideBark": 13,
        "stumpHeight": 0.3
    },
    "sb": {
        "topDiamInsideBark": 7,
        "stumpDiamOutsideBark": 13,
        "stumpHeight": 0.3
    },
    "pl": {
        "topDiamInsideBark": 7,
        "stumpDiamOutsideBark": 13,
        "stumpHeight": 0.3
    }
}

def get_initial_basal_area(current_basal_area):
    initial_basal_area = 0.001

    if initial_basal_area > current_basal_area * 0.5:
        initial_basal_area /= 2

    return initial_basal_area

def get_basal_area_factors_for_all_species(**kwargs):
    logger.debug('Getting factors for all species')

    f_Sb = 0
    f_Aw = 0
    f_Sw = 0
    f_Pl = 0

    if kwargs['N0_Aw'] > 0:
        f_Aw = estimate_basal_area_factor_aw(**kwargs)

    if kwargs['N0_Sb'] > 0:
        f_Sb = estimate_basal_area_factor_sb(**kwargs)

    if kwargs['N0_Sw'] > 0:
        f_Sw = estimate_basal_area_factor_sw(**kwargs)

    if kwargs['N0_Pl'] > 0:
        f_Pl = estimate_basal_area_factor_pl(**kwargs)

    return {'f_Aw':f_Aw,
            'f_Sb':f_Sb,
            'f_Sw':f_Sw,
            'f_Pl':f_Pl,}

def simulate_forwards_df(plot_df, utiliz_params=None):
    """Run forwards simulation

    Keyword Arguments:
    plot_df -- pandas.DataFrame with plot data

    Return:
    !TODO!
    """
    if utiliz_params is None:
        utiliz_params = DEFAULT_UTILIZATIONS

    output_dict = {}
    logger.debug('Starting forwards simulation')
    n_rows = plot_df.shape[0]

    for _, row in plot_df.iterrows():
        start = datetime.datetime.now()
        _log_loop_progress(_, n_rows)
        plot_id = str(int(row['id_l1']))

        logger.info('Starting simulation for plot: %s', plot_id)
        SI_bh_Aw = row.at['SI_Aw']
        SI_bh_Sw = row.at['SI_Sw']
        SI_bh_Pl = row.at['SI_Pl']
        SI_bh_Sb = row.at['SI_Sb']
        N_bh_AwT = row.at['N_Aw']
        N_bh_SwT = row.at['N_Sw']
        N_bh_PlT = row.at['N_Pl']
        N_bh_SbT = row.at['N_Sb']
        y2bh_Aw = row.at['y2bh_Aw']
        y2bh_Sw = row.at['y2bh_Sw']
        y2bh_Sb = row.at['y2bh_Sb']
        y2bh_Pl = row.at['y2bh_Pl']
        tage_AwT = row.at['tage_Aw']
        tage_SwT = row.at['tage_Sw']
        tage_PlT = row.at['tage_Pl']
        tage_SbT = row.at['tage_Sb']
        BA_AwT = row.at['BA_Aw']
        BA_SwT = row.at['BA_Sw']
        BA_PlT = row.at['BA_Pl']
        BA_SbT = row.at['BA_Sb']
        SDF_Aw0 = row.at['SDF_Aw']
        SDF_Sw0 = row.at['SDF_Sw']
        SDF_Pl0 = row.at['SDF_Pl']
        SDF_Sb0 = row.at['SDF_Sb']
        N0_Aw = row.at['N0_Aw']
        N0_Sw = row.at['N0_Sw']
        N0_Pl = row.at['N0_Pl']
        N0_Sb = row.at['N0_Sb']

        # TODO: don't we already have this in the data frame from data prep?
        BA_Aw0 = get_initial_basal_area(BA_AwT)
        BA_Sw0 = get_initial_basal_area(BA_SwT)
        BA_Sb0 = get_initial_basal_area(BA_SbT)
        BA_Pl0 = get_initial_basal_area(BA_PlT)

        # TODO: don't we already have this in the data frame from data prep?
        SC_Aw, SC_Sw, SC_Sb, SC_Pl = estimate_species_composition(N0_Aw, N0_Sb, N0_Sw, N0_Pl)

        tageData = [tage_AwT, tage_SwT, tage_PlT, tage_SbT]
        startTageAw = tageData[0]
        startTageSw = tageData[1]
        startTagePl = tageData[2]
        startTageSb = tageData[3]

        tageData = sorted(tageData, reverse=True)
        startTage = int(tageData[0])

        densities = densities_speciescomp_topheight_to_250(
            startTage=startTage,
            startTageAw=startTageAw,
            y2bh_Aw=y2bh_Aw,
            startTageSw=startTageSw,
            y2bh_Sw=y2bh_Sw,
            startTageSb=startTageSb,
            y2bh_Sb=y2bh_Sb,
            startTagePl=startTagePl,
            y2bh_Pl=y2bh_Pl,
            SDF_Aw0=SDF_Aw0,
            SDF_Sw0=SDF_Sw0,
            SDF_Pl0=SDF_Pl0,
            SDF_Sb0=SDF_Sb0,
            SI_bh_Aw=SI_bh_Aw,
            SI_bh_Sw=SI_bh_Sw,
            SI_bh_Sb=SI_bh_Sb,
            SI_bh_Pl=SI_bh_Pl
        )

        # estimating correction factor to fit BA at t0 and BA at t and
        # choosing whether simulating with multiplication factor
        # or starting at t recalculating the densities and SC
        species_factors = get_basal_area_factors_for_all_species(
            startTage=startTage,
            startTageAw=startTageAw,
            y2bh_Aw=y2bh_Aw,
            SC_Aw=SC_Aw,
            SI_bh_Aw=SI_bh_Aw,
            N_bh_AwT=N_bh_AwT,
            N0_Aw=N0_Aw,
            BA_Aw0=BA_Aw0,
            BA_AwT=BA_AwT,
            startTageSb=startTageSb,
            y2bh_Sb=y2bh_Sb,
            SC_Sb=SC_Sb,
            SI_bh_Sb=SI_bh_Sb,
            N_bh_SbT=N_bh_SbT,
            N0_Sb=N0_Sb,
            BA_Sb0=BA_Sb0,
            BA_SbT=BA_SbT,
            startTageSw=startTageSw,
            y2bh_Sw=y2bh_Sw,
            SC_Sw=SC_Sw,
            SI_bh_Sw=SI_bh_Sw,
            N_bh_SwT=N_bh_SwT,
            N0_Sw=N0_Sw,
            SDF_Pl0=SDF_Pl0,
            SDF_Sb0=SDF_Sb0,
            BA_Sw0=BA_Sw0,
            BA_SwT=BA_SwT,
            startTagePl=startTagePl,
            y2bh_Pl=y2bh_Pl,
            SC_Pl=SC_Pl,
            SI_bh_Pl=SI_bh_Pl,
            N_bh_PlT=N_bh_PlT,
            N0_Pl=N0_Pl,
            SDF_Aw0=SDF_Aw0,
            SDF_Sw0=SDF_Sw0,
            BA_Pl0=BA_Pl0,
            BA_PlT=BA_PlT,
            densities=densities,
            printWarnings=True
        )

        f_Aw = species_factors['f_Aw']
        f_Sw = species_factors['f_Sw']
        f_Sb = species_factors['f_Sb']
        f_Pl = species_factors['f_Pl']

        logger.debug('Getting basal area from time 0 to time of data')
        # use_correction_factor_future here is True because aspen was peculiar,
        # empirically demonstrated that using the factor for the whole
        # simulation yielded better results, probably because of variability in
        # density and basal area unique to aspen we can't do simulation choice
        # = no here for sb, sw, pl because the factor should not be applied for
        # them sw, sb, pl use the factor until the time of data. the subsequent
        # years use the regular basal area increment formula julianno sambatti,
        # november 10, 2016
        # further inspection of those functions indicates they can be improved
        # because the increment only dependds on other values, so we can
        # calculate it as a vector operation using the other arrays then
        # the actual values are just the cumulative sums of the
        # increments that's something for another day; because it is
        # complicated by the whole factor business
        BA_0_to_data_Aw_arr = sim_basal_area_aw(
            startTage, SI_bh_Aw,  N0_Aw, BA_Aw0, SDF_Aw0, f_Aw, densities,
            use_correction_factor_future=True, stop_at_initial_age=False
        )
        BA_0_to_data_Sb_arr = sim_basal_area_sb(
            startTage, SI_bh_Sb, N0_Sb,
            BA_Sb0, f_Sb, densities, stop_at_initial_age=False
        )
        BA_0_to_data_Sw_arr = sim_basal_area_sw(
            startTage, SI_bh_Sw, N0_Sw, SDF_Aw0, SDF_Pl0, SDF_Sb0, BA_Sw0,
            f_Sw, densities, stop_at_initial_age=False
        )
        BA_0_to_data_Pl_arr = sim_basal_area_pl(
            startTage, SI_bh_Pl, N0_Pl, SDF_Aw0, SDF_Sw0, SDF_Sb0, BA_Pl0,
            f_Pl, densities, stop_at_initial_age=False
        )

        output_DF_Aw = pd.DataFrame(BA_0_to_data_Aw_arr, columns=['BA_Aw'])
        output_DF_Sw = pd.DataFrame(BA_0_to_data_Sw_arr, columns=['BA_Sw'])
        output_DF_Sb = pd.DataFrame(BA_0_to_data_Sb_arr, columns=['BA_Sb'])
        output_DF_Pl = pd.DataFrame(BA_0_to_data_Pl_arr, columns=['BA_Pl'])

        densities_DF = pd.DataFrame(densities)
        output_DF = pd.concat(
            [densities_DF, output_DF_Aw, output_DF_Sw, output_DF_Sb, output_DF_Pl],
            axis=1
        )

        import ipdb; ipdb.set_trace()
        for spec in SPECIES:
            output_DF['Gross_Total_Volume_%s' % spec] = gross_total_volume(
                spec,
                output_DF['BA_%s' % spec],
                output_DF['topHeight_%s' % spec]
            )

        output_DF['Gross_Total_Volume_Con'] = output_DF['Gross_Total_Volume_Sw'] \
                                              + output_DF['Gross_Total_Volume_Sb'] \
                                              + output_DF['Gross_Total_Volume_Pl']
        output_DF['Gross_Total_Volume_Dec'] = output_DF['Gross_Total_Volume_Aw']
        output_DF['Gross_Total_Volume_Tot'] = output_DF['Gross_Total_Volume_Con'] \
                                              + output_DF['Gross_Total_Volume_Dec']

        # this could go in the loop above, but is left here for now since
        # the tests are sensitive to column order
        for spec in SPECIES:
            output_DF['MerchantableVolume%s' % spec] = merchantable_volume(
                spec,
                output_DF['N_bh_%sT' % spec],
                output_DF['BA_%s' % spec],
                output_DF['topHeight_%s' % spec],
                output_DF['Gross_Total_Volume_%s' % spec],
                top_dib=utiliz_params[spec.lower()]['topDiamInsideBark'],
                stump_dob=utiliz_params[spec.lower()]['stumpDiamOutsideBark'],
                stump_height=utiliz_params[spec.lower()]['stumpHeight']
            )

        output_DF['MerchantableVolume_Con'] = output_DF['MerchantableVolumeSw'] \
                                              + output_DF['MerchantableVolumeSb'] \
                                              + output_DF['MerchantableVolumePl']
        output_DF['MerchantableVolume_Dec'] = output_DF['MerchantableVolumeAw']
        output_DF['MerchantableVolume_Tot'] = output_DF['MerchantableVolume_Con'] \
                                              + output_DF['MerchantableVolume_Dec']

        output_dict[plot_id] = output_DF

        end = datetime.datetime.now()
        duration = end - start
        logger.info('plot %s took %f seconds', plot_id, duration.total_seconds())

    return output_dict
<commit_msg>Remove comments - confirmed with julianno<commit_after># -*- coding: utf-8 -*-
"""Simulation"""
from __future__ import division

import logging
import datetime
import numpy as np
import pandas as pd

import basal_area_increment as incr
from GYPSYNonSpatial import densities_speciescomp_topheight_to_250
from utils import _log_loop_progress, estimate_species_composition
from pygypsy.basal_area_factor import (
    estimate_basal_area_factor_aw,
    estimate_basal_area_factor_sb,
    estimate_basal_area_factor_sw,
    estimate_basal_area_factor_pl,
)
from pygypsy.basal_area_simulate import (
    sim_basal_area_aw,
    sim_basal_area_sw,
    sim_basal_area_sb,
    sim_basal_area_pl,
)
from pygypsy.volume import(
    merchantable_volume,
    gross_total_volume,
)

logger = logging.getLogger(__name__)

SPECIES = ('Aw', 'Sw', 'Sb', 'Pl')

DEFAULT_UTILIZATIONS = {
    "aw": {
        "topDiamInsideBark": 7,
        "stumpDiamOutsideBark": 13,
        "stumpHeight": 0.3
    },
    "sw": {
        "topDiamInsideBark": 7,
        "stumpDiamOutsideBark": 13,
        "stumpHeight": 0.3
    },
    "sb": {
        "topDiamInsideBark": 7,
        "stumpDiamOutsideBark": 13,
        "stumpHeight": 0.3
    },
    "pl": {
        "topDiamInsideBark": 7,
        "stumpDiamOutsideBark": 13,
        "stumpHeight": 0.3
    }
}

def get_initial_basal_area(current_basal_area):
    initial_basal_area = 0.001

    if initial_basal_area > current_basal_area * 0.5:
        initial_basal_area /= 2

    return initial_basal_area

def get_basal_area_factors_for_all_species(**kwargs):
    logger.debug('Getting factors for all species')

    f_Sb = 0
    f_Aw = 0
    f_Sw = 0
    f_Pl = 0

    if kwargs['N0_Aw'] > 0:
        f_Aw = estimate_basal_area_factor_aw(**kwargs)

    if kwargs['N0_Sb'] > 0:
        f_Sb = estimate_basal_area_factor_sb(**kwargs)

    if kwargs['N0_Sw'] > 0:
        f_Sw = estimate_basal_area_factor_sw(**kwargs)

    if kwargs['N0_Pl'] > 0:
        f_Pl = estimate_basal_area_factor_pl(**kwargs)

    return {'f_Aw':f_Aw,
            'f_Sb':f_Sb,
            'f_Sw':f_Sw,
            'f_Pl':f_Pl,}

def simulate_forwards_df(plot_df, utiliz_params=None):
    """Run forwards simulation

    Keyword Arguments:
    plot_df -- pandas.DataFrame with plot data

    Return:
    !TODO!
    """
    if utiliz_params is None:
        utiliz_params = DEFAULT_UTILIZATIONS

    output_dict = {}
    logger.debug('Starting forwards simulation')
    n_rows = plot_df.shape[0]

    for _, row in plot_df.iterrows():
        start = datetime.datetime.now()
        _log_loop_progress(_, n_rows)
        plot_id = str(int(row['id_l1']))

        logger.info('Starting simulation for plot: %s', plot_id)
        SI_bh_Aw = row.at['SI_Aw']
        SI_bh_Sw = row.at['SI_Sw']
        SI_bh_Pl = row.at['SI_Pl']
        SI_bh_Sb = row.at['SI_Sb']
        N_bh_AwT = row.at['N_Aw']
        N_bh_SwT = row.at['N_Sw']
        N_bh_PlT = row.at['N_Pl']
        N_bh_SbT = row.at['N_Sb']
        y2bh_Aw = row.at['y2bh_Aw']
        y2bh_Sw = row.at['y2bh_Sw']
        y2bh_Sb = row.at['y2bh_Sb']
        y2bh_Pl = row.at['y2bh_Pl']
        tage_AwT = row.at['tage_Aw']
        tage_SwT = row.at['tage_Sw']
        tage_PlT = row.at['tage_Pl']
        tage_SbT = row.at['tage_Sb']
        BA_AwT = row.at['BA_Aw']
        BA_SwT = row.at['BA_Sw']
        BA_PlT = row.at['BA_Pl']
        BA_SbT = row.at['BA_Sb']
        SDF_Aw0 = row.at['SDF_Aw']
        SDF_Sw0 = row.at['SDF_Sw']
        SDF_Pl0 = row.at['SDF_Pl']
        SDF_Sb0 = row.at['SDF_Sb']
        N0_Aw = row.at['N0_Aw']
        N0_Sw = row.at['N0_Sw']
        N0_Pl = row.at['N0_Pl']
        N0_Sb = row.at['N0_Sb']

        BA_Aw0 = get_initial_basal_area(BA_AwT)
        BA_Sw0 = get_initial_basal_area(BA_SwT)
        BA_Sb0 = get_initial_basal_area(BA_SbT)
        BA_Pl0 = get_initial_basal_area(BA_PlT)

        SC_Aw, SC_Sw, SC_Sb, SC_Pl = estimate_species_composition(N0_Aw, N0_Sb, N0_Sw, N0_Pl)

        tageData = [tage_AwT, tage_SwT, tage_PlT, tage_SbT]
        startTageAw = tageData[0]
        startTageSw = tageData[1]
        startTagePl = tageData[2]
        startTageSb = tageData[3]

        tageData = sorted(tageData, reverse=True)
        startTage = int(tageData[0])

        densities = densities_speciescomp_topheight_to_250(
            startTage=startTage,
            startTageAw=startTageAw,
            y2bh_Aw=y2bh_Aw,
            startTageSw=startTageSw,
            y2bh_Sw=y2bh_Sw,
            startTageSb=startTageSb,
            y2bh_Sb=y2bh_Sb,
            startTagePl=startTagePl,
            y2bh_Pl=y2bh_Pl,
            SDF_Aw0=SDF_Aw0,
            SDF_Sw0=SDF_Sw0,
            SDF_Pl0=SDF_Pl0,
            SDF_Sb0=SDF_Sb0,
            SI_bh_Aw=SI_bh_Aw,
            SI_bh_Sw=SI_bh_Sw,
            SI_bh_Sb=SI_bh_Sb,
            SI_bh_Pl=SI_bh_Pl
        )

        # estimating correction factor to fit BA at t0 and BA at t and
        # choosing whether simulating with multiplication factor
        # or starting at t recalculating the densities and SC
        species_factors = get_basal_area_factors_for_all_species(
            startTage=startTage,
            startTageAw=startTageAw,
            y2bh_Aw=y2bh_Aw,
            SC_Aw=SC_Aw,
            SI_bh_Aw=SI_bh_Aw,
            N_bh_AwT=N_bh_AwT,
            N0_Aw=N0_Aw,
            BA_Aw0=BA_Aw0,
            BA_AwT=BA_AwT,
            startTageSb=startTageSb,
            y2bh_Sb=y2bh_Sb,
            SC_Sb=SC_Sb,
            SI_bh_Sb=SI_bh_Sb,
            N_bh_SbT=N_bh_SbT,
            N0_Sb=N0_Sb,
            BA_Sb0=BA_Sb0,
            BA_SbT=BA_SbT,
            startTageSw=startTageSw,
            y2bh_Sw=y2bh_Sw,
            SC_Sw=SC_Sw,
            SI_bh_Sw=SI_bh_Sw,
            N_bh_SwT=N_bh_SwT,
            N0_Sw=N0_Sw,
            SDF_Pl0=SDF_Pl0,
            SDF_Sb0=SDF_Sb0,
            BA_Sw0=BA_Sw0,
            BA_SwT=BA_SwT,
            startTagePl=startTagePl,
            y2bh_Pl=y2bh_Pl,
            SC_Pl=SC_Pl,
            SI_bh_Pl=SI_bh_Pl,
            N_bh_PlT=N_bh_PlT,
            N0_Pl=N0_Pl,
            SDF_Aw0=SDF_Aw0,
            SDF_Sw0=SDF_Sw0,
            BA_Pl0=BA_Pl0,
            BA_PlT=BA_PlT,
            densities=densities,
            printWarnings=True
        )

        f_Aw = species_factors['f_Aw']
        f_Sw = species_factors['f_Sw']
        f_Sb = species_factors['f_Sb']
        f_Pl = species_factors['f_Pl']

        logger.debug('Getting basal area from time 0 to time of data')
        # use_correction_factor_future here is True because aspen was peculiar,
        # empirically demonstrated that using the factor for the whole
        # simulation yielded better results, probably because of variability in
        # density and basal area unique to aspen we can't do simulation choice
        # = no here for sb, sw, pl because the factor should not be applied for
        # them sw, sb, pl use the factor until the time of data. the subsequent
        # years use the regular basal area increment formula julianno sambatti,
        # november 10, 2016
        # further inspection of those functions indicates they can be improved
        # because the increment only dependds on other values, so we can
        # calculate it as a vector operation using the other arrays then
        # the actual values are just the cumulative sums of the
        # increments that's something for another day; because it is
        # complicated by the whole factor business
        BA_0_to_data_Aw_arr = sim_basal_area_aw(
            startTage, SI_bh_Aw,  N0_Aw, BA_Aw0, SDF_Aw0, f_Aw, densities,
            use_correction_factor_future=True, stop_at_initial_age=False
        )
        BA_0_to_data_Sb_arr = sim_basal_area_sb(
            startTage, SI_bh_Sb, N0_Sb,
            BA_Sb0, f_Sb, densities, stop_at_initial_age=False
        )
        BA_0_to_data_Sw_arr = sim_basal_area_sw(
            startTage, SI_bh_Sw, N0_Sw, SDF_Aw0, SDF_Pl0, SDF_Sb0, BA_Sw0,
            f_Sw, densities, stop_at_initial_age=False
        )
        BA_0_to_data_Pl_arr = sim_basal_area_pl(
            startTage, SI_bh_Pl, N0_Pl, SDF_Aw0, SDF_Sw0, SDF_Sb0, BA_Pl0,
            f_Pl, densities, stop_at_initial_age=False
        )

        output_DF_Aw = pd.DataFrame(BA_0_to_data_Aw_arr, columns=['BA_Aw'])
        output_DF_Sw = pd.DataFrame(BA_0_to_data_Sw_arr, columns=['BA_Sw'])
        output_DF_Sb = pd.DataFrame(BA_0_to_data_Sb_arr, columns=['BA_Sb'])
        output_DF_Pl = pd.DataFrame(BA_0_to_data_Pl_arr, columns=['BA_Pl'])

        densities_DF = pd.DataFrame(densities)
        output_DF = pd.concat(
            [densities_DF, output_DF_Aw, output_DF_Sw, output_DF_Sb, output_DF_Pl],
            axis=1
        )

        import ipdb; ipdb.set_trace()
        for spec in SPECIES:
            output_DF['Gross_Total_Volume_%s' % spec] = gross_total_volume(
                spec,
                output_DF['BA_%s' % spec],
                output_DF['topHeight_%s' % spec]
            )

        output_DF['Gross_Total_Volume_Con'] = output_DF['Gross_Total_Volume_Sw'] \
                                              + output_DF['Gross_Total_Volume_Sb'] \
                                              + output_DF['Gross_Total_Volume_Pl']
        output_DF['Gross_Total_Volume_Dec'] = output_DF['Gross_Total_Volume_Aw']
        output_DF['Gross_Total_Volume_Tot'] = output_DF['Gross_Total_Volume_Con'] \
                                              + output_DF['Gross_Total_Volume_Dec']

        # this could go in the loop above, but is left here for now since
        # the tests are sensitive to column order
        for spec in SPECIES:
            output_DF['MerchantableVolume%s' % spec] = merchantable_volume(
                spec,
                output_DF['N_bh_%sT' % spec],
                output_DF['BA_%s' % spec],
                output_DF['topHeight_%s' % spec],
                output_DF['Gross_Total_Volume_%s' % spec],
                top_dib=utiliz_params[spec.lower()]['topDiamInsideBark'],
                stump_dob=utiliz_params[spec.lower()]['stumpDiamOutsideBark'],
                stump_height=utiliz_params[spec.lower()]['stumpHeight']
            )

        output_DF['MerchantableVolume_Con'] = output_DF['MerchantableVolumeSw'] \
                                              + output_DF['MerchantableVolumeSb'] \
                                              + output_DF['MerchantableVolumePl']
        output_DF['MerchantableVolume_Dec'] = output_DF['MerchantableVolumeAw']
        output_DF['MerchantableVolume_Tot'] = output_DF['MerchantableVolume_Con'] \
                                              + output_DF['MerchantableVolume_Dec']

        output_dict[plot_id] = output_DF

        end = datetime.datetime.now()
        duration = end - start
        logger.info('plot %s took %f seconds', plot_id, duration.total_seconds())

    return output_dict
<|endoftext|>


<commit_before>import unittest
import time
import threading
from aggregator import MetricsAggregator
from dogstatsd import Dogstatsd, init, Server
from util import PidFile
from config import start_jmx_connector
import os
import signal

STATSD_PORT = 8129
class DummyReporter(threading.Thread):
    def __init__(self, metrics_aggregator):
        threading.Thread.__init__(self)
        self.finished = threading.Event()
        self.metrics_aggregator = metrics_aggregator
        self.interval = 4
        self.metrics = None
        self.finished = False
        self.start()

    def run(self):
        while not self.finished:
            time.sleep(self.interval)
            self.flush()

    def flush(self):
        metrics = self.metrics_aggregator.flush()
        if metrics:
            self.metrics = metrics

class JMXTestCase(unittest.TestCase):
    def setUp(self):
        aggregator = MetricsAggregator("test_host")
        self.server = Server(aggregator, "localhost", STATSD_PORT)
        pid_file = PidFile('dogstatsd')
        self.reporter = DummyReporter(aggregator)

        self.t1 = threading.Thread(target=self.server.start)
        self.t1.start()

        confd_path = os.path.realpath(os.path.join(os.path.abspath(__file__), "..", "jmx_yamls"))
        self.jmxfetch_pid = start_jmx_connector(confd_path, {}, statsd_port=STATSD_PORT)

    def tearDown(self):
        self.server.stop()
        self.reporter.finished = True
        os.kill(self.jmxfetch_pid, signal.SIGKILL)

    def testCustomJMXMetric(self):
        count = 0
        while self.reporter.metrics is None:
            time.sleep(1)
            count += 1
            if count > 20:
                raise Exception("No metrics were received in 20 seconds")

        metrics = self.reporter.metrics

        self.assertTrue(type(metrics) == type([]))
        self.assertTrue(len(metrics) > 0)
        self.assertEquals(len([t for t in metrics if t['metric'] == "my.metric.buf" and "instance:jmx_instance1" in t['tags']]), 2, metrics)
        self.assertTrue(len([t for t in metrics if 'type:ThreadPool' in t['tags'] and "instance:jmx_instance1" in t['tags'] and "jmx.catalina" in t['metric']]) > 8, metrics)
        self.assertTrue(len([t for t in metrics if "jmx.java.lang" in t['metric'] and "instance:jmx_instance1" in t['tags']]) > 40, metrics)
        self.assertTrue(len([t for t in metrics if "jvm." in t['metric'] and "instance:jmx_instance1" in t['tags']]) == 6, metrics)

if __name__ == "__main__":
    unittest.main()
<commit_msg>Fix java jmx travis test<commit_after>import unittest
import time
import threading
from aggregator import MetricsAggregator
from dogstatsd import Dogstatsd, init, Server
from util import PidFile
from config import start_jmx_connector
import os
import signal

STATSD_PORT = 8129
class DummyReporter(threading.Thread):
    def __init__(self, metrics_aggregator):
        threading.Thread.__init__(self)
        self.finished = threading.Event()
        self.metrics_aggregator = metrics_aggregator
        self.interval = 4
        self.metrics = None
        self.finished = False
        self.start()

    def run(self):
        while not self.finished:
            time.sleep(self.interval)
            self.flush()

    def flush(self):
        metrics = self.metrics_aggregator.flush()
        if metrics:
            self.metrics = metrics

class JMXTestCase(unittest.TestCase):
    def setUp(self):
        aggregator = MetricsAggregator("test_host")
        self.server = Server(aggregator, "localhost", STATSD_PORT)
        pid_file = PidFile('dogstatsd')
        self.reporter = DummyReporter(aggregator)

        self.t1 = threading.Thread(target=self.server.start)
        self.t1.start()

        confd_path = os.path.realpath(os.path.join(os.path.abspath(__file__), "..", "jmx_yamls"))
        self.jmxfetch_pid = start_jmx_connector(confd_path, {}, statsd_port=STATSD_PORT)

    def tearDown(self):
        self.server.stop()
        self.reporter.finished = True
        os.kill(self.jmxfetch_pid, signal.SIGKILL)

    def testCustomJMXMetric(self):
        count = 0
        while self.reporter.metrics is None:
            time.sleep(1)
            count += 1
            if count > 20:
                raise Exception("No metrics were received in 20 seconds")

        metrics = self.reporter.metrics

        self.assertTrue(type(metrics) == type([]))
        self.assertTrue(len(metrics) > 0)
        self.assertEquals(len([t for t in metrics if t['metric'] == "my.metric.buf" and "instance:jmx_instance1" in t['tags']]), 2, metrics)
        self.assertTrue(len([t for t in metrics if 'type:ThreadPool' in t['tags'] and "instance:jmx_instance1" in t['tags'] and "jmx.catalina" in t['metric']]) > 8, metrics)
        self.assertTrue(len([t for t in metrics if "jmx.java.lang" in t['metric'] and "instance:jmx_instance1" in t['tags']]) > 30, metrics)
        self.assertTrue(len([t for t in metrics if "jvm." in t['metric'] and "instance:jmx_instance1" in t['tags']]) == 6, metrics)

if __name__ == "__main__":
    unittest.main()
<|endoftext|>


<commit_before>#!/usr/bin/env python3
#
# marathon_group.py: generate a Marathon Service Group out of a list of containers
#
# Author: Fernando Sanchez [ fernando at mesosphere.com ]

import json
import sys
import argparse
import subprocess
import os
import uuid

def create_group ( name, containers ):
	"""
	Creates a marathon group taking a list of containers as a parameter.
	If the list has a single member if returns the member.
	"""
	output = '{ 			\
	  "id": "'+name+'",		\
	  "apps": '+containers+'\
	  }'

	return str(output)

def create_external_volume( external_volume_name ):
	"""
	Create an RBD external volume. Assumes RBD is working in the host
	"""

	#check that the volume exists
	#command = "rbd ls | grep "+external_volume_name
	command = "docker volume ls | grep "+external_volume_name
	proc = subprocess.Popen( [command], stdout=subprocess.PIPE, shell=True)
	(out, err) = proc.communicate()
	if external_volume_name in out.decode('utf-8'):
		print('**INFO: volume {0} already exists'.format( external_volume_name ))
		return out.decode('utf-8')
	else:
		command = "docker volume create --driver=rexray --name="+external_volume_name
		proc = subprocess.Popen( [command], stdout=subprocess.PIPE, shell=True)
		(out, err) = proc.communicate()

	#check whether the volume is already mapped
	command = "rbd showmapped | grep "+external_volume_name
	proc = subprocess.Popen( [command], stdout=subprocess.PIPE, shell=True)
	(out, err) = proc.communicate()

	if external_volume_name in out.decode('utf-8'):
		external_volume_device=out.decode('utf-8').split(' ')[-1]
		print('**INFO: volume {0} already mapped to {1}'.format( external_volume_name, external_volume_device  ))
	else:
		#if not, map it
		command = "rbd map "+external_volume_name
		proc = subprocess.Popen( command, stdout=subprocess.PIPE, shell=True)
		(out, err) = proc.communicate()
		#TODO error checking
		external_volume_device = out.decode('utf-8')

	print("**DEBUG: external_volume_MAPPING {}".format(out.decode('utf-8')))

	#format the volume
	command = "mkfs.xfs -f "+external_volume_device
	proc = subprocess.Popen( [command], stdout=subprocess.PIPE, shell=True)
	(out, err) = proc.communicate()

	#unmap rbd
	command = "rbd unmap "+external_volume_name
	proc = subprocess.Popen( [command], stdout=subprocess.PIPE, shell=True)
	(out, err) = proc.communicate()
	#print("**DEBUG: output of rbd map is {}".format(out.decode('utf-8')))

	return out.decode('utf-8')

def	copy_content_to_external_volume( external_volume_name, source_path, mount_path, dest_path ):
	"""
	Copy recursively the content in localhost under "path" to the external volume.
	Assumes the path exists. Assumes the external_volume exists.
	Also creates a mount path which is where the image mounts the copied files (docker volume mount)
	"""

	#check that the volume exists
	command = "rbd ls | grep "+external_volume_name
	proc = subprocess.Popen( [command], stdout=subprocess.PIPE, shell=True)
	(out, err) = proc.communicate()
	if not external_volume_name in out.decode('utf-8'):
		print('**ERROR: volume {0} to copy content into not found'.format( external_volume_name ))
		return False

	#check that the local path exists
	if not os.path.isdir( source_path ):
		print('**ERROR: directory {0} to copy content from not found'.format( source_path ))
		return False

	#check whether the volume is already mapped
	command = "rbd showmapped | grep "+external_volume_name
	proc = subprocess.Popen( [command], stdout=subprocess.PIPE, shell=True)
	(out, err) = proc.communicate()

	if external_volume_name in out.decode('utf-8'):
		output_list=out.decode('utf-8').split()
		print('**DEBUG: output list {}'.format(output_list))
		external_volume_device=output_list[-1]
		print('**INFO: volume {0} already mapped to {1}'.format( external_volume_name, external_volume_device  ))
	else:
		#if not, map it
		print('**INFO: mapping volume {0} for copy'.format( external_volume_name, out.decode('utf-8')  ))
		command = "rbd map "+external_volume_name
		proc = subprocess.Popen( [command], stdout=subprocess.PIPE, shell=True)
		(out, err) = proc.communicate()
		#TODO error checking
		external_volume_device = out.decode('utf-8')

	#TODO error checking
	print('**INFO: volume {0} mapped to {1} for copy'.format( external_volume_name, external_volume_device ))

	#print("**DEBUG: external_volume_device {}".format(external_volume_device))
	#create mount point
	mount_point="/tmp/"+external_volume_name
	command = "mkdir -p "+mount_point
	proc = subprocess.Popen( command, stdout=subprocess.PIPE, shell=True)
	(out, err) = proc.communicate()
	print("**DEBUG: mount_point {}".format(mount_point))

	#mount the volume
	command = ("mount "+external_volume_device+" "+mount_point).replace( '\n', ' ')
	print("**DEBUG: MOUNT COMMAND {}".format(command))

	proc = subprocess.Popen( command, stdout=subprocess.PIPE, shell=True)
	(out, err) = proc.communicate()

	print("**DEBUG: MOUNT command result {}".format(out.decode('utf-8')))

	#create path
	#command = "mkdir -p "+mount_point+"/"+dest_path
	#proc = subprocess.Popen( [command], stdout=subprocess.PIPE, shell=True)
	#(out, err) = proc.communicate()

	#create mount path
	print("**DEBUG: MOUNT path to be created is: {}".format(mount_point+"/"+mount_path))
	command = "mkdir -p "+mount_point+"/"+mount_path
	proc = subprocess.Popen( [command], stdout=subprocess.PIPE, shell=True)
	(out, err) = proc.communicate()

	#recursively copy the content
	print("**DEBUG: COPY from {0} to {1}".format(source_path, mount_point+"/"+source_path[2:]))

	#copy source to src_path
	command = "cp -R "+source_path+" "+mount_point+"/"+source_path[2:]
	proc = subprocess.Popen( command, stdout=subprocess.PIPE, shell=True)
	(out, err) = proc.communicate()

	#recursively copy the content
	print("**DEBUG: COPY from {0} to {1}".format(source_path, mount_point+"/"+mount_path+"/"+source_path[2:]))

	#copy source to mount_path/src_path
	command = "cp -R "+source_path+" "+mount_point+"/"+mount_path+"/"+source_path[2:]
	proc = subprocess.Popen( command, stdout=subprocess.PIPE, shell=True)
	(out, err) = proc.communicate()

	data_dir="data"
	print("**DEBUG: mkdir {0} ".format(mount_point+"/"+mount_path+"/"+data_dir))
	#mkdir "/data"
	command = "mkdir "+mount_point+"/"+mount_path+"/"+data_dir
	proc = subprocess.Popen( command, stdout=subprocess.PIPE, shell=True)
	(out, err) = proc.communicate()

	#copy source to mount_path - not needed, just need to CREATE mount_path
	#command = "cp -R "+source_path+" "+mount_point+"/"+mount_path
	#proc = subprocess.Popen( command, stdout=subprocess.PIPE, shell=True)
	#(out, err) = proc.communicate()

	#umount the volume
	command = "umount "+external_volume_device
	proc = subprocess.Popen( [command], stdout=subprocess.PIPE, shell=True)
	(out, err) = proc.communicate()

	#unmap rbd
	command = "rbd unmap "+external_volume_name
	proc = subprocess.Popen( [command], stdout=subprocess.PIPE, shell=True)
	(out, err) = proc.communicate()
	#print("**DEBUG: output of rbd map is {}".format(out.decode('utf-8')))

	#delete temp mount point
	command = "rm -Rf "+mount_point+"/"+mount_path
	proc = subprocess.Popen( [command], stdout=subprocess.PIPE, shell=True)
	(out, err) = proc.communicate()
	print("**DEBUG: remove mount point (should be empty) {}".format(mount_point))

	return out.decode('utf-8')

def modify_volume_for_external ( volume, app_name ):
	"""
	Receives a Marathon volume definition as a dictionary that includes a local directory to load.
	Creates an external persistent volume in a docker external volume (assumed to be mounted).
	Creates the local volume structure on the external volume, and copies the contents.
	Modifies the volume to use the newly create external volume.
	Returns the new volume dictionary.
	"""

	#get firstPartOfHostPath, etc.
	host_path = volume['hostPath'] 					#./app
	#first_part_of_host_path = host_path.split('/' , 1)[0]	#app
	#last_part_of_host_path = host_path.split('/' , 1)[1]		#NULL
	container_path = volume['containerPath']							#/src/app
	first_part_of_container_path = container_path[1:].split('/', 1)[0]	#src
	last_part_of_container_path = container_path[1:].split('/', 1)[1]	#app
	#create a volume
	external_volume_name = app_name+'-'+host_path[2:].replace('/','_')
	create_external_volume( external_volume_name ) #nextcloud_apps_UUID
	#copy content from volume[hostPath] to volume
	copy_content_to_external_volume( external_volume_name, host_path, \
		first_part_of_container_path, last_part_of_container_path )
	#modify volume
	volume['external'] = { 						#mount it as external volume
		'name': external_volume_name,
		'provider': 'dvdi',
		'options': {
		'dvdi/driver': 'rexray'
		}
	}
	#change containerPath to the firstPiece only
	volume['containerPath'] = first_part_of_container_path	#src
	del( volume['hostPath'] )							#external volumes do not use hostpath

	return volume

def modify_group ( group ):
	"""
	Modifies a marathon group received as a printable string to adapt the apps inside it
	to the desired parameters.
	Adds AcceptedResourceRoles = "*"
	Deletes any hostPort values to have Marathon assign them automatically.
	Adds a label HAPROXY_GROUP=external for all hostPort values.
	If the apps mounts any volume that is local for compose (localhost), it mounts it as "external/rexray"
	Returns the group as a modified string.
	"""

	group_dict = json.loads( group )

	for app in group_dict['apps']:
		app['acceptedResourceRoles']=["*"]
		for portMapping in app.get('container',{}).get('docker',{}).get('portMappings',{}):
			if portMapping.get('hostPort',{}): 	#delete ANY hostPort values, use them for VIP
				#create a VIP for every app, with a known pattern: group_dict['id']+'-'+app['id']:hostPort
				vip = "/"+group_dict['id']+'-'+app['id']+":"+str(portMapping['hostPort'])
				portMapping['labels'] = { "VIP_0": vip }
				portMapping['hostPort'] = 0
				#make the app available in MarathonLB
				if 'labels' in app:
					app['labels'].update( {"HAPROXY_GROUP": "external"} )# if there was a hostPort add to MLB
				else:
					app['labels'] = { "HAPROXY_GROUP": "external" }

		#modify all volumes in the groups apps so that "this directory" volumes become external
		for volume in app.get('container', {}).get('volumes', {}):
			if volume['hostPath'][:2] == "./":			#if the volume is "this dir" for compose
				volume = modify_volume_for_external( volume, group_dict['id']+'-'+app['id'] )
						#modify it so that the local files are reachable via external volume

	return json.dumps( group_dict )

if __name__ == "__main__":

	#parse command line arguments
	parser = argparse.ArgumentParser(description='Convert a list of containers to a Marathon Service Group.', \
		usage='marathon_group.py -i [container_list_filename] -n [group_name]'
		)
	parser.add_argument('-i', '--input', help='name of the file including the list of containers', required=True)
	parser.add_argument('-n', '--name', help='name to be given to the Marathon Service Group', required=True)
	args = vars( parser.parse_args() )

	#remove the trailing \n from file
	containers = ""
	for line in open( args['input'], 'r' ):
		containers += line.rstrip()
	#detect if it's just one app - if so, get in list
	if containers[0]=="{":
		containers="["+containers+"]"
	group = create_group( args['name'], containers )
	modified_group = modify_group( group )
	output_file=open( "./group.json", "w")
	print( modified_group, file=output_file )
	input( "***DEBUG: Press ENTER to continue...")
	sys.exit(0)

<commit_msg>mkdir data_dir<commit_after>#!/usr/bin/env python3
#
# marathon_group.py: generate a Marathon Service Group out of a list of containers
#
# Author: Fernando Sanchez [ fernando at mesosphere.com ]

import json
import sys
import argparse
import subprocess
import os
import uuid

def create_group ( name, containers ):
	"""
	Creates a marathon group taking a list of containers as a parameter.
	If the list has a single member if returns the member.
	"""
	output = '{ 			\
	  "id": "'+name+'",		\
	  "apps": '+containers+'\
	  }'

	return str(output)

def create_external_volume( external_volume_name ):
	"""
	Create an RBD external volume. Assumes RBD is working in the host
	"""

	#check that the volume exists
	#command = "rbd ls | grep "+external_volume_name
	command = "docker volume ls | grep "+external_volume_name
	proc = subprocess.Popen( [command], stdout=subprocess.PIPE, shell=True)
	(out, err) = proc.communicate()
	if external_volume_name in out.decode('utf-8'):
		print('**INFO: volume {0} already exists'.format( external_volume_name ))
		return out.decode('utf-8')
	else:
		command = "docker volume create --driver=rexray --name="+external_volume_name
		proc = subprocess.Popen( [command], stdout=subprocess.PIPE, shell=True)
		(out, err) = proc.communicate()

	#check whether the volume is already mapped
	command = "rbd showmapped | grep "+external_volume_name
	proc = subprocess.Popen( [command], stdout=subprocess.PIPE, shell=True)
	(out, err) = proc.communicate()

	if external_volume_name in out.decode('utf-8'):
		external_volume_device=out.decode('utf-8').split(' ')[-1]
		print('**INFO: volume {0} already mapped to {1}'.format( external_volume_name, external_volume_device  ))
	else:
		#if not, map it
		command = "rbd map "+external_volume_name
		proc = subprocess.Popen( command, stdout=subprocess.PIPE, shell=True)
		(out, err) = proc.communicate()
		#TODO error checking
		external_volume_device = out.decode('utf-8')

	print("**DEBUG: external_volume_MAPPING {}".format(out.decode('utf-8')))

	#format the volume
	command = "mkfs.xfs -f "+external_volume_device
	proc = subprocess.Popen( [command], stdout=subprocess.PIPE, shell=True)
	(out, err) = proc.communicate()

	#unmap rbd
	command = "rbd unmap "+external_volume_name
	proc = subprocess.Popen( [command], stdout=subprocess.PIPE, shell=True)
	(out, err) = proc.communicate()
	#print("**DEBUG: output of rbd map is {}".format(out.decode('utf-8')))

	return out.decode('utf-8')

def	copy_content_to_external_volume( external_volume_name, source_path, mount_path, dest_path ):
	"""
	Copy recursively the content in localhost under "path" to the external volume.
	Assumes the path exists. Assumes the external_volume exists.
	Also creates a mount path which is where the image mounts the copied files (docker volume mount)
	"""

	#check that the volume exists
	command = "rbd ls | grep "+external_volume_name
	proc = subprocess.Popen( [command], stdout=subprocess.PIPE, shell=True)
	(out, err) = proc.communicate()
	if not external_volume_name in out.decode('utf-8'):
		print('**ERROR: volume {0} to copy content into not found'.format( external_volume_name ))
		return False

	#check that the local path exists
	if not os.path.isdir( source_path ):
		print('**ERROR: directory {0} to copy content from not found'.format( source_path ))
		return False

	#check whether the volume is already mapped
	command = "rbd showmapped | grep "+external_volume_name
	proc = subprocess.Popen( [command], stdout=subprocess.PIPE, shell=True)
	(out, err) = proc.communicate()

	if external_volume_name in out.decode('utf-8'):
		output_list=out.decode('utf-8').split()
		print('**DEBUG: output list {}'.format(output_list))
		external_volume_device=output_list[-1]
		print('**INFO: volume {0} already mapped to {1}'.format( external_volume_name, external_volume_device  ))
	else:
		#if not, map it
		print('**INFO: mapping volume {0} for copy'.format( external_volume_name, out.decode('utf-8')  ))
		command = "rbd map "+external_volume_name
		proc = subprocess.Popen( [command], stdout=subprocess.PIPE, shell=True)
		(out, err) = proc.communicate()
		#TODO error checking
		external_volume_device = out.decode('utf-8')

	#TODO error checking
	print('**INFO: volume {0} mapped to {1} for copy'.format( external_volume_name, external_volume_device ))

	#print("**DEBUG: external_volume_device {}".format(external_volume_device))
	#create mount point
	mount_point="/tmp/"+external_volume_name
	command = "mkdir -p "+mount_point
	proc = subprocess.Popen( command, stdout=subprocess.PIPE, shell=True)
	(out, err) = proc.communicate()
	print("**DEBUG: mount_point {}".format(mount_point))

	#mount the volume
	command = ("mount "+external_volume_device+" "+mount_point).replace( '\n', ' ')
	print("**DEBUG: MOUNT COMMAND {}".format(command))

	proc = subprocess.Popen( command, stdout=subprocess.PIPE, shell=True)
	(out, err) = proc.communicate()

	print("**DEBUG: MOUNT command result {}".format(out.decode('utf-8')))

	#create path
	#command = "mkdir -p "+mount_point+"/"+dest_path
	#proc = subprocess.Popen( [command], stdout=subprocess.PIPE, shell=True)
	#(out, err) = proc.communicate()

	#create mount path
	print("**DEBUG: MOUNT path to be created is: {}".format(mount_point+"/"+mount_path))
	command = "mkdir -p "+mount_point+"/"+mount_path
	proc = subprocess.Popen( [command], stdout=subprocess.PIPE, shell=True)
	(out, err) = proc.communicate()

	#recursively copy the content
	print("**DEBUG: COPY from {0} to {1}".format(source_path, mount_point+"/"+source_path[2:]))

	#copy source to src_path
	command = "cp -R "+source_path+" "+mount_point+"/"+source_path[2:]
	proc = subprocess.Popen( command, stdout=subprocess.PIPE, shell=True)
	(out, err) = proc.communicate()

	#recursively copy the content
	print("**DEBUG: COPY from {0} to {1}".format(source_path, mount_point+"/"+mount_path+"/"+source_path[2:]))

	#copy source to mount_path/src_path
	command = "cp -R "+source_path+" "+mount_point+"/"+mount_path+"/"+source_path[2:]
	proc = subprocess.Popen( command, stdout=subprocess.PIPE, shell=True)
	(out, err) = proc.communicate()

	data_dir="data"
	print("**DEBUG: mkdir {0} ".format(mount_point+"/"+mount_path+"/"+data_dir))
	#mkdir "/data"
	command = "mkdir "+mount_point+"/"+mount_path+"/"+data_dir
	proc = subprocess.Popen( command, stdout=subprocess.PIPE, shell=True)
	(out, err) = proc.communicate()

	#chmod 777 *
	print("**DEBUG: COPY from {0} to {1}".format(source_path, mount_point+"/"+mount_path+"/"+source_path[2:]))

	#copy source to mount_path/src_path
	command = "chmod -R 777 "+source_path+" "+mount_point
	proc = subprocess.Popen( command, stdout=subprocess.PIPE, shell=True)
	(out, err) = proc.communicate()

	#copy source to mount_path - not needed, just need to CREATE mount_path
	#command = "cp -R "+source_path+" "+mount_point+"/"+mount_path
	#proc = subprocess.Popen( command, stdout=subprocess.PIPE, shell=True)
	#(out, err) = proc.communicate()

	#umount the volume
	command = "umount "+external_volume_device
	proc = subprocess.Popen( [command], stdout=subprocess.PIPE, shell=True)
	(out, err) = proc.communicate()

	#unmap rbd
	command = "rbd unmap "+external_volume_name
	proc = subprocess.Popen( [command], stdout=subprocess.PIPE, shell=True)
	(out, err) = proc.communicate()
	#print("**DEBUG: output of rbd map is {}".format(out.decode('utf-8')))

	#delete temp mount point
	command = "rm -Rf "+mount_point+"/"+mount_path
	proc = subprocess.Popen( [command], stdout=subprocess.PIPE, shell=True)
	(out, err) = proc.communicate()
	print("**DEBUG: remove mount point (should be empty) {}".format(mount_point))

	return out.decode('utf-8')

def modify_volume_for_external ( volume, app_name ):
	"""
	Receives a Marathon volume definition as a dictionary that includes a local directory to load.
	Creates an external persistent volume in a docker external volume (assumed to be mounted).
	Creates the local volume structure on the external volume, and copies the contents.
	Modifies the volume to use the newly create external volume.
	Returns the new volume dictionary.
	"""

	#get firstPartOfHostPath, etc.
	host_path = volume['hostPath'] 					#./app
	#first_part_of_host_path = host_path.split('/' , 1)[0]	#app
	#last_part_of_host_path = host_path.split('/' , 1)[1]		#NULL
	container_path = volume['containerPath']							#/src/app
	first_part_of_container_path = container_path[1:].split('/', 1)[0]	#src
	last_part_of_container_path = container_path[1:].split('/', 1)[1]	#app
	#create a volume
	external_volume_name = app_name+'-'+host_path[2:].replace('/','_')
	create_external_volume( external_volume_name ) #nextcloud_apps_UUID
	#copy content from volume[hostPath] to volume
	copy_content_to_external_volume( external_volume_name, host_path, \
		first_part_of_container_path, last_part_of_container_path )
	#modify volume
	volume['external'] = { 						#mount it as external volume
		'name': external_volume_name,
		'provider': 'dvdi',
		'options': {
		'dvdi/driver': 'rexray'
		}
	}
	#change containerPath to the firstPiece only
	volume['containerPath'] = first_part_of_container_path	#src
	del( volume['hostPath'] )							#external volumes do not use hostpath

	return volume

def modify_group ( group ):
	"""
	Modifies a marathon group received as a printable string to adapt the apps inside it
	to the desired parameters.
	Adds AcceptedResourceRoles = "*"
	Deletes any hostPort values to have Marathon assign them automatically.
	Adds a label HAPROXY_GROUP=external for all hostPort values.
	If the apps mounts any volume that is local for compose (localhost), it mounts it as "external/rexray"
	Returns the group as a modified string.
	"""

	group_dict = json.loads( group )

	for app in group_dict['apps']:
		app['acceptedResourceRoles']=["*"]
		for portMapping in app.get('container',{}).get('docker',{}).get('portMappings',{}):
			if portMapping.get('hostPort',{}): 	#delete ANY hostPort values, use them for VIP
				#create a VIP for every app, with a known pattern: group_dict['id']+'-'+app['id']:hostPort
				vip = "/"+group_dict['id']+'-'+app['id']+":"+str(portMapping['hostPort'])
				portMapping['labels'] = { "VIP_0": vip }
				portMapping['hostPort'] = 0
				#make the app available in MarathonLB
				if 'labels' in app:
					app['labels'].update( {"HAPROXY_GROUP": "external"} )# if there was a hostPort add to MLB
				else:
					app['labels'] = { "HAPROXY_GROUP": "external" }

		#modify all volumes in the groups apps so that "this directory" volumes become external
		for volume in app.get('container', {}).get('volumes', {}):
			if volume['hostPath'][:2] == "./":			#if the volume is "this dir" for compose
				volume = modify_volume_for_external( volume, group_dict['id']+'-'+app['id'] )
						#modify it so that the local files are reachable via external volume

	return json.dumps( group_dict )

if __name__ == "__main__":

	#parse command line arguments
	parser = argparse.ArgumentParser(description='Convert a list of containers to a Marathon Service Group.', \
		usage='marathon_group.py -i [container_list_filename] -n [group_name]'
		)
	parser.add_argument('-i', '--input', help='name of the file including the list of containers', required=True)
	parser.add_argument('-n', '--name', help='name to be given to the Marathon Service Group', required=True)
	args = vars( parser.parse_args() )

	#remove the trailing \n from file
	containers = ""
	for line in open( args['input'], 'r' ):
		containers += line.rstrip()
	#detect if it's just one app - if so, get in list
	if containers[0]=="{":
		containers="["+containers+"]"
	group = create_group( args['name'], containers )
	modified_group = modify_group( group )
	output_file=open( "./group.json", "w")
	print( modified_group, file=output_file )
	input( "***DEBUG: Press ENTER to continue...")
	sys.exit(0)

<|endoftext|>


<commit_before>import os
import nipype.pipeline.engine as pe
from nipype.interfaces.dcmstack import DcmStack, MergeNifti, CopyMeta
from nipype.interfaces.utility import IdentityInterface, Function
from ..interfaces import XNATDownload, XNATUpload, Fastfit

def run(*input_specs, **opts):
    """
    Executes the PK mapping workflow built in :meth:`create_workflow`.

    Each input specification identifies an XNAT resource container. The
    specification is a (subject, session, container type, container name)
    tuple, where the container type is a :meth:`XNATDownload.download`
    container type and the container name is a XNAT resource container
    label. There is one resource per session.

    The NiFTI session mask and resource files are downloaded from XNAT
    into the ``input`` subdirectory of the ``base_dir`` option value
    (default is the current directory).

    The PK mapping workflow is then built as described in
    :meth:`create_workflow` and executed. The result is uploaded to an XNAT
    analysis resource.

    :param input_specs: the XNAT (subject, session, resource) image tuples
    :param opts: the workflow options
    :return: the PK mapping XNAT (subject, session, analysis) tuples
    """

    # The work directory.
    base_dir = opts.get('base_dir') or os.getcwd()
    # The image download location.
    work = os.path.join(base_dir, 'input')

    # Run the workflow on each session.
    outputs = {spec: _analyze_session_images(*spec, work=work, **opts)
        for spec in input_specs}

    return outputs

def _analyze_session_images(subject, session, ctr_type, ctr_name, work, **opts):
    # The work directory.
    base_dir = opts.get('base_dir') or os.getcwd()
    # The image download location.
    dest = os.path.join(base_dir, 'data', subject, session)

    # Download the mask.
    dl_mask = XNATDownload(project='QIN', subject=subject, session=session,
        reconstruction='mask', dest=dest)
    mask = dl_mask.run().result.outputs.out_files[0]

    # Download the images.
    ctr_opt = {ctr_type: ctr_name}
    dl_images = XNATDownload(project='QIN', subject=subject, session=session,
        dest=dest, **ctr_opt)
    images = dl_images.run().result.outputs.out_files

    # Make the workflow.
    wf = create_workflow(in_files=images)

    # Execute the workflow.
    result = wf.run()

    return result.outputs.output_spec.get()

def create_workflow(**opts):
    """
    Creates the Nipype workflow.

    .. |H2O| replace:: H\ :sub:`2`\ O
    .. |R10| replace:: R1\ :sub:`0`

    The workflow calculates the PK mapping parameters for the input as follows:

    - Compute the |R10| value, if it is not given in the options

    - Convert the DCE time series to a R1 map series

    - Determine the AIF and R1 fit parameters from the time series

    - Perform the BOLERO model pharmacokinetic mapping

    - Upload the PK mapping result to XNAT

    The workflow inputs are defined in the ``input_spec`` node.
    image specification identifies an XNAT resource. The
    specification is a *(subject, session, resource)* tuple, where
    *resource* is a :meth:`XNATDownload.download`
    The NiFTI images are downloaded from XNAT into the
    ``input`` subdirectory of the ``base_dir`` specified in the options
    (default is the current directory).

    If the R1:sub:`0` option is not set, then it is computed from
    the proton density weighted scans and DCE series baseline image.

    The outputs are collected in the ``output_spec`` node with the following
    fields:

    - ``r1_series``: the R1 series files

    - ``params_csv``: the AIF and R1 parameter CSV file

    - ``k_trans``: the K\ :sup:`trans` extra/intravasation transfer rate

    - ``v_e``: the v\ :sub:`e` interstitial volume fraction

    - ``tau_i``: the intracellular |H2O| mean lifetime

    In addition, if R1:sub:`0` is computed, then the output includes the
    following fields:

    - ``pdw_image``: the proton density weighted image

    - ``dce_baseline``: the DCE series baseline image

    - ``r1_0``: the computed |R10| value

    This workflow is adapted from https://everett.ohsu.edu/hg/qin_dce.

    :param opts: the optional workflow inputs
    :keyword r1_0_val: the optional |R10| value
    :keyword pd_dir: the proton density weighted scan directory,
        if the R1:sub:`0` option is not set
    :keyword max_r1_0: the maximum computed |R10| value,
        if the |R10| option is not set
    :keyword in_files: the input images
    :keyword mask_file: the constraining mask file
    :keyword baseline_end_idx: the number of images to merge into a baseline image
    """
    workflow = pe.Workflow(name='pk_mapping')

    # Set up the input node.
    inputs = ['in_files', 'mask_file', 'baseline_end_idx']
    if opts.get(use_fixed_r1_0):
        inputs += ['r1_0_val']
    else:
        inputs += ['pd_dir', 'max_r1_0']
    input_spec = pe.Node(IdentityInterface(fields=inputs),
                        name='input_spec')

    # Merge the DCE data to 4D.
    dce_merge = pe.Node(MergeNifti(), name='dce_merge')
    dce_merge.inputs.out_format = 'dce_series'
    workflow.connect(input_spec, 'in_files', dce_merge, 'in_files')

    # If we are not using a fixed r1_0 value, then compute a map from a
    # proton density weighted scan and the baseline of the DCE series.
    if not use_fixed_r1_0:
        # Convert each series dir to a Nifti.
        pd_stack = pe.Node(DcmStack(), name='pd_stack')
        pd_stack.inputs.embed_meta = True
        workflow.connect(input_spec, 'pd_dir',
                         pd_stack, 'dicom_files')

        # Create the DCE baseline image.
        make_base = pe.Node(Function(input_names=['dce_images',
                                                  'baseline_end_idx'],
                                     output_names=['baseline_nii'],
                                     function=_make_baseline),
                            name='make_base')
        workflow.connect(input_spec, 'in_files',
                         make_base, 'dce_images')
        workflow.connect(input_spec, 'baseline_end_idx',
                         make_base, 'baseline_end_idx')

        # Create the R1_0 map.
        get_r1_0 = pe.Node(Function(input_names=['pdw_image',
                                                 't1w_image',
                                                 'max_r1_0',
                                                 'mask_file',
                                                ],
                                    output_names=['r1_0_map'],
                                    function=_make_r1_0),
                           name='get_r1_0')
        workflow.connect(pd_stack, 'out_file', get_r1_0, 'pdw_image')
        workflow.connect(make_base, 'baseline_nii', get_r1_0, 't1w_image')
        workflow.connect(input_spec, 'max_r1_0', get_r1_0, 'max_r1_0')
        workflow.connect(input_spec, 'mask_file', get_r1_0, 'mask_file')

    # Convert DCE time series to series of R1 maps.
    get_r1_series = pe.Node(Function(input_names=['time_series',
                                                  'r1_0',
                                                  'baseline_end',
                                                  'mask_file',
                                                 ],
                                     output_names=['r1_series'],
                                     function=_make_r1_series),
                            name='get_r1_series')
    workflow.connect(dce_merge, 'out_file',
                     get_r1_series, 'time_series')
    workflow.connect(input_spec, 'baseline_end_idx',
                     get_r1_series, 'baseline_end')
    workflow.connect(input_spec, 'mask_file', get_r1_series, 'mask_file')
    if use_fixed_r1_0:
        workflow.connect(input_spec, 'r1_0_val', get_r1_series, 'r1_0')
    else:
        workflow.connect(get_r1_0, 'r1_0_map', get_r1_series, 'r1_0')

    # Copy meta data to the R1 series.
    copy_meta = pe.Node(CopyMeta(), name='copy_meta')
    copy_meta.inputs.include_classes = [('global', 'const'), ('time', 'samples')]
    workflow.connect(dce_merge, 'out_file', copy_meta, 'src_file')
    workflow.connect(get_r1_series, 'r1_series', copy_meta, 'dest_file')

    # Get the pharmacokinetic mapping parameters.
    get_params = pe.Node(Function(input_names=['time_series'],
                                  output_names=['params_csv'],
                                  function=_get_fit_params),
                         name='get_params')
    workflow.connect(dce_merge, 'out_file', get_params, 'time_series')

    # Perform the pharmacokinetic mapping.
    pk_map = pe.Node(Fastfit(), name='pk_map')
    pk_map.inputs.model_name = 'bolero'
    workflow.connect(copy_meta, 'dest_file',
                     pk_map, 'target_data')
    workflow.connect(input_spec, 'mask_file',
                     pk_map, 'mask')
    workflow.connect(get_params, 'params_csv',
                     pk_map, 'params_csv')

    # Set up the outputs.
    outputs = ['r1_series',
               'params_csv',
               'k_trans',
               'v_e',
               'tau_i'
              ]
    if not use_fixed_r1_0:
        outputs += ['pdw_image', 'dce_baseline', 'r1_0']
    output_spec = pe.Node(IdentityInterface(fields=outputs), name='output_spec')
    workflow.connect(copy_meta, 'dest_file', output_spec, 'r1_series')
    workflow.connect(get_params, 'params_csv', output_spec, 'params_csv')
    workflow.connect(pk_map, 'k_trans', output_spec, 'k_trans')
    workflow.connect(pk_map, 'v_e', output_spec, 'v_e')
    workflow.connect(pk_map, 'tau_i', output_spec, 'tau_i')

    if not use_fixed_r1_0:
        workflow.connect(pd_stack, 'out_file', output_spec, 'pdw_image')
        workflow.connect(make_base, 'baseline_nii',
                         output_spec, 'dce_baseline')
        workflow.connect(get_r1_0, 'r1_0_map', output_spec, 'r1_0')

    # Upload the outputs to XNAT.
    for field in outputs:
        pass # TODO

    return workflow

def _make_baseline(dce_images, baseline_end_idx):
    from nipype.interfaces.dcmstack import MergeNifti
    from nipype.interfaces import fsl

    assert baseline_end_idx > 0
    if baseline_end_idx > 1:
        baseline_images = dce_images[:baseline_end_idx]
        merge_nii = MergeNifti()
        merge_nii.inputs.in_files = baseline_images
        merge_nii.inputs.out_format = 'baseline_merged'
        merged = merge_nii.run().outputs.out_file
        mean_image = fsl.MeanImage(name='mean_image')
        mean_image.inputs.in_file = merged

        return mean_image.run().outputs.out_file
    else:
        return dce_images[0]

def _make_r1_0(pdw_image, t1w_image, max_r1_0, mask_file, **kwargs):
    import os
    from os import path
    import nibabel as nb
    import numpy as np
    from pdw_t1w_to_r1 import pdw_t1w_to_r1
    from dcmstack.dcmmeta import NiftiWrapper

    pdw_nw = NiftiWrapper(nb.load(pdw_image), make_empty=True)
    t1w_nw = NiftiWrapper(nb.load(t1w_image), make_empty=True)
    r1_space = np.arange(0.01, max_r1_0, 0.01)
    mask = nb.load(mask_file).get_data()
    r1_0 = pdw_t1w_to_r1(pdw_nw, t1w_nw, r1_space=r1_space, mask=mask,
                         **kwargs)

    cwd = os.getcwd()
    out_nii = nb.Nifti1Image(r1_0, pdw_nw.nii_img.get_affine())
    out_fn = path.join(cwd, 'r1_0_map.nii.gz')
    nb.save(out_nii, out_fn)
    return out_fn

def _make_r1_series(time_series, r1_0, mask_file, **kwargs):
    import os
    from os import path
    import nibabel as nb
    import numpy as np
    from dce_to_r1 import dce_to_r1
    from dcmstack.dcmmeta import NiftiWrapper

    dce_nw = NiftiWrapper(nb.load(time_series), make_empty=True)
    if not isinstance(r1_0, float):
        r1_0 = nb.load(r1_0).get_data()
    mask = nb.load(mask_file).get_data()
    r1_series = dce_to_r1(dce_nw, r1_0, mask=mask, **kwargs)

    cwd = os.getcwd()
    out_nii = nb.Nifti1Image(r1_series, dce_nw.nii_img.get_affine())
    out_fn = path.join(cwd, 'r1_series.nii.gz')
    nb.save(out_nii, out_fn)
    return out_fn

def _get_fit_params(time_series):
    import os, csv
    from os import path
    import nibabel as nb
    import numpy as np
    from dcmstack.dcmmeta import NiftiWrapper
    from dcmstack import dcm_time_to_sec

    nii = nb.load(time_series)
    data = nii.get_data()
    n_vols = data.shape[-1]
    signal_means = np.array([np.mean(data[:,:,:,idx])
                             for idx in xrange(n_vols)])
    signal_diffs = np.diff(signal_means)

    # If we see a difference in average signal larger than double the .
    # difference from first two points, take that as bolus arrival.
    base_diff = np.abs(signal_diffs[0])
    for idx, diff_val in enumerate(signal_diffs[1:]):
        if diff_val > 2 * base_diff:
            bolus_idx = idx + 1
            break
    else:
        raise ValueError("Unable to determine bolus arrival")

    # Figure out the time in between bolus_idx and bolus_idx+1.
    nw = NiftiWrapper(nii)
    acq_time0 = dcm_time_to_sec(nw.get_meta('AcquisitionTime', (0, 0, 0, 0)))
    acq_time1 = dcm_time_to_sec(nw.get_meta('AcquisitionTime', (0, 0, 0, bolus_idx)))
    acq_time2 = dcm_time_to_sec(nw.get_meta('AcquisitionTime', (0, 0, 0, bolus_idx+1)))
    aif_shift = ((acq_time1 + acq_time2) / 2.0) - acq_time0

    # Create parameter CSV.
    with open('params.csv', 'w') as csv_file:
        csv_writer = csv.writer(csv_file)
        csv_writer.writerow(['aif_params', '0.674', '0.4', '2.2',
                             '0.23', '1.3', '0.09', '0.0013', '0.0'])
        csv_writer.writerow(['aif_delta_t', '1.5'])
        csv_writer.writerow(['aif_shift', str(aif_shift)])
        csv_writer.writerow(['r1_cr', '3.8'])
        csv_writer.writerow(['r1_b_pre', '0.71'])
    return path.join(os.getcwd(), 'params.csv')
<commit_msg>Add reST substitutions.<commit_after>import os
import nipype.pipeline.engine as pe
from nipype.interfaces.dcmstack import DcmStack, MergeNifti, CopyMeta
from nipype.interfaces.utility import IdentityInterface, Function
from ..interfaces import XNATDownload, XNATUpload, Fastfit

def run(*input_specs, **opts):
    """
    Executes the PK mapping workflow built in :meth:`create_workflow`.

    Each input specification identifies an XNAT resource container. The
    specification is a (subject, session, container type, container name)
    tuple, where the container type is a :meth:`XNATDownload.download`
    container type and the container name is a XNAT resource container
    label. There is one resource per session.

    The NiFTI session mask and resource files are downloaded from XNAT
    into the ``input`` subdirectory of the ``base_dir`` option value
    (default is the current directory).

    The PK mapping workflow is then built as described in
    :meth:`create_workflow` and executed. The result is uploaded to an XNAT
    analysis resource.

    :param input_specs: the XNAT (subject, session, resource) image tuples
    :param opts: the workflow options
    :return: the PK mapping XNAT (subject, session, analysis) tuples
    """

    # The work directory.
    base_dir = opts.get('base_dir') or os.getcwd()
    # The image download location.
    work = os.path.join(base_dir, 'input')

    # Run the workflow on each session.
    outputs = {spec: _analyze_session_images(*spec, work=work, **opts)
        for spec in input_specs}

    return outputs

def _analyze_session_images(subject, session, ctr_type, ctr_name, work, **opts):
    # The work directory.
    base_dir = opts.get('base_dir') or os.getcwd()
    # The image download location.
    dest = os.path.join(base_dir, 'data', subject, session)

    # Download the mask.
    dl_mask = XNATDownload(project='QIN', subject=subject, session=session,
        reconstruction='mask', dest=dest)
    mask = dl_mask.run().result.outputs.out_files[0]

    # Download the images.
    ctr_opt = {ctr_type: ctr_name}
    dl_images = XNATDownload(project='QIN', subject=subject, session=session,
        dest=dest, **ctr_opt)
    images = dl_images.run().result.outputs.out_files

    # Make the workflow.
    wf = create_workflow(in_files=images)

    # Execute the workflow.
    result = wf.run()

    return result.outputs.output_spec.get()

def create_workflow(**opts):
    """
    Creates the Nipype workflow.

    .. |H2O| replace:: H\ :sub:`2`\ O
    .. |Ktrans| replace:: K\ :sup:`trans`
    .. |ve| replace:: v\ :sub:`e`
    .. |R10| replace:: R1\ :sub:`0`

    The workflow calculates the PK mapping parameters for the input as follows:

    - Compute the |R10| value, if it is not given in the options

    - Convert the DCE time series to a R1 map series

    - Determine the AIF and R1 fit parameters from the time series

    - Perform the BOLERO model pharmacokinetic mapping

    - Upload the PK mapping result to XNAT

    The workflow inputs are defined in the ``input_spec`` node.
    image specification identifies an XNAT resource. The
    specification is a *(subject, session, resource)* tuple, where
    *resource* is a :meth:`XNATDownload.download`
    The NiFTI images are downloaded from XNAT into the
    ``input`` subdirectory of the ``base_dir`` specified in the options
    (default is the current directory).

    If the |R10| option is not set, then it is computed from
    the proton density weighted scans and DCE series baseline image.

    The outputs are collected in the ``output_spec`` node with the following
    fields:

    - ``r1_series``: the R1 series files

    - ``params_csv``: the AIF and R1 parameter CSV file

    - ``k_trans``: the |Ktrans| extra/intravasation transfer rate

    - ``v_e``: the |ve| interstitial volume fraction

    - ``tau_i``: the intracellular |H2O| mean lifetime

    In addition, if |R10| is computed, then the output includes the
    following fields:

    - ``pdw_image``: the proton density weighted image

    - ``dce_baseline``: the DCE series baseline image

    - ``r1_0``: the computed |R10| value

    This workflow is adapted from https://everett.ohsu.edu/hg/qin_dce.

    :param opts: the optional workflow inputs
    :keyword r1_0_val: the optional |R10| value
    :keyword pd_dir: the proton density weighted scan directory,
        if the |R10| option is not set
    :keyword max_r1_0: the maximum computed |R10| value,
        if the |R10| option is not set
    :keyword in_files: the input images
    :keyword mask_file: the constraining mask file
    :keyword baseline_end_idx: the number of images to merge into a baseline image
    """
    workflow = pe.Workflow(name='pk_mapping')

    # Set up the input node.
    inputs = ['in_files', 'mask_file', 'baseline_end_idx']
    if opts.get(use_fixed_r1_0):
        inputs += ['r1_0_val']
    else:
        inputs += ['pd_dir', 'max_r1_0']
    input_spec = pe.Node(IdentityInterface(fields=inputs),
                        name='input_spec')

    # Merge the DCE data to 4D.
    dce_merge = pe.Node(MergeNifti(), name='dce_merge')
    dce_merge.inputs.out_format = 'dce_series'
    workflow.connect(input_spec, 'in_files', dce_merge, 'in_files')

    # If we are not using a fixed r1_0 value, then compute a map from a
    # proton density weighted scan and the baseline of the DCE series.
    if not use_fixed_r1_0:
        # Convert each series dir to a Nifti.
        pd_stack = pe.Node(DcmStack(), name='pd_stack')
        pd_stack.inputs.embed_meta = True
        workflow.connect(input_spec, 'pd_dir',
                         pd_stack, 'dicom_files')

        # Create the DCE baseline image.
        make_base = pe.Node(Function(input_names=['dce_images',
                                                  'baseline_end_idx'],
                                     output_names=['baseline_nii'],
                                     function=_make_baseline),
                            name='make_base')
        workflow.connect(input_spec, 'in_files',
                         make_base, 'dce_images')
        workflow.connect(input_spec, 'baseline_end_idx',
                         make_base, 'baseline_end_idx')

        # Create the R1_0 map.
        get_r1_0 = pe.Node(Function(input_names=['pdw_image',
                                                 't1w_image',
                                                 'max_r1_0',
                                                 'mask_file',
                                                ],
                                    output_names=['r1_0_map'],
                                    function=_make_r1_0),
                           name='get_r1_0')
        workflow.connect(pd_stack, 'out_file', get_r1_0, 'pdw_image')
        workflow.connect(make_base, 'baseline_nii', get_r1_0, 't1w_image')
        workflow.connect(input_spec, 'max_r1_0', get_r1_0, 'max_r1_0')
        workflow.connect(input_spec, 'mask_file', get_r1_0, 'mask_file')

    # Convert DCE time series to series of R1 maps.
    get_r1_series = pe.Node(Function(input_names=['time_series',
                                                  'r1_0',
                                                  'baseline_end',
                                                  'mask_file',
                                                 ],
                                     output_names=['r1_series'],
                                     function=_make_r1_series),
                            name='get_r1_series')
    workflow.connect(dce_merge, 'out_file',
                     get_r1_series, 'time_series')
    workflow.connect(input_spec, 'baseline_end_idx',
                     get_r1_series, 'baseline_end')
    workflow.connect(input_spec, 'mask_file', get_r1_series, 'mask_file')
    if use_fixed_r1_0:
        workflow.connect(input_spec, 'r1_0_val', get_r1_series, 'r1_0')
    else:
        workflow.connect(get_r1_0, 'r1_0_map', get_r1_series, 'r1_0')

    # Copy meta data to the R1 series.
    copy_meta = pe.Node(CopyMeta(), name='copy_meta')
    copy_meta.inputs.include_classes = [('global', 'const'), ('time', 'samples')]
    workflow.connect(dce_merge, 'out_file', copy_meta, 'src_file')
    workflow.connect(get_r1_series, 'r1_series', copy_meta, 'dest_file')

    # Get the pharmacokinetic mapping parameters.
    get_params = pe.Node(Function(input_names=['time_series'],
                                  output_names=['params_csv'],
                                  function=_get_fit_params),
                         name='get_params')
    workflow.connect(dce_merge, 'out_file', get_params, 'time_series')

    # Perform the pharmacokinetic mapping.
    pk_map = pe.Node(Fastfit(), name='pk_map')
    pk_map.inputs.model_name = 'bolero'
    workflow.connect(copy_meta, 'dest_file',
                     pk_map, 'target_data')
    workflow.connect(input_spec, 'mask_file',
                     pk_map, 'mask')
    workflow.connect(get_params, 'params_csv',
                     pk_map, 'params_csv')

    # Set up the outputs.
    outputs = ['r1_series',
               'params_csv',
               'k_trans',
               'v_e',
               'tau_i'
              ]
    if not use_fixed_r1_0:
        outputs += ['pdw_image', 'dce_baseline', 'r1_0']
    output_spec = pe.Node(IdentityInterface(fields=outputs), name='output_spec')
    workflow.connect(copy_meta, 'dest_file', output_spec, 'r1_series')
    workflow.connect(get_params, 'params_csv', output_spec, 'params_csv')
    workflow.connect(pk_map, 'k_trans', output_spec, 'k_trans')
    workflow.connect(pk_map, 'v_e', output_spec, 'v_e')
    workflow.connect(pk_map, 'tau_i', output_spec, 'tau_i')

    if not use_fixed_r1_0:
        workflow.connect(pd_stack, 'out_file', output_spec, 'pdw_image')
        workflow.connect(make_base, 'baseline_nii',
                         output_spec, 'dce_baseline')
        workflow.connect(get_r1_0, 'r1_0_map', output_spec, 'r1_0')

    # Upload the outputs to XNAT.
    for field in outputs:
        pass # TODO

    return workflow

def _make_baseline(dce_images, baseline_end_idx):
    from nipype.interfaces.dcmstack import MergeNifti
    from nipype.interfaces import fsl

    assert baseline_end_idx > 0
    if baseline_end_idx > 1:
        baseline_images = dce_images[:baseline_end_idx]
        merge_nii = MergeNifti()
        merge_nii.inputs.in_files = baseline_images
        merge_nii.inputs.out_format = 'baseline_merged'
        merged = merge_nii.run().outputs.out_file
        mean_image = fsl.MeanImage(name='mean_image')
        mean_image.inputs.in_file = merged

        return mean_image.run().outputs.out_file
    else:
        return dce_images[0]

def _make_r1_0(pdw_image, t1w_image, max_r1_0, mask_file, **kwargs):
    import os
    from os import path
    import nibabel as nb
    import numpy as np
    from pdw_t1w_to_r1 import pdw_t1w_to_r1
    from dcmstack.dcmmeta import NiftiWrapper

    pdw_nw = NiftiWrapper(nb.load(pdw_image), make_empty=True)
    t1w_nw = NiftiWrapper(nb.load(t1w_image), make_empty=True)
    r1_space = np.arange(0.01, max_r1_0, 0.01)
    mask = nb.load(mask_file).get_data()
    r1_0 = pdw_t1w_to_r1(pdw_nw, t1w_nw, r1_space=r1_space, mask=mask,
                         **kwargs)

    cwd = os.getcwd()
    out_nii = nb.Nifti1Image(r1_0, pdw_nw.nii_img.get_affine())
    out_fn = path.join(cwd, 'r1_0_map.nii.gz')
    nb.save(out_nii, out_fn)
    return out_fn

def _make_r1_series(time_series, r1_0, mask_file, **kwargs):
    import os
    from os import path
    import nibabel as nb
    import numpy as np
    from dce_to_r1 import dce_to_r1
    from dcmstack.dcmmeta import NiftiWrapper

    dce_nw = NiftiWrapper(nb.load(time_series), make_empty=True)
    if not isinstance(r1_0, float):
        r1_0 = nb.load(r1_0).get_data()
    mask = nb.load(mask_file).get_data()
    r1_series = dce_to_r1(dce_nw, r1_0, mask=mask, **kwargs)

    cwd = os.getcwd()
    out_nii = nb.Nifti1Image(r1_series, dce_nw.nii_img.get_affine())
    out_fn = path.join(cwd, 'r1_series.nii.gz')
    nb.save(out_nii, out_fn)
    return out_fn

def _get_fit_params(time_series):
    import os, csv
    from os import path
    import nibabel as nb
    import numpy as np
    from dcmstack.dcmmeta import NiftiWrapper
    from dcmstack import dcm_time_to_sec

    nii = nb.load(time_series)
    data = nii.get_data()
    n_vols = data.shape[-1]
    signal_means = np.array([np.mean(data[:,:,:,idx])
                             for idx in xrange(n_vols)])
    signal_diffs = np.diff(signal_means)

    # If we see a difference in average signal larger than double the .
    # difference from first two points, take that as bolus arrival.
    base_diff = np.abs(signal_diffs[0])
    for idx, diff_val in enumerate(signal_diffs[1:]):
        if diff_val > 2 * base_diff:
            bolus_idx = idx + 1
            break
    else:
        raise ValueError("Unable to determine bolus arrival")

    # Figure out the time in between bolus_idx and bolus_idx+1.
    nw = NiftiWrapper(nii)
    acq_time0 = dcm_time_to_sec(nw.get_meta('AcquisitionTime', (0, 0, 0, 0)))
    acq_time1 = dcm_time_to_sec(nw.get_meta('AcquisitionTime', (0, 0, 0, bolus_idx)))
    acq_time2 = dcm_time_to_sec(nw.get_meta('AcquisitionTime', (0, 0, 0, bolus_idx+1)))
    aif_shift = ((acq_time1 + acq_time2) / 2.0) - acq_time0

    # Create parameter CSV.
    with open('params.csv', 'w') as csv_file:
        csv_writer = csv.writer(csv_file)
        csv_writer.writerow(['aif_params', '0.674', '0.4', '2.2',
                             '0.23', '1.3', '0.09', '0.0013', '0.0'])
        csv_writer.writerow(['aif_delta_t', '1.5'])
        csv_writer.writerow(['aif_shift', str(aif_shift)])
        csv_writer.writerow(['r1_cr', '3.8'])
        csv_writer.writerow(['r1_b_pre', '0.71'])
    return path.join(os.getcwd(), 'params.csv')
<|endoftext|>


<commit_before><commit_msg>Add baseline convnet<commit_after><|endoftext|>


<commit_before>import os
import sys
import time
import logging
import logging.config

import click
import websocket  # Depedency of slackclient, needed for exception handling
from slackclient import SlackClient

def bail(msg_type: str, color: str, text: str) -> str:
    return f"{click.style(msg_type, fg=color)}: {text}"

def setup_logging(verbose: int) -> None:
    logging.config.dictConfig({
        'version': 1,
        'disable_existing_loggers': False,
        'formatters': {
            'standard': {
                'format': "%(asctime)s [%(levelname)s] %(message)s",
                'datefmt': "[%Y-%m-%d %H:%M:%S %z]",
            },
        },
        'handlers': {
            'console': {
                'class': 'logging.StreamHandler',
                'formatter': 'standard',
            },
        },
        'loggers': {
            '': {
                'handlers': ['console'],
                'level': logging.INFO if verbose < 1 else logging.DEBUG,
                'propagate': True,
            },
            'requests.packages.urllib3': {  # Oh, do shut up, requests.
                'handlers': ['console'],
                'level': logging.CRITICAL,
            },
        },
    })

def find_channel_id(channel: str, sc: SlackClient) -> str:
    channels_list = sc.api_call("channels.list").get('channels')
    groups_list = sc.api_call("groups.list").get('groups')

    if not channels_list and not groups_list:
        sys.exit(bail('fatal', 'red', "Couldn't enumerate channels/groups"))

    # Is there a better way to search a list of dictionaries? Probably.
    channel_ids = [c['id'] for c in channels_list + groups_list if c['name'] == channel]

    if not channel_ids:
        sys.exit(bail('fatal', 'red', f"Couldn't find #{channel}"))

    return channel_ids[0]

def handle_message(message: str, channel: str, channel_id: str, sc: SlackClient,
                   logger: logging.Logger) -> None:
    logger.debug(message)

    subtype = message.get('subtype')
    user = message.get('user')

    if subtype in ('group_join', 'channel_join') and user:

        message_channel_id = message.get('channel')
        user_profile = message.get('user_profile')
        username = user_profile.get('name')

        if message_channel_id == channel_id:
            try:
                sc.rtm_send_message(channel, f"Welcome, <@{user}>! :wave:")
                logger.info(f"Welcomed {username} to #{channel}")
            except AttributeError:
                logger.setLevel(logging.ERROR)
                logger.error(f"Couldn't send message to #{channel}")

@click.command()
@click.option('-c', '--channel', envvar='WB2K_CHANNEL', default='general',
              show_default=True, metavar='CHANNEL',
              help='The channel to welcome users to.')
@click.option('-v', '--verbose', count=True, help='It goes to 11.')
@click.option('-r', '--retries', envvar='WB2K_RETRIES', default=8, type=(int), metavar='max_retries',
              help='The maximum number of times to attempt to reconnect on websocket connection errors')
@click.version_option()
def cli(channel: str, verbose: int, retries: int) -> None:

    if verbose > 11:
        sys.exit(bail('fatal', 'red', "It doesn't go beyond 11"))

    # Get our logging in order.
    logger = logging.getLogger(__name__)
    setup_logging(verbose)

    # Make sure we have an API token.
    api_token = os.environ.get('WB2K_TOKEN')
    if not api_token:
        sys.exit(bail('fatal', 'red', 'WB2K_TOKEN envvar undefined'))

    # Instantiate and connect!
    sc = SlackClient(api_token)
    if sc.rtm_connect():
        logger.info("Connected to Slack")

        channel_id = find_channel_id(channel, sc)
        logger.debug(f"Found channel ID {channel_id} for #{channel}")

        logger.info(f"Listening for joins in #{channel}")

        retry_count = 0
        backoff = 0.5

        while True:
            try:
                # Handle dem messages!
                for message in sc.rtm_read():
                    handle_message(message, channel, channel_id, sc, logger)

                # Reset exponential backoff retry strategy every time we
                # successfully loop. Failure would have happened in rtm_read()
                retry_count = 0

                time.sleep(0.5)

            # This is necessary to handle an error caused by a bug in Slack's
            # Python client. For more information see
            # https://github.com/slackhq/python-slackclient/issues/127
            #
            # The TimeoutError could be more elegantly resolved by making a PR
            # to the websocket-client library and letting them coerce that
            # exception to a WebSocketTimeoutException.
            except (websocket.WebSocketConnectionClosedException, TimeoutError):
                logger.error("Lost connection to Slack, reconnecting...")
                if not sc.rtm_connect():
                    logger.info("Failed to reconnect to Slack")
                    if retry_count >= retries:
                        sys.exit(bail(
                            'fatal',
                            'red',
                            "Too many failed reconnect attempts, shutting down")
                        )
                    time.sleep((backoff ** 2) / 4)
                else:
                    logger.info("Reconnected to Slack")

                retry_count += 1

    else:
        sys.exit(bail('fatal', 'red', "Couldn't connect to Slack"))

if __name__ == '__main__':
    cli()
<commit_msg>Pretty format messages for easier debugging<commit_after>import os
import sys
import time
import logging
import logging.config
from pprint import pformat

import click
import websocket  # Depedency of slackclient, needed for exception handling
from slackclient import SlackClient

def bail(msg_type: str, color: str, text: str) -> str:
    return f"{click.style(msg_type, fg=color)}: {text}"

def setup_logging(verbose: int) -> None:
    logging.config.dictConfig({
        'version': 1,
        'disable_existing_loggers': False,
        'formatters': {
            'standard': {
                'format': "%(asctime)s [%(levelname)s] %(message)s",
                'datefmt': "[%Y-%m-%d %H:%M:%S %z]",
            },
        },
        'handlers': {
            'console': {
                'class': 'logging.StreamHandler',
                'formatter': 'standard',
            },
        },
        'loggers': {
            '': {
                'handlers': ['console'],
                'level': logging.INFO if verbose < 1 else logging.DEBUG,
                'propagate': True,
            },
            'requests.packages.urllib3': {  # Oh, do shut up, requests.
                'handlers': ['console'],
                'level': logging.CRITICAL,
            },
        },
    })

def find_channel_id(channel: str, sc: SlackClient) -> str:
    channels_list = sc.api_call("channels.list").get('channels')
    groups_list = sc.api_call("groups.list").get('groups')

    if not channels_list and not groups_list:
        sys.exit(bail('fatal', 'red', "Couldn't enumerate channels/groups"))

    # Is there a better way to search a list of dictionaries? Probably.
    channel_ids = [c['id'] for c in channels_list + groups_list if c['name'] == channel]

    if not channel_ids:
        sys.exit(bail('fatal', 'red', f"Couldn't find #{channel}"))

    return channel_ids[0]

def handle_message(message: str, channel: str, channel_id: str, sc: SlackClient,
                   logger: logging.Logger) -> None:
    pretty_message = pformat(message)
    logger.debug(f"New message received:\n{pretty_message}")

    subtype = message.get('subtype')
    user = message.get('user')

    if subtype in ('group_join', 'channel_join') and user:

        message_channel_id = message.get('channel')
        user_profile = message.get('user_profile')
        username = user_profile.get('name')

        if message_channel_id == channel_id:
            try:
                sc.rtm_send_message(channel, f"Welcome, <@{user}>! :wave:")
                logger.info(f"Welcomed {username} to #{channel}")
            except AttributeError:
                logger.setLevel(logging.ERROR)
                logger.error(f"Couldn't send message to #{channel}")

@click.command()
@click.option('-c', '--channel', envvar='WB2K_CHANNEL', default='general',
              show_default=True, metavar='CHANNEL',
              help='The channel to welcome users to.')
@click.option('-v', '--verbose', count=True, help='It goes to 11.')
@click.option('-r', '--retries', envvar='WB2K_RETRIES', default=8, type=(int), metavar='max_retries',
              help='The maximum number of times to attempt to reconnect on websocket connection errors')
@click.version_option()
def cli(channel: str, verbose: int, retries: int) -> None:

    if verbose > 11:
        sys.exit(bail('fatal', 'red', "It doesn't go beyond 11"))

    # Get our logging in order.
    logger = logging.getLogger(__name__)
    setup_logging(verbose)

    # Make sure we have an API token.
    api_token = os.environ.get('WB2K_TOKEN')
    if not api_token:
        sys.exit(bail('fatal', 'red', 'WB2K_TOKEN envvar undefined'))

    # Instantiate and connect!
    sc = SlackClient(api_token)
    if sc.rtm_connect():
        logger.info("Connected to Slack")

        channel_id = find_channel_id(channel, sc)
        logger.debug(f"Found channel ID {channel_id} for #{channel}")

        logger.info(f"Listening for joins in #{channel}")

        retry_count = 0
        backoff = 0.5

        while True:
            try:
                # Handle dem messages!
                for message in sc.rtm_read():
                    handle_message(message, channel, channel_id, sc, logger)

                # Reset exponential backoff retry strategy every time we
                # successfully loop. Failure would have happened in rtm_read()
                retry_count = 0

                time.sleep(0.5)

            # This is necessary to handle an error caused by a bug in Slack's
            # Python client. For more information see
            # https://github.com/slackhq/python-slackclient/issues/127
            #
            # The TimeoutError could be more elegantly resolved by making a PR
            # to the websocket-client library and letting them coerce that
            # exception to a WebSocketTimeoutException.
            except (websocket.WebSocketConnectionClosedException, TimeoutError):
                logger.error("Lost connection to Slack, reconnecting...")
                if not sc.rtm_connect():
                    logger.info("Failed to reconnect to Slack")
                    if retry_count >= retries:
                        sys.exit(bail(
                            'fatal',
                            'red',
                            "Too many failed reconnect attempts, shutting down")
                        )
                    time.sleep((backoff ** 2) / 4)
                else:
                    logger.info("Reconnected to Slack")

                retry_count += 1

    else:
        sys.exit(bail('fatal', 'red', "Couldn't connect to Slack"))

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


<commit_before>#!/usr/bin/python2.7

# Copyright 2011 Ajay Roopakalu
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import twitter
import datetime
import rfc822
import ConfigParser
import os

class TweetRc():
	def __init__(self):
		self._config = None

	def GetConsumerKey(self):
		return self._GetOption('consumer_key')

	def GetConsumerSecret(self):
		return self._GetOption('consumer_secret')

	def GetAccessKey(self):
		return self._GetOption('access_token_key')

	def GetAccessSecret(self):
		return self._GetOption('access_token_secret')

	def _GetOption(self, option):
		try:
			return self._GetConfig().get('CommanderTweet', option)
		except:
			return None

	def _GetConfig(self):
		if not self._config:
		 	self._config = ConfigParser.ConfigParser()
		 	self._config.read(os.path.expanduser('~/.ctweetrc'))
		return self._config

def bold(x):
	return '\x1b[1m' + x + '\x1b[0m'

def prettydate(d):
	diff = datetime.datetime.utcnow() - d
	s = diff.seconds

	if diff.days > 7 or diff.days < 0:
		return d.strftime('%d %b %y')
	elif diff.days == 1:
		return '1 day ago'
	elif diff.days > 1:
		return '{} days ago'.format(diff.days)
	elif s <= 1:
		return 'just now'
	elif s < 60:
		return '{} seconds ago'.format(s)
	elif s < 120:
		return '1 minute ago'
	elif s < 3600:
		return '{} minutes ago'.format(s/60)
	elif s < 7200:
		return '1 hour ago'
	else:
		return '{} hours ago'.format(s/3600)

def main():
	rc = TweetRc()
	api = twitter.Api(consumer_key=rc.GetConsumerKey(),
					  consumer_secret=rc.GetConsumerSecret(),
					  access_token_key=rc.GetAccessKey(),
					  access_token_secret=rc.GetAccessSecret())
	home_timeline = api.GetHomeTimeline()

	for tweet in home_timeline:
		created_at = datetime.datetime(*rfc822.parsedate(tweet.created_at)[:-2])
		print bold(tweet.user.name), '@{0}'.format(tweet.user.screen_name),
		print '({0})'.format(prettydate(created_at))
		print '  ', tweet.text

if __name__ == '__main__':
	main()
<commit_msg>Changed now reflected in main.py as well.<commit_after>#!/usr/bin/python2.7

# Copyright 2011 Ajay Roopakalu
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import tweepy
import ConfigParser
import os
import time
import utils
import locale

from threading import Thread

class TweetStream():
	def __init__(self, tweets):
		self.ut = utils.Utils()
		self.tweets = tweets
		locale.setlocale(locale.LC_ALL, '')
		self.code = locale.getpreferredencoding()

	def disp(self):
		for tweet in self.tweets:
			print self.ut.bold(tweet.user.name.encode(self.code)),
			print '@{}'.format(tweet.user.screen_name.encode(self.code)),
			print '({})'.format(self.ut.human_date(tweet.created_at))
			print '  {}'.format(tweet.text.encode(self.code))

	def update(self, tweet):
		self.tweets.insert(0, tweet)
		self.tweets = self.tweets[:-1]
		self.disp()

class TweetRc():
	def __init__(self):
		self._config = None

	def GetConsumerKey(self):
		return self._GetOption('consumer_key')

	def GetConsumerSecret(self):
		return self._GetOption('consumer_secret')

	def GetAccessKey(self):
		return self._GetOption('access_token_key')

	def GetAccessSecret(self):
		return self._GetOption('access_token_secret')

	def _GetOption(self, option):
		try:
			return self._GetConfig().get('CommanderTweet', option)
		except:
			return None

	def _GetConfig(self):
		if not self._config:
		 	self._config = ConfigParser.ConfigParser()
		 	self._config.read(os.path.expanduser('~/.ctweetrc'))
		return self._config

class StreamListener(tweepy.StreamListener):
	def on_status(self, status):
		try:
			TS.update(status)
		except Exception, e:
			print e

def main():
	global TS
	rc = TweetRc()

	auth = tweepy.OAuthHandler(rc.GetConsumerKey(), rc.GetConsumerSecret())
	auth.set_access_token(rc.GetAccessKey(), rc.GetAccessSecret())

	api = tweepy.API(auth)
	os.system('clear')
	TS = TweetStream(api.home_timeline())
	TS.disp()

	stream = StreamListener()
	streamer = tweepy.Stream(auth, stream, secure=True)
	streamer.userstream()

if __name__ == '__main__':
	TS = None
	try:
		main()
	except KeyboardInterrupt:
		pass

<|endoftext|>