<commit_before>#!/usr/bin/env python
from __future__ import print_function

import sys,json
import html2text
from bs4 import BeautifulSoup
import re
from collections import OrderedDict
import math
import ipdb

def info(type, value, tb):
	ipdb.pm()

sys.excepthook = info

text_maker = html2text.HTML2Text()
text_maker.body_width = 0 #disable random line-wrapping from html2text

levels_count = dict()
nodes_lookup = dict()
fixups_queue = list()
objective_node = None

def get_node_children(node):
	return OrderedDict(sorted(node.get('ideas', dict()).iteritems(), key=lambda t: float(t[0]))).values()

def is_node_a_leaf(node):
	return len(get_node_children(node)) == 0

def get_raw_description(node):
	#prefer the mindmup 2.0 'note' to the 1.0 'attachment'
	description = node.get('attr', dict()).get('note', dict()).get('text', '')
	if description is '':
		description = node.get('attr', dict()).get('attachment', dict()).get('content', '')

	return description

def detect_html(text):
	return bool(BeautifulSoup(text, "html.parser").find())

def get_description(node):
	global text_maker

	description = get_raw_description(node)

	if detect_html(description):
		description = text_maker.handle(description)

	return description

def parse_evita_raps(node):
	if not 'EVITA::' in get_raw_description(node):
		raise ValueError("couldn't find EVITA:: tag in leaf node", node)

	for line in get_description(node).splitlines():
		if not 'EVITA::' in line:
			continue

		evita_line = line.strip().split('|')

		if node.get('attr', None) is None:
		    node.update({'attr': dict()})

		attr = node.get('attr')

		attr.update({'evita_et': float(evita_line[5])})
		attr.update({'evita_e': float(evita_line[6])})
		attr.update({'evita_k': float(evita_line[7])})
		attr.update({'evita_wo': float(evita_line[8])})
		attr.update({'evita_eq': float(evita_line[9])})

	return

def derive_evita_apt(node):
	attrs = node.get('attr')

	total_rap = attrs.get('evita_et') + attrs.get('evita_e') + attrs.get('evita_k') + attrs.get('evita_wo') + attrs.get('evita_eq')
	if total_rap < 0:
		raise ValueError('encountered negative Total Required Attack Potential', node.get('attr'))
	elif total_rap < 10:
		apt = 5
	elif total_rap < 14:
		apt = 4
	elif total_rap < 20:
		apt = 3
	elif total_rap < 25:
		apt = 2
	else:
		apt = 1
	#TODO support non-zero controllability

	attrs.update({'evita_apt': apt})
	return

def parse_evita_severities(node):
	if not 'EVITA::' in get_raw_description(node):
		raise ValueError("couldn't find EVITA:: tag in leaf node", node)

	for line in get_description(node).splitlines():
		if not 'EVITA::' in line:
			continue

		evita_line = line.strip().split('|')
		attr = node.get('attr')

		attr.update({'evita_fs': float(evita_line[1])})
		attr.update({'evita_os': float(evita_line[2])})
		attr.update({'evita_ps': float(evita_line[3])})
		attr.update({'evita_ss': float(evita_line[4])})

	return

evita_security_risk_table=[
	[0,0,0,0,0],
	[0,0,1,2,3],
	[0,1,2,3,4],
	[1,2,3,4,5],
	[2,3,4,5,6]
]
def get_evita_security_risk_level(non_safety_severity, combined_attack_probability):
	global evita_security_risk_table

	if non_safety_severity < 0 or non_safety_severity > 4:
		raise ValueError('encountered an invalid non-safety severity', non_safety_severity)
	return evita_security_risk_table[int(non_safety_severity)][int(combined_attack_probability)-1]

def derive_evita_risks(this_node, objective_node):
	these_attrs = this_node.get('attr')
	objective_attrs = objective_node.get('attr')

	these_attrs.update({'evita_fr': get_evita_security_risk_level(objective_attrs.get('evita_fs'), these_attrs.get('evita_apt'))})
	these_attrs.update({'evita_or': get_evita_security_risk_level(objective_attrs.get('evita_os'), these_attrs.get('evita_apt'))})
	these_attrs.update({'evita_pr': get_evita_security_risk_level(objective_attrs.get('evita_ps'), these_attrs.get('evita_apt'))})
	these_attrs.update({'evita_sr': get_evita_security_risk_level(objective_attrs.get('evita_ss'), these_attrs.get('evita_apt'))})
	return

def is_objective(node):
	raw_description = get_raw_description(node)
	return 'OBJECTIVE::' in raw_description

def is_riskpoint(node):
	raw_description = get_raw_description(node)
	return 'RISK_HERE::' in raw_description

def do_children_firstpass(node):
	for child in get_node_children(node):
		do_node_firstpass(child)
	return

def do_node_firstpass(node):
	global nodes_lookup

	do_children_firstpass(node)

	node_title = node.get('title', '')

	if node_title == 'AND':
		return

	if node_title == '...':
		return

	if is_node_a_reference(node):
		return

	if nodes_lookup.get(node_title, None) is None:
		nodes_lookup.update({node_title: node})

	return

def get_node_referent_title(node):
	title = node.get('title', '')

	if '(*)' in node.get('title'):
		wip_referent_title = title.replace('(*)','').strip()
	else:
		referent_coords = re.search(r'\((\d+\..*?)\)', title).groups()[0]
		wip_referent_title = "%s %s" % (referent_coords, re.sub(r'\(\d+\..*?\)', '', title).strip())
	return wip_referent_title

def is_node_a_reference(node):
	title = node.get('title', '')

	return (not title.find('(*)') == -1) or (not re.search(r'\(\d+\..*?\)', title) is None)

def is_node_weigthed(node):
	apt = get_node_apt(node)
	return (not apt is None) and (not math.isnan(apt)) and (not math.isinf(apt))

def update_node_apt(node, apt):
	if node.get('attr', None) is None:
		node.update({'attr': dict()})

	node.get('attr').update({'evita_apt': apt})

def pos_infs_of_children(node):
	for child in get_node_children(node):
		if get_node_apt(child) == float('-inf'):
			update_node_apt(child, float('inf'))

def neg_infs_of_children(node):
	for child in get_node_children(node):
		if get_node_apt(child) == float('inf'):
			update_node_apt(child, float('-inf'))

def get_max_apt_of_children(node):
	child_maximum = float('-inf')

	for child in get_node_children(node):
		child_maximum = max(child_maximum, get_node_apt(child))

	return child_maximum

def get_min_apt_of_children(node):
	child_minimum = float('inf')

	for child in get_node_children(node):
		child_minimum = min(child_minimum, get_node_apt(child))

	return child_minimum

def get_node_apt(node):
	return node.get('attr', dict()).get('evita_apt', None)

def get_node_referent(node, nodes_lookup):
	node_referent_title = get_node_referent_title(node)
	node_referent = nodes_lookup.get(node_referent_title, None)

	if node_referent is None:
		raise ValueError("ERROR missing node referent: %s" % node_referent_title)
		return node
	else:
		return node_referent

def get_node_title(node):
	return node.get('title', '')

def do_children_secondpass(node, nodes_context):
	for child in get_node_children(node):
		do_node_secondpass(child, nodes_context)
	return

def do_node_secondpass(node, nodes_context):
	global nodes_lookup
	global fixups_queue

	if is_node_weigthed(node):
		return

	update_node_apt(node, float('nan'))

	if is_node_a_reference(node):
		node_referent = get_node_referent(node, nodes_lookup)
		node_referent_title=get_node_title(node_referent)

		if (not get_node_apt(node_referent) is None) and (math.isnan(get_node_apt(node_referent))):
			#is referent in-progress? then we have a loop. update the reference node with the identity of the tree reduction operation and return
			update_node_apt(node, float('-inf'))
		else:
			#otherwise, descend through referent's children

			#do all on the referent and copy the node apt back
			do_node_secondpass(node_referent, nodes_context)

			update_node_apt(node,get_node_apt(node_referent))
	else:
		if is_node_a_leaf(node):
			parse_evita_raps(node)
			derive_evita_apt(node)

		else:
			nodes_context.append(get_node_title(node))
			do_children_secondpass(node, nodes_context)
			nodes_context.pop()

			if node.get('title', None) == 'AND':
				pos_infs_of_children(node)
				update_node_apt(node, get_min_apt_of_children(node))
			else:
				neg_infs_of_children(node)
				update_node_apt(node, get_max_apt_of_children(node))

	if math.isinf(get_node_apt(node)):
		fixups_queue.append(node)
	return

def do_fixups(nodes_context):
    global fixups_queue
    fixups_len = len(fixups_queue)

    while len(fixups_queue) > 0:
	fixups_this_time = list(fixups_queue)
	fixups_queue = list()
	for node in fixups_this_time:
		do_node_secondpass(node, nodes_context)

	if len(fixups_queue) >= fixups_len:
		raise ValueError("ERROR couldn't resolve remaining infs %s" % fixups_queue)
		break
	else:
		fixups_len = len(fixups_queue)

def do_children_riskspass(node, nodes_context):
	for child in get_node_children(node):
		do_node_riskspass(child, nodes_context)
	return

def do_node_riskspass(node, nodes_context):
	global nodes_lookup
	global objective_node

	saved_objective = objective_node
	if is_objective(node):
		parse_evita_severities(node)
		objective_node = node

	if is_riskpoint(node):
		derive_evita_risks(node, objective_node)
		return

	do_children_riskspass(node, nodes_context)
	objective_node = saved_objective

if len(sys.argv) < 2:
	fd_in=sys.stdin
else:
	fd_in=open(sys.argv[1], 'r')

data = json.load(fd_in)

if len(sys.argv) < 2:
	fd_out = sys.stdout
else:
	fd_in.close()
	fd_out=open(sys.argv[1],'w')

nodes_context=list()

if 'id' in data and data['id'] == 'root':
	#version 2 mindmup
	root_node = data['ideas']['1']
else:
	root_node = data

do_children_firstpass(root_node)
do_node_secondpass(root_node, nodes_context)
do_fixups(nodes_context)
do_node_riskspass(root_node, nodes_context)

str = json.dumps(data, indent=2, sort_keys=True)
str = re.sub(r'\s+$', '', str, 0, re.M)
str = re.sub(r'\s+$', '', str, flags=re.M)

fd_out.write(str)

if len(sys.argv) >= 1:
	fd_out.close()

<commit_msg>clamp the values that get serialized in the JSON -- infinities are illegal there<commit_after>#!/usr/bin/env python
from __future__ import print_function

import sys,json
import html2text
from bs4 import BeautifulSoup
import re
from collections import OrderedDict
import math
import ipdb

def info(type, value, tb):
	ipdb.pm()

sys.excepthook = info

text_maker = html2text.HTML2Text()
text_maker.body_width = 0 #disable random line-wrapping from html2text

levels_count = dict()
nodes_lookup = dict()
fixups_queue = list()
objective_node = None

def get_node_children(node):
	return OrderedDict(sorted(node.get('ideas', dict()).iteritems(), key=lambda t: float(t[0]))).values()

def is_node_a_leaf(node):
	return len(get_node_children(node)) == 0

def get_raw_description(node):
	#prefer the mindmup 2.0 'note' to the 1.0 'attachment'
	description = node.get('attr', dict()).get('note', dict()).get('text', '')
	if description is '':
		description = node.get('attr', dict()).get('attachment', dict()).get('content', '')

	return description

def detect_html(text):
	return bool(BeautifulSoup(text, "html.parser").find())

def get_description(node):
	global text_maker

	description = get_raw_description(node)

	if detect_html(description):
		description = text_maker.handle(description)

	return description

def clamp_to_json_values(val):
    return max(-1 * sys.float_info.max, min(val, sys.float_info.max))

