<commit_before>import numpy as np
import pandas as pd
import re

class Data:
    def __init__(self, filepath):
        table_names = ['all_encounter_data', 'demographics', 'encounters',
                       'family_hist_for_Enc', 'family_hist_list',
                       'ICD_for_Enc', 'macula_findings_for_Enc',
                       'SL_Lens_for_Enc', 'SNOMED_problem_list',
                       'systemic_disease_for_Enc', 'systemic_disease_list']
        self.__data = {}
        self.__normdata = {"all_encounter_data" : None,
                           "all_person_data" : None}
        for name in table_names:
            self.__data[name] = pd.read_pickle(filepath + name + '.pickle')

    def __getitem__(self,name):
        if self.__normdata.has_key(name):
            if self.__normdata.get(name) is None:
                if name == "all_encounter_data": self.create_enc_table()
                if name == "all_person_data": self.create_person_table()
            return self.__normdata[name]

        return self.__data.__getitem__(name)

    def get_underlying(self):
        return self.__data

    def create_enc_table(self):
        # Helper function to convert string values to floats
        # 87 Glucose column have BP values which this function ignores and assigns NaN
        # TODO: Swap the erroneous column values. Filter out outliers
        NaN = float("NaN")
        def to_float(x):
            try:
                return float(x)
            except ValueError:
                return NaN

        # Split up the BP field into Systolic and Diastolic readings
        d_enc = self.__data["all_encounter_data"].drop(["Enc_ID","Person_ID"], axis=1)
        pattern = re.compile("(?P<BP_Systolic>\d+)\s*\/\s*(?P<BP_Diastolic>\d+)")
        d_enc = pd.merge(d_enc, d_enc["BP"].str.extract(pattern, expand=True),
                         left_index=True, right_index=True).drop("BP", axis=1)

        # Drop columns with multiple values for a single Enc_Nbr
        # When multiple values are observed, take the mean and merge them back into the complete table
        columns = ["A1C","BMI","BP_Systolic","BP_Diastolic","Glucose"]
        self.__normdata["all_encounter_data"] = \
            pd.merge(d_enc.drop(columns, axis=1).drop_duplicates().set_index("Enc_Nbr"),
                     d_enc.groupby("Enc_Nbr")[columns].agg(lambda x: np.mean(x.map(to_float))),
                     left_index=True, right_index=True)

        # Add diagnoses to table
        dEI = self.__data["ICD_for_Enc"].loc[:,["Enc_Nbr", "Diagnosis_Code_ID"]]

        # Manually fix erroneous values
        dEI.loc[90899]="367.4"
        dEI.loc[168442]="362.3"

        # Diabetes is under 250.* and 362.0.* for ICD9 and E08,E09,E10,E11,E13 for ICD10
        dEI["DM"]=dEI["Diagnosis_Code_ID"].str.contains("^250.*|^362\.0.*|^E(?:0[89]|1[013])(?:\.[A-Z0-9]{1,4})?")

        # Macular edema is under 362.07 for ICD9 and E(08|09|10|11|13).3([1-5]1|7) for ICD10
        dEI["ME"]=dEI["Diagnosis_Code_ID"].str.contains("^362\.07|^E(?:0[89]|1[013])\.3(?:[1-5]1|7).*")

        # Mild Nonproliferative Diabetic Retinopathy is under 362.04 for ICD9 and E(08|09|10|11|13).32
        dEI["mNPDR"]=dEI["Diagnosis_Code_ID"].str.contains("^362\.04|^E(?:0[89]|1[013])\.32.*")

        # Moderate Nonproliferative Diabetic Retinopathy is under 362.05 for ICD9 and E(08|09|10|11|13).33
        dEI["MNPDR"]=dEI["Diagnosis_Code_ID"].str.contains("^362\.05|^E(?:0[89]|1[013])\.33.*")

        # Severe Nonproliferative Diabetic Retinopathy is under 362.06 for ICD9 and E(08|09|10|11|13).34
        dEI["SNPDR"]=dEI["Diagnosis_Code_ID"].str.contains("^362\.06|^E(?:0[89]|1[013])\.34.*")

        # Proliferative Diabetic Retinopathy is under 362.02 for ICD9 and E(08|09|10|11|13).35
        dEI["PDR"]=dEI["Diagnosis_Code_ID"].str.contains("^362\.02|^E(?:0[89]|1[013])\.35.*")

        # Diabetic Retinopathy
        dEI["DR"]=dEI["mNPDR"] & dEI["MNPDR"] & dEI["SNPDR"] & dEI["PDR"]

        # Glaucoma Suspect is under 365.0 for ICD9 and H40.0 for ICD10
        dEI["Glaucoma_Suspect"]=dEI["Diagnosis_Code_ID"].str.contains("^365\.0.*|^H40\.0.*")

        # Open-angle Glaucoma is under 365.1 for ICD9 and H40.1 for ICD10
        dEI["Open_angle_Glaucoma"]=dEI["Diagnosis_Code_ID"].str.contains("^365\.1.*|^H40\.1.*")

        # Cataract is under 366 for ICD9 and H25 and H26 for ICD10
        dEI["Cataract"]=dEI["Diagnosis_Code_ID"].str.contains("^366(?:\.\d{1,2})?|^H2[56](?:\.[A-Z0-9]{1,4})?")

        columns = ["DM","DR","ME","MNPDR","PDR","SNPDR","mNPDR","Glaucoma_Suspect",
                   "Open_angle_Glaucoma","Cataract"]
        self.__normdata["all_encounter_data"] = \
            pd.merge(self.__normdata["all_encounter_data"],
                     dEI.groupby("Enc_Nbr")[columns].any(),
                     left_index=True, right_index=True)

    def create_person_table(self):
        d_enc = self["all_encounter_data"]

        # Average the past year of data
        f = lambda x: x[x["Enc_Date"]>=x["Enc_Date"].max() - pd.DateOffset(years=1)].drop(["Person_Nbr","Enc_Date"],axis=1).mean()
        columns = ["Enc_Date", "Person_Nbr", "A1C", "BMI", "Glucose", "BP_Systolic", "BP_Diastolic"]
        self.__normdata["all_person_data"] = \
            pd.merge(self.__data["demographics"].set_index("Person_Nbr"),
                     d_enc.loc[:,columns].groupby("Person_Nbr").apply(f),
                     left_index=True, right_index=True)

        # Combine all diagnoses
        columns = ["DM","DR","ME","MNPDR","PDR","SNPDR","mNPDR","Glaucoma_Suspect",
                   "Open_angle_Glaucoma","Cataract"]
        self.__normdata["all_person_data"] = \
            pd.merge(self.__normdata["all_person_data"],
                     d_enc.groupby("Person_Nbr")[columns].any(),
                     left_index=True, right_index=True)
<commit_msg>Further code cleanup. Removed DR column<commit_after>import numpy as np
import pandas as pd
import re

class Data:
    def __init__(self, filepath):
        table_names = ['all_encounter_data', 'demographics', 'encounters',
                       'family_hist_for_Enc', 'family_hist_list',
                       'ICD_for_Enc', 'macula_findings_for_Enc',
                       'SL_Lens_for_Enc', 'SNOMED_problem_list',
                       'systemic_disease_for_Enc', 'systemic_disease_list']
        self.__data = {}
        self.__normdata = {"all_encounter_data" : None,
                           "all_person_data" : None}
        for name in table_names:
            self.__data[name] = pd.read_pickle(filepath + name + '.pickle')

    def __getitem__(self,name):
        if self.__normdata.has_key(name):
            if self.__normdata.get(name) is None:
                if name == "all_encounter_data": self.create_enc_table()
                if name == "all_person_data": self.create_person_table()
            return self.__normdata[name]

        return self.__data.__getitem__(name)

    def get_underlying(self):
        return self.__data

    def create_enc_table(self):
        # Helper function to convert string values to floats
        # 87 Glucose column have BP values which this function ignores and assigns NaN
        # TODO: Swap the erroneous column values. Filter out outliers
        NaN = float("NaN")
        def to_float(x):
            try:
                return float(x)
            except ValueError:
                return NaN

        # Split up the BP field into Systolic and Diastolic readings
        d_enc = self.__data["all_encounter_data"].drop(["Enc_ID","Person_ID"], axis=1)
        pattern = re.compile("(?P<BP_Systolic>\d+)\s*\/\s*(?P<BP_Diastolic>\d+)")
        d_enc = pd.merge(d_enc, d_enc["BP"].str.extract(pattern, expand=True),
                         left_index=True, right_index=True).drop("BP", axis=1)

        # Drop columns with multiple values for a single Enc_Nbr
        # When multiple values are observed, take the mean and merge them back into the complete table
        columns = ["A1C","BMI","BP_Systolic","BP_Diastolic","Glucose"]
        self.__normdata["all_encounter_data"] = \
            pd.merge(d_enc.drop(columns, axis=1).drop_duplicates().set_index("Enc_Nbr"),
                     d_enc.groupby("Enc_Nbr")[columns].agg(lambda x: np.mean(x.map(to_float))),
                     left_index=True, right_index=True)

        # Add diagnoses to table
        dEI = self.__data["ICD_for_Enc"].loc[:,["Enc_Nbr", "Diagnosis_Code_ID"]]

        # Manually fix erroneous values
        dEI.loc[90899]="367.4"
        dEI.loc[168442]="362.3"

        diagnoses = {
            # Diabetes is under 250.* and 362.0.* for ICD9 and E08,E09,E10,E11,E13 for ICD10
            "DM" : "^250.*|^362\.0.*|^E(?:0[89]|1[013])(?:\.[A-Z0-9]{1,4})?",
            # Macular edema is under 362.07 for ICD9 and E(08|09|10|11|13).3([1-5]1|7) for ICD10
            "ME" : "^362\.07|^E(?:0[89]|1[013])\.3(?:[1-5]1|7).*",
            # Mild Nonproliferative Diabetic Retinopathy is under 362.04 for ICD9 and E(08|09|10|11|13).32
            "mNPDR" : "^362\.04|^E(?:0[89]|1[013])\.32.*",
            # Moderate Nonproliferative Diabetic Retinopathy is under 362.05 for ICD9 and E(08|09|10|11|13).33
            "MNPDR" : "^362\.05|^E(?:0[89]|1[013])\.33.*",
            # Severe Nonproliferative Diabetic Retinopathy is under 362.06 for ICD9 and E(08|09|10|11|13).34
            "SNPDR" : "^362\.06|^E(?:0[89]|1[013])\.34.*",
            # Proliferative Diabetic Retinopathy is under 362.02 for ICD9 and E(08|09|10|11|13).35
            "PDR" : "^362\.02|^E(?:0[89]|1[013])\.35.*",
            # Glaucoma Suspect is under 365.0 for ICD9 and H40.0 for ICD10
            "Glaucoma_Suspect" : "^365\.0.*|^H40\.0.*",
            # Open-angle Glaucoma is under 365.1 for ICD9 and H40.1 for ICD10
            "Open_angle_Glaucoma" : "^365\.1.*|^H40\.1.*",
            # Cataract is under 366 for ICD9 and H25 and H26 for ICD10
            "Cataract" : "^366(?:\.\d{1,2})?|^H2[56](?:\.[A-Z0-9]{1,4})?"
        }
        for diagnosis, pattern in diagnoses.iteritems():
            dEI[diagnosis]=dEI["Diagnosis_Code_ID"].str.contains(pattern)

        self.__normdata["all_encounter_data"] = \
            pd.merge(self.__normdata["all_encounter_data"],
                     dEI.groupby("Enc_Nbr")[list(diagnoses)].any(),
                     left_index=True, right_index=True)

    def create_person_table(self):
        d_enc = self["all_encounter_data"]

        # Average the past year of data
        f = lambda x: x[x["Enc_Date"]>=x["Enc_Date"].max() - pd.DateOffset(years=1)].drop(["Person_Nbr","Enc_Date"],axis=1).mean()
        columns = ["Enc_Date", "Person_Nbr", "A1C", "BMI", "Glucose", "BP_Systolic", "BP_Diastolic"]
        self.__normdata["all_person_data"] = \
            pd.merge(self.__data["demographics"].set_index("Person_Nbr"),
                     d_enc.loc[:,columns].groupby("Person_Nbr").apply(f),
                     left_index=True, right_index=True)

        # Combine all diagnoses
        columns = ["DM","ME","MNPDR","PDR","SNPDR","mNPDR","Glaucoma_Suspect",
                   "Open_angle_Glaucoma","Cataract"]
        self.__normdata["all_person_data"] = \
            pd.merge(self.__normdata["all_person_data"],
                     d_enc.groupby("Person_Nbr")[columns].any(),
                     left_index=True, right_index=True)
