text:
# -*- coding: cp936 -*-
"""
The main purpose of this script is to extract Process Tail from single end Read.

=============================
Usage: python SAM2SE_Tail_length_obj_edited8_v1.2.py
-h help

-i files to processed                           *[No default value]

-r refrence dictionary file(.dict)              *[No default value]

-P primer (fasta format)                        *[No default value]

-R maximum Read length                          *[No default value]
   should be read length-1

-p minimum mapped percent of the read            [default: 0.00]

-m Minimum size range based on Bioanalyzer      [default: 0]

-M maximum size range based on Bioanalyzer      [default: 10000]

-c cigar string count                           [default:2]
    this means it allows cigar string to have equal or less that 2 fragments

-o prefix of output files                       [default:"precursor_length"]

-s suffix for the files to be processed         [default value "sam"]

-S output SUFFIX                                [default: "txt"]

-d sep_char within among columns                [default value '\t']

-j skip header lines                            [default value 1]
    if skip header lines = 0 indicates that there is no header line
    if skip header lines = n indicates that first n lines are header lines

-I input file path              [default: current folder]

-O output file path             [default: current folder]

-L unique_id length for the infile      [default value 2]

===================
input description:
input files:
1. SAM file

======================
output files:
1. overall precursor length
2. each transcript length , count

============================

Python & Module requirement:
Versions 2.x : 2.4 or above
Module: No additional Python Module is required.

============================
Library file requirement:

Not Standalone version, few library file is required.

============================

External Tools requirement:

============================
command line example:

============================
versions update
1.3:
fix incorrect handle of reverse strand
1.2:
fix the H cigar string incorrect Tail retrieve and tail assignment
1.1:
fix over-estimate problem on ATail and underestimate Utail
due to incorrectly use round and int function

"""

##Copyright
##By Liye Zhang
##Contact: bioliyezhang@gmail.com
##Compatible Python Version:2.4 or above

###Code Framework

### Specific Functions definiation
def detect_longest_shortU_tail(seq):

    U_count=0
    U_percentage_list=[]
    for index in range(len(seq)):
        nt=seq[index]
        if nt=="T":
            U_count+=1
        U_percentage = 100.00 * U_count / (index+1)
        if U_percentage<50.00:
            break
        else:
            U_percentage_list.append(U_percentage)
    ##
    if len(U_percentage_list)==0:
        return 0
    else:
        max_U_percent = max(U_percentage_list)
        for index in range(len(U_percentage_list)):
            U_percentage=U_percentage_list[index]
            if U_percentage==max_U_percent:
                final_index=index+1
        return final_index

def tail_type_call(tail_seq):
    #tail_type="unknown"

    ## pattern definition
    #pattern_dict = dict()

    tail_length = len(tail_seq)
    U_count = tail_seq.count("T")
    if tail_seq.startswith("TT") and U_count>=round(0.90*tail_length):
        tail_type="Utail"
    elif tail_seq=="T":
        tail_type="Utail"
    else:
        ## (optional U + polyA) tail condition
        shortU_begin_len= detect_longest_shortU_tail(tail_seq)
        trim_Utail_tail_seq= tail_seq[shortU_begin_len:]
        trim_tail_len = len(trim_Utail_tail_seq)
        trim_A_count = trim_Utail_tail_seq.count("A")
        if trim_A_count>=round(0.90*trim_tail_len):
            tail_type="Atail"
        else:
            trim_U_count = trim_Utail_tail_seq.count("T")
            trim_AU_count = trim_A_count+trim_U_count
            if trim_AU_count>=round(0.80*trim_tail_len):
                tail_type="AUtail"
            else:
                tail_type="unknown"

    return tail_type

def count_SH_in_cigar(cigar_data_list):
    SH_count = 0
    for cigar_data in cigar_data_list:
        if cigar_data[-1] in ["S","H"]:
            SH_count+=1
    return SH_count

def QC_filter(seq,qc_score,MIN_MEAN_QC=20,MAX_LOW_QC_COUNT=3,LOW_QC=10,MAX_N=0):
    result = True
    seq_length = len(seq)

    #module I: max N
    N_count = seq.count("N")+seq.count("n")
    if N_count>MAX_N:
        result=False

    #module II: Mean QC score
    qc_score_total=0
    qc_score_low_count=0
    for qc in qc_score:
        qc_score_number = ord(qc)-32
        qc_score_total+=qc_score_number
        if qc_score_number<=LOW_QC:
            qc_score_low_count+=1
    qc_score_mean = qc_score_total/seq_length

    if qc_score_mean<MIN_MEAN_QC:
        result=False
    if qc_score_low_count>=MAX_LOW_QC_COUNT:
        result=False
    return result

