text:
import numpy as np
import sys
from scipy.integrate import odeint
import time
from numpy.random import *

class ca1_spine:
    """Computational biophysical model of CA1 spine"""
    #### Global constants:
    rtol = 1e-6
    atol = 1e-6
    F = 96485.33 ## Coulomb/mole
    Nav = 6.022e23
    e = 2.718
    ############## Defining constant parameters ##################
    ## Spine compartment and ER size:
    #Vspine = 0.06 ## um^3, nominal spine volume
    Vr = 0.1 ## Ver/Vspine

    def __init__(self, g_N, nRyR, nip3, tau_refill, Vspine): #add ksk if required
        """initialize the spine parameters and volume/surface area-dependent constants:
        params: (nmda conductance (pS), #ryr, tau_refill (s), Vspine(um^3))"""
        self.g_N = float(g_N) * 1e-12 #S
        self.nRyR = int(nRyR)
        self.ip3r_tot = int(nip3)
        self.tau_refill = float(tau_refill)
        #self.Vscale = float(Vscale)
        #self.k_sk = float(k_sk)
        self.Vspine = float(Vspine) ## in um^3
        self.d_spine = (6*self.Vspine/3.14)**0.333  ## um
        self.Aspine = 3.14 * self.d_spine**2 * 1e-8  ## cm^2
        self.Vspine = self.Vspine * 1e-15 ## liter
        self.Aer = 0.1 * self.Aspine  ## cm^2
        self.Ver = 0.1 * self.Vspine ## liter
        self.Vspine -= self.Ver   ## liter
        #self.Vspine_um - self.Vspine * 1e15 ##um^3 for scaling AMPA conductance
            ## Total concentration of PMCA and NCX pumps in the spine head (uM):
        self.pHtot = (1e14) * 1000 * self.Aspine/(self.Nav * self.Vspine)
        self.pLtot = (1e14) * 140 * self.Aspine/(self.Nav * self.Vspine)
                ######################################################################################
        self.g_N_Ca = 0.1 * (self.g_N/(2*self.F*78.0*self.ca_ext)) * 1e6   ## Ca conductance of NMDAR channels; liter/sec
        self.k_erleak = self.Vmax_serca * (self.ca_0**2)/((self.Kd_serca**2 + self.ca_0**2)*(self.ca_er_0 - self.ca_0))  ## /s

    ## Parameters for endogenous immobile buffer (CBP):
    kbuff_f = 247 ## /uM/s
    kbuff_b = 524 ## /s

    ## Parameters for endogenous slow buffer:
    kslow_f = 24.7 ## /uM/s
    kslow_b = 52.4 ## /s

    ## Parameters for calbindin-Ca2+ kinetics:
    km0m1=174 ## /uM/s
    km1m2=87 ## /uM/s
    km1m0=35.8 ## /s
    km2m1=71.6 ## /s
    kh0h1=22 ## /uM/s
    kh1h2=11 ## /uM/s
    kh1h0=2.6 ## /s
    kh2h1=5.2 ## /s

    ## Parameters for PMCA and NCX pumps:
    k1H,k2H,k3H,kH_leak = [150,15,12,3.33]  ## (/uM/s, /s, /s, /s)
    k1L,k2L,k3L,kL_leak = [300,300,600,10]  ## (/uM/s, /s, /s, /s)

    ## Parameters for CaM-Ca2+ interaction:
    k1c_on = 6.8  ## /uM/s
    k1c_off = 68  ## /s
    k2c_on = 6.8 ## /uM/s
    k2c_off = 10 ## /s
    k1n_on = 108 ## /uM/s
    k1n_off = 4150 ## /s
    k2n_on = 108 ## /uM/s
    k2n_off = 800 ## /s

    ## Membrane and leak parameters:
    Cmem = 1e-6 ##  F/cm^2
    g_L = 2e-4  ## S/cm^2
    E_L = -70   ## mV

    ## AMPA receptor parameters:
    tau_A1 = 0.2e-3 ## s
    tau_A2 = 2e-3  ## s
    E_A = 0  ## mV
    g_A = 0.5e-9  ## S

    ## NMDA receptor parameters:
    tau_N1 = 5e-3 ## s
    tau_N2 = 50e-3 ## s
    E_N = 0  ## mV
    #g_N = 65e-12 ## S
    # insert NMDA conductance #

    ## L-VGCC parameters:
    um = -20 ## mV
    kmv = 5  ## mV
    tau_mv = 0.08e-3 ## sec
    uh = -65  ## mV
    khv = -7 ## mV
    tau_hv = 300e-3  ## sec

    ## Spine neck parameters:
    Rneck = 1e8  ## Ohm
    gc = 1.0/Rneck ## S
    rho_spines = 0  ## Surface density of co-active synaptic inputs on dendritic compartment (cm^-2)

    ## SERCA kinetic parameters:
    Vmax_serca = 1  ## uM/sec
    Kd_serca = 0.2 ## uM

    ## Parameters for Ca2+-based plasticity model:
    P1,P2,P3,P4 = [5.0,10.0,0.001,2]
    beta1,beta2 = [0.091,0.31]  ## /uM
    alpha1,alpha2 = [0.20,34.25] ## uM
    eta1, eta2 = [2, 10]
    omega_0 = 1.8 ## offset

    ## Thresholds set for canonical spine of volume 0.06 um^3

    ######################### ER Buffer Calreticulin Parameters ####################################
    #### Site C ####
    Kdiss_c = 2e3 #uM
    Scer = 20 #mol Ca/mol
    kcerf = 0.1 #uM^-1 s^-1

    #### Site P ####
    Kdiss_p = 10 #uM
    Sp = 1 #mol Ca/mol
    kpf = 100 #uM^-1 s^-1

    kcerb = kcerf * Kdiss_c
    kpb = kpf * Kdiss_p

    ################# ER Buffer ###########################
    totBer_0 = 3.6e3 #uM
    totBcer = Scer * totBer_0 #uM
    totBper = Sp * totBer_0 #uM
    #buff_flag = 1

    ###############################################################################################
    ############################################ SOCC Model ##########################################
    ############################################################### Defining parameters ######################################################
    V_socc_max = 1.57 #um/s per channel max flux
    K_socc = 50 #um/s
    tau_socc = 30 #s
    n_socc = 4 #Hill coefficient
    socc_clus = 10 # no of socc channels
    ###########################################################################################################################################

    #########################################################
    ########### Concentrations of various species ###########
    #########################################################

    ## External Ca (uM):
    ca_ext = 2e3

    ## Resting cytosolic Ca (uM):ip3
    ca_0 = 0.05

    ## Resting Ca in ER (uM):
    ca_er_0 = 250#300#460#1.21E4#1.182e4#250

    ## Total calbindin concentration in spine (uM):
    calb_tot = 45

    ## Total CBP concentration in the spine (uM):
    cbp_tot = 80

    ## Total slow buffer concentration in the spine (uM):
    Bslow_tot = 40

    ## Total concentration of CaM in spine (uM):
    cam_tot = 50

    #########################################################################################
    ################ RyR ########################################
    #### The model consists of a gating scheme consisting of 5 closed states and 2 open states ####
    #################### Rate Constants ###################################
    kon = 712.0 #1/(uM.s)
    koff = 3000.0 #### units of all the quantities hereafter is s^-1 ######
    kc4o1 = 10000.0
    kc4o2 = 1
    ko1c4 = 500
    ko2c4 = 0.5
    ko1c5 = 2.0
    ko2c5 = 3000.0
    kc5o1 = 0.6666
    kc5o2 = 100.0
    kc5I  = 0.5
    kIc5 = 1.5
    kRc1 = 4*kon
    kc1R = koff
    kc1c2 = 3*kon
    kc2c1 = 2*koff
    kc2c3 = 2*kon
    kc3c2 = 3*koff
    kc3c4 = kon
    kc4c3 = 4*koff
    ##########################################################
    ######## RyR parameters (play around these to make sure everyting works well together) ##########
    gRyR = 18.5 #s^-1  ## Johenning's data -> see analysis
    ####################################################################################
    ################### SK Channel parameters ##########################################
    # tau_sk = 6.3e-3 #s
    # n_sk = 4
    # g_sk = 25e-12 #S
    # E_sk = -90 #mV
    ###################################################################################
    ################################ IP3 params #######################################
    ############################### Reaction parameters for mGluR_IP3 pathway:
    PIP2 = 4000  ## uM
    a1f = 11.1 ## /uM/s
    a1b = 2  ## 100 /s
    a2f = 11.1 ## /uM/s
    a2b = 2  ## 100 /s
    a3f = 2 ## /uM/s
    a3b = 100 ## /s
    a4f = 2 ## /uM/s
    a4b = 100 ## /s
    a5 = 116  ## 116 /s
    a6 = 0.001 ## /s
    a7 = 0.02 ## /s
    a8 = 6 ## /s
    b1f = 300 ## /uM/s
    b1b = 100 ## /s
    b2f = 900 ## /uM/s
    b2b = 30 ## /s
    b3f = 800 ## /uM/s
    b3b = 40 ## /s
    b4f = 1200 ## /uM/s
    b4b = 6 ## /s
    b5f = 1200 ## /uM/s
    b5b = 6 ## /s
    b6 = 2 ## /s
    b7 = 160 ## /s
    b8f = 1 ## /uM/s
    b8b = 170 ## /s
    b9f = 1 ## /uM/s
    b9b = 170 ## /s
    b10 = 8 ## /s
    b11 = 2  ## 8 /s
    b12 = 8 ## /s
    DAGdegrate = 0.15 ## /s

    ## Parameters for IP3R model (Fink et al., 2000 and Vais et al., 2010):
    Kinh = 0.2  ## uM
    d1 = 0.8 ## uM
    d5 = 0.3 ## uM
    a2 = 2.7 ## /uM/s
    alpha_ip3r = (0.15/3.2)*(1e7)*(1e6)/500.0  ## /uM/sec

        ## Total concentrations of IP3 3-kinase and IP3 5-phosphatase in the spine (uM):
    ip5pconc = 1
    ip3kconc = 0.9

    ## Number of IP3R:
#     ip3r_tot = int(float(30))

    ############################################################################
    ##############################################################################################
    #### Temporal profile of glutamate availability:
    def glu(self, t, flag):
        """temporal profile of glutamate availability"""
        #print("in glu function")
        if flag == 0:
            return 0
        else:
            self.tau_glu = 1e-3 #sec
            self.glu_0 = 2.718 * 300 #um

            total = 0

            for tpuff in self.tpre:
                if t > tpuff: total += self.glu_0 * np.exp(-(t-tpuff)/self.tau_glu) * ((t-tpuff)/self.tau_glu)
            return total

        ##############################################################################################

        ##############################################################################################
        #### Voltage profile of BAP at the dendritic compartment:
        ##Use this to get the voltage of the spine whenever required
    def u_bpap(self, t):
        """generates bAPs (for STDP). Not used for RDP/theta-burst"""

        V0 = 67
        total = 0

        for tbp in tpost:
            if t > tbp: total += V0 * (0.7 * np.exp(-(t - tbp)/0.003) + 0.3 * np.exp(-(t - tbp)/0.04))
        return self.E_L + total

    ##############################################################################################
    #### AMPAR conductance profile:
    def I_A(self,flag,u,t):
        """use this function to get AMPA current"""
        #print("in AMPA current")

        if flag==0:
            return 0
        else:
            total = 0
            for tpuff in self.tpre:
                if t>tpuff: total += self.g_A * (np.exp(-(t-tpuff)/self.tau_A2) - np.exp(-(t-tpuff)/self.tau_A1))
            return total * (u - self.E_A)

    def I_N(self, flag, u, t):
        """use this function to get the NMDA receptor current"""
        #print("in NMDA current")

        if flag == 0:
            return 0

        else:
            total = 0
            for tpuff in self.tpre:
                if t > tpuff : total+= self.g_N * (np.exp(-(t - tpuff)/self.tau_N2) - np.exp(-(t - tpuff)/self.tau_N1))
            return total * (u - self.E_N)/(1 + 0.28 * np.exp(-0.062 * u))

#############################################################################################################33
############### SK Channel Hill function ###########################################
    def sk_inf(self, x):
        """SK-Hill function. args: cytosolic_calcium """
        return np.power(x, self.n_sk)/(np.power(x, self.n_sk) + np.power(self.k_sk, self.n_sk))

##############################################################################################
#### Plasticity model, Omega function

    def omega(self, x):
        """defines the Omega function to classify LTP/LTD based on induction var x"""
        U = -self.beta2 * (x - self.alpha2)
        V = -self.beta1 * (x - self.alpha1)

        if U > 100: U = 100
        if V > 100: V = 100
        return self.omega_0 + (self.eta2/(1. + np.exp(U))) - (self.eta1/(1 + np.exp(V)))

#############################################################################################
#### Plasticity model Omega function

    def omega_tau(self, x):
        """defines the tau function for synaptic weight change"""
        return self.P1 + (self.P2/(self.P3 + (2 * x /(self.alpha1 + self.alpha2))**self.P4))