def parse_evita_raps(node):
	if not 'EVITA::' in get_raw_description(node):
		raise ValueError("couldn't find EVITA:: tag in leaf node", node)

	for line in get_description(node).splitlines():
		if not 'EVITA::' in line:
			continue

		evita_line = line.strip().split('|')

		if node.get('attr', None) is None:
		    node.update({'attr': dict()})

		attr = node.get('attr')

		attr.update({'evita_et': clamp_to_json_values(float(evita_line[5]))})
		attr.update({'evita_e':  clamp_to_json_values(float(evita_line[6]))})
		attr.update({'evita_k':  clamp_to_json_values(float(evita_line[7]))})
		attr.update({'evita_wo': clamp_to_json_values(float(evita_line[8]))})
		attr.update({'evita_eq': clamp_to_json_values(float(evita_line[9]))})

	return

def derive_evita_apt(node):
	attrs = node.get('attr')

	total_rap = attrs.get('evita_et') + attrs.get('evita_e') + attrs.get('evita_k') + attrs.get('evita_wo') + attrs.get('evita_eq')
	if total_rap < 0:
		raise ValueError('encountered negative Total Required Attack Potential', node.get('attr'))
	elif total_rap < 10:
		apt = 5
	elif total_rap < 14:
		apt = 4
	elif total_rap < 20:
		apt = 3
	elif total_rap < 25:
		apt = 2
	else:
		apt = 1
	#TODO support non-zero controllability

	attrs.update({'evita_apt': apt})
	return

def parse_evita_severities(node):
	if not 'EVITA::' in get_raw_description(node):
		raise ValueError("couldn't find EVITA:: tag in leaf node", node)

	for line in get_description(node).splitlines():
		if not 'EVITA::' in line:
			continue

		evita_line = line.strip().split('|')
		attr = node.get('attr')

		attr.update({'evita_fs': clamp_to_json_values(float(evita_line[1]))})
		attr.update({'evita_os': clamp_to_json_values(float(evita_line[2]))})
		attr.update({'evita_ps': clamp_to_json_values(float(evita_line[3]))})
		attr.update({'evita_ss': clamp_to_json_values(float(evita_line[4]))})

	return

evita_security_risk_table=[
	[0,0,0,0,0],
	[0,0,1,2,3],
	[0,1,2,3,4],
	[1,2,3,4,5],
	[2,3,4,5,6]
]
def get_evita_security_risk_level(non_safety_severity, combined_attack_probability):
	global evita_security_risk_table

	if non_safety_severity < 0 or non_safety_severity > 4:
		raise ValueError('encountered an invalid non-safety severity', non_safety_severity)
	return evita_security_risk_table[int(non_safety_severity)][int(combined_attack_probability)-1]

def derive_evita_risks(this_node, objective_node):
	these_attrs = this_node.get('attr')
	objective_attrs = objective_node.get('attr')

	these_attrs.update({'evita_fr': get_evita_security_risk_level(objective_attrs.get('evita_fs'), these_attrs.get('evita_apt'))})
	these_attrs.update({'evita_or': get_evita_security_risk_level(objective_attrs.get('evita_os'), these_attrs.get('evita_apt'))})
	these_attrs.update({'evita_pr': get_evita_security_risk_level(objective_attrs.get('evita_ps'), these_attrs.get('evita_apt'))})
	these_attrs.update({'evita_sr': get_evita_security_risk_level(objective_attrs.get('evita_ss'), these_attrs.get('evita_apt'))})
	return

def is_objective(node):
	raw_description = get_raw_description(node)
	return 'OBJECTIVE::' in raw_description

def is_riskpoint(node):
	raw_description = get_raw_description(node)
	return 'RISK_HERE::' in raw_description

def do_children_firstpass(node):
	for child in get_node_children(node):
		do_node_firstpass(child)
	return

def do_node_firstpass(node):
	global nodes_lookup

	do_children_firstpass(node)

	node_title = node.get('title', '')

	if node_title == 'AND':
		return

	if node_title == '...':
		return

	if is_node_a_reference(node):
		return

	if nodes_lookup.get(node_title, None) is None:
		nodes_lookup.update({node_title: node})

	return

def get_node_referent_title(node):
	title = node.get('title', '')

	if '(*)' in node.get('title'):
		wip_referent_title = title.replace('(*)','').strip()
	else:
		referent_coords = re.search(r'\((\d+\..*?)\)', title).groups()[0]
		wip_referent_title = "%s %s" % (referent_coords, re.sub(r'\(\d+\..*?\)', '', title).strip())
	return wip_referent_title

def is_node_a_reference(node):
	title = node.get('title', '')

	return (not title.find('(*)') == -1) or (not re.search(r'\(\d+\..*?\)', title) is None)

def is_node_weigthed(node):
	apt = get_node_apt(node)
	return (not apt is None) and (not math.isnan(apt)) and (not math.isinf(apt))

def update_node_apt(node, apt):
	if node.get('attr', None) is None:
		node.update({'attr': dict()})

	node.get('attr').update({'evita_apt': apt})

def pos_infs_of_children(node):
	for child in get_node_children(node):
		if get_node_apt(child) == float('-inf'):
			update_node_apt(child, float('inf'))

def neg_infs_of_children(node):
	for child in get_node_children(node):
		if get_node_apt(child) == float('inf'):
			update_node_apt(child, float('-inf'))

def get_max_apt_of_children(node):
	child_maximum = float('-inf')

	for child in get_node_children(node):
		child_maximum = max(child_maximum, get_node_apt(child))

	return child_maximum

def get_min_apt_of_children(node):
	child_minimum = float('inf')

	for child in get_node_children(node):
		child_minimum = min(child_minimum, get_node_apt(child))

	return child_minimum

def get_node_apt(node):
	return node.get('attr', dict()).get('evita_apt', None)

def get_node_referent(node, nodes_lookup):
	node_referent_title = get_node_referent_title(node)
	node_referent = nodes_lookup.get(node_referent_title, None)

	if node_referent is None:
		raise ValueError("ERROR missing node referent: %s" % node_referent_title)
		return node
	else:
		return node_referent

def get_node_title(node):
	return node.get('title', '')

def do_children_secondpass(node, nodes_context):
	for child in get_node_children(node):
		do_node_secondpass(child, nodes_context)
	return

def do_node_secondpass(node, nodes_context):
	global nodes_lookup
	global fixups_queue

	if is_node_weigthed(node):
		return

	update_node_apt(node, float('nan'))

	if is_node_a_reference(node):
		node_referent = get_node_referent(node, nodes_lookup)
		node_referent_title=get_node_title(node_referent)

		if (not get_node_apt(node_referent) is None) and (math.isnan(get_node_apt(node_referent))):
			#is referent in-progress? then we have a loop. update the reference node with the identity of the tree reduction operation and return
			update_node_apt(node, float('-inf'))
		else:
			#otherwise, descend through referent's children

			#do all on the referent and copy the node apt back
			do_node_secondpass(node_referent, nodes_context)

			update_node_apt(node,get_node_apt(node_referent))
	else:
		if is_node_a_leaf(node):
			parse_evita_raps(node)
			derive_evita_apt(node)

		else:
			nodes_context.append(get_node_title(node))
			do_children_secondpass(node, nodes_context)
			nodes_context.pop()

			if node.get('title', None) == 'AND':
				pos_infs_of_children(node)
				update_node_apt(node, get_min_apt_of_children(node))
			else:
				neg_infs_of_children(node)
				update_node_apt(node, get_max_apt_of_children(node))

	if math.isinf(get_node_apt(node)):
		fixups_queue.append(node)
	return

def do_fixups(nodes_context):
    global fixups_queue
    fixups_len = len(fixups_queue)

    while len(fixups_queue) > 0:
	fixups_this_time = list(fixups_queue)
	fixups_queue = list()
	for node in fixups_this_time:
		do_node_secondpass(node, nodes_context)

	if len(fixups_queue) >= fixups_len:
		raise ValueError("ERROR couldn't resolve remaining infs %s" % fixups_queue)
		break
	else:
		fixups_len = len(fixups_queue)

def do_children_riskspass(node, nodes_context):
	for child in get_node_children(node):
		do_node_riskspass(child, nodes_context)
	return

def do_node_riskspass(node, nodes_context):
	global nodes_lookup
	global objective_node

	saved_objective = objective_node
	if is_objective(node):
		parse_evita_severities(node)
		objective_node = node

	if is_riskpoint(node):
		derive_evita_risks(node, objective_node)
		return

	do_children_riskspass(node, nodes_context)
	objective_node = saved_objective

if len(sys.argv) < 2:
	fd_in=sys.stdin
else:
	fd_in=open(sys.argv[1], 'r')

data = json.load(fd_in)

if len(sys.argv) < 2:
	fd_out = sys.stdout
else:
	fd_in.close()
	fd_out=open(sys.argv[1],'w')

nodes_context=list()

if 'id' in data and data['id'] == 'root':
	#version 2 mindmup
	root_node = data['ideas']['1']
else:
	root_node = data

do_children_firstpass(root_node)
do_node_secondpass(root_node, nodes_context)
do_fixups(nodes_context)
do_node_riskspass(root_node, nodes_context)

str = json.dumps(data, indent=2, sort_keys=True)
str = re.sub(r'\s+$', '', str, 0, re.M)
str = re.sub(r'\s+$', '', str, flags=re.M)

fd_out.write(str)

if len(sys.argv) >= 1:
	fd_out.close()

<|endoftext|>


<commit_before>import os
import shutil
import textwrap
from collections import defaultdict
from multiprocessing import cpu_count
from subprocess import call, Popen, PIPE

import numpy as np
import pandas as pd
from scipy.stats import pearsonr

from rnaseq_lib.docker import fix_permissions, get_base_call
from rnaseq_lib.tissues import map_genes
from rnaseq_lib.utils import mkdir_p

def log2fc(a, b, pad=0.001):
    """
    Calculate the log2 Fold Change between two arrays, floats, or integers
    a and b cannot be, nor contain, values less than 0

    :param int|float|np.array a: Value or array
    :param int|float|np.array b: Value or array
    :param int|float pad: Buffer to add to value before log2 calculation
    :return: L2FC array or value
    :rtype: int|float|np.array
    """
    return np.log2(a + pad) - np.log2(b + pad)

def de_pearson_dataframe(df, genes, pair_by='type', gtex=True, tcga=True):
    """
    PearsonR scores of gene differential expression between tumor and normal types.

    1. Calculate log2FC of genes for TCGA tumor samples with matching TCGA normal types
    2. Compare log2fc to tumor type compared to all other normal types
    3. Calculate PearsonR and save

    :param pd.DataFrame df: Exp/TPM dataframe containing "type"/"tissue/tumor/label" metadata columns
    :param list genes: Genes to use in differential expression calculation
    :param str pair_by: How to pair tumors/normals. Either by "type" or "tissue"
    :param bool gtex: If True, includes GTEx in normal set
    :param bool tcga: If True, includes TCGA in normal set
    :return: PearsonR dataframe
    :rtype: pd.DataFrame
    """
    # Subset by Tumor/Normal
    tumor = df[df.label == 'tcga-tumor']
    tcga_n = df[df.label == 'tcga-normal']
    if gtex:
        normal = df[df.tumor == 'no']
    else:
        normal = tcga_n

    # Identify tumor types with paired tcga-normal
    tum_types = [x for x in sorted(tumor[pair_by].unique())
                 if x in sorted(df[df.label == 'tcga-normal'][pair_by].unique())]
    norm_types = []

    # For all paired tumor_types, calculate l2fc, then PearsonR of l2fc to all normal tumor types
    pearson_l2fc = defaultdict(list)
    for tum_type in tum_types:

        # First calculate TCGA tumor/normal prior for comparison
        t = tumor[tumor[pair_by] == tum_type]
        t_med = t[genes].median()
        n = tcga_n[tcga_n[pair_by] == tum_type]
        prior_l2fc = log2fc(t_med, n[genes].median())

        # For every normal type, calculate pearsonR correlation
        for (norm_type, label), _ in normal.groupby(pair_by).label.value_counts().iteritems():
            if tum_type == norm_type:
                l2fc = prior_l2fc
            else:
                n = normal[normal[pair_by] == norm_type]
                l2fc = log2fc(t_med, n[genes].median())

            # Calculate PearsonR Save l2fc and comparison tissue/type
            pearson_r = round(pearsonr(prior_l2fc, l2fc)[0], 2)
            pearson_l2fc[tum_type[:20]].append(pearson_r)
            norm_label = '{}_{}'.format(label, norm_type[:20])
            if norm_label not in norm_types:
                norm_types.append(norm_label)

    return pd.DataFrame(pearson_l2fc, index=norm_types)