<|endoftext|>


<commit_before># coding=utf-8

import sys
import time

from config import MSG_DATA
from config import LOG_FILE_HANDLE

def save_to_log_file(message):
    LOG_FILE_HANDLE.write(message)

def console(message):
    sys.stdout.write(message + "\n")

def write(message):

    is_msg_data = False

    for msg in MSG_DATA:
        if msg == message:
            is_msg_data = True
            break

    if is_msg_data:
        message = MSG_DATA[message]

    message = time.strftime('%Y-%m-%d %H:%M:%S') + ": " + message + "\n"

    if is_msg_data:
        sys.stdout.write(message)

    save_to_log_file(message)

def generate_log_file_name():

    file_name = "network-diagnostic-"
    date_str = time.strftime('%Y-%m-%d-%H_%M_%S')

    return file_name + date_str + ".log"

def open_file(file_name):
    global LOG_FILE_HANDLE
    LOG_FILE_HANDLE = open(file_name, "a+")

def close_file():
    LOG_FILE_HANDLE.close()

<commit_msg>Improve log.<commit_after># coding=utf-8

import sys
import time

from config import *

def save_to_log_file(message):
    global LOG_FILE_HANDLE
    LOG_FILE_HANDLE.write(message)

def console(message):
    sys.stdout.write(message + "\n")

def write(message):

    is_msg_data = False

    for msg in MSG_DATA:
        if msg == message:
            is_msg_data = True
            break

    if is_msg_data:
        message = MSG_DATA[message]

    message = time.strftime('%Y-%m-%d %H:%M:%S') + ": " + message + "\n"

    if is_msg_data:
        sys.stdout.write(message)

    save_to_log_file(message)

def generate_log_file_name():

    file_name = "network-diagnostic-"
    date_str = time.strftime('%Y-%m-%d-%H_%M_%S')

    return file_name + date_str + ".log"

def open_file(file_name):
    global LOG_FILE_HANDLE
    LOG_FILE_HANDLE = open(file_name, "a+")

def close_file():
    global LOG_FILE_HANDLE
    LOG_FILE_HANDLE.close()
    LOG_FILE_HANDLE = None

<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
"""
functions to use GEE within Qgis python script
"""
import math
import ee

from qgis.core import QgsCoordinateReferenceSystem, QgsCoordinateTransform, QgsProject, QgsPointXY, QgsRectangle
from qgis.utils import iface

import ee_plugin.utils
import ee_plugin.ee_auth

# init the Google Earth Engine user authorization system
ee_plugin.ee_auth.init()

def addLayer(eeObject, visParams=None, name=None, shown=True, opacity=1.0):
    """
        Adds a given EE object to the map as a layer.

        https://developers.google.com/earth-engine/api_docs#mapaddlayer

        Uses:
            >>> from ee_plugin import Map
            >>> Map.addLayer(.....)
    """

    ee_plugin.utils.add_or_update_ee_layer(eeObject, visParams, name, shown, opacity)

def centerObject(feature, zoom=None):
    """
        Centers the map view on a given object.

        https://developers.google.com/earth-engine/api_docs#mapcenterobject

        Uses:
            >>> from ee_plugin import Map
            >>> Map.centerObject(feature)
    """

    feature = ee.Feature(feature)

    if not zoom:
        # make sure our geometry is in geo
        rect = feature.geometry().transform(ee.Projection('EPSG:4326'), 1)

        # get coordinates
        coords = rect.bounds().getInfo()['coordinates'][0]
        xmin = coords[0][0]
        ymin = coords[0][1]
        xmax = coords[2][0]
        ymax = coords[2][1]

        # construct QGIS geometry
        rect = QgsRectangle(xmin, ymin, xmax, ymax)

        # transform rect to a crs used by current project
        crs_src = QgsCoordinateReferenceSystem(4326)
        crs_dst = QgsCoordinateReferenceSystem(QgsProject.instance().crs())
        geo2proj = QgsCoordinateTransform(crs_src, crs_dst, QgsProject.instance())
        rect_proj = geo2proj.transform(rect)

        # center geometry
        iface.mapCanvas().zoomToFeatureExtent(rect_proj)
    else:
        # set map center to feature centroid at a specified zoom
        center = feature.geometry().centroid().coordinates().getInfo()
        setCenter(center[0], center[1], zoom)

def getBounds(asGeoJSON=False):
    """
        Returns the bounds of the current map view, as a list in the format [west, south, east, north] in degrees.

        https://developers.google.com/earth-engine/api_docs#mapgetbounds

        Uses:
            >>> from ee_plugin import Map
            >>> bounds = Map.getBounds(True)
            >>> Map.addLayer(bounds, {}, 'bounds')
    """
    ex = iface.mapCanvas().extent()
    # return ex
    xmax = ex.xMaximum()
    ymax = ex.yMaximum()
    xmin = ex.xMinimum()
    ymin = ex.yMinimum()

    # return as [west, south, east, north]
    if not asGeoJSON:
        return [xmin, ymin, xmax, ymax]

    # return as geometry
    crs = iface.mapCanvas().mapSettings().destinationCrs().authid()

    return ee.Geometry.Rectangle([xmin, ymin, xmax, ymax], crs, False)

def getCenter():
    """
        Returns the coordinates at the center of the map.

        https://developers.google.com/earth-engine/api_docs#mapgetcenter

        Uses:
            >>> from ee_plugin import Map
            >>> center = Map.getCenter()
            >>> Map.addLayer(center, { 'color': 'red' }, 'center')
    """
    center = iface.mapCanvas().center()

    crs = iface.mapCanvas().mapSettings().destinationCrs().authid()

    return ee.Geometry.Point([center.x(), center.y()], crs)

def setCenter(lon, lat, zoom=None):
    """
        Centers the map view at the given coordinates with the given zoom level. If no zoom level is provided, it uses the most recent zoom level on the map.

        https://developers.google.com/earth-engine/api_docs#mapsetcenter

        Uses:
            >>> from ee_plugin import Map
            >>> Map.setCenter(lon, lat, zoom)
    """

    ### center
    center_point_in = QgsPointXY(lon, lat)
    # convert coordinates
    crsSrc = QgsCoordinateReferenceSystem(4326)  # WGS84
    crsDest = QgsCoordinateReferenceSystem(QgsProject.instance().crs())
    xform = QgsCoordinateTransform(crsSrc, crsDest, QgsProject.instance())
    # forward transformation: src -> dest
    center_point = xform.transform(center_point_in)
    iface.mapCanvas().setCenter(center_point)

    ### zoom
    if zoom is not None:
        # transform the zoom level to scale
        scale_value = 591657550.5 / 2 ** (zoom - 1)
        iface.mapCanvas().zoomScale(scale_value)

def getScale():
    """
        Returns the approximate pixel scale of the current map view, in meters.

        https://developers.google.com/earth-engine/api_docs#mapgetscale

        Uses:
            >>> from ee_plugin import Map
            >>> print(Map.getScale())
    """

    return iface.mapCanvas().scale() / 1000

def getZoom():
    """
        Returns the current zoom level of the map.

        https://developers.google.com/earth-engine/api_docs#mapgetzoom

        Note that in QGIS zoom is a floating point number

        Uses:
            >>> from ee_plugin import Map
            >>> print(Map.getZoom())
    """

    # from https://gis.stackexchange.com/questions/268890/get-current-zoom-level-from-qgis-map-canvas
    scale = iface.mapCanvas().scale()
    dpi = iface.mainWindow().physicalDpiX()
    maxScalePerPixel = 156543.04
    inchesPerMeter = 39.37
    zoom = math.log((dpi * inchesPerMeter * maxScalePerPixel / scale), 2)

    return zoom

def setZoom(zoom):
    """
        Sets the zoom level of the map.

        https://developers.google.com/earth-engine/api_docs#mapsetzoom

        Uses:
            >>> from ee_plugin import Map
            >>> Map.setZoom(15)
    """

    center = getCenter()
    centerObject(ee.Feature(center), zoom)
<commit_msg>Fix the issue to center the object for some specific EE types<commit_after># -*- coding: utf-8 -*-
"""
functions to use GEE within Qgis python script
"""
import math
import ee

from qgis.core import QgsCoordinateReferenceSystem, QgsCoordinateTransform, QgsProject, QgsPointXY, QgsRectangle
from qgis.utils import iface

import ee_plugin.utils
import ee_plugin.ee_auth

# init the Google Earth Engine user authorization system
ee_plugin.ee_auth.init()

def addLayer(eeObject, visParams=None, name=None, shown=True, opacity=1.0):
    """
        Adds a given EE object to the map as a layer.

        https://developers.google.com/earth-engine/api_docs#mapaddlayer

        Uses:
            >>> from ee_plugin import Map
            >>> Map.addLayer(.....)
    """

    ee_plugin.utils.add_or_update_ee_layer(eeObject, visParams, name, shown, opacity)

def centerObject(feature, zoom=None):
    """
        Centers the map view on a given object.

        https://developers.google.com/earth-engine/api_docs#mapcenterobject

        Uses:
            >>> from ee_plugin import Map
            >>> Map.centerObject(feature)
    """

    if not hasattr(feature, 'geometry'):
        feature = ee.Feature(feature)

    if not zoom:
        # make sure our geometry is in geo
        rect = feature.geometry().transform(ee.Projection('EPSG:4326'), 1)

        # get coordinates
        coords = rect.bounds().getInfo()['coordinates'][0]
        xmin = coords[0][0]
        ymin = coords[0][1]
        xmax = coords[2][0]
        ymax = coords[2][1]

        # construct QGIS geometry
        rect = QgsRectangle(xmin, ymin, xmax, ymax)

        # transform rect to a crs used by current project
        crs_src = QgsCoordinateReferenceSystem(4326)
        crs_dst = QgsCoordinateReferenceSystem(QgsProject.instance().crs())
        geo2proj = QgsCoordinateTransform(crs_src, crs_dst, QgsProject.instance())
        rect_proj = geo2proj.transform(rect)

        # center geometry
        iface.mapCanvas().zoomToFeatureExtent(rect_proj)
    else:
        # set map center to feature centroid at a specified zoom
        center = feature.geometry().centroid().coordinates().getInfo()
        setCenter(center[0], center[1], zoom)

def getBounds(asGeoJSON=False):
    """
        Returns the bounds of the current map view, as a list in the format [west, south, east, north] in degrees.

        https://developers.google.com/earth-engine/api_docs#mapgetbounds

        Uses:
            >>> from ee_plugin import Map
            >>> bounds = Map.getBounds(True)
            >>> Map.addLayer(bounds, {}, 'bounds')
    """
    ex = iface.mapCanvas().extent()
    # return ex
    xmax = ex.xMaximum()
    ymax = ex.yMaximum()
    xmin = ex.xMinimum()
    ymin = ex.yMinimum()

    # return as [west, south, east, north]
    if not asGeoJSON:
        return [xmin, ymin, xmax, ymax]

    # return as geometry
    crs = iface.mapCanvas().mapSettings().destinationCrs().authid()

    return ee.Geometry.Rectangle([xmin, ymin, xmax, ymax], crs, False)

def getCenter():
    """
        Returns the coordinates at the center of the map.

        https://developers.google.com/earth-engine/api_docs#mapgetcenter

        Uses:
            >>> from ee_plugin import Map
            >>> center = Map.getCenter()
            >>> Map.addLayer(center, { 'color': 'red' }, 'center')
    """
    center = iface.mapCanvas().center()

    crs = iface.mapCanvas().mapSettings().destinationCrs().authid()

    return ee.Geometry.Point([center.x(), center.y()], crs)

