text:
#***************************************************************************************************
# Copyright 2015, 2019 National Technology & Engineering Solutions of Sandia, LLC (NTESS).
# Under the terms of Contract DE-NA0003525 with NTESS, the U.S. Government retains certain rights
# in this software.
# Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except
# in compliance with the License.  You may obtain a copy of the License at
# http://www.apache.org/licenses/LICENSE-2.0 or in the LICENSE file in the root pyGSTi directory.
#***************************************************************************************************

import numpy as _np
from . import rpeconfig as _rpeconfig

rpeconfig_GxPi2_GyPi2_00_dict = {}
rpeconfig_GxPi2_GyPi2_00_dict['fixed_axis_gate_label'] = 'Gx'
rpeconfig_GxPi2_GyPi2_00_dict['fixed_axis_label'] = 'X'
rpeconfig_GxPi2_GyPi2_00_dict['fixed_axis_target'] = [0, 1, 0, 0]
rpeconfig_GxPi2_GyPi2_00_dict['loose_axis_gate_label'] = 'Gy'
rpeconfig_GxPi2_GyPi2_00_dict['loose_axis_label'] = 'Y'
rpeconfig_GxPi2_GyPi2_00_dict['loose_axis_target'] = [0, 0, 1, 0]
rpeconfig_GxPi2_GyPi2_00_dict['auxiliary_axis_gate_label'] = 'Gz'
rpeconfig_GxPi2_GyPi2_00_dict['auxiliary_axis_label'] = 'Z'
rpeconfig_GxPi2_GyPi2_00_dict['rhoExpressions'] = ["0"]
rpeconfig_GxPi2_GyPi2_00_dict['EExpressions'] = ["0", "1"]
rpeconfig_GxPi2_GyPi2_00_dict['ELabels'] = ["0", "1"]
#rpeconfig_GxPi2_GyPi200n_dict['spamLabelDict'] = {'plus': (0,0), 'minus': (0,-1) }
rpeconfig_GxPi2_GyPi2_00_dict['spamLabelDict'] = {'0': ('rho0', 'E0'), '1': ('rho0', 'remainder')}
rpeconfig_GxPi2_GyPi2_00_dict['dn_labels'] = ['0']
rpeconfig_GxPi2_GyPi2_00_dict['up_labels'] = ['1']
rpeconfig_GxPi2_GyPi2_00_dict['alpha'] = _np.pi / 2
rpeconfig_GxPi2_GyPi2_00_dict['epsilon'] = _np.pi / 2
rpeconfig_GxPi2_GyPi2_00_dict['theta'] = 0  # This should always be 0.
rpeconfig_GxPi2_GyPi2_00_dict['new_epsilon_func'] = lambda epsilon: (epsilon / (_np.pi / 2)) - 1
rpeconfig_GxPi2_GyPi2_00_dict['alpha_hat_func'] = lambda xhat, yhat, Nx, Ny: _np.arctan2(
    (xhat - Nx / 2.) / Nx, -(yhat - Ny / 2.) / Ny)
rpeconfig_GxPi2_GyPi2_00_dict['epsilon_hat_func'] = lambda xhat, yhat, Nx, Ny: _np.arctan2(
    (xhat - Nx / 2.) / Nx, -(yhat - Ny / 2.) / Ny)
rpeconfig_GxPi2_GyPi2_00_dict['Phi_hat_func'] = lambda xhat, yhat, Nx, Ny: _np.arctan2(
    (xhat - Nx / 2.) / Nx, -(yhat - Ny / 2.) / Ny)
rpeconfig_GxPi2_GyPi2_00_dict['alpha_cos_prep_tuple'] = ()
rpeconfig_GxPi2_GyPi2_00_dict['alpha_cos_prep_str'] = ''
rpeconfig_GxPi2_GyPi2_00_dict['alpha_cos_germ_tuple'] = ('Gx',)
rpeconfig_GxPi2_GyPi2_00_dict['alpha_cos_germ_str'] = 'Gx'
rpeconfig_GxPi2_GyPi2_00_dict['alpha_cos_meas_tuple'] = ()
rpeconfig_GxPi2_GyPi2_00_dict['alpha_cos_meas_str'] = ''
rpeconfig_GxPi2_GyPi2_00_dict['alpha_sin_prep_tuple'] = ()
rpeconfig_GxPi2_GyPi2_00_dict['alpha_sin_prep_str'] = ''
rpeconfig_GxPi2_GyPi2_00_dict['alpha_sin_germ_tuple'] = ('Gx',)
rpeconfig_GxPi2_GyPi2_00_dict['alpha_sin_germ_str'] = 'Gx'
rpeconfig_GxPi2_GyPi2_00_dict['alpha_sin_meas_tuple'] = ('Gx',)
rpeconfig_GxPi2_GyPi2_00_dict['alpha_sin_meas_str'] = 'Gx'
rpeconfig_GxPi2_GyPi2_00_dict['epsilon_cos_prep_tuple'] = ()
rpeconfig_GxPi2_GyPi2_00_dict['epsilon_cos_prep_str'] = ''
rpeconfig_GxPi2_GyPi2_00_dict['epsilon_cos_germ_tuple'] = ('Gy',)
rpeconfig_GxPi2_GyPi2_00_dict['epsilon_cos_germ_str'] = 'Gy'
rpeconfig_GxPi2_GyPi2_00_dict['epsilon_cos_meas_tuple'] = ()
rpeconfig_GxPi2_GyPi2_00_dict['epsilon_cos_meas_str'] = ''
rpeconfig_GxPi2_GyPi2_00_dict['epsilon_sin_prep_tuple'] = ()
rpeconfig_GxPi2_GyPi2_00_dict['epsilon_sin_prep_str'] = ''
rpeconfig_GxPi2_GyPi2_00_dict['epsilon_sin_germ_tuple'] = ('Gy',)
rpeconfig_GxPi2_GyPi2_00_dict['epsilon_sin_germ_str'] = 'Gy'
rpeconfig_GxPi2_GyPi2_00_dict['epsilon_sin_meas_tuple'] = ('Gy',)
rpeconfig_GxPi2_GyPi2_00_dict['epsilon_sin_meas_str'] = 'Gy'
rpeconfig_GxPi2_GyPi2_00_dict['theta_cos_prep_tuple'] = ()
rpeconfig_GxPi2_GyPi2_00_dict['theta_cos_prep_str'] = ''
rpeconfig_GxPi2_GyPi2_00_dict['theta_cos_germ_tuple'] = ('Gx', 'Gy', 'Gy', 'Gx', 'Gx', 'Gy', 'Gy', 'Gx')
rpeconfig_GxPi2_GyPi2_00_dict['theta_cos_germ_str'] = 'GxGyGyGxGxGyGyGx'
rpeconfig_GxPi2_GyPi2_00_dict['theta_cos_meas_tuple'] = ()
rpeconfig_GxPi2_GyPi2_00_dict['theta_cos_meas_str'] = ''
rpeconfig_GxPi2_GyPi2_00_dict['theta_sin_prep_tuple'] = ()
rpeconfig_GxPi2_GyPi2_00_dict['theta_sin_prep_str'] = ''
rpeconfig_GxPi2_GyPi2_00_dict['theta_sin_germ_tuple'] = ('Gx', 'Gy', 'Gy', 'Gx', 'Gx', 'Gy', 'Gy', 'Gx')
rpeconfig_GxPi2_GyPi2_00_dict['theta_sin_germ_str'] = 'GxGyGyGxGxGyGyGx'
rpeconfig_GxPi2_GyPi2_00_dict['theta_sin_meas_tuple'] = ('Gy',)
rpeconfig_GxPi2_GyPi2_00_dict['theta_sin_meas_str'] = 'Gy'

rpeconfig_GxPi2_GyPi2_00 = _rpeconfig(rpeconfig_GxPi2_GyPi2_00_dict)

#rpeconfig_GxPi2_GyPi2_UpDn_dict['make_rpe_gs_func'] = lambda alphaTrue, epsilonTrue, Zrot, withId=True :\
#    if withId:
#        outputModel = _setc.build_explicit_model(
#            [('Q0',)],['Gi','Gy','Gx'],
#            [ "I(Q0)", "Y(%s,Q0)" % epsilonTrue, "X(%s,Q0)" % alphaTrue],
#            rhoExpressions=["0"], EExpressions=["1"],
#            spamLabelDict={'plus': (0,0), 'minus': (0,-1) } )
#    else:
#        outputModel = _setc.build_explicit_model(
#            [('Q0',)],['Gy','Gx'],
#            [ "Y(%s,Q0)" % epsilonTrue, "X(%s,Q0)" % alphaTrue],
#            rhoExpressions=["0"], EExpressions=["1"],
#            spamLabelDict={'plus': (0,0), 'minus': (0,-1) })
#
#    if Zrot != 0:
#        modelAux1 = _setc.build_explicit_model(
#            [('Q0',)],['Gi','Gz','Gx'],
#            [ "I(Q0)", "Z(%s,Q0)" % Zrot, "X(pi/2,Q0)"],
#            rhoExpressions=["0"], EExpressions=["1"],
#            spamLabelDict={'plus': (0,0), 'minus': (0,-1) })
#
#        outputModel.set_gate('Gy',_objs.FullDenseOp(
#                _np.dot( _np.dot(_np.linalg.inv(modelAux1['Gz']),
#                               outputModel['Gy']),modelAux1['Gz'])) )
#
#
#
#    SPAMdepol, gateDepol=None,

max_stars_repo_name: drewrisinger/pyGSTi

max_issues_repo_name: drewrisinger/pyGSTi

max_forks_repo_name: drewrisinger/pyGSTi


text:
import os
import numpy as np
import np_tif
from stack_registration import stack_registration
from stack_registration import bucket
import matplotlib.pyplot as plt