def run_deseq2(df_path, tissue, output_dir, gtex=True, cores=None):
    """
    Runs DESeq2 on a specific tissue

    :param str df_path: Path to samples by genes dataframe
    :param str tissue: Tissue to run
    :param str output_dir: Full path to output directory
    :param bool gtex: If True uses GTEx as normal tissue. Otherwise uses TCGA Normal
    :param int cores: Number of cores to use. Defaults to # of cores on machine.
    """
    # Make workspace directories
    work_dir = os.path.join(output_dir, 'work_dir')
    mkdir_p(work_dir)

    # Get samples for tissue
    tumor = get_tumor_samples(tissue)
    normal = get_gtex_samples(tissue) if gtex else get_normal_samples(tissue)

    # Write out vectors
    tissue_vector = os.path.join(work_dir, 'tissue.vector')
    with open(tissue_vector, 'w') as f:
        f.write('\n'.join(tumor + normal))

    disease_vector = os.path.join(work_dir, 'disease.vector')
    with open(disease_vector, 'w') as f:
        f.write('\n'.join(['T' if x in tumor else 'N' for x in tumor + normal]))

    # Write out script
    cores = cores if cores else int(cpu_count())
    script_path = os.path.join(work_dir, 'deseq2.R')
    with open(script_path, 'w') as f:
        f.write(
            textwrap.dedent("""
            library('DESeq2'); library('data.table'); library('BiocParallel')
            register(MulticoreParam({cores}))

            # Argument parsing
            args <- commandArgs(trailingOnly = TRUE)
            df_path <- args[1]
            tissue_path <- args[2]
            disease_path <- args[3]
            tissue <- '{tissue}'
            output_dir <- '/data/'

            # Read in vectors
            tissue_vector <- read.table(tissue_path)$V1
            disease_vector <- read.table(disease_path)$V1

            # Read in table and process
            n <- read.table(df_path, sep='\\t', header=1, row.names=1, check.names=FALSE)
            sub <- n[, colnames(n)%in%tissue_vector]
            setcolorder(sub, as.character(tissue_vector))

            # Preprocessing
            countData <- round(sub)
            colData <- data.frame(disease=disease_vector, row.names=colnames(countData))
            y <- DESeqDataSetFromMatrix(countData = countData, colData = colData, design = ~ disease)

            # Run DESeq2
            y <- DESeq(y, parallel=TRUE)
            res <- results(y, parallel=TRUE)
            summary(res)

            # Write out table
            resOrdered <- res[order(res$padj),]
            res_name <- paste(tissue, '.tsv', sep='')
            res_path <- paste(output_dir, res_name, sep='/')
            write.table(as.data.frame(resOrdered), file=res_path, col.names=NA, sep='\\t',  quote=FALSE)

            # MA Plot
            ma_name <- paste(tissue, '-MA.pdf', sep='')
            ma_path <- paste(output_dir, ma_name, sep='/')
            pdf(ma_path, width=7, height=7)
            plotMA(res, main='DESeq2')
            dev.off()

            # Dispersion Plot
            disp_name <- paste(tissue, '-dispersion.pdf', sep='')
            disp_path <- paste(output_dir, disp_name, sep='/')
            pdf(disp_path, width=7, height=7)
            plotDispEsts( y, ylim = c(1e-6, 1e1) )
            dev.off()

            # PVal Hist
            hist_name <- paste(tissue, '-pval-hist.pdf', sep='')
            hist_path <- paste(output_dir, hist_name, sep='/')
            pdf(hist_path, width=7, height=7)
            hist( res$pvalue, breaks=20, col="grey" )
            dev.off()

            # Ratios plots
            qs <- c( 0, quantile( res$baseMean[res$baseMean > 0], 0:7/7 ) )
            bins <- cut( res$baseMean, qs )
            levels(bins) <- paste0("~",round(.5*qs[-1] + .5*qs[-length(qs)]))
            ratios <- tapply( res$pvalue, bins, function(p) mean( p < .01, na.rm=TRUE ) )
            ratio_name <- paste(tissue, '-ratios.pdf', sep='')
            ratio_path <- paste(output_dir, ratio_name, sep='/')
            pdf(ratio_path, width=7, height=7)
            barplot(ratios, xlab="mean normalized count", ylab="ratio of small $p$ values")
            dev.off()
            """.format(cores=cores, tissue=tissue)))

    # Call DESeq2
    docker_parameters = ['docker', 'run',
                         '-v', '{}:/data'.format(output_dir),
                         '-v', '{}:/df'.format(os.path.dirname(df_path)),
                         'jvivian/deseq2']

    parameters = ['/data/work_dir/deseq2.R',
                  '/df/{}'.format(os.path.basename(df_path)),
                  '/data/{}'.format(os.path.join('work_dir', 'tissue.vector')),
                  '/data/{}'.format(os.path.join('work_dir', 'disease.vector'))]

    print '\nCalling: {}\n'.format(' '.join(docker_parameters + parameters))
    p = Popen(docker_parameters + parameters, stderr=PIPE, stdout=PIPE)
    out, err = p.communicate()
    if out or err:
        print out
        print err

    # Fix output of files
    fix_permissions(tool='jvivian/deseq2', work_dir=output_dir)

    # Add gene names to output
    output_tsv = os.path.join(output_dir, '{}.tsv'.format(tissue))
    df = pd.read_csv(output_tsv, index_col=0, sep='\t')
    df.index = map_genes(df.index)
    df.to_csv(output_tsv, sep='\t')

    # Clean up
    shutil.rmtree(work_dir)

def deseq2_normalize(df_path,
                     output_dir='.',
                     map_gene_names=True,
                     clean_workdir=True,
                     normalize_fn=None,
                     suffix='.deseq2-normalized.tsv'):
    """
    Accepts a gene by sample expression matrix normalized values with DESeq2
    Output filename: <INPUT>.deseq2-normalized.tsv

    :param str df_path: Path to input expression gene by sample dataframe
    :param str output_dir: Output directory
    :param bool map_gene_names: If True, maps gene IDs to gene names
    :param bool clean_workdir: If True, deletes temporary work directory
    :param fn normalize_fn: Pass a function to apply to the dataframe before normalization. e.g. lambda x: 2**x + 1
    :param str suffix: Suffix added after df_paths basename
    :return: Path to normalized dataframe
    :rtype: str
    """
    # Make workspace directory
    output_dir = os.path.abspath(output_dir)
    work_dir = os.path.join(output_dir, 'work_dir')
    mkdir_p(work_dir)

    # Write out pseudo-vector
    # Tested that results are the same regardless of vector grouping
    samples = [x.strip() for x in open(df_path, 'r').readline().split()]
    tissue_vector = os.path.join(work_dir, 'tissue.vector')
    with open(tissue_vector, 'w') as f:
        f.write('\n'.join(samples))

    if normalize_fn:
        df = pd.read_csv(df_path, sep='\t', index_col=0)
        df = df.apply(normalize_fn)
        df_path = os.path.join(os.path.dirname(df_path), 'processed.' + os.path.basename(df_path))
        df.to_csv(df_path, sep='\t')

    # Write normalization script
    output_path = os.path.join(output_dir, os.path.basename(df_path).split('.')[0] + suffix)
    script_path = os.path.join(work_dir, 'deseq2.R')
    with open(script_path, 'w') as f:
        f.write(
            textwrap.dedent("""
            suppressMessages(library('DESeq2'))

            # Argument parsing
            output_dir <- '/data/'

            # Read in pseudo-vector
            tissue_vector <- read.table('/data/work_dir/tissue.vector')$V1

            # Read in table and process
            print("Reading in dataframe")
            n <- read.table('/df/{df_path}', sep='\\t', header=1, row.names=1, check.names=FALSE)

            # Preprocessing
            print("Rounding data to integers")
            countData <- round(n)

            print("Creating DESeq2 Dataset Object")
            colData <- data.frame(tissue=tissue_vector, row.names=colnames(countData))
            dds <- DESeqDataSetFromMatrix(countData = countData, colData = colData, design = ~ tissue)

            # Estimate size factors
            print("Estimating Size Factors")
            dds <- estimateSizeFactors(dds)

            # Extract Normalized Counts and Write
            print("Extracting normalized counts")
            norm <- counts(dds, normalized = TRUE)

            print("Writing output: {output_name}")
            write.table(norm, file="/data/{output_name}", sep='\\t', quote=F, dec='.', col.names=NA)
            """.format(df_path=os.path.basename(df_path), output_name=os.path.basename(output_path))))

    # Call Docker
    base_params = get_base_call(os.path.dirname(output_path))
    parameters = base_params + ['-v', '{}:/df'.format(os.path.abspath(os.path.dirname(df_path))),
                                'jvivian/deseq2',
                                '/data/work_dir/deseq2.R']

    print '\nCalling: {}\n'.format(' '.join(parameters))
    call(parameters)

    # Fix output of files
    fix_permissions(tool='jvivian/deseq2', work_dir=output_dir)

    # Map gene IDs to gene names
    if map_gene_names:
        df = pd.read_csv(output_path, index_col=0, sep='\t')
        df.index = map_genes(df.index)
        df.to_csv(output_path, sep='\t')

    # Clean up
    if clean_workdir:
        shutil.rmtree(work_dir)

    return output_path
<commit_msg>Refactor DESeq2 for any arbitrary groups<commit_after>import os
import shutil
import textwrap
from collections import defaultdict
from multiprocessing import cpu_count
from subprocess import call, Popen, PIPE

import numpy as np
import pandas as pd
from rnaseq_lib.data import map_genes
from rnaseq_lib.docker import fix_permissions, get_base_call
from rnaseq_lib.utils import mkdir_p
from scipy.stats import pearsonr

def log2fc(a, b, pad=0.001):
    """
    Calculate the log2 Fold Change between two arrays, floats, or integers
    a and b cannot be, nor contain, values less than 0

    :param int|float|np.array a: Value or array
    :param int|float|np.array b: Value or array
    :param int|float pad: Buffer to add to value before log2 calculation
    :return: L2FC array or value
    :rtype: int|float|np.array
    """
    return np.log2(a + pad) - np.log2(b + pad)