def specific_function(infiles,reference,primer_infile):
    ## infile: are SAM file
    ## reference: dict file generated for fasta file
    ## primer_infile: fasta for primer

    from numpy import median
    global pattern_list
    cmd_records=record_command_line()  ##record the command line

    ## Part 1: Process the primer infile
    primer_dict = read_fasta(primer_infile)
    primer_list = primer_dict.values()

    ## Part 2: Process the dict file for reference fasta
    infile_obj=GeneralFile_class(reference)  ##create file obj(class)
    infile_obj.SKIP_HEADER=1    ##setup up the manual skip header if necessary
    infile_obj.SAMPLE_ID_LEN=unique_id_length  ##unique ID length
    infile_reader=infile_obj.reader_gen()  ##create the file reader to process infile
    reference_length_dict= dict()
    for row in infile_reader:
        ID=row[1][3:]
        length=int(row[2][3:])
        reference_length_dict[ID]=length
    reference_list = reference_length_dict.keys()

    ## Part 3: Process SAM file
    for infile in infiles:

        print "Processing infile:", infile
        final_type_list= ["unknown","Utail","Atail","AUtail","notail",""]
        final_type_count_dict = dict()
        for tail_type in final_type_list:
            final_type_count_dict[tail_type]=0 ## number

        ID_list = []
        tail_type_ID_dict = dict()
        ID_type_dict = dict() # Per ID record the Tail type
        for pattern in pattern_list:
            tail_type_ID_dict[pattern] = [] ##? not sure about this

        tail_data_dict = dict()
        notail_f_dict = dict()
        notail_r_dict = dict()
        for reference_ID in reference_list:
            length=reference_length_dict[reference_ID]
            notail_r_dict[reference_ID]=[0]*length ## record the no tail readcount
            notail_f_dict[reference_ID]=[0]*length

        ##Set up infile object
        infile_obj=SAM_File_class(infile)  ##create file obj(class)
        infile_obj.SKIP_HEADER=infile_skip    ##setup up the manual skip header if necessary
        infile_obj.SAMPLE_ID_LEN=unique_id_length  ##unique ID length
        infile_reader=infile_obj.reader_gen()

        ## round 1 Process (initilize the Dictionary)
        ## dictiontionary for tail length
        tail_length_dict =dict()
        full_seq_dict = dict() ##
        for row in infile_reader:
            ID = row[infile_obj.QNAME_COLUMN]
            tail_length_dict[ID] = [1000,1000] # two value, one for maxcircle, one for edited
            ID_type_dict[ID] = "" ## make sure no empty value
            full_seq_dict[ID] = ""

        ## Round 2 record the different data (design to handle RPS12)
        infile_obj=SAM_File_class(infile)  ##create file obj(class)
        infile_obj.SKIP_HEADER=infile_skip    ##setup up the manual skip header if necessary
        infile_obj.SAMPLE_ID_LEN=unique_id_length  ##unique ID length
        infile_reader=infile_obj.reader_gen()
        for row in infile_reader:

            ID = row[infile_obj.QNAME_COLUMN]
            ref_ID = row[infile_obj.CHRO_COLUMN]

            ## design to capture which side
            if ref_ID=="max_circle":
                tail_index=0
            elif ref_ID.startswith("Tb"):
                tail_index=1
            else:
                tail_index=-1 ## non-maxcircle and rps12

            ## will make this a common function to handle flag
            flag = int(row[infile_obj.FLG_COLUMN])
            flag_binary = bin(flag)
            flag_binary_len = len(flag_binary)-2
            final_binary = "0"*(12-flag_binary_len)+flag_binary[2:]

            ## design on the potential tail length

            cigar_data = row[infile_obj.CIGAR_COLUMN]

            if cigar_data!="*":

                map_start_coor = int(row[infile_obj.COOR_COLUMN])
                cigar_extracted_list, read_length, match_length = read_cigar_coor_offset_withlist(cigar_data)
                min_length = len(row[infile_obj.SEQ_COLUMN])*MATCH_PERCENT_LIMIT/100

                ## only look into the Cigar that is not terribly complicated, current set to 10
                if len(cigar_extracted_list)<=CIGAR_UPPER_LIMIT and match_length>=min_length:

                    full_seq = row[infile_obj.SEQ_COLUMN]
                    if final_binary[7]=="0":  ## 0 means forward, no reverse complement
                        direction="+"
                        if len(full_seq)>len(full_seq_dict[ID]):
                            full_seq_dict[ID]=full_seq
                        fragment_start = map_start_coor
                        fragment_end = map_start_coor + match_length -1

                        first_cigar_string = cigar_extracted_list[-1]
                        if first_cigar_string[-1] in ["S","H"]:
                            tail_length = int(first_cigar_string[:-1])
                            if tail_index>=0:
                                if tail_length<tail_length_dict[ID][tail_index]:
                                    tail_length_dict[ID][tail_index]=tail_length

                    else:
                        if len(full_seq)>len(full_seq_dict[ID]):
                            full_seq_dict[ID]=reverse_complementary(full_seq)
                        direction="-"
                        fragment_end = map_start_coor
                        fragment_start = map_start_coor+match_length-1
                        first_cigar_string = cigar_extracted_list[0]

                        if first_cigar_string[-1] in ["S","H"]:
                            if tail_index>=0:
                                tail_length = int(first_cigar_string[:-1])
                                if tail_length<tail_length_dict[ID][tail_index]:
                                    tail_length_dict[ID][tail_index]=tail_length

        ## round 2 Process
        ##Set up infile object
        infile_obj=SAM_File_class(infile)  ##create file obj(class)
        infile_obj.SKIP_HEADER=infile_skip    ##setup up the manual skip header if necessary
        infile_obj.SAMPLE_ID_LEN=unique_id_length  ##unique ID length
        infile_reader=infile_obj.reader_gen()

        for row in infile_reader:
            ID = row[infile_obj.QNAME_COLUMN]
            ID_type_dict[ID] = "" ## what does this actually mean in here?
            ref_ID = row[infile_obj.CHRO_COLUMN]

            ## will make this a common function to handle flag
            flag = int(row[infile_obj.FLG_COLUMN])
            flag_binary = bin(flag)
            flag_binary_len = len(flag_binary)-2
            final_binary = "0"*(12-flag_binary_len)+flag_binary[2:]

            map_start_coor = int(row[infile_obj.COOR_COLUMN])
            cigar_data = row[infile_obj.CIGAR_COLUMN]
            if cigar_data!="*":

                cigar_extracted_list, read_length, match_length = read_cigar_coor_offset_withlist(cigar_data)
                min_length = len(row[infile_obj.SEQ_COLUMN])*MATCH_PERCENT_LIMIT/100

                ## only look into the reads that are shorter than the total fragment
                if len(cigar_extracted_list)<=CIGAR_UPPER_LIMIT and match_length>=min_length:
                    start_with_designed_primer = False

                    if len(ID_list)!=0:
                        if ID==ID_list[-1]:
                            pass
                        else:
                            ID_list.append(ID)
                    else:
                        ID_list.append(ID)

                    full_seq = full_seq_dict[ID]
                    #full_QC = row[infile_obj.QUAL_COLUMN]

                    tail_unmapped=False
                    if final_binary[7]=="0":
                        ## forward strand
                        for primer_seq in primer_list:
                            if full_seq.startswith(primer_seq):
                                start_with_designed_primer=True
                                break

                        direction="+"
                        fragment_start = map_start_coor
                        fragment_end = map_start_coor + match_length -1

                        first_cigar_string = cigar_extracted_list[-1]
                        if first_cigar_string[-1] in ["S","H"]:
                            tail_length = int(first_cigar_string[:-1])
                            tail_unmapped=True
                            sense_tail_seq = full_seq[tail_length*(-1):]
                            #tail_QC = full_QC[tail_length*(-1):]
                            #QC_pass = QC_filter(sense_tail_seq,tail_QC)

                    else:

                        direction="-"
                        fragment_end = map_start_coor
                        fragment_start = map_start_coor+match_length-1

                        first_cigar_string = cigar_extracted_list[0]
                        for primer_seq in primer_list:
                            if full_seq.startswith(primer_seq):
                                start_with_designed_primer=True
                                break

                        if first_cigar_string[-1] in ["S","H"]:
                            tail_length = int(first_cigar_string[:-1])
                            tail_unmapped=True
                            sense_tail_seq = full_seq[tail_length*(-1):]

                    ## check the S,H count in cigar string
                    SH_fragment_count = count_SH_in_cigar(cigar_extracted_list)
                    #fragment_start = map_start_coor
                    if start_with_designed_primer==True:
                        if tail_unmapped==True and SH_fragment_count<3:
                            ##
                            min_tail_length = min(tail_length_dict[ID])
                            if min_tail_length>0:
                                if len(sense_tail_seq)==min_tail_length:
                                    tail_type= tail_type_call(sense_tail_seq)
                                    ID_type_dict[ID] = tail_type
                                    tail_type_ID_dict[tail_type].append(ID)
                                    U_count = sense_tail_seq.count("T")
                                    A_count = sense_tail_seq.count("A")
                                    final_type_count_dict[tail_type]+=1
                                    tail_data_dict[ID] = [full_seq,sense_tail_seq,U_count,A_count,ref_ID,fragment_start,fragment_end,direction,tail_type]
                            else:
                                if read_length<=MAX_READLENGTH and first_cigar_string[-1]=="M":
                                    final_type_count_dict["notail"]+=1
                                    ID_type_dict[ID] = "notail"
                                    if direction=="+":
                                        notail_f_dict[ref_ID][fragment_end-1]+=1
                                    elif direction=="-":
                                        notail_r_dict[ref_ID][fragment_end-1]+=1
                                    else:
                                        print "direction unrecognized", direction
                                        sys.exit(0)
                        else:
                            ## with no tail read, require the READ length is smaller than the 300nt (max length)

                            if read_length<=MAX_READLENGTH and first_cigar_string[-1]=="M":
                                final_type_count_dict["notail"]+=1
                                ID_type_dict[ID] = "notail"
                                if direction=="+":
                                    notail_f_dict[ref_ID][fragment_end-1]+=1
                                elif direction=="-":
                                    notail_r_dict[ref_ID][fragment_end-1]+=1
                                else:
                                    print "direction unrecognized", direction
                                    sys.exit(0)
                    else:
                        pass

        outfile_name=infile_obj.outputfilename_gen(prefix,OUTPUT_SUFFIX) ##create output file
        outfile_path=OUTPUT_PATH+"/"+outfile_name
        outfile_obj=GeneralFile_class(outfile_path)              ##create output obj
        outfile_obj.RECORD=cmd_records
        outfile_obj.output_handle_gen()    ##generate output handle
        outfile_obj.handle.write("#IDs\n")
        ID_set = set(ID_list)
        for seq_ID in ID_set:
            outfile_obj.handle.write(seq_ID+'\n')
        outfile_obj.handle.close()

        outfile_name=infile_obj.outputfilename_gen("Tail",OUTPUT_SUFFIX) ##create output file
        outfile_path=OUTPUT_PATH+"/"+outfile_name
        outfile_obj=GeneralFile_class(outfile_path)              ##create output obj
        outfile_obj.RECORD=cmd_records
        outfile_obj.output_handle_gen()    ##generate output handle
        outfile_obj.handle.write("#fastq_ID\tfull_Rseq\ttail_seq\tTail_length\tUcount\t")
        outfile_obj.handle.write("Acount\tChro\tFragment_start\tFragment_end\tdirection\t")
        outfile_obj.handle.write("Type\n")

        for seq_ID in tail_data_dict.keys():
            #tail_count+=1
            outfile_obj.handle.write(seq_ID+'\t')
            data_list = tail_data_dict[seq_ID]
            outfile_obj.handle.write(data_list[0]+'\t')
            outfile_obj.handle.write(data_list[1]+'\t')
            outfile_obj.handle.write(str(len(data_list[1]))+'\t')
            outfile_obj.handle.write(str(data_list[2])+'\t') #U count
            outfile_obj.handle.write(str(data_list[3])+'\t') #A count
            outfile_obj.handle.write(str(data_list[4])+'\t') #Chromosome
            outfile_obj.handle.write(str(data_list[5])+'\t') #fragment start
            outfile_obj.handle.write(str(data_list[6])+'\t') #fragment end
            outfile_obj.handle.write(str(data_list[7])+'\t') #direction of transcript
            outfile_obj.handle.write(str(data_list[8])+'\n') #tail type
        outfile_obj.handle.close()

        outfile_name=infile_obj.outputfilename_gen("Notail_Stat_BED6","bed") ##create output file
        outfile_path=OUTPUT_PATH+"/"+outfile_name
        outfile_obj=GeneralFile_class(outfile_path)              ##create output obj
        outfile_obj.RECORD=cmd_records
        outfile_obj.output_handle_gen()    ##generate output handle
        outfile_obj.handle.write("#ref_ID\tStart\tEND\tName\tReadCount\tStrand\n")
        ## output forward strand first
        for reference_ID in reference_list:
            length=reference_length_dict[reference_ID]
            for index in range(length):
                readcount=notail_f_dict[reference_ID][index]
                ID=reference_ID+"_"+str(index+1)+"_"+"p"
                outfile_obj.handle.write(reference_ID+'\t'+str(index)+'\t')
                outfile_obj.handle.write(str(index+1)+'\t'+ID+"\t"+str(readcount))
                outfile_obj.handle.write('\t+\n')
        ## output reverse strand first
        for reference_ID in reference_list:
            length=reference_length_dict[reference_ID]
            for index in range(length):
                readcount=notail_r_dict[reference_ID][index]
                ID=reference_ID+"_"+str(index+1)+"_"+"m"
                outfile_obj.handle.write(reference_ID+'\t'+str(index)+'\t')
                outfile_obj.handle.write(str(index+1)+'\t'+ID+"\t"+str(readcount))
                outfile_obj.handle.write('\t-\n')
        outfile_obj.handle.close()

        ##  generate the separate SAM file

        outfile_obj_dict = dict()
        for pattern in final_type_list:
            if pattern!="":
                outfile_name = infile_obj.outputfilename_gen(pattern,"sam")
            else:
                outfile_name = infile_obj.outputfilename_gen("unprocessed","sam")

            outfile_path = OUTPUT_PATH+"/"+outfile_name
            outfile_obj_dict[pattern] = GeneralFile_class(outfile_path)
            outfile_obj_dict[pattern].output_handle_gen()

        infile_obj=SAM_File_class(infile)  ##create file obj(class)
        infile_obj.SKIP_HEADER=infile_skip    ##setup up the manual skip header if necessary
        infile_obj.SAMPLE_ID_LEN=unique_id_length  ##unique ID length
        infile_obj.AUTOSKIP_HEADER = False
        infile_reader=infile_obj.reader_gen()
        for row in infile_reader:
            if row[0][0]=="@":
                for pattern in final_type_list:
                    output_row(outfile_obj_dict[pattern].handle,row)
            else:
                ID = row[infile_obj.QNAME_COLUMN]
                ID_type = ID_type_dict[ID]
                output_row(outfile_obj_dict[ID_type].handle,row)

        for pattern in final_type_list:
            outfile_obj_dict[pattern].handle.close()

        ##
        outfile_name=infile_obj.outputfilename_gen("stat",OUTPUT_SUFFIX) ##create output file
        outfile_path=OUTPUT_PATH+"/"+outfile_name
        outfile_obj=GeneralFile_class(outfile_path)              ##create output obj
        outfile_obj.RECORD=cmd_records
        outfile_obj.output_handle_gen()    ##generate output handle
        outfile_obj.handle.write("#type\tCount\tPercentage\n")
        total_count=0
        for tail_type in final_type_list[:-1]:
            total_count+=final_type_count_dict[tail_type]

        for tail_type in final_type_list[:-1]:
            tail_count = final_type_count_dict[tail_type]
            tail_percent = 100.00 * tail_count/total_count
            outfile_obj.handle.write(tail_type+'\t'+str(tail_count)+'\t')
            outfile_obj.handle.write(str(tail_percent)+'\n')
        outfile_obj.handle.write("all\t"+str(total_count)+'\t100.00\n')
        outfile_obj.handle.close()

if __name__ == "__main__":
    ###Python General Module Import
    import sys, csv, getopt, re
    import os
    import math
    from itertools import ifilter

    ##Liye own common function,class loading
    from Constant_Library import *
    from General_Library import *
    from File_Class import *  ###
    from Sequencing_Library import *

    OUTPUT_SEP_CHAR='\t'

    #exit if not enough arguments
    if len(sys.argv) < 3:
        print __doc__
        sys.exit(0)

    ###set default value
    suffix="sam"
    infile=None
    infile_skip=0
    sep_char='\t'
    sep_gene=','
    header_file=None
    unique_id_length=2
    primer_infile=None
    INPUT_PATH=os.getcwd()
    OUTPUT_PATH=os.getcwd()
    prefix="proper_PE_IDs"
    OUTPUT_SUFFIX="txt"
    CIGAR_UPPER_LIMIT=4
    MATCH_PERCENT_LIMIT=0.00
    MIN_SIZE = 0
    MAX_SIZE = 1000
    MAX_READLENGTH = 299
    reference = None

    global pattern_list
    global original_list
    # original_list = ["polyAG_30_pattern","polyAU_30_pattern","polyA_30_pattern",
    # "oligoAG_5_29_pattern","oligoAU_5_29_pattern","oligoA_5_29_pattern",
    # "polyU_5_200_pattern","shortA_2_4_pattern","shortU_2_4_pattern","polyUA_pattern"]
    # pattern_list = []
    # for pattern in original_list:
    #     pattern_list.append(pattern)
    #     ligation_type = pattern + "-ligation"
    #     pattern_list.append(ligation_type)
    pattern_list=["unknown","Utail","Atail","AUtail","notail",""]

    ###get arguments(parameters)
    optlist, cmd_list = getopt.getopt(sys.argv[1:], 'hi:s:S:r:c:m:M:d:P:R:D:j:I:t:p:L:o:O:z',["test="])
    for opt in optlist:
        if opt[0] == '-h':
            print __doc__; sys.exit(0)
        elif opt[0] == '-i': infile = opt[1]
        elif opt[0] == '-I': INPUT_PATH = opt[1]
        elif opt[0] == '-O': OUTPUT_PATH = opt[1]
        elif opt[0] == '-S': OUTPUT_SUFFIX = opt[1]
        elif opt[0] == '-s': suffix = opt[1]
        elif opt[0] == '-d': sep_char =opt[1]
        elif opt[0] == '-r': reference =opt[1]
        elif opt[0] == '-R': MAX_READLENGTH = int(opt[1])
        elif opt[0] == '-P': primer_infile = opt[1]
        elif opt[0] == '-m': MIN_SIZE = int(opt[1])
        elif opt[0] == '-M': MAX_SIZE = int(opt[1])
        elif opt[0] == '-p': MATCH_PERCENT_LIMIT = float(opt[1])
        elif opt[0] == '-c': CIGAR_UPPER_LIMIT = int(opt[1])
        elif opt[0] == '-j': infile_skip= int(opt[1])
        elif opt[0] == '-o': prefix = opt[1]
        elif opt[0] == '-L': unique_id_length = int(opt[1])
        elif opt[0] == '--test': long_input = opt[1]

    #print "Test long input", long_input
    if infile==None:
        infiles=CurrentFolder_to_Infiles(INPUT_PATH, suffix)
    else:
        infiles=[infile]

    specific_function(infiles,reference,primer_infile)

reason: from numpy


text:
import os
import sys
import argparse
import pickle
import json

import gym
import numpy as np

from core.evaluator import ALEvaluator
from core.basis import GaussianKernel, RBFKernel2
from core.feature_expectation import MCMuEstimator, LSTDMuEstimator, DeepMuEstimator
from core.policy import RandomPolicy2
from core.data_util import GymDataset