def setCenter(lon, lat, zoom=None):
    """
        Centers the map view at the given coordinates with the given zoom level. If no zoom level is provided, it uses the most recent zoom level on the map.

        https://developers.google.com/earth-engine/api_docs#mapsetcenter

        Uses:
            >>> from ee_plugin import Map
            >>> Map.setCenter(lon, lat, zoom)
    """

    ### center
    center_point_in = QgsPointXY(lon, lat)
    # convert coordinates
    crsSrc = QgsCoordinateReferenceSystem(4326)  # WGS84
    crsDest = QgsCoordinateReferenceSystem(QgsProject.instance().crs())
    xform = QgsCoordinateTransform(crsSrc, crsDest, QgsProject.instance())
    # forward transformation: src -> dest
    center_point = xform.transform(center_point_in)
    iface.mapCanvas().setCenter(center_point)

    ### zoom
    if zoom is not None:
        # transform the zoom level to scale
        scale_value = 591657550.5 / 2 ** (zoom - 1)
        iface.mapCanvas().zoomScale(scale_value)

def getScale():
    """
        Returns the approximate pixel scale of the current map view, in meters.

        https://developers.google.com/earth-engine/api_docs#mapgetscale

        Uses:
            >>> from ee_plugin import Map
            >>> print(Map.getScale())
    """

    return iface.mapCanvas().scale() / 1000

def getZoom():
    """
        Returns the current zoom level of the map.

        https://developers.google.com/earth-engine/api_docs#mapgetzoom

        Note that in QGIS zoom is a floating point number

        Uses:
            >>> from ee_plugin import Map
            >>> print(Map.getZoom())
    """

    # from https://gis.stackexchange.com/questions/268890/get-current-zoom-level-from-qgis-map-canvas
    scale = iface.mapCanvas().scale()
    dpi = iface.mainWindow().physicalDpiX()
    maxScalePerPixel = 156543.04
    inchesPerMeter = 39.37
    zoom = math.log((dpi * inchesPerMeter * maxScalePerPixel / scale), 2)

    return zoom

def setZoom(zoom):
    """
        Sets the zoom level of the map.

        https://developers.google.com/earth-engine/api_docs#mapsetzoom

        Uses:
            >>> from ee_plugin import Map
            >>> Map.setZoom(15)
    """

    center = getCenter()
    centerObject(ee.Feature(center), zoom)
<|endoftext|>


<commit_before>import hashlib

## DATATYPES

# Unlike Git, which has Blobs and Trees; Rev has Atoms, Dictionaries, Lists
# and Tuples (the latter three corresponding to their Python equivalents)
#
# Other datatypes can be defined just by subclassing NodeBase and implementing
# shrink and expand.

class NodeBase:

    def __init__(self, repo, content):
        self.repo = repo
        self.content = self.shrink(content)

    def __bytes__(self):
        return ("%s\n%r" % (self.__class__, self.content)).encode("utf-8")

class Atom(NodeBase):

    def shrink(self, content):
        return content

    def expand(self):
        return self.content

class Dictionary(NodeBase):

    def shrink(self, content):
        return {
            self.repo.shrink(key): self.repo.shrink(value)
            for key, value in content.items()
        }

    def expand(self):
        return {
            self.repo.expand(key): self.repo.expand(value)
            for key, value in self.content.items()
        }

class List(NodeBase):

    def shrink(self, content):
        return [self.repo.shrink(item) for item in content]

    def expand(self):
        return [self.repo.expand(item) for item in self.content]

class Tuple(NodeBase):

    def shrink(self, content):
        return tuple(self.repo.shrink(item) for item in content)

    def expand(self):
        return tuple(self.repo.expand(item) for item in self.content)

## END DATATYPES

class Commit:

    def __init__(self, repo, obj_sha, message, parents):
        # @@@ doesn't store author/committer/timestamp yet
        self.repo = repo
        self.obj_sha = obj_sha
        self.message = message
        self.parents = parents
        self.content = repo.shrink((obj_sha, parents, message))

    def __bytes__(self):
        return ("%s\n%r" % (self.__class__, self.content)).encode("utf-8")

class Repo:

    def __init__(self):
        self.objects = {}
        self.refs = {}
        self.HEAD = "master"

    def store(self, obj):
        sha = hashlib.sha1(bytes(obj)).hexdigest()
        self.objects[sha] = obj

        return sha

    def shrink(self, content):
        if isinstance(content, dict):
            return self.store(Dictionary(self, content))
        elif isinstance(content, list):
            return self.store(List(self, content))
        elif isinstance(content, tuple):
            return self.store(Tuple(self, content))
        else:
            return self.store(Atom(self, content))

        return store(obj)

    def get_object(self, sha):
        """
        get the shallow version of the object with the given SHA.

        In other words, the components of the given object (if non-atomic)
        will be further SHA references.
        """
        return self.objects[sha]

    def expand(self, sha):
        """
        get the deep version of the object with the given SHA.

        In other words, all components of the given object will be expanded
        recursively to their actual values with no SHA references left.
        """
        return self.get_object(sha).expand()

    def create_commit(self, obj_sha, message, parents=None):
        if parents is None:
            parents = []

        return self.store(Commit(self, obj_sha, message, parents))

    def commit(self, obj, message):
        obj_sha = self.shrink(obj)
        commit_sha = self.create_commit(obj_sha, message, parents=[self.refs[self.HEAD]])
        self.refs[self.HEAD] = commit_sha

        return commit_sha

    def retrieve_commit(self, commit_sha=None):
        if commit_sha is None:
            commit_sha = self.refs[self.HEAD]

        return self.expand(self.get_object(commit_sha).obj_sha)

    def create_branch(self, branch_name, commit_sha=None):
        if commit_sha is None:
            commit_sha = self.refs[self.HEAD]

        self.refs[branch_name] = commit_sha or self.refs[self.HEAD]

    def checkout_branch(self, branch_name):
        self.HEAD = branch_name
        # @@@ do anything to index?

<commit_msg>refactored commit a little and added note about thread safety<commit_after>import hashlib

## DATATYPES

# Unlike Git, which has Blobs and Trees; Rev has Atoms, Dictionaries, Lists
# and Tuples (the latter three corresponding to their Python equivalents)
#
# Other datatypes can be defined just by subclassing NodeBase and implementing
# shrink and expand.

class NodeBase:

    def __init__(self, repo, content):
        self.repo = repo
        self.content = self.shrink(content)

    def __bytes__(self):
        return ("%s\n%r" % (self.__class__, self.content)).encode("utf-8")

class Atom(NodeBase):

    def shrink(self, content):
        return content

    def expand(self):
        return self.content

class Dictionary(NodeBase):

    def shrink(self, content):
        return {
            self.repo.shrink(key): self.repo.shrink(value)
            for key, value in content.items()
        }

    def expand(self):
        return {
            self.repo.expand(key): self.repo.expand(value)
            for key, value in self.content.items()
        }

class List(NodeBase):

    def shrink(self, content):
        return [self.repo.shrink(item) for item in content]

    def expand(self):
        return [self.repo.expand(item) for item in self.content]

class Tuple(NodeBase):

    def shrink(self, content):
        return tuple(self.repo.shrink(item) for item in content)

    def expand(self):
        return tuple(self.repo.expand(item) for item in self.content)

## END DATATYPES

class Commit:

    def __init__(self, repo, obj_sha, message, parents):
        # @@@ doesn't store author/committer/timestamp yet
        self.repo = repo
        self.obj_sha = obj_sha
        self.message = message
        self.parents = parents
        self.content = repo.shrink((obj_sha, parents, message))

    def __bytes__(self):
        return ("%s\n%r" % (self.__class__, self.content)).encode("utf-8")

class Repo:

    def __init__(self):
        self.objects = {}
        self.refs = {}
        self.HEAD = "master"

    def store(self, obj):
        sha = hashlib.sha1(bytes(obj)).hexdigest()
        self.objects[sha] = obj

        return sha

    def shrink(self, content):
        if isinstance(content, dict):
            return self.store(Dictionary(self, content))
        elif isinstance(content, list):
            return self.store(List(self, content))
        elif isinstance(content, tuple):
            return self.store(Tuple(self, content))
        else:
            return self.store(Atom(self, content))

        return store(obj)

    def get_object(self, sha):
        """
        get the shallow version of the object with the given SHA.

        In other words, the components of the given object (if non-atomic)
        will be further SHA references.
        """
        return self.objects[sha]

    def expand(self, sha):
        """
        get the deep version of the object with the given SHA.

        In other words, all components of the given object will be expanded
        recursively to their actual values with no SHA references left.
        """
        return self.get_object(sha).expand()

    def create_commit(self, obj_sha, message, parents=None):
        if parents is None:
            parents = []

        return self.store(Commit(self, obj_sha, message, parents))

    def commit(self, obj, message):
        obj_sha = self.shrink(obj)
        old_head = self.refs[self.HEAD]
        commit_sha = self.create_commit(obj_sha, message, parents=[old_head])
        # @@@ make thread-safe? HEAD could change during commit
        self.refs[self.HEAD] = commit_sha

        return commit_sha

    def retrieve_commit(self, commit_sha=None):
        if commit_sha is None:
            commit_sha = self.refs[self.HEAD]

        return self.expand(self.get_object(commit_sha).obj_sha)

    def create_branch(self, branch_name, commit_sha=None):
        if commit_sha is None:
            commit_sha = self.refs[self.HEAD]

        self.refs[branch_name] = commit_sha or self.refs[self.HEAD]

    def checkout_branch(self, branch_name):
        self.HEAD = branch_name
        # @@@ do anything to index?

<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
import numpy as np
import matplotlib.pyplot as plt

h = 6.626e-34 # J*s
k_B = 1.381e-23 #J/K

T_CMB = 2.725 #K
sigma_T = 6.6524e-25 #cm^2
m_e = 0.511e3 #keV/c^2

H_0 = 68. #km/s/Mpc
h_70 = H_0/70.

Omega_m = 0.31
Omega_Lambda = 0.69

[P_0,c_500,gamma,alpha,beta] = [8.403*h_70**(-3./2.), 1.177,
                                0.3081, 1.0510, 5.4905]

def lp(x):
    '''
    lp: dimensionless universal profile

    Returns lp(x), x = r/R_500
    '''
    return P_0/((c_500*x)*(1+(c_500*x)**alpha)**((beta-gamma)/alpha))

def plot_lp():
    x = np.linspace(0.03,1,100)
    y = lp(x)/P_0

    plt.plot(x,y)
    plt.xscale('log')
    plt.yscale('log')
    plt.show()

def f_sz(freq):
    '''
    frequency component factor of sunyaev zel'dovich effect
    freq in GHz
    '''
    x = 10.**9*freq*h/(k_B*T_CMB)
    return x*(np.exp(x)+1)/(np.expm1(x))-4.

def plot_f_sz():
    x = np.linspace(1.,500.,1000)
    y = [f_sz(f) for f in x]
    plt.plot(x,y)
    plt.show()

def f_sz_rel(freq,T_e):
    '''
    relativistic frequency component factor of Sunyaev Zel'dovich effect
    freq en GHz

    needs and electron temperature T_e in keV/k_B (a 5 value means K_B*T_e=5keV)

    Typical values ranges 1 to 5
    '''
    theta_e = float(T_e)/0.511e3 #K_B*T_e in keV, since m_e*c^2 = 0.511 MeV
    x = 10.**9*float(freq)*h/(k_B*T_CMB)
    X = x/np.tanh(x/2)
    S = x/np.sinh(x/2)
    Y_0 = -4.0+X
    Y_1 = -10.0+(47.0/2.0)*X-(42.0/5.0)*X**2+(7.0/10.0)*X**3+(S**2)*(-21.0/5.0
          +(7.0/5.0)*X)
    Y_2 = (-15.0/2.0+(1023.0/8.0)*X-(868.0/5.0)*X**2+(329.0/5.0)*X**3
          -(44.0/5.0)*X**4+(11.0/30.0)*X**5+(S**2)*(-434.0/5.0+(658.0/5.0)*X
          -(242.0/5.0)*X**2+(143.0/30.0)*X**3)+(S**4)*(-44.0/5.0+(187.0/60.0)*X))
    return (Y_0+Y_1*theta_e+Y_2*theta_e**2)