def de_pearson_dataframe(df, genes, pair_by='type', gtex=True, tcga=True):
    """
    PearsonR scores of gene differential expression between tumor and normal types.

    1. Calculate log2FC of genes for TCGA tumor samples with matching TCGA normal types
    2. Compare log2fc to tumor type compared to all other normal types
    3. Calculate PearsonR and save

    :param pd.DataFrame df: Exp/TPM dataframe containing "type"/"tissue/tumor/label" metadata columns
    :param list genes: Genes to use in differential expression calculation
    :param str pair_by: How to pair tumors/normals. Either by "type" or "tissue"
    :param bool gtex: If True, includes GTEx in normal set
    :param bool tcga: If True, includes TCGA in normal set
    :return: PearsonR dataframe
    :rtype: pd.DataFrame
    """
    # Subset by Tumor/Normal
    tumor = df[df.label == 'tcga-tumor']
    tcga_n = df[df.label == 'tcga-normal']
    if gtex:
        normal = df[df.tumor == 'no']
    else:
        normal = tcga_n

    # Identify tumor types with paired tcga-normal
    tum_types = [x for x in sorted(tumor[pair_by].unique())
                 if x in sorted(df[df.label == 'tcga-normal'][pair_by].unique())]
    norm_types = []

    # For all paired tumor_types, calculate l2fc, then PearsonR of l2fc to all normal tumor types
    pearson_l2fc = defaultdict(list)
    for tum_type in tum_types:

        # First calculate TCGA tumor/normal prior for comparison
        t = tumor[tumor[pair_by] == tum_type]
        t_med = t[genes].median()
        n = tcga_n[tcga_n[pair_by] == tum_type]
        prior_l2fc = log2fc(t_med, n[genes].median())

        # For every normal type, calculate pearsonR correlation
        for (norm_type, label), _ in normal.groupby(pair_by).label.value_counts().iteritems():
            if tum_type == norm_type:
                l2fc = prior_l2fc
            else:
                n = normal[normal[pair_by] == norm_type]
                l2fc = log2fc(t_med, n[genes].median())

            # Calculate PearsonR Save l2fc and comparison tissue/type
            pearson_r = round(pearsonr(prior_l2fc, l2fc)[0], 2)
            pearson_l2fc[tum_type[:20]].append(pearson_r)
            norm_label = '{}_{}'.format(label, norm_type[:20])
            if norm_label not in norm_types:
                norm_types.append(norm_label)

    return pd.DataFrame(pearson_l2fc, index=norm_types)

def run_deseq2(df_path, group_a, group_b, output_dir, gtex=True, cores=None):
    """
    Runs DESeq2 standard comparison between group A and group B

    :param str df_path: Path to samples by genes dataframe
    :param list(str) group_a: List of samples in group A
    :param list(str) group_b: List of samples in group B
    :param str output_dir: Full path to output directory
    :param bool gtex: If True uses GTEx as normal tissue. Otherwise uses TCGA Normal
    :param int cores: Number of cores to use. Defaults to # of cores on machine.
    """
    # Make workspace directories
    work_dir = os.path.join(output_dir, 'work_dir')
    mkdir_p(work_dir)

    # Write out vectors
    tissue_vector = os.path.join(work_dir, 'tissue.vector')
    with open(tissue_vector, 'w') as f:
        f.write('\n'.join(group_a + group_b))

    disease_vector = os.path.join(work_dir, 'disease.vector')
    with open(disease_vector, 'w') as f:
        f.write('\n'.join(['A' if x in group_a else 'B' for x in group_a + group_b]))

    # Write out script
    cores = cores if cores else int(cpu_count())
    script_path = os.path.join(work_dir, 'deseq2.R')
    with open(script_path, 'w') as f:
        f.write(
            textwrap.dedent("""
            library('DESeq2'); library('data.table'); library('BiocParallel')
            register(MulticoreParam({cores}))

            # Argument parsing
            args <- commandArgs(trailingOnly = TRUE)
            df_path <- args[1]
            tissue_path <- args[2]
            disease_path <- args[3]
            output_dir <- '/data/'

            # Read in vectors
            tissue_vector <- read.table(tissue_path)$V1
            disease_vector <- read.table(disease_path)$V1

            # Read in table and process
            n <- read.table(df_path, sep='\\t', header=1, row.names=1, check.names=FALSE)
            sub <- n[, colnames(n)%in%tissue_vector]
            setcolorder(sub, as.character(tissue_vector))

            # Preprocessing
            countData <- round(sub)
            colData <- data.frame(disease=disease_vector, row.names=colnames(countData))
            y <- DESeqDataSetFromMatrix(countData = countData, colData = colData, design = ~ disease)

            # Run DESeq2
            y <- DESeq(y, parallel=TRUE)
            res <- results(y, parallel=TRUE)
            summary(res)

            # Write out table
            resOrdered <- res[order(res$padj),]
            res_name <- 'results.tsv'
            res_path <- paste(output_dir, res_name, sep='/')
            write.table(as.data.frame(resOrdered), file=res_path, col.names=NA, sep='\\t',  quote=FALSE)

            # MA Plot
            ma_name <- 'MA.pdf'
            ma_path <- paste(output_dir, ma_name, sep='/')
            pdf(ma_path, width=7, height=7)
            plotMA(res, main='DESeq2')
            dev.off()

            # Dispersion Plot
            disp_name <- 'dispersion.pdf'
            disp_path <- paste(output_dir, disp_name, sep='/')
            pdf(disp_path, width=7, height=7)
            plotDispEsts( y, ylim = c(1e-6, 1e1) )
            dev.off()

            # PVal Hist
            hist_name <- 'pval-hist.pdf'
            hist_path <- paste(output_dir, hist_name, sep='/')
            pdf(hist_path, width=7, height=7)
            hist( res$pvalue, breaks=20, col="grey" )
            dev.off()

            # Ratios plots
            qs <- c( 0, quantile( res$baseMean[res$baseMean > 0], 0:7/7 ) )
            bins <- cut( res$baseMean, qs )
            levels(bins) <- paste0("~",round(.5*qs[-1] + .5*qs[-length(qs)]))
            ratios <- tapply( res$pvalue, bins, function(p) mean( p < .01, na.rm=TRUE ) )
            ratio_name <- 'ratios.pdf'
            ratio_path <- paste(output_dir, ratio_name, sep='/')
            pdf(ratio_path, width=7, height=7)
            barplot(ratios, xlab="mean normalized count", ylab="ratio of small $p$ values")
            dev.off()
            """.format(cores=cores)))

    # Call DESeq2
    docker_parameters = ['docker', 'run',
                         '-v', '{}:/data'.format(output_dir),
                         '-v', '{}:/df'.format(os.path.dirname(df_path)),
                         'jvivian/deseq2']

    parameters = ['/data/work_dir/deseq2.R',
                  '/df/{}'.format(os.path.basename(df_path)),
                  '/data/{}'.format(os.path.join('work_dir', 'tissue.vector')),
                  '/data/{}'.format(os.path.join('work_dir', 'disease.vector'))]

    print '\nCalling: {}\n'.format(' '.join(docker_parameters + parameters))
    p = Popen(docker_parameters + parameters, stderr=PIPE, stdout=PIPE)
    out, err = p.communicate()
    if out or err:
        print out
        print err

    # Fix output of files
    fix_permissions(tool='jvivian/deseq2', work_dir=output_dir)

    # Add gene names to output
    output_tsv = os.path.join(output_dir, 'results.tsv')
    df = pd.read_csv(output_tsv, index_col=0, sep='\t')
    df.index = map_genes(df.index)
    df.to_csv(output_tsv, sep='\t')

    # Clean up
    shutil.rmtree(work_dir)

def deseq2_normalize(df_path,
                     output_dir='.',
                     map_gene_names=True,
                     clean_workdir=True,
                     normalize_fn=None,
                     suffix='.deseq2-normalized.tsv'):
    """
    Accepts a gene by sample expression matrix normalized values with DESeq2
    Output filename: <INPUT>.deseq2-normalized.tsv

    :param str df_path: Path to input expression gene by sample dataframe
    :param str output_dir: Output directory
    :param bool map_gene_names: If True, maps gene IDs to gene names
    :param bool clean_workdir: If True, deletes temporary work directory
    :param fn normalize_fn: Pass a function to apply to the dataframe before normalization. e.g. lambda x: 2**x + 1
    :param str suffix: Suffix added after df_paths basename
    :return: Path to normalized dataframe
    :rtype: str
    """
    # Make workspace directory
    output_dir = os.path.abspath(output_dir)
    work_dir = os.path.join(output_dir, 'work_dir')
    mkdir_p(work_dir)

    # Write out pseudo-vector
    # Tested that results are the same regardless of vector grouping
    samples = [x.strip() for x in open(df_path, 'r').readline().split()]
    tissue_vector = os.path.join(work_dir, 'tissue.vector')
    with open(tissue_vector, 'w') as f:
        f.write('\n'.join(samples))

    if normalize_fn:
        df = pd.read_csv(df_path, sep='\t', index_col=0)
        df = df.apply(normalize_fn)
        df_path = os.path.join(os.path.dirname(df_path), 'processed.' + os.path.basename(df_path))
        df.to_csv(df_path, sep='\t')

    # Write normalization script
    output_path = os.path.join(output_dir, os.path.basename(df_path).split('.')[0] + suffix)
    script_path = os.path.join(work_dir, 'deseq2.R')
    with open(script_path, 'w') as f:
        f.write(
            textwrap.dedent("""
            suppressMessages(library('DESeq2'))

            # Argument parsing
            output_dir <- '/data/'

            # Read in pseudo-vector
            tissue_vector <- read.table('/data/work_dir/tissue.vector')$V1

            # Read in table and process
            print("Reading in dataframe")
            n <- read.table('/df/{df_path}', sep='\\t', header=1, row.names=1, check.names=FALSE)

            # Preprocessing
            print("Rounding data to integers")
            countData <- round(n)

            print("Creating DESeq2 Dataset Object")
            colData <- data.frame(tissue=tissue_vector, row.names=colnames(countData))
            dds <- DESeqDataSetFromMatrix(countData = countData, colData = colData, design = ~ tissue)

            # Estimate size factors
            print("Estimating Size Factors")
            dds <- estimateSizeFactors(dds)

            # Extract Normalized Counts and Write
            print("Extracting normalized counts")
            norm <- counts(dds, normalized = TRUE)

            print("Writing output: {output_name}")
            write.table(norm, file="/data/{output_name}", sep='\\t', quote=F, dec='.', col.names=NA)
            """.format(df_path=os.path.basename(df_path), output_name=os.path.basename(output_path))))

    # Call Docker
    base_params = get_base_call(os.path.dirname(output_path))
    parameters = base_params + ['-v', '{}:/df'.format(os.path.abspath(os.path.dirname(df_path))),
                                'jvivian/deseq2',
                                '/data/work_dir/deseq2.R']

    print '\nCalling: {}\n'.format(' '.join(parameters))
    call(parameters)

    # Fix output of files
    fix_permissions(tool='jvivian/deseq2', work_dir=output_dir)

    # Map gene IDs to gene names
    if map_gene_names:
        df = pd.read_csv(output_path, index_col=0, sep='\t')
        df.index = map_genes(df.index)
        df.to_csv(output_path, sep='\t')

    # Clean up
    if clean_workdir:
        shutil.rmtree(work_dir)

    return output_path
<|endoftext|>


<commit_before>"""
    This a garbo discovery service for AWS EC2 resources
"""
import logging

import dateutil.parser
import boto.ec2
import boto.ec2.elb
import boto.exception

from garbo import config
from garbo.model.aws.ec2 import EBSSnapshot, EBSVolume, Image, Instance, LoadBalancer, SecurityGroup
from garbo.model import Relation

__author__ = 'nati'

_resource_collectors = set()

def aws_collector(f):
    """
    Collect AWS collectors to a global module set

    :param f:  AWS collector function, getting boto EC2 connection and yielding resources and/or relations
    """
    _resource_collectors.add(f)
    return f