from contrib.scirl import train_scirl, train_scirl_v2, train_scirl_v3
from contrib.baselines import bench, logger
from contrib.baselines.common import set_global_seeds
from contrib.baselines.common.misc_util import boolean_flag
from contrib.baselines.gail import mlp_policy
from contrib.baselines.common import set_global_seeds, tf_util as U

from model.mma import train_mma
from model.bc import train_bc
from model.lstd import LSTDMu, LSPI
from model.dqn import DQNSepsis

file_path = os.path.dirname(os.path.realpath(__file__))
root_path = os.path.abspath(os.path.join(file_path, os.pardir))
sys.path.insert(0, root_path)

import logging
from logging import config

with open(os.path.join(root_path, "context.json"), "r") as f:
    context = json.load(f)

ENV_LIST = context["env"]
MODEL_LIST = context["model"]

def argsparser():
    parser = argparse.ArgumentParser()

    parser.add_argument("task", type=str, choices=ENV_LIST)
    parser.add_argument("--model_id", type=str, choices=MODEL_LIST)
    parser.add_argument("--n_e", type=int, help="demo expert sample size")
    parser.add_argument("--n_ne", type=int, help="demo no expert sample size")
    parser.add_argument("--save_path", type=str, help="save to txt file")
    parser.add_argument("--seed", type=int, default=0, help="seed to keep experiment results in sync")
    parser.add_argument("--overwrite", action="store_true",
                        help="overwrite best param if available")

    parser.add_argument('--env_id', help='environment ID', default='MountainCar-v0')
    parser.add_argument('--expert_path', type=str, default="demo.deterministic.expert.npz")
    parser.add_argument('--task_mode', type=str, default="train", choices=["train", "evaluate"])
    parser.add_argument('--checkpoint_dir', help='the directory to save model', default='task')
    parser.add_argument('--horizon', help='Finite Horizon for a traj', type=int,
            default=200)
    parser.add_argument('--log_dir', help='the directory to save log file', default='log')
    parser.add_argument('--traj_limitation', type=int, default=1)
    # Network Configuration (Using MLP Policy)
    parser.add_argument('--policy_hidden_size', type=int, default=128)
    parser.add_argument('--dim_phi', type=int, default=64)
    # for evaluatation
    parser.add_argument('--BC_max_iter', help='Max iteration for training BC',
            type=int, default=3e4)

    boolean_flag(parser, 'stochastic_policy', default=False,
            help='use stochastic/deterministic policy to evaluate')
    boolean_flag(parser, 'save_sample', default=False,
            help='save the trajectories or not')
    boolean_flag(parser, 'pretrain', default=True,
            help='pretrain a poilcy and phi with imitation learning')

    return parser.parse_args()

def main(args):
    set_global_seeds(args.seed)

    task_path = os.path.join(root_path, "task", args.task)
    data_path = os.path.join(task_path, "data")

    task_desc_path = os.path.join(task_path, "task.json")

    with open(task_desc_path, "r") as f:
        task_desc = json.load(f)

    e_path_async = os.path.join(data_path, args.expert_path)
    D_e_bc = np.load(e_path_async)

    data = {}
    data["irl"] = D_e_bc

    model_path = os.path.join(task_path, "model", "{}.json".format(args.model_id))
    with open(model_path, "r") as f:
        params = json.load(f)

    if args.task_mode == "train":
        train(data, task_desc, params, args, task_path)

    elif args.task_mode == "evaluate":
        evaluate(data, task_desc, params, args, task_path)

def make_gym_env(env_id, seed):
    import gym
    env = gym.make(env_id)
    env = bench.Monitor(env, logger.get_dir() and
                        os.path.join(logger.get_dir(), "monitor.json"),
                        allow_early_resets=True)

    env.seed(seed)
    gym.logger.setLevel(logging.WARN)
    return env

def train(data, task_desc, params, args, task_path):
    if args.pretrain:
        model_path = os.path.join(root_path, "task", args.task, "model")
        fname = "ckpt.bc.{}.{}".format(args.traj_limitation, args.seed)
        ckpt_dir = os.path.join(model_path, fname)
        pretrained_path = os.path.join(ckpt_dir, fname)

        if not os.path.exists(os.path.join(ckpt_dir, "checkpoint")):
            env = make_gym_env(task_desc["env_id"], args.seed)
            ob_space = env.observation_space
            ac_space = env.action_space

            pretrained_path = train_bc(task_desc,
                                     params,
                                     ob_space,
                                     ac_space,
                                     args,
                                     env)
            env.close()

        env = make_gym_env(task_desc["env_id"], args.seed)

        ob_space = env.observation_space
        ac_space = env.action_space
        U.make_session(num_cpu=1).__enter__()
        set_global_seeds(args.seed)

        def mlp_pi_wrapper(name, ob_space, ac_space, reuse=False):
            return mlp_policy.MlpPolicy(name=name,
                                        ob_space=ob_space,
                                        ac_space=ac_space,
                                        reuse=reuse,
                                        hid_size_phi=args.policy_hidden_size,
                                        num_hid_layers_phi=2,
                                        dim_phi=args.dim_phi)

        env_name = task_desc["env_id"]
        scope_name = "pi.{}.{}".format(env_name.lower().split("-")[0], args.traj_limitation)

        # pi_0
        # phi_0
        pi_bc = mlp_pi_wrapper(scope_name, ob_space, ac_space)
        U.initialize()
        U.load_state(pretrained_path)
        phi_bc = pi_bc.featurize

        def phi_old(s, a):
            """
            TODO: if action is discrete
            one hot encode action and concatenate with phi(s)
            """
            # expect phi(s) -> (N, state_dim)
            # expect a -> (N, action_dim)
            phi_s = phi_bc(s)

            if len(phi_s.shape) == 1:
                # s -> (1, state_dim)
                phi_s = np.expand_dims(phi_s, axis=0)

            # if a = 5
            try:
                if a == int(a):
                    a = [a]
            except:
                pass

            a = np.array(a)
            # if a = [5]
            if len(a.shape) == 1:
                a = np.expand_dims(a, axis=1)
            # otherwise if a = [[5], [3]]
            phi_sa = np.hstack((phi_s, a))
            return phi_sa

        def phi_discrete_action(n_action):
            def f(s, a):
                # expect phi(s) -> (N, state_dim)
                # expect a -> (N, action_dim)
                phi_s = phi_bc(s)

                try:
                    if a == int(a):
                        a = [a]
                except:
                    pass
                a = np.array(a)

                a_onehot = np.eye(n_action)[a.astype(int)]

                if len(phi_s.shape) == 1:
                    # s -> (1, state_dim)
                    phi_s = np.expand_dims(phi_s, axis=0)

                try:
                    phi_sa = np.hstack((phi_s, a_onehot))
                except:
                    a_onehot = a_onehot.reshape(a_onehot.shape[0],
                            a_onehot.shape[2])
                    phi_sa = np.hstack((phi_s, a_onehot))

                return phi_sa
            return f

        if isinstance(ac_space, gym.spaces.Discrete):
            phi = phi_discrete_action(ac_space.n)
        elif isinstance(ac_space, gym.spaces.Box):
            phi = phi_continuous_action
        else:
            raise NotImplementedError

        env.close()

    D = data["irl"]
    obs = D["ob_list"].reshape(-1, D["ob_list"].shape[-1])
    obs_p1 = D["ob_next_list"].reshape(-1, D["ob_next_list"].shape[-1])
    #assuming action dof of 1
    acs = D["ac_list"].reshape(-1)

    if "new" in D:
        new = D["new"].reshape(-1)
    else:
        # no terminal
        new = np.zeros(acs.shape)

    data = {}
    data["s"] = obs
    data["a"] = acs
    data["s_next"] = obs_p1
    data["done"] = data["absorb"] = new

    data["phi_sa"] = phi(obs, acs)
    data["phi_fn"] = phi
    data["phi_fn_s"] = phi_bc

    data["psi_sa"] = data["phi_sa"]
    data["psi_fn"] = phi

    env = make_gym_env(task_desc["env_id"], args.seed)

    evaluator = ALEvaluator(data, task_desc["gamma"], env)
    data_path = os.path.join(task_path, "data")
    expert_log_path = os.path.join(data_path, "expert.log")

    with open(expert_log_path, 'r') as f:
        log = f.read()
        evaluator.perf_exp = float(log.split()[1])
    env.close()

    pi_0 = pi_bc

    phi_dim = data["phi_sa"].shape[1]

    model_id = "{}.{}".format(params["id"], params["version"])

    if model_id == "mma.0":
        result = train_mma(pi_0, phi_dim, task_desc, params, data, evaluator)
    elif model_id == "mma.1":
        result = train_mma(pi_0, phi_dim, task_desc, params, data, evaluator)
    elif model_id == "mma.2":
        #result = train_scirl_v2(data, phi_bc, evaluator, phi_dim, task_desc, params)
        #result = train_scirl_v3(data, phi_bc, evaluator)
        result = train_scirl(data, phi_bc, evaluator)
    else:
        raise NotImplementedError

    # also saves checkpoints for eval
    name = "{}.{}.{}".format(model_id, args.n_e, args.seed)
    result_path = os.path.join(args.save_path, name + "train.log")

    with open(result_path, "w") as f:
        for step in range(params["n_iteration"] + 1):
            data_points = [step,
                           #round(result["perf"][step],2),
                           #round(result["margin_mu"][step],2),
                           #round(result["margin_v"][step],2),
                           round(result["a_match"][step],2)
                           ]
            f.write("{}\t{}\t{}\t{}\n".format(*data_points))

    with open(os.path.join(args.save_path, name + ".pkl") , "wb") as f:
        pickle.dump(result, f, pickle.HIGHEST_PROTOCOL)

if __name__ == "__main__":
    args = argsparser()
    main(args)

ext: py

max_stars_repo_name: dtak/batch-apprenticeship-learning

max_issues_repo_name: dtak/batch-apprenticeship-learning

max_forks_repo_name: dtak/batch-apprenticeship-learning

reason: import numpy


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

def calculate_pi(n):

	x_and_y = np.random.rand(n,2)
	print("This is the array of x and y values \n",x_and_y)

	x=x_and_y[:,0]
	y=x_and_y[:,1]
	#np.x_and_y.shape

	r= (x_and_y[:,0]**2 + x_and_y[:,1]**2)**0.5

	print("This is the array of r values \n", r)

	r_gt =sum(float(gt) >= 1 for gt in r)

	r_lt=sum(float(lt) <= 1 for lt in r)
	print(r_gt,r_lt)

	pi=4*(r_lt/(r_gt+r_lt))
	print(pi)

	fig=plt.figure()
	plt.scatter(x[r<1],y[r<1],color='red')
	plt.scatter(x[r>=1],y[r>=1],color='blue')

	plt.show()


text:
# Copyright (C) 2021 Intel Corporation
# SPDX-License-Identifier: Apache-2.0

import pytest
import numpy as np
import os
from sys import platform
from pathlib import Path

import openvino.opset8 as ov
from openvino import Core, IENetwork, ExecutableNetwork, blob_from_file
from openvino.impl import Function, Shape, Type
from openvino.impl.op import Parameter
from openvino import TensorDesc, Blob

from ..conftest import model_path, model_onnx_path, plugins_path

test_net_xml, test_net_bin = model_path()
test_net_onnx = model_onnx_path()
plugins_xml, plugins_win_xml, plugins_osx_xml = plugins_path()

def test_blobs():
    input_shape = [1, 3, 4, 4]
    input_data_float32 = (np.random.rand(*input_shape) - 0.5).astype(np.float32)

    td = TensorDesc("FP32", input_shape, "NCHW")

    input_blob_float32 = Blob(td, input_data_float32)

    assert np.all(np.equal(input_blob_float32.buffer, input_data_float32))

    input_data_int16 = (np.random.rand(*input_shape) + 0.5).astype(np.int16)

    td = TensorDesc("I16", input_shape, "NCHW")

    input_blob_i16 = Blob(td, input_data_int16)

    assert np.all(np.equal(input_blob_i16.buffer, input_data_int16))