#######################################################################
##### Coupled ODEs describing ER bearing spine head ###################
    def spine_model(self,x,t, vsp, ga):
        """ODE-based CA1 spine model"""

        pH, pL, cbp, Bslow, calb, calb_m1, calb_h1, calb_m2, calb_h2, calb_m1h1, calb_m2h1, calb_m1h2, c1n0, c2n0, c0n1, c0n2, c1n1, c2n1, c1n2, c2n2, R_Gq, Gact,IP3,ca_Gact_PLC_PIP2,DAGdegr,\
        PLC_PIP2,DAG,IP3_IP5P,IP3degr,glu_R_Gq,Gbc,ca_PLC,IP3_IP3K_2ca,R,ca_PLC_PIP2,IP3K_2ca,Gact_PLC_PIP2,Gq,IP5P,GaGDP,ca_Gact_PLC,glu_R,IP3K, mv, hv, w, u, ud, h,\
        ca_er, ca, R, c1, c2, c3, c4, c5, o1, o2, Bcer, Bper, p_socc = x
        ###########|----------- RyR ----------|   |-Buffer-|  |-socc-|##############

        ########################### changing AMPA conductance and volume with weights ####################################
        vsp *= w
        ga *= ga
        aspine = np.power(vsp * 1e15, float(2/3)) * np.power(36 * np.pi, float(1/3)) * 1e-8 ## cm^2
        #print(w)

        nt = self.glu(t,self.flag)
        #print(self.flag)

#         if (self.flag>0 and self.input_pattern=='stdp'):
#             ud = self.u_bpap(t)

        p_socc_inf = np.power(self.K_socc,self.n_socc)/(np.power(self.K_socc,self.n_socc) + np.power(ca_er, self.n_socc))
        dPdt = (p_socc_inf - p_socc)/self.tau_socc

        ## RyR Kinetics:#################################################
        Jrel = self.gRyR  * (o1 + o2) * (ca_er - ca)
        ######################### Channel Ca Release ########################################
        ca_eq = self.nRyR * Jrel * self.Vr + self.V_socc_max * self.socc_clus * p_socc
        ca_er_eq = - self.nRyR * Jrel  + (self.ca_er_0 - ca_er)/self.tau_refill
        ########################################################################################################################
        #print("ca: ryr {}".format(ca))
        ##########################################################################################################
        ################################### IP3 ###################################
        R_Gq_eq = -self.a2f*R_Gq*nt + self.a2b*glu_R_Gq + self.a3f*R*Gq-self.a3b*R_Gq
        Gact_eq = +self.a5*glu_R_Gq+self.a6*Gq - self.a7*Gact-self.b3f*Gact*PLC_PIP2 + self.b3b*Gact_PLC_PIP2 - self.b4f*Gact*ca_PLC_PIP2 + self.b4b*ca_Gact_PLC_PIP2 - self.b5f*ca_PLC*Gact + self.b5b*ca_Gact_PLC
        IP3_eq = +self.b6*ca_PLC_PIP2 + self.b7*ca_Gact_PLC_PIP2 - 100*IP3K_2ca*IP3 + 80*IP3_IP3K_2ca - 12.5*IP5P*IP3 + 72*IP3_IP5P ## change multiplier to IP5P*IP3 term to 9 for Gaurang's model
        ca_Gact_PLC_PIP2_eq = +self.b2f*ca*Gact_PLC_PIP2 - self.b2b*ca_Gact_PLC_PIP2 + self.b4f*Gact*ca_PLC_PIP2 - self.b4b*ca_Gact_PLC_PIP2 - self.b11*ca_Gact_PLC_PIP2 + self.b9f*ca_Gact_PLC*self.PIP2 - self.b9b*ca_Gact_PLC_PIP2 - self.b7*ca_Gact_PLC_PIP2
        DAGdegr_eq = +self.DAGdegrate*DAG
        PLC_PIP2_eq = -self.b1f*ca*PLC_PIP2 + self.b1b*ca_PLC_PIP2 - self.b3f*Gact*PLC_PIP2+self.b3b*Gact_PLC_PIP2 + self.b10*Gact_PLC_PIP2
        DAG_eq = +self.b6*ca_PLC_PIP2 + self.b7*ca_Gact_PLC_PIP2 - self.DAGdegrate*DAG
        IP3_IP5P_eq = +9*IP5P*IP3 - 72*IP3_IP5P - 18*IP3_IP5P
        IP3degr_eq = +20*IP3_IP3K_2ca + 18*IP3_IP5P
        glu_R_Gq_eq = +self.a2f*R_Gq*nt - self.a2b*glu_R_Gq + self.a4f*glu_R*Gq - self.a4b*glu_R_Gq-self.a5*glu_R_Gq
        Gbc_eq = +self.a5*glu_R_Gq + self.a6*Gq - self.a8*GaGDP*Gbc
        ca_PLC_eq = -self.b8f*ca_PLC*self.PIP2 + self.b8b*ca_PLC_PIP2 + self.b6*ca_PLC_PIP2 - self.b5f*ca_PLC*Gact + self.b5b*ca_Gact_PLC + self.b12*ca_Gact_PLC
        IP3_IP3K_2ca_eq = +100*IP3K_2ca*IP3-80*IP3_IP3K_2ca-20*IP3_IP3K_2ca
        R_eq = -self.a1f*R*nt + self.a1b*glu_R - self.a3f*R*Gq + self.a3b*R_Gq
        ca_PLC_PIP2_eq = +self.b1f*ca*PLC_PIP2 - self.b1b*ca_PLC_PIP2 - self.b4f*Gact*ca_PLC_PIP2 + self.b4b*ca_Gact_PLC_PIP2 + self.b11*ca_Gact_PLC_PIP2 + self.b8f*ca_PLC*self.PIP2 - self.b8b*ca_PLC_PIP2 - self.b6*ca_PLC_PIP2
        IP3K_2ca_eq = +1111*IP3K*ca*ca - 100*IP3K_2ca - 100*IP3K_2ca*IP3 + 80*IP3_IP3K_2ca + 20*IP3_IP3K_2ca
        Gact_PLC_PIP2_eq = -self.b2f*ca*Gact_PLC_PIP2 + self.b2b*ca_Gact_PLC_PIP2 + self.b3f*Gact*PLC_PIP2 - self.b3b*Gact_PLC_PIP2 - self.b10*Gact_PLC_PIP2
        Gq_eq = -self.a3f*R*Gq + self.a3b*R_Gq - self.a4f*glu_R*Gq + self.a4b*glu_R_Gq - self.a6*Gq + self.a8*GaGDP*Gbc
        IP5P_eq = -9*IP5P*IP3+72*IP3_IP5P+18*IP3_IP5P
        GaGDP_eq = +self.a7*Gact - self.a8*GaGDP*Gbc + self.b10*Gact_PLC_PIP2 + self.b11*ca_Gact_PLC_PIP2 + self.b12*ca_Gact_PLC
        ca_Gact_PLC_eq = -self.b9f*ca_Gact_PLC*self.PIP2 + self.b9b*ca_Gact_PLC_PIP2 + self.b7*ca_Gact_PLC_PIP2 + self.b5f*ca_PLC*Gact - self.b5b*ca_Gact_PLC - self.b12*ca_Gact_PLC
        glu_R_eq = +self.a1f*R*nt - self.a1b*glu_R - self.a4f*glu_R*Gq + self.a4b*glu_R_Gq + self.a5*glu_R_Gq
        IP3K_eq = -1111*IP3K*ca*ca+100*IP3K_2ca

        ca_eq += (-self.b1f*ca*PLC_PIP2 - self.b2f*ca*Gact_PLC_PIP2 - 1111*IP3K*ca*ca - 1111*IP3K*ca*ca + (self.b1b*ca_PLC_PIP2 + self.b2b*ca_Gact_PLC_PIP2 + 100*IP3K_2ca + 100*IP3K_2ca))

        ## IP3 receptor kinetics:

        x = IP3/(IP3 + self.d1)
        y = ca/(ca + self.d5)
        Q2 = self.Kinh
        h_eq = self.a2*(Q2 - (Q2+ca)*h)

        ca_eq += self.ip3r_tot * ((x*y*h)**3) * self.alpha_ip3r * (ca_er - ca)/(self.Nav * vsp)

        ca_er_eq += -self.alpha_ip3r * self.ip3r_tot * ((x*y*h)**3) * (ca_er - ca)/(self.Nav * vsp * 0.1) ## vsp * 0.1 = ver

        ## Buffer equations:

        Bslow_eq = -self.kslow_f*Bslow*ca + self.kslow_b*(self.Bslow_tot - Bslow)
        ca_eq += -self.kslow_f*Bslow*ca + self.kslow_b*(self.Bslow_tot - Bslow)

        cbp_eq = -self.kbuff_f*ca*cbp + self.kbuff_b*(self.cbp_tot - cbp)
        ca_eq += -self.kbuff_f*ca*cbp + self.kbuff_b*(self.cbp_tot - cbp)

        calb_m2h2 = self.calb_tot - calb - calb_m1 - calb_h1 - calb_m2 - calb_h2 - calb_m1h1 - calb_m2h1 - calb_m1h2
        calb_eqs = [ -ca*calb*(self.km0m1 + self.kh0h1) + self.km1m0*calb_m1 + self.kh1h0*calb_h1,\
                         ca*calb*self.km0m1 - self.km1m0*calb_m1 + calb_m2*self.km2m1 - ca*calb_m1*self.km1m2 + calb_m1h1*self.kh1h0 - ca*calb_m1*self.kh0h1,\
                         ca*calb*self.kh0h1 - self.kh1h0*calb_h1 + calb_h2*self.kh2h1 - ca*calb_h1*self.kh1h2 + calb_m1h1*self.km1m0 - ca*calb_h1*self.km0m1,\
                         ca*calb_m1*self.km1m2 - self.km2m1*calb_m2 + self.kh1h0*calb_m2h1 - ca*self.kh0h1*calb_m2,\
                         ca*calb_h1*self.kh1h2 - self.kh2h1*calb_h2 + self.km1m0*calb_m1h2 - ca*self.km0m1*calb_h2,\
                         ca*(calb_h1*self.km0m1 + calb_m1*self.kh0h1) - (self.km1m0+self.kh1h0)*calb_m1h1 - ca*calb_m1h1*(self.km1m2+self.kh1h2) + self.kh2h1*calb_m1h2 + self.km2m1*calb_m2h1,\
                         ca*self.km1m2*calb_m1h1 - self.km2m1*calb_m2h1 + self.kh2h1*calb_m2h2 - self.kh1h2*ca*calb_m2h1 + self.kh0h1*ca*calb_m2 - self.kh1h0*calb_m2h1,\
                         ca*self.kh1h2*calb_m1h1 - self.kh2h1*calb_m1h2 + self.km2m1*calb_m2h2 - self.km1m2*ca*calb_m1h2 + self.km0m1*ca*calb_h2 - self.km1m0*calb_m1h2 ]
        ca_eq += -ca*(self.km0m1*(calb+calb_h1+calb_h2) + self.kh0h1*(calb+calb_m1+calb_m2) + self.km1m2*(calb_m1+calb_m1h1+calb_m1h2) + self.kh1h2*(calb_h1+calb_m1h1+calb_m2h1))+\
                    self.km1m0*(calb_m1+calb_m1h1+calb_m1h2) + self.kh1h0*(calb_h1+calb_m1h1+calb_m2h1) + self.km2m1*(calb_m2+calb_m2h1+calb_m2h2) + self.kh2h1*(calb_h2+calb_m1h2+calb_m2h2)
        #print("ca: calb {}".format(ca))
        #print(calb)
        ## Ca2+/calmodulin kinetics:

        c0n0 = self.cam_tot - c1n0 - c2n0 - c0n1 - c0n2 - c1n1 - c2n1 - c1n2 - c2n2
        c1n0_eq = -(self.k2c_on*ca + self.k1c_off + self.k1n_on*ca)*c1n0 + self.k1c_on*ca*c0n0 + self.k2c_off*c2n0 + self.k1n_off*c1n1
        c2n0_eq = -(self.k2c_off + self.k1n_on*ca)*c2n0 + self.k2c_on*ca*c1n0 + self.k1n_off*c2n1
        c0n1_eq = -(self.k2n_on*ca + self.k1n_off + self.k1c_on*ca)*c0n1 + self.k1n_on*ca*c0n0 + self.k2n_off*c0n2 + self.k1c_off*c1n1
        c0n2_eq = -(self.k2n_off + self.k1c_on*ca)*c0n2 + self.k2n_on*ca*c0n1 + self.k1c_off*c1n2
        c1n1_eq = -(self.k2c_on*ca + self.k1c_off + self.k1n_off + self.k2n_on*ca)*c1n1 + self.k1c_on*ca*c0n1 + self.k1n_on*ca*c1n0 + self.k2c_off*c2n1 + self.k2n_off*c1n2
        c2n1_eq = -(self.k2c_off + self.k2n_on*ca)*c2n1 + self.k2c_on*ca*c1n1 + self.k2n_off*c2n2 + self.k1n_on*ca*c2n0 - self.k1n_off*c2n1
        c1n2_eq = -(self.k2n_off + self.k2c_on*ca)*c1n2 + self.k2n_on*ca*c1n1 + self.k2c_off*c2n2 + self.k1c_on*ca*c0n2 - self.k1c_off*c1n2
        c2n2_eq = -(self.k2c_off + self.k2n_off)*c2n2 + self.k2c_on*ca*c1n2 + self.k2n_on*ca*c2n1
        cam_eqs = [c1n0_eq, c2n0_eq, c0n1_eq, c0n2_eq, c1n1_eq, c2n1_eq, c1n2_eq, c2n2_eq]
        ca_eq += -ca*(self.k1c_on*(c0n0+c0n1+c0n2) + self.k1n_on*(c0n0+c1n0+c2n0) + self.k2c_on*(c1n0+c1n1+c1n2) + self.k2n_on*(c0n1+c1n1+c2n1)) + self.k1c_off*(c1n0+c1n1+c1n2) + self.k1n_off*(c0n1+c1n1+c2n1) \
        + self.k2c_off*(c2n0+c2n1+c2n2) + self.k2n_off*(c0n2+c1n2+c2n2)
        #print("ca: cam {}".format(ca))

        ############################ ER Buffer kinetics #############################################################
        if self.buff_flag:
            dBcer_dt = -self.kcerf*(ca_er * Bcer) + self.kcerb * (self.totBcer - Bcer)
            dBper_dt = -self.kpf * (ca_er * Bper) + self.kpb * (self.totBper - Bper)
            ca_er_eq += -self.kcerf*(ca_er * Bcer) + self.kcerb * (self.totBcer - Bcer) -self.kpf * (ca_er * Bper) + self.kpb * (self.totBper - Bper)
        else:
            dBcer_dt = 0
            dBper_dt = 0
            ca_er_eq += 0

        ## PMCA/NCX kinetics:

        ca_eq += pH*self.kH_leak - ca*pH*self.k1H + self.k2H*(self.pHtot - pH)  +  pL*self.kL_leak - ca*pL*self.k1L + self.k2L*(self.pLtot - pL)
        pH_eq = self.k3H*(self.pHtot - pH) - ca*pH*self.k1H + self.k2H*(self.pHtot - pH)
        pL_eq = self.k3L*(self.pLtot - pL) - ca*pL*self.k1L + self.k2L*(self.pLtot - pL)
        #print("ca: pmca {}".format(ca))

          ## SERCA kinetics:

        ca_eq += -self.Vmax_serca * ca**2/(self.Kd_serca**2 + ca**2) + self.k_erleak*(ca_er - ca)

        ## VGCC equatiosn:

        mv_eq = ((1.0/(1 + np.exp(-(u-self.um)/self.kmv))) - mv)/self.tau_mv
        hv_eq = ((1.0/(1 + np.exp(-(u-self.uh)/self.khv))) - hv)/self.tau_hv
        I_vgcc = -0.001 * self.Nav * 3.2e-19 * self.g_vgcc * (mv**2) * hv * 0.078 * u * (ca - self.ca_ext*np.exp(-0.078*u))/(1 - np.exp(-0.078*u))

        ## Spine and dendrite voltage eqns:

        sp_hh_eq = -(1/self.Cmem) * ( self.g_L*(u - self.E_L) + self.I_A(self.flag,u,t)/aspine + self.I_N(self.flag,u,t)/aspine - (self.gc/aspine)*(ud - u) - I_vgcc/aspine) #- self.g_sk * s_sk * (u - self.E_sk)
        dend_hh_eq = -(1/self.Cmem) * ( self.g_L*(ud - self.E_L) + self.rho_spines*self.gc*(ud - u))
        #print(self.I_A(self.flag,u,t)/self.Aspine)

        ## Ca2+ influx through NMDAR and VGCC:
        #print(u)

        ca_eq += -(self.g_N_Ca/vsp) * (self.I_N(self.flag,u,t)/(self.g_N*(u - self.E_N))) * 0.078 * u * (ca - self.ca_ext*np.exp(-0.078*u))/(1 - np.exp(-0.078*u)) \
                -(self.g_vgcc/vsp) * (mv**2) * hv * 0.078 * u * (ca - self.ca_ext*np.exp(-0.078*u))/(1 - np.exp(-0.078*u))
        #print("ca: nmda/vgcc {}".format(ca))

        ################################## RyR Dynamics #################################################
        ryr_eqs = [-R*ca*self.kRc1 + c1*self.kc1R, -(self.kc1c2*ca+ self.kc1R)*c1 + c2*self.kc2c1 + self.kRc1*ca*R, -(self.kc2c3*ca + self.kc2c1)*c2 + self.kc3c2*c3 + self.kc1c2*ca*c1,\
                    -(self.kc3c4*ca + self.kc3c2)*c3 + self.kc2c3*ca*c2 + self.kc4c3*c4, -(self.kc4o1 + self.kc4o2 + self.kc4c3)*c4 + self.kc3c4*ca*c3 + self.ko1c4*o1 + self.ko2c4*o2,\
                    -(self.kc5o1 + self.kc5o2 + self.kc5I)*c5 + self.ko1c5*o1 + self.kIc5*(1-R-c1-c2-c3-c4-c5-o1-o2) + self.ko2c5*o2, -(self.ko1c5 + self.ko1c4)*o1 + self.kc4o1*c4 + self.kc5o1*c5,\
                     -(self.ko2c4 + self.ko2c5)*o2 + self.kc4o2*c4 + self.kc5o2*c5]

        ##################################################################################################

        ########################## SK - channel equations ################################################
        #dskdt = (self.sk_inf(ca) - s_sk)/self.tau_sk

        ## Equation for plasticity variable w:
        acam = self.cam_tot - c0n0
        w_eq = (1.0/self.omega_tau(acam))*(self.omega(acam) - w)
        #print(ca)

        #print('') [dRdt, dc1dt, dc2dt, dc3dt, dc4dt, dc5dt, do1dt, do2dt]
        #print([pH_eq, pL_eq, cbp_eq, Bslow_eq] + calb_eqs + cam_eqs + [mv_eq, hv_eq] + [w_eq] + [sp_hh_eq, dend_hh_eq,\
               # ca_er_eq, ca_eq] + ryr_eqs + [dBcer_dt, dBper_dt] + [dPdt] + [dskdt])

        return [pH_eq, pL_eq, cbp_eq, Bslow_eq] + calb_eqs + cam_eqs + [R_Gq_eq,Gact_eq,\
                IP3_eq,ca_Gact_PLC_PIP2_eq,DAGdegr_eq,PLC_PIP2_eq,DAG_eq,IP3_IP5P_eq,IP3degr_eq,glu_R_Gq_eq,Gbc_eq,ca_PLC_eq,IP3_IP3K_2ca_eq,\
                                                                        R_eq,ca_PLC_PIP2_eq,IP3K_2ca_eq,Gact_PLC_PIP2_eq,Gq_eq,IP5P_eq,GaGDP_eq,ca_Gact_PLC_eq,glu_R_eq,IP3K_eq] + [mv_eq, hv_eq] + [w_eq] + [sp_hh_eq, dend_hh_eq, h_eq, ca_er_eq, ca_eq] + ryr_eqs + [dBcer_dt, dBper_dt] + [dPdt] #+ [dskdt]

    ################################################################################################################
    ############################################### Experiments ##################################################
    ################### simulate ER+ spine without inputs to cnverge to steady state #############
    def get_resting_params(self):
        """simulates the spine in absence of inputs and returns the resting state params"""
        self.flag = 0
        self.buff_flag = 1 ## set buff_flag for using ER buffer
        self.g_vgcc = 0
        self.input_pattern="rdp"

        if (self.nRyR == 0 and self.ip3r_tot == 0):
            self.buff_flag = 0
            self.Vmax_serca = 0
            self.k_erleak = 0
            self.V_socc_max = 0
        ##########################################################################################################
        ######################## Initializing all variables:######################################################
        ########################################################################################################

        pumps_init = [self.pHtot, self.pLtot]
        buff_init =  [self.cbp_tot, self.Bslow_tot] + [self.calb_tot,0,0,0,0,0,0,0]
        CaM_init = [0]*8
        vgcc_init = [0,1]
        w_init = [1]
        voltage_init = [self.E_L, self.E_L]
        #ip3r_init = [1]
        ca_init = [self.ca_er_0, self.ca_0]
        RyR_init = [1/8]*8
        calreticulin_init = [63840,135]#[self.totBcer, self.totBper]#
        SOCC_init = [0.001]
        sk_init = [0.001]
        mGluR_init = [0,0,0,0,0,0.8,0,0,0,0,0,0,0,0.3,0,0,0,1.0,self.ip5pconc,0,0,0,self.ip3kconc]
        ip3r_init = [1]

        xinit0 = mGluR_init + pumps_init + buff_init + CaM_init + vgcc_init + w_init + voltage_init + ip3r_init + ca_init + RyR_init + calreticulin_init + SOCC_init

        #print(xinit0)

        ################ solving #################################################
        t0 = np.arange(0., 500., 0.001)
        sol0 = odeint(self.spine_model, xinit0, t0, rtol=self.rtol, atol=self.atol, args=(self.Vspine, self.g_A))
        print("initial ER calcium = {}".format(sol0[-1,-13]))
        print("initial cyto calcium = {}".format(sol0[-1,-12]))

        return sol0[-1,:]