@aws_collector
def snapshots(conn):
    """
    Collect EC2 EBS Snapshots created by this account

    :param conn: boto EC2 connection
    """
    volume_ids = [v.id for v in conn.get_all_volumes()]
    for snapshot in conn.get_all_snapshots(owner='self'):
        snapshot_resource = EBSSnapshot(region=conn.region.name, resource_id=snapshot.id,
                                        created=dateutil.parser.parse(snapshot.start_time))
        yield snapshot_resource
        if snapshot.volume_id in volume_ids:
            yield Relation(snapshot_resource,
                           EBSVolume(region=conn.region.name, resource_id=snapshot.volume_id))

@aws_collector
def load_balancers(conn):
    """
    Collect EC2 Elastic Load Balances (ELBs)

    :param conn: boto EC2 connection
    """
    # Convert EC2 connection to ELB connection
    conn = boto.ec2.elb.connect_to_region(conn.region.name,
                                          aws_access_key_id=conn.provider.access_key,
                                          aws_secret_access_key=conn.provider.secret_key)
    for elb in conn.get_all_load_balancers():
        lb_resource = LoadBalancer(region=conn.region.name, resource_id=elb.dns_name,
                                   created=dateutil.parser.parse(elb.created_time))
        yield lb_resource
        # security groups relations
        for group_id in elb.security_groups:
            yield Relation(lb_resource,
                           SecurityGroup(region=conn.region.name, resource_id=group_id),
                           dependency=True)
        # instances relations
        for instance_id in {i.id for i in elb.instances}:
            yield Relation(lb_resource,
                           Instance(region=conn.region.name, resource_id=instance_id),
                           dependency=True)

@aws_collector
def security_groups(conn):
    """
    Collect EC2 Security Groups

    :param conn: boto EC2 connection
    """
    for group in conn.get_all_security_groups():
        # Fun fact: there is no creation/modification date associated with AWS security groups
        yield SecurityGroup(region=conn.region.name, resource_id=group.id)
        # TODO: add cross-group dependency

@aws_collector
def images(conn):
    """
    Collect EC2 Amazon Machine Images (AMIs) created by this account

    :param conn: boto EC2 connection
    """
    snapshot_ids = [s.id for s in conn.get_all_snapshots(owner='self')]
    for image in conn.get_all_images(owners='self'):
        image_resource = Image(region=conn.region.name, resource_id=image.id,
                               created=dateutil.parser.parse(image.creationDate))
        yield image_resource
        # Images can have mapping to an EBS snapshot (stored in S3)
        for snapshot_id in {v.snapshot_id for v in image.block_device_mapping.values()
                            if v.snapshot_id in snapshot_ids}:
            yield Relation(image_resource,
                           EBSSnapshot(region=conn.region.name, resource_id=snapshot_id),
                           dependency=True)

@aws_collector
def instances(conn):
    """
    Collect EC2 Instances

    :param conn: boto EC2 connection
    """
    image_ids = [i.id for i in conn.get_all_images(owners='self')]
    for instance in conn.get_only_instances():
        instance_resource = Instance(region=conn.region.name, resource_id=instance.id,
                                     created=dateutil.parser.parse(instance.launch_time),
                                     used=Instance.is_used(instance))
        yield instance_resource
        # only yield relations for self-owned images
        if instance.image_id in image_ids:
            yield Relation(instance_resource,
                           Image(region=conn.region.name, resource_id=instance.image_id))
        # security groups relations
        for group_id in {sg.id for sg in instance.groups}:
            yield Relation(instance_resource,
                           SecurityGroup(region=conn.region.name, resource_id=group_id),
                           dependency=True)

@aws_collector
def ebs_volumes(conn):
    """
    Collect EC2 EBS Volumes

    :param conn: boto EC2 connection
    """
    snapshot_ids = [s.id for s in conn.get_all_snapshots(owner='self')]
    for volume in conn.get_all_volumes():
        ebs_volume_resource = EBSVolume(region=conn.region.name, resource_id=volume.id,
                                        created=dateutil.parser.parse(volume.create_time))
        yield ebs_volume_resource
        # If volume is attached, yield a relation to the instance
        if volume.attach_data.status in ('attaching', 'attached'):
            yield Relation(Instance(region=conn.region.name,
                                    resource_id=volume.attach_data.instance_id),
                           ebs_volume_resource,
                           dependency=True)
        # Which snapshot was this volume created from
        if volume.snapshot_id in snapshot_ids:
            yield Relation(ebs_volume_resource,
                           EBSSnapshot(region=conn.region.name, resource_id=volume.snapshot_id))

def collect_all(aws_access_key=None, aws_secret_key=None):
    """
    Yield all EC2 resources and relations associated with an AWS account

    :param aws_access_key:
    :param aws_secret_key:
    """
    aws_access_key = aws_access_key or config.aws.access_key
    aws_secret_key = aws_secret_key or config.aws.secret_access_key

    regions = config.aws.regions or (r.name for r in boto.ec2.regions())
    for region in regions:
        conn = boto.ec2.connect_to_region(region,
                                          aws_access_key_id=aws_access_key,
                                          aws_secret_access_key=aws_secret_key)
        for collector in _resource_collectors:
            try:
                for item in collector(conn):
                    yield item
            except boto.exception.BotoServerError as e:
                logging.warn('unable to run collector %s for region %s: %s',
                             collector.__name__, region, e.message)
<commit_msg>add cross-security-group dependencies<commit_after>"""
    This a garbo discovery service for AWS EC2 resources
"""
import logging

import dateutil.parser
import boto.ec2
import boto.ec2.elb
import boto.exception

from garbo import config
from garbo.model.aws.ec2 import EBSSnapshot, EBSVolume, Image, Instance, LoadBalancer, SecurityGroup
from garbo.model import Relation

__author__ = 'nati'

_resource_collectors = set()

def aws_collector(f):
    """
    Collect AWS collectors to a global module set

    :param f:  AWS collector function, getting boto EC2 connection and yielding resources and/or relations
    """
    _resource_collectors.add(f)
    return f

@aws_collector
def snapshots(conn):
    """
    Collect EC2 EBS Snapshots created by this account

    :param conn: boto EC2 connection
    """
    volume_ids = [v.id for v in conn.get_all_volumes()]
    for snapshot in conn.get_all_snapshots(owner='self'):
        snapshot_resource = EBSSnapshot(region=conn.region.name, resource_id=snapshot.id,
                                        created=dateutil.parser.parse(snapshot.start_time))
        yield snapshot_resource
        if snapshot.volume_id in volume_ids:
            yield Relation(snapshot_resource,
                           EBSVolume(region=conn.region.name, resource_id=snapshot.volume_id))

@aws_collector
def load_balancers(conn):
    """
    Collect EC2 Elastic Load Balances (ELBs)

    :param conn: boto EC2 connection
    """
    # Convert EC2 connection to ELB connection
    conn = boto.ec2.elb.connect_to_region(conn.region.name,
                                          aws_access_key_id=conn.provider.access_key,
                                          aws_secret_access_key=conn.provider.secret_key)
    for elb in conn.get_all_load_balancers():
        lb_resource = LoadBalancer(region=conn.region.name, resource_id=elb.dns_name,
                                   created=dateutil.parser.parse(elb.created_time))
        yield lb_resource
        # security groups relations
        for group_id in elb.security_groups:
            yield Relation(lb_resource,
                           SecurityGroup(region=conn.region.name, resource_id=group_id),
                           dependency=True)
        # instances relations
        for instance_id in {i.id for i in elb.instances}:
            yield Relation(lb_resource,
                           Instance(region=conn.region.name, resource_id=instance_id),
                           dependency=True)

@aws_collector
def security_groups(conn):
    """
    Collect EC2 Security Groups

    :param conn: boto EC2 connection
    """
    for group in conn.get_all_security_groups():
        # Fun fact: there is no creation/modification date associated with AWS security groups
        sg_resource = SecurityGroup(region=conn.region.name, resource_id=group.id)
        yield sg_resource
        # add cross-group dependency
        for group_id in {g.group_id for r in group.rules + group.rules_egress for g in r.grants
                         if g.group_id and g.owner_id == group.owner_id}:
            yield Relation(sg_resource,
                           SecurityGroup(region=conn.region.name, resource_id=group_id),
                           dependency=True)

@aws_collector
def images(conn):
    """
    Collect EC2 Amazon Machine Images (AMIs) created by this account

    :param conn: boto EC2 connection
    """
    snapshot_ids = [s.id for s in conn.get_all_snapshots(owner='self')]
    for image in conn.get_all_images(owners='self'):
        image_resource = Image(region=conn.region.name, resource_id=image.id,
                               created=dateutil.parser.parse(image.creationDate))
        yield image_resource
        # Images can have mapping to an EBS snapshot (stored in S3)
        for snapshot_id in {v.snapshot_id for v in image.block_device_mapping.values()
                            if v.snapshot_id in snapshot_ids}:
            yield Relation(image_resource,
                           EBSSnapshot(region=conn.region.name, resource_id=snapshot_id),
                           dependency=True)

@aws_collector
def instances(conn):
    """
    Collect EC2 Instances

    :param conn: boto EC2 connection
    """
    image_ids = [i.id for i in conn.get_all_images(owners='self')]
    for instance in conn.get_only_instances():
        instance_resource = Instance(region=conn.region.name, resource_id=instance.id,
                                     created=dateutil.parser.parse(instance.launch_time),
                                     used=Instance.is_used(instance))
        yield instance_resource
        # only yield relations for self-owned images
        if instance.image_id in image_ids:
            yield Relation(instance_resource,
                           Image(region=conn.region.name, resource_id=instance.image_id))
        # security groups relations
        for group_id in {sg.id for sg in instance.groups}:
            yield Relation(instance_resource,
                           SecurityGroup(region=conn.region.name, resource_id=group_id),
                           dependency=True)

@aws_collector
def ebs_volumes(conn):
    """
    Collect EC2 EBS Volumes

    :param conn: boto EC2 connection
    """
    snapshot_ids = [s.id for s in conn.get_all_snapshots(owner='self')]
    for volume in conn.get_all_volumes():
        ebs_volume_resource = EBSVolume(region=conn.region.name, resource_id=volume.id,
                                        created=dateutil.parser.parse(volume.create_time))
        yield ebs_volume_resource
        # If volume is attached, yield a relation to the instance
        if volume.attach_data.status in ('attaching', 'attached'):
            yield Relation(Instance(region=conn.region.name,
                                    resource_id=volume.attach_data.instance_id),
                           ebs_volume_resource,
                           dependency=True)
        # Which snapshot was this volume created from
        if volume.snapshot_id in snapshot_ids:
            yield Relation(ebs_volume_resource,
                           EBSSnapshot(region=conn.region.name, resource_id=volume.snapshot_id))

def collect_all(aws_access_key=None, aws_secret_key=None):
    """
    Yield all EC2 resources and relations associated with an AWS account

    :param aws_access_key:
    :param aws_secret_key:
    """
    aws_access_key = aws_access_key or config.aws.access_key
    aws_secret_key = aws_secret_key or config.aws.secret_access_key

    regions = config.aws.regions or (r.name for r in boto.ec2.regions())
    for region in regions:
        conn = boto.ec2.connect_to_region(region,
                                          aws_access_key_id=aws_access_key,
                                          aws_secret_access_key=aws_secret_key)
        for collector in _resource_collectors:
            try:
                for item in collector(conn):
                    yield item
            except boto.exception.BotoServerError as e:
                logging.warn('unable to run collector %s for region %s: %s',
                             collector.__name__, region, e.message)
<|endoftext|>


<commit_before>import logging
from ConfigParser import ConfigParser
from calendar import monthrange
from datetime import datetime
from glob import glob
from os import getcwd
from os.path import join, basename, splitext, abspath, dirname
import numpy
import pandas
from scipy.interpolate import griddata
from utils import *

LOG_FILE_NAME = 'process.log'

configure_logging(LOG_FILE_NAME)