@pytest.mark.skip(reason="Fix")
def test_ie_core_class():
    input_shape = [1, 3, 4, 4]
    param = ov.parameter(input_shape, np.float32, name="parameter")
    relu = ov.relu(param, name="relu")
    func = Function([relu], [param], "test")
    func.get_ordered_ops()[2].friendly_name = "friendly"

    cnn_network = IENetwork(func)

    ie_core = Core()
    ie_core.set_config({}, device_name="CPU")
    executable_network = ie_core.load_network(cnn_network, "CPU", {})

    td = TensorDesc("FP32", input_shape, "NCHW")

    # from IPython import embed; embed()

    request = executable_network.create_infer_request()
    input_data = np.random.rand(*input_shape) - 0.5

    expected_output = np.maximum(0.0, input_data)

    input_blob = Blob(td, input_data)

    request.set_input({"parameter": input_blob})
    request.infer()

    result = request.get_blob("relu").buffer

    assert np.allclose(result, expected_output)

def test_load_network(device):
    ie = Core()
    net = ie.read_network(model=test_net_xml, weights=test_net_bin)
    exec_net = ie.load_network(net, device)
    assert isinstance(exec_net, ExecutableNetwork)

def test_read_network():
    ie_core = Core()
    net = ie_core.read_network(model=test_net_xml, weights=test_net_bin)
    assert isinstance(net, IENetwork)

    net = ie_core.read_network(model=test_net_xml)
    assert isinstance(net, IENetwork)

def test_read_network_from_blob():
    ie_core = Core()
    model = open(test_net_xml).read()
    blob = blob_from_file(test_net_bin)
    net = ie_core.read_network(model=model, blob=blob)
    assert isinstance(net, IENetwork)

def test_read_network_from_blob_valid():
    ie_core = Core()
    model = open(test_net_xml).read()
    blob = blob_from_file(test_net_bin)
    net = ie_core.read_network(model=model, blob=blob)
    ref_net = ie_core.read_network(model=test_net_xml, weights=test_net_bin)
    assert net.name == ref_net.name
    assert net.batch_size == ref_net.batch_size
    ii_net = net.input_info
    ii_net2 = ref_net.input_info
    o_net = net.outputs
    o_net2 = ref_net.outputs
    assert ii_net.keys() == ii_net2.keys()
    assert o_net.keys() == o_net2.keys()

def test_read_network_as_path():
    ie_core = Core()
    net = ie_core.read_network(model=Path(test_net_xml), weights=Path(test_net_bin))
    assert isinstance(net, IENetwork)

    net = ie_core.read_network(model=test_net_xml, weights=Path(test_net_bin))
    assert isinstance(net, IENetwork)

    net = ie_core.read_network(model=Path(test_net_xml))
    assert isinstance(net, IENetwork)

def test_read_network_from_onnx():
    ie_core = Core()
    net = ie_core.read_network(model=test_net_onnx)
    assert isinstance(net, IENetwork)

def test_read_network_from_onnx_as_path():
    ie_core = Core()
    net = ie_core.read_network(model=Path(test_net_onnx))
    assert isinstance(net, IENetwork)

def test_read_net_from_buffer():
    ie_core = Core()
    with open(test_net_bin, "rb") as f:
        bin = f.read()
    with open(model_path()[0], "rb") as f:
        xml = f.read()
    net = ie_core.read_network(model=xml, weights=bin)
    assert isinstance(net, IENetwork)

def test_net_from_buffer_valid():
    ie_core = Core()
    with open(test_net_bin, "rb") as f:
        bin = f.read()
    with open(model_path()[0], "rb") as f:
        xml = f.read()
    net = ie_core.read_network(model=xml, weights=bin)
    ref_net = ie_core.read_network(model=test_net_xml, weights=test_net_bin)
    assert net.name == ref_net.name
    assert net.batch_size == ref_net.batch_size
    ii_net = net.input_info
    ii_net2 = ref_net.input_info
    o_net = net.outputs
    o_net2 = ref_net.outputs
    assert ii_net.keys() == ii_net2.keys()
    assert o_net.keys() == o_net2.keys()

def test_get_version(device):
    ie = Core()
    version = ie.get_versions(device)
    assert isinstance(version, dict), "Returned version must be a dictionary"
    assert device in version, "{} plugin version wasn't found in versions"
    assert hasattr(version[device], "major"), "Returned version has no field 'major'"
    assert hasattr(version[device], "minor"), "Returned version has no field 'minor'"
    assert hasattr(version[device], "description"), "Returned version has no field 'description'"
    assert hasattr(version[device], "build_number"), "Returned version has no field 'build_number'"

def test_available_devices(device):
    ie = Core()
    devices = ie.available_devices
    assert device in devices, f"Current device '{device}' is not listed in " \
                              f"available devices '{', '.join(devices)}'"

def test_get_config():
    ie = Core()
    conf = ie.get_config("CPU", "CPU_BIND_THREAD")
    assert conf == "YES"

@pytest.mark.skipif(os.environ.get("TEST_DEVICE", "CPU") != "CPU",
                    reason=f"Cannot run test on device {os.environ.get('TEST_DEVICE')}, Plugin specific test")
def test_get_metric_list_of_str():
    ie = Core()
    param = ie.get_metric("CPU", "OPTIMIZATION_CAPABILITIES")
    assert isinstance(param, list), "Parameter value for 'OPTIMIZATION_CAPABILITIES' " \
                                    f"metric must be a list but {type(param)} is returned"
    assert all(isinstance(v, str) for v in param), \
        "Not all of the parameter values for 'OPTIMIZATION_CAPABILITIES' metric are strings!"

@pytest.mark.skipif(os.environ.get("TEST_DEVICE", "CPU") != "CPU",
                    reason=f"Cannot run test on device {os.environ.get('TEST_DEVICE')}, Plugin specific test")
def test_get_metric_tuple_of_two_ints():
    ie = Core()
    param = ie.get_metric("CPU", "RANGE_FOR_STREAMS")
    assert isinstance(param, tuple), "Parameter value for 'RANGE_FOR_STREAMS' " \
                                     f"metric must be tuple but {type(param)} is returned"
    assert all(isinstance(v, int) for v in param), \
        "Not all of the parameter values for 'RANGE_FOR_STREAMS' metric are integers!"

@pytest.mark.skipif(os.environ.get("TEST_DEVICE", "CPU") != "CPU",
                    reason=f"Cannot run test on device {os.environ.get('TEST_DEVICE')}, Plugin specific test")
def test_get_metric_tuple_of_three_ints():
    ie = Core()
    param = ie.get_metric("CPU", "RANGE_FOR_ASYNC_INFER_REQUESTS")
    assert isinstance(param, tuple), "Parameter value for 'RANGE_FOR_ASYNC_INFER_REQUESTS' " \
                                     f"metric must be tuple but {type(param)} is returned"
    assert all(isinstance(v, int) for v in param), "Not all of the parameter values for " \
                                                   "'RANGE_FOR_ASYNC_INFER_REQUESTS' metric are integers!"

@pytest.mark.skipif(os.environ.get("TEST_DEVICE", "CPU") != "CPU",
                    reason=f"Cannot run test on device {os.environ.get('TEST_DEVICE')}, Plugin specific test")
def test_get_metric_str():
    ie = Core()
    param = ie.get_metric("CPU", "FULL_DEVICE_NAME")
    assert isinstance(param, str), "Parameter value for 'FULL_DEVICE_NAME' " \
                                   f"metric must be string but {type(param)} is returned"

def test_query_network(device):
    ie = Core()
    net = ie.read_network(model=test_net_xml, weights=test_net_bin)
    query_res = ie.query_network(network=net, device_name=device)
    func_net = net.get_function()
    ops_net = func_net.get_ordered_ops()
    ops_net_names = [op.friendly_name for op in ops_net]
    assert [key for key in query_res.keys() if key not in ops_net_names] == [], \
        "Not all network layers present in query_network results"
    assert next(iter(set(query_res.values()))) == device, "Wrong device for some layers"

@pytest.mark.skipif(os.environ.get("TEST_DEVICE", "CPU") != "CPU", reason="Device independent test")
def test_register_plugin():
    ie = Core()
    ie.register_plugin("MKLDNNPlugin", "BLA")
    net = ie.read_network(model=test_net_xml, weights=test_net_bin)
    exec_net = ie.load_network(net, "BLA")
    assert isinstance(exec_net, ExecutableNetwork), \
        "Cannot load the network to the registered plugin with name 'BLA'"

@pytest.mark.skipif(os.environ.get("TEST_DEVICE", "CPU") != "CPU", reason="Device independent test")
def test_register_plugins():
    ie = Core()
    if platform == "linux" or platform == "linux2":
        ie.register_plugins(plugins_xml)
    elif platform == "darwin":
        ie.register_plugins(plugins_osx_xml)
    elif platform == "win32":
        ie.register_plugins(plugins_win_xml)

    net = ie.read_network(model=test_net_xml, weights=test_net_bin)
    exec_net = ie.load_network(net, "CUSTOM")
    assert isinstance(exec_net,
                      ExecutableNetwork), "Cannot load the network to " \
                                          "the registered plugin with name 'CUSTOM' " \
                                          "registred in the XML file"

def test_create_IENetwork_from_nGraph():
    element_type = Type.f32
    param = Parameter(element_type, Shape([1, 3, 22, 22]))
    relu = ov.relu(param)
    func = Function([relu], [param], "test")
    cnnNetwork = IENetwork(func)
    assert cnnNetwork is not None
    func2 = cnnNetwork.get_function()
    assert func2 is not None
    assert len(func2.get_ops()) == 3

max_issues_repo_name: thomas-yanxin/openvino

max_forks_repo_name: thomas-yanxin/openvino


text:
import math as m

import numpy as np
from scipy import ndimage

if False:
    from typing import List

def image_to_zoom_tensor(image,  # type: np.ndarray
                         num_colors,  # type: int
                         center_dimensions,  # type: List[int]
                         scale  # type: float
                         ):  # type: (...)->np.ndarray
    """Converts an image to image pyramid.

    Each image in the image pyramid will have the same number of pixels. They will take larger and larger sections
     of the input image and scale them down to the center_dimensions size.

    Example:

    _____________________________________________________________________________________
    |                    |                    |                    |                    |
    |     **********     |        ****        |                    |                    |
    |   *            *   |      *      *      |        ****        |                    |
    | *                * |   *            *   |      *      *      |         **         |
    | *                * |   *            *   |      *      *      |         **         |
    |   *            *   |      *      *      |        ****        |                    |
    |     **********     |        ****        |                    |                    |
    |____________________|____________________|____________________|____________________|

    :param image: The input image.
    :param num_colors: The number of colors in the input image.
    :param center_dimensions: The size of the images used in the image pyramid.
    :param scale: The consecutive scales of the regions covered by the images in the image pyramid.
    :return: A tensor containing the image pyramid. Format: Pyramid[num_images, center_width, center_height, colors]
    """
    assert scale > 1, "Scale must be greater than one."
    assert num_colors > 0, "Number of colors must be greater than zero."
    for d in center_dimensions:
        assert d > 0, "Each dimension must be larger than zero."

    image_dimensions = image.shape[:-1]  # take out colors from n-dimensional image
    center_dimensions = list(reversed(center_dimensions))  # because OpenCV likes doing things backwards
    num_scales = int(
        m.ceil(max([m.log(x_i / x_c, scale) for x_i, x_c in zip(image_dimensions, center_dimensions)])))
    pyramid_tensor = np.empty([num_scales] + center_dimensions + [num_colors])
    for s in range(num_scales):
        slice_dimensions = [c * (scale ** s) for c in center_dimensions]
        image_slicers = [slice(int(max((i - c) / 2, 0)), int((i + c) / 2)) for i, c in
                         zip(image_dimensions, slice_dimensions)]
        image_slice = image[image_slicers]
        scaled_image_slice = np.empty(center_dimensions + [num_colors])
        for c in range(num_colors):
            zoomed_image = ndimage.zoom(
                np.squeeze(
                    image_slice[[slice(None)] * len(image_dimensions) + [slice(c, c + 1)]]
                ),
                1.0 / (scale ** s), prefilter=False, order=5)[
                [slice(None)] * len(center_dimensions) + [np.newaxis]]
            x_max = min(scaled_image_slice.shape[0], zoomed_image.shape[0])
            y_max = min(scaled_image_slice.shape[1], zoomed_image.shape[1])
            scaled_image_slice[[slice(None)] * len(center_dimensions) + [slice(c, c + 1)]][0:x_max, 0:y_max] = \
                zoomed_image[0:x_max, 0:y_max]

        pyramid_tensor_slicers = [slice(None) for _ in image_dimensions] + [slice(None)]
        pyramid_tensor[[slice(s, s + 1)] + pyramid_tensor_slicers] = scaled_image_slice[[np.newaxis] +
                                                                                pyramid_tensor_slicers]
    return pyramid_tensor