######################################### RDP ################################################################

    def do_rdp(self, f_input, n_inputs):
        """performs RDP
        parameters: (frequency of stimulation, no. of presynaptic inputs)"""
        if (self.nRyR == 0 and self.ip3r_tot == 0):
            self.buff_flag = 0
            self.Vmax_serca = 0
            self.k_erleak = 0
            self.V_socc_max = 0

        xinit = self.get_resting_params()
        self.flag = 1
        self.buff_flag = 1
        self.tpre = [i/float(f_input) for i in range(n_inputs)]
        t = np.linspace(min(self.tpre), max(self.tpre), 10000)
        sol = odeint(self.spine_model,xinit, t, args=(self.Vspine, self.g_A), rtol = self.rtol, atol = self.atol)
        ################ saving in a file ######################
        fname = "rdp_out_vinit{}_nryr{}_f{}.csv".format(round(self.Vspine * 1e15, 2), self.nRyR, f_input)
        return sol

    ########################################## theta-burst #################################################################
    def theta_burst(self, lam, mu, sd, n_turns):
        """performs theta-burst protocol: gives spine 100Hz stimulation. The no. of inputs for each burst are sampled from poisson distr with mean lam and the ISI's are sampled from a normal distribution of mean = mu and SD=sd
        params: lam: mean no. of inputs from poisson distr; mu, sd: mean and SD of normal distr, used to sample ISI; n_turns: no. of times the stim is applied"""
        seed(234534)

        if (self.nRyR == 0 and self.ip3r_tot == 0):
            self.buff_flag = 0
            self.Vmax_serca = 0
            self.k_erleak = 0
            self.V_socc_max = 0
            self.ip3r_tot = 0

        xinit = self.get_resting_params()
        self.flag = 1
        t_rest = normal(mu, sd, n_turns) ## gives resting time after each stim
        t_burst = np.cumsum(t_rest)
        t_burst = np.insert(t_burst, 0, 0) ## produce a spike at t = 0s
        t_burst = t_burst[:-1]
        nipArr = np.zeros(t_burst.shape[0]) ## array to store the no. of inputs during each burst

        print(t_rest)
        print(t_burst)

        self.flag = 1

        ###### simulating #####
        for i in range(int(t_rest.shape[0])):
            n_inputs = poisson(lam)
            nipArr[i] = n_inputs
            self.tpre = [t_burst[i] + (k/100.0) for k in range(int(n_inputs))]
            self.tpre = np.array(self.tpre)
            print(self.tpre)
            t = np.arange(np.amin(self.tpre), np.amax(self.tpre) + t_rest[i], 0.0001)
            print(t)
            sol = odeint(self.spine_model, xinit, t, args=(self.Vspine, self.g_A))
            xinit = sol[-1,:]
            sol = sol[::250,:]
            camTrace = [np.sum(sol[j, 13:20]) for j in range(sol.shape[0])]
            ################# saving in a file ###############
            ## File names:
            if self.nRyR: data = "real_erSpine_encoding_states_vinit{}.csv".format(round(self.Vspine * 1e15, 2))
            else: data = "real_NoErSpine_encoding_states_vinit{}.csv".format(round(self.Vspine * 1e15, 2))
            print(data)

            for x in range(sol.shape[0]):
                with open(data, "a+") as f:
                    f.write("{},{},{} \n".format(float(sol[x,-16]), float(sol[x, -12]), float(camTrace[x])))
                    f.close()
        return sol

reason: import numpy,from numpy,from scipy


text:
from __future__ import division, print_function, absolute_import, with_statement
import numpy as np
import scipy as sp
from scipy import constants
import math
import time
import os
import build
import sys
from lammps import lammps