DATA_FRAMES_DIR = r'DATA_FRAMES'
DAILY_RESULTS_DIR = r'DAILY'
MONTHLY_RESULTS_DIR = r'MONTHLY'

def create_weather_reports():
    logging.info('Start parsing weather data')
    config = ConfigParser()
    config.read(join(dirname(abspath(__file__)), 'paths_config.cfg'))

    csv_data_path = config.get('paths', 'csv_data_path')
    section_files_path = config.get('paths', 'section_files_path')
    basein_files_path = config.get('paths', 'basein_files_path')

    output_dir_path = config.get('paths', 'output_dir_path')

    data_frames_dir_path = join(output_dir_path, DATA_FRAMES_DIR)
    daily_results_dir_path = join(output_dir_path, DAILY_RESULTS_DIR)
    monthly_results_dir_path = join(output_dir_path, MONTHLY_RESULTS_DIR)

    clear_dir(data_frames_dir_path)
    open(join(getcwd(), LOG_FILE_NAME), 'w').close()
    shape_file_and_correspondent_stations = get_shape_file_and_correspondent_stations(csv_data_path, section_files_path)
    run_process_of_making_data_frames(shape_file_and_correspondent_stations, data_frames_dir_path)

    clear_dir(daily_results_dir_path)
    data_files = get_files_for_getting_daily_metrics(basein_files_path, section_files_path, data_frames_dir_path)
    run_processing_daily_metrics(data_files, daily_results_dir_path)

    clear_dir(monthly_results_dir_path)
    run_processing_monthly_metrics(daily_results_dir_path, monthly_results_dir_path)
    logging.info('Reports created. You can find them inside output folder')

def run_process_of_making_data_frames(shape_file_and_correspondent_stations, output_path):
    logging.info('Starting making data frame tasks')
    tasks_desc = []
    for shape_file, stations in shape_file_and_correspondent_stations.items():
        task_desc = {'target': make_data_frames, 'args': (shape_file, stations, output_path)}
        tasks_desc.append(task_desc)

    run_tasks(tasks_desc)

def make_data_frames(station_shape_file, stations_data, output_path):
    logging.info('Making data frame for shapefile: %s', basename(station_shape_file))
    prcp, tmax, tmin = aggregate_data_frames(stations_data)

    logging.info("joining {0:d} dataframes".format(len(prcp)))
    prcp_df, tmax_df, tmin_df = join_data_aggregated_data_frames(prcp, tmax, tmin)

    output_base = join(output_path, basename(station_shape_file).rstrip('.shp'))
    dump_data_frames(output_base, prcp_df, tmax_df, tmin_df)

def aggregate_data_frames(stations_data):
    prcp, tmin, tmax = [], [], []
    for csv_file_path, stations in stations_data.items():
        logging.info('Reading file: %s', csv_file_path)
        csv_file = pandas.read_csv(csv_file_path, na_values=-9999,
                                   parse_dates=["DATE"], usecols=[0, 5, 8, 11, 12])
        for group_info in csv_file.groupby("STATION").groups.items():
            station, _ = group_info
            if station in stations:
                logging.info("--processing station: " + station)
                prcp.append(get_specific_data_frame_for_group(csv_file, group_info, "PRCP"))
                tmin.append(get_specific_data_frame_for_group(csv_file, group_info, "TMIN"))
                tmax.append(get_specific_data_frame_for_group(csv_file, group_info, "TMAX"))

    return prcp, tmax, tmin

def get_specific_data_frame_for_group(global_data_frame, group_info, data_name):
    group_id, indexes = group_info
    data_frame = global_data_frame.loc[indexes, ["DATE", data_name]]
    data_frame[group_id] = data_frame.pop(data_name)
    data_frame.index = data_frame.pop("DATE")
    return data_frame

def join_data_aggregated_data_frames(prcp, tmax, tmin):
    prcp_df = join_data_frames(prcp)
    tmin_df = join_data_frames(tmin)
    tmax_df = join_data_frames(tmax)
    return prcp_df, tmax_df, tmin_df

def join_data_frames(data_frames):
    result_df = data_frames[0]
    for df in data_frames[1:]:
        result_df = result_df.merge(df, how="outer", left_index=True, right_index=True)
    return result_df

def dump_data_frames(output_base, prcp_df, tmax_df, tmin_df):
    prcp_df.to_csv(output_base + '_prcp.csv')
    tmin_df.to_csv(output_base + '_tmin.csv')
    tmax_df.to_csv(output_base + '_tmax.csv')

def run_processing_daily_metrics(data_files, output_path):
    logging.info('Start processing daily')
    tasks_desc = []
    for basein_file_path, section_file_path, csv_files_paths in data_files:
        task_desc = {'target': process_daily, 'args': (basein_file_path, section_file_path,
                                                       csv_files_paths, output_path)}
        tasks_desc.append(task_desc)

    run_tasks(tasks_desc)

def process_daily(basein_file_path, section_file_path, csv_files_paths, output_path):
    logging.info('Process daily for %s', basename(basein_file_path))
    x, y, masks_array = load_basein_file(basein_file_path)
    section_stations = load_section_file(section_file_path)

    for csv_file_path in csv_files_paths:
        csv_file = pandas.read_csv(csv_file_path, parse_dates=True, index_col=0)
        output_file_name = splitext(basename(csv_file_path))[0] + '_processed.csv'
        output_file_path = join(output_path, output_file_name)

        with open(output_file_path, 'w', 0) as output_file:
            output_file.write(','.join(("date", "area-weighted", "max", "min", "count", "gauge_pairs")) + '\n')

            for date in csv_file.index:
                points, stations, values = prepare_data_raw_data(csv_file, date, section_stations)

                count = len(values)
                if len(values) == 0:
                    awa = numpy.NaN
                    upr = numpy.NaN
                    lwr = numpy.NaN
                else:
                    values = numpy.array(values)
                    points = numpy.array(points)
                    if not numpy.all(values == values[0]):
                        grid = griddata(points, values, (x, y), method="nearest")
                        grid = numpy.ma.masked_where(masks_array == 0, grid)
                        contrib = []
                        unique_vals = numpy.unique(grid)
                        count = 0
                        for station, value in zip(stations, values):
                            if value in unique_vals:
                                mask = numpy.zeros_like(grid)
                                mask = numpy.ma.masked_where(masks_array == 0, mask)
                                mask[numpy.where(grid == value)] = 1
                                unique_vals_count = numpy.sum(mask)
                                contrib.append(unique_vals_count / masks_array.sum())
                                count += 1
                        awa = grid.mean()
                        upr = max(unique_vals)
                        lwr = min(unique_vals)
                        values = contrib
                    else:
                        awa = values[0]
                        upr = values[0]
                        lwr = values[0]

                line = "{0:s},{1:E},{2:E},{3:E},{4:d},".format(str(date), float(awa), float(upr), float(lwr), count)
                for site, val in zip(stations, values):
                    line += ' ' + site + "|{0:E}".format(val)
                line += '\n'
                output_file.write(line)

def prepare_data_raw_data(csv_file, date, section_stations):
    date_df = csv_file.loc[date, :].dropna()
    points = []
    values = []
    stations = []
    if date_df.shape[0] == 1:
        values = date_df.values
    elif date_df.shape[0] != 0:
        for station, value in zip(date_df.index, date_df.values):
            if station in stations:
                raise Exception("Same station: %s found again on date %s" % (str(station), str(date)))

            points.append(section_stations[station])
            values.append(value)
            stations.append(station)
    return points, stations, values

def run_processing_monthly_metrics(data_files_path, output_path):
    tasks_desc = []
    for file_path in glob(join(data_files_path, '*.csv')):
        task_desc = {'target': calculate_monthly_values, 'args': (file_path, output_path)}
        tasks_desc.append(task_desc)

    run_tasks(tasks_desc)

def calculate_monthly_values(data_file_path, output_path):
    logging.info('Calculating monthly metrics for file: %s', data_file_path)
    daily_data_frame = pandas.read_csv(data_file_path, parse_dates=True, index_col=0,
                                       usecols=["date", "area-weighted"], na_values=["NAN"])

    if "_prcp_" in data_file_path.lower():
        groups = daily_data_frame.groupby([lambda x: x.year, lambda x: x.month]).sum()
    else:
        groups = daily_data_frame.groupby([lambda x: x.year, lambda x: x.month]).mean()

    datetimes_collection = []
    for year, month in groups.index:
        datetimes_collection.append(datetime(year=year, month=month, day=monthrange(year, month)[1],
                                             hour=23, minute=59, second=59))
    groups.index = datetimes_collection
    groups.to_csv(join(output_path, basename(data_file_path)), index_label="datetime")

if __name__ == '__main__':
    try:
        create_weather_reports()
    except Exception as e:
        logging.exception('Error on executing main function')
<commit_msg>!U add logging<commit_after>import logging
from ConfigParser import ConfigParser
from calendar import monthrange
from datetime import datetime
from glob import glob
from os import getcwd
from os.path import join, basename, splitext, abspath, dirname
import numpy
import pandas
from scipy.interpolate import griddata
from utils import *

LOG_FILE_NAME = 'process.log'

configure_logging(LOG_FILE_NAME)

DATA_FRAMES_DIR = r'DATA_FRAMES'
DAILY_RESULTS_DIR = r'DAILY'
MONTHLY_RESULTS_DIR = r'MONTHLY'

def create_weather_reports():
    logging.info('Start parsing weather data')
    config = ConfigParser()
    config.read(join(dirname(abspath(__file__)), 'paths_config.cfg'))

    csv_data_path = config.get('paths', 'csv_data_path')
    section_files_path = config.get('paths', 'section_files_path')
    basein_files_path = config.get('paths', 'basein_files_path')

    output_dir_path = config.get('paths', 'output_dir_path')

    data_frames_dir_path = join(output_dir_path, DATA_FRAMES_DIR)
    daily_results_dir_path = join(output_dir_path, DAILY_RESULTS_DIR)
    monthly_results_dir_path = join(output_dir_path, MONTHLY_RESULTS_DIR)

    clear_dir(data_frames_dir_path)
    open(join(getcwd(), LOG_FILE_NAME), 'w').close()
    shape_file_and_correspondent_stations = get_shape_file_and_correspondent_stations(csv_data_path, section_files_path)
    run_process_of_making_data_frames(shape_file_and_correspondent_stations, data_frames_dir_path)

    clear_dir(daily_results_dir_path)
    data_files = get_files_for_getting_daily_metrics(basein_files_path, section_files_path, data_frames_dir_path)
    run_processing_daily_metrics(data_files, daily_results_dir_path)

    clear_dir(monthly_results_dir_path)
    run_processing_monthly_metrics(daily_results_dir_path, monthly_results_dir_path)
    logging.info('Reports created. You can find them inside output folder')

def run_process_of_making_data_frames(shape_file_and_correspondent_stations, output_path):
    logging.info('Starting making data frame tasks')
    tasks_desc = []
    for shape_file, stations in shape_file_and_correspondent_stations.items():
        task_desc = {'target': make_data_frames, 'args': (shape_file, stations, output_path)}
        tasks_desc.append(task_desc)

    run_tasks(tasks_desc)

def make_data_frames(station_shape_file, stations_data, output_path):
    logging.info('Making data frame for shapefile: %s', basename(station_shape_file))
    prcp, tmax, tmin = aggregate_data_frames(stations_data)

    logging.info("joining {0:d} dataframes".format(len(prcp)))
    prcp_df, tmax_df, tmin_df = join_data_aggregated_data_frames(prcp, tmax, tmin)

    output_base = join(output_path, basename(station_shape_file).rstrip('.shp'))
    dump_data_frames(output_base, prcp_df, tmax_df, tmin_df)