def main():

    assert os.path.isdir('./../images')
    if not os.path.isdir('./../images/figure_A4'):
        os.mkdir('./../images/figure_A4')

    z_voltages = [
        '_34000mV',
        '_36000mV',
        '_38000mV',
        '_40000mV',
        '_42000mV',
        '_44000mV',
        '_46000mV',
        '_48000mV',
        '_50000mV',
        '_52000mV',
        '_54000mV',
        '_56000mV',
        '_58000mV',
        '_60000mV',
        '_62000mV',
        '_64000mV',
        '_66000mV',
        '_68000mV',
        '_70000mV',
        ]

    bg_level = 111.5

    # for each piezo voltage (axial location) we:
    # 1. repetition average the images for each delay
    # 2. register the repetition averaged "green off" image to the
    # corresponding "green on" image

    num_delays = 5
    num_reps = 20
    num_z = len(z_voltages)
    width = 380
    height = 128
    less_rows = 3
    data = np.zeros((num_z, num_delays, height - less_rows * 2, width))
    data_ctrl = np.zeros((num_z, num_delays, height - less_rows * 2, width))

    for z_index, z_v in enumerate(z_voltages):
        print('Piezo voltage', z_v)
        filename = (
            './../../stimulated_emission_imaging-data' +
            '/2016_11_02_STE_z_stack_depletion' +
            '/STE_darkfield_113_green_1500mW_red_300mW' +
            z_v + '_many_delays.tif')
        filename_ctrl = (
            './../../stimulated_emission_imaging-data' +
            '/2016_11_02_STE_z_stack_depletion' +
            '/STE_darkfield_113_green_1500mW_red_0mW' +
            z_v + '_many_delays.tif')
        data_z = np_tif.tif_to_array(filename).astype(np.float64) - bg_level
        data_z_ctrl = np_tif.tif_to_array(filename_ctrl).astype(np.float64) - bg_level
        # reshape data arrays: 0th axis is rep number, 1st is delay number
        data_z = data_z.reshape(num_reps, num_delays, height, width)
        data_z_ctrl = data_z_ctrl.reshape(num_reps, num_delays, height, width)
        # crop to remove overexposed rows (not necessary for low signal
        # measurements like fluorescence but whatever)
        data_z = data_z[:, :, less_rows:height - less_rows, :]
        data_z_ctrl = data_z_ctrl[:, :, less_rows:height - less_rows, :]

        # repetition average both image sets
        data_z_rep_avg = data_z.mean(axis=0)
        data_z_ctrl_rep_avg = data_z_ctrl.mean(axis=0)

        # registration shift control data to match "green on" data
        align_to_this_slice = data_z_rep_avg[0, :, :]
        print("Computing registration shift (no red)...")
        shifts = stack_registration(
            data_z_ctrl_rep_avg,
            align_to_this_slice=align_to_this_slice,
            refinement='integer',
            register_in_place=True,
            background_subtraction='edge_mean')
        print("... done computing shifts.")

        # build arrays with repetition averaged images
        data[z_index, ...] = data_z_rep_avg
        data_ctrl[z_index, ...] = data_z_ctrl_rep_avg

    # from the image where red/green are simultaneous, subtract the
    # average of images taken when the delay magnitude is greatest
    depletion_stack = (
        data[:,2,:,:] - # zero red/green delay
        0.5 * (data[:,0,:,:] + data[:,4,:,:]) # max red/green delay
        )
    crosstalk_stack = (
        data_ctrl[:,2,:,:] - # zero red/green delay
        0.5 * (data_ctrl[:,0,:,:] + data_ctrl[:,4,:,:]) # max red/green delay
        )
    # darkfield image (no STE) stack
    fluorescence_stack = 0.5 * (data[:,0,:,:] + data[:,4,:,:])

    # plot darkfield and stim emission signal
    top = 0
    bot = 106
    left = 92
    right = 240
    fluorescence_cropped = fluorescence_stack[:,top:bot,left:right]
    depletion_cropped = (depletion_stack[:,top:bot,left:right] -
                   crosstalk_stack[:,top:bot,left:right])

    # Our pixels are tiny (8.7 nm/pixel) to give large dynamic range.
    # This is not great for viewing, because fluctuations can swamp the
    # signal. This step bins the pixels into a more typical size.
    bucket_width = 8 # bucket width in pixels
    fluorescence_cropped = bucket(
        fluorescence_cropped, (1, bucket_width, bucket_width)) / bucket_width ** 2
    depletion_cropped = bucket(
        depletion_cropped, (1, bucket_width, bucket_width)) / bucket_width ** 2

    depletion_y_z = np.zeros((
        depletion_cropped.shape[0]*3, depletion_cropped.shape[1]))
    fluorescence_y_z = np.zeros((
        fluorescence_cropped.shape[0]*3, fluorescence_cropped.shape[1]))
    depletion_x_z = np.zeros((
        depletion_cropped.shape[0]*3, depletion_cropped.shape[2]))
    fluorescence_x_z = np.zeros((
        fluorescence_cropped.shape[0]*3, fluorescence_cropped.shape[2]))

    for z_num in range(fluorescence_cropped.shape[0]):
        # get darkfield and STE images and create yz and xz views
        depletion_image = depletion_cropped[z_num,:,:]
        depletion_y_z[z_num*3:(z_num+1)*3,:] = depletion_image.mean(axis=1)
        depletion_x_z[z_num*3:(z_num+1)*3,:] = depletion_image.mean(axis=0)

        fluorescence_image = fluorescence_cropped[z_num,:,:]
        fluorescence_y_z[z_num*3:(z_num+1)*3,:] = fluorescence_image.mean(axis=1)
        fluorescence_x_z[z_num*3:(z_num+1)*3,:] = fluorescence_image.mean(axis=0)

        # generate and save plot

        # scale bar
        fluorescence_image[-2:-1, 1:6] = 353.11
        depletion_image[-2:-1, 1:6] = -34.01

        fig, (ax0, ax1) = plt.subplots(nrows=1,ncols=2,figsize=(14,5))

        cax0 = ax0.imshow(fluorescence_image, cmap=plt.cm.gray,
                          interpolation='nearest', vmax=353.11, vmin=0)
        ax0.axis('off')
        cbar0 = fig.colorbar(cax0,ax=ax0)
        ax0.set_title('A', fontsize=30)

        cax1 = ax1.imshow(depletion_image, cmap=plt.cm.gray,
                          interpolation='nearest', vmax=0.61, vmin=-34.01)
        cbar1 = fig.colorbar(cax1, ax = ax1)
        ax1.set_title('B', fontsize=30)
        ax1.axis('off')
        plt.close()

    # x and y projections of stack of fluorescence images
    fluorescence_y_z[-2:-1, 1:6] =  np.max(fluorescence_y_z,(0,1))# scale bar
    fluorescence_x_z[-2:-1, 1:6] =  np.max(fluorescence_x_z,(0,1))# scale bar
    fig, (ax0, ax1) = plt.subplots(nrows=1,ncols=2, figsize=(14,14))
    cax0 = ax0.imshow(fluorescence_y_z, cmap=plt.cm.gray,
                      interpolation='nearest')
    ax0.axis('off')
    cbar0 = fig.colorbar(cax0,ax=ax0)
    ax0.set_title('A', fontsize=30)

    cax1 = ax1.imshow(fluorescence_x_z, cmap=plt.cm.gray,
                      interpolation='nearest')
    cbar1 = fig.colorbar(cax1, ax = ax1)
    ax1.set_title('B', fontsize=30)
    ax1.axis('off')
    plt.show()

    # x projection of stack of fluorescence and depletion images
    fluorescence_y_z[-2:-1, 1:6] =  np.max(fluorescence_y_z,(0,1))# scale bar
    depletion_y_z[-2:-1, 1:6] =  np.min(depletion_y_z,(0,1))# scale bar
    fig, (ax0, ax1) = plt.subplots(nrows=1,ncols=2, figsize=(14,14))
    cax0 = ax0.imshow(fluorescence_y_z, cmap=plt.cm.gray,
                      interpolation='nearest')
    ax0.axis('off')
    cbar0 = fig.colorbar(cax0,ax=ax0)
    ax0.set_title('A', fontsize=30)

    cax1 = ax1.imshow(depletion_y_z, cmap=plt.cm.gray,
                      interpolation='nearest')
    cbar1 = fig.colorbar(cax1, ax = ax1)
    ax1.set_title('B', fontsize=30)
    ax1.axis('off')
    plt.show()

    # y projection of stack of fluorescence and depletion images
    fluorescence_x_z[-2:-1, 1:6] =  np.max(fluorescence_x_z,(0,1))# scale bar
    depletion_x_z[-2:-1, 1:6] =  np.min(depletion_x_z,(0,1))# scale bar

    fig, (ax0, ax1) = plt.subplots(nrows=1,ncols=2,figsize=(14,14))
    cax0 = ax0.imshow(fluorescence_x_z, cmap=plt.cm.gray,
                      interpolation='nearest')
    ax0.axis('off')
    cbar0 = fig.colorbar(cax0,ax=ax0)
    ax0.set_title('A', fontsize=30)

    cax1 = ax1.imshow(depletion_x_z, cmap=plt.cm.gray,
                      interpolation='nearest')
    cbar1 = fig.colorbar(cax1, ax = ax1)
    ax1.set_title('B', fontsize=30)
    ax1.axis('off')
    plt.savefig('./../images/figure_A4/fluorescence_nd_image_xz.svg')
    plt.show()

    # average points around center lobe of the nanodiamond image to get
    # "average signal level" for darkfield and STE images
    top = 18
    bot = 78
    left = 133
    right = 188
    depletion_signal = (
        depletion_stack[:,top:bot,left:right].mean(axis=2).mean(axis=1))
    crosstalk_signal = (
        crosstalk_stack[:,top:bot,left:right].mean(axis=2).mean(axis=1))
    fluorescence_signal = (
        fluorescence_stack[:,top:bot,left:right].mean(axis=2).mean(axis=1))

    # plot signal v z
    z_list = a = np.arange(-1400,2301,200)
    plt.figure()
    plt.plot(z_list,fluorescence_signal,'.-',color='black')
    plt.title('Fluorescence image main lobe brightness')
    plt.xlabel('Z (nm)')
    plt.ylabel('Average intensity (CMOS pixel counts)')
    plt.grid()
    plt.show()
    plt.figure()
    plt.plot(z_list,depletion_signal,'.-',label='depletion signal',color='blue')
    plt.plot(z_list,crosstalk_signal,'.-',label='AOM crosstalk',color='green')
    plt.title('Depletion signal main lobe intensity')
    plt.xlabel('Z (nm)')
    plt.ylabel('Change in fluorescence light signal (CMOS pixel counts)')
    plt.legend(loc='lower right')
    plt.grid()
    plt.show()

    return None

main()

ext: py

reason: import numpy


text:
import numpy as np
from src.data import Problem, Case, Matter

def fill_somewhere(x_arr: np.array, del_color: int) -> np.array:

    assert min(x_arr.shape) >= 15
    assert x_arr.min() >= 0

    if (x_arr == del_color).sum() == 0:
        return x_arr

    x_arr_copy = x_arr.copy()
    x_arr_copy[x_arr == del_color] = -1

    x_sum = (x_arr == del_color).sum(axis=1)
    y_sum = (x_arr == del_color).sum(axis=0)

    x0 = min([i for i in range(x_arr.shape[0]) if x_sum[i]])
    x1 = max([i for i in range(x_arr.shape[0]) if x_sum[i]]) + 1
    y0 = min([i for i in range(x_arr.shape[1]) if y_sum[i]])
    y1 = max([i for i in range(x_arr.shape[1]) if y_sum[i]]) + 1

    x0_ = max(x0 - 1, 0)
    x1_ = min(x1 + 1, x_arr.shape[0])
    y0_ = max(y0 - 1, 0)
    y1_ = min(y1 + 1, x_arr.shape[1])

    search_arr = x_arr_copy[x0_:x1_, y0_:y1_].copy()

    for t in range(4):
        for i in range(x_arr_copy.shape[0] - search_arr.shape[0] + 1):
            for j in range(x_arr_copy.shape[1] - search_arr.shape[1] + 1):
                # perfect on >= 0
                match_arr = (x_arr_copy[i:i+search_arr.shape[0], j:j+search_arr.shape[1]] == search_arr).astype(int)
                match_arr[x_arr_copy[i:i+search_arr.shape[0], j:j+search_arr.shape[1]] < 0] = 1
                match_arr[search_arr < 0] = 1
                if match_arr.min() == 0:
                    continue
                # no match on negative
                if x_arr_copy[i:i+search_arr.shape[0], j:j+search_arr.shape[1]].min() < 0:
                    continue

                new_v = x_arr_copy[i:i+search_arr.shape[0], j:j+search_arr.shape[1]].copy()
                if t % 2 == 0:
                    new_v = new_v.reshape((x1_ - x0_, y1_ - y0_))
                else:
                    new_v = new_v.reshape((y1_ - y0_, x1_ - x0_))
                x_arr_copy[x0_:x1_, y0_:y1_] = np.rot90(new_v, 4 - t)
                return x_arr_copy

        search_arr = np.rot90(search_arr)

    raise AssertionError

class FillSomewhere:

    def __init__(self):
        pass

    @classmethod
    def case(cls, c: Case) -> Case:
        if c.color_delete is not None:
            new_case = c.copy()
            new_values = fill_somewhere(c.repr_values(), c.color_delete)
            new_case.matter_list = [Matter(new_values, background_color=c.background_color, new=True)]
            return new_case
        else:
            return c

    @classmethod
    def problem(cls, p: Problem) -> Problem:
        q: Problem = p.copy()
        q.train_x_list = [cls.case(c) for c in p.train_x_list]
        q.test_x_list = [cls.case(c) for c in p.test_x_list]
        return q

if __name__ == "__main__":
    xx_arr = np.array([[9, 9, 3, 1, 2], [9, 9, 3, 1, 2], [3, 3, 3, 3, 3]])
    xxx_arr = np.zeros((15, 15), dtype=int)
    xxx_arr[:3, :5] = xx_arr
    yyy_arr = fill_somewhere(xxx_arr, 9)
    print(xxx_arr)
    print(yyy_arr)


text:
import sys
import os
sys.path.append(os.path.split(__file__)[0])
from _andorsdk import AndorSDK

import time
import numpy as np
import matplotlib.pyplot as plt

class Andor():
    def __init__(self, targetT, Npixels, Wpixels, devNr=0, dllpath=""):
        """
        Init highlevel Andor Class.
        All Measurements in Full Vertical Binning (FVB).

        inputs:
          - targetT:    Target CCD Temperature (degrees C)
          - Npixels:    Number of Pixels on CCD
          - Wpixels:    Pixel Width of each Px on CCD (in microns)
          - devNr:      (optional) device number, if several installed. default = 0.
          - dllpath:    (optional) path to andor/shamrock DLLs. default = "".

        """
        self.Npixels = Npixels
        self.Wpixels = Wpixels
        self.devNr   = devNr
        self.dllpath = dllpath
        self.targetT = targetT

        self.ready       = False
        self.status      = ""
        self.Tstable     = False
        self.tempMsg     = ""
        self.temperature = 0
        self.calib       = []
        self.minLambda   = 0
        self.maxLambda   = 0
        self.atZeroOrder = False
        self.wavelength  = 0
        self.grating     = 0
        self.exposure    = 0
        self.accumulate  = 0
        self.kinetic     = 0
        self.wait_time   = 0
        self.lastData    = []

    def loadAndor(self):
        self.sdk = AndorSDK(dllpath=self.dllpath)

        ## Call init functions (get grating, wavelength, calibration...)
        self.device =  self.devNr
        self.setCCDConfig(self.Npixels, self.Wpixels)
        self.getGrating()
        self.gratingLines = self.getGratingInfo()[0]
        self.getWavelength()
        self.setExposureTime(0.1)
        self.getExposureTime()
        self.getStatus()

        self.sdk.SetReadMode(0)              # Full vert. binning

        self.setCooler(True)
        self.setTemperature(self.targetT)
        self.ready = True
        time.sleep(.1)

    def unloadAndor(self):
        del self.sdk
        self.ready = False
        time.sleep(.1)

    ####################################################################
    def _waitForAquisition(self, verbose=False):
        while self.getStatus() in ["DRV_ACQUIRING",
                   "DRV_KINETIC_TIME_NOT_MET","DRV_ERROR_ACK",
                   "DRV_ACQ_BUFFER", "DRV_TIME_NOT_MET"]:
            if verbose: print("not finished, waiting...")
            time.sleep(0.001)

    ## ---
    def setCCDConfig(self, Npixels, Wpixels):
        self.sdk.SetNumberPixelsSR(self.devNr, 1024)
        self.sdk.SetPixelWidthSR(self.devNr, 26)
        self.getCalibration()

    def getCalibration(self):
        _calib = self.sdk.GetCalibrationSR(self.devNr)
        self.calib = np.frombuffer(_calib[1], dtype=np.dtype('f4'))
        self.minLambda = self.calib.min()
        self.maxLambda = self.calib.max()
        return self.calib

    ## ---
    def setWavelength(self, targetWL):
        self._waitForAquisition()
        self.sdk.SetWavelengthSR(self.devNr, targetWL)
        self.getCalibration()
        self.getWavelength()

    def getWavelength(self):
        self.wavelength = self.sdk.GetWavelengthSR(self.devNr)[1]
        self.atZeroOrder = self.sdk.AtZeroOrderSR(self.devNr)[1]
        return self.wavelength

    ## ---
    def setGrating(self, gratingNr):
        self._waitForAquisition()
        self.sdk.SetGratingSR(self.devNr, gratingNr)
        self.getCalibration()
        self.getGrating()
        self.gratingLines = self.getGratingInfo()[0]

    def getGrating(self):
        self.grating = self.sdk.GetGratingSR(self.devNr)[1]
        return self.grating

    def getGratingInfo(self, gratingNr=-1):
        if gratingNr == -1: gratingNr = self.grating
        Lines, Blaze, Home, Offset = self.sdk.GetGratingInfoSR(self.devNr, gratingNr)[1:]
        return Lines, Home, Offset

    ## ---
    def setCooler(self, state):
        self._waitForAquisition()
        if state:
            self.sdk.SetCoolerMode(1)
            self.sdk.CoolerON()
        else:
            self.sdk.CoolerOFF()

    def setTemperature(self, targetT):
        self.setCooler(False)
        self.sdk.SetTemperature(int(targetT))
        self.setCooler(True)
        self.getTemperature()

    def getTemperature(self):
        MSG, T = self.sdk.GetTemperature()
        if MSG == 'DRV_TEMP_STABILIZED':
            self.Tstable = True
        else:
            self.Tstable = False
        self.tempMsg = MSG
        self.temperature = T
        return self.temperature

    ## ---
    def setExposureTime(self, exposure):
        self._waitForAquisition()
        self.sdk.SetExposureTime(exposure)
        self.getExposureTime()

    def getExposureTime(self):
        self.sdk.GetAcquisitionTimings()

        self.exposure = self.sdk._exposure
        self.accumulate = self.sdk._accumulate
        self.kinetic = self.sdk._kinetic
        self.wait_time = self.exposure + 0.02# + self.accumulate + self.kinetic + 0.05
        return self.exposure

    ## ---
    def getStatus(self):
        self.status = self.sdk.GetStatus()
        return self.status

    ## ---
    def startAquisition(self):
        self.sdk.StartAcquisition()

    def getAquisition(self, verbose=False):
        self._waitForAquisition(verbose)
        self.lastData = self.sdk.GetAcquiredData([])
        return np.array(self.lastData)