def calculate_ionic_energy(atoms, energies=np.array([0,0,0]), system=build.NaCl_110_110, label='temp'):
    # This combines the coulombic interactions
    # and some short range repulsion potentiale.
    # Currently the classic 6-12 Lennard Jones
    # we are assuming that the values for atomic positions
    # are in angstroms and that the atomic charges are
    # in elemental charges.
    ################################################
    # First we see if a previous run has checkpointed and aborted.
    # If the atoms contained in a temp_atoms.npy file match the tail end of the
    # array of atoms we have exactly then the simulation must have had the same
    # parameters at creation. So we can skip the leading end of the array.

    test = os.path.isfile(label+'temp_atoms.npy')
    test = test and os.path.isfile(label+'temp_energies.npy')
    if test:
        atoms_from_file = np.load(label+'temp_atoms.npy')
        if np.array_equal(atoms_from_file, atoms[-len(atoms_from_file):]):
            energies = np.load(label+'temp_energies.npy')
            atoms = atoms_from_file

    total_energy = energies[0]
    electrostatic_energy = energies[0]
    repulsive_energy = energies[0]
    timer_start = time.time()

    for n, ref_atom in enumerate(atoms):
        # Here we define a reference atom that we
        # will calculate the interactions for.
        # Once we have defined the reference atom
        # as the first atom in th list we remove the
        # reference atom from the array.
        # This means we don't consider self interactions, and
        # avoid double counting. It also removes the need to deal
        # with dividing by zero in the coulombic calculation.
        ref_charge = ref_atom[3]
        ref_position = ref_atom[:3]
        remaining_atoms = atoms[n+1:]
        step_size = 5000000

        for m in range(0,remaining_atoms.shape[0],step_size):
            # We run through the atoms in
            # slices so we don't fill up memory
            atoms_slice = remaining_atoms[m:m+step_size]
            charges = atoms_slice[:,3]
            positions = atoms_slice[:,:3] # in angstroms
            displacements = positions - ref_position
            seps_squared = np.sum(np.square(displacements),axis=1)
            separations = seps_squared ** 0.5

            masked_seps_squared = seps_squared[separations < 320 * system['r0']]
            # Using only those bonds less than 320 time slarger than the smallest one will place
            # an upper bound on the computation required without affecting the results
            # at 320 times the shortest bond the LJ term is below the machine precision of a
            # 64 bit float.
            masked_seps_to_minus6 = np.power(masked_seps_squared, -3) # as scalars
            charge_products = charges * ref_charge
            masked_charge_products = charge_products[separations < 320 * system['r0']]

            electrostatic_energy += coulomb_energy(separations,
                                                   charge_products)
            repulsive_energy += LJ6_energy(masked_seps_to_minus6,
                                           masked_charge_products,
                                           ref_charge,
                                           system=system)
            masked_seps_to_minus12 = masked_seps_to_minus6**2
            repulsive_energy += LJ12_energy(masked_seps_to_minus12,
                                            masked_charge_products,
                                            ref_charge,
                                            system=system)
            total_energy = electrostatic_energy + repulsive_energy
        # If siginificant time has past then a checkpoint will be created
        # We have found the energy for each atom iterated over so far wrt
        # each other and to the atoms in remaining atoms. Hence saving
        # the energies and the remaining atoms provides a suitable restart point
        if time.time() - timer_start > 1000:
            energies = np.array([total_energy, electrostatic_energy, repulsive_energy])
            np.save(label+'temp_atoms.npy', remaining_atoms)
            np.save(label+'temp_energies.npy', energies)
            timer_start = time.time()

    # Return the energies all in absolute values of joules
    # delete the check point files on a succesful completion
    test = os.path.isfile(label+'temp_atoms.npy')
    test = test and os.path.isfile(label+'temp_energies.npy')
    if test:
        os.remove(label+'temp_atoms.npy')
        os.remove(label+'temp_energies.npy')
    return total_energy, electrostatic_energy, repulsive_energy #in Joules

# The Lennard Jones Parameters are included in the details of the system
def LJ6_energy(seps_to_minus6, charge_products, ref_charge, system=build.NaCl):
    # the classic lennard jones, the 6 term
    energy = 0
    # we always have the cation -anion interactions,
    # and the charge product for these is always -1
    # So mask the like-like interactions where qq > 0

    # we now have a matrix of distances with
    # only the reference cation to anion distances unmasked
    ca_energy = - system['B_ca'] * np.sum(seps_to_minus6[charge_products < 0])
    energy += ca_energy

    if ref_charge > 0: #i.e. a cation
        # cation-cation interaction qq > 0
        # mask where qq < 0
        cc_energy = - system['B_cc'] * np.sum(seps_to_minus6[charge_products>0])
        energy += cc_energy

    if ref_charge < 0: #i.e. an anion
        # anion-anion interaction qq > 0
        # mask where qq < 0
        aa_energy = - system['B_cc'] * np.sum(seps_to_minus6[charge_products>0])
        energy += aa_energy

    return energy #in joules

def LJ12_energy(seps_to_minus12, charge_products, ref_charge, system=build.NaCl):
    # The order 12 term of the Lennard-Jones
    energy = 0
    # we always have the cation-anion interactions,
    # and the charge product for these is always negative

    ca_energy = system['A_ca'] * np.sum(seps_to_minus12[charge_products < 0])
    energy += ca_energy

    if ref_charge > 0: #i.e. a cation
        # cation-cation interaction qq > 0
        # mask where qq < 0
        cc_energy = system['A_cc'] * np.sum(seps_to_minus12[charge_products>0])
        energy += cc_energy

    if ref_charge < 0: #i.e. an anion
        # anion-anion interaction qq > 0
        # mask where qq < 0
        aa_energy = system['A_cc'] * np.sum(seps_to_minus12[charge_products>0])
        energy += aa_energy

    return energy #in joules

def coulomb_energy(separations, charge_products):
    # This is a fairly general way of calculating
    # the coulombic energy of the interactions between
    # two sets of point charges.
    # distances is a matrix such that d_ij is the
    # distance between atom i and atom j
    # similarly charge_products_ij is the product of
    # the charge on atom i with the charge of atom j
    # The total energy function defined above is producing a
    # 1xn matrix, but this works for a general 2d matrix
    # watch out for divide by zero errors though

    energy = np.sum(charge_products / separations)
    coulomb_constant = 1/(4*constants.pi*constants.epsilon_0)
    units = coulomb_constant*(constants.e**2)/constants.angstrom
    # Note we're assuming that charges are elemenetal charges
    # and distances are in angstroms
    if energy:
        energy *= units
    else:
        energy = 0
    return energy # In Joules

def calc_lammps_coloumbic_energy(atoms, slip_system=build.NaCl_110_001):
    x_lo = np.min(atoms[:,0]) - 50
    x_hi = np.max(atoms[:,0]) + 50
    y_lo = np.min(atoms[:,1]) - 50
    y_hi = np.max(atoms[:,1]) + 50
    z_lo = np.min(atoms[:,2]) - slip_system['latt_params'][2]/4
    z_hi = np.max(atoms[:,2]) + slip_system['latt_params'][2]/4

    #build the data file for lamps.
    data_file_header = 'RPT26 on a LAMMPS adventure'
    data_file_header += '\n'
    data_file_header += '\n'
    data_file_header += '\n'
    data_file_header += '{:d} atoms\n\n'.format(len(atoms))
    data_file_header += (' 2 atom types\n')
    data_file_header += '{:} {:} xlo xhi\n'.format(x_lo, x_hi)
    data_file_header += '{:} {:} ylo yhi\n'.format(y_lo, y_hi)
    data_file_header += '{:} {:} zlo zhi\n'.format(z_lo, z_hi)
    data_file_header += '\n'
    data_file_header += '\n'
    data_file_header += 'Atoms\n'
    data_file_header += '\n'
    atoms_text = ''
    for i in range(len(atoms)):
        x = atoms[i,0]
        y = atoms[i,1]
        z = atoms[i,2]
        charge = atoms[i,-1]
        ID = i + 1
        if charge > 0:
            atom_type = 1
        elif charge < 0:
            atom_type = 2
        atoms_text += (' {:6} {: 2,g} {: 3,g} {: 9,g} {: 9,g} {: 9,g}\n'.format(ID, atom_type, charge, x, y, z))

    data_file_text = data_file_header + atoms_text
    with open('data.electro', 'wt') as file:
        file.write(data_file_text)

    lmp = lammps()
    lmp.file('in.electro')
    electro_energy = lmp.get_thermo('pe')
    lmp.close()

    return electro_energy

def calc_lammps_LJ_energy(atoms, slip_system=build.NaCl_110_001):
    x_lo = 2*np.min(atoms[:,0])
    x_hi = 2*np.max(atoms[:,0])
    y_lo = 2*np.min(atoms[:,1])
    y_hi = 2*np.max(atoms[:,1])
    z_lo = np.min(atoms[:,2]) - slip_system['latt_params'][2]/4
    z_hi = np.max(atoms[:,2]) + slip_system['latt_params'][2]/4

    #build the data file for lamps.
    data_file_header = 'RPT26 on a LAMMPS adventure'
    data_file_header += '\n'
    data_file_header += '\n'
    data_file_header += '\n'
    data_file_header += '{:d} atoms\n\n'.format(len(atoms))
    data_file_header += (' 2 atom types\n')
    data_file_header += '{:} {:} xlo xhi\n'.format(x_lo, x_hi)
    data_file_header += '{:} {:} ylo yhi\n'.format(y_lo, y_hi)
    data_file_header += '{:} {:} zlo zhi\n'.format(z_lo, z_hi)
    data_file_header += '\n'
    data_file_header += '\n'
    data_file_header += 'Atoms\n'
    data_file_header += '\n'
    atoms_text = ''
    for i in range(len(atoms)):
        x = atoms[i,0]
        y = atoms[i,1]
        z = atoms[i,2]
        charge = atoms[i,-1]
        ID = i + 1
        if charge > 0:
            atom_type = 1
        elif charge < 0:
            atom_type = 2
        atoms_text += (' {:6} {: 2,g} {: 3,g} {: 9,g} {: 9,g} {: 9,g}\n'.format(ID, atom_type, charge, x, y, z))

    data_file_text = data_file_header + atoms_text
    with open('data.LJ', 'wt') as file:
        file.write(data_file_text)

    lmp = lammps()
    lmp.file('in.LJ')

    LJ_energy = lmp.get_thermo('pe')
    lmp.close()

    return LJ_energy

def calc_lammps_energy(atoms, slip_system=build.NaCl_110_001):
    energy = calc_lammps_coloumbic_energy(atoms, slip_system=slip_system)
    energy += calc_lammps_LJ_energy(atoms, slip_system=slip_system)
    return energy

def scaled_disloc_energy(disloc_variables, sim_params):

    fn = sys.argv[1] + 'ref_energy.csv'
    if os.path.isfile(fn):
        ref_energy = np.loadtxt(fn)
    else:
        ref_energy = disloc_energy(np.ones_like(disloc_variables), sim_params)
        np.savetxt(fn, np.array([ref_energy]))

    energy = disloc_energy(disloc_variables, sim_params)
    scaled_energy = - math.fabs(energy / ref_energy)
    return scaled_energy

def disloc_energy(disloc_params, sim_params):
    disloc_energy = calc_disloc_energies(disloc_params, sim_params)[0]
    return disloc_energy

def calc_disloc_energies(disloc_params, sim_params):
    shape = sim_params[0]
    alpha = sim_params[1]
    slip_system = sim_params[2]

    full_disloc_params = np.zeros(4)
    full_disloc_params[:len(disloc_params)] = disloc_params

    # The build module expects 8 arguments however there are good
    # reasons to set some of these equal to each other.
    atoms = build.dislocation(shape=shape,
                              disloc_params=full_disloc_params,
                              alpha=alpha,
                              slip_system=slip_system)
    disloc_energies = e.calculate_ionic_energy(atoms, system=slip_system, label=sys.argv[1])
    return disloc_energies

ext: py

reason: import numpy,import scipy,from scipy


text:
#!/usr/bin/env python
# -*- coding:utf-8 -*-
# author:juzphy
# datetime:2020/2/29 6:05 下午
from config import config
import pandas as pd
import numpy as np
from tqdm import tqdm
from sklearn.metrics import classification_report, f1_score
from sklearn.model_selection import StratifiedKFold, KFold
from sklearn.decomposition import TruncatedSVD
from sklearn.feature_extraction.text import TfidfVectorizer, CountVectorizer
from gensim.models import Word2Vec
import lightgbm as lgb
import os
import warnings

warnings.filterwarnings('ignore')
os.environ['PYTHONHASHSEED'] = '0'

def hashfxn(astring):
    return ord(astring[0])

