text:
#!/usr/bin/env pythonw
# -*- coding: utf-8 -*-
from __future__ import print_function
from builtins import str
from builtins import range
import wx
import sys
import os
import scipy
from scipy import *

#------------------------------------------------------------------------
#   def main():
#------------------------------------------------------------------------

"""
NAME
    tdt_magic.py.py

DESCRIPTION
    converts TDT formatted files to magic_measurements format files

SYNTAX
    tdt_magic.py -WD <PATH>

INPUT:
    TDT formatted files with suffix .tdt

OUTPUT:
    combined measurement file saved in <PATH>

Log:
    Initial revision 4/24/2014
    some bug fix 06/12/2015
"""

#===========================================
# GUI
#===========================================

class convert_tdt_files_to_MagIC(wx.Frame):
    """"""
    title = "Convert tdt files to MagIC format"

    def __init__(self,WD):
        wx.Frame.__init__(self, None, wx.ID_ANY, self.title)
        self.panel = wx.Panel(self)
        self.max_files=10

        os.chdir(WD)
        self.WD=os.getcwd()+"/"
        self.create_menu()
        self.InitUI()

    def InitUI(self):

        pnl = self.panel

        #---sizer infor ----

        TEXT1="Instructions:\n"
        TEXT2="1. Put all individual tdt files from the same location in one folder.\n"
        TEXT3="   Each tdt file file should end with '.tdt'\n"
        TEXT4="2. If there are more than one location use multiple folders. One folder for each location.\n"
        TEXT5="3. If the magnetization in in units are mA/m (as in the original TT program) volume is required to convert to moment.\n\n"
        TEXT6="For more information check the help menubar option.\n"

        TEXT7="(for support contact ron.shaar@mail.huji.ac.il)"

        TEXT=TEXT1+TEXT2+TEXT3+TEXT4+TEXT5+TEXT6+TEXT7
        bSizer_info = wx.StaticBoxSizer( wx.StaticBox( self.panel, wx.ID_ANY, "" ), wx.HORIZONTAL )
        bSizer_info.Add(wx.StaticText(pnl,label=TEXT),wx.ALIGN_LEFT)

        #---sizer 0 ----
        TEXT="output file:"
        bSizer0 = wx.StaticBoxSizer( wx.StaticBox( self.panel, wx.ID_ANY, "" ), wx.VERTICAL )
        bSizer0.Add(wx.StaticText(self.panel,label=TEXT),wx.ALIGN_LEFT)
        bSizer0.AddSpacer(5)
        self.output_file_path = wx.TextCtrl(self.panel, id=-1, size=(1000,25))
        #self.output_file_path.SetEditable(False)
        bSizer0.Add(self.output_file_path,wx.ALIGN_LEFT)
        self.output_file_path.SetValue(os.path.join(self.WD, "magic_measurements.txt"))
        #---sizer 1 ----
        TEXT="\n choose a path\n with no spaces in name"
        bSizer1 = wx.StaticBoxSizer( wx.StaticBox( self.panel, wx.ID_ANY, "" ), wx.VERTICAL )
        bSizer1.Add(wx.StaticText(pnl,label=TEXT),wx.ALIGN_TOP)
        bSizer1.AddSpacer(5)
        for i in range(self.max_files):
            command= "self.dir_path_%i = wx.TextCtrl(self.panel, id=-1, size=(100,25), style=wx.TE_READONLY)"%i
            exec(command)
            command= "self.add_dir_button_%i =  wx.Button(self.panel, id=-1, label='add',name='add_%i')"%(i,i)
            exec(command)
            command= "self.Bind(wx.EVT_BUTTON, self.on_add_dir_button_i, self.add_dir_button_%i)"%i
            #print command
            exec(command)
            command="bSizer1_%i = wx.BoxSizer(wx.HORIZONTAL)"%i
            exec(command)
            command="bSizer1_%i.Add(wx.StaticText(pnl,label=('%i  '[:2])),wx.ALIGN_LEFT)"%(i,i+1)
            exec(command)

            command="bSizer1_%i.Add(self.dir_path_%i,wx.ALIGN_LEFT)" %(i,i)
            exec(command)
            command="bSizer1_%i.Add(self.add_dir_button_%i,wx.ALIGN_LEFT)" %(i,i)
            exec(command)
            command="bSizer1.Add(bSizer1_%i,wx.ALIGN_TOP)" %i
            exec(command)
            bSizer1.AddSpacer(5)

        #---sizer 1a ----

        TEXT="\n\nexperiment:"
        bSizer1a = wx.StaticBoxSizer( wx.StaticBox( self.panel, wx.ID_ANY, "" ), wx.VERTICAL )
        bSizer1a.Add(wx.StaticText(pnl,label=TEXT),wx.ALIGN_TOP)
        self.experiments_names=['Thellier','ATRM 6 positions','NLT']
        bSizer1a.AddSpacer(5)
        for i in range(self.max_files):
            command="self.protocol_info_%i = wx.ComboBox(self.panel, -1, self.experiments_names[0], size=(100,25), choices=self.experiments_names, style=wx.CB_DROPDOWN|wx.CB_READONLY)"%i
            exec(command)
            command="bSizer1a.Add(self.protocol_info_%i,wx.ALIGN_TOP)"%i
            exec(command)
            bSizer1a.AddSpacer(5)

        #---sizer 1b ----

        TEXT="\nBlab direction\n dec, inc: "
        bSizer1b = wx.StaticBoxSizer( wx.StaticBox( self.panel, wx.ID_ANY, "" ), wx.VERTICAL )
        bSizer1b.Add(wx.StaticText(pnl,label=TEXT),wx.ALIGN_TOP)
        bSizer1b.AddSpacer(5)
        for i in range(self.max_files):
            #command= "self.file_info_Blab_%i = wx.TextCtrl(self.panel, id=-1, size=(40,25))"%i
            #exec command
            command= "self.file_info_Blab_dec_%i = wx.TextCtrl(self.panel, id=-1, size=(40,25))"%i
            exec(command)
            command= "self.file_info_Blab_dec_%i.SetValue('0')"%i
            exec(command)
            command= "self.file_info_Blab_inc_%i = wx.TextCtrl(self.panel, id=-1, size=(40,25))"%i
            exec(command)
            command= "self.file_info_Blab_inc_%i.SetValue('90')"%i
            exec(command)
            command="bSizer_blab%i = wx.BoxSizer(wx.HORIZONTAL)"%i
            exec(command)
            #command="bSizer_blab%i.Add(self.file_info_Blab_%i ,wx.ALIGN_LEFT)" %(i,i)
            #exec command
            command="bSizer_blab%i.Add(self.file_info_Blab_dec_%i,wx.ALIGN_LEFT)" %(i,i)
            exec(command)
            command="bSizer_blab%i.Add(self.file_info_Blab_inc_%i,wx.ALIGN_LEFT)" %(i,i)
            exec(command)
            command="bSizer1b.Add(bSizer_blab%i,wx.ALIGN_TOP)" %i
            exec(command)
            bSizer1b.AddSpacer(5)

        #---sizer 1c ----

        TEXT="\nmoment\nunits:"
        bSizer1c = wx.StaticBoxSizer( wx.StaticBox( self.panel, wx.ID_ANY, "" ), wx.VERTICAL )
        bSizer1c.Add(wx.StaticText(pnl,label=TEXT),wx.ALIGN_TOP)
        self.moment_units_names=['mA/m','emu','Am^2']
        bSizer1c.AddSpacer(5)
        for i in range(self.max_files):
            command="self.moment_units_%i = wx.ComboBox(self.panel, -1, self.moment_units_names[0], size=(80,25), choices=self.moment_units_names, style=wx.CB_DROPDOWN|wx.CB_READONLY)"%i
            exec(command)
            command="bSizer1c.Add(self.moment_units_%i,wx.ALIGN_TOP)"%i
            exec(command)
            bSizer1c.AddSpacer(5)

        #---sizer 1d ----

        TEXT="\nvolume\n[cubic m]:"
        bSizer1d = wx.StaticBoxSizer( wx.StaticBox( self.panel, wx.ID_ANY, "" ), wx.VERTICAL )
        bSizer1d.Add(wx.StaticText(pnl,label=TEXT),wx.ALIGN_TOP)
        bSizer1d.AddSpacer(5)
        for i in range(self.max_files):
            command= "self.sample_volume_%i = wx.TextCtrl(self.panel, id=-1, size=(80,25))"%i
            exec(command)
            command= "self.sample_volume_%i.SetValue('1.287555e-5')"%i
            exec(command)
            command="bSizer1d.Add(self.sample_volume_%i,wx.ALIGN_TOP)"%i
            exec(command)
            bSizer1d.AddSpacer(5)

        #---sizer 1e ----

        TEXT="\nuser\nname:"
        bSizer1e = wx.StaticBoxSizer( wx.StaticBox( self.panel, wx.ID_ANY, "" ), wx.VERTICAL )
        bSizer1e.Add(wx.StaticText(pnl,label=TEXT),wx.ALIGN_TOP)
        bSizer1e.AddSpacer(5)
        for i in range(self.max_files):
            command= "self.file_info_user_%i = wx.TextCtrl(self.panel, id=-1, size=(60,25))"%i
            exec(command)
            command="bSizer1e.Add(self.file_info_user_%i,wx.ALIGN_TOP)" %i
            exec(command)
            bSizer1e.AddSpacer(5)

        #---sizer 2 ----

        TEXT="\nlocation\nname:"
        bSizer2 = wx.StaticBoxSizer( wx.StaticBox( self.panel, wx.ID_ANY, "" ), wx.VERTICAL )
        bSizer2.Add(wx.StaticText(pnl,label=TEXT),wx.ALIGN_TOP)
        bSizer2.AddSpacer(5)
        for i in range(self.max_files):
            command= "self.file_location_%i = wx.TextCtrl(self.panel, id=-1, size=(60,25))"%i
            exec(command)
            command="bSizer2.Add(self.file_location_%i,wx.ALIGN_TOP)" %i
            exec(command)
            bSizer2.AddSpacer(5)