max_stars_repo_name: SimLeek/pySILEnT

max_forks_repo_name: SimLeek/pySILEnT

reason: import numpy,from scipy


text:
# Python code to run Sparse Regression Code (SPARCs) simulations
#
# Copyright (c) 2020 Kuan Hsieh

import numpy as np
from sparc import sparc_encode, sparc_decode

def sparc_sim(code_params, decode_params, awgn_var, rand_seed=None):

    """
    End-to-end simulation of Sparse Regression Code (SPARC) encoding/decoding
    in AWGN channel.
    """

    # Currently cheating as encoder directly passes fast transforms Ab and Az to
    # the deocder. (Decoder doesn't use random seed to generate fast transform.)

    # Simulation
    bits_i, beta0, x, Ab, Az = sparc_encode(code_params, awgn_var, rand_seed)
    y                        = awgn_channel(x, awgn_var, rand_seed)
    bits_o, beta, T, nmse, expect = sparc_decode(y, code_params, decode_params,
                                                 awgn_var, rand_seed, beta0, Ab, Az)

    # Error analysis
    ber     = calc_ber(bits_i, bits_o)
    cer     = 1.0*(ber>0)
    detect  = 1.0*(not (ber>0)^expect) # Frame error detected
    results = {'ber':ber, 'cer':cer, 't_final':T, 'nmse':nmse, 'detect':detect}

    # Section-wise error analysis
    if not code_params['modulated']:
        ser, loc_of_sec_errs, num_of_sec_errs = calc_ser(beta0, beta, code_params['L'])
        results_append = {'ser': ser,
                          'loc_of_sec_errs': loc_of_sec_errs,
                          'num_of_sec_errs': num_of_sec_errs}
    else:
        err_rates, loc_of_errs, num_of_errs = calc_ler_ver(beta0, beta, code_params['L'],
                                                                        code_params['K'])
        ler, ver, ser = err_rates
        loc_of_loc_errs, loc_of_val_errs, loc_of_sec_errs = loc_of_errs
        num_of_loc_errs, num_of_val_errs, num_of_sec_errs = num_of_errs
        results_append = {'ser': ser, 'ler': ler, 'ver': ver,
                          'loc_of_sec_errs': loc_of_sec_errs,
                          'loc_of_loc_errs': loc_of_loc_errs,
                          'loc_of_val_errs': loc_of_val_errs,
                          'num_of_sec_errs': num_of_sec_errs,
                          'num_of_loc_errs': num_of_loc_errs,
                          'num_of_val_errs': num_of_val_errs}
    results.update(results_append)

    # SNR parameters
    #snr = P/awgn_var              # Signal-to-noise ratio
    #C = 0.5 * np.log2(1 + snr)    # Channel capacity
    #EbN0 = 1/(2*R) * (P/awgn_var) # Energy per bit/noise PSD

    return results

######## Error analysis ########

def calc_ber(true_bin_array, est_bin_array):
    '''
    Calculate the bit error rate (BER)
    '''
    assert true_bin_array.dtype == 'bool'
    assert est_bin_array.dtype  == 'bool'
    k = true_bin_array.size
    assert k == est_bin_array.size
    return np.count_nonzero(np.bitwise_xor(true_bin_array, est_bin_array))/k

def calc_ser(beta0, beta, L):
    '''
    Find the section error rate of the estimated message vector.

    beta0 : true message vector
    beta  : estimated message vector
    L     : number of sections

    Returns
        ser: section error rate, i.e., # of sections decoded in error / L
        loc_of_sec_errs: list of integers specifying sections decoded in error
    '''
    assert beta.size == beta0.size,  'beta and beta0 are of different size'
    assert beta.dtype == beta0.dtype,'beta and beta0 are of different type'
    assert type(L)==int and L>0
    assert beta.size % L == 0

    M = beta.size // L
    error_array = np.zeros(L, dtype=bool)
    for l in range(L):
        error_array[l] = not np.array_equal(beta[l*M:(l+1)*M], beta0[l*M:(l+1)*M])

    num_of_sec_errs = np.count_nonzero(error_array)
    ser             = num_of_sec_errs / L
    loc_of_sec_errs = np.flatnonzero(error_array)

    return ser, loc_of_sec_errs, num_of_sec_errs

def calc_ler_ver(beta0, beta, L, K):
    '''
    Find the location error rate and value error rate of
    the estimated message vector.

    [Parameters]

    beta0 : true message vector
    beta  : estimated message vector
    L     : number of sections
    K     : number of K-PSK constellation symbols

    [Returns]

    ler: location error rate, i.e., # of nonzero locations in error / L
    ver: value error rate, i.e., # of nonzero values in error / L
    loc_of_loc_errs: list of integers specifying locations of location error
    loc_of_val_errs: list of integers specifying locations of value error

    ----------

    For modulated SPARCs, L*LogM bits are encoded in the location of the
    non-zero entries of the message vector whilst the other L*logK bits are
    encoded in the value of those non-zero entries.

    In addition to finding the fraction of message vector sections decoded
    in error, i.e. the section error rate (SER), we would like to find the
    fraction of non-zero entry locations were decoded in error, i.e.
    location error rate (LER), and also the fraction of non-zero entry values
    decoded in error, i.e. value error rate (VER).

    The LER is the same as the SER for non-modulated SPARCs – when a location
    is decoded in error, (on average) half of the bits are in error.
    The VER is different – if the non-zero value is decoded in error, it
    depends on whether the error is small or large. A small decoding error
    will result in a small bit error, e.g. decoding a symbol to its
    neighbouring constellation symbol would result in a single bit error with
    the use of Gray coding.

    '''
    assert beta.size == beta0.size,  'beta and beta0 are of different size'
    assert beta.dtype == beta0.dtype,'beta and beta0 are of different type'
    assert type(L)==int and L>0
    assert beta.size % L == 0

    M = beta.size // L

    ## Probably don't need for loop, reshape to L-by-M matrix
    beta0_reshape  = beta0.reshape(L,M)
    beta_reshape   = beta.reshape(L,M)
    idxs11, idxs12 = np.nonzero(beta0_reshape)
    idxs21, idxs22 = np.nonzero(beta_reshape)
    assert np.array_equal(idxs11, np.arange(L)) # Exactly 1 nonzero each row
    assert np.array_equal(idxs21, np.arange(L)) # Exactly 1 nonzero each row
    vals1 = beta0_reshape[(idxs11,idxs12)]
    vals2 = beta_reshape[(idxs21,idxs22)]

    loc_err = np.logical_not(idxs12 == idxs22)
    val_err = np.logical_not(vals1 == vals2)
    sec_err = np.logical_or(loc_err, val_err)

    num_of_loc_errs = np.count_nonzero(loc_err)
    num_of_val_errs = np.count_nonzero(val_err)
    num_of_sec_errs = np.count_nonzero(sec_err)
    ler = num_of_loc_errs / L
    ver = num_of_val_errs / L
    ser = num_of_sec_errs / L
    loc_of_loc_errs = np.flatnonzero(loc_err)
    loc_of_val_errs = np.flatnonzero(val_err)
    loc_of_sec_errs = np.flatnonzero(sec_err)

    num_of_errs = (num_of_loc_errs, num_of_val_errs, num_of_sec_errs)
    error_rates = (ler, ver, ser)
    loc_of_errs = (loc_of_loc_errs, loc_of_val_errs, loc_of_sec_errs)

    return error_rates, loc_of_errs, num_of_errs

######## Channel models ########

def awgn_channel(input_array, awgn_var, rand_seed):
    '''
    Adds Gaussian noise to input array

    Real input_array:
        Add Gaussian noise of mean 0 variance awgn_var.

    Complex input_array:
        Add complex Gaussian noise. Indenpendent Gaussian noise of mean 0
        variance awgn_var/2 to each dimension.
    '''

    assert input_array.ndim == 1, 'input array must be one-dimensional'
    assert awgn_var >= 0

    rng = np.random.RandomState(rand_seed)
    n   = input_array.size

    if input_array.dtype == np.float:
        return input_array + np.sqrt(awgn_var)*rng.randn(n)

    elif input_array.dtype == np.complex:
        return input_array + np.sqrt(awgn_var/2)*(rng.randn(n)+1j* rng.randn(n))

    else:
        raise Exception("Unknown input type '{}'".format(input_array.dtype))


text:
# glottal voice source as input of Two Tubes Model of vocal tract
# Glottal Volume Velocity
# based on A.E.Rosenberg's formula as Glottal Volume Velocity

import numpy as np

class Glottal(object):
    def __init__(self, tclosed, trise, tfall, sampling_rate):
        # initalize
        self.tclosed = tclosed  # duration time of close state [mSec]
        self.trise = trise  # duration time of opening [mSec]
        self.tfall = tfall  # duration time of closing [mSec]
        self.sr = sampling_rate
        self.yg = self.make_one_plus()
        self.yg_repeat = self.make_N_repeat()

    def make_one_plus(
        self,
    ):
        # output yg
        self.N1 = int((self.tclosed / 1000.0) * self.sr)
        self.N2 = int((self.trise / 1000.0) * self.sr)
        self.N3 = int((self.tfall / 1000.0) * self.sr)
        self.LL = self.N1 + self.N2 + self.N3
        yg = np.zeros(self.LL)
        # print ('Length= ', self.LL)
        for t0 in range(self.LL):
            if t0 < self.N1:
                pass
            elif t0 <= (self.N2 + self.N1):
                yg[t0] = 0.5 * (1.0 - np.cos((np.pi / self.N2) * (t0 - self.N1)))
            else:
                yg[t0] = np.cos((np.pi / (2.0 * self.N3)) * (t0 - (self.N2 + self.N1)))
        return yg

    def make_N_repeat(self, repeat_num=200):
        self.yg_repeat = np.zeros(len(self.yg) * repeat_num)
        for loop in range(repeat_num):
            self.yg_repeat[len(self.yg) * loop : len(self.yg) * (loop + 1)] = self.yg
        return self.yg_repeat

    def get_output(self):
        return self.yg_repeat

    def fone(self, f):
        # calculate one point of frequecny response
        xw = 2.0 * np.pi * f / self.sr
        yi = 0.0
        yb = 0.0
        for v in range(0, (self.N2 + self.N3)):
            yi += self.yg[self.N1 + v] * np.exp(-1j * xw * v)
            yb += self.yg[self.N1 + v]
        val = yi / yb
        return np.sqrt(val.real ** 2 + val.imag ** 2)

    def H0(self, freq_low=100, freq_high=5000, Band_num=256):
        # get Log scale frequecny response, from freq_low to freq_high, Band_num points
        amp = []
        freq = []
        bands = np.zeros(Band_num + 1)
        fcl = freq_low * 1.0  # convert to float
        fch = freq_high * 1.0  # convert to float
        delta1 = np.power(fch / fcl, 1.0 / (Band_num))  # Log Scale
        bands[0] = fcl
        # print ("i,band = 0", bands[0])
        for i in range(1, Band_num + 1):
            bands[i] = bands[i - 1] * delta1
            # print ("i,band =", i, bands[i])
        for f in bands:
            amp.append(self.fone(f))
        return np.log10(amp) * 20, bands  # = amp value, freq list

max_stars_repo_name: yashlala/vocal-tube-model

max_issues_repo_name: yashlala/vocal-tube-model

max_forks_repo_name: yashlala/vocal-tube-model


text:
def get_samples_metadata_columns():
    import pandas as pd
    samples_meta_df = pd.read_csv('Belly_Button_Biodiversity_Metadata.csv')
    return list(samples_meta_df.columns)