def geohash_encode(latitude, longitude, precision=12):
    """
    Encode a position given in float arguments latitude, longitude to
    a geohash which will have the character count precision.
    """
    lat_interval, lon_interval = (-90.0, 90.0), (-180.0, 180.0)
    base32 = '0123456789bcdefghjkmnpqrstuvwxyz'
    geohash = []
    bits = [16, 8, 4, 2, 1]
    bit = 0
    ch = 0
    even = True
    while len(geohash) < precision:
        if even:
            mid = (lon_interval[0] + lon_interval[1]) / 2
            if longitude > mid:
                ch |= bits[bit]
                lon_interval = (mid, lon_interval[1])
            else:
                lon_interval = (lon_interval[0], mid)
        else:
            mid = (lat_interval[0] + lat_interval[1]) / 2
            if latitude > mid:
                ch |= bits[bit]
                lat_interval = (mid, lat_interval[1])
            else:
                lat_interval = (lat_interval[0], mid)
        even = not even
        if bit < 4:
            bit += 1
        else:
            geohash += base32[ch]
            bit = 0
            ch = 0
    return ''.join(geohash)

def get_geohash_tfidf(df, group_id, group_target, num):
    df[group_target] = df.apply(
        lambda x: geohash_encode(x['lat'], x['lon'], 7),
        axis=1)
    tmp = df.groupby(group_id)[group_target].agg(list).reset_index()
    tmp[group_target] = tmp[group_target].apply(lambda x: ' '.join(x))

    tfidf_enc_tmp = TfidfVectorizer()
    tfidf_vec_tmp = tfidf_enc_tmp.fit_transform(tmp[group_target])
    svd_tag_tmp = TruncatedSVD(n_components=num, n_iter=20, random_state=1024)
    tag_svd_tmp = svd_tag_tmp.fit_transform(tfidf_vec_tmp)
    tag_svd_tmp = pd.DataFrame(tag_svd_tmp)
    tag_svd_tmp.columns = ['{}_tfidf_{}'.format(group_target, i)
                           for i in range(num)]

    countvec = CountVectorizer()
    count_vec_tmp = countvec.fit_transform(tmp[group_target])
    svd_tmp = TruncatedSVD(n_components=num, n_iter=20, random_state=1024)
    svd_tmp = svd_tmp.fit_transform(count_vec_tmp)
    svd_tmp = pd.DataFrame(svd_tmp)
    svd_tmp.columns = ['{}_countvec_{}'.format(group_target, i)
                       for i in range(num)]

    return pd.concat([tmp[[group_id]], tag_svd_tmp, svd_tmp], axis=1)

def get_grad_tfidf(df, group_id, group_target, num):
    grad_df = df.groupby(group_id)['lat'].apply(
        lambda x: np.gradient(x)).reset_index()
    grad_df['lon'] = df.groupby(group_id)['lon'].apply(
        lambda x: np.gradient(x)).reset_index()['lon']
    feat_list = [group_id]
    grad_df['lat'] = grad_df['lat'].apply(lambda x: np.round(x, 4))
    grad_df['lon'] = grad_df['lon'].apply(lambda x: np.round(x, 4))
    grad_df[group_target] = grad_df.apply(lambda x: ' '.join([
        '{}_{}'.format(z[0], z[1]) for z in zip(x['lat'], x['lon'])
    ]),
                                          axis=1)

    tfidf_enc_tmp = TfidfVectorizer()
    tfidf_vec_tmp = tfidf_enc_tmp.fit_transform(grad_df[group_target])
    svd_tag_tmp = TruncatedSVD(n_components=num, n_iter=20, random_state=1024)
    tag_svd_tmp = svd_tag_tmp.fit_transform(tfidf_vec_tmp)
    tag_svd_tmp = pd.DataFrame(tag_svd_tmp)
    tag_svd_tmp.columns = ['{}_tfidf_{}'.format(group_target, i)
                           for i in range(num)]
    return pd.concat([grad_df[feat_list], tag_svd_tmp], axis=1)

def get_sample_tfidf(df, group_id, group_target, num):
    tmp = df.groupby(group_id)['lat_lon'].apply(
        lambda x: x.sample(frac=0.1,
                           random_state=1)).reset_index()
    del tmp['level_1']
    tmp.columns = [group_id, group_target]
    tmp = tmp.groupby(group_id)[group_target].agg(list).reset_index()
    tmp[group_target] = tmp[group_target].apply(lambda x: ' '.join(x))

    tfidf_enc_tmp = TfidfVectorizer()
    tfidf_vec_tmp = tfidf_enc_tmp.fit_transform(tmp[group_target])
    svd_tag_tmp = TruncatedSVD(n_components=num, n_iter=20, random_state=1024)
    tag_svd_tmp = svd_tag_tmp.fit_transform(tfidf_vec_tmp)
    tag_svd_tmp = pd.DataFrame(tag_svd_tmp)
    tag_svd_tmp.columns = ['{}_tfidf_{}'.format(group_target, i)
                           for i in range(num)]

    return pd.concat([tmp[[group_id]], tag_svd_tmp], axis=1)

def w2v_feat(df, group_id, feat, length):
    print('start word2vec ...')
    data_frame = df.groupby(group_id)[feat].agg(list).reset_index()
    model = Word2Vec(data_frame[feat].values,
                     size=length,
                     window=5,
                     min_count=1,
                     workers=1,
                     iter=10,
                     seed=1,
                     hashfxn=hashfxn)
    data_frame[feat] = data_frame[feat].apply(lambda x: pd.DataFrame([
        model[c] for c in x
    ]))
    for m in range(length):
        data_frame['w2v_{}_mean'.format(m)] = data_frame[feat].apply(
            lambda x: x[m].mean())
    del data_frame[feat]
    return data_frame

def q10(x):
    return x.quantile(0.1)

def q20(x):
    return x.quantile(0.2)

def q30(x):
    return x.quantile(0.3)

def q40(x):
    return x.quantile(0.4)

def q60(x):
    return x.quantile(0.6)

def q70(x):
    return x.quantile(0.7)

def q80(x):
    return x.quantile(0.8)

def q90(x):
    return x.quantile(0.9)

def gen_feat(df):
    df.sort_values(['ID', 'time'], inplace=True)

    df['time'] = df['time'].apply(lambda x: '2019-' + x.split(
        ' ')[0][:2] + '-' + x.split(' ')[0][2:] + ' ' + x.split(' ')[1])
    df['time'] = pd.to_datetime(df['time'])

    df['lat_diff'] = df.groupby('ID')['lat'].diff(1)
    df['lon_diff'] = df.groupby('ID')['lon'].diff(1)
    df['speed_diff'] = df.groupby('ID')['speed'].diff(1)
    df['diff_minutes'] = df.groupby('ID')['time'].diff(1).dt.seconds // 60
    df['anchore'] = df.apply(lambda x: 1
                             if x['lat_diff'] < 0.01 and x['lon_diff'] < 0.01
                             and x['speed'] < 0.1 and x['diff_minutes'] <= 10
                             else 0,
                             axis=1)

    gt_zero = df[(df['lat_diff'] != 0) & (df['lat_diff'] != 0)]
    speed_gt_zero = df[df['speed_diff'] != 0]

    df['type'] = df['type'].map({'围网': 0, '刺网': 1, '拖网': 2, 'unknown': -1})
    group_df = df.groupby('ID')['type'].agg(
        {'label': 'mean',
         'cnt': 'count'}).reset_index()

    anchore_df = df.groupby('ID')['anchore'].agg('sum').reset_index()
    anchore_df.columns = ['ID', 'anchore_cnt']

    group_df = group_df.merge(anchore_df, on='ID', how='left')
    group_df['anchore_ratio'] = group_df['anchore_cnt'] / group_df['cnt']

    stat_functions = ['min', 'max', 'mean', 'median', 'nunique', q10, q20, q30,
                      q40, q60, q70, q80, q90]
    stat_ways = ['min', 'max', 'mean', 'median', 'nunique', 'q_10', 'q_20',
                 'q_30', 'q_40', 'q_60', 'q_70', 'q_80', 'q_90']

    stat_cols = ['lat', 'lon', 'speed', 'direction']
    group_tmp = df.groupby('ID')[stat_cols].agg(stat_functions).reset_index()
    group_tmp.columns = ['ID'] + ['{}_{}'.format(i, j)
                                  for i in stat_cols for j in stat_ways]

    gt_zero_group = gt_zero.groupby(
        'ID')[stat_cols].agg(stat_functions).reset_index()
    gt_zero_group.columns = ['ID'] + ['pos_gt_zero_{}_{}'.format(i, j)
                                      for i in stat_cols for j in stat_ways]

    speed_gt_zero_group = speed_gt_zero.groupby(
        'ID')[stat_cols].agg(stat_functions).reset_index()
    speed_gt_zero_group.columns = ['ID'] + ['speed_gt_zero_{}_{}'.format(
        i, j) for i in stat_cols for j in stat_ways]

    group_df = group_df.merge(group_tmp, on='ID', how='left')
    group_df = group_df.merge(gt_zero_group, on='ID', how='left')
    group_df = group_df.merge(speed_gt_zero_group, on='ID', how='left')

    mode_df = df.groupby(['ID', 'lat', 'lon'])['time'].agg(
        {'mode_cnt': 'count'}).reset_index()
    mode_df['rank'] = mode_df.groupby('ID')['mode_cnt'].rank(method='first',
                                                             ascending=False)
    for i in range(1, 4):
        tmp_df = mode_df[mode_df['rank'] == i]
        del tmp_df['rank']
        tmp_df.columns = ['ID', 'rank{}_mode_lat'.format(i),
                          'rank{}_mode_lon'.format(i),
                          'rank{}_mode_cnt'.format(i)]
        group_df = group_df.merge(tmp_df, on='ID', how='left')

    tfidf_df = get_geohash_tfidf(df, 'ID', 'lat_lon', 30)
    group_df = group_df.merge(tfidf_df, on='ID', how='left')
    print('geohash tfidf finished.')

    grad_tfidf = get_grad_tfidf(df, 'ID', 'grad', 30)
    group_df = group_df.merge(grad_tfidf, on='ID', how='left')
    print('gradient tfidf finished.')

    sample_tfidf = get_sample_tfidf(df, 'ID', 'sample', 30)
    group_df = group_df.merge(sample_tfidf, on='ID', how='left')
    print('sample tfidf finished.')

    w2v_df = w2v_feat(df, 'ID', 'lat_lon', 30)
    group_df = group_df.merge(w2v_df, on='ID', how='left')
    print('word2vec finished.')

    return group_df

def f1_score_eval(preds, valid_df):
    labels = valid_df.get_label()
    preds = np.argmax(preds.reshape(3, -1), axis=0)
    scores = f1_score(y_true=labels, y_pred=preds, average='macro')
    return 'f1_score', scores, True

def sub_on_line_lgb(train_, test_, pred, label, cate_cols, split,
                    is_shuffle=True,
                    use_cart=False,
                    get_prob=False):
    n_class = 3
    train_pred = np.zeros((train_.shape[0], n_class))
    test_pred = np.zeros((test_.shape[0], n_class))
    n_splits = 5

    assert split in ['kf', 'skf'
                    ], '{} Not Support this type of split way'.format(split)

    if split == 'kf':
        folds = KFold(n_splits=n_splits, shuffle=is_shuffle, random_state=1024)
        kf_way = folds.split(train_[pred])
    else:
        folds = StratifiedKFold(n_splits=n_splits,
                                shuffle=is_shuffle,
                                random_state=1024)
        kf_way = folds.split(train_[pred], train_[label])

    print('Use {} features ...'.format(len(pred)))

    params = {
        'learning_rate': 0.05,
        'boosting_type': 'gbdt',
        'objective': 'multiclass',
        'metric': 'None',
        'num_leaves': 63,
        'feature_fraction': 0.8,
        'bagging_fraction': 0.8,
        'bagging_freq': 5,
        'seed': 1,
        'bagging_seed': 1,
        'feature_fraction_seed': 7,
        'min_data_in_leaf': 20,
        'num_class': n_class,
        'nthread': 8,
        'verbose': -1,
    }
    for n_fold, (train_idx, valid_idx) in enumerate(kf_way, start=1):
        print('the {} training start ...'.format(n_fold))
        train_x, train_y = train_[pred].iloc[train_idx
                                            ], train_[label].iloc[train_idx]
        valid_x, valid_y = train_[pred].iloc[valid_idx
                                            ], train_[label].iloc[valid_idx]

        if use_cart:
            dtrain = lgb.Dataset(train_x,
                                 label=train_y,
                                 categorical_feature=cate_cols)
            dvalid = lgb.Dataset(valid_x,
                                 label=valid_y,
                                 categorical_feature=cate_cols)
        else:
            dtrain = lgb.Dataset(train_x, label=train_y)
            dvalid = lgb.Dataset(valid_x, label=valid_y)

        clf = lgb.train(params=params,
                        train_set=dtrain,
                        num_boost_round=3000,
                        valid_sets=[dvalid],
                        early_stopping_rounds=100,
                        verbose_eval=100,
                        feval=f1_score_eval)
        train_pred[valid_idx] = clf.predict(valid_x,
                                            num_iteration=clf.best_iteration)
        test_pred += clf.predict(test_[pred],
                                 num_iteration=clf.best_iteration) / folds.n_splits
    print(classification_report(train_[label], np.argmax(train_pred,
                                                         axis=1),
                                digits=4))
    if get_prob:
        sub_probs = ['qyxs_prob_{}'.format(q) for q in ['围网', '刺网', '拖网']]
        prob_df = pd.DataFrame(test_pred, columns=sub_probs)
        prob_df['ID'] = test_['ID'].values
        return prob_df
    else:
        test_['label'] = np.argmax(test_pred, axis=1)
        return test_[['ID', 'label']]