##        #---sizer 3 ----
##
##        missing

        #---sizer 4 ----

        TEXT="\nsample-specimen\nnaming convention:"
        bSizer4 = wx.StaticBoxSizer( wx.StaticBox( self.panel, wx.ID_ANY, "" ), wx.VERTICAL )
        bSizer4.Add(wx.StaticText(pnl,label=TEXT),wx.ALIGN_TOP)
        self.sample_naming_conventions=['sample=specimen','no. of terminate characters','charceter delimited']
        bSizer4.AddSpacer(5)
        for i in range(self.max_files):
            command="self.sample_naming_convention_%i = wx.ComboBox(self.panel, -1, self.sample_naming_conventions[0], size=(150,25), choices=self.sample_naming_conventions, style=wx.CB_DROPDOWN|wx.CB_READONLY)"%i
            exec(command)
            command="self.sample_naming_convention_char_%i = wx.TextCtrl(self.panel, id=-1, size=(40,25))"%i
            exec(command)
            command="bSizer4_%i = wx.BoxSizer(wx.HORIZONTAL)"%i
            exec(command)
            command="bSizer4_%i.Add(self.sample_naming_convention_%i,wx.ALIGN_LEFT)" %(i,i)
            exec(command)
            command="bSizer4_%i.Add(self.sample_naming_convention_char_%i,wx.ALIGN_LEFT)" %(i,i)
            exec(command)
            command="bSizer4.Add(bSizer4_%i,wx.ALIGN_TOP)"%i
            exec(command)

            bSizer4.AddSpacer(5)

        #---sizer 5 ----

        TEXT="\nsite-sample\nnaming convention:"
        bSizer5 = wx.StaticBoxSizer( wx.StaticBox( self.panel, wx.ID_ANY, "" ), wx.VERTICAL )
        bSizer5.Add(wx.StaticText(pnl,label=TEXT),wx.ALIGN_TOP)
        self.site_naming_conventions=['site=sample','no. of terminate characters','charceter delimited']
        bSizer5.AddSpacer(5)
        for i in range(self.max_files):
            command="self.site_naming_convention_char_%i = wx.TextCtrl(self.panel, id=-1, size=(40,25))"%i
            exec(command)
            command="self.site_naming_convention_%i = wx.ComboBox(self.panel, -1, self.site_naming_conventions[0], size=(150,25), choices=self.site_naming_conventions, style=wx.CB_DROPDOWN|wx.CB_READONLY)"%i
            exec(command)
            command="bSizer5_%i = wx.BoxSizer(wx.HORIZONTAL)"%i
            exec(command)
            command="bSizer5_%i.Add(self.site_naming_convention_%i,wx.ALIGN_LEFT)" %(i,i)
            exec(command)
            command="bSizer5_%i.Add(self.site_naming_convention_char_%i,wx.ALIGN_LEFT)" %(i,i)
            exec(command)
            command="bSizer5.Add(bSizer5_%i,wx.ALIGN_TOP)"%i
            exec(command)
            bSizer5.AddSpacer(5)

        #------------------

        self.okButton = wx.Button(self.panel, wx.ID_OK, "&OK")
        self.Bind(wx.EVT_BUTTON, self.on_okButton, self.okButton)

        self.cancelButton = wx.Button(self.panel, wx.ID_CANCEL, '&Cancel')
        self.Bind(wx.EVT_BUTTON, self.on_cancelButton, self.cancelButton)

        hbox1 = wx.BoxSizer(wx.HORIZONTAL)
        #hbox1.Add(self.add_file_button)
        #hbox1.Add(self.remove_file_button )

        hbox2 = wx.BoxSizer(wx.HORIZONTAL)
        hbox2.Add(self.okButton)
        hbox2.Add(self.cancelButton )

        #------

        vbox=wx.BoxSizer(wx.VERTICAL)

        hbox = wx.BoxSizer(wx.HORIZONTAL)
        hbox.AddSpacer(1)
        hbox.Add(bSizer1, flag=wx.ALIGN_LEFT)
        hbox.AddSpacer(1)
        hbox.Add(bSizer1a, flag=wx.ALIGN_LEFT)
        hbox.AddSpacer(1)
        hbox.Add(bSizer1b, flag=wx.ALIGN_LEFT)
        hbox.AddSpacer(1)
        hbox.Add(bSizer1c, flag=wx.ALIGN_LEFT)
        hbox.AddSpacer(1)
        hbox.Add(bSizer1d, flag=wx.ALIGN_LEFT)
        hbox.AddSpacer(1)
        hbox.Add(bSizer1e, flag=wx.ALIGN_LEFT)
        hbox.AddSpacer(1)
        hbox.Add(bSizer2, flag=wx.ALIGN_LEFT)
        hbox.AddSpacer(1)