########################################################################
if __name__ == "__main__":
    andorTest = Andor(targetT=-20, Npixels=1024, Wpixels=26, devNr=0, dllpath="")
    for i in range(2):
        print("     Ready: ",andorTest.ready)
        andorTest.loadAndor()
        print( "     Ready: ",andorTest.ready)

        print(andorTest.grating)
        print( andorTest.getGratingInfo())
        print( andorTest.wavelength)
        print( andorTest.getStatus())
        print( '-----')
        print( andorTest.getTemperature())
        print( andorTest.tempMsg)
        print( '-----')

        andorTest.setExposureTime(0.1)
        print( "Total expected time per measurement (s):", andorTest.wait_time)
        plt.ion()
        for i in range(20):
            andorTest.startAquisition()
            DATA = andorTest.getAquisition()
            print( " T:", andorTest.getTemperature(), ' ({})'.format(andorTest.tempMsg))

            plt.clf()
            plt.title("MEASUREMENT #{}".format(i+1))
            plt.plot(andorTest.calib, DATA)
            plt.draw()
            plt.pause(.01)
        plt.ioff()
        plt.show()

        print( "     Ready: ",andorTest.ready)
        print( 30*"-" + "\nUNLOAD ANDOR...\n\n\n\n" + 30*"-")
        andorTest.unloadAndor()
        print( "     Ready: ",andorTest.ready)


text:
"""
import numpy as np
import cv2
import matplotlib.pyplot as plt

#generate black square to display
img = np.zeros(shape=(512,512,3))
plt.imshow(img)

#drawing shapes on images
#cv2.shape(line,rectangle, etc)(image, pt1,pt2, color, thickness)
linediag_red = cv2.line(img,(0,0),(511,511),(255,0,0),5)
plt.imshow(linediag_red)

linediag_green = cv2.line(img,(0,511),(511,0),(0,255,0),8)
plt.imshow(linediag_green)

#drawing blue rectangle
rectangle = cv2.rectangle(img,(384,0),(510,128),(0,0,255),5)
plt.imshow(rectangle)

#circle, need center and radius
circle_blue = cv2.circle(img,(447,63),63,(0,0,255),-1)
plt.imshow(circle_blue)

#writing on images
font = cv2.FONT_HERSHEY_SIMPLEX
text = cv2.putText(img,'OpenCV',(10,500), font, 4,(255,255,255),2,cv2.LINE_AA)
plt.imshow(text)

cv2.imshow('tutorial',img)
"""

max_stars_repo_name: dcstang/posture-detection

max_issues_repo_name: dcstang/posture-detection

max_forks_repo_name: dcstang/posture-detection


text:
from satlasaddon import RateModel, RateModelPolar
import satlas as sat
sat.set(['standard'])
import numpy as np
import matplotlib.pyplot as plt
import scipy.constants as csts

EV_TO_MHZ = csts.physical_constants['electron volt-hertz relationship'][0] * 1e-6

ABC = [[-520, 0, 0], [-3070, 0, 0]]
A_array = np.array([[0, (3.854e-9)**-1], [0, 0]])
energies = [4.422440762145948, 0]
I = 0.5
J = [0.5, 0.5]
L = [0, 1]
laser_intensity = 50
scale = -0.15

laser_mode = -1
centroids = [5852]
# centroids = [0]

f_laser = 1071455771.3
f_trans = 1069341427.47
iscool = 50019.34633
A=0.075760925
B=0.99935566
C=-1.7894284e-009
iscool = A+B*iscool+C*iscool**2

fluke = -6.086265000000001
A=999.97869
B=0.0031695015
if fluke>=0:
    fluke = A*fluke+B*fluke*fluke
else:
    fluke = A*fluke-B*fluke*fluke

data = np.loadtxt('minus.dat')
line_voltage = data[:, 0]
total_voltage = iscool - fluke - 50.428133 * line_voltage
mass = 30.996546

import collaps
doppler = collaps.dopplerfactor(mass, total_voltage)
frequencies = f_laser * doppler
asymmetry = data[:, 1]

background = -.02

model1 = RateModelPolar(I, J, L, ABC, centroids, energies, A_array, laser_intensity=laser_intensity, scale=scale, laser_mode=laser_mode, interaction_time=1e-5, background=background, shape='lorentzian', field=0)
# model2 = RateModelPolar(I, J, L, ABC, centroids, energies, A_array, laser_intensity=laser_intensity, scale=scale, laser_mode=laser_mode, interaction_time=1e-5, background=background, shape='voigt', field=0)
variation = {'Laser_intensity': False, 'Interaction_time': False, 'Transition_strength_0_to_1': False, 'FWHMG': True, 'Asymmetry_parameter': True, 'Field': False}
model1.set_variation(variation)

fig, ax = plt.subplots(1, 1)
ax.set_xlabel('Relative frequency [MHz]')
ax.set_ylabel('Asymmetry')
ax.set_title(r'$^{31}$Mg')
ax.errorbar(frequencies - f_trans, asymmetry, yerr=data[:, 2], fmt='o')
sat.chisquare_fit(model1, frequencies, asymmetry, yerr=data[:, 2])
# sat.chisquare_fit(model2, frequencies, asymmetry, yerr=data[:, 2])
from contextlib import redirect_stdout
with open('results_for_agi.txt', 'w') as f:
    with redirect_stdout(f):
        model1.display_chisquare_fit()
model1.save('test_model.spectrum')
# model2.display_chisquare_fit()
# frequencies = np.array([0])
ax.plot(frequencies - f_trans, model1(frequencies))
# fig.savefig('Test_31Mg_voigt.pdf', bbox_inches='tight')

# figc, axc, cbar = sat.generate_correlation_map(model1, frequencies, asymmetry, method='chisquare', fit_kws={'yerr': data[:, 2]}, filter=['A_level', 'Centroid', 'Scale'])
# figc.savefig('Test_31Mg_voigt_correlation.pdf', bbox_inches='tight')
plt.show()

max_stars_repo_name: woutergins/polarization

max_issues_repo_name: woutergins/polarization

max_forks_repo_name: woutergins/polarization

reason: import numpy,import scipy


text:
#!/usr/bin/python
'''
Sript to ingest historical gdelt data and save to s3.
'''
# pickle.dump(gkg_columns, open('~/gdelt_distributed_architecture/data/interim/gkg_columns.pkl', 'wb'))

# Import Dependencies

from __future__ import print_function
from __future__ import unicode_literals
import os
import io
import sys
import time
import zipfile
import requests
import argparse
import datetime
import pandas as pd
import requests
import lxml.html as lh
import numpy as np
import re
from tqdm import tqdm
import pickle
import fastparquet
import datetime
import dotenv
import logging

import boto3
import boto
# from boto.s3.key import Key
# from boto.s3.connection import S3Connection

def get_list_of_urls():
    '''
    Gets a list of gdelt archives in the form of URLs.
    INPUT:
        - NA
    OUTPUT:
        - masterlist_urls: List
                           URLs to obtain zipped CSVs.
    '''
    # Base directory for gdelt version 2 files
    gdelt_2_events_base_url = 'http://data.gdeltproject.org/gdeltv2/'

    # Directory containing all urls of historical data
    masterList_page = requests.get(gdelt_2_events_base_url+'masterfilelist.txt')
    print(type(masterList_page))

    # Requests output as text
    masterList = masterList_page.text
    print(len(masterList))

    # Extract URLs from text with regex
    masterList_urls = re.findall(r'(http|ftp|https)(://)([\w_-]+(?:(?:\.[\w_-]+)+))([\w.,@?^=%&:/~+#-]*[\w@?^=%&/~+#-])?',
                                 masterList)
    print(len(masterList_urls))

    # Concatenate regex groups to unified url strings
    masterList_urls = list(map(lambda x: ''.join(x), masterList_urls))
    print(masterList_urls[0])

    return masterList_urls

def download_zip_in_chunks(directory, url):
    '''
    Downloads a zipped file in chunks.
    INPUT:
        - directory: String
                     Directory to write downloaded file.
        - url: String
               URL of the file to download.
    OUTPUT:
        - local_file: String
                      Filepath where contents were written.

    '''
    # Unique name of file at the end of a URL
    base_file = os.path.basename(url)
    print('\n\nDOWNLOADING {}'.format(base_file))

    # Path to save output
    temp_path = directory

    if base_file in os.listdir(temp_path):
        print("FILE EXISTS, CONTINUING...")
    else:
        try:
            local_file = os.path.join(temp_path, base_file)
            # Connect to url
            req = requests.get(url, stream=True)

            # Open local file and write content to it
            with io.open(local_file, 'wb') as fp:
                # Stream data from requests generator in units of 1000
                for chunk in tqdm(req.iter_content(chunk_size=1000)):
                    # Write chunk content if found
                    if chunk:
                        fp.write(chunk)
        # Print errors and corresponding urls
        except Exception as e:
            print("ERROR: {}; {}".format(e, url))

def unzip_file(directory, zipped_file):
    '''
    Unzips files that have been downloaded locally and stores them in a directory.
    INPUT:
        - directory: String
                     Destination to save output.
        - zipped_file: String
                       Filepath to find zipped file for unzipping
    OUTPUT:
        - out_path: String
                    Filepath where content was written
    '''
    print('\n\nUnzipping {}'.format(zipped_file))
    # Try to unzip file
    try:
        # instatiate ZipFile object
        z = zipfile.ZipFile(zipped_file)
        # for each unit in z
        for name in tqdm(z.namelist()):
            # Open file and set output path
            f = z.open(name)
            out_path = os.path.join(directory, name)
            # Open out_path file for writing in utf-8
            with io.open(out_path, 'w', encoding='utf-8') as out_file:
                # Read unzipped content content as utf-8
                content = f.read().decode('utf-8')
                # Write content to output file
                out_file.write(content)

        print('Finished Unzippling {}'.format(zipped_file))
        return out_path

    # Print errors
    except zipfile.BadZipfile:
        print('Bad zip file for {}, passing on...'.format(zipped_file))

def make_gdelt_dataframe(csv_path):
    '''
    Takes a filepath, determines what type of gdelt data it is, and returns
    a pandas DataFrame.
    INPUT:
        - csv_path : string
    OUTPUT:
        - pandas dataframe
    '''

    if re.search(r'gkg', csv_path):
        columns = pickle.load(open('./gkg_columns.pkl', 'rb'))
        # columns = pickle.load(open(os.path.join(unzip_filepath,
        #                                         'gkg_columns.pkl'
        #                                         ),
        #                            'rb')
        #                       )
    elif re.search(r'export', csv_path):
        columns = pickle.load(open('./events_columns.pkl', 'rb'))

        # columns = pickle.load(open(os.path.join(unzip_filepath,
        #                                         'events_columns.pkl'
        #                                         ),
        #                            'rb')
        #                       )

    elif re.search(r'mentions', csv_path):
        columns = pickle.load(open('./mentions_columns.pkl', 'rb'))

        # columns = pickle.load(open(os.path.join(unzip_filepath,
        #                                         'mentions_columns.pkl'
        #                                         ),
        #                            'rb')
        #                       )
    else:
        return 'NOT A VALID GDELT FILE'

    df = pd.read_csv(csv_path, sep='\t')
    df.columns = columns

    return df