def get_data(file_path, model):
    assert model in ['train',
                     'test'], '{} Not Support this type of file'.format(model)
    paths = os.listdir(file_path)
    tmp = []
    for t in tqdm(range(len(paths))):
        p = paths[t]
        with open('{}/{}'.format(file_path, p), encoding='utf-8') as f:
            next(f)
            for line in f.readlines():
                tmp.append(line.strip().split(','))
    tmp_df = pd.DataFrame(tmp)
    if model == 'train':
        tmp_df.columns = ['ID', 'lat', 'lon', 'speed', 'direction', 'time',
                          'type']
    else:
        tmp_df['type'] = 'unknown'
        tmp_df.columns = ['ID', 'lat', 'lon', 'speed', 'direction', 'time',
                          'type']
    tmp_df['lat'] = tmp_df['lat'].astype(float)
    tmp_df['lon'] = tmp_df['lon'].astype(float)
    tmp_df['speed'] = tmp_df['speed'].astype(float)
    tmp_df['direction'] = tmp_df['direction'].astype(int)
    return tmp_df

if __name__ == "__main__":
    TRAIN_PATH = config.train_dir
    TEST_PATH = config.test_dir
    PROB_PATH = config.prob_qyxs
    use_prob = True
    train = get_data(TRAIN_PATH, 'train')
    test = get_data(TEST_PATH, 'test')
    train = train.append(test)
    all_df = gen_feat(train)
    del train, test
    use_train = all_df[all_df['label'] != -1]
    use_test = all_df[all_df['label'] == -1]
    use_feats = [c for c in use_train.columns if c not in ['ID', 'label']]
    sub = sub_on_line_lgb(use_train, use_test, use_feats, 'label', [], 'kf',
                          is_shuffle=True,
                          use_cart=False,
                          get_prob=use_prob)
    if use_prob:
        sub.to_csv(PROB_PATH, encoding='utf-8', index=False)
    else:
        sub['label'] = sub['label'].map({0: '围网', 1: '刺网', 2: '拖网'})
        sub.to_csv(config.save_path,
                   encoding='utf-8',
                   header=None,
                   index=False)

reason: import numpy


text:
# import numpy as np
# import PIL.Image

# from dream_utils import *

# resize_in = (224, 224)
# resize_out = (700, 700)

# Athena squared
# image = np.float32(PIL.Image.open('images/athena_louvre_700px.jpg'))
# image_mask = PIL.Image.open('images/athena_louvre_700px_face_mask.png')

# Louise
# image = np.float32(PIL.Image.open('images/louise.jpg'))
# image_mask = PIL.Image.open('images/louise_crop_mask.png')

TIME_FORMAT = "%Y-%m-%d_%H:%M:%S.%f"

SOLVERS = [
    # {'name': '0010', 'snapshot':  1,  'max_iter':   10, 'base_lr': 0.0001, 'test_interval': 10},  # 0
    {'snapshot':   1, 'max_iter':   20, 'base_lr': 0.0001, 'test_interval': 20},  # 1
    {'snapshot':   1, 'max_iter':   40, 'base_lr': 0.001,  'test_interval': 20},  # 2
    {'snapshot':   1, 'max_iter':   80, 'base_lr': 0.01,   'test_interval': 40},  # 3
    {'snapshot':   1, 'max_iter':  120, 'base_lr': 0.01,   'test_interval': 20},  # 4
    {'snapshot':   1, 'max_iter':  160, 'base_lr': 0.01,   'test_interval': 20},  # 5
    {'snapshot':   1, 'max_iter':  200, 'base_lr': 0.01,   'test_interval': 20},  # 6
    {'snapshot':   1, 'max_iter':  220, 'base_lr': 0.02,   'test_interval': 10},  # 7
    {'snapshot':   1, 'max_iter':  240, 'base_lr': 0.02,   'test_interval': 10},  # 8
    {'snapshot':   1, 'max_iter':  250, 'base_lr': 0.03,   'test_interval': 10},  # 9
    {'snapshot':   1, 'max_iter':  260, 'base_lr': 0.01,   'test_interval': 10},  # 10
    {'snapshot':   1, 'max_iter':  270, 'base_lr': 0.0005, 'test_interval': 10},  # 11
    {'snapshot':   1, 'max_iter':  290, 'base_lr': 0.001,  'test_interval': 10},  # 12
    {'snapshot':   5, 'max_iter':  340, 'base_lr': 0.01,   'test_interval': 10},  # 13
    {'snapshot':  20, 'max_iter':  480, 'base_lr': 0.01,   'test_interval': 100}, # 14
    {'snapshot':  20, 'max_iter':  680, 'base_lr': 0.01,   'test_interval': 100}, # 15
    {'snapshot':  20, 'max_iter':  880, 'base_lr': 0.01,   'test_interval': 100}, # 16
    {'snapshot':  20, 'max_iter': 1080, 'base_lr': 0.01,   'test_interval': 100}, # 17
    {'snapshot':  40, 'max_iter': 1200, 'base_lr': 0.01,   'test_interval': 100}, # 18
    {'snapshot':  40, 'max_iter': 1400, 'base_lr': 0.005,  'test_interval': 100}, # 19
    {'snapshot': 100, 'max_iter': 2200, 'base_lr': 0.005,  'test_interval': 200}, # 20
    {'snapshot': 100, 'max_iter': 3000, 'base_lr': 0.005,  'test_interval': 200}, # 21
    {'snapshot':   1, 'max_iter': 3020, 'base_lr': 0.001,  'test_interval': 20, 'gamma': 1.0}, # 22
    {'snapshot':   1, 'max_iter': 3040, 'base_lr': 0.01,   'test_interval': 20, 'gamma': 1.0}, # 23
    {'snapshot':   1, 'max_iter': 3060, 'base_lr': 0.02,   'test_interval': 20, 'gamma': 1.0}, # 24
    {'snapshot':   1, 'max_iter': 3080, 'base_lr': 0.015,  'test_interval': 20, 'gamma': 1.0}, # 25
    {'snapshot':   1, 'max_iter': 3100, 'base_lr': 0.01,   'test_interval': 20, 'gamma': 1.0}, # 26
    {'snapshot':   1, 'max_iter': 3150, 'base_lr': 0.005,  'test_interval': 25, 'gamma': 1.0}, # 27
    {'snapshot':   1, 'max_iter': 3170, 'base_lr': 0.005,  'test_interval': 20, 'gamma': 1.0}, # 28
    {'snapshot':   1, 'max_iter': 3200, 'base_lr': 0.005,  'test_interval': 20, 'gamma': 1.0}, # 29
    {'snapshot':  10, 'max_iter': 3300, 'base_lr': 0.001,  'test_interval': 20, 'gamma': 1.0}, # 30
    {'snapshot':  25, 'max_iter': 3500, 'base_lr': 0.0005, 'test_interval': 100, 'gamma': 1.0}, # 31
]
# default solver name equals to max_iter
for solver in SOLVERS:
    if not solver.has_key('name'):
        solver['name'] = '%s' % solver['max_iter']

RIA_MODEL_DIR = 'models/Ria_Gurtow/'
RIA_MODEL_SNAPSHOTS_PREFIX = 'models/Ria_Gurtow/generations/ria_gurtow_iter_'
EMOTIONS_MODEL = 'models/VGG_S_rgb/EmotiW_VGG_S.caffemodel'

SITUATIONS_FRAMES = 'situations/data/frames/'

max_stars_repo_name: Sitin/neuromatriarchy

max_issues_repo_name: Sitin/neuromatriarchy

max_forks_repo_name: Sitin/neuromatriarchy


text:
"""
Problem 3.

calculate the time series

yt = 5 + .05 * t + Et (Where E is epsilon)

for years 1960, 1961, ..., 2001 assuming Et independently and
identically distributed with mean 0 and sigma 0.2.
"""
from random import uniform

from matplotlib.pyplot import plot, show
from numpy import array, polyfit, poly1d

def create_distribution(size):
    """
    Create a distribution, identically distributed, with mean 0 and
    sigma 0.2
    """
    # Shit it's way easier to just do some uniform distribution
    # This is a bit over my head, and not possible for me without
    # pen and paper
    return array([uniform(-0.2, .2) for _ in xrange(size)])

def create_time_series(start_year, end_year):
    """
    Create the time series, yt, then perform a regress on yt, plot yt and the
    its trendline
    """
    t_array = array(range(start_year, end_year + 1))
    epsilon_t = create_distribution(len(t_array))
    yt = array([5 + .05 * t_i + epsilon_t[i] for i, t_i in enumerate(t_array)])
    fit = polyfit(t_array, yt, 1)
    fit_func = poly1d(fit)
    plot(t_array, yt, "yo", t_array, fit_func(t_array), "--k")
    show()

def main():
    create_time_series(1960, 2001)

if __name__ == "__main__":
    main()

max_stars_repo_name: hahnicity/ace

max_issues_repo_name: lnsongxf/ace

max_forks_repo_name: lnsongxf/ace

reason: from numpy


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

# Copyright 2018 Kyoto University (Hirofumi Inaguma)
#  Apache 2.0  (http://www.apache.org/licenses/LICENSE-2.0)