##        hbox.Add(bSizer3, flag=wx.ALIGN_LEFT)
##        hbox.AddSpacer(5)
        hbox.Add(bSizer4, flag=wx.ALIGN_LEFT)
        hbox.AddSpacer(1)
        hbox.Add(bSizer5, flag=wx.ALIGN_LEFT)
        hbox.AddSpacer(1)

        #-----

        vbox.AddSpacer(5)
        vbox.Add(bSizer_info,flag=wx.ALIGN_CENTER_HORIZONTAL)
        vbox.AddSpacer(2)
        vbox.Add(hbox)
        vbox.AddSpacer(5)
        vbox.Add(hbox1,flag=wx.ALIGN_CENTER_HORIZONTAL)
        #vbox.AddSpacer(20)
        vbox.AddSpacer(5)
        vbox.Add(bSizer0, flag=wx.ALIGN_CENTER_HORIZONTAL)
        vbox.Add(hbox2,flag=wx.ALIGN_CENTER_HORIZONTAL)
        vbox.AddSpacer(5)

        self.panel.SetSizer(vbox)
        vbox.Fit(self)
        self.Show()
        self.Centre()

    def create_menu(self):
        """ Create menu
        """
        self.menubar = wx.MenuBar()

        menu_about = wx.Menu()
        menu_help = menu_about.Append(-1, "&Some notes", "")
        self.Bind(wx.EVT_MENU, self.on_menu_help, menu_help)

        self.menubar.Append(menu_about, "& Instructions")

        self.SetMenuBar(self.menubar)

    def on_menu_help (self,event):

        dia = message_box("Help")
        dia.Show()
        dia.Center()

    def on_add_dir_button_i(self,event):

        dlg = wx.DirDialog(
            None,message="choose directtory with tdt files",
            defaultPath ="./",
            style=wx.FD_OPEN | wx.FD_CHANGE_DIR
            )
        if dlg.ShowModal() == wx.ID_OK:
            FILE = dlg.GetPath()
        # fin=open(FILE,'r')
        button = event.GetEventObject()
        name=button.GetName()
        i=int((name).split("_")[-1])
        #print "The button's name is " + button.GetName()

        command="self.dir_path_%i.SetValue(FILE)"%i
        exec(command)

    def read_generic_file(self,path):
        Data={}
        Fin=open(path,'r')
        header=Fin.readline().strip('\n').split('\t')

        for line in Fin.readlines():
            tmp_data={}
            l=line.strip('\n').split('\t')
            if len(l)<len(header):
                continue
            else:
                for i in range(len(header)):
                    tmp_data[header[i]]=l[i]
                specimen=tmp_data['Specimen']
                if specimen not in list(Data.keys()):
                    Data[specimen]=[]
                # check dupliactes
                if len(Data[specimen]) >0:
                    if tmp_data['Treatment (aka field)']==Data[specimen][-1]['Treatment (aka field)']:
                        print("-W- WARNING: duplicate measurements specimen %s, Treatment %s. keeping onlt the last one"%(tmp_data['Specimen'],tmp_data['Treatment (aka field)']))
                        Data[specimen].pop()

                Data[specimen].append(tmp_data)
        return(Data)

    def on_okButton(self,event):

        DIRS_data={}

        for i in range(self.max_files):

            # read directiory path
            dirpath=""
            command="dirpath=self.dir_path_%i.GetValue()"%i
            exec(command)
            if dirpath!="":
                dir_name=str(dirpath.split("/")[-1])
                DIRS_data[dir_name]={}
                DIRS_data[dir_name]['path']=str(dirpath)
            else:
                continue

            # get experiment
            command="experiment=self.protocol_info_%i.GetValue()"%i
            exec(command)
            DIRS_data[dir_name]['experiment']=str(experiment)

            # get location
            user_name=""
            command="location_name=self.file_location_%i.GetValue()"%i
            exec(command)
            DIRS_data[dir_name]['er_location_name']=str(location_name)

            # get Blab direction
            labfield_DI=["0.","90."]
            command="labfield_DI[0]=self.file_info_Blab_dec_%i.GetValue()"%i
            exec(command)
            command="labfield_DI[1]=self.file_info_Blab_inc_%i.GetValue()"%i
            exec(command)
            DIRS_data[dir_name]['labfield_DI']=labfield_DI

            # get Moment units
            command="moment_units=self.moment_units_%i.GetValue()"%i
            exec(command)
            DIRS_data[dir_name]['moment_units']=moment_units

            # get sample volume
            command="sample_volume=self.sample_volume_%i.GetValue()"%i
            exec(command)
            DIRS_data[dir_name]['sample_volume']=sample_volume

            # get User_name
            user_name=""
            command="user_name=self.file_info_user_%i.GetValue()"%i
            exec(command)
            DIRS_data[dir_name]['user_name']=user_name

            # get sample-specimen naming convention

            sample_naming_convenstion=["",""]
            command="sample_naming_convenstion[0]=str(self.sample_naming_convention_%i.GetValue())"%i
            exec(command)
            command="sample_naming_convenstion[1]=str(self.sample_naming_convention_char_%i.GetValue())"%i
            exec(command)
            DIRS_data[dir_name]["sample_naming_convenstion"]=sample_naming_convenstion

            # get site-sample naming convention

            site_naming_convenstion=["",""]
            command="site_naming_convenstion[0]=str(self.site_naming_convention_%i.GetValue())"%i
            exec(command)
            command="site_naming_convenstion[1]=str(self.site_naming_convention_char_%i.GetValue())"%i
            exec(command)
            DIRS_data[dir_name]["site_naming_convenstion"]=site_naming_convenstion

        #print "DIRS_data",DIRS_data
        self.convert_2_magic(DIRS_data)

    def on_cancelButton(self,event):
        self.Destroy()

    def get_sample_name(self,specimen,sample_naming_convenstion):
        if sample_naming_convenstion[0]=="sample=specimen":
            sample=specimen
        elif sample_naming_convenstion[0]=="no. of terminate characters":
            n=int(sample_naming_convenstion[1])*-1
            sample=specimen[:n]
        elif sample_naming_convenstion[0]=="charceter delimited":
            d=sample_naming_convenstion[1]
            sample_splitted=specimen.split(d)
            if len(sample_splitted)==1:
                sample=sample_splitted[0]
            else:
                sample=d.join(sample_splitted[:-1])
        return sample

    def get_site_name(self,sample,site_naming_convenstion):
        if site_naming_convenstion[0]=="site=sample":
            site=sample
        elif site_naming_convenstion[0]=="no. of terminate characters":
            n=int(site_naming_convenstion[1])*-1
            site=sample[:n]
        elif site_naming_convenstion[0]=="charceter delimited":
            d=site_naming_convenstion[1]
            site_splitted=sample.split(d)
            if len(site_splitted)==1:
                site=site_splitted[0]
            else:
                site=d.join(site_splitted[:-1])
        return site

    #===========================================
    # Convert to MagIC format
    #===========================================

    def convert_2_magic(self,DIRS_data):
        #--------------------------------------
        # Read the files
        #
        # Database structure
        # Thellier_type experiment:
        #
        # 1) Each file contains the data one specimen
        # 2) First line is the header: "Thellier-tdt"
        # 3) Second line in header inlucdes 4 fields:
        #    [Blab] ,[unknown_1] , [unknown_2] , [unknown_3] , [unknown_4]
        # 4) Body includes 5 fields
        #    [specimen_name], [treatments], [moment],[meas_dec],[meas_dec
        # Tretment: XXX.0 (zerofield)
        #           XXX.1 (infield)
        #           XXX.2 (pTRM check)
        #           XXX.3 (Tail check)
        #           XXX.4 (Additivity check; Krasa et al., 2003)
        #           XXX.5 (Original Thellier-Thellier protocol. )
        #                 (where .5 is for the second direction and .1 in the first)
        # XXX = temperature in degrees
        #
        #
        # IMPORTANT ASSUMPTION:
        # (1) lab field is always in Z direction (theta=0, phi=90)
        # (2) Thermal demagnetization - NO MICROWAVE
        # (3) if if XXX <50 then assuming that this is NRM (273K)
        #
        # -------------------------------------
        #
        #   ATRM in six positions
        #
        # Tretment: XXX.0 zerofield
        #           XXX.1 +x
        #           XXX.2 +y
        #           XXX.3 +z
        #           XXX.4 -x
        #           XXX.5 -y
        #           XXX.6 -z
        #           XXX.7 alteration check
        #   IMPORTANT REMARKS:
        #
        # (1) If the program check if the direction of the magnetization fits the coding above
        # if not, an error message will appear
        # (2) Alteration ckeck can be in any direction
        # (3) the order of the measurements is not important
        #
        # For questions and support: rshaar@ucsd.edu
        # -------------------------------------------------------------

        magic_measurements_headers=[]
        er_specimens_headers=[]
        MagRecs=[]
        ErRecs=[]
        Data={}

        for dir_name in list(DIRS_data.keys()):

            #-----------------------------------
            # First, read all files and sort data by specimen and by Experiment type
            #-----------------------------------

            for files in os.listdir(DIRS_data[dir_name]["path"]):
                if files.endswith(".tdt"):
                    print("Open file: ", DIRS_data[dir_name]["path"]+"/"+files)
                    fin=open(DIRS_data[dir_name]["path"]+"/"+files,'r')
                    header_codes=['labfield','core_azimuth','core_plunge','bedding_dip_direction','bedding_dip']
                    body_codes=['specimen_name','treatment','moment','dec','inc']
                    tmp_body=[]
                    tmp_header_data={}
                    line_number=0
                    continue_reading=True
                    line=fin.readline() # ignore first line
                    for line in fin.readlines():

                        if "END" in line:
                            break

                        if line.strip('\n') =="":
                            break

                        this_line=line.strip('\n').split()

                        if len(this_line)<5:
                            continue

                        #---------------------------------------------------
                        # fix muxworthy funky data format
                        #---------------------------------------------------
                        if len(this_line)<5 and line_number!=0:
                            new_line=[]
                            for i in range(len(this_line)):
                                if i>1 and "-" in this_line[i]:
                                    tmp=this_line[i].replace("-"," -")
                                    tmp1=tmp.split()
                                    for i in range(len(tmp1)):
                                        new_line.append(tmp1[i])
                                else:
                                    new_line.append(this_line[i])
                            this_line=list(copy(new_line))

                        #-------------------------------
                        # Read infromation from Header and body
                        # The data is stored in a dictionary:
                        # Data[specimen][Experiment_Type]['header_data']=tmp_header_data  --> a dictionary with header data
                        # Data[specimen][Experiment_Type]['meas_data']=[dict1, dict2, ...] --> a list of dictionaries with measurement data
                        #-------------------------------

                        #---------------------------------------------------
                        # header
                        #---------------------------------------------------
                        if  line_number==0:

                            for i in range(len(this_line)):
                                tmp_header_data[header_codes[i]]=this_line[i]

                            line_number+=1

                        #---------------------------------------------------
                        # body
                        #---------------------------------------------------

                        else:
                            tmp_data={}
                            for i in range(min(len(this_line),len(body_codes))):
                                tmp_data[body_codes[i]]=this_line[i]
                            tmp_body.append(tmp_data)

                            #------------

                            specimen=tmp_body[0]['specimen_name']
                            line_number+=1

                    if specimen not in list(Data.keys()):
                        Data[specimen]={}
                    Experiment_Type=DIRS_data[dir_name]['experiment']
                    if Experiment_Type not in list(Data[specimen].keys()):
                        Data[specimen][Experiment_Type]={}
                    Data[specimen][Experiment_Type]['meas_data']=tmp_body
                    Data[specimen][Experiment_Type]['header_data']=tmp_header_data
                    Data[specimen][Experiment_Type]['sample_naming_convenstion']=DIRS_data[dir_name]['sample_naming_convenstion']
                    Data[specimen][Experiment_Type]['site_naming_convenstion']=DIRS_data[dir_name]['site_naming_convenstion']
                    Data[specimen][Experiment_Type]['er_location_name']=DIRS_data[dir_name]['er_location_name']
                    Data[specimen][Experiment_Type]['user_name']=DIRS_data[dir_name]['user_name']
                    Data[specimen][Experiment_Type]['sample_volume']=DIRS_data[dir_name]['sample_volume']
                    Data[specimen][Experiment_Type]['moment_units']=DIRS_data[dir_name]['moment_units']
                    Data[specimen][Experiment_Type]['labfield_DI']=DIRS_data[dir_name]['labfield_DI']

        #-----------------------------------
        # Convert Data{} to MagIC
        #-----------------------------------

        specimens_list=list(Data.keys())
        specimens_list.sort()
        for specimen in specimens_list:
            Experiment_Types_list=list(Data[specimen].keys())
            Experiment_Types_list.sort()
            for Experiment_Type in Experiment_Types_list:
                if Experiment_Type in ["Thellier"]:

                    tmp_MagRecs=[]

                    # IMORTANT:
                    # phi and theta of lab field are not defined
                    # defaults are defined here:
                    phi,theta='0.','90.'

                    header_line=Data[specimen][Experiment_Type]['header_data']
                    experiment_treatments=[]
                    measurement_running_number=0
                    methcodes=["LP-PI-TRM"] # start to make a list of the methcodes. and later will merge it to one string

                    for i in range(len(Data[specimen][Experiment_Type]['meas_data'])):
                        meas_line=Data[specimen][Experiment_Type]['meas_data'][i]

                        #------------------
                        # check if the same treatment appears more than once. If yes, assuming that the measurements is repeated twice,
                        # ignore the first, and take only the second one
                        #------------------

                        if i< (len(Data[specimen][Experiment_Type]['meas_data'])-2) :
                            Repeating_measurements=True
                            for key in ['treatment','specimen_name']:
                                if Data[specimen][Experiment_Type]['meas_data'][i][key]!=Data[specimen][Experiment_Type]['meas_data'][i+1][key]:
                                    Repeating_measurements=False
                            if Repeating_measurements==True:
                                "Found a repeating measurement at line %i, sample %s. taking the last one"%(i,specimen)
                                continue
                        #------------------
                        # Special treatment for first line (NRM data).
                        #------------------

                        if i==0:
                           if "." not in meas_line['treatment']:
                               meas_line['treatment']="0.0"
                           elif meas_line['treatment'].split(".")[0]=="" and meas_line['treatment'].split(".")[1]=='0': # if NRM is in the form of ".0" instead of "0.0"
                               meas_line['treatment']="0.0"
                           elif  float(meas_line['treatment'].split(".")[0])<50 and float(meas_line['treatment'].split(".")[-1])==0: # if NRM is in the form of "20.0" instead of "0.0"
                               meas_line['treatment']="0.0"

                        #------------------
                        # fix line in format of XX instead of XX.YY
                        #------------------
                        if "." not in meas_line['treatment']:
                            meas_line['treatment']=meas_line['treatment']+".0"
                        if meas_line['treatment'].split(".")[1]=="":
                            meas_line['treatment']=meas_line['treatment']+"0"

                        #------------------
                        # header data
                        #------------------

                        MagRec={}
                        MagRec['er_citation_names']="This study"
                        labfield=float(header_line['labfield'])*1e-6 # convert from microT to Tesla
                        MagRec["magic_experiment_name"]=""
                        #------------------
                        # Body data
                        #------------------

                        MagRec["er_specimen_name"]=specimen
                        MagRec["er_sample_name"]=self.get_sample_name(MagRec["er_specimen_name"],Data[specimen][Experiment_Type]['sample_naming_convenstion'])
                        MagRec["er_site_name"]=self.get_site_name(MagRec["er_sample_name"],Data[specimen][Experiment_Type]['site_naming_convenstion'])
                        MagRec['er_location_name']=Data[specimen][Experiment_Type]['er_location_name']
                        MagRec['er_analyst_mail_names']=Data[specimen][Experiment_Type]['user_name']

                        MagRec["measurement_flag"]='g'
                        MagRec["measurement_standard"]='u'
                        MagRec["measurement_number"]="%i"%measurement_running_number
                        MagRec["measurement_dec"]=meas_line['dec']
                        MagRec["measurement_inc"]=meas_line['inc']
                        if Data[specimen][Experiment_Type]['moment_units']=='mA/m':
                            MagRec["measurement_magn_moment"]="%5e"%(float(meas_line['moment'])*1e-3*float(Data[specimen][Experiment_Type]['sample_volume'])) # converted to Am^2
                        if Data[specimen][Experiment_Type]['moment_units']=='emu':
                            MagRec["measurement_magn_moment"]="%5e"%(float(meas_line['moment'])*1e-3) # converted to Am^2
                        if Data[specimen][Experiment_Type]['moment_units']=='Am^2':
                            MagRec["measurement_magn_moment"]="%5e"%(float(meas_line['moment'])) # converted to Am^2
                        MagRec["measurement_temp"]='273.' # room temp in kelvin

                        # Date and time