def plot_f_sz_rel():
    x = np.linspace(100,300.,10000)
    y = [f_sz(f) for f in x]
    y_rel_5 = [f_sz_rel(f,5.) for f in x]
    y_rel_1 = [f_sz_rel(f,1.) for f in x]
    plt.plot(x,y_rel_5, color = 'Blue')
    plt.plot(x,y_rel_1, color = 'Red')
    plt.plot(x,y, color = 'Black')
    plt.show()

def h(z):
    '''
    Hubble scaled redshift function
    h(z) = H(z)/H_0 = (Omega_m*(1+z)^3+Omega_Lambda)^1/2
    '''
    return np.sqrt(Omega_m*(1+z)**3+Omega_Lambda)

def P(r,z,M_500, R_500):
    '''
    Universal Pressure Profile

    M_500 in 10^14 M_sun. typical from 0.1 to 10
    R_500 in Mpc. typical from 0.1 to 1

    P(r) = 1.65x10^-3 h(z)^8/3 [M_500 / 3x10^14 h_70^-1]^(2/3+α_p+α'_p(x))
           x lp(x) h_70^2 keV cm^-3
    '''
    x = r/R_500
    alpha_p = 0.12
    def alphap_p(x):
        #not using this correction for now, because it's bugged
        #return (0.10-(alpha_p+0.10)*(  ((x/0.5)**3)/(1.+ ((x/0.5)**3) )  ))
        return 0

    P_500 = (1.65e-3)*(h(z)**(8./3.))*((M_500/(3.e14*h_70**(-1)))**(2./3.))*h_70**2.
    P = P_500*((M_500/(3.e14*h_70**(-1)))**(alpha_p+alphap_p(x)))*lp(x)
    return P

def plot_P():
    x = np.linspace(0.01,1,1000)
    y = [P(r,1.,1.,1.)[0]/P(r,1.,1.,1.)[1] for r in x]
    plt.plot(x,y)
    plt.xlim([0.01,1.])
    plt.xscale('log')
    plt.yscale('log')
    plt.show()

def P_ss(r,z,M_500):
    '''
    Self similar universal pressure profile

    r in Mpc, typical values are 0.01 to 1
    M_500 in solar masses
    '''
    #first we pass all to mks
    G = 6.67408e-11 # m^3 kg^-1 s^-2

    #for the Hubble constant we divide km/Mpc
    #1 pc = 3.0857 × 10^16 m
    H_0_s = H_0*(1.e3/(1.e6*3.0857e16)) #s^-1
    def H(z):
        return H_0_s*np.sqrt(Omega_m*(1+z)**3+Omega_Lambda)

    f_B = 0.175
    mu = 0.59
    mu_e = 1.14

    #passing M_500 to kg
    M_500_kg = 1.989e30*M_500

    P_500 = (3./(8.*np.pi))*((500.*G**(-1./4)*H(z)**2)/2.)**(4./3.)*(mu/mu_e)*f_B*((M_500_kg)**(2./3.))
    #P_500 is in kg m^-1 s^-2 = kg m^2/s^2 m^-3 = J m^-3
    #we pass this to keV cm^-3
    #keV = 1000*1.602176565e-19 J => J = (1/1.602176565e-16) keV
    #1m = 100 cm -> m^-3 = 10^-6 cm
    #=> J m^-3 = 1.e10/1.602176565
    P_500 *= 1.e10/1.602176565

    def rho_c(z):
        '''
        critical density in kg m^-3
        '''
        return (3*(H(z))**2)/(8*np.pi*G)

    R_500 = (3*M_500_kg/(4*np.pi*500*rho_c(z)))**(1./3.) #in m

    #1 m = 1/(3.0857e16) pc
    R_500 = R_500/3.0857e22 #in Mpc
    x = r/R_500
    P = P_500*lp(x)
    return P

def plot_bothP():
    x = np.linspace(0.01,1,1000)
    y = [P(r,0.1,1.e13,0.3232) for r in x]
    y_ss = [P_ss(r,0.1,1.e13) for r in x]
    plt.plot(x,y, color = 'Blue')
    plt.plot(x,y_ss, color = 'Red')
    #plt.xlim([0.01,1.])
    plt.xscale('log')
    plt.yscale('log')
    plt.show()<commit_msg>Change function h(z) to h_(z) to avoid conflict with the variable h<commit_after># -*- coding: utf-8 -*-
import numpy as np
import matplotlib.pyplot as plt

h = 6.626e-34 # J*s
k_B = 1.381e-23 #J/K

T_CMB = 2.725 #K
sigma_T = 6.6524e-25 #cm^2
m_e = 0.511e3 #keV/c^2

H_0 = 68. #km/s/Mpc
h_70 = H_0/70.

Omega_m = 0.31
Omega_Lambda = 0.69

[P_0,c_500,gamma,alpha,beta] = [8.403*h_70**(-3./2.), 1.177,
                                0.3081, 1.0510, 5.4905]

def lp(x):
    '''
    lp: dimensionless universal profile

    Returns lp(x), x = r/R_500
    '''
    return P_0/((c_500*x)*(1+(c_500*x)**alpha)**((beta-gamma)/alpha))

def plot_lp():
    x = np.linspace(0.03,1,100)
    y = lp(x)/P_0

    plt.plot(x,y)
    plt.xscale('log')
    plt.yscale('log')
    plt.show()

def f_sz(freq):
    '''
    frequency component factor of sunyaev zel'dovich effect
    freq in GHz
    '''
    x = 10.**9*freq*h/(k_B*T_CMB)
    return x*(np.exp(x)+1)/(np.expm1(x))-4.

def plot_f_sz():
    x = np.linspace(1.,500.,1000)
    y = [f_sz(f) for f in x]
    plt.plot(x,y)
    plt.show()

def f_sz_rel(freq,T_e):
    '''
    relativistic frequency component factor of Sunyaev Zel'dovich effect
    freq en GHz

    needs and electron temperature T_e in keV/k_B (a 5 value means K_B*T_e=5keV)

    Typical values ranges 1 to 5
    '''
    theta_e = float(T_e)/0.511e3 #K_B*T_e in keV, since m_e*c^2 = 0.511 MeV
    x = 10.**9*float(freq)*h/(k_B*T_CMB)
    X = x/np.tanh(x/2)
    S = x/np.sinh(x/2)
    Y_0 = -4.0+X
    Y_1 = -10.0+(47.0/2.0)*X-(42.0/5.0)*X**2+(7.0/10.0)*X**3+(S**2)*(-21.0/5.0
          +(7.0/5.0)*X)
    Y_2 = (-15.0/2.0+(1023.0/8.0)*X-(868.0/5.0)*X**2+(329.0/5.0)*X**3
          -(44.0/5.0)*X**4+(11.0/30.0)*X**5+(S**2)*(-434.0/5.0+(658.0/5.0)*X
          -(242.0/5.0)*X**2+(143.0/30.0)*X**3)+(S**4)*(-44.0/5.0+(187.0/60.0)*X))
    return (Y_0+Y_1*theta_e+Y_2*theta_e**2)

def plot_f_sz_rel():
    x = np.linspace(100,300.,10000)
    y = [f_sz(f) for f in x]
    y_rel_5 = [f_sz_rel(f,5.) for f in x]
    y_rel_1 = [f_sz_rel(f,1.) for f in x]
    plt.plot(x,y_rel_5, color = 'Blue')
    plt.plot(x,y_rel_1, color = 'Red')
    plt.plot(x,y, color = 'Black')
    plt.show()

def h_(z):
    '''
    Hubble scaled redshift function
    h(z) = H(z)/H_0 = (Omega_m*(1+z)^3+Omega_Lambda)^1/2
    '''
    return np.sqrt(Omega_m*(1+z)**3+Omega_Lambda)

def P(r,z,M_500, R_500):
    '''
    Universal Pressure Profile

    M_500 in 10^14 M_sun. typical from 0.1 to 10
    R_500 in Mpc. typical from 0.1 to 1

    P(r) = 1.65x10^-3 h(z)^8/3 [M_500 / 3x10^14 h_70^-1]^(2/3+α_p+α'_p(x))
           x lp(x) h_70^2 keV cm^-3
    '''
    x = r/R_500
    alpha_p = 0.12
    def alphap_p(x):
        #not using this correction for now, because it's bugged
        #return (0.10-(alpha_p+0.10)*(  ((x/0.5)**3)/(1.+ ((x/0.5)**3) )  ))
        return 0

    P_500 = (1.65e-3)*(h_(z)**(8./3.))*((M_500/(3.e14*h_70**(-1)))**(2./3.))*h_70**2.
    P = P_500*((M_500/(3.e14*h_70**(-1)))**(alpha_p+alphap_p(x)))*lp(x)
    return P

def plot_P():
    x = np.linspace(0.01,1,1000)
    y = [P(r,1.,1.,1.)[0]/P(r,1.,1.,1.)[1] for r in x]
    plt.plot(x,y)
    plt.xlim([0.01,1.])
    plt.xscale('log')
    plt.yscale('log')
    plt.show()

def P_ss(r,z,M_500):
    '''
    Self similar universal pressure profile

    r in Mpc, typical values are 0.01 to 1
    M_500 in solar masses
    '''
    #first we pass all to mks
    G = 6.67408e-11 # m^3 kg^-1 s^-2

    #for the Hubble constant we divide km/Mpc
    #1 pc = 3.0857 × 10^16 m
    H_0_s = H_0*(1.e3/(1.e6*3.0857e16)) #s^-1
    def H(z):
        return H_0_s*np.sqrt(Omega_m*(1+z)**3+Omega_Lambda)

    f_B = 0.175
    mu = 0.59
    mu_e = 1.14

    #passing M_500 to kg
    M_500_kg = 1.989e30*M_500

    P_500 = (3./(8.*np.pi))*((500.*G**(-1./4)*H(z)**2)/2.)**(4./3.)*(mu/mu_e)*f_B*((M_500_kg)**(2./3.))
    #P_500 is in kg m^-1 s^-2 = kg m^2/s^2 m^-3 = J m^-3
    #we pass this to keV cm^-3
    #keV = 1000*1.602176565e-19 J => J = (1/1.602176565e-16) keV
    #1m = 100 cm -> m^-3 = 10^-6 cm
    #=> J m^-3 = 1.e10/1.602176565
    P_500 *= 1.e10/1.602176565

    def rho_c(z):
        '''
        critical density in kg m^-3
        '''
        return (3*(H(z))**2)/(8*np.pi*G)

    R_500 = (3*M_500_kg/(4*np.pi*500*rho_c(z)))**(1./3.) #in m

    #1 m = 1/(3.0857e16) pc
    R_500 = R_500/3.0857e22 #in Mpc
    x = r/R_500
    P = P_500*lp(x)
    return P

def plot_bothP():
    x = np.linspace(0.01,1,1000)
    y = [P(r,0.1,1.e13,0.3232) for r in x]
    y_ss = [P_ss(r,0.1,1.e13) for r in x]
    plt.plot(x,y, color = 'Blue')
    plt.plot(x,y_ss, color = 'Red')
    #plt.xlim([0.01,1.])
    plt.xscale('log')
    plt.yscale('log')
    plt.show()<|endoftext|>


<commit_before>import chainer
from chainer import cuda
import chainer.functions as F
from chainer import optimizers
import numpy as np
from PIL import Image, ImageChops
import time
import joblib
import json
import os
#import pickle
from tqdm import *

if not os.path.exists('cache'):
    os.mkdir('cache')
mem = joblib.Memory('cache')