def aggregate_data_frames(stations_data):
    prcp, tmin, tmax = [], [], []
    for csv_file_path, stations in stations_data.items():
        logging.info('Reading file: %s', csv_file_path)
        try:
            csv_file = pandas.read_csv(csv_file_path, na_values=-9999,
                                       parse_dates=["DATE"], usecols=[0, 5, 8, 11, 12])
        except ValueError:
            logging.exception('Failed to parse file: %s', csv_file_path)
        else:
            for group_info in csv_file.groupby("STATION").groups.items():
                station, _ = group_info
                if station in stations:
                    logging.info("--processing station: " + station)
                    prcp.append(get_specific_data_frame_for_group(csv_file, group_info, "PRCP"))
                    tmin.append(get_specific_data_frame_for_group(csv_file, group_info, "TMIN"))
                    tmax.append(get_specific_data_frame_for_group(csv_file, group_info, "TMAX"))
    return prcp, tmax, tmin

def get_specific_data_frame_for_group(global_data_frame, group_info, data_name):
    group_id, indexes = group_info
    data_frame = global_data_frame.loc[indexes, ["DATE", data_name]]
    data_frame[group_id] = data_frame.pop(data_name)
    data_frame.index = data_frame.pop("DATE")
    return data_frame

def join_data_aggregated_data_frames(prcp, tmax, tmin):
    prcp_df = join_data_frames(prcp)
    tmin_df = join_data_frames(tmin)
    tmax_df = join_data_frames(tmax)
    return prcp_df, tmax_df, tmin_df

def join_data_frames(data_frames):
    result_df = data_frames[0]
    for df in data_frames[1:]:
        result_df = result_df.merge(df, how="outer", left_index=True, right_index=True)
    return result_df

def dump_data_frames(output_base, prcp_df, tmax_df, tmin_df):
    prcp_df.to_csv(output_base + '_prcp.csv')
    tmin_df.to_csv(output_base + '_tmin.csv')
    tmax_df.to_csv(output_base + '_tmax.csv')

def run_processing_daily_metrics(data_files, output_path):
    logging.info('Start processing daily')
    tasks_desc = []
    for basein_file_path, section_file_path, csv_files_paths in data_files:
        task_desc = {'target': process_daily, 'args': (basein_file_path, section_file_path,
                                                       csv_files_paths, output_path)}
        tasks_desc.append(task_desc)

    run_tasks(tasks_desc)

def process_daily(basein_file_path, section_file_path, csv_files_paths, output_path):
    logging.info('Process daily for %s', basename(basein_file_path))
    x, y, masks_array = load_basein_file(basein_file_path)
    section_stations = load_section_file(section_file_path)

    for csv_file_path in csv_files_paths:
        logging.info('Processing %s', basename(csv_file_path))
        csv_file = pandas.read_csv(csv_file_path, parse_dates=True, index_col=0)
        output_file_name = splitext(basename(csv_file_path))[0] + '_processed.csv'
        output_file_path = join(output_path, output_file_name)

        with open(output_file_path, 'w', 0) as output_file:
            output_file.write(','.join(("date", "area-weighted", "max", "min", "count", "gauge_pairs")) + '\n')

            for date in csv_file.index:
                points, stations, values = prepare_data_raw_data(csv_file, date, section_stations)

                count = len(values)
                if len(values) == 0:
                    awa = numpy.NaN
                    upr = numpy.NaN
                    lwr = numpy.NaN
                else:
                    values = numpy.array(values)
                    points = numpy.array(points)
                    if not numpy.all(values == values[0]):
                        grid = griddata(points, values, (x, y), method="nearest")
                        grid = numpy.ma.masked_where(masks_array == 0, grid)
                        contrib = []
                        unique_vals = numpy.unique(grid)
                        count = 0
                        for station, value in zip(stations, values):
                            if value in unique_vals:
                                mask = numpy.zeros_like(grid)
                                mask = numpy.ma.masked_where(masks_array == 0, mask)
                                mask[numpy.where(grid == value)] = 1
                                unique_vals_count = numpy.sum(mask)
                                contrib.append(unique_vals_count / masks_array.sum())
                                count += 1
                        awa = grid.mean()
                        upr = max(unique_vals)
                        lwr = min(unique_vals)
                        values = contrib
                    else:
                        awa = values[0]
                        upr = values[0]
                        lwr = values[0]

                line = "{0:s},{1:E},{2:E},{3:E},{4:d},".format(str(date), float(awa), float(upr), float(lwr), count)
                for site, val in zip(stations, values):
                    line += ' ' + site + "|{0:E}".format(val)
                line += '\n'
                output_file.write(line)

def prepare_data_raw_data(csv_file, date, section_stations):
    date_df = csv_file.loc[date, :].dropna()
    points = []
    values = []
    stations = []
    if date_df.shape[0] == 1:
        values = date_df.values
    elif date_df.shape[0] != 0:
        for station, value in zip(date_df.index, date_df.values):
            if station in stations:
                raise Exception("Same station: %s found again on date %s" % (str(station), str(date)))

            points.append(section_stations[station])
            values.append(value)
            stations.append(station)
    return points, stations, values

def run_processing_monthly_metrics(data_files_path, output_path):
    tasks_desc = []
    for file_path in glob(join(data_files_path, '*.csv')):
        task_desc = {'target': calculate_monthly_values, 'args': (file_path, output_path)}
        tasks_desc.append(task_desc)

    run_tasks(tasks_desc)

def calculate_monthly_values(data_file_path, output_path):
    logging.info('Calculating monthly metrics for file: %s', data_file_path)
    daily_data_frame = pandas.read_csv(data_file_path, parse_dates=True, index_col=0,
                                       usecols=["date", "area-weighted"], na_values=["NAN"])

    if "_prcp_" in data_file_path.lower():
        groups = daily_data_frame.groupby([lambda x: x.year, lambda x: x.month]).sum()
    else:
        groups = daily_data_frame.groupby([lambda x: x.year, lambda x: x.month]).mean()

    datetimes_collection = []
    for year, month in groups.index:
        datetimes_collection.append(datetime(year=year, month=month, day=monthrange(year, month)[1],
                                             hour=23, minute=59, second=59))
    groups.index = datetimes_collection
    groups.to_csv(join(output_path, basename(data_file_path)), index_label="datetime")

if __name__ == '__main__':
    try:
        create_weather_reports()
    except Exception as e:
        logging.exception('Error on executing main function')
<|endoftext|>


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

# Copyright (c) 2009-2013, The Linux Foundation. All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#     * Redistributions of source code must retain the above copyright
#       notice, this list of conditions and the following disclaimer.
#     * Redistributions in binary form must reproduce the above copyright
#       notice, this list of conditions and the following disclaimer in the
#       documentation and/or other materials provided with the distribution.
#     * Neither the name of The Linux Foundation nor
#       the names of its contributors may be used to endorse or promote
#       products derived from this software without specific prior written
#       permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
# NON-INFRINGEMENT ARE DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT OWNER OR
# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS;
# OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY,
# WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR
# OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF
# ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

# Build the kernel for all targets using the Android build environment.
#
# TODO: Accept arguments to indicate what to build.

import glob
from optparse import OptionParser
import subprocess
import os
import os.path
import re
import shutil
import sys

version = 'build-all.py, version 0.01'

build_dir = '../all-kernels'
make_command = ["vmlinux", "modules", "dtbs"]
make_env = os.environ
make_env.update({
        'ARCH': 'arm',
        'KCONFIG_NOTIMESTAMP': 'true' })
make_env.setdefault('CROSS_COMPILE', 'arm-none-linux-gnueabi-')
all_options = {}

def error(msg):
    sys.stderr.write("error: %s\n" % msg)

def fail(msg):
    """Fail with a user-printed message"""
    error(msg)
    sys.exit(1)

def check_kernel():
    """Ensure that PWD is a kernel directory"""
    if (not os.path.isfile('MAINTAINERS') or
        not os.path.isfile('arch/arm/mach-msm/Kconfig')):
        fail("This doesn't seem to be an MSM kernel dir")

def check_build():
    """Ensure that the build directory is present."""
    if not os.path.isdir(build_dir):
        try:
            os.makedirs(build_dir)
        except OSError as exc:
            if exc.errno == errno.EEXIST:
                pass
            else:
                raise

def update_config(file, str):
    print 'Updating %s with \'%s\'\n' % (file, str)
    defconfig = open(file, 'a')
    defconfig.write(str + '\n')
    defconfig.close()

def scan_configs():
    """Get the full list of defconfigs appropriate for this tree."""
    names = {}
    arch_pats = (
        r'[fm]sm[0-9]*_defconfig',
        r'apq*_defconfig',
        r'qsd*_defconfig',
	r'msmkrypton*_defconfig',
        )
    for p in arch_pats:
        for n in glob.glob('arch/arm/configs/' + p):
            names[os.path.basename(n)[:-10]] = n
    return names

class Builder:
    def __init__(self, logname):
        self.logname = logname
        self.fd = open(logname, 'w')

    def run(self, args):
        devnull = open('/dev/null', 'r')
        proc = subprocess.Popen(args, stdin=devnull,
                env=make_env,
                bufsize=0,
                stdout=subprocess.PIPE,
                stderr=subprocess.STDOUT)
        count = 0
        # for line in proc.stdout:
        rawfd = proc.stdout.fileno()
        while True:
            line = os.read(rawfd, 1024)
            if not line:
                break
            self.fd.write(line)
            self.fd.flush()
            if all_options.verbose:
                sys.stdout.write(line)
                sys.stdout.flush()
            else:
                for i in range(line.count('\n')):
                    count += 1
                    if count == 64:
                        count = 0
                        print
                    sys.stdout.write('.')
                sys.stdout.flush()
        print
        result = proc.wait()

        self.fd.close()
        return result

failed_targets = []

def build(target):
    dest_dir = os.path.join(build_dir, target)
    log_name = '%s/log-%s.log' % (build_dir, target)
    print 'Building %s in %s log %s' % (target, dest_dir, log_name)
    if not os.path.isdir(dest_dir):
        os.mkdir(dest_dir)
    defconfig = 'arch/arm/configs/%s_defconfig' % target
    dotconfig = '%s/.config' % dest_dir
    savedefconfig = '%s/defconfig' % dest_dir
    shutil.copyfile(defconfig, dotconfig)

    staging_dir = 'install_staging'
    modi_dir = '%s' % staging_dir
    hdri_dir = '%s/usr' % staging_dir
    shutil.rmtree(os.path.join(dest_dir, staging_dir), ignore_errors=True)

    devnull = open('/dev/null', 'r')
    subprocess.check_call(['make', 'O=%s' % dest_dir,
        '%s_defconfig' % target], env=make_env, stdin=devnull)
    devnull.close()

    if not all_options.updateconfigs:
        # Build targets can be dependent upon the completion of previous
        # build targets, so build them one at a time.
        cmd_line = ['make',
            'INSTALL_HDR_PATH=%s' % hdri_dir,
            'INSTALL_MOD_PATH=%s' % modi_dir,
            'O=%s' % dest_dir]
        build_targets = []
        for c in make_command:
            if re.match(r'^-{1,2}\w', c):
                cmd_line.append(c)
            else:
                build_targets.append(c)
        for t in build_targets:
            build = Builder(log_name)

            result = build.run(cmd_line + [t])
            if result != 0:
                if all_options.keep_going:
                    failed_targets.append(target)
                    fail_or_error = error
                else:
                    fail_or_error = fail
                fail_or_error("Failed to build %s, see %s" %
                              (target, build.logname))

    # Copy the defconfig back.
    if all_options.configs or all_options.updateconfigs:
        devnull = open('/dev/null', 'r')
        subprocess.check_call(['make', 'O=%s' % dest_dir,
            'savedefconfig'], env=make_env, stdin=devnull)
        devnull.close()
        shutil.copyfile(savedefconfig, defconfig)