def csv_to_parquet(file_path, output_path):
    '''
    Takes a csv path and makes a pandas DataFrame out of it so it can be
    compressed as a parqeut file.
    '''
    print('\n\nCONVERTING {} TO PARQUET'.format(file_path))

    df = make_gdelt_dataframe(file_path)
    basename = os.path.basename(file_path)
    output_filename = '{}.gzip.parquet'.format(basename)
    if basename in os.listdir(output_path):
        print("FILE EXISTS, CONTINUING...")
        return output_filename
    else:
        output_filepath = os.path.join(output_path, output_filename)
        fastparquet.write(output_filepath, df, compression='GZIP')

        return output_filepath

def parquet_to_s3(parquet_path):

    pf = fastparquet.ParquetFile(parquet_path)

    # Setup firehose connection
    firehose = boto3.client('firehose', region_name='us-west-2')

    firehose.put_record(DeliveryStreamName='gdelt-firehose',\
    Record={'Data': pf})

def upload_file_to_s3(csv_path):
    print('\n\nUPLOADING {} TO S3'.format(csv_path))

    #s3 = boto.connect_s3()

    if re.search(r'gkg', csv_path):
        bucket_name = 'gdelt-streaming'
    elif re.search(r'export', csv_path):
        bucket_name = 'gdelt-streaming-events'
    elif re.search(r'mentions', csv_path):
        bucket_name = 'gdelt-streaming-mentions'
    else:
        bucket_name = 'gdelt-streaming'

    bucket = s3.lookup(bucket_name)

    current_keys = bucket.get_all_keys()
    current_keys = [k.key for k in current_keys]

    if os.path.basename(csv_path) in current_keys:
        print('FILE FOUND IN S3, CONTINUING...')
        return

    else:

        key_name = os.path.basename(csv_path)

        key = bucket.new_key(key_name)

        now = datetime.datetime.now()

        if re.search(r'gkg', csv_path):
            dataset = 'gdelt_gkg'
        elif re.search(r'export', csv_path):
            dataset = 'gdelt_events'
        elif re.search(r'mentions', csv_path):
            dataset = 'gdelt_mentions'
        else:
            dataset = 'other'

        metadata = {'dataset': dataset, 'upload_time': str(now)}

        key.metadata.update(metadata)

        key.set_contents_from_filename(csv_path)

        print('{} UPLOAD SUCCESSFUL'.format(key_name))
        #time.sleep(300)

def get_most_recent_files_to_s3():
    gdelt_last_15 = requests.get('http://data.gdeltproject.org/gdeltv2/lastupdate.txt').text

    urls = re.findall(r'(?P<url>https?://[^\s]+)', gdelt_last_15)

    for f in tqdm(urls):

        f_zip_path = download_zip_in_chunks(download_filepath, f)

        f_unzip_path = unzip_file(unzip_filepath,
                                    os.path.join(download_filepath,
                                                 os.path.basename(f)))

        df = make_gdelt_dataframe(os.path.join(f_unzip_path))

        parquet = csv_to_parquet(f_unzip_path, parquet_filepath)

        upload_status = upload_file_to_s3(parquet)
        print(upload_status)

    delete_data_after_s3_upload()
def delete_data_after_s3_upload():
    raw_path = os.path.join(PROJ_ROOT, 'data/raw')
    interim_path = os.path.join(PROJ_ROOT, 'data/interim')
    parquet_path = os.path.join(PROJ_ROOT, 'data/parquet')

    file_paths = [raw_path, interim_path, parquet_path]

    print("REMOVING LEFTOVER FILES")
    for p in tqdm(file_paths):
        for f in os.listdir(p):
            os.remove(os.path.join(p, f))

if __name__ == '__main__':

    log = logging.getLogger(__name__)
    if len(log.handlers) == 0:
        hdlr = logging.FileHandler('get_gdelt_data.log')
        formatter = logging.Formatter('[%(asctime)s] p%(process)s {%(pathname)s:%(lineno)d} %(levelname)s - %(message)s','%m-%d %H:%M:%S')
        hdlr.setFormatter(formatter)
        log.addHandler(hdlr)
        log.setLevel(logging.INFO)

    log.info("started")
    PROJ_ROOT = '.'
    #PROJ_ROOT = os.path.join(os.pardir, os.pardir)
    download_filepath = os.path.join(PROJ_ROOT, 'data/raw')
    unzip_filepath = os.path.join(PROJ_ROOT, 'data/interim')
    parquet_filepath = os.path.join(PROJ_ROOT, 'data/parquet')

    dotenv_path = os.path.join(PROJ_ROOT, '.env')
    dotenv.load_dotenv(dotenv_path)

    AWS_ACCESS_KEY = os.environ.get('AWS_ACCESS_KEY')
    AWS_SECRET_ACCESS_KEY = os.environ.get('AWS_SECRET_ACCESS_KEY')
    host = 's3-us-west-2.amazonaws.com'

    s3 = boto.connect_s3(AWS_ACCESS_KEY, AWS_SECRET_ACCESS_KEY)

    # url_list = get_list_of_urls()
    # gkg_url = 'http://data.gdeltproject.org/gdeltv2/20150220183000.gkg.csv.zip'
    # gkg_zip_path = download_zip_in_chunks(download_filepath, gkg_url)
    # gkg_unzip_path = unzip_file(unzip_filepath,
    #                             os.path.join(download_filepath,
    #                                          os.path.basename(gkg_url)
    #                                          )
    #                             )
    #
    # mentions_url = 'http://data.gdeltproject.org/gdeltv2/20150220184500.mentions.CSV.zip'
    # mentions_zip_path = download_zip_in_chunks(download_filepath, mentions_url)
    # mentions_unzip_path = unzip_file(unzip_filepath,
    #                                  os.path.join(download_filepath,
    #                                               os.path.basename(mentions_url)
    #                                               )
    #                                  )
    #
    # events_url = 'http://data.gdeltproject.org/gdeltv2/20150220184500.export.CSV.zip'
    # events_zip_path = download_zip_in_chunks(download_filepath, events_url)
    # mentions_unzip_path = unzip_file(unzip_filepath,
    #                                  os.path.join(download_filepath,
    #                                               os.path.basename(events_url)
    #                                               )
    #                                 )
    #
    #
    gkg_columns = ['GKGRECORDID', 'DATE', 'SourceCollectionIdentifier',
              'SourceCommonName', 'DocumentIdentifier', 'Counts',
              'V2Counts', 'Themes', 'V2Themes', 'Locations',
              'V2Locations', 'Persons', 'V2Persons', 'Organizations',
              'V2Organizations', 'V2Tone', 'Dates', 'GCAM',
               'SharingImage', 'RelatedImages', 'SocialImageEmbeds',
              'SocialVideoEmbeds', 'Quotations', 'AllNames', 'Amounts',
              'TranslationInfo', 'Extras']

    events_columns = ['GLOBALEVENTID', 'SQLDATE', 'MonthYear', 'Year', 'FractionDate',
       'Actor1Code', 'Actor1Name', 'Actor1CountryCode',
       'Actor1KnownGroupCode', 'Actor1EthnicCode', 'Actor1Religion1Code',
       'Actor1Religion2Code', 'Actor1Type1Code', 'Actor1Type2Code',
       'Actor1Type3Code', 'Actor2Code', 'Actor2Name', 'Actor2CountryCode',
       'Actor2KnownGroupCode', 'Actor2EthnicCode', 'Actor2Religion1Code',
       'Actor2Religion2Code', 'Actor2Type1Code', 'Actor2Type2Code',
       'Actor2Type3Code', 'IsRootEvent', 'EventCode', 'EventBaseCode',
       'EventRootCode', 'QuadClass', 'GoldsteinScale', 'NumMentions',
       'NumSources', 'NumArticles', 'AvgTone', 'Actor1Geo_Type',
       'Actor1Geo_FullName', 'Actor1Geo_CountryCode', 'Actor1Geo_ADM1Code',
       'Actor1Geo_ADM2Code',
       'Actor1Geo_Lat', 'Actor1Geo_Long', 'Actor1Geo_FeatureID',
       'Actor2Geo_Type', 'Actor2Geo_FullName', 'Actor2Geo_CountryCode',
       'Actor2Geo_ADM1Code',
       'Actor2Geo_ADM2Code',
        'Actor2Geo_Lat', 'Actor2Geo_Long',
       'Actor2Geo_FeatureID', 'ActionGeo_Type', 'ActionGeo_FullName',
       'ActionGeo_CountryCode', 'ActionGeo_ADM1Code',
       'ActionGeo_ADM2Code',
        'ActionGeo_Lat',
       'ActionGeo_Long', 'ActionGeo_FeatureID', 'DATEADDED', 'SOURCEURL']

    mentions_columns = ['GLOBALEVENTID', 'EventTimeDate', 'MentionTimeDate',
                           'MentionType', 'MentionSourceName', 'MentionIdentifier',
                           'SentenceID', 'Actor1CharOffset', 'Actor2CharOffset',
                           'ActionCharOffset', 'InRawText', 'Confidence',
                           'MentionDocLen', 'MentionDocTone',
                           'MentionDocTranslationInfo', 'Extras']

    pickle.dump(gkg_columns, open('./gkg_columns.pkl', 'wb'))

    pickle.dump(events_columns, open('./events_columns.pkl', 'wb'))

    pickle.dump(mentions_columns, open('./mentions_columns.pkl', 'wb'))
    #
    # gkg = make_gdelt_dataframe(os.path.join(PROJ_ROOT, 'data/raw/20150220183000.gkg.csv'))
    # events = make_gdelt_dataframe(os.path.join(PROJ_ROOT, 'data/raw/20150220184500.export.CSV'))
    # mentions = make_gdelt_dataframe(os.path.join(PROJ_ROOT, 'data/raw/20150220184500.mentions.CSV'))

    # csv_to_parquet(os.path.join(PROJ_ROOT,'data/raw/20150220183000.gkg.csv'),
    #                             parquet_filepath)
    #
    # p = '../../data/parquet/20150220183000.gkg.csv.gzip.parquet'
    #
    # csv = '../../data/interim/20150220183000.gkg.csv'

    #upload_file_to_s3(p,'gdelt-streaming')

    while True:
        log.info("started")
        print('BEGINNING UPLOAD ########### {}'.format(datetime.datetime.now()))
        get_most_recent_files_to_s3()
        print('SLEEPING')
        time.sleep(60)


text:
import pandas as pd
import numpy as np
import os
from utils.utils import *
from pathlib import Path

base_path = os.path.dirname(os.path.realpath(__file__))
base_path = str(Path(base_path).parent.absolute())

with open(base_path + '/datasets/TS1_ids') as f:
    TS1_ids = f.read().splitlines()
with open(base_path + '/datasets/TS2_ids') as f:
    TS2_ids = f.read().splitlines()
with open(base_path + '/datasets/TS3_ids') as f:
    TS3_ids = f.read().splitlines()

n = 1  # top L/n contacts

#predictors = ['GREMLIN', 'plmDCA', 'mfDCA', 'PLMC', 'RNAfold', 'LinearPartition', 'SPOT-RNA-2D-Single', 'SPOT-RNA-2D']
predictors = ['RNAfold', 'LinearPartition', 'SPOT-RNA-2D-Single', 'SPOT-RNA-2D']
sets = [TS1_ids, TS2_ids, TS3_ids]

print('\n \t \t\t TS1 \t      TS2 \t    TS3\n' )
print('\n \t \t     ss   non-ss   ss   non-ss   ss   non-ss' )

for number, pred in enumerate(predictors[0:]):

	for num, ids in enumerate(sets[0:]):

		# initialize empty count and lists to count no. of RNAs and save f1 scores
		count = 0;		save_base_pairs_f1 = [];	save_non_base_pairs_f1 = []

		for K,k in enumerate(ids[0:]):

			# read RNA sequence
			with open(base_path + '/datasets/all_sequences/' + str(k)) as f:
				temp_1 = pd.read_csv(f, comment='#', delim_whitespace=True, header=None, skiprows=[0]).values
			seq = [j.upper() for j in temp_1[0, 0]]

			# read list of native secondary structure base pairs (canonical, non-canonical, pseudoknots, multiplets, lone-pairs)
			native_ss_base_pairs = np.loadtxt(base_path + '/datasets/ss_base_pairs_labels/' + k + '.bps')
			native_ss_base_pairs = [ [int(i[0]), int(i[1])] for i in native_ss_base_pairs]  # convert float base-pairs [1.0, 67.0] to int base-pair [1, 67]
			native_ss_base_pairs = [i for i in native_ss_base_pairs if abs(i[0]-i[1]) >= 24]  # extract long-range base-pairs only

			# read native distance-based native contacts LxL matrix
			native_dist_contacts = np.loadtxt(base_path + '/datasets/distance_contact_labels/' + k + '.true')
			native_dist_contacts[native_dist_contacts==-1] = 0
			tri_inds = np.where(native_dist_contacts==1)
			native_dist_contacts = [[i,j] for i,j in zip(tri_inds[0], tri_inds[1]) if abs(i-j) >= 24] # extract long-range base-pairs only

			native_non_ss_pairs = [i for i in native_dist_contacts if i not in native_ss_base_pairs]

			# read predicted contacts LxL matrix from different predictors
			if pred == 'RNAfold':
				pred_contacts_matrix = RNAfold_bp_prob(base_path, k, seq)
			elif pred == 'LinearPartition':
				pred_contacts_matrix = LinearPartition(base_path, k, seq)
			elif pred == 'CentroidFold':
				pred_contacts_matrix = CentroidFold_mfe_bp_prob(base_path, k, seq)
			elif pred == 'SPOT-RNA':
				pred_contacts_matrix = spotrna(base_path, k, seq)
			elif pred == 'SPOT-RNA2':
				pred_contacts_matrix = spotrna2(base_path, k, seq)
			elif pred == 'GREMLIN':
				pred_contacts_matrix = GREMLIN_prob(base_path, k, seq)
			elif pred == 'PLMC':
				pred_contacts_matrix = plmc_prob(base_path, k, seq)
			elif pred == 'mfDCA':
				pred_contacts_matrix = mfdca_prob(base_path, k, seq)
			elif pred == 'plmDCA':
				pred_contacts_matrix = plmdca_prob(base_path, k, seq)
			elif pred == 'SPOT-RNA-2D-Single':
				pred_contacts_matrix = spotrna_2d_single(base_path, k, seq)
			elif pred == 'SPOT-RNA-2D':
				pred_contacts_matrix = spotrna_2d(base_path, k, seq)

			# extract upper triangular top L long-range contacts predicted proability LxL matrix