##                                    date=meas_line['Measurement Date'].strip("\"").split('-')
##                                    yyyy=date[2];dd=date[1];mm=date[0]
##                                    hour=meas_line['Measurement Time'].strip("\"")
##                                    MagRec["measurement_date"]=yyyy+':'+mm+":"+dd+":"+hour

                        # lab field data: distinguish between PI experiments to AF/Thermal
                        treatments=meas_line['treatment'].split(".")
                        if float(treatments[1])==0:
                            MagRec["treatment_dc_field"]='0'
                            MagRec["treatment_dc_field_phi"]='0'
                            MagRec["treatment_dc_field_theta"]='0'
                        else:
                            MagRec["treatment_dc_field"]='%8.3e'%(labfield)
                            MagRec["treatment_dc_field_phi"]=Data[specimen][Experiment_Type]['labfield_DI'][0]
                            MagRec["treatment_dc_field_theta"]=Data[specimen][Experiment_Type]['labfield_DI'][1]

                        #------------------
                        # Lab Treatments
                        #------------------

                        # NRM
                        if float(treatments[0])==0 and float(treatments[1])==0:
                                MagRec["magic_method_codes"]="LT-NO"
                                experiment_treatments.append('0')
                                MagRec["treatment_temp"]='273.'
                                IZorZI=""

                        # Zerofield step
                        elif float(treatments[1])==0:
                                MagRec["magic_method_codes"]="LT-T-Z"
                                MagRec["treatment_temp"]='%8.3e' % (float(treatments[0])+273.) # temp in kelvin

                                #  check if this is ZI or IZ:
                                for j in range (0,i):
                                    previous_lines=Data[specimen][Experiment_Type]['meas_data'][j]
                                    if previous_lines['treatment'].split(".")[0] == meas_line['treatment'].split(".")[0]:
                                        if float(previous_lines['treatment'].split(".")[1]) == 1 or float(previous_lines['treatment'].split(".")[1]) == 10:
                                            if "LP-PI-TRM-IZ" not in methcodes:
                                                methcodes.append("LP-PI-TRM-IZ")
                                            IZorZI=""
                                        else:
                                            IZorZI="Z"

                        # Infield step
                        elif float(treatments[1])==1 or float(treatments[1])==10:
                                MagRec["magic_method_codes"]="LT-T-I"
                                MagRec["treatment_temp"]='%8.3e' % (float(treatments[0])+273.) # temp in kelvin

                                # check if this is ZI,IZ:
                                for j in range (0,i):
                                    previous_lines=Data[specimen][Experiment_Type]['meas_data'][j]
                                    if previous_lines['treatment'].split(".")[0] == meas_line['treatment'].split(".")[0]:
                                        if float(previous_lines['treatment'].split(".")[1]) == 0:
                                            if "LP-PI-TRM-ZI" not in methcodes:
                                                methcodes.append("LP-PI-TRM-ZI")
                                                IZorZI=""
                                        else:
                                                IZorZI="I"
                        # pTRM check step
                        elif float(treatments[1])==2 or float(treatments[1])==20:
                                MagRec["magic_method_codes"]="LT-PTRM-I"
                                MagRec["treatment_temp"]='%8.3e' % (float(treatments[0])+273.) # temp in kelvin
                                if "LP-PI-ALT" not in methcodes:
                                    methcodes.append("LP-PI-ALT")

                        # Tail check step
                        elif float(treatments[1])==3 or float(treatments[1])==30:
                                MagRec["magic_method_codes"]="LT-PTRM-MD"
                                MagRec["treatment_temp"]='%8.3e' % (float(treatments[0])+273.) # temp in kelvin
                                if "LP-PI-BT-MD" not in methcodes:
                                    methcodes.append("LP-PI-BT-MD")
                                    MagRec["treatment_dc_field"]="0"
                                    MagRec["treatment_dc_field_phi"]="0"
                                    MagRec["treatment_dc_field_theta"]="0"

                        # Additivity check step
                        elif float(treatments[1])==4 or float(treatments[1])==40:
                                MagRec["magic_method_codes"]="LT-PTRM-AC"
                                MagRec["treatment_temp"]='%8.3e' % (float(treatments[0])+273.) # temp in kelvin
                                if "LP-PI-BT" not in methcodes:
                                    methcodes.append("LP-PI-BT")

                        # Thellier Thellier protocol (1 for one direction and 5 for the antiparallel)
                        # Lab field direction of 1 is as put in the GUI dialog box
                        # Lab field direction of 5 is the anti-parallel direction of 1

                        elif float(treatments[1])==5 or float(treatments[1])==50:
                                MagRec["magic_method_codes"]="LT-T-I"
                                MagRec["treatment_temp"]='%8.3e' % (float(treatments[0])+273.) # temp in kelvin
                                MagRec["treatment_dc_field_phi"]="%.2f"%((float(Data[specimen][Experiment_Type]['labfield_DI'][0])+180.)%360.)
                                MagRec["treatment_dc_field_theta"]="%.2f"%(float(Data[specimen][Experiment_Type]['labfield_DI'][1])*-1.)
                                if "LP-PI-II" not in methcodes:
                                    methcodes.append("LP-PI-II")

                        else:
                                print("-E- ERROR in file %s"%Experiment_Type)
                                print("-E- ERROR in treatment ",meas_line['treatment'])
                                print("... exiting until you fix the problem")

                        #-----------------------------------

                        #MagRec["magic_method_codes"]=lab_treatment+":"+lab_protocols_string
                        #MagRec["magic_experiment_name"]=specimen+":"+lab_protocols_string

                        tmp_MagRecs.append(MagRec)
                        measurement_running_number+=1
                        headers=list(MagRec.keys())
                        for key in headers:
                            if key not in magic_measurements_headers:
                                magic_measurements_headers.append(key)

                    # arrange magic_method_codes and magic_experiment_name:
                    magic_method_codes="LP-PI-TRM"
                    # Coe mothod
                    if "LP-PI-TRM-ZI" in methcodes and "LP-PI-TRM-IZ" not in methcodes and "LP-PI-II" not in methcodes:
                        magic_method_codes=magic_method_codes+":LP-PI-TRM-ZI"
                    if "LP-PI-TRM-ZI" not in methcodes and "LP-PI-TRM-IZ"  in methcodes and "LP-PI-II" not in methcodes:
                        magic_method_codes=magic_method_codes+":LP-PI-TRM-IZ"
                    if "LP-PI-TRM-ZI"  in methcodes and "LP-PI-TRM-IZ"  in methcodes and "LP-PI-II" not in methcodes:
                        magic_method_codes=magic_method_codes+":LP-PI-BT-IZZI"
                    if "LP-PI-II"  in methcodes:
                        magic_method_codes=magic_method_codes+":LP-PI-II"
                    if "LP-PI-ALT"  in methcodes:
                        magic_method_codes=magic_method_codes+":LP-PI-ALT"
                    if "LP-PI-BT-MD"  in methcodes:
                        magic_method_codes=magic_method_codes+":LP-PI-BT-MD"
                    if "LP-PI-BT"  in methcodes:
                        magic_method_codes=magic_method_codes+":LP-PI-BT"
                    for i in range(len(tmp_MagRecs)):
                        STRING=":".join([tmp_MagRecs[i]["magic_method_codes"],magic_method_codes])
                        tmp_MagRecs[i]["magic_method_codes"]=STRING
                        STRING=":".join([tmp_MagRecs[i]["er_specimen_name"],magic_method_codes])
                        tmp_MagRecs[i]["magic_experiment_name"]=STRING
                        MagRecs.append(tmp_MagRecs[i])

                elif Experiment_Type in ["ATRM 6 positions"]:

                    tmp_MagRecs=[]

                    header_line=Data[specimen][Experiment_Type]['header_data']
                    experiment_treatments=[]
                    measurement_running_number=0
                    methcodes=["LP-AN-TRM"] # start to make a list of the methcodes. and later will merge it to one string

                    for i in range(len(Data[specimen][Experiment_Type]['meas_data'])):
                        meas_line=Data[specimen][Experiment_Type]['meas_data'][i]

                        #------------------
                        # check if the same treatment appears more than once. If yes, assuming that the measurements is repeated twice,
                        # ignore the first, and take only the second one
                        #------------------

                        if i< (len(Data[specimen][Experiment_Type]['meas_data'])-2) :
                            Repeating_measurements=True
                            for key in ['treatment','specimen_name']:
                                if Data[specimen][Experiment_Type]['meas_data'][i][key]!=Data[specimen][Experiment_Type]['meas_data'][i+1][key]:
                                    Repeating_measurements=False
                            if Repeating_measurements==True:
                                "Found a repeating measurement at line %i, sample %s. taking the last one"%(i,specimen)
                                continue

                        #------------------
                        # fix line in format of XX instead of XX.0
                        #------------------
                        if "." not in meas_line['treatment']:
                            meas_line['treatment']=meas_line['treatment']+".0"
                        if meas_line['treatment'].split(".")[1]=="":
                            meas_line['treatment']=meas_line['treatment']+"0"

                        #------------------
                        # header data
                        #------------------

                        MagRec={}
                        MagRec['er_citation_names']="This study"
                        labfield=float(header_line['labfield'])*1e-6 # convert from microT to Tesal
                        MagRec["magic_experiment_name"]=""

                        MagRec["er_specimen_name"]=specimen
                        #MagRec["magic_method_codes"]="LP-AN-TRM"
                        MagRec["magic_experiment_name"]=MagRec["er_specimen_name"]+":LP-AN-TRM"

                        #------------------
                        # Body data
                        #------------------

                        MagRec["er_specimen_name"]=specimen
                        MagRec["er_sample_name"]=self.get_sample_name(MagRec["er_specimen_name"],Data[specimen][Experiment_Type]['sample_naming_convenstion'])
                        MagRec["er_site_name"]=self.get_site_name(MagRec["er_sample_name"],Data[specimen][Experiment_Type]['site_naming_convenstion'])
                        MagRec['er_location_name']=Data[specimen][Experiment_Type]['er_location_name']
                        MagRec['er_analyst_mail_names']=Data[specimen][Experiment_Type]['user_name']

                        MagRec["measurement_flag"]='g'
                        MagRec["measurement_standard"]='u'
                        MagRec["measurement_number"]="%i"%measurement_running_number
                        MagRec["measurement_dec"]=meas_line['dec']
                        MagRec["measurement_inc"]=meas_line['inc']
                        MagRec["measurement_magn_moment"]="%5e"%(float(meas_line['moment'])*1e-3*float(Data[specimen][Experiment_Type]['sample_volume'])) # converted to Am^2
                        MagRec["measurement_temp"]='273.' # room temp in kelvin

                        treatments=meas_line['treatment'].split(".")
                        if len(treatments[1])>1:
                            treatments[1]=treatments[1][0]

                        MagRec["treatment_temp"]='%8.3e' % (float(treatments[0])+273.) # temp in kelvin

                        # labfield direction
                        if float(treatments[1])==0:
                            MagRec["treatment_dc_field"]='0'
                            MagRec["treatment_dc_field_phi"]='0'
                            MagRec["treatment_dc_field_theta"]='0'
                            MagRec["magic_method_codes"]="LT-T-Z:LP-AN-TRM"
                        else:

                            MagRec["treatment_dc_field"]='%8.3e'%(labfield)

                            if float(treatments[1])==7 or float(treatments[1])==70: # alteration check as final measurement
                                    MagRec["magic_method_codes"]="LT-PTRM-I:LP-AN-TRM"
                            else:
                                    MagRec["magic_method_codes"]="LT-T-I:LP-AN-TRM"

                            # find the direction of the lab field in two ways:
                            # (1) using the treatment coding (XX.1=+x, XX.2=+y, XX.3=+z, XX.4=-x, XX.5=-y, XX.6=-z)
                            tdec=[0,90,0,180,270,0,0,90,0] # atrm declination/inlclination order
                            tinc=[0,0,90,0,0,-90,0,0,90] # atrm declination/inlclination order

                            ipos_code=int(treatments[1])-1
                            # (2) using the magnetization
                            DEC=float(MagRec["measurement_dec"])
                            INC=float(MagRec["measurement_inc"])
                            if INC < 45 and INC > -45:
                                if DEC>315  or DEC<45: ipos_guess=0
                                if DEC>45 and DEC<135: ipos_guess=1
                                if DEC>135 and DEC<225: ipos_guess=3
                                if DEC>225 and DEC<315: ipos_guess=4
                            else:
                                if INC >45: ipos_guess=2
                                if INC <-45: ipos_guess=5
                            # prefer the guess over the code
                            ipos=ipos_guess
                            MagRec["treatment_dc_field_phi"]='%7.1f' %(tdec[ipos])
                            MagRec["treatment_dc_field_theta"]='%7.1f'% (tinc[ipos])
                            # check it
                            if ipos_guess!=ipos_code and treatments[1]!='7':
                                print("-E- ERROR: check specimen %s step %s, ATRM measurements, coding does not match the direction of the lab field!"%(MagRec["er_specimen_name"],".".join(list(treatments))))

                        tmp_MagRecs.append(MagRec)
                        measurement_running_number+=1
                        headers=list(MagRec.keys())
                        for key in headers:
                            if key not in magic_measurements_headers:
                                magic_measurements_headers.append(key)

                    for i in range(len(tmp_MagRecs)):
                        MagRecs.append(tmp_MagRecs[i])

                else:
                    print("-E- ERROR. sorry, file format %s is not supported yet. Please contact rshaar@ucsd.edu"%Experiment_Type)

        #-------------------------------------------
        #  magic_measurements.txt
        #-------------------------------------------

        #fout=open("magic_measurements.txt",'w')
        fout=open(self.output_file_path.GetValue(), 'w')
        fout.write("tab\tmagic_measurements\n")
        header_string=""
        for i in range(len(magic_measurements_headers)):
            header_string=header_string+magic_measurements_headers[i]+"\t"
        fout.write(header_string[:-1]+"\n")

        for MagRec in MagRecs:
            line_string=""
            for i in range(len(magic_measurements_headers)):
                if magic_measurements_headers[i] in list(MagRec.keys()):
                    line_string=line_string+MagRec[magic_measurements_headers[i]]+"\t"
                else:
                    line_string=line_string+"\t"

            fout.write(line_string[:-1]+"\n")

        #-------------------------------------------

        dlg1 = wx.MessageDialog(None,caption="Message:", message="file converted to {}\n you can try running thellier gui...\n".format(self.output_file_path.GetValue()) ,style=wx.OK|wx.ICON_INFORMATION)
        dlg1.ShowModal()
        dlg1.Destroy()
        self.Destroy()