"""Criterions."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import math
import numpy as np
import torch
import torch.nn.functional as F

class MBR(torch.autograd.Function):
    """Minimum Bayes Risk (MBR) training.

    Args:
        vocab (int): number of nodes in softmax layer

    """
    @staticmethod
    def forward(ctx, log_probs, hyps, exp_risk, grad):
        """Forward pass.

        Args:
            log_probs (FloatTensor): `[N_best, L, vocab]`
            hyps (LongTensor): `[N_best, L]`
            exp_risk (FloatTensor): `[1]` (for forward)
            grad (FloatTensor): `[1]` (for backward)
        Returns:
            loss (FloatTensor): `[1]`

        """
        device_id = torch.cuda.device_of(log_probs).idx
        onehot = torch.eye(log_probs.size(-1)).cuda(device_id)[hyps]
        grads = grad * onehot  # mask out other classes
        log_probs = log_probs.requires_grad_()
        ctx.save_for_backward(log_probs, grads)
        return exp_risk

    @staticmethod
    def backward(ctx, grad_output):
        input, grads, = ctx.saved_tensors
        input.grad = grads
        return input, None, None, None

def cross_entropy_lsm(logits, ys, lsm_prob, ignore_index, training, normalize_length=False):
    """Compute cross entropy loss for label smoothing of sequence-to-sequence models.

    Args:
        logits (FloatTensor): `[B, T, vocab]`
        ys (LongTensor): Indices of labels. `[B, L]`
        lsm_prob (float): label smoothing probability
        ignore_index (int): index for padding
        normalize_length (bool): normalize XE loss by target sequence length
    Returns:
        loss_mean (FloatTensor): `[1]`
        ppl (float): perplexity

    """
    bs, _, vocab = logits.size()
    ys = ys.view(-1)
    logits = logits.view((-1, logits.size(2)))

    if lsm_prob == 0 or not training:
        loss = F.cross_entropy(logits, ys,
                               ignore_index=ignore_index, reduction='mean')
        ppl = np.exp(loss.item())
        if not normalize_length:
            loss *= (ys != ignore_index).sum() / bs
    else:
        with torch.no_grad():
            target_dist = logits.new_zeros(logits.size())
            target_dist.fill_(lsm_prob / (vocab - 1))
            mask = (ys == ignore_index)
            ys_masked = ys.masked_fill(mask, 0)
            target_dist.scatter_(1, ys_masked.unsqueeze(1), 1 - lsm_prob)

        log_probs = torch.log_softmax(logits, dim=-1)
        loss_sum = -torch.mul(target_dist, log_probs)
        n_tokens = len(ys) - mask.sum().item()
        denom = n_tokens if normalize_length else bs
        loss = loss_sum.masked_fill(mask.unsqueeze(1), 0).sum() / denom

        ppl = np.exp(loss.item()) if normalize_length else np.exp(loss.item() * bs / n_tokens)

    return loss, ppl

def distillation(logits_student, logits_teacher, ylens, temperature=5.0):
    """Compute cross entropy loss for knowledge distillation of sequence-to-sequence models.

    Args:
        logits_student (FloatTensor): `[B, T, vocab]`
        logits_teacher (FloatTensor): `[B, T, vocab]`
        ylens (IntTensor): `[B]`
        temperature (float):
    Returns:
        loss_mean (FloatTensor): `[1]`

    """
    bs, _, vocab = logits_student.size()

    log_probs_student = torch.log_softmax(logits_student, dim=-1)
    probs_teacher = torch.softmax(logits_teacher / temperature, dim=-1).data
    loss = -torch.mul(probs_teacher, log_probs_student)
    loss_mean = np.sum([loss[b, :ylens[b], :].sum() for b in range(bs)]) / ylens.sum()
    return loss_mean

def kldiv_lsm_ctc(logits, ylens):
    """Compute KL divergence loss for label smoothing of CTC and Transducer models.

    Args:
        logits (FloatTensor): `[B, T, vocab]`
        ylens (IntTensor): `[B]`
    Returns:
        loss_mean (FloatTensor): `[1]`

    """
    bs, _, vocab = logits.size()

    log_uniform = logits.new_zeros(logits.size()).fill_(math.log(1 / (vocab - 1)))
    probs = torch.softmax(logits, dim=-1)
    log_probs = torch.log_softmax(logits, dim=-1)
    loss = torch.mul(probs, log_probs - log_uniform)
    loss_mean = np.sum([loss[b, :ylens[b], :].sum() for b in range(bs)]) / ylens.sum()
    # assert loss_mean >= 0
    return loss_mean

def focal_loss(logits, ys, ylens, alpha, gamma):
    """Compute focal loss.

    Args:
        logits (FloatTensor): `[B, T, vocab]`
        ys (LongTensor): Indices of labels. `[B, L]`
        ylens (IntTensor): `[B]`
        alpha (float):
        gamma (float):
    Returns:
        loss_mean (FloatTensor): `[1]`

    """
    bs = ys.size(0)

    log_probs = torch.log_softmax(logits, dim=-1)
    probs_inv = -torch.softmax(logits, dim=-1) + 1
    loss = -alpha * torch.mul(torch.pow(probs_inv, gamma), log_probs)
    loss_mean = np.sum([loss[b, :ylens[b], :].sum() for b in range(bs)]) / ylens.sum()
    return loss_mean


text:
import numpy as np
from scipy.linalg import solve, eig
import mayavi.mlab as mlab
from project4d import *
from helpers import *
from js.geometry.rotations import *
from js.data.plyParse import PlyParse
from js.utils.plot.colors import colorScheme

colors = colorScheme("label")

cfgBunnyZipper = {"name":"bun_zipper", "lambdaS3": [60], "lambdaR3": 0.001}
cfg = cfgBunnyZipper
if cfg["name"] == "bun_zipper":
  scans = ['../data/bunny_rnd/bun_zipper.ply',
      '../data/bunny_rnd/bun_zipper_angle_90_translation_0.3.ply']
  gt = ['../data/bunny_rnd/bun_zipper_angle_90_translation_0.3.ply_TrueTransformation_angle_90_translation_0.3.csv']

scanApath = scans[0]
scanBpath = scans[1]
nameA = os.path.splitext(os.path.split(scanApath)[1])[0]
nameB = os.path.splitext(os.path.split(scanBpath)[1])[0]

plyA = PlyParse(); plyA.parse(scanApath); pcA = plyA.getPc()
plyB = PlyParse(); plyB.parse(scanBpath); pcB = plyB.getPc()

ranges = [-np.abs(pcA[:,0]).max()*2, np.abs(pcA[:,0]).max()*2,
    -np.abs(pcA[:,1]).max()*2, np.abs(pcA[:,1]).max()*2,
    -np.abs(pcA[:,2]).max()*2, np.abs(pcA[:,2]).max()*2]

#mlab.points3d(pcA[:,0], pcA[:,1], pcA[:,2], mode="point",
#    color=colors[0])
#mlab.view(azimuth=90, elevation=30, distance=0.6, roll=0, focalpoint=
#    np.array([ 0.0017717 ,  0.00829888, -0.00625299]))
#print mlab.view()
#mlab.show(stop=True)

path = "./bb_nodes_per_iteration_R3_bunnyFull.csv"
ts, tetras, props = LoadNodesPerIteration(path)

q_star = np.array([0.707107, 0.381209, -0.0583624, 0.592683])
q = Quaternion(q_star[0], q_star[1], q_star[2], q_star[3])

figm = mlab.figure(bgcolor=(1,1,1))
for n in range(len(ts)):
  # plot max UB rotation
  idMax = np.argmax(props[n][:,0])
  T = np.zeros((3,8))
  for j in range(8):
    T[:,j] = ts[n][tetras[n][idMax, j],:]
  t = T.sum(axis=1)/8.

  a = 2.*n/float(len(ts))*2.*np.pi
  Ry = np.array([[np.cos(a), 0, -np.sin(a)],
    [0,1,0],
    [np.sin(a),0, np.cos(a)]])

#  R = q.toRot().R.T
#  pcB_T = (1.001*R.dot(pcB.T)).T - R.dot(t)
#  pcA_T = pcA

  R = Ry.dot(q.toRot().R.T)
  pcB_T = (1.001*R.dot(pcB.T)).T - R.dot(t)
  pcA_T = (Ry.dot(pcA.T)).T

  mlab.points3d(pcA_T[:,0], pcA_T[:,1], pcA_T[:,2], mode="point",
      color=colors[0])
  mlab.points3d(pcB_T[:,0], pcB_T[:,1], pcB_T[:,2], mode="point",
        color=colors[1])
  mlab.view(azimuth=90,
      elevation=30, distance=0.9, roll=0, focalpoint=
      np.array([ 0.0017717 ,  0.00829888, -0.00625299]))
  print mlab.view()
  mlab.axes(extent=ranges)
  mlab.savefig("frame_pcR3_{:05}.png".format(n), figure=figm)
#  mlab.show(stop=True)
  mlab.clf(figm)

max_stars_repo_name: jstraub/dpOptTrans

max_issues_repo_name: jstraub/dpOptTrans

max_forks_repo_name: jstraub/dpOptTrans

reason: import numpy,from scipy


text:
"""Cacher for any matcher, which caches the output on disk in the top level folder `cache`.

This class provides the caching functionality to a GTSFM matcher. To use this cacher, initialize it with the matcher obj
you want to apply the cache on.

Example: To cache output of `TwoWayMatcher`, use `MatcherCacher(matcher_obj=TwoWayMatcher())`.

Authors: Ayush Baid
"""
from pathlib import Path
from typing import List, Optional, Tuple

import numpy as np

import gtsfm.utils.cache as cache_utils
import gtsfm.utils.io as io_utils
import gtsfm.utils.logger as logger_utils
from gtsfm.common.keypoints import Keypoints
from gtsfm.frontend.matcher.matcher_base import MatcherBase

logger = logger_utils.get_logger()

CACHE_ROOT_PATH = Path(__file__).resolve().parent.parent.parent.parent / "cache"
NUM_KEYPOINTS_TO_SAMPLE_FOR_HASH = 10

class MatcherCacher(MatcherBase):
    """Cacher for matcher output on disk, keyed on the input."""

    def __init__(self, matcher_obj: MatcherBase) -> None:
        super().__init__()
        self._matcher = matcher_obj
        # TODO(ayushbaid): make the obj cache key dependent on the code
        self._matcher_obj_key = type(self._matcher).__name__

    def __get_cache_path(self, cache_key: str) -> Path:
        """Gets the file path to the cache bz2 file from the cache key."""
        return CACHE_ROOT_PATH / "matcher" / "{}.pbz2".format(cache_key)

    def __generate_cache_key(
        self,
        keypoints_i1: Keypoints,
        keypoints_i2: Keypoints,
        descriptors_i1: np.ndarray,
        descriptors_i2: np.ndarray,
        im_shape_i1: Tuple[int, int],
        im_shape_i2: Tuple[int, int],
    ) -> str:
        """Generates the cache key from the input detections, image shapes, and underlying matcher.

        This function uses the first few keypoints and descriptors to generate a key, and hence is not guaranteed to
        be unique. However, since its unlikely that keypoints/descriptors exactly match in the selected few indices,
        and differ in the remaining, its a good choice for a key.
        """
        numpy_arrays_to_hash: List[np.ndarray] = []

        for keypoints_i, descriptors_i in zip([keypoints_i1, keypoints_i2], [descriptors_i1, descriptors_i2]):
            # subsample and concatenate keypoints and descriptors
            numpy_arrays_to_hash.append(keypoints_i.coordinates[:NUM_KEYPOINTS_TO_SAMPLE_FOR_HASH].flatten())
            if keypoints_i.responses is not None:
                numpy_arrays_to_hash.append(keypoints_i.responses[:NUM_KEYPOINTS_TO_SAMPLE_FOR_HASH].flatten())
            if keypoints_i.scales is not None:
                numpy_arrays_to_hash.append(keypoints_i.scales[:NUM_KEYPOINTS_TO_SAMPLE_FOR_HASH].flatten())
            numpy_arrays_to_hash.append(descriptors_i[:NUM_KEYPOINTS_TO_SAMPLE_FOR_HASH].flatten())

        # add the shapes as a numpy array
        h1, w1 = im_shape_i1
        h2, w2 = im_shape_i2
        numpy_arrays_to_hash.append(np.array([h1, w1, h2, w2]))

        # hash the concatenation of all the numpy arrays
        input_key = cache_utils.generate_hash_for_numpy_array(np.concatenate(numpy_arrays_to_hash))

        return "{}_{}".format(self._matcher_obj_key, input_key)

    def __load_result_from_cache(
        self,
        keypoints_i1: Keypoints,
        keypoints_i2: Keypoints,
        descriptors_i1: np.ndarray,
        descriptors_i2: np.ndarray,
        im_shape_i1: Tuple[int, int],
        im_shape_i2: Tuple[int, int],
    ) -> Optional[np.ndarray]:
        """Load cached result, if it exists. The cached result will be a 2D numpy array with 2 columns."""
        cache_path = self.__get_cache_path(
            cache_key=self.__generate_cache_key(
                keypoints_i1=keypoints_i1,
                keypoints_i2=keypoints_i2,
                descriptors_i1=descriptors_i1,
                descriptors_i2=descriptors_i2,
                im_shape_i1=im_shape_i1,
                im_shape_i2=im_shape_i2,
            )
        )
        return io_utils.read_from_bz2_file(cache_path)

    def __save_result_to_cache(
        self,
        keypoints_i1: Keypoints,
        keypoints_i2: Keypoints,
        descriptors_i1: np.ndarray,
        descriptors_i2: np.ndarray,
        im_shape_i1: Tuple[int, int],
        im_shape_i2: Tuple[int, int],
        match_indices: np.ndarray,
    ) -> None:
        """Save the results (match indice) to the cache."""
        cache_path = self.__get_cache_path(
            cache_key=self.__generate_cache_key(
                keypoints_i1=keypoints_i1,
                keypoints_i2=keypoints_i2,
                descriptors_i1=descriptors_i1,
                descriptors_i2=descriptors_i2,
                im_shape_i1=im_shape_i1,
                im_shape_i2=im_shape_i2,
            )
        )
        io_utils.write_to_bz2_file(match_indices, cache_path)

    def match(
        self,
        keypoints_i1: Keypoints,
        keypoints_i2: Keypoints,
        descriptors_i1: np.ndarray,
        descriptors_i2: np.ndarray,
        im_shape_i1: Tuple[int, int],
        im_shape_i2: Tuple[int, int],
    ) -> np.ndarray:
        """Match descriptor vectors.

        If the results are in the cache, they are fetched and returned. Otherwise, the `detect_and_describe` of the
        underlying object's API is called and the results are cached.

        # Some matcher implementations (such as SuperGlue) utilize keypoint coordinates as
        # positional encoding, so our matcher API provides them for optional use.

        Output format:
        1. Each row represents a match.
        2. First column represents keypoint index from image #i1.
        3. Second column represents keypoint index from image #i2.
        4. Matches are sorted in descending order of the confidence (score), if possible.

        Args:
            keypoints_i1: keypoints for image #i1, of length N1.
            keypoints_i2: keypoints for image #i2, of length N2.
            descriptors_i1: descriptors corr. to keypoints_i1.
            descriptors_i2: descriptors corr. to keypoints_i2.
            im_shape_i1: shape of image #i1, as (height,width).
            im_shape_i2: shape of image #i2, as (height,width).

        Returns:
            Match indices (sorted by confidence), as matrix of shape (N, 2), where N < min(N1, N2).
        """
        cached_data = self.__load_result_from_cache(
            keypoints_i1=keypoints_i1,
            keypoints_i2=keypoints_i2,
            descriptors_i1=descriptors_i1,
            descriptors_i2=descriptors_i2,
            im_shape_i1=im_shape_i1,
            im_shape_i2=im_shape_i2,
        )

        if cached_data is not None:
            return cached_data

        match_indices = self._matcher.match(
            keypoints_i1=keypoints_i1,
            keypoints_i2=keypoints_i2,
            descriptors_i1=descriptors_i1,
            descriptors_i2=descriptors_i2,
            im_shape_i1=im_shape_i1,
            im_shape_i2=im_shape_i2,
        )

        self.__save_result_to_cache(
            keypoints_i1=keypoints_i1,
            keypoints_i2=keypoints_i2,
            descriptors_i1=descriptors_i1,
            descriptors_i2=descriptors_i2,
            im_shape_i1=im_shape_i1,
            im_shape_i2=im_shape_i2,
            match_indices=match_indices,
        )

        return match_indices

max_stars_repo_name: swershrimpy/gtsfm

max_issues_repo_name: swershrimpy/gtsfm

max_forks_repo_name: swershrimpy/gtsfm


text:
# Copyright (c) 2015,2017 MetPy Developers.
# Distributed under the terms of the BSD 3-Clause License.
# SPDX-License-Identifier: BSD-3-Clause
r"""Module to provide unit support.