# TODO: unit test
#   TODO: unit test for pickle / depickle, c/gpu
#   TODO: test that before training the RMSE is as random as you'd expect
#   TODO: deccode / encode on an easy example
# TODO: setup.py & PyPI Instllation
# TODO: Downloader for images
# TODO: docstrings for public functions
# TODO: private funcs tend to have underscores
# TODO: watch the 80 character limit

class ImageAutoEncoder():
    '''Variational Auto-encoder Class for image data.

    Using linear transformations with ReLU activations, this class peforms
    an encoding and then decoding step to form a full generative model for
    image data. Images can thus be encoded to a latent representation-space or
    decoded/generated from latent vectors.

    See  Kingma, Diederik and Welling, Max; "Auto-Encoding Variational Bayes" (2013)

    Given a set of input images train an artificial neural network, respampling at the latent
    stage from an approximate posterior multivariate gaussian distribution
    with unit covariance with mean and variance trained by the encoding step.

    Attributes
    ----------
    encode_sizes : List[int]
        List of layer sizes for hiden linear encoding layers of the model.
    decode_sizes : List[int]
        List of layer sizes for hiden linear decoding layers of the model.
    latent_dim : int
        Dimension of latent encoding space.
    img_width : int
        Width of the desired image representation.
    img_height : int
        Height of the desired image representation.
    color_channels : int
        Number of color channels in the input images.
    img_length : int
        Total dimensionality of input image `img_length = img_width *
        img_height * color_channels`.
    flag_gpu : bool
        Flag to toggle GPU training functionality.
    flag_dropout : bool
        Flag to toggle image dropout functionality.
    flag_autocrop : bool
        Flag to toggle image autocropping from background functionality.
    flag_grayscale : bool
        Flag to toggle image color channel averaging to grayscale.
    model : chainer.FunctionSet
        FunctionSet of chainer model layers for encoding and decoding.
    optimizer : chainer.Optimizer
        Chiner optimizer used to do backpropagation.
    x_all : numpy.array
        Numpy array used to hold training image data.

    '''

    def __init__(self, img_width=75, img_height=100, color_channels=3,
                 encode_layers=[1000, 600, 300],
                 decode_layers=[300, 800, 1000],
                 latent_width=100, flag_gpu=None, flag_dropout=False,
                 flag_autocrop=False, flag_grayscale=False):
        '''Setup for the variational auto-encoder.

        Inititalizes the layer setup for the NN to the defined dimensions.

        Parameters
        ----------

        img_width : int
            Width of the desired image representation.
        img_height : int
            Height of the desired image representation.
        color_channels : int
            Number of color channels in the input images (set to 1 automatically
                if `flag_grayscale = True`).
        encode_layers : List[int]
            List of layer sizes for hiden linear encoding layers of the model.
        decode_layers : List[int]
            List of layer sizes for hiden linear decoding layers of the model.
        latent_width : int
            Dimension of latent encoding space.
        flag_gpu : bool
            Flag to toggle GPU training functionality.
        flag_dropout : bool
            Flag to toggle image dropout functionality.
        flag_autocrop : bool
            Flag to toggle image autocropping from background functionality.
        flag_grayscale : bool
            Flag to toggle image color channel averaging to grayscale.

        '''
        self.encode_sizes = encode_layers
        self.decode_sizes = decode_layers
        self.latent_dim = latent_width
        self.img_width = img_width
        self.img_height = img_height
        if flag_grayscale:
            slef.color_channels = 1
        else:
            self.color_channels = color_channels
        self.img_len = self.img_width*self.img_height*self.color_channels
        self.flag_gpu = flag_gpu
        self.flag_dropout = flag_dropout
        self.flag_autocrop = flag_autocrop
        self.flag_grayscale = flag_grayscale
        self.model = self._layer_setup()
        self.optimizer = optimizers.Adam()
        self.optimizer.setup(self.model)
        self.x_all = np.array([], dtype=np.float32)

    def _layer_setup(self):
        # setup chainer layers for NN.

        layers = {}

        # Encoding Steps
        encode_layer_pairs = [(self.img_len, self.encode_sizes[0])]
        encode_layer_pairs += zip(self.encode_sizes[:-1],
                                  self.encode_sizes[1:])
        encode_layer_pairs += [(self.encode_sizes[-1], self.latent_dim * 2)]
        for i, (n_in, n_out) in enumerate(encode_layer_pairs):
            layers['encode_%i' % i] = F.Linear(n_in, n_out)

        # Decoding Steps
        decode_layer_pairs = [(self.latent_dim, self.decode_sizes[0])]
        decode_layer_pairs += zip(self.decode_sizes[:-1],
                                  self.decode_sizes[1:])
        decode_layer_pairs += [(self.decode_sizes[-1], self.img_len)]

        for i, (n_in, n_out) in enumerate(decode_layer_pairs):
            layers['decode_%i' % i] = F.Linear(n_in, n_out)
        model = chainer.FunctionSet(**layers)

        if self.flag_gpu:
            cuda.init()
            model.to_gpu()

        return model

    def _encode(self, img_batch):

        batch = img_batch

        if self.flag_dropout:
            batch = F.dropout(batch)
        n_layers = len(self.encode_sizes)
        for i in xrange(n_layers):
            batch = F.relu(getattr(self.model, 'encode_%i' % i)(batch))
        batch = F.relu(getattr(self.model, 'encode_%i' % n_layers)(batch))

        return batch

    def _decode(self, latent_vec):

        batch = latent_vec
        n_layers = len(self.decode_sizes)
        for i in xrange(n_layers):
            batch = F.relu(getattr(self.model, 'decode_%i' % i)(batch))
        batch = F.sigmoid(getattr(self.model, 'decode_%i' % n_layers)(batch))
        return batch

    def forward(self, img_batch):
        # TODO: underscore this method if it isn't public
        batch = chainer.Variable(img_batch / 255.)
        encoded = self._encode(batch)

        #Split latent space into `\mu` and `\sigma` parameters
        mean, std = F.split_axis(encoded, 2, 1)

        #create `latent_dim` N(0,1) normal samples.
        samples = np.random.standard_normal(mean.data.shape).astype('float32')
        if self.flag_gpu:
            samples = cuda.to_gpu(samples)
        samples = chainer.Variable(samples)

        #Scale samples to model trained parameters.
        sample_set = samples * F.exp(0.5*std) + mean
        output = self._decode(sample_set)
        reconstruction_loss = F.mean_squared_error(output, batch)

        #construct and scale KL Divergence loss.
        kl_div = -0.5 * F.sum(1 + std - mean ** 2 - F.exp(std))
        kl_div /= (img_batch.shape[1] * img_batch.shape[0])

        return reconstruction_loss, kl_div, output

    #@mem.cache
    def load_images(self, files):

        # TODO: docstring for all public functions
        # TODO: use logger instead of print
        # https://docs.python.org/2/library/logging.html
        shape = (self.img_width, self.img_height)
        print("Loading Image Files...")

        #helper function to crop images from background
        def trim(im):
            bg = Image.new(im.mode, im.size, im.getpixel((0, 0)))
            diff = ImageChops.difference(im, bg)
            diff = ImageChops.add(diff, diff, 2.0, -50)
            bbox = diff.getbbox()
            if bbox:
                return im.crop(bbox)

        #helper function to quickly resize images
        def resize(fname):
            im = Image.open(fname)
            #on the fly image cropping
            if self.flag_autocrop:
                im = trim(im)
            im.thumbnail(shape, Image.ANTIALIAS)
            im = im.resize(shape, Image.ANTIALIAS)
            im = np.float32(im)
            if self.flag_grayscale:
                im = im.mean(axis=2)
            return im
        self.x_all = np.array([resize(fname) for fname in tqdm(files)]).astype('float32')
        print("Image Files Loaded!")

    def fit(self, n_epochs=200, batch_size=100):
        # TODO: docstring for all public functions
        # TODO: rename to use sklearn's fit / predict / transform convention
        if self.x_all.shape[0] == 0:
            raise ValueError('Images not yet loaded, call load_images() first.')
        n_samp = self.x_all.shape[0]
        x_train = self.x_all.flatten().reshape((n_samp, -1))
        # Train Model#
        print("\n Training for %i epochs. \n" % n_epochs)
        for epoch in xrange(1, n_epochs + 1):
            print('epoch: %i' % epoch)
            t1 = time.time()
            indexes = np.random.permutation(x_train.shape[0])
            last_loss_kl = 0.
            last_loss_rec = 0.
            q = len(range(0, x_train.shape[0], batch_size))
            for i in tqdm(xrange(0, x_train.shape[0], batch_size)):
                if self.flag_gpu:
                    x_batch = cuda.to_gpu(x_train[indexes[i: i + batch_size]])
                else:
                    x_batch = x_train[indexes[i: i + batch_size]]
                self.optimizer.zero_grads()
                r_loss, kl_div, out = self.forward(x_batch)
                loss = r_loss + kl_div
                loss.backward()
                self.optimizer.update()
                last_loss_kl += kl_div.data
                last_loss_rec += r_loss.data

            msg = "r_loss = {0} , kl_div = {1}"
            print(msg.format(last_loss_rec/q, last_loss_kl/q))
            t_diff = time.time()-t1
            print("time: %f\n\n" % t_diff)

    def dump(self, filepath):
        # TODO: docstring for all public functions
        if not os.path.exists(os.path.dirname(filepath)):
            os.makedirs(os.path.dirname(filepath))
        print("Dumping model to file: %s " % filepath)
        cls_data = self.__dict__
        model = cls_data.pop('model')
        model.to_cpu()
        cls_data.pop('optimizer')
        cls_data.pop('x_all')
        meta = json.dumps(cls_data)
        with open(filepath+'_meta.json', 'wb') as f:
            f.write(meta)
        joblib.dump(model, filepath)

    @classmethod
    def load(cls, filepath, flag_gpu=False):
        # TODO: docstring for all public functions
        modpath = filepath
        metapath = filepath + '_meta.json'

        mess = "Model file {0} does not exist. Please check the file path."
        if not os.path.exists(modpath):
            print(mess.format(modpath))
        elif not os.path.exists(metapath):
            print(mess.format(metapath))
        else:
            with open(metapath, 'r') as f:
                meta = json.load(f)
            new_vae = cls(img_width=meta['img_width'],
                          img_height=meta['img_height'],
                          color_channels=meta['color_channels'],
                          encode_layers=meta['encode_sizes'],
                          decode_layers=meta['decode_sizes'],
                          latent_width=meta['latent_dim'],
                          flag_gpu=flag_gpu,
                          flag_dropout=meta['flag_dropout'],
                          flag_autocrop=meta['flag_autocrop'],
                          flag_grayscale=meta['flag_grayscale'])
            new_vae.model = joblib.load(modpath)
            new_vae.optimizer = optimizers.Adam()
            new_vae.optimizer.setup(new_vae.model)
            return new_vae
<commit_msg>Completed transform and inverse_transform methods as well as docstrings.<commit_after>import chainer
from chainer import cuda
import chainer.functions as F
from chainer import optimizers
import numpy as np
from PIL import Image, ImageChops
import time
import joblib
import json
import os
from tqdm import *

# TODO: unit test
#   TODO: unit test for pickle / depickle, c/gpu
#   TODO: test that before training the RMSE is as random as you'd expect
#   TODO: deccode / encode on an easy example
# TODO: setup.py & PyPI Instllation
# TODO: Downloader for images