class message_box(wx.Frame):
    """"""

    #----------------------------------------------------------------------
    def __init__(self,title):
        wx.Frame.__init__(self, parent=None,size=(1000,500))

        self.panel = wx.Panel(self)
        self.sizer = wx.BoxSizer(wx.VERTICAL)
        self.text_log = wx.TextCtrl(self.panel, id=-1, style=wx.TE_MULTILINE | wx.TE_READONLY  | wx.HSCROLL)
        self.sizer.Add(self.text_log, 1, wx.EXPAND)
        TEXT='''
        # -------------------------------------
        #
        # Programs assumptions:
        #
        # 1) Each file contains the data one specimen
        # 2) First line is the header: "Thellier-tdt"
        # 3) Second line in header inlucdes 4 fields:
        #    [Blab] ,['core_azimuth'] , ['core_plunge'] , ['bedding_dip_direction'] , ['bedding_dip']
        # 4) Body includes 5 fields
        #    [specimen_name], [treatments], [moment],[meas_dec],[meas_inc]
        # -------------------------------------
        # Thellier experiment:
        #
        # Tretments: XXX.0 (zerofield)
        #            XXX.1 (infield)
        #            XXX.2 (pTRM check)
        #            XXX.3 (Tail check)
        #            XXX.4 (Additivity check; Krasa et al., 2003)
        #            XXX.5 (Original Thellier-Thellier protocol. )
        #                 (where .5 is for the second direction and .1 in the first)
        # XXX = temperature in degrees
        #
        #
        # 1) If if XXX <50 then assuming that this is NRM (273K)
        # 2) Lab field defaul is Z direction (theta=0, phi=90)
        # 3) The program does not support Thermal demagnetization
        #
        # -------------------------------------
        #
        #   ATRM in six positions
        #
        # Tretments: XXX.0 zerofield
        #            XXX.1 +x
        #            XXX.2 +y
        #            XXX.3 +z
        #            XXX.4 -x
        #            XXX.5 -y
        #            XXX.6 -z
        #            XXX.7 alteration check
        #
        #
        # 1) The program checks if the direction of the magnetization fits the coding above.
        #    If not: an error message will appear
        # 2) Alteration check can be in any direction
        # 3) The order of the measurements is not important
        #
        # For questions and support: rshaar@ucsd.edu
        # -------------------------------------'''

        self.text_log.AppendText(TEXT)
##        fin =open(file_path,'r')
##        for line in fin.readlines():
##            if "-E-" in line :
##                self.text_log.SetDefaultStyle(wx.TextAttr(wx.RED))
##                self.text_log.AppendText(line)
##            if "-W-" in line:
##                self.text_log.SetDefaultStyle(wx.TextAttr(wx.BLACK))
##                self.text_log.AppendText(line)
##        fin.close()
        #sizer.Fit(self)
        self.panel.SetSizer(self.sizer)

def main(command_line=True, wd=None):
    if command_line:
        import sys
        args=sys.argv
        if "-WD" in args:
            ind=args.index("-WD")
            WD=args[ind+1]
        else:
            print("please specify working directory for output file WD")
            return False
    if not command_line:
        if not wd:
            import os
            WD = os.getcwd()
        else:
            WD = wd

    app = wx.App()
    app.frame = convert_tdt_files_to_MagIC(WD)
    app.frame.Show()
    app.frame.Center()
    app.MainLoop()

if __name__ == '__main__':
    main()

#main()

reason: import scipy,from scipy


text:
import warnings

import numpy as np
import scipy.linalg as spla
import scipy.special as spsp