def get_samples_metadata_values():
    import pandas as pd
    samples_meta_df = pd.read_csv('Belly_Button_Biodiversity_Metadata.csv')
    return samples_meta_df.values.tolist()

def get_otu_pie_labels():
    import numpy as np
    import pandas as pd
    import matplotlib.pyplot as plt
    import sqlalchemy
    from sqlalchemy.ext.automap import automap_base
    from sqlalchemy.orm import Session
    from sqlalchemy import create_engine, inspect
    from sqlalchemy import func, desc
    from matplotlib.ticker import NullFormatter
    import matplotlib.dates as mdates
    from datetime import datetime, timedelta
    import seaborn as sns
    from flask import Flask, jsonify
    import datetime as dt
    engine = create_engine("sqlite:///belly_button_biodiversity.sqlite", echo=False)
    Base = automap_base()
    Base.prepare(engine, reflect=True)
    Otu = Base.classes.otu
    session = Session(engine)
    otu_id_list = session.query(Otu.otu_id).all()
    otu_taxonomic_list = session.query(Otu.lowest_taxonomic_unit_found).all()
    otu_id = pd.DataFrame(otu_id_list)
    otu_taxonomic = pd.DataFrame(otu_taxonomic_list)
    otu_df = otu_id.join(otu_taxonomic)
    new_otu_df = otu_df.groupby('lowest_taxonomic_unit_found').count().sort_values(by=['otu_id'], ascending=False).reset_index()
    final_new_otu_df = new_otu_df.rename(columns={'otu_id' : "count"})
    otu_id_df = otu_df.groupby('lowest_taxonomic_unit_found').max().reset_index()
    final_otu_df = final_new_otu_df.merge(otu_id_df, how='inner', on='lowest_taxonomic_unit_found')
    named_labels = list(final_otu_df[:10]['otu_id'])
    return named_labels

def get_otu_pie_values():
    import numpy as np
    import pandas as pd
    import matplotlib.pyplot as plt
    import sqlalchemy
    from sqlalchemy.ext.automap import automap_base
    from sqlalchemy.orm import Session
    from sqlalchemy import create_engine, inspect
    from sqlalchemy import func, desc
    from matplotlib.ticker import NullFormatter
    import matplotlib.dates as mdates
    from datetime import datetime, timedelta
    import seaborn as sns
    from flask import Flask, jsonify
    import datetime as dt
    engine = create_engine("sqlite:///belly_button_biodiversity.sqlite", echo=False)
    Base = automap_base()
    Base.prepare(engine, reflect=True)
    Otu = Base.classes.otu
    session = Session(engine)
    otu_id_list = session.query(Otu.otu_id).all()
    otu_taxonomic_list = session.query(Otu.lowest_taxonomic_unit_found).all()
    otu_id = pd.DataFrame(otu_id_list)
    otu_taxonomic = pd.DataFrame(otu_taxonomic_list)
    otu_df = otu_id.join(otu_taxonomic)
    new_otu_df = otu_df.groupby('lowest_taxonomic_unit_found').count().sort_values(by=['otu_id'], ascending=False).reset_index()
    named_otu_pie = list(new_otu_df[:10]['otu_id'].values)

    final_named_otu_pie = []
    for i in named_otu_pie:
        final_named_otu_pie.append(int(i))

    return final_named_otu_pie

    return named_otu_pie

def get_samples():
    import pandas as pd
    samples_df = pd.read_csv('belly_button_biodiversity_samples.csv')
    return list(samples_df.columns[1:])


text:
import numpy as np
from imagepy.core.engine import Simple, Filter
from numpy.fft import fft2, ifft2, fftshift, ifftshift
from imagepy import IPy

class FFT(Simple):
	title = 'FFT'
	note = ['8-bit', '16-bit', 'int', 'float']
	para = {'shift':True, 'slice':False}
	view = [(bool, 'shift', 'zero center'),
			(bool, 'slice', 'slices')]

	def run(self, ips, imgs, para = None):
		if not para['slice']: imgs = [ips.img]
		shift = fftshift if para['shift'] else lambda x:x
		rst = []
		for i in range(len(imgs)):
			rst.append(shift(fft2(imgs[i])))
			self.progress(i, len(imgs))
		IPy.show_img(rst, '%s-fft'%ips.title)

class LogPower(Simple):
	title = 'Log Power'
	note = ['complex']
	para = {'slice':False, 'type':'float', 'log':2.718}
	view = [(float, 'log', (2,30), 3, 'log', ''),
			(list, 'type', ['uint8', 'int', 'float'], str, 'type', ''),
			(bool, 'slice', 'slices')]

	def run(self, ips, imgs, para = None):
		if not para['slice']: imgs = [ips.img]
		tp = {'uint8':np.uint8, 'int':np.int32, 'float':np.float32}
		rst, tp = [], tp[para['type']]
		for i in range(len(imgs)):
			zs = np.log(np.abs(imgs[i]))
			zs /= np.log(para['log'])
			rst.append(zs.astype(tp))
			self.progress(i, len(imgs))
		IPy.show_img(rst, '%s-fft'%ips.title)

class IFFT(Simple):
	title = 'Inverse FFT'
	note = ['complex']
	para = {'shift':True, 'slice':False, 'type':'float'}
	view = [(list, 'type', ['uint8', 'int', 'float'], str, 'type', ''),
			(bool, 'shift', 'zero center'),
			(bool, 'slice', 'slices')]

	def run(self, ips, imgs, para = None):
		if not para['slice']: imgs = [ips.img]
		shift = ifftshift if para['shift'] else lambda x:x
		tp = {'uint8':np.uint8, 'int':np.int32, 'float':np.float32}
		rst, tp = [], tp[para['type']]
		for i in range(len(imgs)):
			rst.append(ifft2(shift(ips.img)).astype(tp))
			self.progress(i, len(imgs))
		IPy.show_img(rst, '%s-ifft'%ips.title)

class Shift(Filter):
	title = 'Zero Center'
	note = ['complex']

	def run(self, ips, snap, img, para = None):
		return fftshift(img)

class IShift(Filter):
	title = 'Zero Edge'
	note = ['complex']

	def run(self, ips, snap, img, para = None):
		return ifftshift(img)

plgs = [FFT, IFFT, '-', Shift, IShift, LogPower]

max_stars_repo_name: BioinfoTongLI/imagepy

max_issues_repo_name: BioinfoTongLI/imagepy

max_forks_repo_name: BioinfoTongLI/imagepy

reason: import numpy,from numpy


text:
#!/usr/bin/env python
# coding: utf-8

# # Data analysis of a 2AFC task

# In[ ]:

# Modifiable variables when changing computer location:
# Directory where the images will be stored:
directory_images = '/home/emma/github/TFG/results/fit_analysis/'
directory_functions = '/home/emma/github/TFG/functions'
directory_data = '/home/emma/github/TFG/data/'
directory_results = '/home/emma/github/TFG/results/'

# In[ ]:

# Importing all the libraries that will be used
import numpy as np
import statsmodels.api as sm
import matplotlib.pyplot as plt
from scipy.optimize import curve_fit
import json
from sklearn.linear_model import LinearRegression
import sys
import datetime

# Insert path where the functions are stored
sys.path.insert(1, directory_functions)

# Importing the functions that will be used
import rat_functions1 as rf
import help_plot as hp

# ## Storing the data into variables

# In[ ]:

# Name of the rats
rats = ["Rat 24","Rat 25","Rat 35", "Rat 36", "Rat 37"]

# ## Open fit stimulus

# In[ ]:

# Opening JSON file
f = open(directory_data+'ordered_processed_data.json',)
# returns JSON object as
# a dictionary
data = json.load(f)
f.close()

stimulus = data[0]
coherences = data[1]
target_sigmas = data[5]
rewards = data[2]
performances = data[4]
choices = data[3]
n_total_stimulus = [len(a) for a in stimulus]
dates = data[6]

# In[ ]:

for rat in range(len(rats)):
    print(rats[rat]+": "+str(len(stimulus[rat])))

# ## Segment data

# In[ ]:

# i = 1
# t = 2
# # end = len(stimulus)-1
# stimulus = [stimulus[a][1000:len(stimulus[a])-1] for a in range(len(stimulus))]
# # stimulus_test = stimulus[:][8000:10000]
# coherences = [coherences[a][1000:len(stimulus[a])-1] for a in range(len(coherences))]
# # coherences_test = coherences[:][8000:10000]
# target_sigmas = [target_sigmas[a][1000:len(stimulus[a])-1] for a in range(len(target_sigmas))]
# # target_sigmas_test = target_sigmas[:][8000:10000]
# rewards =[rewards[a][1000:len(stimulus[a])-1] for a in range(len(rewards))]
# # rewards_test = rewards[:][8000:10000]
# performances =[performances[a][1000:len(stimulus[a])-1] for a in range(len(performances))]
# # performances_test = performances[:][8000:10000]
# choices = [choices[a][1000:len(stimulus[a])-1] for a in range(len(choices))]
# # choices_test = choices[:][8000:10000]
# choices_fit = [choices_fit[a][1000:len(stimulus[a])-1] for a in range(len(choices_fit))]

# In[ ]:

# i = 1
# t = 2
# # end = len(stimulus)-1
# stimulus = [stimulus[a][15000:3000] for a in range(len(stimulus))]
# # stimulus_test = stimulus[:][8000:10000]
# coherences = [coherences[a][10000:30000] for a in range(len(coherences))]
# # coherences_test = coherences[:][8000:10000]
# target_sigmas = [target_sigmas[a][10000:30000] for a in range(len(target_sigmas))]
# # target_sigmas_test = target_sigmas[:][8000:10000]
# rewards =[rewards[a][10000:30000] for a in range(len(rewards))]
# # rewards_test = rewards[:][8000:10000]
# performances =[performances[a][10000:30000] for a in range(len(performances))]
# # performances_test = performances[:][8000:10000]
# choices = [choices[a][10000:30000] for a in range(len(choices))]
# # choices_test = choices[:][8000:10000]
# # choices_fit = [choices_fit[a][10000:30000] for a in range(len(choices_fit))]

# ## Separing the data by coherences (real and target)
# ### Separing the data by target coherences

# In[ ]:

# Will compute a vector with the target coherences used for each rat
coherence_vectors = [] # coherence_vectors[rat][coherence]

for coh in coherences:
    coh_vec = rf.return_coherences_vector(coh)
    coherence_vectors.append(coh_vec)

# In[ ]:

coherence_vectors

# In[ ]:

# The stimulus, coherences, rewards, decisions, performances and target sigmas will be divided by coherences
results_divided_coherence = [] # The vector will be organized as follows:
                                                  # results_divided_coherence[rat][coherence]=[rewards, decisions, performances, stimulus, target_sigmas]

for rat in range(len(rats)):
    divided_coherence = []
    for coherence in coherence_vectors[rat]:
        divided_coherence.append(rf.divide_coh(coherences[rat],rewards[rat],choices[rat],performances[rat],stimulus[rat],target_sigmas[rat],coherence))
    results_divided_coherence.append(divided_coherence)

# In[ ]:

len(results_divided_coherence)

# In[ ]:

# The percentages, the number of trials and the coherence vectors avobe x trials will be stored into the following variables for each rat:
percentages = [] # percentages[rat][strong_coherence]
total_number_trials = [] # total_number_trials[rat][strong_coherence]
strong_coherence_vectors = [] # strong_coherence_vectors[rat] coherences which have more than 2500 trials

n_trials_threshold = 1000 # Only the coherences which have avobe n_trials_threshold will be used

for rat in range(len(rats)):
    perc = []
    number_trials = []
    coh_vector = []
    # for each coherence which has more than n_trials_threshold, the percentage of right choice will be computed
    for i in range(len(results_divided_coherence[rat])):
        if len(results_divided_coherence[rat][i][1]) > n_trials_threshold:
            perc.append(sum(results_divided_coherence[rat][i][1])/len(results_divided_coherence[rat][i][1]))
            coh_vector.append(coherence_vectors[rat][i])
            number_trials.append(len(results_divided_coherence[rat][i][1]))
    percentages.append(perc)
    total_number_trials.append(number_trials)
    strong_coherence_vectors.append(coh_vector)

# In[ ]:

# For each rat, plot the psychometric curve with the coherence values which have above x trials
# fig, axs = plt.subplots(figsize = (8,5))
for rat in range(len(rats)):
    fig, axs = plt.subplots(figsize = (8,5))
    # Computing the confidence intervals (95%)
    z = 1.96
    conf_int = [[percentages[rat][i]-z*(np.sqrt((percentages[rat][i]*(1-percentages[rat][i]))/total_number_trials[rat][i])) for i in range(len(total_number_trials[rat]))],
                [percentages[rat][i]+z*(np.sqrt((percentages[rat][i]*(1-percentages[rat][i]))/total_number_trials[rat][i])) for i in range(len(total_number_trials[rat]))]]
    conf_int = [[percentages[rat][i]-conf_int[0][i] for i in range(len(percentages[rat]))],[conf_int[1][i]-percentages[rat][i] for i in range(len(percentages[rat]))]]
    axs.errorbar(strong_coherence_vectors[rat],percentages[rat],conf_int,marker='o',label = "Real data")
#     conf_int = [[percentages_fit[rat][i]-z*(np.sqrt((percentages_fit[rat][i]*(1-percentages_fit[rat][i]))/total_number_trials_fit[rat][i])) for i in range(len(total_number_trials_fit[rat]))],
#                 [percentages_fit[rat][i]+z*(np.sqrt((percentages_fit[rat][i]*(1-percentages_fit[rat][i]))/total_number_trials_fit[rat][i])) for i in range(len(total_number_trials_fit[rat]))]]
#     conf_int = [[percentages_fit[rat][i]-conf_int[0][i] for i in range(len(percentages_fit[rat]))],[conf_int[1][i]-percentages_fit[rat][i] for i in range(len(percentages_fit[rat]))]]
# #     axs.errorbar(strong_coherence_vectors_fit[rat],percentages_fit[rat],conf_int,marker='o',label = "Fitted model")
#     conf_int = [[percentages_fit[rat][i]-z*(np.sqrt((percentages_fit[rat][i]*(1-percentages_fit[rat][i]))/total_number_trials_fit[rat][i])) for i in range(len(total_number_trials_fit[rat]))],
#                 [percentages_fit[rat][i]+z*(np.sqrt((percentages_fit[rat][i]*(1-percentages_fit[rat][i]))/total_number_trials_fit[rat][i])) for i in range(len(total_number_trials_fit[rat]))]]
#     axs.fill_between(strong_coherence_vectors_fit[rat],conf_int[0],conf_int[1],alpha=0.2,facecolor='#089FFF', linewidth=4, linestyle='dashdot', antialiased=True,label = "Fitted model")

    fig.suptitle(" Psychometric curve "+rats[rat],fontsize = 18)
    axs.legend(loc = "upper left",fontsize = 12)
    fig.text(0.5, 0.02, "Stimulus evidence (coherence)", ha='center',fontsize = 16)
    fig.text(0.03, 0.5, "Probability to choose right", va='center', rotation='vertical',fontsize = 18)
    axs.axis([-1.1,1.1,0,1])
    hp.xticks(axs,[-1,-0.75,-0.5,-0.25,0,0.25,0.5,0.75,1],xticklabels=None,fontsize=14)
    hp.yticks(axs,[0,0.2,0.4,0.6,0.8,1],yticklabels = None,fontsize =14)

    hp.remove_axis(axs)
    plt.hlines(0.5,-1,1,linestyle='--',color='lightgrey')
    plt.vlines(0,0.05,0.95,linestyle='--',color='lightgrey')

# fig.savefig(directory_images+'psychometric_curves/psychometric_curves_target_coherences.png', bbox_inches = 'tight')

# ## Separing the data by sigma and coherence

# In[ ]:

# Will compute a vector with the sigma used for each rat and each coherence
sigma_lists = [] #sigma_lists[rat][sigma]

for rat in range(len(rats)):
    sigma_list = []
    for i in range(len(results_divided_coherence[rat])):
        for a in rf.count(results_divided_coherence[rat][i][4]):
            sigma_list.append(a)
    sigma_list = rf.count(sigma_list)
    sigma_list = sorted(sigma_list)
    sigma_lists.append(sigma_list)

# In[ ]:

# Data for each sigma will be stored
results_divided_sigma = [] # results_divided_sigma[rat][coherence][sigma][reward, decision, performance, stim]

for rat in range(len(rats)):
    divided_sigma = []
    for i in range(len(results_divided_coherence[rat])):
        results = []
        for sigma in sigma_lists[rat]:
            results.append(rf.divide_sigma(results_divided_coherence[rat][i][0],results_divided_coherence[rat][i][1],results_divided_coherence[rat][i][2],
                                             results_divided_coherence[rat][i][3],results_divided_coherence[rat][i][4],sigma))
        divided_sigma.append(results)
    results_divided_sigma.append(divided_sigma)

# In[ ]:

# Compute the percentages of right choice for the results_divided_sigma
percentages_lists, coherences_lists, length_lists, length_all_rats = rf.compute_percentages(rats,results_divided_sigma,coherence_vectors)

# ### Comparing coherence distribution for each sigma

# In[ ]:

# Computing the coherence vector for all rats together
coherences_all_rats = []
for rat in range(len(rats)):
    coherences_all_rats.append(coherence_vectors[rat])
coherences_all_rats = rf.return_vector(coherences_all_rats)

# Computing the sigma vector for all rats tohether

sigmas_all_rats = rf.return_vector(sigma_lists)

# In[ ]:

# Generates a file with a table of the number of trials for each coherence and sigma for all rats
rf.recount_trials_all_rats(sigma_lists,coherences_all_rats,length_all_rats,"all rats",coherence_vectors,rats,sigmas_all_rats,directory_images)
# Will make the chart with the recount of trials for each coherence and each sigma for each rat
for rat in range(len(rats)):
    rf.recount_trials(sigma_lists[rat],coherence_vectors[rat],length_lists[rat],rats[rat],coherences_lists[rat],directory_images)

# In[ ]:

# Compute sensitivity and bias using curve_fit
all_sensitivity, all_bias = rf.compute_sensitivity(rats,coherences_lists,percentages_lists,sigma_lists,directory_images)

# In[ ]:

# Performing bootstrap to compute confidence intervals
n_trials = 10

all_percentages_lists_boot = []
all_coherences_lists_boot = []
all_length_lists_boot = []
all_length_all_rats_boot  = []
for i in range(n_trials):
    new_divided_sigma = rf.bootstrap(results_divided_sigma)
    percentages_lists_boot, coherences_lists_boot, length_lists_boot, length_all_rats_boot = rf.compute_percentages(rats,new_divided_sigma,coherence_vectors)
    all_percentages_lists_boot.append(percentages_lists_boot)
    all_coherences_lists_boot.append(coherences_lists_boot)
    all_length_lists_boot.append(length_lists_boot)
    all_length_all_rats_boot.append(length_all_rats_boot)

# In[ ]:

# Compute sensitivity and bias using curve_fit for the bootstrap data
all_sensitivity_boot = []
all_bias_boot = []
for i in range(len(all_percentages_lists_boot)):
    sensitivity_boot, bias_boot = rf.compute_sensitivity_boot(rats,all_coherences_lists_boot[i],all_percentages_lists_boot[i],sigma_lists)
    all_sensitivity_boot.append(sensitivity_boot)
    all_bias_boot.append(bias_boot)

# In[ ]:

# Computing the confidence intervals for the sensitivity
sensitivity_sigma = [np.zeros((len(sigma_lists[0]),n_trials)),np.zeros((len(sigma_lists[1]),n_trials)),np.zeros((len(sigma_lists[2]),n_trials)),np.zeros((len(sigma_lists[3]),n_trials)),np.zeros((len(sigma_lists[4]),n_trials))]

for trial in range(len(all_sensitivity_boot)):
    for rat in range(len(all_sensitivity_boot[trial])):
        for sigm in range(len(all_sensitivity_boot[trial][rat])):
             sensitivity_sigma[rat][sigm][trial] = all_sensitivity_boot[trial][rat][sigm]

conf_intervals_sensitivity = []
for rat in range(len(rats)):
    sensitivity_rat = []
    for a in range(len(all_sensitivity_boot[0][rat])):
        sens = sensitivity_sigma[rat][a]

        z = 1.96
        conf_int = [np.mean(sens)-z*np.std(sens),np.mean(sens)+z*np.std(sens)]
        sensitivity_rat.append(conf_int)

    conf_intervals_sensitivity.append(sensitivity_rat)

# In[ ]:

# Computing the bias confidence intervals
bias_sigma = [np.zeros((len(sigma_lists[0]),n_trials)),np.zeros((len(sigma_lists[1]),n_trials)),np.zeros((len(sigma_lists[2]),n_trials)),np.zeros((len(sigma_lists[3]),n_trials)),np.zeros((len(sigma_lists[4]),n_trials))]

for trial in range(len(all_bias_boot)):
    for rat in range(len(all_bias_boot[trial])):
        for sigm in range(len(all_bias_boot[trial][rat])):
             bias_sigma[rat][sigm][trial] = all_bias_boot[trial][rat][sigm]

conf_intervals_bias = []
for rat in range(len(rats)):
    bias_rat = []
    for a in range(len(all_bias_boot[0][rat])):
        bias = bias_sigma[rat][a]

        z = 1.96
        conf_int = [np.mean(bias)-z*np.std(bias),np.mean(bias)+z*np.std(bias)]
        bias_rat.append(conf_int)

    conf_intervals_bias.append(bias_rat)

# In[ ]:

# Plot of the sensitivity over sigma
for rat in range(len(rats)):
    fig, axs = plt.subplots(figsize = (8,5))
    conf_int = [[conf_intervals_sensitivity[rat][i][0] for i in range(len(sigma_lists[rat]))],[conf_intervals_sensitivity[rat][i][1] for i in range(len(sigma_lists[rat]))]]
    conf_int_0 = []
    conf_int_1 = []
    new_sigma_list = []
    new_sensitivity = []
    for i in range(len(conf_int[0])):
        if str(conf_int[0][i]) != "nan":
            if all_sensitivity[rat][i] != None:

                conf_int_0.append(conf_int[0][i]-all_sensitivity[rat][i])
                conf_int_1.append(all_sensitivity[rat][i]-conf_int[1][i])
                new_sigma_list.append(sigma_lists[rat][i])
                new_sensitivity.append(all_sensitivity[rat][i])
    conf_int = [conf_int_0,conf_int_1]
    plt.errorbar(new_sigma_list,new_sensitivity,conf_int,marker='o',color= "skyblue")
    plt.axis([-0.1,0.9,-1,50])

    hp.remove_axis(axs)
    fig.suptitle(rats[rat]+" sensitivity",fontsize = 18)
    fig.text(0.5, 0.02, "Sigma", ha='center',fontsize = 16)
    fig.text(0.03, 0.5, "Sensitivity", va='center', rotation='vertical',fontsize = 18)
    fig.savefig(directory_images+'sensitivity/sensitivity_'+rats[rat]+'.png', bbox_inches = 'tight')

# In[ ]:

# Plot of the bias over sigma
for rat in range(len(rats)):
    fig, axs = plt.subplots(figsize = (8,5))
    conf_int = [[conf_intervals_bias[rat][i][0] for i in range(len(sigma_lists[rat]))],[conf_intervals_bias[rat][i][1] for i in range(len(sigma_lists[rat]))]]
    conf_int_0 = []
    conf_int_1 = []
    new_sigma_list = []
    new_bias = []
    for i in range(len(conf_int[0])):
        if str(conf_int[0][i]) != "nan":
            if all_bias[rat][i] != None:
                conf_int_0.append(conf_int[0][i]-all_bias[rat][i])
                conf_int_1.append(all_bias[rat][i]-conf_int[1][i])
                new_sigma_list.append(sigma_lists[rat][i])
                new_bias.append(all_bias[rat][i])
    conf_int = [conf_int_0,conf_int_1]
    plt.errorbar(new_sigma_list,new_bias,conf_int,marker='o',color="skyblue")

    plt.hlines(0,-0.1,0.9,color = "grey",linestyles = "--")
    plt.axis([-0.1,0.9,-0.8,0.3])
    hp.remove_axis(axs)
    fig.suptitle(rats[rat]+" bias",fontsize = 18)
    fig.text(0.5, 0.02, "Sigma", ha='center',fontsize = 16)
    fig.text(0.03, 0.5, "Bias", va='center', rotation='vertical',fontsize = 18)
    fig.savefig(directory_images+'bias/bias'+rats[rat]+'.png', bbox_inches = 'tight')