class ImageAutoEncoder():
    '''Variational Auto-encoder Class for image data.

    Using linear transformations with ReLU activations, this class peforms
    an encoding and then decoding step to form a full generative model for
    image data. Images can thus be encoded to a latent representation-space or
    decoded/generated from latent vectors.

    See  Kingma, Diederik and Welling, Max; "Auto-Encoding Variational Bayes" (2013)

    Given a set of input images train an artificial neural network, respampling
    at the latent stage from an approximate posterior multivariate gaussian
    distribution with unit covariance with mean and variance trained by the
    encoding step.

    Attributes
    ----------
    encode_sizes : List[int]
        List of layer sizes for hiden linear encoding layers of the model.
    decode_sizes : List[int]
        List of layer sizes for hiden linear decoding layers of the model.
    latent_dim : int
        Dimension of latent encoding space.
    img_width : int
        Width of the desired image representation.
    img_height : int
        Height of the desired image representation.
    color_channels : int
        Number of color channels in the input images.
    img_length : int
        Total dimensionality of input image `img_length = img_width *
        img_height * color_channels`.
    rec_kl_ratio : float
            Ratio of relative importance between reconstruction loss and KL
            Divergence terms.
    flag_gpu : bool
        Flag to toggle GPU training functionality.
    flag_dropout : bool
        Flag to toggle image dropout functionality.
    flag_autocrop : bool
        Flag to toggle image autocropping from background functionality.
    flag_grayscale : bool
        Flag to toggle image color channel averaging to grayscale.
    model : chainer.FunctionSet
        FunctionSet of chainer model layers for encoding and decoding.
    optimizer : chainer.Optimizer
        Chiner optimizer used to do backpropagation.
    x_all : numpy.array
        Numpy array used to hold training image data.

    '''

    def __init__(self, img_width=75, img_height=100, color_channels=3,
                 encode_layers=[1000, 600, 300],
                 decode_layers=[300, 800, 1000],
                 latent_width=100, rec_kl_ratio=1.0, flag_gpu=None, flag_dropout=False,
                 flag_autocrop=False, flag_grayscale=False):
        '''Setup for the variational auto-encoder.

        Inititalizes the layer setup for the NN to the defined dimensions.

        Parameters
        ----------

        img_width : int
            Width of the desired image representation.
        img_height : int
            Height of the desired image representation.
        color_channels : int
            Number of color channels in the input images (set to 1 automatically
                if `flag_grayscale = True`).
        encode_layers : List[int]
            List of layer sizes for hiden linear encoding layers of the model.
        decode_layers : List[int]
            List of layer sizes for hiden linear decoding layers of the model.
        latent_width : int
            Dimension of latent encoding space.
        rec_kl_ratio : float
            Ratio of relative importance between reconstruction loss and KL
            Divergence terms.
        flag_gpu : bool
            Flag to toggle GPU training functionality.
        flag_dropout : bool
            Flag to toggle image dropout functionality.
        flag_autocrop : bool
            Flag to toggle image autocropping from background functionality.
        flag_grayscale : bool
            Flag to toggle image color channel averaging to grayscale.

        '''
        self.encode_sizes = encode_layers
        self.decode_sizes = decode_layers
        self.latent_dim = latent_width
        self.img_width = img_width
        self.img_height = img_height
        if flag_grayscale:
            slef.color_channels = 1
        else:
            self.color_channels = color_channels
        self.img_len = self.img_width*self.img_height*self.color_channels
        self.rec_kl_ratio = rec_kl_ratio
        self.flag_gpu = flag_gpu
        self.flag_dropout = flag_dropout
        self.flag_autocrop = flag_autocrop
        self.flag_grayscale = flag_grayscale
        self.model = self._layer_setup()
        self.optimizer = optimizers.Adam()
        self.optimizer.setup(self.model)
        self.x_all = np.array([], dtype=np.float32)

    def _layer_setup(self):
        # setup chainer layers for NN.

        layers = {}

        # Encoding Steps
        encode_layer_pairs = [(self.img_len, self.encode_sizes[0])]
        encode_layer_pairs += zip(self.encode_sizes[:-1],
                                  self.encode_sizes[1:])
        encode_layer_pairs += [(self.encode_sizes[-1], self.latent_dim * 2)]
        for i, (n_in, n_out) in enumerate(encode_layer_pairs):
            layers['encode_%i' % i] = F.Linear(n_in, n_out)

        # Decoding Steps
        decode_layer_pairs = [(self.latent_dim, self.decode_sizes[0])]
        decode_layer_pairs += zip(self.decode_sizes[:-1],
                                  self.decode_sizes[1:])
        decode_layer_pairs += [(self.decode_sizes[-1], self.img_len)]

        for i, (n_in, n_out) in enumerate(decode_layer_pairs):
            layers['decode_%i' % i] = F.Linear(n_in, n_out)
        model = chainer.FunctionSet(**layers)

        if self.flag_gpu:
            cuda.init()
            model.to_gpu()

        return model

    def _encode(self, img_batch):

        batch = img_batch

        if self.flag_dropout:
            batch = F.dropout(batch)
        n_layers = len(self.encode_sizes)
        for i in xrange(n_layers):
            batch = F.relu(getattr(self.model, 'encode_%i' % i)(batch))
        batch = F.relu(getattr(self.model, 'encode_%i' % n_layers)(batch))

        return batch

    def _decode(self, latent_vec):

        batch = latent_vec
        n_layers = len(self.decode_sizes)
        for i in xrange(n_layers):
            batch = F.relu(getattr(self.model, 'decode_%i' % i)(batch))
        batch = F.sigmoid(getattr(self.model, 'decode_%i' % n_layers)(batch))
        return batch

    def _forward(self, img_batch):

        batch = chainer.Variable(img_batch / 255.)
        encoded = self._encode(batch)

        #Split latent space into `\mu` and `\sigma` parameters
        mean, std = F.split_axis(encoded, 2, 1)

        #create `latent_dim` N(0,1) normal samples.
        samples = np.random.standard_normal(mean.data.shape).astype('float32')
        if self.flag_gpu:
            samples = cuda.to_gpu(samples)
        samples = chainer.Variable(samples)

        #Scale samples to model trained parameters.
        sample_set = samples * F.exp(0.5*std) + mean
        output = self._decode(sample_set)
        reconstruction_loss = F.mean_squared_error(output, batch)

        #construct and scale KL Divergence loss.
        kl_div = -0.5 * F.sum(1 + std - mean ** 2 - F.exp(std))
        kl_div /= (img_batch.shape[1] * img_batch.shape[0])

        return reconstruction_loss, kl_div, output

    def load_images(self, files):
        '''Load in image files from list of paths.

        Parameters
        ----------

        files : List[str]
            List of file paths of images to be loaded.

        '''

        # TODO: use logger instead of print
        # https://docs.python.org/2/library/logging.html
        shape = (self.img_width, self.img_height)
        print("Loading Image Files...")

        #helper function to crop images from background
        def trim(im):
            bg = Image.new(im.mode, im.size, im.getpixel((0, 0)))
            diff = ImageChops.difference(im, bg)
            diff = ImageChops.add(diff, diff, 2.0, -50)
            bbox = diff.getbbox()
            if bbox:
                return im.crop(bbox)

        #helper function to quickly resize images
        def resize(fname):
            im = Image.open(fname)
            #on the fly image cropping
            if self.flag_autocrop:
                im = trim(im)
            im.thumbnail(shape, Image.ANTIALIAS)
            im = im.resize(shape, Image.ANTIALIAS)
            im = np.float32(im)
            if self.flag_grayscale:
                im = im.mean(axis=2)
            return im
        self.x_all = np.array([resize(fname) for fname in tqdm(files)]).astype('float32')
        print("Image Files Loaded!")

    def fit(self, n_epochs=200, batch_size=100):
        '''Fit the VAE model to the image data.

        Parameters
        ----------

        n_epochs [optional] : int
            Gives the number of training epochs to run through for the fitting
            process.
        batch_size [optional] : int
            The size of the batch to use when training. Note: generally larger
            batch sizes will result in fater epoch iteration, but at the const
            of lower granulatity when updating the layer weights.

        '''

        if self.x_all.shape[0] == 0:
            raise ValueError('Images not yet loaded, call load_images() first.')
        n_samp = self.x_all.shape[0]
        x_train = self.x_all.flatten().reshape((n_samp, -1))
        # Train Model#
        print("\n Training for %i epochs. \n" % n_epochs)
        for epoch in xrange(1, n_epochs + 1):
            print('epoch: %i' % epoch)
            t1 = time.time()
            indexes = np.random.permutation(x_train.shape[0])
            last_loss_kl = 0.
            last_loss_rec = 0.
            q = len(range(0, x_train.shape[0], batch_size))
            for i in tqdm(xrange(0, x_train.shape[0], batch_size)):
                if self.flag_gpu:
                    x_batch = cuda.to_gpu(x_train[indexes[i: i + batch_size]])
                else:
                    x_batch = x_train[indexes[i: i + batch_size]]
                self.optimizer.zero_grads()
                r_loss, kl_div, out = self._forward(x_batch)
                loss = r_loss + kl_div*self.rec_kl_ratio
                loss.backward()
                self.optimizer.update()
                last_loss_kl += kl_div.data
                last_loss_rec += r_loss.data

            msg = "r_loss = {0} , kl_div = {1}"
            print(msg.format(last_loss_rec/q, last_loss_kl/q))
            t_diff = time.time()-t1
            print("time: %f\n\n" % t_diff)

    def transform(self, images, normalized=False):
        '''Transform image data to latent space.

        Parameters
        ----------
        images : array-like shape (n_images, image_width, image_height,
                                   n_colors)
            Input numpy array of images.
        normalized [optional] : bool
            Normalization flag that specifies whether pixel data is normalized
            to a [0,1] scale.

        Returns
        -------

        latent_vec : array-like shape (n_images, latent_dim)

        '''

        n_samp = images.shape[0]
        x_encoding = images.flatten().reshape((n_samp, -1))
        x_encoding = chainer.Variable(x_encoding)
        if not normalized:
            x_encoding /= 255.
        x_encoded = self._encode(x_encoding)

        mean, std = F.split_axis(x_encoded, 2, 1)

        #create `latent_dim` N(0,1) normal samples.
        samples = np.random.standard_normal(mean.data.shape).astype('float32')
        if self.flag_gpu:
            samples = cuda.to_gpu(samples)
        samples = chainer.Variable(samples)

        #Scale samples to model trained parameters.
        sample_set = samples * F.exp(0.5*std) + mean

        return sample_set.data

    def inverse_transform(self, encoded, normalized=True):
        '''Takes a latent space vector and transforms it into an image.

        Parameters
        ----------
        images : array-like shape (n_images, image_width, image_height,
                                   n_colors)
            Input numpy array of images.
        normalized [optional] : bool
            Normalization flag that specifies whether pixel data is normalized
            to a [0,1] scale.

        '''
        encoded = chainer.Variable(encoded)
        output = self._decode(encoded)
        return output.data.reshape((-1, self.img_height, self.img_width, 3))

    def dump(self, filepath):

        '''Saves model as a sequence of files.

        Parameters
        ----------
        filepath : str
            The path of the file you wish to save the model to. Note: the
            model will also contain files with paths '{filepath}_{number}.npy'
            and '{filepath}_meta.json'.

        '''
        if not os.path.exists(os.path.dirname(filepath)):
            os.makedirs(os.path.dirname(filepath))
        print("Dumping model to file: %s " % filepath)
        cls_data = self.__dict__
        model = cls_data.pop('model')
        model.to_cpu()
        cls_data.pop('optimizer')
        cls_data.pop('x_all')
        meta = json.dumps(cls_data)
        with open(filepath+'_meta.json', 'wb') as f:
            f.write(meta)
        joblib.dump(model, filepath)

    @classmethod
    def load(cls, filepath, flag_gpu=False):

        '''Loads in model as a class instance.

        Parameters
        ----------
        filepath : str
            Path to the first file that contains model information (will be
            without the '_{number}.npy'  or '_meta.json' tags at the end)
        flag_gpu : bool
            Specifies whether to load the model to use gpu capabilities.

        Returns
        -------

        class instance of self.

        '''
        modpath = filepath
        metapath = filepath + '_meta.json'

        mess = "Model file {0} does not exist. Please check the file path."
        if not os.path.exists(modpath):
            print(mess.format(modpath))
        elif not os.path.exists(metapath):
            print(mess.format(metapath))
        else:
            with open(metapath, 'r') as f:
                meta = json.load(f)
            new_vae = cls(img_width=meta['img_width'],
                          img_height=meta['img_height'],
                          color_channels=meta['color_channels'],
                          encode_layers=meta['encode_sizes'],
                          decode_layers=meta['decode_sizes'],
                          latent_width=meta['latent_dim'],
                          flag_gpu=flag_gpu,
                          flag_dropout=meta['flag_dropout'],
                          flag_autocrop=meta['flag_autocrop'],
                          flag_grayscale=meta['flag_grayscale'])
            new_vae.model = joblib.load(modpath)
            new_vae.optimizer = optimizers.Adam()
            new_vae.optimizer.setup(new_vae.model)
            return new_vae