class CovMatAdapt:
    """ Covariance Matrix Adaption Evolution Strategy

    Implements the Covariance Matrix Adaption Evolution Strategy
    as described in 'The CMA Evolution Strategy: A Tutorial' by
    Nikolaus Hansen. For exact details and meaning of the parameters,
    please see the tutorial.

    Args:
        func (numpy array -> float): Function to optimize.
        mean_vec (numpy array): Initiale locations to start the optimization.
        step_size (float > 0): Overall variance step size. Rule of thumb:
            Choose it in a way that the optimum is within mean_vec +/-
            3*step_size.
        pop_size (int > 0): How many samples should be generated in each
            generation?
        elite_size (int > 0): How many of the top samples are used to
            reproduce?
        weights (numpy array): To weight the contribution of the different
            elites.
        c_sig (float > 0): Learning rate for the cumulation for the step-size
            control.
        d_sig (float > 0): Dampening parameter for step-size updates.
        c_c (float > 0): Learning rate for cumulation for rank-one updates.
        c_1 (float > 0): Learning rate for rank-one updates.
        c_mu (float > 0): Learning rate for rank-mu updates.
        c_m (float > 0): Learning rate for mean vector updates.
        abs_tol (float > 0): Absolute change criterion for early stopping.
        rel_tol (float > 0): Relative change criterion for early stopping.
        maxgen (int > 0): Maximum number of generations/iterations before
            termination.
        verbose (bool): Provides extra feedback for the user.
    """

    def __init__(self, func, mean_vec, step_size, pop_size=None,
                 elite_size=None, weights=None,
                 c_sig=None, d_sig=None, c_c=None,
                 c_1=None, c_mu=None, c_m=None, abs_tol=1e-5,
                 rel_tol=1e-4, maxgen=50, verbose=True):

        self.func = func
        self.mean_vec = mean_vec
        self.step_size = step_size
        self.abs_tol = abs_tol
        self.rel_tol = rel_tol
        self.maxgen = maxgen
        self.ndim = mean_vec.size

        # Default values for population sizes and weights
        if pop_size:
            self.pop_size = pop_size
        else:
            self.pop_size = int(4 + np.floor(3 * np.log(self.ndim)))

        if elite_size:
            self.elite_size = elite_size
        else:
            self.elite_size = self.pop_size // 2

        if weights:
            self.weights = weights
        else:
            self.weights = (np.log(0.5 * (self.pop_size + 1)) -
                            np.log(np.arange(1, self.pop_size + 1)))

        norm_pos_weights = (self.weights[:self.elite_size] /
                            sum(self.weights[:self.elite_size]))
        self.mu_eff = (sum(norm_pos_weights)**2 /
                       sum(norm_pos_weights**2))

        # Initializes the step size parameters
        if c_sig:
            self.c_sig = c_sig
        else:
            self.c_sig = (self.mu_eff + 2) / (self.ndim + self.mu_eff + 5)

        if d_sig:
            self.d_sig = d_sig
        else:
            self.d_sig = (1 + 2 * max(0, np.sqrt((self.mu_eff - 1) /
                                                 (self.ndim + 1)) - 1) +
                          self.c_sig)

        # Initializes parameters for covariance adaption
        if c_c:
            self.c_c = c_c
        else:
            self.c_c = ((4 + self.mu_eff / self.ndim) /
                        (self.ndim + 4 + 2 * self.mu_eff / self.ndim))

        if c_1:
            self.c_1 = c_1
        else:
            self.c_1 = 2 / ((self.ndim + 1.3)**2 + self.mu_eff)

        if c_mu:
            self.c_mu = c_mu
        else:
            self.c_mu = min(1 - self.c_1, 2 *
                            ((self.mu_eff - 2 + 1 / self.mu_eff) /
                             ((self.ndim + 2)**2 + self.mu_eff)))

        if c_m:
            self.c_m = c_m
        else:
            self.c_m = 1

        # Set the negative weights, if not provided:
        if not weights:
            neg_weights = self.weights[self.elite_size:]
            mu_eff_neg = (sum(neg_weights)**2 /
                          sum(neg_weights**2))

            alpha_mu_neg = 1 + self.c_1 / self.c_mu
            alpha_mu_eff_neg = 1 + 2 * mu_eff_neg / (self.mu_eff + 2)
            alpha_pos_def_neg = ((1 - self.c_1 - self.c_mu) /
                                 (self.ndim * self.c_mu))
            self.weights[:self.elite_size] = norm_pos_weights
            denom = min(alpha_mu_neg, alpha_mu_eff_neg, alpha_pos_def_neg)
            self.weights[self.elite_size:] = (denom * neg_weights /
                                              -sum(neg_weights))

        self.cov_mat = np.identity(len(self.mean_vec))

        # Set the evolution paths
        self.p_sig = 0
        self.p_c = 0

        self.verbose = verbose

    def sample_and_evaluate(self, func, ndim, mean_vec,
                            cov_mat, pop_size, step_size):
        """ Samples the function evaluations

        Used to sample the candidate values and evaluates them.

        Args:
            func (numpy array -> float): Function to evaluate.
            ndim (int > 0): Number of dimensions.
            mean_vec (numpy array): Mean vector of the multivariate Gaussian
                distribution.
            cov_mat (numpy array): Covariance matrix of the multivariate
                Gaussian distribution.
            pop_size (int > 0): Number of samples.
            step_size (float > 0): Overall variance step size.

        Returns:
            Numpy array containing the function values, the mean-centered and
            step_size ajusted input values, and zero-centered and unadjusted
            input values as columns.
        """
        with warnings.catch_warnings():
            if not self.verbose:
                warnings.filterwarnings('ignore')

            y = np.random.multivariate_normal(np.zeros(ndim),
                                              cov_mat,
                                              size=pop_size)

        x = mean_vec + step_size * y
        func_values = list(map(func, x))
        conc_matrix = np.c_[func_values, x, y]
        conc_matrix = conc_matrix[conc_matrix[:, 0].argsort()]
        return conc_matrix

    def heaviside(self, gen, norm_p_sig, c_sig, ndim, expec_norm_gaussian):
        """Heaviside function

        Prevents too fast increases of axes of the covariance matrix.

        Args:
            gen (int > 0): The current generation the minimization procedure is
                in.
            norm_p_sig(float > 0): The norm of p_sig.
            c_sig (float > 0): Learning rate for the cumulation for the
                step-size control.
            ndim (int > 0): Number of dimensions.
            expec_norm_gaussian (float > 0): Expectation of the norm of a
                multivariate standard normal distribution.

        Returns:
            0 or 1.

        """
        a = norm_p_sig / (np.sqrt(1 - (1 - c_sig)**(2 * (gen + 1))))
        b = (1.4 + 2 / (ndim + 1)) * expec_norm_gaussian

        if a < b:
            return 1
        else:
            return 0

    def minimize(self):
        """Minimize objective function

        Returns:
            Dictionary containing following keys:
                'best fvalue' (float): Best function value that has been
                    reached.
                'best param' (numpy array): Parameters corresponding to best
                    function value.
                'mean vector' (numpy array): Last mean vector obtained during
                    the minimization.
                'cov matrix' (numpy array): Last covariance matrix obtained
                    during minimization.
                'converged in' (int): Generation in which the procedure
                    converged.
                'abs/rel conv': Information about absolute and relative
                    convergence.

        """
        # Matrix to fix some numerical issues with the cholesky decomposition.
        offset_matrix = np.identity(self.ndim) * 0.1
        expec_norm_gaussian = (np.sqrt(2) * spsp.gamma(0.5 * (self.ndim + 1)) /
                               spsp.gamma(0.5 * self.ndim))
        pop_matrix = self.sample_and_evaluate(self.func, self.ndim,
                                              self.mean_vec, self.cov_mat,
                                              self.pop_size, self.step_size)

        y = pop_matrix[:, (self.ndim + 1):]
        x = pop_matrix[:, 1:(self.ndim + 1)]
        func_values = pop_matrix[:, 0]
        best_fvalue = pop_matrix[0, 0]
        best_param = x[0, :]

        if self.verbose:
            print('Initial value:', best_fvalue)

        for gen in range(1, self.maxgen + 1):
            y_weighted = np.sum(y[:self.elite_size] *
                                self.weights[:self.elite_size, np.newaxis],
                                axis=0)
            upd_mean_vec = (self.mean_vec + self.c_m * self.step_size *
                            y_weighted)
            diff_vec = abs(upd_mean_vec - self.mean_vec)

            if (max(diff_vec) < self.abs_tol or
                    max(diff_vec / abs(self.mean_vec)) < self.rel_tol):
                result_dict = {'best fvalue': best_fvalue,
                               'best param': best_param,
                               'mean vector': self.mean_vec,
                               'cov matrix': self.cov_mat,
                               'converged in': gen,
                               'abs/rel conv': (max(diff_vec),
                                                max(diff_vec /
                                                    abs(self.mean_vec)))}
                return result_dict

            self.mean_vec = upd_mean_vec
            try:
                chol_covm_L = spla.cholesky(self.cov_mat, lower=True)
            except np.linalg.linalg.LinAlgError as err:
                if self.verbose:
                    print(('We have hit an linalg exception',
                           'in the cholesky decomposition'))
                    print(repr(err))
                try:
                    chol_covm_L = spla.cholesky(
                        self.cov_mat + offset_matrix, lower=True)
                except BaseException:
                    if self.verbose:
                        print(('Even the matrix with increased diagonal',
                               'elements failed.'))

            c_y_vec = spla.solve_triangular(
                chol_covm_L, y_weighted, lower=True)

            # Updating the step sizes
            self.p_sig = ((1 - self.c_mu) * self.p_sig +
                          np.sqrt(self.c_sig * (2 - self.c_sig) *
                                  self.mu_eff) * c_y_vec)
            norm_p_sig = np.linalg.norm(self.p_sig)
            self.step_size = (self.step_size *
                              np.exp(self.c_sig / self.d_sig *
                                     norm_p_sig /
                                     expec_norm_gaussian - 1))

            # Updating the covariance matrix
            h_sig = self.heaviside(gen, norm_p_sig, self.c_sig,
                                   self.ndim, expec_norm_gaussian)

            self.p_c = ((1 - self.c_c) * self.p_c +
                        h_sig * np.sqrt(self.c_c * (2 - self.c_c) *
                                        self.mu_eff) * y_weighted)
            c_y_mat = spla.solve_triangular(chol_covm_L,
                                            y[self.elite_size:, :].T,
                                            lower=True).T
            c_y_mat_norm = np.linalg.norm(c_y_mat, axis=1)

            w_adj = self.weights.copy()
            w_adj[self.elite_size:] = (w_adj[self.elite_size:] *
                                       (self.ndim / c_y_mat_norm))

            matrix_sum = (w_adj * y.T).dot(y)
            self.cov_mat = ((1 + self.c_1 * (1 - h_sig) * self.c_c *
                             (2 - self.c_c) - self.c_1 - self.c_mu *
                             sum(self.weights)) * self.cov_mat + self.c_1 *
                            np.outer(self.p_c, self.p_c)
                            + self.c_mu * matrix_sum)

            # Create new generations
            pop_matrix = self.sample_and_evaluate(self.func,
                                                  self.ndim,
                                                  self.mean_vec,
                                                  self.cov_mat,
                                                  self.pop_size,
                                                  self.step_size)

            y = pop_matrix[:, (self.ndim + 1):]
            x = pop_matrix[:, 1:(self.ndim + 1)]
            func_values = pop_matrix[:, 0]

            if func_values[0] < best_fvalue:
                best_fvalue = pop_matrix[0, 0]
                best_param = x[0, :]

            if self.verbose and gen % 5 == 0:
                print('Best value after', gen, 'generations:', best_fvalue)

        else:
            print('Failed to converge within maximum generation bound.')
            result_dict = {'best fvalue': best_fvalue,
                           'best param': best_param,
                           'mean vector': self.mean_vec,
                           'cov matrix': self.cov_mat,
                           'converged in': self.maxgen}
            return result_dict

ext: py

max_stars_repo_name: RoetGer/CMAES

max_issues_repo_name: RoetGer/CMAES

max_forks_repo_name: RoetGer/CMAES

reason: import numpy,import scipy


text:
import os

import numpy as np
from tensorflow.keras import Input, Sequential
from tensorflow.keras.layers import Dense

from kitt.callbacks import ModelCheckpoint

def create_model():
    return Sequential([Input(shape=(1,)), Dense(4), Dense(1)])

def metric_return_y(y_true, y_pred):
    return y_true

def test_checkpoint_best_n_max(tmpdir):
    checkpointer = ModelCheckpoint(
        os.path.join(tmpdir, "{epoch:02d}.{metric}.hdf5"), "metric", save_n_best=2
    )
    train(checkpointer, [1, 2, 1, 0, 5, 3, 2, 4, 6])
    assert sorted(os.listdir(tmpdir)) == ["05.5.0.hdf5", "09.6.0.hdf5", "final.hdf5"]

def test_checkpoint_best_n_min(tmpdir):
    checkpointer = ModelCheckpoint(
        os.path.join(tmpdir, "{epoch:02d}.{metric}.hdf5"),
        "metric",
        mode="min",
        save_n_best=2,
    )
    train(checkpointer, [-1, -2, -1, 0, -5, -3, -2, -4, -6])
    assert sorted(os.listdir(tmpdir)) == ["05.-5.0.hdf5", "09.-6.0.hdf5", "final.hdf5"]