#			tri_inds = np.triu_indices(pred_contacts_matrix.shape[0], k=1)
#			all_pred_contacts = np.array([[i,j,pred_contacts_matrix[i,j]] for i,j in zip(tri_inds[0], tri_inds[1]) if abs(i-j) >= 24])

			tri_inds = np.triu_indices(pred_contacts_matrix.shape[0], k=24)
			all_pred_contacts = np.array([[i,j,pred_contacts_matrix[i,j]] for i,j in zip(tri_inds[0], tri_inds[1]) if [i, j] not in native_non_ss_pairs])

			all_pred_contacts_sorted = all_pred_contacts[all_pred_contacts[:,2].argsort()[::-1]]
			pred_contacts = [[int(i[0]), int(i[1])] for i in all_pred_contacts_sorted[0:int(len(seq)/n)]]

########################## seconadary structure base pair perfornace measure ##################
			true_pairs = native_ss_base_pairs

			pred_correctly = [i for i in pred_contacts if i in true_pairs]; #print(len(pred_correctly))
			pred_wrongly = [i for i in pred_contacts if i not in true_pairs]; #print(len(pred_wrongly))
#			pred_wrongly = [i for i in pred_contacts if i not in native_dist_contacts]; #print(len(pred_wrongly))

			tp = len(pred_correctly)
			fp = len(pred_wrongly)
			fn = len([i for i in true_pairs if i not in pred_contacts])

			try:
				pre = tp / (tp + fp)
				sen = tp / (tp + fn)
				f1 = 2*((pre*sen)/(pre + sen))
			except:
				f1 = 0

			save_base_pairs_f1.append(f1)

########################## non-base pair perfornace measure ##################

			true_pairs = native_non_ss_pairs

			# read predicted contacts LxL matrix from different predictors
			if pred == 'RNAfold':
				pred_contacts_matrix = RNAfold_bp_prob(base_path, k, seq)
			elif pred == 'LinearPartition':
				pred_contacts_matrix = LinearPartition(base_path, k, seq)
			elif pred == 'CentroidFold':
				pred_contacts_matrix = CentroidFold_mfe_bp_prob(base_path, k, seq)
			elif pred == 'SPOT-RNA':
				pred_contacts_matrix = spotrna(base_path, k, seq)
			elif pred == 'SPOT-RNA2':
				pred_contacts_matrix = spotrna2(base_path, k, seq)
			elif pred == 'GREMLIN':
				pred_contacts_matrix = GREMLIN_prob(base_path, k, seq)
			elif pred == 'PLMC':
				pred_contacts_matrix = plmc_prob(base_path, k, seq)
			elif pred == 'mfDCA':
				pred_contacts_matrix = mfdca_prob(base_path, k, seq)
			elif pred == 'plmDCA':
				pred_contacts_matrix = plmdca_prob(base_path, k, seq)
			elif pred == 'SPOT-RNA-2D-Single':
				pred_contacts_matrix = spotrna_2d_single(base_path, k, seq)
			elif pred == 'SPOT-RNA-2D':
				pred_contacts_matrix = spotrna_2d(base_path, k, seq)

			tri_inds = np.triu_indices(pred_contacts_matrix.shape[0], k=24)
			all_pred_contacts = np.array([[i,j,pred_contacts_matrix[i,j]] for i,j in zip(tri_inds[0], tri_inds[1]) if [i, j] not in native_ss_base_pairs])

			all_pred_contacts_sorted = all_pred_contacts[all_pred_contacts[:,2].argsort()[::-1]]
			pred_contacts = [[int(i[0]), int(i[1])] for i in all_pred_contacts_sorted[0:int(len(seq)/n)]]

			pred_correctly = [i for i in pred_contacts if i in true_pairs]; #print(len(pred_correctly))
			pred_wrongly = [i for i in pred_contacts if i not in true_pairs]; #print(len(pred_wrongly))
#			pred_wrongly = [i for i in pred_wrongly if i not in native_ss_base_pairs]; #print(len(pred_wrongly))

			tp = len(pred_correctly)
			fp = len(pred_wrongly)
			fn = len([i for i in true_pairs if i not in pred_contacts])

			try:
				pre = tp / (tp + fp)
				sen = tp / (tp + fn)
				f1 = 2*((pre*sen)/(pre + sen))
			except:
				f1 = 0

			save_non_base_pairs_f1.append(f1)

			count += 1

		if num==0:
			print(pred + " "*(len('SPOT-RNA-2D-Single')-len(pred)),end=' ')
		print(' {0:.3f}'.format(np.mean(save_base_pairs_f1)), end=' ')
		print(' {0:.3f}'.format(np.mean(save_non_base_pairs_f1)), end=' ')
		#print(' {}'.format(count), end='\t')
	print()

print()


text:
import gtracr.lib.igrf_utils as iuf
from gtracr.lib.constants import EARTH_RADIUS, G10, DEG_PER_RAD
'''
Library that controls the equation for the Earth's magnetic field.
'''

import os
import sys
import os.path as p
import numpy as np
from scipy.interpolate import interp1d

CURRENT_DIR = os.path.dirname(os.path.realpath(__file__))
PARENT_DIR = os.path.dirname(CURRENT_DIR)

class MagneticField:
    '''
    The base class that contains the properties of
    Earth's geomagnetic field defined in spherical coordinates.

    The base class uses the ideal dipole approximation
    in which the first-order approximation of the series expansion
    of the magnetic potential of Earth's geomagnetic field
    is used.

    Notes
    ------
    - The Earth's magnetic field is assumed to consist of
    no external currents, i.e. curl(B) = 0

    '''
    def __init__(self):
        pass

    def values(self, r, theta, phi):
        '''
        The values of the magnetic field in spherical
        coordinates at a given location (r, theta, phi)

        Parameters
        ----------
        - r : float
            the radial component in meters
        - theta : float
            the polar component in radians
        - phi : float
            the phi component in radians

        Returns
        -------
        - Br, Btheta, Bphi: np.array(float), size 3
            The spherical components of the magnetic field
        '''

        Br = -2. * (EARTH_RADIUS / r)**3. * G10 * np.cos(theta)
        Btheta = -(EARTH_RADIUS / r)**3. * G10 * np.sin(theta)
        Bphi = 0.

        return np.array([Br, Btheta, Bphi])

class IGRF13(MagneticField):
    '''
    Derived class from MagneticField class that
    describes Earth's geomagnetic field using the
    International Geomagnetic Reference Field (IGRF)
    of the 13th Edition (2020).

    The scripts that produces the IGRF magnetic field
    are obtained from the Python scripts that are
    officially constructed by the V-MOD group in the
    International Union of Geodesy and Geophysics.

    More information about the IGRF-13 model can be obtained
    here: https://www.ngdc.noaa.gov/IAGA/vmod/igrf.html

    Members
    -------
    - curr_year : float
        the current year in years
    - igrf_coeffs : np.array, float
        the Gauss coefficients for the year specified in curr_year.
    - nmax : int
        the number of truncation

    Initialization Parameters
    -------------------------
    - curr_year : float
        the current year in years. This is required as the
        magnetic field values vary in a linear fashion per year.

    - nmax : int, optional
        the number to truncate the series expansion of the
        magnetic field. The default value is obtained from
        the .shc file.

    '''
    def __init__(self, curr_year, nmax=None):
        # override MagneticField __init__ function
        # not necessary, but for decorative sake
        # super().__init__(self)

        self.curr_year = curr_year
        fpath = os.path.join(PARENT_DIR, "data", "IGRF13.shc")

        leap_year = False

        # check if current year is a leap year or not
        if self.curr_year % 4 == 0:
            leap_year = True

        # import the coefficients from .shc file
        igrf = iuf.load_shcfile(fpath, leap_year=leap_year)

        # interpolate to account for any year
        # linear interpolation since time variation in
        # coefficients is a linear one
        interp_coeffs = interp1d(
            igrf.time,
            igrf.coeffs,
            kind="linear",
        )

        # get the number of truncation based on the
        # igrf .shc file
        self.nmax = igrf.parameters["nmax"] if nmax is None else nmax

        self.igrf_coeffs = interp_coeffs(self.curr_year).T

    def values(self, r, theta, phi):
        # transform theta, phi to degrees and r into kilometers
        r *= 1e-3
        theta *= DEG_PER_RAD
        phi *= DEG_PER_RAD

        # check if theta / phi is > 180 or > 360 respectively
        # this is to prevent errors to be raised by the igrf evaluator
        if theta > 180.:
            theta = (theta % 180.)  # normalize by nearest remainder from 180

        if phi > 360.:
            phi = (phi % 360.)  # normalize by nearest remainder from 360

        # obtaining the values can easily be done by using
        # the function synth_values from igrf_utils.py
        Br, Btheta, Bphi = iuf.synth_values(self.igrf_coeffs,
                                            r,
                                            theta,
                                            phi,
                                            nmax=self.nmax)

        # take first index since synth_values return arrays, not scalars
        return np.array([Br, Btheta, Bphi]) \
                                            * (1e-9)  # convert to teslas

reason: import numpy,from scipy


text:
import os
import subprocess
import sys
sys.path.append(os.path.join(os.environ['HOME'], 'bin/VaspCZ/sourcecode'))
import zzdlib
import numpy as np

import VaspCZ

pys_path = os.path.join(os.environ['HOME'], 'bin/VaspCZ/sourcecode')

def decode_vib_frequency(data):
	# data 是shell代码出来的结果，是一个列表
	frequncies = []
	for i in range(len(data)):
		line = data[i]
		if 'THz' in line:
			frequncy = line.split('=')[-1].split()[0]
			frequncy = float(f'{float(frequncy):.4f}')  # 保留4位有效数字
			frequncies.append(frequncy)
	frequncies_np = np.array(frequncies)
	return frequncies_np

def grep_fre(path):
	code = f'grep THz {path}/OUTCAR'
	result = zzdlib.getshellResult(code)
	fre = decode_vib_frequency(result)
	return fre

def run_Sb(ele):
	vac_1nn = grep_fre(f'{ele}-Energies/Fe107V1/vib_analysis/1nn')
	vac_2nn = grep_fre(f'{ele}-Energies/Fe107V1/vib_analysis/2nn')
	print(f'vac_1nn: {vac_1nn}  vac_2nn: {vac_2nn}')
	item1 = np.power(np.prod(vac_1nn), 12) * np.power(np.prod(vac_2nn), 6)

	pure = grep_fre(f'{ele}-Energies/Fe108/vib_analysis')
	print(f'pure: {pure}')
	item2 = np.power(np.prod(pure), 10)  # 几次方

	sol = grep_fre(f'{ele}-Energies/Fe107{ele}1/vib_analysis/335050')  # 溶质位置335050
	sol_1nn = grep_fre(f'{ele}-Energies/Fe107{ele}1/vib_analysis/505066')
	sol_2nn = grep_fre(f'{ele}-Energies/Fe107{ele}1/vib_analysis/665050')
	print(f'sol: {sol}  sol_1nn: {sol_1nn}  sol_2nn: {sol_2nn}')
	item3 = np.prod(sol) * np.power(np.prod(sol_1nn), 12) * np.power(np.prod(sol_2nn), 6)

	vs = grep_fre(f'{ele}-Energies/Fe106{ele}1V1/vib_analysis/*')
	print(vs.shape)
	item4 = np.prod(vs)

	inner = item1 * item3 / item2 / item4
	Sb = np.log(inner)

	print(f'item1: {item1}\nitem2: {item2} \nitem3: {item3} \nitem4: {item4}\ninner: {inner:.4f} Sb: {Sb:.4f}')

def run_Hf_and_Hb(ele):
	# 4个文件夹下的结构优化的WARNING 检查和能量检查。
	E_ele1V1 = zzdlib.Vasp.check_WARNING_and_Energy(path=f'{ele}-Energies/Fe106{ele}1V1')
	E_pure = zzdlib.Vasp.check_WARNING_and_Energy(path=f'{ele}-Energies/Fe108')
	E_ele1 = zzdlib.Vasp.check_WARNING_and_Energy(path=f'{ele}-Energies/Fe107{ele}1')
	E_V1 = zzdlib.Vasp.check_WARNING_and_Energy(path=f'{ele}-Energies/Fe107V1')

	Hf = E_V1 - 107/108*E_pure
	Hb = E_ele1V1 + E_pure - E_ele1 - E_V1
	print(f'Hf: {Hf:.4f}  Hb: {Hb:.4f}')

def run_Sf(ele):
	vac_1nn = grep_fre(f'{ele}-Energies/Fe107V1/vib_analysis/1nn')
	vac_2nn = grep_fre(f'{ele}-Energies/Fe107V1/vib_analysis/2nn')
	# print(f'vac_1nn: {vac_1nn}  vac_2nn: {vac_2nn}')
	item1 = np.power(np.prod(vac_1nn), 12) * np.power(np.prod(vac_2nn), 6)

	pure = grep_fre(f'{ele}-Energies/Fe108/vib_analysis')
	# print(f'pure: {pure}')
	item2 = np.power(np.prod(pure), 18)  # 几次方

	inner = item2/item1
	Sf = np.log(inner)
	print(f'inner: {inner:.4f} Sf: {Sf:.4f}')