<|endoftext|>


<commit_before>#!/usr/bin/env python
# encoding: utf-8

import re
import sys
import gzip
import pysam
import itertools
import mimetypes
from pypgen.fstats import fstats
from collections import OrderedDict, defaultdict

class VCF(object):
    """docstring for VCF"""
    def __init__(self, input, output=None, populations=None, region=None, window_size=1, step=0, snvs=None):
        super(VCF, self).__init__()

        self.input = input
        self.output = output
        self.step = step
        self.populations = populations
        self.region = region
        self.window_size = window_size
        self.chrms_2_sizes = self._get_chrm_ids_and_sizes_()
        self.empty_vcf_line = self.make_empty_vcf_ordered_dict()

    def __open_vcf__(self):
        """Open vcf file as gzip or as text."""

        if mimetypes.guess_type(self.input)[-1] is 'gzip':
            fin = gzip.open(self.input, 'rb')
        else:
            fin = open(self.input, 'r')
        return fin

    def _get_chrm_ids_and_sizes_(self):
        """ Extract chromosome ids and sizes from vcf file.
            Return as dictionary"""

        chrms_sizes_dict = OrderedDict()
        with self.__open_vcf__() as fin:

            for line in fin:

                if line.startswith("##contig"):
                    chrm_name = re.findall(r'ID=.*,', line)
                    chrm_name = chrm_name[0].strip('ID=').strip(',')

                    chrm_length = re.findall(r'length=.*>', line)
                    chrm_length = int(chrm_length[0].strip('length=').strip('>'))

                    chrms_sizes_dict[chrm_name] = chrm_length
                    break

                if line.startswith("#CHROM") is True:
                    break

        return chrms_sizes_dict

    def make_empty_vcf_ordered_dict(self):
        """Open VCF file and read in #CHROM line as an Ordered Dict"""

        header_dict = None
        with self.__open_vcf__() as fin:
            for line in fin:
                if line.startswith("#CHROM"):
                    header = line.strip("#").strip().split()
                    header_dict = OrderedDict([(item, None) for item in header])
                    break

        return header_dict

    def process_snp(self, snp_call):
        if snp_call == "0/0":
            return (0,0)

        elif snp_call == "1/1":
            return (1,1)

        elif snp_call == '1/0' or \
               snp_call == '0/1':
            return (0,1)

        # skip multiallelic sites
        else:
            return None

    def process_snp_call(self, snp_call, ref, alt, IUPAC_ambiguities=False):
        """Process VCF genotype fields."""

        # IUPAC ambiguity codes
        IUPAC_dict = {('A', 'C'): 'M',
                      ('A', 'G'): 'R',
                      ('A', 'T'): 'W',
                      ('C', 'G'): 'S',
                      ('C', 'T'): 'Y',
                      ('G', 'T'): 'K',
                      ('A', 'C', 'G'): 'V',
                      ('A', 'C', 'T'): 'H',
                      ('A', 'G', 'T'): 'D',
                      ('C', 'G', 'T'): 'B'}

        #called_base = ""

        # process blanks

        if snp_call == None:
            called_base = "-"

        if snp_call["GT"] == "./.":
            called_base = "-"

        else:
            allele1, allele2 = snp_call["GT"].split("/")

            # process "0/0"
            if allele1 == '0' and allele2 == '0':
                called_base = ref

            if allele1 == '1' and allele2 == '1':
                called_base = alt

            # process "0/N"
            if allele1 == '0' and allele2 != '0':

                if IUPAC_ambiguities == False:
                    called_base = 'N'

                else:
                    call = [ref] + [alt.split(',')[int(allele2) - 1]]
                    call.sort()
                    call = tuple(call)
                    called_base = IUPAC_dict[call]

            # process "2/2, 1/2, etc."
            if int(allele1) >= 1 and int(allele2) > 1:

                # deal with homozygotes
                if allele1 == allele2:
                    called_base = alt.split(',')[int(allele1) - 1]

                # deal with heterozygotes
                else:

                    if IUPAC_ambiguities == False:
                        called_base = 'N'

                    else:
                        ref = alt.split(',')[int(allele1) - 1]
                        alt = alt.split(',')[int(allele2) - 1]
                        call = [ref, alt]
                        call.sort()
                        call = tuple(call)
                        called_base = IUPAC_dict[call]

        return called_base

    def count_alleles(self, chunk):

        results = []
        for line in chunk:
            pop_counts = {}
            for pop in self.populations.keys():
                allele_counts = {'REF':0, 'ALT':0}
                for sample in self.populations[pop]:
                    if line[sample] != None:
                        ref, alt = self.process_snp(line[sample]['GT'])
                        allele_counts['REF'] += ref
                        allele_counts['ALT'] += alt

                pop_counts[pop] = allele_counts.copy()

            results.append(pop_counts.copy())

        return results

    def slice_vcf(self, vcf_bgzipped_file, chrm, start, stop):

        tbx = pysam.Tabixfile(vcf_bgzipped_file)

        try:
            vcf_slice = tbx.fetch(chrm, start, stop)
        except ValueError:
            return ()

        else:
            return tuple(row for row in vcf_slice)

    def vcf_file_iterator(self, as_dict=True):

        for line in self.__open_vcf__():
            if line.startswith("#") is not True:
                if as_dict == True:
                    yield self.parse_vcf_line(line, self.empty_vcf_line)
                else:
                    yield line
            else:
                continue

    def parse_info_field(self, info_field):

        info_dict = {}
        for item in info_field.split(';'):
            pair = item.split("=")
            if len(pair) == 2:
                info_dict[pair[0]] = pair[1]   # this could be improved on

        return info_dict

    def get_population_sizes(self, vcfline, populations):

        sample_counts = {}

        for pop in populations.keys():
            sample_count = 0

            for sample in populations[pop]:
                if vcfline[sample] is not None:
                    sample_count += 1

            sample_counts[pop] = sample_count

        return sample_counts

    def parse_vcf_line(self, pos, vcf_line_dict):
        """Read in VCF line and convert it to an OrderedDict"""

        pos_parts = pos.strip().split()

        for count, item in enumerate(vcf_line_dict):
            vcf_line_dict[item] = pos_parts[count]

        sample_format = vcf_line_dict["FORMAT"].split(":")

        for count, item in enumerate(vcf_line_dict):
            if count >= 9:
                genotype = vcf_line_dict[item]

                if  "./." in genotype or genotype == ".":      # "'./.'' for dip, '.' for haploid
                    vcf_line_dict[item] = None

                else:
                    genotype = dict(zip(sample_format, genotype.split(":")))

                    # CONVERT STRINGS TO APPOPRIATE TYPES (INTS, FLOATS, ETC.)
                    if genotype.has_key("GQ"):
                        genotype['GQ'] = float(genotype['GQ'])

                    if genotype.has_key("DP"):
                        genotype['DP'] = int(genotype['DP'])

                    if genotype.has_key("AD"):
                        genotype['AD'] = tuple(int(ad) for ad in genotype['AD'].split(","))

                    if genotype.has_key("PL"):
                        genotype['PL'] = tuple(int(ad) for ad in genotype['PL'].split(","))

                    vcf_line_dict[item] = genotype

        vcf_line_dict['POS'] = int(vcf_line_dict['POS'])

        try:
            vcf_line_dict['QUAL'] = float(vcf_line_dict['QUAL'])
        except ValueError:
            pass

        vcf_line_dict['INFO'] = self.parse_info_field(vcf_line_dict['INFO'])

        return vcf_line_dict.copy()

    def filter_vcf_line(self, vcf_line_dict, filters=None):
        parse_vcf_line
        pass

    def lines_2_dicts(self, chunk):

        vcf_line_dict = self.empty_vcf_line.copy()
        return [self.parse_vcf_line(line, vcf_line_dict) for line in chunk]

    def get_chromosome_lengths(self, regions_to_skip=[]):

        try:
            tbx = pysam.Tabixfile(self.input)  # TODO: create try statement to test that file is actually a VCF
        except:
            print 'Input not Tabix Indexed.'
            sys.exit()

        # PARSE LENGTH INFO FROM HEADER

        chrm_lengths = []
        chrm_lengths_dict = {}
        for line in tbx.header:

            if line.startswith("##contig="):

                chrm_name = re.findall(r'ID=.*,', line)
                chrm_name = chrm_name[0].strip('ID=').strip(',')

                chrm_length = re.findall(r'length=.*>', line)
                chrm_length = int(chrm_length[0].strip('length=').strip('>'))

                if chrm_name in regions_to_skip:
                    print 'skipping', chrm_name
                    continue

                chrm_lengths.append((chrm_name, 1, chrm_length))
                chrm_lengths_dict[chrm_name] = chrm_length

        chrm_lengths = tuple(chrm_lengths)
        tbx.close()

        return chrm_lengths_dict

    def get_slice_indicies(self):
        """Get slice information from VCF file that is tabix indexed (bgzipped). """

        # GENERATE SLICES:
        # Does not yield final partial slice. Not a bug!

        if self.region == [None]:

            # ITERATE OVER CHROMOSOMES (USE ORDERED DICT TO KEEP IN VCF ORDER)
            for chrm, length in self.chrms_2_sizes.iteritems():

                cStart = 0
                cStop = 0
                iCount = 0

                while cStop < length:

                    if iCount == 0:
                        cStart = 1
                        cStop = self.window_size
                        iCount += 1

                    yield (chrm, cStart, cStop)
                    cStart += self.step
                    cStop += self.step

        else:

            chrm, start, stop = self.region

            cStart = 0
            cStop = 0
            iCount = 0

            if self.window_size == None:
                self.window_size = stop - start

            if self.step == None:
                self.step = 0

            while cStop < stop:

                if iCount == 0:
                    cStart = start
                    cStop = start + self.window_size - 1
                    iCount += 1

                yield (chrm, cStart, cStop)

                if self.step == 0:
                    cStart += self.window_size
                else:
                    cStart += self.step

                cStop = cStart + self.window_size + self.step - 1

    def snp_chunks_interator(self, num_snv_per_chunk=3, fillvalue=None):

        def grouper(iterable, n):
            """Yields chunks of SNV of n number of SNVs.
               Properly handles chromsome so that
               chunks don't overlap chromosomes."""

            chrom = None
            chunk = []

            for count, i in enumerate(iterable):

                s = i.split("\t")[:3]
                current_chrom = s[0]

                if count == 0:
                    chrom = current_chrom

                if current_chrom != chrom:
                    yield chunk
                    chunk = [i]
                    chrom = current_chrom
                    continue

                if current_chrom == chrom:

                    chunk.append(i)

                    if len(chunk) >= n:
                        yield chunk
                        chunk = []

                chrom = current_chrom

        chunk_iterator = grouper(self.vcf_file_iterator(), n=num_snv_per_chunk)

        for c, i in enumerate(chunk_iterator):

            if len(i) == 0: # handles edge case
                continue

            if len(i) != num_snv_per_chunk:
                padding =  num_snv_per_chunk - len(i)
                i.extend([fillvalue]*padding)

            yield i

    def calc_allele_counts(self, vcf_line_dict, sample_ids=None):

        #allele_counts = defaultdict({0:0.0,1:0.0,2:0.0,3:0.0,4:0.0})
        allele_counts = dict((key, {0: 0.0, 1: 0.0, 2: 0.0, 3: 0.0}) for key in self.populations.keys())

        for population in self.populations.keys():
            for sample_id in self.populations[population]:

                if vcf_line_dict[sample_id] != None:

                    genotype = vcf_line_dict[sample_id]
                    genotype = genotype["GT"].split("/")   # TODO add phased logic

                    if genotype == [".", "."]:
                        continue

                    genotype = [int(item) for item in genotype]

                    for allele in genotype:
                        allele_counts[population][allele] += 1.0

        return allele_counts