def test_checkpoint_every_n(tmpdir):
    checkpointer = ModelCheckpoint(
        os.path.join(tmpdir, "{epoch:02d}.{metric}.hdf5"),
        "metric",
        save_n_best=False,
        save_every_n_epochs=2,
    )
    train(checkpointer, [1, 2, 3, 4])
    assert sorted(os.listdir(tmpdir)) == ["02.2.0.hdf5", "04.4.0.hdf5", "final.hdf5"]

def test_checkpoint_every_n_best_n(tmpdir):
    checkpointer = ModelCheckpoint(
        os.path.join(tmpdir, "{epoch:02d}.{metric}.hdf5"),
        "metric",
        save_n_best=2,
        save_every_n_epochs=2,
    )
    train(checkpointer, [1, 2, 3, 4, 5, 6])
    assert sorted(os.listdir(tmpdir)) == [
        "02.2.0.hdf5",
        "04.4.0.hdf5",
        "05.5.0.hdf5",
        "06.6.0.hdf5",
        "final.hdf5",
    ]

def test_checkpoint_save_final(tmpdir):
    checkpointer = ModelCheckpoint(
        os.path.join(tmpdir, "{epoch:02d}.{metric}.hdf5"), "metric", save_n_best=False
    )
    train(checkpointer, [1, 2, 3])
    assert os.listdir(tmpdir) == ["final.hdf5"]

def train(checkpointer, y_values):
    metric = metric_return_y
    metric.__name__ = "metric"

    def generator():
        for item in y_values:
            yield np.array([item]), np.array([item])

    model = create_model()
    model.compile(optimizer="adam", loss="mse", metrics=[metric])
    model.fit(
        generator(),
        steps_per_epoch=1,
        shuffle=False,
        epochs=len(y_values),
        callbacks=[checkpointer],
        workers=0,
    )

max_stars_repo_name: David-Ciz/kitt

max_issues_repo_name: David-Ciz/kitt

max_forks_repo_name: David-Ciz/kitt

reason: import numpy


text:
# Copyright 2021 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import abc
import logging

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

import rlkit.torch.pytorch_util as ptu
from rlkit.policies.base import ExplorationPolicy
from rlkit.torch.core import torch_ify, elem_or_tuple_to_numpy
from rlkit.torch.distributions import (
    Delta, TanhNormal, MultivariateDiagonalNormal, GaussianMixture, GaussianMixtureFull,
)
from rlkit.torch.networks import Mlp, CNN
from rlkit.torch.networks.basic import MultiInputSequential
from rlkit.torch.networks.stochastic.distribution_generator import (
    DistributionGenerator
)
from rlkit.torch.lvm.latent_variable_model import LatentVariableModel

class VAEPolicy(LatentVariableModel):
    def __init__(
            self,
            obs_dim,
            action_dim,
            latent_dim,
    ):
        encoder = Encoder(obs_dim, latent_dim, action_dim)
        decoder = Decoder(obs_dim, latent_dim, action_dim)
        super().__init__(encoder, decoder)

        self.latent_dim = latent_dim

    def forward(self, state, action):
        z = F.relu(self.encoder.e1(torch.cat([state, action], 1)))
        z = F.relu(self.encoder.e2(z))

        mean = self.encoder.mean(z)
        # Clamped for numerical stability
        log_std = self.encoder.log_std(z).clamp(-4, 15)
        std = torch.exp(log_std)
        z = mean + std * ptu.from_numpy(
            np.random.normal(0, 1, size=(std.size())))

        u = self.decode(state, z)

        return u, mean, std

    def decode(self, state, z=None):
        if z is None:
            z = ptu.from_numpy(np.random.normal(0, 1, size=(
            state.size(0), self.latent_dim))).clamp(-0.5, 0.5)

        a = F.relu(self.decoder.d1(torch.cat([state, z], 1)))
        a = F.relu(self.decoder.d2(a))
        return torch.tanh(self.decoder.d3(a))

    def decode_multiple(self, state, z=None, num_decode=10):
        if z is None:
            z = ptu.from_numpy(np.random.normal(0, 1, size=(
            state.size(0), num_decode, self.latent_dim))).clamp(-0.5, 0.5)

        a = F.relu(self.decoder.d1(torch.cat(
            [state.unsqueeze(0).repeat(num_decode, 1, 1).permute(1, 0, 2), z],
            2)))
        a = F.relu(self.decoder.d2(a))
        return torch.tanh(self.decoder.d3(a)), self.decoder.d3(a)

class Encoder(nn.Module):
    def __init__(self, obs_dim, latent_dim, action_dim):
        super().__init__()
        self.latent_dim = latent_dim

        self.e1 = torch.nn.Linear(obs_dim + action_dim, 750)
        self.e2 = torch.nn.Linear(750, 750)

        self.mean = torch.nn.Linear(750, self.latent_dim)
        self.log_std = torch.nn.Linear(750, self.latent_dim)

    def forward(self, state, action):
        z = F.relu(self.e1(torch.cat([state, action], 1)))
        z = F.relu(self.e2(z))

        mean = self.mean(z)
        # Clamped for numerical stability
        log_std = self.log_std(z).clamp(-4, 15)
        std = torch.exp(log_std)
        return MultivariateDiagonalNormal(mean, std)

class Decoder(nn.Module):
    def __init__(self, obs_dim, latent_dim, action_dim):
        super().__init__()
        self.latent_dim = latent_dim

        self.d1 = torch.nn.Linear(obs_dim + self.latent_dim, 750)
        self.d2 = torch.nn.Linear(750, 750)
        self.d3 = torch.nn.Linear(750, action_dim)

    def forward(self, state, z):
        a = F.relu(self.d1(torch.cat([state, z], 1)))
        a = F.relu(self.d2(a))
        return Delta(torch.tanh(self.d3(a)))

max_stars_repo_name: google-research/DBAP-algorithm

max_issues_repo_name: google-research/DBAP-algorithm

max_forks_repo_name: google-research/DBAP-algorithm


text:
""" Test the AdaLossScaled wrapper """

import unittest
import numpy as np
import chainer
import chainer.links as L
import chainer.functions as F

from chainer import testing
from chainercv.links import PickableSequentialChain

from ada_loss.chainer_impl.ada_loss_scaled import AdaLossScaled, loss_scaling
from ada_loss.chainer_impl.links.identity_loss_scaling import IdentityLossScalingWrapper
from ada_loss.chainer_impl.links.ada_loss_linear import AdaLossLinear

np.random.seed(42)

class MLP(PickableSequentialChain):
    def __init__(self, n_units, n_out):
        super(MLP, self).__init__()
        with self.init_scope():
            # the size of the inputs to each layer will be inferred
            self.l1 = L.Linear(16, n_units)  # n_in -> n_units
            self.relu1 = lambda x: F.relu(x)  # should be explicit
            self.l2 = L.Linear(n_units, n_units)  # n_units -> n_units
            self.relu2 = lambda x: F.relu(x)  # should be explicit
            self.l3 = L.Linear(n_units, n_out)  # n_units -> n_out

class AdaLossScaledTest(unittest.TestCase):
    """ """

    def test_forward(self):
        """ """
        net1 = MLP(16, 10)
        net2 = AdaLossScaled(net1, init_scale=16.0)

        x = chainer.Variable(np.random.normal(size=(1, 16)).astype("float32"))
        y1 = net1(x)
        y2 = net2(x)
        self.assertTrue(np.allclose(y1.array, y2.array))

    def test_backward(self):
        """ init_scale should be effective """
        net1 = MLP(16, 10)
        net2 = AdaLossScaled(
            net1, init_scale=16.0, transform_functions=False, transforms=None
        )

        x_data = np.random.normal(size=(1, 16)).astype("float32")
        y_data = np.random.normal(size=(1, 10)).astype("float32")

        x1 = chainer.Variable(x_data)
        y1 = net1(x1)
        y1.grad = y_data
        y1.backward()

        x2 = chainer.Variable(x_data)
        y2 = net2(x2)
        y2.grad = y_data
        y2.backward()

        self.assertTrue(np.allclose(x1.grad * 16, x2.grad))

    def test_transform_mlp(self):
        """ Check how transform works """
        np.random.seed(0)

        with chainer.using_config("dtype", "float16"):
            cfg = {
                "loss_scale_method": "fixed",
                "fixed_loss_scale": 2.0,
            }
            net1 = MLP(16, 10)

            x_data = np.random.normal(size=(1, 16)).astype("float16")
            y_data = np.random.normal(size=(1, 10)).astype("float16")

            x = chainer.Variable(x_data)
            y1 = net1(x)
            y1.grad = y_data
            y1.backward()

            x_grad1 = x.grad

            net2 = AdaLossScaled(net1, init_scale=16.0, cfg=cfg, verbose=True)
            x = chainer.Variable(x_data)
            y2 = net2(x)
            self.assertTrue(np.allclose(y1.array, y2.array))

            self.assertTrue(hasattr(net2.link, "l1"))
            self.assertTrue(hasattr(net2.link, "l2"))
            self.assertTrue(hasattr(net2.link, "l3"))
            self.assertIsInstance(getattr(net2.link, "l1"), AdaLossLinear)
            self.assertIsInstance(getattr(net2.link, "l2"), AdaLossLinear)
            self.assertIsInstance(getattr(net2.link, "l3"), AdaLossLinear)

            y2.grad = y_data
            y2.backward()

            self.assertTrue(hasattr(x.grad_var, "loss_scale"))
            # 3 different layers
            self.assertEqual(getattr(x.grad_var, "loss_scale"), 16 * 2 * 2 * 2)

    def test_transform_picked(self):
        """ Test how the transformation performs with picked output """
        np.random.seed(0)

        with chainer.using_config("dtype", "float16"):
            cfg = {
                "loss_scale_method": "fixed",
                "fixed_loss_scale": 2.0,
            }
            net1 = MLP(16, 16)
            net1.pick = ("l1", "l2")

            x_data = np.random.normal(size=(1, 16)).astype("float16")

            x = chainer.Variable(x_data)
            ys1 = net1(x)

            net2 = AdaLossScaled(net1, init_scale=16.0, cfg=cfg, verbose=True)
            x = chainer.Variable(x_data)
            ys2 = net2(x)

            for i in range(len(ys1)):
                self.assertTrue(np.allclose(ys1[i].array, ys2[i].array))

            # try backward propagation
            loss = sum(ys2)
            loss.grad = np.ones_like(loss.array)
            loss.backward()