def get_w(Hm, v, T=1000):
	"""
	w = v*exp(-Hm/(kB*T))
	:return:
	"""
	# 假定T=1000
	kB = 1.380649 * 10 ** -23
	Hm_unit_j = Hm * 1.602 * 10 ** -19
	w = v * 10**12 * np.exp(-Hm_unit_j/(kB*T))
	return w

def get_f2(Hm_and_v, T=1000):
	"""
	f2 = (2*w1 + 7*w3*F)/(2*w1 + 2*w2+ 7*w3*F)
	:param Hm_and_v:
	:return:
	"""
	data = Hm_and_v
	w1 = get_w(Hm_and_v[0,2], Hm_and_v[1,2], T=T)
	w3_2nn = get_w(data[0,6], data[1,6], T=T)
	w3_3nn = get_w(data[0,8], data[1,8], T=T)
	w3_4nn = get_w(data[0,10], data[1,10], T=T)
	w3 = (2*w3_2nn + 4*w3_3nn + w3_4nn)/7
	w2 = get_w(data[0,4], data[1,4], T=T)

	# 获取zeta
	w0 = get_w(data[0,0], data[1,0], T=T)
	w4_2nn = get_w(data[0,7], data[1,7], T=T)
	w4_3nn = get_w(data[0,9], data[1,9], T=T)
	w4_4nn = get_w(data[0,11], data[1,11], T=T)
	w4 = (2*w4_2nn + 4*w4_3nn + w4_4nn)/7
	# print(w4, w0)
	zeta = w4/w0
	fenzi = (10*np.power(zeta,4) + 180.5*np.power(zeta,3) + 927*np.power(zeta,2) + 1341*zeta)
	fenmu = 7*(2*np.power(zeta,4) + 40.2*np.power(zeta,3) + 254*np.power(zeta,2) + 597*zeta + 436)
	F = 1 - fenzi/fenmu
	print(f'zeta: {zeta}')
	print(f'F: {F}')
	f2 = (2* w1 + 7 * w3 * F) / (2 * w1 + 2 * w2 + 7 * w3 * F)
	return f2

def get_Hm_and_v_from_raw():
	cases = 'w0_3nn,w1,w2,w3_2nn,w3_3nn,w3_4nn'.split(',')
	Hm_list = []
	v_list = []
	for i in range(len(cases)):
		case = cases[i]
		os.chdir(case)

		# 获取Hm
		# print(pys_path, os.getcwd())
		res = zzdlib.getshellResult(f'python3 {pys_path}/NEBCheck.py --func=2')
		# print(res)
		line, index = zzdlib.File.getLine(res, 'IMAGE')
		IMAGE = []
		Barrier = []
		# print(res)
		for j in range(index + 1, len(res)):
			if res[j] == '\n':
				break
			img = int(res[j].split()[0])
			barrier = float(res[j].split()[-1])
			IMAGE.append(img)
			Barrier.append(barrier)
		Hm_foreward = np.max(Barrier)
		Hm_backward = np.max(Barrier) - Barrier[-1]
		Hm_list.append(Hm_foreward)
		Hm_list.append(Hm_backward)

		# 获取v
		code = f'python3 {pys_path}/VaspVibAna_Result.py --isprint=True'
		res = zzdlib.getshellResult(code)
		v_foreward = zzdlib.File.getLine(res, keywords='foreward')[0].split()[-1]
		v_backward = zzdlib.File.getLine(res, keywords='backward')[0].split()[-1]
		if v_backward == 'Match':
			v_backward = '0'
		v_list.append(float(v_foreward))
		v_list.append(float(v_backward))
		os.chdir('..')

	Hm_and_v = np.concatenate((np.array(Hm_list).reshape(1, -1), np.array(v_list).reshape(1, -1)), axis=0)
	return Hm_and_v

def run_f2(ele, T=1000):
	# 先获取Hm和v的数据
	os.chdir(f'{ele}-V')

	if os.path.isfile('.Hm_and_v'):
		with open('.Hm_and_v') as f:
			data = f.read()
			Hm_and_v = np.array(eval(data))

	else:
		Hm_and_v = get_Hm_and_v_from_raw()
		with open('.Hm_and_v', 'w') as f:
			f.write(str(Hm_and_v.tolist()))

	os.chdir('..')

	# 算f0
	Hm_self = [1.3943 for i in range(12)]
	v_self = [6.2805 for i in range(12)]

	Hm_and_v_self = np.concatenate((np.array(Hm_self).reshape(1, -1), np.array(v_self).reshape(1, -1)), axis=0)
	# print(Hm_and_v_self)
	f0 = get_f2(Hm_and_v_self, T=T)
	print(f'f0: {f0}')

	# 计算
	f2 = get_f2(Hm_and_v, T=T)
	print(f'f2: {f2}')
	return f2

if __name__ == '__main__':
	os.system(f'source ~/.bashrc')
	ele = os.path.basename(os.getcwd()).split('with')[-1]
	print(f'ele: {ele}')

	run_Sb(ele)
	run_Sf(ele)
	run_Hf_and_Hb(ele)
	run_f2(ele, T=1000)

max_stars_repo_name: zybbigpy/VaspCZ

max_issues_repo_name: yizhiwoniu/VaspCZ

max_forks_repo_name: yizhiwoniu/VaspCZ


text:
import numpy as np

'''
This module contains the line shapes for the NMR signals
In general the parameters taken are

Args:
    x: axis
    intensity: intensity
    loc: location in each dimension (list)
    width: line width in each dimension (list)
    mixing: optional mixing term for combining multiple line shapes

Returns:
    array: array of values taken by that line shape
'''

#=====//===== Checks used in this module =====//=====

def check_2D(x):
    '''
    checks the length of x if it is not 2 then
    it exits python with a message
    '''

    if len(x) > 2:
        print '================================================='
        print 'only 2D cases are implemented with this lineshape'
        print '================================================='
        sys.exit()

def std_to_lwhm(width):
    '''
    this is because the widths of the Lorentzian and
    Gaussian are defined differently
    '''
    return width/2.*np.sqrt(2.*np.log(2.))

def gaussian(x, intensity, loc, width):
    '''
    Gaussian line shape
    '''

    #correction so that the lw it the width at half maximum
    width = std_to_lwhm(width)

    index = ((x-loc)**2.)/(2*(width**2))*(-1.)
    gauss =  np.e**(index)
    return gauss

def lorenzian(x, intensity, loc, width):
    '''
    Lorentzian line shape
    '''
    val = (loc-x)/width
    shape = 1./(1+(val**2))
    return shape

def rotatable_gauss_2D_shape(x, loc, width,theta):

    a1 = np.divide(np.cos(theta)**2.,2*(width[0]**2.))
    a2 = np.divide(np.sin(theta)**2.,2*(width[1]**2.))
    a = a1+ a2

    b1 = -1.*np.divide(np.sin(2.*theta), 4*(width[0]**2.))
    b2 = np.divide(np.sin(2.*theta), 4*(width[1]**2.))
    b = b1 + b2

    c1 = np.divide(np.sin(theta)**2.,2*(width[0]**2.))
    c2 = np.divide(np.cos(theta)**2.,2*(width[1]**2.))
    c = c1+c2

    x1 = np.ndarray.flatten(x[0])
    x2 = np.ndarray.flatten(x[1])

    model = np.zeros((len(x1), len(x2)))

    def rot_gauss(xi, yi,a,b,c, loc):

        term1 = a*(xi-loc[0])**2.
        term2 = 2*b*(xi-loc[0])*(yi-loc[1])
        term3 = c*(yi-loc[1])**2.

        index = -1.*(term1 + term2 + term3)

        model = (np.e**index)
        return model

    model = rot_gauss(x1[:,None], x2[None,:], a,b,c,loc)
    return model

def rotatable_lorren_2D_shape(x, loc, width,theta):

    x0_rot = loc[0]*np.cos(theta) - loc[1]*np.sin(theta)
    y0_rot = loc[0]*np.sin(theta) + loc[1]*np.cos(theta)

    x1 = np.ndarray.flatten(x[0])
    x2 = np.ndarray.flatten(x[1])

    def rot_lorren(xi, yi, x0_rot, y0_rot, theta, width):

	    xi_rot = xi*np.cos(theta) - yi*np.sin(theta)
	    yi_rot = xi*np.sin(theta) + yi*np.cos(theta)

	    frac1 = (x0_rot-xi_rot)/width[0]
	    frac1 = frac1**2.
	    frac1 = 1./(1.+frac1)

	    frac2 = (y0_rot-yi_rot)/width[1]
	    frac2 = frac2**2.
	    frac2 = 1./(1.+frac2)

	    return frac1*frac2

    return rot_lorren(x1[:,None], x2[None,:], x0_rot, y0_rot, theta, width)

def rotatable_glore_2D_shape(x, loc, width,theta, mixing):

    gauss_wahm = [std_to_lwhm(a) for a in width]

    a1 = np.divide(np.cos(theta)**2.,2*(gauss_wahm[0]**2.))
    a2 = np.divide(np.sin(theta)**2.,2*(gauss_wahm[1]**2.))
    a = a1+ a2

    b1 = -1.*np.divide(np.sin(2.*theta), 4*(gauss_wahm[0]**2.))
    b2 = np.divide(np.sin(2.*theta), 4*(gauss_wahm[1]**2.))
    b = b1 + b2

    c1 = np.divide(np.sin(theta)**2.,2*(gauss_wahm[0]**2.))
    c2 = np.divide(np.cos(theta)**2.,2*(gauss_wahm[1]**2.))
    c = c1+c2

    x1 = np.ndarray.flatten(x[0])
    x2 = np.ndarray.flatten(x[1])

    model = np.zeros((len(x1), len(x2)))

    x0_rot = loc[0]*np.cos(theta) - loc[1]*np.sin(theta)
    y0_rot = loc[0]*np.sin(theta) + loc[1]*np.cos(theta)

    def glore(xi, yi, x0_rot, y0_rot, theta, width, mixing):

        #lorrenzian
        xi_rot = xi*np.cos(theta) - yi*np.sin(theta)
        yi_rot = xi*np.sin(theta) + yi*np.cos(theta)

        frac1 = (x0_rot-xi_rot)/width[0]
        frac1 = frac1**2.
        frac1 = 1./(1.+frac1)

        frac2 = (y0_rot-yi_rot)/width[1]
        frac2 = frac2**2.
        frac2 = 1./(1.+frac2)

        lorren = frac1*frac2

        #gaussian
        term1 = a*(xi-loc[0])**2.
        term2 = 2*b*(xi-loc[0])*(yi-loc[1])
        term3 = c*(yi-loc[1])**2.

        index = -1.*(term1 + term2 + term3)
        gauss = (np.e**index)

        return mixing*gauss + (1.-mixing)*lorren

    return glore(x1[:,None], x2[None,:], x0_rot, y0_rot, theta, width, mixing)

def rotatable_x_glore_2D_shape(x, loc, width,theta, mixing):

    model = np.zeros((len(x1), len(x2)))
    x0_rot = loc[0]*np.cos(theta) - loc[1]*np.sin(theta)
    y0_rot = loc[1]
    gauss_wahm = [std_to_lwhm(a) for a in width]

    def glore(xi, yi, x0_rot, y0_rot, theta, width, mixing, gauss_wahm):

        #lorrenzian
        xi_rot = xi*np.cos(theta) - yi*np.sin(theta)
        yi_rot = yi

        frac1 = (x0_rot-xi_rot)/width[0]
        frac1 = frac1**2.
        frac1 = 1./(1.+frac1)

        frac2 = (y0_rot-yi_rot)/width[1]
        frac2 = frac2**2.
        frac2 = 1./(1.+frac2)

        lorren = frac1*frac2

        #gaussian
        index = ((xi_rot-x0_rot)**2.)/(2*(gauss_wahm[0]**2))*(-1.)
        index2 = ((yi_rot-y0_rot)**2.)/(2*(gauss_wahm[1]**2))*(-1.)

        gauss =  np.e**(index+index2)

        return mixing*gauss + (1.-mixing)*lorren

    return glore(x1[:,None], x2[None,:], x0_rot, y0_rot, theta, width, mixing)

def rotatable_lorren_2D(x, intensity, loc, width,theta):
	check_2D(x)
	return intensity*rotatable_lorren_2D_shape(x, loc, width,theta)

def rotatable_gauss_2D(x, intensity, loc, width,theta):
    width = [std_to_lwhm(a) for a in width]
    check_2D(x)
    return intensity*rotatable_gauss_2D_shape(x, loc, width,theta)

def rotatable_glore_2D(x, intensity, loc, width,theta, mixing):
    check_2D(x)
    return rotatable_glore_2D_shape(x, loc, width,theta, mixing)*intensity

def rotatable_x_glore_2D(x, intensity, loc, width,theta, mixing):
    check_2D(x)
    return rotatable_x_glore_2D_shape(x, loc, width,theta, mixing)*intensity

def ndGaussian(x, intensity, loc, width):
    '''
    Gaussian line shape
    '''
    working = 1.

    for i,j,k in zip(x, loc, width):
        working = working*gaussian(i, 1., j, k)

    final = working * intensity
    return final

def ndLorentzian(x, intensity, loc, width):
    '''
    Lorentzian line shape
    '''
    working = 1.
    for i,j,k in zip(x, loc, width):
        working = working*lorenzian(i, 1., j, k)

    return intensity*working