This makes use of the :mod:`pint` library and sets up the default settings
for good temperature support.

Attributes
----------
units : :class:`pint.UnitRegistry`
    The unit registry used throughout the package. Any use of units in MetPy should
    import this registry and use it to grab units.

"""

from __future__ import division

import functools
import logging
import warnings

import numpy as np
import pint
import pint.unit

log = logging.getLogger(__name__)

UndefinedUnitError = pint.UndefinedUnitError
DimensionalityError = pint.DimensionalityError

units = pint.UnitRegistry(autoconvert_offset_to_baseunit=True)

# For pint 0.6, this is the best way to define a dimensionless unit. See pint #185
units.define(pint.unit.UnitDefinition('percent', '%', (),
             pint.converters.ScaleConverter(0.01)))

# Define commonly encountered units not defined by pint
units.define('degrees_north = degree = degrees_N = degreesN = degree_north = degree_N '
             '= degreeN')
units.define('degrees_east = degree = degrees_E = degreesE = degree_east = degree_E = degreeE')

# Alias geopotential meters (gpm) to just meters
try:
    units._units['meter']._aliases = ('metre', 'gpm')
    units._units['gpm'] = units._units['meter']
except AttributeError:
    log.warning('Failed to add gpm alias to meters.')

# Silence UnitStrippedWarning
if hasattr(pint, 'UnitStrippedWarning'):
    warnings.simplefilter('ignore', category=pint.UnitStrippedWarning)

def pandas_dataframe_to_unit_arrays(df, column_units=None):
    """Attach units to data in pandas dataframes and return united arrays.

    Parameters
    ----------
    df : `pandas.DataFrame`
        Data in pandas dataframe.

    column_units : dict
        Dictionary of units to attach to columns of the dataframe. Overrides
        the units attribute if it is attached to the dataframe.

    Returns
    -------
        Dictionary containing united arrays with keys corresponding to the dataframe
        column names.

    """
    if not column_units:
        try:
            column_units = df.units
        except AttributeError:
            raise ValueError('No units attribute attached to pandas '
                             'dataframe and col_units not given.')

    # Iterate through columns attaching units if we have them, if not, don't touch it
    res = {}
    for column in df:
        if column in column_units and column_units[column]:
            res[column] = df[column].values * units(column_units[column])
        else:
            res[column] = df[column].values
    return res

def concatenate(arrs, axis=0):
    r"""Concatenate multiple values into a new unitized object.

    This is essentially a unit-aware version of `numpy.concatenate`. All items
    must be able to be converted to the same units. If an item has no units, it will be given
    those of the rest of the collection, without conversion. The first units found in the
    arguments is used as the final output units.

    Parameters
    ----------
    arrs : Sequence of arrays
        The items to be joined together

    axis : integer, optional
        The array axis along which to join the arrays. Defaults to 0 (the first dimension)

    Returns
    -------
    `pint.Quantity`
        New container with the value passed in and units corresponding to the first item.

    """
    dest = 'dimensionless'
    for a in arrs:
        if hasattr(a, 'units'):
            dest = a.units
            break

    data = []
    for a in arrs:
        if hasattr(a, 'to'):
            a = a.to(dest).magnitude
        data.append(np.atleast_1d(a))

    # Use masked array concatenate to ensure masks are preserved, but convert to an
    # array if there are no masked values.
    data = np.ma.concatenate(data, axis=axis)
    if not np.any(data.mask):
        data = np.asarray(data)

    return units.Quantity(data, dest)

def diff(x, **kwargs):
    """Calculate the n-th discrete difference along given axis.

    Wraps :func:`numpy.diff` to handle units.

    Parameters
    ----------
    x : array-like
        Input data
    n : int, optional
        The number of times values are differenced.
    axis : int, optional
        The axis along which the difference is taken, default is the last axis.

    Returns
    -------
    diff : ndarray
        The n-th differences. The shape of the output is the same as `a`
        except along `axis` where the dimension is smaller by `n`. The
        type of the output is the same as that of the input.

    See Also
    --------
    numpy.diff

    """
    if hasattr(x, 'units'):
        ret = np.diff(x.magnitude, **kwargs)
        # Can't just use units because of how things like temperature work
        it = x.flat
        true_units = (next(it) - next(it)).units
        return ret * true_units
    else:
        return np.diff(x, **kwargs)

def atleast_1d(*arrs):
    r"""Convert inputs to arrays with at least one dimension.

    Scalars are converted to 1-dimensional arrays, whilst other
    higher-dimensional inputs are preserved. This is a thin wrapper
    around `numpy.atleast_1d` to preserve units.

    Parameters
    ----------
    arrs : arbitrary positional arguments
        Input arrays to be converted if necessary

    Returns
    -------
    `pint.Quantity`
        A single quantity or a list of quantities, matching the number of inputs.

    """
    mags = [a.magnitude if hasattr(a, 'magnitude') else a for a in arrs]
    orig_units = [a.units if hasattr(a, 'units') else None for a in arrs]
    ret = np.atleast_1d(*mags)
    if len(mags) == 1:
        if orig_units[0] is not None:
            return units.Quantity(ret, orig_units[0])
        else:
            return ret
    return [units.Quantity(m, u) if u is not None else m for m, u in zip(ret, orig_units)]

def atleast_2d(*arrs):
    r"""Convert inputs to arrays with at least two dimensions.

    Scalars and 1-dimensional arrays are converted to 2-dimensional arrays,
    whilst other higher-dimensional inputs are preserved. This is a thin wrapper
    around `numpy.atleast_2d` to preserve units.

    Parameters
    ----------
    arrs : arbitrary positional arguments
        Input arrays to be converted if necessary

    Returns
    -------
    `pint.Quantity`
        A single quantity or a list of quantities, matching the number of inputs.

    """
    mags = [a.magnitude if hasattr(a, 'magnitude') else a for a in arrs]
    orig_units = [a.units if hasattr(a, 'units') else None for a in arrs]
    ret = np.atleast_2d(*mags)
    if len(mags) == 1:
        if orig_units[0] is not None:
            return units.Quantity(ret, orig_units[0])
        else:
            return ret
    return [units.Quantity(m, u) if u is not None else m for m, u in zip(ret, orig_units)]

def masked_array(data, data_units=None, **kwargs):
    """Create a :class:`numpy.ma.MaskedArray` with units attached.

    This is a thin wrapper around :func:`numpy.ma.masked_array` that ensures that
    units are properly attached to the result (otherwise units are silently lost). Units
    are taken from the ``data_units`` argument, or if this is ``None``, the units on ``data``
    are used.

    Parameters
    ----------
    data : array_like
        The source data. If ``data_units`` is `None`, this should be a `pint.Quantity` with
        the desired units.
    data_units : str or `pint.Unit`, optional
        The units for the resulting `pint.Quantity`
    kwargs
        Arbitrary keyword arguments passed to `numpy.ma.masked_array`, optional

    Returns
    -------
    `pint.Quantity`

    """
    if data_units is None:
        data_units = data.units
    return units.Quantity(np.ma.masked_array(data, **kwargs), data_units)

def _check_argument_units(args, dimensionality):
    """Yield arguments with improper dimensionality."""
    for arg, val in args.items():
        # Get the needed dimensionality (for printing) as well as cached, parsed version
        # for this argument.
        try:
            need, parsed = dimensionality[arg]
        except KeyError:
            # Argument did not have units specified in decorator
            continue

        # See if the value passed in is appropriate
        try:
            if val.dimensionality != parsed:
                yield arg, val.units, need
        # No dimensionality
        except AttributeError:
            # If this argument is dimensionless, don't worry
            if parsed != '':
                yield arg, 'none', need

def check_units(*units_by_pos, **units_by_name):
    """Create a decorator to check units of function arguments."""
    try:
        from inspect import signature

        def dec(func):
            # Match the signature of the function to the arguments given to the decorator
            sig = signature(func)
            bound_units = sig.bind_partial(*units_by_pos, **units_by_name)

            # Convert our specified dimensionality (e.g. "[pressure]") to one used by
            # pint directly (e.g. "[mass] / [length] / [time]**2). This is for both efficiency
            # reasons and to ensure that problems with the decorator are caught at import,
            # rather than runtime.
            dims = {name: (orig, units.get_dimensionality(orig.replace('dimensionless', '')))
                    for name, orig in bound_units.arguments.items()}

            @functools.wraps(func)
            def wrapper(*args, **kwargs):
                # Match all passed in value to their proper arguments so we can check units
                bound_args = sig.bind(*args, **kwargs)
                bad = list(_check_argument_units(bound_args.arguments, dims))

                # If there are any bad units, emit a proper error message making it clear
                # what went wrong.
                if bad:
                    msg = '`{0}` given arguments with incorrect units: {1}.'.format(
                        func.__name__,
                        ', '.join('`{}` requires "{}" but given "{}"'.format(arg, req, given)
                                  for arg, given, req in bad))
                    if 'none' in msg:
                        msg += ('\nAny variable `x` can be assigned a unit as follows:\n'
                                '    from metpy.units import units\n'
                                '    x = x * units.meter / units.second')
                    raise ValueError(msg)
                return func(*args, **kwargs)

            return wrapper

    # signature() only available on Python >= 3.3, so for 2.7 we just do nothing.
    except ImportError:
        def dec(func):
            return func

    return dec

try:
    # Try to enable pint's built-in support
    units.setup_matplotlib()
except (AttributeError, RuntimeError, ImportError):  # Pint's not available, try our own
    import matplotlib.units as munits

    class PintAxisInfo(munits.AxisInfo):
        """Support default axis and tick labeling and default limits."""

        def __init__(self, units):
            """Set the default label to the pretty-print of the unit."""
            super(PintAxisInfo, self).__init__(label='{:P}'.format(units))

    class PintConverter(munits.ConversionInterface):
        """Implement support for pint within matplotlib's unit conversion framework."""

        def __init__(self, registry):
            """Initialize converter for pint units."""
            super(PintConverter, self).__init__()
            self._reg = registry

        def convert(self, value, unit, axis):
            """Convert :`Quantity` instances for matplotlib to use."""
            if isinstance(value, (tuple, list)):
                return [self._convert_value(v, unit, axis) for v in value]
            else:
                return self._convert_value(value, unit, axis)

        def _convert_value(self, value, unit, axis):
            """Handle converting using attached unit or falling back to axis units."""
            if hasattr(value, 'units'):
                return value.to(unit).magnitude
            else:
                return self._reg.Quantity(value, axis.get_units()).to(unit).magnitude

        @staticmethod
        def axisinfo(unit, axis):
            """Return axis information for this particular unit."""
            return PintAxisInfo(unit)

        @staticmethod
        def default_units(x, axis):
            """Get the default unit to use for the given combination of unit and axis."""
            if isinstance(x, (tuple, list)):
                return getattr(x[0], 'units', 'dimensionless')
            else:
                return getattr(x, 'units', 'dimensionless')

    # Register the class
    munits.registry[units.Quantity] = PintConverter(units)
    del munits

del pint


text:
from numpy import size
from matplotlib import pyplot as plt
from PCD_WRITE import pcd_write, pcd_write_color, make_ply_file, make_obj_point
from Read_files import read_file, mat_to_np_arr, read_vtk2arr
from mpl_toolkits.mplot3d import axes3d
from Total import gen_3d_pts

image = read_file('images')             # load images
P = mat_to_np_arr('projective.mat')          # load projection matrices

PTS, colors = gen_3d_pts(image, P)             # generate 3D vertices with colors

# exporting file
length = size(PTS[0])

pcd_write_color(PTS, colors, length, 'points_3D_color.pcd')
make_obj_point('point_cloud.obj', PTS, colors)

# plotting figure
fig = plt.figure()
ax = fig.gca(projection='3d')
ax.plot(PTS[0], PTS[1], PTS[2], 'ko')
plt.axis('equal')
plt.show()


text:
import numpy as np

template = """
#!/bin/bash
#SBATCH -J scaling
#SBATCH -o scaling.o%j
#SBATCH -e scaling.e%j
#SBATCH -N {n_nodes}
#SBATCH -t 08:00:00
#SBATCH -p cca

source ~/.bash_profile

cd /mnt/ceph/users/apricewhelan/projects/thejoker-benchmarks/scripts

module load gcc openmpi2 lib/hdf5/1.10.1 intel/mkl/2019-3

conda activate hq

date

mpirun -n {pool_size} python speed_scaling.py --mpi

date
"""

for pool_size in 2 ** np.arange(2, 8+1, 1):
    size = pool_size + 1
    n_nodes = size // 28
    if size / 28 > 0:
        n_nodes += 1
    run_text = template.format(pool_size=size, n_nodes=n_nodes)

    with open('run-{0:d}.sh'.format(pool_size), 'w') as f:
        f.write(run_text)

max_stars_repo_name: adrn/thejoker-benchmarks

max_issues_repo_name: adrn/thejoker-benchmarks

max_forks_repo_name: adrn/thejoker-benchmarks