class IdentityLossScalingTest(unittest.TestCase):
    """ """

    def test_relu(self):
        """ We assume that this loss scale will be propagated to
            the very beginning.  """
        loss_scale = 16.0
        link = chainer.Sequential(
            IdentityLossScalingWrapper(F.relu),
            IdentityLossScalingWrapper(F.relu),
            lambda x: loss_scaling(x, loss_scale),
        )
        x = chainer.Variable(np.random.normal(size=16).astype("float32"))
        y = link(x)
        y.grad = np.random.normal(size=16).astype("float32")
        y.backward()

        # check the propagated back gradient
        self.assertTrue(hasattr(x.grad_var, "loss_scale"))
        self.assertEqual(getattr(x.grad_var, "loss_scale"), loss_scale)

    def test_concat(self):
        """ Test how concatenation works. """
        loss_scale = 16.0
        link = chainer.Sequential(
            IdentityLossScalingWrapper(lambda xs: F.concat(xs, axis=0)),
            # last output should be scaled
            lambda x: loss_scaling(x, loss_scale),
        )

        xs = [
            chainer.Variable(np.random.normal(size=16).astype("float32")),
            chainer.Variable(np.random.normal(size=16).astype("float32")),
        ]
        y = link(xs)
        y.grad = np.random.normal(size=32).astype("float32")
        y.backward()

testing.run_module(__name__, __file__)

max_stars_repo_name: kumasento/gradient-scaling

max_issues_repo_name: kumasento/gradient-scaling

max_forks_repo_name: kumasento/gradient-scaling


text:
import numpy as np
import pandas as pd

data = pd.read_csv('fer2013.csv')

selected_labels = [0, 2, 3, 4, 6]
images_num = np.zeros(7)

training = []
test = []

for category in data['Usage'].unique():
    if category[0] != 'T':
        continue
    category_data = data[data['Usage'] == category]
    samples = category_data['pixels'].values
    labels = category_data['emotion'].values

    for i in range(len(samples)):
        if labels[i] in selected_labels and images_num[labels[i]] < 3000:
            image = np.fromstring(samples[i], dtype=int, sep=' ').reshape(48, 48)
            training.append((image, labels[i]))
            images_num[labels[i]] += 1
        elif labels[i] in selected_labels and images_num[labels[i]] < 3300:
            image = np.fromstring(samples[i], dtype=int, sep=' ').reshape(48, 48)
            test.append((image, labels[i]))
            images_num[labels[i]] += 1

np.save('data/training.npy', training)
np.save('data/test.npy', test)

max_stars_repo_name: Yuxiang-Wei/DeepLearning

max_issues_repo_name: csyxwei/DeepLearning

max_forks_repo_name: csyxwei/DeepLearning


text:
import numpy as np
import pandas as pd

import time
import datetime
import math
import pickle

import ffn

from sklearn import linear_model
from numpy.lib.stride_tricks import as_strided

class Algo(object):
    def __init__(self):
        pass

class CalculateMomentum(Algo):
    def __init__(self, window = 120):
        super(CalculateMomentum, self).__init__()
        self.window = window

    def __call__(self, states):
        # NOTE: MODIFYING STATES HERE WILL BE MODIFYING ORIGINAL VALUES
        # THIS IS DELIBERATE
        states['score'] = states['data'].pct_change(self.window)
        return states

class CalculateInvVol(Algo):
    def __init__(self, window = 120):
        super(CalculateInvVol, self).__init__()
        self.window = window

    def __call__(self, states):

        returns = states['data'].pct_change().fillna(method='ffill')
        vol = 1.0 / returns.rolling(self.window).std()
        vol[np.isinf(vol)] = np.NaN
        vol[vol>1000000] = np.NaN # @HACK/TODO: for too small volatility?
        vols = vol.sum(axis=1)
        weights = vol.divide(vols, axis = 0)

        states['score'] = weights

        return states

class SelectN(Algo):
    def __init__(self, n, ascending = False):
        super(SelectN, self).__init__()
        self.n = n
        self.ascending = ascending

    def __call__(self, states):
        states['selected'] = states['score'].rank(axis=1, ascending=self.ascending) <= self.n
        return states

class SelectNPct(Algo):
    def __init__(self, pct, ascending = False):
        super(SelectNPct, self).__init__()
        self.pct = pct
        self.ascending = ascending

    def __call__(self, states):
        states['selected'] = states['score'].rank(axis=1, ascending=self.ascending).le(states['score'].count(axis=1) * self.pct, axis=0)
        return states

class SelectAll(Algo):
    def __init__(self):
        super(SelectAll, self).__init__()

    def __call__(self, states):
        states['selected'] = ~states['data'].isna()
        return states

class SelectWhere(Algo):
    def __init__(self, cond):
        super(SelectWhere, self).__init__()
        self.cond = cond

    def __call__(self, states):
        states['selected'] = states['data'][self.cond]
        return states

class WeighEqually(Algo):
    def __init__(self):
        super(WeighEqually, self).__init__()

    def __call__(self, states):

        selected = states['selected']
        data = states['data']

        counts = data[selected].count(axis=1)
        counts[counts==0] = np.NaN
        rolling_weights = pd.DataFrame(index=data.index, columns=data.columns)
        rolling_weights[selected] = 1

        rolling_weights = rolling_weights.divide(counts, axis=0)

        states['weights'] = rolling_weights

        return states

class WeighSpecific(Algo):
    def __init__(self, specific_weights):
        super(WeighSpecific, self).__init__()
        self.specific_weights = specific_weights

    def __call__(self, states):

        data = states['data']

        rolling_weights = pd.DataFrame(index=data.index, columns=data.columns)

        for k, w in self.specific_weights.items():
            rolling_weights[k] = w

        states['weights'] = rolling_weights

        return states

class Weigh(Algo):
    def __init__(self, weights):
        super(Weigh, self).__init__()
        self.weights = weights

    def __call__(self, states):

        states['weights'] = self.weights

        return states

class WeighInvVol(Algo):
    def __init__(self, window = 120):
        super(WeighInvVol, self).__init__()
        self.window = window

    def __call__(self, states):

        selected = states['selected']

        returns = states['data'].pct_change().fillna(method='ffill')
        vol = 1.0 / returns.rolling(self.window).std()

        # Below lines to make sure that if an asset is not selected (NaN in data),
        # it will be excluded from calculation.
        vol[np.isnan(selected)] = np.NaN
        vol[selected==False] = np.NaN
        vol[np.isinf(vol)] = np.NaN
        vol[vol>1000000] = np.NaN # @HACK/TODO: for too small volatility?
        vols = vol.sum(axis=1)
        weights = vol.divide(vols, axis = 0)

        states['weights'] = weights

        return states

class WeighBetaHedge(Algo):
    def __init__(self, window = 30, major = '', hedge = ''):
        super(WeighBetaHedge, self).__init__()
        self.window = window
        self.major = major
        self.hedge = hedge

    def __call__(self, states):

        returns = states['data'].pct_change().fillna(method='ffill')
        major = returns[self.major]
        hedge = returns[self.hedge]

        beta = major.rolling(self.window).cov(hedge) / hedge.rolling(self.window).var()

        rolling_weights = pd.DataFrame(index=returns.index, columns=returns.columns)
        rolling_weights[self.major] = 1
        rolling_weights[self.hedge] = -beta

        states['weights'] = rolling_weights

        return states

class TargetVol(Algo):
    def __init__(self, target = 0.2, window = 250, annualization_factor = 252):
        super(TargetVol, self).__init__()
        self.target = target
        self.window = window
        self.annualization_factor = annualization_factor

    def __call__(self, states):

        selected = states['selected']
        data = states['data']

        rolling_weights = pd.DataFrame(index=data.index, columns=data.columns)
        rolling_weights[:] = 1

        returns = data.pct_change()
        vol = returns.rolling(self.window).std() * np.sqrt(self.annualization_factor)

        resize = self.target / vol

        resize[np.isinf(resize)] = np.NaN
        resize[resize>10000] = np.NaN # @HACK/TODO: for too small volatility?

        rolling_weights = rolling_weights * resize

        rolling_weights[np.isnan(selected)] = np.NaN
        rolling_weights[selected==False] = np.NaN

        states['weights'] = rolling_weights

        return states

class WeightCap(Algo):
    def __init__(self, cap = 0.2):
        super(WeightCap, self).__init__()
        self.cap = cap

    def __call__(self, states):
        """
        This only deals with rolling target weights, not the actual weights
        that might changed due to market returns
        """

        w = states['weights'].copy()
        w[w > self.cap] = self.cap
        w[w < -self.cap] = -self.cap

        states['weights'] = w

        return states

class WeightRounding(Algo):
    def __init__(self, threshold = 0.1):
        super(WeightRounding, self).__init__()
        self.threshold = threshold

    def __call__(self, states):
        """
        This only deals with rolling target weights, not the actual weights
        that might changed due to market returns
        """

        w = states['weights'].copy()
        w[w.abs() < self.threshold] = 0
        states['weights'] = w

        return states

class CloseAt(Algo):
    def __init__(self, dates, next_reopen = False):
        super(CloseAt, self).__init__()
        self.dates = dates
        self.next_reopen = next_reopen

    def __call__(self, states):

        weights = states['weights'].copy()

        for date in self.dates:
            date_id = weights.index.get_loc(date)
            for c in weights.columns:
                if self.next_reopen:
                    weights[c].iloc[date_id+1] = weights[c].loc[date]
                weights[c].loc[date] = 0

        states['weights'] = weights

        return states

class Rebalance(Algo):
    def __init__(self, freq = 'M'):
        super(Rebalance, self).__init__()
        self.freq = freq

    def _get_rebalance_dates(self, df, freq='M'):
        if freq == 'M':
            dates1 = pd.Series(df.index.month)
            dates2 = pd.Series(df.index.month).shift(1)
        elif freq == 'Q':
            dates1 = pd.Series(df.index.quarter)
            dates2 = pd.Series(df.index.quarter).shift(1)
        elif freq == 'W':
            dates1 = pd.Series(df.index.week)
            dates2 = pd.Series(df.index.week).shift(1)
        else:
            dates1 = pd.Series(df.index.day)
            dates2 = pd.Series(df.index.day).shift(1)

        is_rebalance = (dates1 != dates2)
        rebalance_dates = list(df.index[is_rebalance])

        return rebalance_dates

    def __call__(self, states):

        rebalance_dates = self._get_rebalance_dates(states['weights'], freq=self.freq)

        weights = pd.DataFrame(index=states['weights'].index, columns=states['weights'].columns)
        weights.loc[rebalance_dates] = states['weights']

        states['weights'] = weights

        return states

max_stars_repo_name: exted/matrix

max_issues_repo_name: exted/matrix

max_forks_repo_name: exted/matrix

reason: import numpy,from numpy