def ndGlore(x, intensity, loc, width, mixing):
    '''
    Gaussian  + Lorentzian line shape
    '''
    working = 1.
    for i,j,k,m in zip(x, loc, width, mixing):

        val = (m*gaussian(i, 1., j, k))+((1-m)*lorenzian(i, 1., j, k))
        working = working*val

    return working*intensity

max_stars_repo_name: L-Siemons/FrequentFitter

max_issues_repo_name: L-Siemons/FrequentFitter

max_forks_repo_name: L-Siemons/FrequentFitter


text:
import os
import matplotlib
import matplotlib.pyplot as plt
import matplotlib.collections as mc
import numpy as np
from myutils.geo import get_shp_points

palettes = {
        'saturated': ['#e41a1c','#377eb8','#4daf4a','#984ea3','#ff7f00','#ffff33','#a65628','#f781bf','#999999'],
        'pastel': ['#e41a1c','#377eb8','#4daf4a','#984ea3','#ff7f00','#ffff33','#a65628','#f781bf','#999999'],
        }

##########################################################
def export_subplots(ax, fig, labels, outdir, pad=0.3, prefix='', fmt='pdf'):
    n = ax.shape[0]*ax.shape[1]
    for k in range(n):
        i = k // ax.shape[1]
        j = k  % ax.shape[1]
        ax[i, j].set_title('')

    for k in range(n):
        i = k // ax.shape[1]
        j = k  % ax.shape[1]
        coordsys = fig.dpi_scale_trans.inverted()
        extent = ax[i, j].get_window_extent().transformed(coordsys)
        x0, y0, x1, y1 = extent.extents

        if isinstance(pad, list):
            x0 -= pad[0]; y0 -= pad[1]; x1 += pad[2]; y1 += pad[3];
        else:
            x0 -= pad; y0 -= pad; x1 += pad; y1 += pad;

        bbox =  matplotlib.transforms.Bbox.from_extents(x0, y0, x1, y1)
        fig.savefig(os.path.join(outdir, prefix + labels[k] + '.' + fmt),
                      bbox_inches=bbox)

##########################################################
def hex2rgb(hexcolours, normalize=False, alpha=None):
    n = len(hexcolours)
    rgbcolours = np.zeros((n, 3), dtype=float)

    for i, h in enumerate(hexcolours):
        rgbcolours[i, :] = np.array([int(h.lstrip('#')[i:i+2], 16) for i in (0, 2, 4)])

    if alpha != None:
        rgbcolours = np.concatenate([rgbcolours, np.ones((n, 1))*alpha], axis=1)

    if normalize:
        rgbcolours[:, :3] = rgbcolours[:, :3] / 255

    return rgbcolours

##########################################################
def create_meshgrid(x, y, nx=100, ny=100, relmargin=.1):
    """Create a meshgrid considering @x and @y bounds with @nx, @ny tiles
    and relative margins @relmargins"""

    marginx = (max(x) - min(x)) * relmargin
    marginy = (max(y) - min(y)) * relmargin
    xrange = [np.min(x) - marginx, np.max(x) + marginx]
    yrange = [np.min(y) - marginy - .15, np.max(y) + marginy]
    dx = (xrange[1] - xrange[0]) / nx
    dy = (yrange[1] - yrange[0]) / ny
    xx, yy = np.mgrid[xrange[0]:xrange[1]:(nx*1j), yrange[0]:yrange[1]:(ny*1j)]
    return xx, yy, dx, dy

#############################################################
def plot_graph_coords(vcoords, ecoords, ax, shppath):
    """Plot the grpah, with vertices colored by accessibility."""

    if not ax: fig, ax = plt.subplots(figsize=(7, 7))

    sc = ax.scatter(vcoords[:, 0], vcoords[:, 1], c='k',
            linewidths=0, alpha=.8, s=3, zorder=10) # vertices

    segs = mc.LineCollection(ecoords, colors='k', linewidths=.5, alpha=.5) # edges
    ax.add_collection(segs)

    if shppath: # border
        mapx, mapy = get_shp_points(shppath)
        ax.plot(mapx, mapy, c='dimgray')

    return ax

##########################################################
def plot_graph(gin, plotpath='/tmp/mygraph.png', inverty=False, ax=None, shppath=''):
    """Plot the graph @gin which can be an igraph object or a graphml file path. If @shppath is provided, plot the border given by @shppath."""

    import igraph
    if type(gin) == str: g = igraph.Graph.Read(gin)
    else: g = gin

    for attrs in [('lon', 'lat'), ('posx', 'posy'), ('x', 'y')]:
        if attrs[0] in g.vertex_attributes():
            xattr = attrs[0]; yattr = attrs[1]

    if inverty: ax.invert_yaxis()

    vcoords = np.array([(x, y) for x, y in zip(g.vs[xattr], g.vs[yattr])])
    vcoords = vcoords.astype(float)

    ecoords = []
    for e in g.es:
        ecoords.append([ [float(g.vs[e.source][xattr]), float(g.vs[e.source][yattr])],
                [float(g.vs[e.target][xattr]), float(g.vs[e.target][yattr])], ])

    ax = plot_graph_coords(vcoords, ecoords, ax, shppath)
    ax.axis('off')
    plt.tight_layout()
    plt.savefig(plotpath)

max_stars_repo_name: tokudaek/myutils

max_issues_repo_name: tokudaek/myutils

max_forks_repo_name: tokudaek/myutils


text:
import argparse
from scipy.optimize import curve_fit
from scipy.special import erfc
import numpy as np

# Our amino acid conversion utilities
three = ["ALA", "ARG", "ASN", "ASP", "CYS", "GLN", "GLU", "GLY", "HIS",
         "ILE", "LEU", "LYS", "MET", "PHE", "PRO", "SER", "THR", "TRP", "TYR", "VAL"]
one = ["A", "R", "N", "D", "C", "Q", "E", "G", "H", "I",
       "L", "K", "M", "F", "P", "S", "T", "W", "Y", "V"]
aadict = {}

for i in range(len(three)):
    aadict[three[i]] = one[i]
    aadict[one[i]] = three[i]

def one2three2one(aa):
    """
    Amino acid 1to3 conversion
    """
    return(aadict[aa.upper()])

def get_diagonal(matrix, n=0):
    z = np.zeros(matrix.shape, int)
    for i in range(len(z) - n):
        z[i, i + n] = True
    return matrix[z > 0]

def smallmatrixlimits(ires, cutoff, len):
    ileft = max(1, ires - cutoff)
    iright = min(ileft + 2 * cutoff, len)
    if iright == len:
        ileft = max(1, iright - 2 * cutoff)
    return (ileft, iright)

def smallmatrixpos(ires, cutoff, len):
    resi = cutoff + 1
    if ires < cutoff + 1:
        resi = ires
    if ires > len - cutoff:
        resi = min(len, 2 * cutoff + 1) - (len - ires)
    return resi

def hasselbalch(pH, pK, nH):
    return (10 ** (pK - pH)) ** nH / (1. + (10 ** (pK - pH)) ** nH)

def k(ion, T, eps):
    return np.sqrt((8.0 * np.pi * ion * 0.2003457539870666) / (eps * T * 0.0019872041))

def W1(r, ion, T, eps):
    kappa = k(ion, T, eps)
    x = kappa * r / np.sqrt(6.0)
    return 332.286 * np.sqrt(6.0 / np.pi) * (1 - np.sqrt(np.pi) * x * np.exp(x ** 2) * erfc(x)) / (eps * r)

def WP(ze, c, ion, T, eps):
    kappa = k(ion, T, eps)
    r = 11.8 / pow(c, 1.0 / 3.0)  # in Angstroms
    return ze * ze * np.exp(-kappa * r) / (eps * r)

def W(ds, ion, eps, weight=None):
    k = np.sqrt(ion) / 3.08
    ds[ds == 0] = 1
    tmp = 332.286 / (eps * ds) * np.exp(- k * ds)
    return np.average(tmp, axis=2, weights=weight)

def w2logp(x, R, T):
    return x * 4181.2 / (R * T * np.log(10))

# Argument parsing
parser = argparse.ArgumentParser()
parser.add_argument("--sequence", default='nMDVFMKGLSKAKEGVVAAAEKTKQGVAEAAGKTKEGVLYVGSKTKEGVVHGVATVAEKTKEQVTNVGGAVVTGVTAVAQKTVEGAGSIAAATGFVKKDQLGKNEEGAPQEGILEDMPVDPDNEAYEMPSEEGYQDYEPEAc', help="Protein sequence in one-letter FASTA format. [n] denotes N-terminus. [c] denotes C-terminus. Default [nMDVFMKGLSKAKEGVVAAAEKTKQGVAEAAGKTKEGVLYVGSKTKEGVVHGVATVAEKTKEQVTNVGGAVVTGVTAVAQKTVEGAGSIAAATGFVKKDQLGKNEEGAPQEGILEDMPVDPDNEAYEMPSEEGYQDYEPEAc].")
parser.add_argument("--temperature", default=283.15, help="Temperature in [K]. Default [283.15].")
parser.add_argument("--ionicstrength", default=0.0, help="Ionic strength in [M]. Default [0.0].")
parser.add_argument("--epsilon", default=83.83, help="Dielecttric permittivity. Default [83.83].")
parser.add_argument("--gca", default=5.0, help="Charge shift distance due to side chain [A]. Default [5.0].")
parser.add_argument("--gcb", default=7.5, help="The effective residue separation [A]. Default [7.5].")
parser.add_argument("--cutoff", default=2, help="Explicit calculation cutoff. Default [2].")
parser.add_argument("--ncycles", default=3, help="The number of super-cycles. Default [3].")
parser.add_argument("--nooutput", action='store_true', default=False, help="Do not generate output file. Default [False].")
parser.add_argument("--silent", action='store_true', default=False, help="Do not print calculation output. Default [False].")
args = parser.parse_args()

if not args.silent:
    print (args)

# Variables init
seq = args.sequence
temp = float(args.temperature)
ion = float(args.ionicstrength)
eps = float(args.epsilon)
gca = float(args.gca)
gcb = float(args.gcb)
cutoff = int(args.cutoff)
ncycles = int(args.ncycles)

pK0 = {"n": 7.5, "C": 8.6, "D": 4.0, "E": 4.4,
       "H": 6.6, "K": 10.4, "R": 12.0, "Y": 9.6, "c": 3.5}

q0 = {"n": 0.0, "C": -1.0, "D": -1.0, "E": -1.0,
      "H": 0.0, "K": 0.0, "R": 0.0, "Y": -1.0, "c": -1.0}

pos = np.array([i for i in xrange(len(seq)) if seq[i] in pK0.keys()])
sites = ''.join([seq[i] for i in pos])

l = np.array([abs(pos - pos[i]) for i in range(len(pos))])
d = gca + np.sqrt(l) * gcb
tmp = W1(d, ion, temp, eps)

R = 8.314472
pHstep = 0.2
pHs = np.arange(1.0, 13.01, pHstep)

neg = np.array([i for i in range(len(pos)) if seq[pos[i]] in 'CDEYc'])

#print('distancematrix: ', ds.shape)
I = np.diag(np.ones(len(pos)))
tmp[I == 1] = 0
ww = w2logp(tmp, R, temp) / 2

chargesempty = np.zeros(pos.shape[0])
if len(neg):
    chargesempty[neg] = -1

pK0s = [pK0[c] for c in sites]
nH0s = [0.9 for c in sites]

titration = np.zeros((len(pos), len(pHs)))
smallN = min(2 * cutoff + 1, len(pos))
smallI = np.diag(np.ones(smallN))

alltuples = [[int(c) for c in np.binary_repr(i, smallN)]
             for i in xrange(2 ** (smallN))]
gmatrix = [np.zeros((smallN, smallN)) for p in range(len(pHs))]

if len(sites) <= 2 * cutoff + 1:
    ncycles = 1

titration = np.array([[hasselbalch(pHs[p], pK0s[i], nH0s[i])
                       for i in range(len(pos))] for p in range(len(pHs))]).transpose()

Qs = titration

for icycle in range(ncycles):
    # print icycle + 1
    fractionhold = titration.transpose()

    for ires in range(1, len(pos) + 1):
        (ileft, iright) = smallmatrixlimits(ires, cutoff, len(pos))
        resi = smallmatrixpos(ires, cutoff, len(pos))

        for p in range(len(pHs)):
            fraction = fractionhold[p].copy()
            fraction[ileft - 1: iright] = 0
            charges = chargesempty + fraction
            ww0 = np.diag(np.dot(ww, charges) * 2)
            gmatrixfull = ww + ww0 + pHs[p] * I - np.diag(pK0s)
            gmatrix[p] = gmatrixfull[ileft - 1: iright, ileft - 1: iright]

        E_all = np.array([sum([10 ** -(gmatrix[p] * np.outer(c, c)).sum()
                               for c in alltuples]) for p in range(len(pHs))])
        E_sel = np.array([sum([10 ** -(gmatrix[p] * np.outer(c, c)).sum()
                               for c in alltuples if c[resi - 1] == 1]) for p in range(len(pHs))])
        titration[ires - 1] = E_sel / E_all

        Qs[ires-1] = titration[ires - 1] + q0[sites[ires-1]]

    sol = np.array([curve_fit(hasselbalch, pHs, titration[i], [
                   pK0s[i], nH0s[i]])[0] for i in range(len(pK0s))])
    (pKs, nHs) = sol.transpose()

first = 0
if not seq[0] == 'n':
    first = 1

# Compute the overall charge
Q = np.sum(Qs, axis=0)
G = -1.0 * w2logp(Q, R, temp)