<commit_msg>Added better STDIN processing for VCF files<commit_after>#!/usr/bin/env python
# encoding: utf-8

import re
import sys
import gzip
import pysam
import itertools
import mimetypes
from pypgen.fstats import fstats
from collections import OrderedDict, defaultdict

class VCF(object):
    """docstring for VCF"""
    def __init__(self, input, output=None, populations=None, region=None, window_size=1, step=0, snvs=None):
        super(VCF, self).__init__()

        self.input = input
        self.output = output
        self.step = step
        self.populations = populations
        self.region = region
        self.window_size = window_size
        self.chrms_2_sizes = self._get_chrm_ids_and_sizes_()
        self.empty_vcf_line = self.make_empty_vcf_ordered_dict()

    def __open_vcf__(self):
        """Open vcf file as gzip or as text."""

        if type(self.input) is file:
            fin = self.input

        elif mimetypes.guess_type(self.input)[-1] is 'gzip':
            fin = gzip.open(self.input, 'rb')

        else:
            fin = open(self.input, 'r')

        return fin

    def _get_chrm_ids_and_sizes_(self):
        """ Extract chromosome ids and sizes from vcf file.
            Return as dictionary"""

        chrms_sizes_dict = OrderedDict()
        # with self.__open_vcf__() as fin:

        for line in self.__open_vcf__():

            if line.startswith("##contig"):
                chrm_name = re.findall(r'ID=.*,', line)
                chrm_name = chrm_name[0].strip('ID=').strip(',')

                chrm_length = re.findall(r'length=.*>', line)
                chrm_length = int(chrm_length[0].strip('length=').strip('>'))

                chrms_sizes_dict[chrm_name] = chrm_length
                break

            if line.startswith("#CHROM") is True:
                break

        return chrms_sizes_dict

    def make_empty_vcf_ordered_dict(self):
        """Open VCF file and read in #CHROM line as an Ordered Dict"""

        header_dict = None

        for line in self.__open_vcf__():
            if line.startswith("#CHROM"):
                header = line.strip("#").strip().split()
                header_dict = OrderedDict([(item, None) for item in header])
                break

        return header_dict

    def process_snp(self, snp_call):
        if snp_call == "0/0":
            return (0,0)

        elif snp_call == "1/1":
            return (1,1)

        elif snp_call == '1/0' or \
               snp_call == '0/1':
            return (0,1)

        # skip multiallelic sites
        else:
            return None

    def process_snp_call(self, snp_call, ref, alt, IUPAC_ambiguities=False):
        """Process VCF genotype fields."""

        # IUPAC ambiguity codes
        IUPAC_dict = {('A', 'C'): 'M',
                      ('A', 'G'): 'R',
                      ('A', 'T'): 'W',
                      ('C', 'G'): 'S',
                      ('C', 'T'): 'Y',
                      ('G', 'T'): 'K',
                      ('A', 'C', 'G'): 'V',
                      ('A', 'C', 'T'): 'H',
                      ('A', 'G', 'T'): 'D',
                      ('C', 'G', 'T'): 'B'}

        if snp_call == None:
            called_base = "-"

        if snp_call["GT"] == "./.":
            called_base = "-"

        else:
            allele1, allele2 = snp_call["GT"].split("/")

            # process "0/0"
            if allele1 == '0' and allele2 == '0':
                called_base = ref

            if allele1 == '1' and allele2 == '1':
                called_base = alt

            # process "0/N"
            if allele1 == '0' and allele2 != '0':

                if IUPAC_ambiguities == False:
                    called_base = 'N'

                else:
                    call = [ref] + [alt.split(',')[int(allele2) - 1]]
                    call.sort()
                    call = tuple(call)
                    called_base = IUPAC_dict[call]

            # process "2/2, 1/2, etc."
            if int(allele1) >= 1 and int(allele2) > 1:

                # deal with homozygotes
                if allele1 == allele2:
                    called_base = alt.split(',')[int(allele1) - 1]

                # deal with heterozygotes
                else:

                    if IUPAC_ambiguities == False:
                        called_base = 'N'

                    else:
                        ref = alt.split(',')[int(allele1) - 1]
                        alt = alt.split(',')[int(allele2) - 1]
                        call = [ref, alt]
                        call.sort()
                        call = tuple(call)
                        called_base = IUPAC_dict[call]

        return called_base

    def count_alleles(self, chunk):

        results = []
        for line in chunk:
            pop_counts = {}
            for pop in self.populations.keys():
                allele_counts = {'REF':0, 'ALT':0}
                for sample in self.populations[pop]:
                    if line[sample] != None:
                        ref, alt = self.process_snp(line[sample]['GT'])
                        allele_counts['REF'] += ref
                        allele_counts['ALT'] += alt

                pop_counts[pop] = allele_counts.copy()

            results.append(pop_counts.copy())

        return results

    def slice_vcf(self, vcf_bgzipped_file, chrm, start, stop):

        tbx = pysam.Tabixfile(vcf_bgzipped_file)

        try:
            vcf_slice = tbx.fetch(chrm, start, stop)
        except ValueError:
            return ()

        else:
            return tuple(row for row in vcf_slice)

    def vcf_file_iterator(self, as_dict=True):

        for line in self.__open_vcf__():
            if line.startswith("#") is not True:
                if as_dict == True:
                    yield self.parse_vcf_line(line, self.empty_vcf_line)
                else:
                    yield line
            else:
                continue

    def parse_info_field(self, info_field):

        info_dict = {}
        for item in info_field.split(';'):
            pair = item.split("=")
            if len(pair) == 2:
                info_dict[pair[0]] = pair[1]   # this could be improved on

        return info_dict

    def get_population_sizes(self, vcfline, populations):

        sample_counts = {}

        for pop in populations.keys():
            sample_count = 0

            for sample in populations[pop]:
                if vcfline[sample] is not None:
                    sample_count += 1

            sample_counts[pop] = sample_count

        return sample_counts

    def parse_vcf_line(self, pos, vcf_line_dict):
        """Read in VCF line and convert it to an OrderedDict"""

        pos_parts = pos.strip().split()

        for count, item in enumerate(vcf_line_dict):
            vcf_line_dict[item] = pos_parts[count]

        sample_format = vcf_line_dict["FORMAT"].split(":")

        for count, item in enumerate(vcf_line_dict):
            if count >= 9:
                genotype = vcf_line_dict[item]

                if  "./." in genotype or genotype == ".":      # "'./.'' for dip, '.' for haploid
                    vcf_line_dict[item] = None

                else:
                    genotype = dict(zip(sample_format, genotype.split(":")))

                    # CONVERT STRINGS TO APPOPRIATE TYPES (INTS, FLOATS, ETC.)
                    if genotype.has_key("GQ"):
                        genotype['GQ'] = float(genotype['GQ'])

                    if genotype.has_key("DP"):
                        genotype['DP'] = int(genotype['DP'])

                    if genotype.has_key("AD"):
                        genotype['AD'] = tuple(int(ad) for ad in genotype['AD'].split(","))

                    if genotype.has_key("PL"):
                        genotype['PL'] = tuple(int(ad) for ad in genotype['PL'].split(","))

                    vcf_line_dict[item] = genotype

        vcf_line_dict['POS'] = int(vcf_line_dict['POS'])

        try:
            vcf_line_dict['QUAL'] = float(vcf_line_dict['QUAL'])
        except ValueError:
            pass

        vcf_line_dict['INFO'] = self.parse_info_field(vcf_line_dict['INFO'])

        return vcf_line_dict.copy()

    def filter_vcf_line(self, vcf_line_dict, filters=None):
        parse_vcf_line
        pass

    def lines_2_dicts(self, chunk):

        vcf_line_dict = self.empty_vcf_line.copy()
        return [self.parse_vcf_line(line, vcf_line_dict) for line in chunk]

    def get_chromosome_lengths(self, regions_to_skip=[]):

        try:
            tbx = pysam.Tabixfile(self.input)  # TODO: create try statement to test that file is actually a VCF
        except:
            print 'Input not Tabix Indexed.'
            sys.exit()

        # PARSE LENGTH INFO FROM HEADER

        chrm_lengths = []
        chrm_lengths_dict = {}
        for line in tbx.header:

            if line.startswith("##contig="):

                chrm_name = re.findall(r'ID=.*,', line)
                chrm_name = chrm_name[0].strip('ID=').strip(',')

                chrm_length = re.findall(r'length=.*>', line)
                chrm_length = int(chrm_length[0].strip('length=').strip('>'))

                if chrm_name in regions_to_skip:
                    print 'skipping', chrm_name
                    continue

                chrm_lengths.append((chrm_name, 1, chrm_length))
                chrm_lengths_dict[chrm_name] = chrm_length

        chrm_lengths = tuple(chrm_lengths)
        tbx.close()

        return chrm_lengths_dict

    def get_slice_indicies(self):
        """Get slice information from VCF file that is tabix indexed (bgzipped). """

        # GENERATE SLICES:
        # Does not yield final partial slice. Not a bug!

        if self.region == [None]:

            # ITERATE OVER CHROMOSOMES (USE ORDERED DICT TO KEEP IN VCF ORDER)
            for chrm, length in self.chrms_2_sizes.iteritems():

                cStart = 0
                cStop = 0
                iCount = 0

                while cStop < length:

                    if iCount == 0:
                        cStart = 1
                        cStop = self.window_size
                        iCount += 1

                    yield (chrm, cStart, cStop)
                    cStart += self.step
                    cStop += self.step

        else:

            chrm, start, stop = self.region

            cStart = 0
            cStop = 0
            iCount = 0

            if self.window_size == None:
                self.window_size = stop - start

            if self.step == None:
                self.step = 0

            while cStop < stop:

                if iCount == 0:
                    cStart = start
                    cStop = start + self.window_size - 1
                    iCount += 1

                yield (chrm, cStart, cStop)

                if self.step == 0:
                    cStart += self.window_size
                else:
                    cStart += self.step

                cStop = cStart + self.window_size + self.step - 1

    def snp_chunks_interator(self, num_snv_per_chunk=3, fillvalue=None):

        def grouper(iterable, n):
            """Yields chunks of SNV of n number of SNVs.
               Properly handles chromsome so that
               chunks don't overlap chromosomes."""

            chrom = None
            chunk = []

            for count, i in enumerate(iterable):

                s = i.split("\t")[:3]
                current_chrom = s[0]

                if count == 0:
                    chrom = current_chrom

                if current_chrom != chrom:
                    yield chunk
                    chunk = [i]
                    chrom = current_chrom
                    continue

                if current_chrom == chrom:

                    chunk.append(i)

                    if len(chunk) >= n:
                        yield chunk
                        chunk = []

                chrom = current_chrom

        chunk_iterator = grouper(self.vcf_file_iterator(), n=num_snv_per_chunk)

        for c, i in enumerate(chunk_iterator):

            if len(i) == 0: # handles edge case
                continue

            if len(i) != num_snv_per_chunk:
                padding =  num_snv_per_chunk - len(i)
                i.extend([fillvalue]*padding)

            yield i

    def calc_allele_counts(self, vcf_line_dict, sample_ids=None):

        #allele_counts = defaultdict({0:0.0,1:0.0,2:0.0,3:0.0,4:0.0})
        allele_counts = dict((key, {0: 0.0, 1: 0.0, 2: 0.0, 3: 0.0}) for key in self.populations.keys())

        for population in self.populations.keys():
            for sample_id in self.populations[population]:

                if vcf_line_dict[sample_id] != None:

                    genotype = vcf_line_dict[sample_id]
                    genotype = genotype["GT"].split("/")   # TODO add phased logic

                    if genotype == [".", "."]:
                        continue

                    genotype = [int(item) for item in genotype]

                    for allele in genotype:
                        allele_counts[population][allele] += 1.0

        return allele_counts
<|endoftext|>