def build_many(allconf, targets):
    print "Building %d target(s)" % len(targets)
    for target in targets:
        if all_options.updateconfigs:
            update_config(allconf[target], all_options.updateconfigs)
        build(target)
    if failed_targets:
        fail('\n  '.join(["Failed targets:"] +
            [target for target in failed_targets]))

def main():
    global make_command

    check_kernel()
    check_build()

    configs = scan_configs()

    usage = ("""
           %prog [options] all                 -- Build all targets
           %prog [options] target target ...   -- List specific targets
           %prog [options] perf                -- Build all perf targets
           %prog [options] noperf              -- Build all non-perf targets""")
    parser = OptionParser(usage=usage, version=version)
    parser.add_option('--configs', action='store_true',
            dest='configs',
            help="Copy configs back into tree")
    parser.add_option('--list', action='store_true',
            dest='list',
            help='List available targets')
    parser.add_option('-v', '--verbose', action='store_true',
            dest='verbose',
            help='Output to stdout in addition to log file')
    parser.add_option('--oldconfig', action='store_true',
            dest='oldconfig',
            help='Only process "make oldconfig"')
    parser.add_option('--updateconfigs',
            dest='updateconfigs',
            help="Update defconfigs with provided option setting, "
                 "e.g. --updateconfigs=\'CONFIG_USE_THING=y\'")
    parser.add_option('-j', '--jobs', type='int', dest="jobs",
            help="Number of simultaneous jobs")
    parser.add_option('-l', '--load-average', type='int',
            dest='load_average',
            help="Don't start multiple jobs unless load is below LOAD_AVERAGE")
    parser.add_option('-k', '--keep-going', action='store_true',
            dest='keep_going', default=False,
            help="Keep building other targets if a target fails")
    parser.add_option('-m', '--make-target', action='append',
            help='Build the indicated make target (default: %s)' %
                 ' '.join(make_command))

    (options, args) = parser.parse_args()
    global all_options
    all_options = options

    if options.list:
        print "Available targets:"
        for target in configs.keys():
            print "   %s" % target
        sys.exit(0)

    if options.oldconfig:
        make_command = ["oldconfig"]
    elif options.make_target:
        make_command = options.make_target

    if options.jobs:
        make_command.append("-j%d" % options.jobs)
    if options.load_average:
        make_command.append("-l%d" % options.load_average)

    if args == ['all']:
        build_many(configs, configs.keys())
    elif args == ['perf']:
        targets = []
        for t in configs.keys():
            if "perf" in t:
                targets.append(t)
        build_many(configs, targets)
    elif args == ['noperf']:
        targets = []
        for t in configs.keys():
            if "perf" not in t:
                targets.append(t)
        build_many(configs, targets)
    elif len(args) > 0:
        targets = []
        for t in args:
            if t not in configs.keys():
                parser.error("Target '%s' not one of %s" % (t, configs.keys()))
            targets.append(t)
        build_many(configs, targets)
    else:
        parser.error("Must specify a target to build, or 'all'")

if __name__ == "__main__":
    main()
<commit_msg>build-all.py: Add support for mpq*_defconfig<commit_after>#! /usr/bin/env python

# Copyright (c) 2009-2013, The Linux Foundation. All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#     * Redistributions of source code must retain the above copyright
#       notice, this list of conditions and the following disclaimer.
#     * Redistributions in binary form must reproduce the above copyright
#       notice, this list of conditions and the following disclaimer in the
#       documentation and/or other materials provided with the distribution.
#     * Neither the name of The Linux Foundation nor
#       the names of its contributors may be used to endorse or promote
#       products derived from this software without specific prior written
#       permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
# NON-INFRINGEMENT ARE DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT OWNER OR
# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS;
# OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY,
# WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR
# OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF
# ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

# Build the kernel for all targets using the Android build environment.
#
# TODO: Accept arguments to indicate what to build.

import glob
from optparse import OptionParser
import subprocess
import os
import os.path
import re
import shutil
import sys

version = 'build-all.py, version 0.01'

build_dir = '../all-kernels'
make_command = ["vmlinux", "modules", "dtbs"]
make_env = os.environ
make_env.update({
        'ARCH': 'arm',
        'KCONFIG_NOTIMESTAMP': 'true' })
make_env.setdefault('CROSS_COMPILE', 'arm-none-linux-gnueabi-')
all_options = {}

def error(msg):
    sys.stderr.write("error: %s\n" % msg)

def fail(msg):
    """Fail with a user-printed message"""
    error(msg)
    sys.exit(1)

def check_kernel():
    """Ensure that PWD is a kernel directory"""
    if (not os.path.isfile('MAINTAINERS') or
        not os.path.isfile('arch/arm/mach-msm/Kconfig')):
        fail("This doesn't seem to be an MSM kernel dir")

def check_build():
    """Ensure that the build directory is present."""
    if not os.path.isdir(build_dir):
        try:
            os.makedirs(build_dir)
        except OSError as exc:
            if exc.errno == errno.EEXIST:
                pass
            else:
                raise

def update_config(file, str):
    print 'Updating %s with \'%s\'\n' % (file, str)
    defconfig = open(file, 'a')
    defconfig.write(str + '\n')
    defconfig.close()

def scan_configs():
    """Get the full list of defconfigs appropriate for this tree."""
    names = {}
    arch_pats = (
        r'[fm]sm[0-9]*_defconfig',
        r'apq*_defconfig',
        r'qsd*_defconfig',
	r'msmkrypton*_defconfig',
	r'mpq*_defconfig',
        )
    for p in arch_pats:
        for n in glob.glob('arch/arm/configs/' + p):
            names[os.path.basename(n)[:-10]] = n
    return names

class Builder:
    def __init__(self, logname):
        self.logname = logname
        self.fd = open(logname, 'w')

    def run(self, args):
        devnull = open('/dev/null', 'r')
        proc = subprocess.Popen(args, stdin=devnull,
                env=make_env,
                bufsize=0,
                stdout=subprocess.PIPE,
                stderr=subprocess.STDOUT)
        count = 0
        # for line in proc.stdout:
        rawfd = proc.stdout.fileno()
        while True:
            line = os.read(rawfd, 1024)
            if not line:
                break
            self.fd.write(line)
            self.fd.flush()
            if all_options.verbose:
                sys.stdout.write(line)
                sys.stdout.flush()
            else:
                for i in range(line.count('\n')):
                    count += 1
                    if count == 64:
                        count = 0
                        print
                    sys.stdout.write('.')
                sys.stdout.flush()
        print
        result = proc.wait()

        self.fd.close()
        return result

failed_targets = []

def build(target):
    dest_dir = os.path.join(build_dir, target)
    log_name = '%s/log-%s.log' % (build_dir, target)
    print 'Building %s in %s log %s' % (target, dest_dir, log_name)
    if not os.path.isdir(dest_dir):
        os.mkdir(dest_dir)
    defconfig = 'arch/arm/configs/%s_defconfig' % target
    dotconfig = '%s/.config' % dest_dir
    savedefconfig = '%s/defconfig' % dest_dir
    shutil.copyfile(defconfig, dotconfig)

    staging_dir = 'install_staging'
    modi_dir = '%s' % staging_dir
    hdri_dir = '%s/usr' % staging_dir
    shutil.rmtree(os.path.join(dest_dir, staging_dir), ignore_errors=True)

    devnull = open('/dev/null', 'r')
    subprocess.check_call(['make', 'O=%s' % dest_dir,
        '%s_defconfig' % target], env=make_env, stdin=devnull)
    devnull.close()

    if not all_options.updateconfigs:
        # Build targets can be dependent upon the completion of previous
        # build targets, so build them one at a time.
        cmd_line = ['make',
            'INSTALL_HDR_PATH=%s' % hdri_dir,
            'INSTALL_MOD_PATH=%s' % modi_dir,
            'O=%s' % dest_dir]
        build_targets = []
        for c in make_command:
            if re.match(r'^-{1,2}\w', c):
                cmd_line.append(c)
            else:
                build_targets.append(c)
        for t in build_targets:
            build = Builder(log_name)

            result = build.run(cmd_line + [t])
            if result != 0:
                if all_options.keep_going:
                    failed_targets.append(target)
                    fail_or_error = error
                else:
                    fail_or_error = fail
                fail_or_error("Failed to build %s, see %s" %
                              (target, build.logname))

    # Copy the defconfig back.
    if all_options.configs or all_options.updateconfigs:
        devnull = open('/dev/null', 'r')
        subprocess.check_call(['make', 'O=%s' % dest_dir,
            'savedefconfig'], env=make_env, stdin=devnull)
        devnull.close()
        shutil.copyfile(savedefconfig, defconfig)

def build_many(allconf, targets):
    print "Building %d target(s)" % len(targets)
    for target in targets:
        if all_options.updateconfigs:
            update_config(allconf[target], all_options.updateconfigs)
        build(target)
    if failed_targets:
        fail('\n  '.join(["Failed targets:"] +
            [target for target in failed_targets]))

def main():
    global make_command

    check_kernel()
    check_build()

    configs = scan_configs()

    usage = ("""
           %prog [options] all                 -- Build all targets
           %prog [options] target target ...   -- List specific targets
           %prog [options] perf                -- Build all perf targets
           %prog [options] noperf              -- Build all non-perf targets""")
    parser = OptionParser(usage=usage, version=version)
    parser.add_option('--configs', action='store_true',
            dest='configs',
            help="Copy configs back into tree")
    parser.add_option('--list', action='store_true',
            dest='list',
            help='List available targets')
    parser.add_option('-v', '--verbose', action='store_true',
            dest='verbose',
            help='Output to stdout in addition to log file')
    parser.add_option('--oldconfig', action='store_true',
            dest='oldconfig',
            help='Only process "make oldconfig"')
    parser.add_option('--updateconfigs',
            dest='updateconfigs',
            help="Update defconfigs with provided option setting, "
                 "e.g. --updateconfigs=\'CONFIG_USE_THING=y\'")
    parser.add_option('-j', '--jobs', type='int', dest="jobs",
            help="Number of simultaneous jobs")
    parser.add_option('-l', '--load-average', type='int',
            dest='load_average',
            help="Don't start multiple jobs unless load is below LOAD_AVERAGE")
    parser.add_option('-k', '--keep-going', action='store_true',
            dest='keep_going', default=False,
            help="Keep building other targets if a target fails")
    parser.add_option('-m', '--make-target', action='append',
            help='Build the indicated make target (default: %s)' %
                 ' '.join(make_command))

    (options, args) = parser.parse_args()
    global all_options
    all_options = options

    if options.list:
        print "Available targets:"
        for target in configs.keys():
            print "   %s" % target
        sys.exit(0)

    if options.oldconfig:
        make_command = ["oldconfig"]
    elif options.make_target:
        make_command = options.make_target

    if options.jobs:
        make_command.append("-j%d" % options.jobs)
    if options.load_average:
        make_command.append("-l%d" % options.load_average)

    if args == ['all']:
        build_many(configs, configs.keys())
    elif args == ['perf']:
        targets = []
        for t in configs.keys():
            if "perf" in t:
                targets.append(t)
        build_many(configs, targets)
    elif args == ['noperf']:
        targets = []
        for t in configs.keys():
            if "perf" not in t:
                targets.append(t)
        build_many(configs, targets)
    elif len(args) > 0:
        targets = []
        for t in args:
            if t not in configs.keys():
                parser.error("Target '%s' not one of %s" % (t, configs.keys()))
            targets.append(t)
        build_many(configs, targets)
    else:
        parser.error("Must specify a target to build, or 'all'")

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