# ## Dividing the data by sigma (without taking coherence into account)

# In[ ]:

results_divided_sigma_nocoh_fit = [] # results_divided_sigma_nocoh[rat][sigma][reward,decision,performance,stim]
perc_sigma_nocoh_fit = []
# Dividing the data over sigma (not by coherence)
results_divided_sigma_nocoh = [] # results_divided_sigma_nocoh[rat][sigma][reward,decision,performance,stim]
perc_sigma_nocoh = []

for rat in range(len(rats)):
    divided_sigma = []
    for sigma in sigma_lists[rat]:
        divided_sigma.append(rf.divide_sigma_1(rewards[rat],choices[rat],performances[rat],stimulus[rat],target_sigmas[rat],sigma))
    perc = []
    number_trials = []
    results_divided_sigma_nocoh.append(divided_sigma)
    for i in range(len(divided_sigma)):
        perc.append(sum(results_divided_coherence[rat][i][2])/len(results_divided_coherence[rat][i][2]))
        number_trials.append(len(results_divided_coherence[rat][i][1]))
    perc_sigma_nocoh.append(perc)

    fig, axs = plt.subplots(figsize = (8,5))
    plt.plot(sigma_lists[rat],perc,marker = 'o',color="skyblue")
    plt.axis([0,0.6,0,1])
    hp.remove_axis(axs)
    fig.suptitle(rats[rat]+" accuracy",fontsize = 18)
    fig.text(0.5, 0.02, "Sigma", ha='center',fontsize = 16)
    fig.text(0.03, 0.5, "Accuracy", va='center', rotation='vertical',fontsize = 18)
#     fig.savefig(directory_images+'accuracy_vs_sigma/accuracy_vs_sigma'+rats[rat]+'.png', bbox_inches = 'tight')
#     divided_sigma = []

#     for sigma in sigma_lists[rat]:
#         divided_sigma.append(rf.divide_sigma_1(rewards[rat],choices_fit[rat],performances[rat],stimulus[rat],target_sigmas[rat],sigma))
#     perc = []
#     number_trials = []
#     results_divided_sigma_nocoh_fit.append(divided_sigma)
#     for i in range(len(divided_sigma)):
#         perc.append(sum(results_divided_coherence_fit[rat][i][2])/len(results_divided_coherence_fit[rat][i][2]))
#         number_trials.append(len(results_divided_coherence_fit[rat][i][1]))
#     perc_sigma_nocoh_fit.append(perc)

#     plt.plot(sigma_lists[rat],perc,marker = 'o',color="red")
#     plt.axis([0,1,0,1])
#     hp.remove_axis(axs)
#     fig.suptitle(rats[rat]+" accuracy",fontsize = 18)
#     fig.text(0.5, 0.02, "Sigma", ha='center',fontsize = 16)
#     fig.text(0.03, 0.5, "Accuracy", va='center', rotation='vertical',fontsize = 18)
#     fig.savefig(directory_images+'accuracy_vs_sigma/accuracy_vs_sigma'+rats[rat]+'.png', bbox_inches = 'tight')

# PODRIA FER ELS CONFIDENCE INTERVALS AQUÍ

# In[ ]:

len(results_divided_sigma_nocoh[0][0][3])

# ## Kernel

# In[ ]:

for rat in range(len(rats)):
    for sig in range(len(results_divided_sigma_nocoh[rat])):
        for sti in range(len(results_divided_sigma_nocoh[rat][sig][3])):
            if len(results_divided_sigma_nocoh[rat][sig][3][sti]) == 20:
                group = 2

                results_divided_sigma_nocoh[rat][sig][3][sti] = np.ndarray.tolist(np.asarray(results_divided_sigma_nocoh[rat][sig][3][sti]).reshape(-1, group).mean(axis=1))

# In[ ]:

group = 2
a = np.ndarray.tolist(np.asarray(results_divided_sigma_nocoh[0][2][3][2]).reshape(-1, group).mean(axis=1))
print(results_divided_sigma_nocoh[0][2][3][2])
print(a)

# In[ ]:

def divided_time_PK(rats,results_divided_sigma_nocoh,results_divided_sigma_nocoh_fit,sigma_lists,rats_list,directory_images):
     # results_divided_sigma_coh[rat][sigma][reward,decision,performance,stim]
        fig, ((axs1,axs2,axs3)) = plt.subplots(1,3,figsize = (18,4))

        axs = [0,axs1,axs2,axs3,0,0]
        color = [0,"maroon","brown","red","orangered","orange","yellow"]
        index = 1
        aranges = [np.arange(1,4,1),np.arange(2,5,1),np.arange(2,5,1),np.arange(2,5,1),np.arange(1,4,1),np.arange(2,5,1),np.arange(2,5,1)]
        for rat in rats:
            for a in aranges[rat]:
                print(rat,len(results_divided_sigma_nocoh[rat][a][3][0]))
                try:
                    conf_int = PK(results_divided_sigma_nocoh[rat][a][1],results_divided_sigma_nocoh[rat][a][3],axs[index],color[index])
                    conf_int_fit = PK_fit(results_divided_sigma_nocoh_fit[rat][a][1],results_divided_sigma_nocoh_fit[rat][a][3],axs[index],color[index])
                    axs[index].set_title("Sigma = "+str(sigma_lists[rat][a]))
                    axs[index].spines['right'].set_visible(False)
                    axs[index].spines['top'].set_visible(False)
                except:
                    print("another error")
                index +=1
#         plt.ylim(-1.5,1.5)
        fig.suptitle("PK for "+rats_list[rat],fontsize = 18)
        fig.text(0.5, 0.001, 'Frame nº', ha='center',fontsize = 16)
        fig.text(0.08, 0.5, 'Impact', va='center', rotation='vertical',fontsize = 18)

        fig.savefig(directory_images+'PK/temporal_PK'+rats_list[rat]+'.png', bbox_inches = 'tight')

# In[ ]:

def PK(decision,stim,axs,col):
    logit_mod = sm.Logit(decision, stim)
    result = logit_mod.fit()

    pars_list =result.params
    confidence_interval = result.conf_int()
    conf_int = [[i[0] for i in confidence_interval],[i[1] for i in confidence_interval]]
    conf_int = [pars_list-conf_int[0],conf_int[1]-pars_list]
    axs.errorbar(np.arange(0,len(pars_list),1),pars_list,conf_int,color = col,marker='s')

    return(confidence_interval)

# In[ ]:

def PK_fit(decision,stim,axs,col):
    logit_mod = sm.Logit(decision, stim)
    result = logit_mod.fit()

    pars_list =result.params
    confidence_interval = result.conf_int()
    conf_int = [[i[0] for i in confidence_interval],[i[1] for i in confidence_interval]]

    axs.fill_between(np.arange(0,len(pars_list),1),conf_int[0],conf_int[1],alpha=0.2, facecolor='#089FFF',
    linewidth=4, linestyle='dashdot', antialiased=True)

    return(confidence_interval)

# In[ ]:

rats_sec = [1,2,4]
plt.figure()
divided_time_PK([0],results_divided_sigma_nocoh,results_divided_sigma_nocoh_fit,sigma_lists,rats,directory_images)

# In[ ]:

rats_sec = [1,2,4]
divided_time_PK([1],results_divided_sigma_nocoh,results_divided_sigma_nocoh_fit,sigma_lists,rats,directory_images)

# In[ ]:

rats_sec = [1,2,4]
divided_time_PK([2],results_divided_sigma_nocoh,results_divided_sigma_nocoh_fit,sigma_lists,rats,directory_images)

# In[ ]:

rats_half = [0,3]
divided_time_PK([3],results_divided_sigma_nocoh,results_divided_sigma_nocoh_fit,sigma_lists,rats,directory_images)

# In[ ]:

rats_half = [0,3]
divided_time_PK([4],results_divided_sigma_nocoh,results_divided_sigma_nocoh_fit,sigma_lists,rats,directory_images)

# In[ ]:

# Spatial Kernel for all sigmas for each rat

# For each stimulus, how many stimulus are from -1 to -0.9, etc to get a transformed stimulus, then make a logistic regression
color = [0,"maroon","brown","red","orangered","orange","yellow"]

new_transf_stim = rf.transformed_stimulus([stimulus[1],stimulus[2],stimulus[4]])
new_transf_sec_stim  = rf.transformed_stimulus([stimulus[0],stimulus[3]])
new_dec_half = [choices[1],choices[2],choices[4]]
new_dec_sec = [choices[0],choices[3]]

for rat in range(len([1,2,4])):
#     try:
    rats_half = [1,2,4]

    fig, axs = plt.subplots(figsize = (8,5))
    plt.xticks([0,2.5,5,7.5,10],["-1","-0.5","0","0.5","1"])

    logit_mod = sm.Logit(new_dec_half[rat], new_transf_stim[rat])
    result = logit_mod.fit()
    pars_list =result.params
    confidence_interval = result.conf_int()
    conf_int = [[i[0] for i in confidence_interval],[i[1] for i in confidence_interval]]
    conf_int = [pars_list-conf_int[0],conf_int[1]-pars_list]
    x_values = np.arange(0,len(pars_list),1)
    plt.errorbar(x_values,pars_list,conf_int,marker='o')

    slope, intercept = np.polyfit([x_values[4],x_values[5],x_values[6]], [pars_list[4],pars_list[5],pars_list[6]], 1)
    abline_values = [slope * i + intercept for i in x_values]
    plt.plot(x_values, abline_values,'grey', linestyle = "dashed")

    fig.suptitle("Spatial Kernel "+rats[rats_half[rat]],fontsize = 18)
    fig.text(0.5, 0.02, "Coherence value", ha='center',fontsize = 16)
    fig.text(0.03, 0.5, "Spatial Kernel", va='center', rotation='vertical',fontsize = 18)

    plt.ylim(-0.5,0.5)
    hp.remove_axis(axs)
    fig.savefig(directory_images+'PK/spatial_PK_'+rats[rats_half[rat]]+'_half_sec.png', bbox_inches = 'tight')
#     except:
#         print("no ha funcionat"+str(rat))

rats_sec = [0,3]

for rat in range(len(rats_sec)):
#     try:
#         print(rat)
    fig, axs = plt.subplots(figsize = (8,5))

    plt.xticks([0,2.5,5,7.5,10], ["-1","-0.5","0","0.5","1"])
    logit_mod = sm.Logit(new_dec_sec[rat], new_transf_sec_stim[rat])
    result = logit_mod.fit()
    pars_list =result.params
    confidence_interval = result.conf_int()
    conf_int = [[i[0] for i in confidence_interval],[i[1] for i in confidence_interval]]
    conf_int = [pars_list-conf_int[0],conf_int[1]-pars_list]
    plt.errorbar(np.arange(0,len(pars_list),1),pars_list,conf_int,marker='o')
    plt.ylim(-0.5,0.5)

    x_values = np.arange(0,len(pars_list),1)

    slope, intercept = np.polyfit([x_values[4],x_values[5],x_values[6]], [pars_list[4],pars_list[5],pars_list[6]], 1)
    abline_values = [slope * i + intercept for i in x_values]
    plt.plot(x_values, abline_values,'grey', linestyle = "dashed")
    plt.ylim(-0.5,0.5)

    fig.suptitle("Spatial Kernel "+rats[rats_sec[rat]],fontsize = 18)
    fig.text(0.5, 0.02, "Coherence value", ha='center',fontsize = 16)
    fig.text(0.03, 0.5, "Spatial Kernel", va='center', rotation='vertical',fontsize = 18)

    hp.remove_axis(axs)
    fig.savefig(directory_images+'PK/spatial_PK_'+rats[rats_sec[rat]]+'_sec.png', bbox_inches = 'tight')
#     except:
#         print("no ha funcionat"+str(rat))

max_stars_repo_name: emmafdz/FinalDegreeProject

max_issues_repo_name: emmafdz/FinalDegreeProject

max_forks_repo_name: emmafdz/FinalDegreeProject

reason: import numpy,from scipy,import statsmodels