# Print the header
if not args.silent:
    print '%s, %s, %s, %s' % ('RES','pKa', 'dpKa', 'n')

for i in range(len(sol)):
    if not args.silent:
        print '%s, %f, %f, %f' % (seq[pos[i]] + str(pos[i] + first), pKs[i], pKs[i] - pK0s[i], nHs[i]) #pKs[i]#, pKs[i] - pK0s[i], nHs[i]
    if not args.nooutput:
        outfile = open(seq[pos[i]] + str(pos[i] + first) + '_titration.dat', 'w')
        outfile.close()

if not args.nooutput:
    outfile2 = open('Total_Q.dat', 'w')
    outfile3 = open('Total_G.dat', 'w')
    for p in range(len(pHs)):
        outfile2.write('%7.7f, %7.7f\n' % (pHs[p], Q[p]))
        outfile3.write('%7.7f, %7.7f\n' % (pHs[p], G[p]))
    outfile2.close()
    outfile3.close()

max_stars_repo_name: ktamiola/pepkalc

max_issues_repo_name: ktamiola/pepKalc

max_forks_repo_name: ktamiola/pepKalc


text:
###############
#
# Transform R to Python Copyright (c) 2019 Masahiro Imai Released under the MIT license
#
###############

import os

import numpy as np
import pystan
import pandas
import pickle
import seaborn as sns
import matplotlib.pyplot as plt
from sklearn.preprocessing import LabelEncoder

fish_num_climate_4 = pandas.read_csv('4-3-1-fish-num-4.csv')
print(fish_num_climate_4.head())
print(fish_num_climate_4.describe())

sns.scatterplot(
    x='temperature',
    y='fish_num',
    hue='human',
    data=fish_num_climate_4
)
plt.show()

fish_num_climate_4_d = pandas.get_dummies(fish_num_climate_4, columns=["human"])
print(fish_num_climate_4_d.head())

fish_num = fish_num_climate_4_d['fish_num']
sample_num = len(fish_num)
temperature = fish_num_climate_4_d['temperature']

# creating teamID
le = LabelEncoder()
le = le.fit(fish_num_climate_4['human'])
fish_num_climate_4['human'] = le.transform(fish_num_climate_4['human'])

sns.scatterplot(
    x='temperature',
    y='fish_num',
    hue='human',
    legend="full",
    data=fish_num_climate_4
)
plt.show()

human_id = fish_num_climate_4['human'].values
human_id = human_id + 1
human_num = len(np.unique(human_id))

stan_data = {
    'N': sample_num,
    'fish_num': fish_num,
    'temp': temperature,
    'human_id': human_id,
    'N_human': human_num
}

if os.path.exists('4-3-2-poisson-glmm.pkl'):
    sm = pickle.load(open('4-3-2-poisson-glmm.pkl', 'rb'))
    # sm = pystan.StanModel(file='4-3-1-poisson-glmm.stan')
else:
    # a model using prior for mu and sigma.
    sm = pystan.StanModel(file='4-3-2-poisson-glmm.stan')

control = {
    'adapt_delta': 0.9999,
    'max_treedepth': 16
}

mcmc_result = sm.sampling(
    data=stan_data,
    seed=1,
    chains=4,
    iter=2000,
    warmup=1000,
    control=control,
    thin=1
)

print(mcmc_result)
mcmc_result.plot()
plt.show()

# saving compiled model
if not os.path.exists('4-3-2-poisson-glmm.pkl'):
    with open('4-3-2-poisson-glmm.pkl', 'wb') as f:
        pickle.dump(sm, f)

mcmc_sample = mcmc_result.extract()
print(mcmc_sample)

# visualization
label_temp = np.arange(10,20)
df = pandas.DataFrame(mcmc_sample, columns=['Intercept', 'b_temp'])
print(df.head())

r = mcmc_sample['r']
df_r = pandas.DataFrame(r)
print(df_r.head())

t = mcmc_sample['t']
df_t = pandas.DataFrame(t)
print(df_t.head())

for h in np.arange(human_num):
    y = []
    for i in label_temp:
        y.append(np.exp(df['Intercept'].mean() +
                        (df['b_temp'].mean() + df_t[h].mean()) * i +
                        df_r[h].mean()))

    plt.plot(label_temp, np.array(y), 'red', label='%d' % (h + 1))
    plt.scatter(fish_num_climate_4.query('human == %d' % h)["temperature"],
                fish_num_climate_4.query('human == %d' % h)["fish_num"],
                c='r')
    plt.legend()
    plt.show()

max_stars_repo_name: MasazI/python-r-stan-bayesian-model

max_issues_repo_name: MasazI/python-r-stan-bayesian-model

max_forks_repo_name: MasazI/python-r-stan-bayesian-model


text:
import os as os
import numpy as np
import json
import pickle
import gzip

import matplotlib.pyplot as plt

class Benchmark():
    """API for TabularBench."""

    def __init__(self, data_dir, cache=False, cache_dir="cached/"):
        """Initialize dataset (will take a few seconds-minutes).

        Keyword arguments:
        bench_data -- str, the raw benchmark data directory
        """
        if not os.path.isfile(data_dir) or not data_dir.endswith(".json"):
            raise ValueError("Please specify path to the bench json file.")

        self.data_dir = data_dir
        self.cache_dir = cache_dir
        self.cache = cache

        print("==> Loading data...")
        self.data = self._read_data(data_dir)
        self.dataset_names = list(self.data.keys())
        print("==> Done.")

    def query(self, dataset_name, tag, config_id):
        """Query a run.

        Keyword arguments:
        dataset_name -- str, the name of the dataset in the benchmark
        tag -- str, the tag you want to query
        config_id -- int, an identifier for which run you want to query, if too large will query the last run
        """
        config_id = str(config_id)
        if dataset_name not in self.dataset_names:
            raise ValueError("Dataset name not found.")

        if config_id not in self.data[dataset_name].keys():
            raise ValueError("Config nr %s not found for dataset %s." % (config_id, dataset_name))

        if tag in self.data[dataset_name][config_id]["log"].keys():
            return self.data[dataset_name][config_id]["log"][tag]

        if tag in self.data[dataset_name][config_id]["results"].keys():
            return self.data[dataset_name][config_id]["results"][tag]

        if tag in self.data[dataset_name][config_id]["config"].keys():
            return self.data[dataset_name][config_id]["config"][tag]

        if tag == "config":
            return self.data[dataset_name][config_id]["config"]

        raise ValueError("Tag %s not found for config %s for dataset %s" % (tag, config_id, dataset_name))

    def query_best(self, dataset_name, tag, criterion, position=0):
        """Query the n-th best run. "Best" here means achieving the largest value at any epoch/step,

        Keyword arguments:
        dataset_name -- str, the name of the dataset in the benchmark
        tag -- str, the tag you want to query
        criterion -- str, the tag you want to use for the ranking
        position -- int, an identifier for which position in the ranking you want to query
        """
        performances = []
        for config_id in self.data[dataset_name].keys():
            performances.append((config_id, max(self.query(dataset_name, criterion, config_id))))

        performances.sort(key=lambda x: x[1]*1000, reverse=True)
        desired_position = performances[position][0]

        return self.query(dataset_name, tag, desired_position)

    def get_queriable_tags(self, dataset_name=None, config_id=None):
        """Returns a list of all queriable tags"""
        if dataset_name is None or config_id is None:
            dataset_name = list(self.data.keys())[0]
            config_id = list(self.data[dataset_name].keys())[0]
        else:
            config_id = str(config_id)
        log_tags = list(self.data[dataset_name][config_id]["log"].keys())
        result_tags = list(self.data[dataset_name][config_id]["results"].keys())
        config_tags = list(self.data[dataset_name][config_id]["config"].keys())
        additional = ["config"]
        return log_tags+result_tags+config_tags+additional

    def get_dataset_names(self):
        """Returns a list of all availabe dataset names like defined on openml"""
        return self.dataset_names

    def get_openml_task_ids(self):
        """Returns a list of openml task ids"""
        task_ids = []
        for dataset_name in self.dataset_names:
            task_ids.append(self.query(dataset_name, "OpenML_task_id", 1))
        return task_ids

    def get_number_of_configs(self, dataset_name):
        """Returns the number of configurations for a dataset"""
        if dataset_name not in self.dataset_names:
            raise ValueError("Dataset name not found.")
        return len(self.data[dataset_name].keys())

    def get_config(self, dataset_name, config_id):
        """Returns the configuration of a run specified by dataset name and config id"""
        if dataset_name not in self.dataset_names:
            raise ValueError("Dataset name not found.")
        return self.data[dataset_name][config_id]["config"]

    def plot_by_name(self, dataset_names, x_col, y_col, n_configs=10, show_best=False, xscale='linear', yscale='linear', criterion=None):
        """Plot multiple datasets and multiple runs.

        Keyword arguments:
        dataset_names -- list
        x_col -- str, tag to plot on x-axis
        y_col -- str, tag to plot on y-axis
        n_configs -- int, number of configs to plot for each dataset
        show_best -- bool, weather to show the n_configs best (according to query_best())
        xscale -- str, set xscale, options as in matplotlib: "linear", "log", "symlog", "logit", ...
        yscale -- str, set yscale, options as in matplotlib: "linear", "log", "symlog", "logit", ...
        criterion -- str, tag used as criterion for query_best()
        """
        if isinstance(dataset_names, str):
            dataset_names = [dataset_names]
        if not isinstance(dataset_names, (list, np.ndarray)):
            raise ValueError("Please specify a dataset name or a list list of dataset names.")

        n_rows = len(dataset_names)
        fig, axes = plt.subplots(n_rows, 1, sharex=False, sharey=False, figsize=(10,7*n_rows))

        if criterion is None:
            criterion = y_col

        loop_arg = enumerate(axes.flatten()) if len(dataset_names)>1 else [(0,axes)]

        for ind_ax, ax in loop_arg:
            for ind in range(n_configs):
                try:
                    if ind==0:
                        instances = int(self.query(dataset_names[ind_ax], "instances", 0))
                        classes = int(self.query(dataset_names[ind_ax], "classes", 0))
                        features = int(self.query(dataset_names[ind_ax], "features", 0))

                    if show_best:
                        x = self.query_best(dataset_names[ind_ax], x_col, criterion, ind)
                        y = self.query_best(dataset_names[ind_ax], y_col, criterion, ind)
                    else:
                        x = self.query(dataset_names[ind_ax], x_col, ind+1)
                        y = self.query(dataset_names[ind_ax], y_col, ind+1)

                    ax.plot(x, y, 'p-')
                    ax.set_xscale(xscale)
                    ax.set_yscale(yscale)
                    ax.set(xlabel="step", ylabel=y_col)
                    title_str = ", ".join([dataset_names[ind_ax],
                                          "features: " + str(features),
                                          "classes: " + str(classes),
                                          "instances: " + str(instances)])
                    ax.title.set_text(title_str)
                except ValueError:
                    print("Run %i not found for dataset %s" %(ind, dataset_names[ind_ax]))
                except Exception as e:
                    raise e

    def _cache_data(self, data, cache_file):
        os.makedirs(self.cache_dir, exist_ok=True)
        with gzip.open(cache_file, 'wb') as f:
            pickle.dump(data, f)

    def _read_cached_data(self, cache_file):
        with gzip.open(cache_file, 'rb') as f:
            data = pickle.load(f)
        return data

    def _read_file_string(self, path):
        """Reads a large json string from path. Python file handler has issues with large files so it has to be chunked."""
        # Shoutout to https://stackoverflow.com/questions/48122798/oserror-errno-22-invalid-argument-when-reading-a-huge-file
        file_str = ''
        with open(path, 'r') as f:
            while True:
                block = f.read(64 * (1 << 20)) # Read 64 MB at a time
                if not block:                  # Reached EOF
                    break
                file_str += block
        return file_str

    def _read_data(self, path):
        """Reads cached data if available. If not, reads json and caches the data as .pkl.gz"""
        cache_file = os.path.join(self.cache_dir, os.path.basename(self.data_dir).replace(".json", ".pkl.gz"))
        if os.path.exists(cache_file) and self.cache:
            print("==> Found cached data, loading...")
            data = self._read_cached_data(cache_file)
        else:
            print("==> No cached data found or cache set to False.")
            print("==> Reading json data...")
            data = json.loads(self._read_file_string(path))
            if self.cache:
                print("==> Caching data...")
                self._cache_data(data, cache_file)
        return data


text:
from pyrep.errors import IKError
import numpy as np

class cartesian_position_controller():
    def __init__(self, panda):
        self.panda = panda
        self.target_p = np.zeros(3)
        self.target_r = np.zeros(4)

    def set_target(self, target):
        self.target_p = target.get_position()
        target.set_orientation([0,3.14,0]) # Set orientation to be upright
        self.target_r = target.get_quaternion()

    def compute_action(self, gain=0.03):
        current_p = self.panda.get_tip().get_position()
        e_pos = (current_p - self.target_p) * gain

        current_r = self.panda.get_tip().get_quaternion()
        e_rot = (current_r - self.target_r) * gain

        try:
            q_vel = np.array(self.panda.solve_ik_via_jacobian(e_pos, quaternion=e_rot))
        except IKError:
            # So let's swap to an alternative IK method...
            # This returns 'max_configs' number of joint positions
            print("Trying sampling...")
            q_vel = np.array(self.panda.solve_ik_via_sampling(e_pos, quaternion=e_rot)[0])

        v = np.append(q_vel, 1.0)
        return v

    def recompute_action(self, current_q, target_q, gain=0.03):
        err = current_q - target_q
        v = -err * gain
        v = np.append(v, 1.0)
        return v