text:
# -*- coding: utf-8 -*-
"""
Created on Thu Mar  8 19:45:39 2018

@author: alnoah
"""

import os
import requests
import numpy as np
import time

from bs4 import BeautifulSoup
from io import BytesIO
from tqdm import trange, tqdm

from astropy.table import Table, join, unique, vstack, hstack
from astropy.coordinates import SkyCoord
from astroquery.irsa_dust import IrsaDust
from astroquery.sdss import SDSS

import casjobs

# Table from Schlafly and Finkbeiner (2011) + IRSA data
extdata = Table.read('sfextdata.tex', format='latex')

def pstarrs(srcids):

    casjobs.add_table(Table([srcids]), 'xmatch_objids')

    filters = ['g', 'r', 'i', 'z', 'y']
    cols_mflags = ''.join(['m.{}Flags, '.format(f) for f in filters])
    cols_mPSFMag = ''.join(['m.{0}MeanPSFMag, m.{0}MeanPSFMagErr, '.format(f)
                            for f in filters])
    cols_mKronMag = ''.join(['m.{0}MeanKronMag, m.{0}MeanKronMagErr, '.format(f)
                             for f in filters])
    cols_mApMag = 'm.rMeanApMag, m.rMeanApMagErr, m.iMeanApMag, m.iMeanApMagErr'

    cols_infoflags = ''.join(['t.{}infoFlag, '.format(f) for f in filters])
    cols_PSFMag = ''.join(['t.{0}PSFMag, t.{0}PSFMagErr, '.format(f)
                           for f in filters])
    cols_KronMag = ''.join(['t.{0}KronMag, t.{0}KronMagErr, '.format(f)
                            for f in filters])
    cols_ApMag = 't.rApMag, t.rApMagErr, t.iApMag, t.iApMagErr'

    # Get stack photometry
    cols = 't.objid, t.bestDetection, {}{}{}{}'
    cols = cols.format(cols_infoflags, cols_PSFMag, cols_KronMag, cols_ApMag)

    qry = 'select {} from MyDB.xmatch_objids s '.format(cols)
    qry += 'inner join StackObjectThin t on s.PSobjID=t.objid '
    qry += 'where t.primaryDetection=1 '
    qry += 'into xmatch_objids_stackPhoto'
    casjobs.run_qry(qry, 'stackPhoto')

    casjobs.get_table('xmatch_objids_stackPhoto')

    photo_stack = Table.read('xmatch_objids_stackPhoto.fits', memmap=True)

    # remove duplicate sources
    photo_stack.add_column(-photo_stack['rPSFMag'], name='mr')
    photo_stack.sort('mr')
    photo_stack.remove_column('mr')
    photo_stack = unique(photo_stack, keys='objid')

    # Get mean photometry
    cols = 'o.objid, o.raMean, o.decMean, o.surveyID, '
    cols += 'o.objInfoFlag, o.qualityFlag, {}{}{}{}'
    cols = cols.format(cols_mflags, cols_mPSFMag, cols_mKronMag, cols_mApMag)

    qry = 'select {} from MyDB.xmatch_objids s '.format(cols)
    qry += 'inner join ObjectThin o on o.objid=s.PSobjID '
    qry += 'inner join MeanObject m on o.objid=m.objid and o.uniquePspsOBid=m.uniquePspsOBid '
    qry += 'into xmatch_objids_meanPhoto'
    casjobs.run_qry(qry, 'meanPhoto')

    casjobs.get_table('xmatch_objids_meanPhoto')

    photo_mean = Table.read('xmatch_objids_meanPhoto.fits', memmap=True)

    # Clean temp files
    casjobs.drop_table('xmatch_objids_stackPhoto')
    casjobs.drop_table('xmatch_objids_meanPhoto')
    casjobs.drop_table('xmatch_objids')
    os.remove('xmatch_objids_meanPhoto.fits')
    os.remove('xmatch_objids_stackPhoto.fits')

    return join(photo_mean, photo_stack, join_type='left', keys='objid')

def pstarrs_clean(photo):
    filters = ['g', 'r', 'i', 'z', 'y']
    msk_ptl = (photo['objInfoFlag'] & 0x00800000) == 0

    good_stack = np.zeros((len(photo), len(filters)))
    good_mean = np.zeros((len(photo), len(filters)))

    for i, f in enumerate(filters) :
        stack_colname = 'stack_{}Mag'.format(f)
        mean_colname = 'mean_{}Mag'.format(f)

        mag = photo[f + 'KronMag']
        mag[msk_ptl] = photo[msk_ptl][f + 'PSFMag']
        magErr = photo[f + 'KronMagErr']
        magErr[msk_ptl] = photo[msk_ptl][f + 'PSFMagErr']

        photo[stack_colname] = mag
        photo[stack_colname + 'Err'] = magErr
        #photo.add_column(Table.Column(mag), name=stack_colname)
        #photo.add_column(Table.Column(magErr), name=stack_colname + 'Err')

        mag = photo[f + 'MeanKronMag']
        mag[msk_ptl] = photo[msk_ptl][f + 'MeanPSFMag']
        magErr = photo[f + 'MeanKronMagErr']
        magErr[msk_ptl] = photo[msk_ptl][f + 'MeanPSFMagErr']

        photo[mean_colname] = mag
        photo[mean_colname + 'Err'] = magErr
        #photo.add_column(Table.Column(mag), name=mean_colname)
        #photo.add_column(Table.Column(magErr), name=mean_colname + 'Err')

        msk_good = np.logical_and(photo[stack_colname] > 0,
                                  photo[stack_colname + 'Err'] > 0)
        good_stack[msk_good,i] = [1]*len(np.where(msk_good))

        msk_good = np.logical_and(photo[mean_colname] > 0,
                                  photo[mean_colname + 'Err'] > 0)
        good_mean[msk_good,i] = [1]*len(np.where(msk_good))

    n_stack = np.sum(good_stack, axis=1)
    n_mean = np.sum(good_mean, axis=1)
    msk_stack = np.logical_and(~photo['bestDetection'].mask,
                               n_stack >= n_mean)
    photo['stack'] = msk_stack
    #photo.add_column(Table.Column(msk_stack), name='stack')

    for f in filters:
        mag = photo['mean_{}Mag'.format(f)]
        mag[msk_stack] = photo[msk_stack]['stack_{}Mag'.format(f)]
        magErr = photo['mean_{}MagErr'.format(f)]
        magErr[msk_stack] = photo[msk_stack]['stack_{}MagErr'.format(f)]

        photo['{}Mag'.format(f)] = mag
        photo['{}MagErr'.format(f)] = magErr
        #photo.add_column(Table.Column(mag), name='{}Mag'.format(f))
        #photo.add_column(Table.Column(magErr), name='{}MagErr'.format(f))

    photo.keep_columns(['objid', 'objInfoFlag', 'qualityFlag', 'stack',
                        'gMag', 'rMag', 'iMag', 'zMag', 'yMag',
                        'gMagErr', 'rMagErr', 'iMagErr', 'zMagErr', 'yMagErr'])

    return photo

def query_wsa(srcids, columns='ra,dec', constraints='sourceID=u.NIRobjID',
              table='lasSource', database='UKIDSSDR10PLUS',
              url='http://wsa.roe.ac.uk:8080/wsa/WSASQL'):

    Table([srcids]).write('temp.fits', format='fits', overwrite=True)
    qry = 'SELECT {} FROM {}, #userTable as u WHERE {}'
    qry = qry.format(columns, table, constraints)
    payload = {'database': database,
               'formaction': 'freeform',
               'uploadSQLFile': '',
               'sqlstmt': qry,
               'iFmt': 'FITS',
               'emailAddress': '',
               'format': 'FITS',
               'compress': 'NONE',
               'rows': 2,
               'timeout': 10800}
    upfile = {'uploadFileToTable': open('temp.fits','rb')}
    r = requests.post(url, files=upfile, data=payload)
    os.remove('temp.fits')

    soup = BeautifulSoup(r.text, "html5lib")
    dl = soup.find('a', id='dl_id')
    r = requests.get(dl['href'])

    return Table.read(BytesIO(r.content), format='fits')

def sdss(srcids, columns='cntr,w1mpro,w2mpro', wsid_column='SDSSID'):

    photo = SDSS.query_sql('select top 10 z, ra, dec, bestObjID')
    print(photo)

def wise(srcids, columns='cntr,w1mpro,w2mpro', wsid_column='WSID'):

    columns = ''.join(['w.{}, '.format(s.replace(' ', ''))
                       for s in columns.split(',')]).rstrip(', ')

#    photo = query_wsa(srcids, columns=columns, constraints='w.cntr=u.WSID',
#                      table='WISE..allwise_sc as w')
    photo = query_wsa(srcids, columns=columns,
                      constraints='w.cntr=u.{}'.format(wsid_column),
                      table='WISE..allwise_sc as w')

    return photo

def tmass(srcids, columns='cntr,w1mpro,w2mpro'):

    columns = ''.join(['t.{}, '.format(s.replace(' ', ''))
                       for s in columns.split(',')]).rstrip(', ')

    photo = query_wsa(srcids, columns=columns, constraints='t.pts_key=u.NIRobjID',
                      table='TWOMASS..twomass_psc as t')
    return photo

def ukidss(srcids, columns='ra,dec', table='lasSource'):

    photo = query_wsa(srcids, columns=columns,
                      constraints='sourceID=u.NIRobjID',
                      table=table)
    return photo

def vista(srcids, columns='ra,dec', table='vhsSource'):

    photo = query_wsa(srcids, columns=columns,
                      constraints='sourceID=u.NIRobjID',
                      table=table, database='VHSDR4',
                      url='http://horus.roe.ac.uk:8080/vdfs/WSASQL')
    return photo

def merge(catalogue, photo_ps, photo_ws, photo_tm, photo_uk, photo_vt):

    # Pan-STARRS photometry
    photo_ps.rename_column('objid', 'PSobjID')
    photo = join(catalogue, photo_ps, join_type='left', keys='PSobjID')

    # NIR photometry
    photo_nir_tm = Table([photo_tm['PTS_KEY'],
                          photo_tm['J_M'], photo_tm['J_MSIGCOM'],
                          photo_tm['H_M'], photo_tm['H_MSIGCOM'],
                          photo_tm['K_M'], photo_tm['K_MSIGCOM'],
                          photo_tm['CC_FLG'], photo_tm['PH_QUAL']],
                         names=['NIRobjID',
                               'JMag', 'JMagErr', 'HMag', 'HMagErr',
                               'KMag', 'KMagErr', 'nir_cc_flags', 'nir_ph_qual'])

    photo_nir_uk = Table([photo_uk['SOURCEID'], photo_uk['J_1PPERRBITS'],
                          photo_uk['HPPERRBITS'], photo_uk['KPPERRBITS'],
                          photo_uk['J_1APERMAG3'], photo_uk['J_1APERMAG3ERR'],
                          photo_uk['HAPERMAG3'], photo_uk['HAPERMAG3ERR'],
                          photo_uk['KAPERMAG3'], photo_uk['KAPERMAG3ERR']],
                         names=['NIRobjID',
                                'Jpperrbit', 'Hpperrbit', 'Kpperrbit',
                                'JMag', 'JMagErr', 'HMag', 'HMagErr',
                                'KMag', 'KMagErr'])

    photo_nir_vt = Table([photo_vt['SOURCEID'], photo_vt['JPPERRBITS'],
                          photo_vt['HPPERRBITS'], photo_vt['KSPPERRBITS'],
                          photo_vt['JAPERMAG3'], photo_vt['JAPERMAG3ERR'],
                          photo_vt['HAPERMAG3'], photo_vt['HAPERMAG3ERR'],
                          photo_vt['KSAPERMAG3'], photo_vt['KSAPERMAG3ERR']],
                         names=['NIRobjID',
                                'Jpperrbit', 'Hpperrbit', 'Kpperrbit',
                                'JMag', 'JMagErr', 'HMag', 'HMagErr',
                                'KMag', 'KMagErr'])

    photo_nir = vstack([photo_nir_tm, photo_nir_uk, photo_nir_vt])

    photo_goodnir = photo[~photo['NIRobjID'].mask]
    photo_badnir = photo[photo['NIRobjID'].mask]
    photo = join(photo_goodnir, photo_nir, join_type='left', keys='NIRobjID')
    photo = vstack([photo, photo_badnir])

    # WISE photometry
    photo_ws.rename_column('CNTR', 'WSID')
    photo_ws.rename_column('W1MPRO', 'W1Mag')
    photo_ws.rename_column('W1SIGMPRO', 'W1MagErr')
    photo_ws.rename_column('W2MPRO', 'W2Mag')
    photo_ws.rename_column('W2SIGMPRO', 'W2MagErr')
    photo_ws.rename_column('CC_FLAGS', 'ws_cc_flags')
    photo_ws.rename_column('PH_QUAL', 'ws_ph_qual')
    photo_ws.remove_columns(['W1SNR', 'W2SNR'])

    photo_goodws = photo[~photo['WSID'].mask]
    photo_badws = photo[photo['WSID'].mask]
    photo = join(photo_goodws, photo_ws, join_type='left', keys='WSID')
    photo = vstack([photo, photo_badws])

    return photo

def get_extinction(coords, filters=['PS1_g', 'PS1_r', 'PS1_i']):

    global extdata
    filters = Table([filters], names=['filter'])
    kk = join(extdata, filters, keys=['filter'])
    AEBV = np.array(kk['AEBV2'])

    t = IrsaDust.get_query_table(coords, section='ebv')
    ebv = np.array(t['ext SFD ref'])

    return ebv*AEBV

def add_extinction(catalogue):

    nsrcs = len(catalogue)
    A = np.full((nsrcs,10), np.nan)
    filters_opt = ['PS1_g', 'PS1_r', 'PS1_i', 'PS1_z', 'PS1_y']
    filters_2mass = ['2MASS_J', '2MASS_H', '2MASS_Ks']
    filters_ukidss = ['UKIRT_J', 'UKIRT_H', 'UKIRT_K']
    filters_wise = ['WISE1', 'WISE2']

    filters_2mass = filters_opt + filters_2mass + filters_wise
    filters_ukidss = filters_opt + filters_ukidss + filters_wise

    coords = SkyCoord(ra=catalogue['posRA'], dec=catalogue['posDec'])
    nirsurvey = catalogue['NIR_SURVEY']

    with open('A.dat', 'w') as f:
        f.write ('#Ag Ar Ai Az Ay AJ AH AK AW1 AW2\n')

        for c,s in tqdm(zip(coords, nirsurvey), total=nsrcs,
                        desc='Adding extinction'):
            if s == '2MASS':
                filters = filters_2mass
            else:
                filters = filters_ukidss

            Al = get_extinction(c, filters=filters)
            src_ext = ' '.join(['{:.3f}'.format(A) for A in Al])
            f.write(src_ext + '\n')

    A = Table.read('A.dat', format='ascii.commented_header')
    catalogue = hstack([catalogue, A], join_type='exact')

    return catalogue

def add_extinction_old(catalogue, getOptical=True, getIR=True,
                       url_ned='http://ned.ipac.caltech.edu/cgi-bin/calc'):

    nsrcs = len(catalogue)
    A = np.full((nsrcs,10), np.nan)
    coords = SkyCoord(ra=catalogue['posRA'], dec=catalogue['posDec'])

    start=0
    for i in trange(nsrcs):
        # Get optical extinction from NED
        if getOptical:
            payload = {'in_csys': 'Equatorial',
                       'in_equinox': 'J2000.0',
                       'obs_epoch': '2000.0',
                       'out_csys': 'Equatorial',
                       'out_equinox': 'J2000.0',
                       'lon': '{:.6f}d'.format(coords[i].ra.value),
                       'lat': '{:.6f}d'.format(coords[i].dec.value)}

            end = time.time()
            if end-start < 1.0 :
                time.sleep(1.0)
            else :
                r = requests.get(url_ned, params=payload)
            start = time.time()

            soup = BeautifulSoup(r.text, 'html5lib')
            ned_table = soup.find('div', id='moreBANDS')
            ned_table = ned_table.find('table')

            j = 0
            for row in ned_table.findAll("tr"):
                cells = row.findAll("td")
                if len(cells) :
                    if cells[0].contents[0] == 'PS1' and j < 5:
                        A[i,j] = cells[3].contents[0]
                        j += 1

        if getIR:
            t = IrsaDust.get_extinction_table(coords[i])

            # NIR extinction
            if catalogue['NIRobjID'][i] > 0:
                if catalogue['NIR_SURVEY'][i] == '2MASS':
                    A[i,5:8] = t['A_SandF'][16:19]

                elif (catalogue['NIR_SURVEY'][i] == 'UKIDSS' or
                      catalogue['NIR_SURVEY'][i]=='VISTA'):
                    A[i,5:8] = t['A_SandF'][13:16]

            # WISE extinction
            if catalogue['WSID'][i] > 0:
                A[i,8:] = t['A_SandF'][23:25]

    A = Table(A, names=['Ag', 'Ar', 'Ai', 'Az', 'Ay',
                        'AJ', 'AH', 'AK', 'AW1', 'AW2'])

    catalogue = hstack([catalogue, A], join_type='exact')

    return catalogue

def define_samples(cat):

    filledcat = cat.filled(-999.)
    # good_opt
    msk_g = np.logical_and(filledcat['gMag'] > 0, filledcat['gMagErr'] > 0)
    msk_r = np.logical_and(filledcat['rMag'] > 0, filledcat['rMagErr'] > 0)
    msk_i = np.logical_and(filledcat['iMag'] > 0, filledcat['iMagErr'] > 0)
    msk_z = np.logical_and(filledcat['zMag'] > 0, filledcat['zMagErr'] > 0)
    msk_y = np.logical_and(filledcat['yMag'] > 0, filledcat['yMagErr'] > 0)
    msk_opt = np.logical_and(np.logical_and(msk_g, msk_r),
                             np.logical_and(msk_i, msk_z))
    msk_opt = np.logical_and(msk_opt, msk_y)

    # good nir
    msk_j = np.logical_and(filledcat['JMag'] > 0, filledcat['JMagErr'] > 0)
    msk_h = np.logical_and(filledcat['HMag'] > 0, filledcat['HMagErr'] > 0)
    msk_k = np.logical_and(filledcat['KMag'] > 0, filledcat['KMagErr'] > 0)
    msk_nir = np.logical_and(msk_j, msk_h)

    # good_mir
    msk_w1 = np.logical_and(filledcat['W1Mag'] > 0, filledcat['W1MagErr'] > 0)
    msk_w2 = np.logical_and(filledcat['W2Mag'] > 0, filledcat['W2MagErr'] > 0)
    msk_mir = np.logical_and(msk_w1, msk_w2)

    sample_10filters = np.logical_and(np.logical_and(msk_opt, msk_mir),
                                      np.logical_and(msk_nir, msk_k))
    sample_10filters = np.logical_and(sample_10filters, cat['sampleXPWN'])

    sample_7filters = np.logical_and(np.logical_or(cat['sampleXPWN'],
                                                   cat['sampleXPW']),
                                     ~sample_10filters)
    sample_7filters = np.logical_and(np.logical_and(msk_opt, msk_mir),
                                     sample_7filters)

    sample_8filters = np.logical_and(np.logical_or(cat['sampleXPWN'],
                                                   cat['sampleXPN']),
                                     np.logical_and(~sample_10filters,
                                                   ~sample_7filters))
    sample_8filters = np.logical_and(np.logical_and(msk_opt, msk_nir),
                                     sample_8filters)

    sample_5filters = np.logical_or(np.logical_or(sample_7filters,
                                                  sample_8filters),
                                    sample_10filters)
    sample_5filters = np.logical_and(msk_opt, ~sample_5filters)

    samples = Table([msk_g, msk_r, msk_i, msk_z, msk_y,
                     msk_j, msk_h, msk_k, msk_w1, msk_w2,
                     msk_opt, msk_nir, msk_mir,
                     sample_5filters, sample_7filters,
                     sample_8filters, sample_10filters],
                    names=['good_g', 'good_r', 'good_i', 'good_z', 'good_y',
                           'good_J', 'good_H', 'good_K', 'good_W1', 'good_W2',
                           'good_opt', 'good_nir', 'good_mir',
                           'sample_5mags', 'sample_7mags',
                           'sample_8mags', 'sample_10mags'])

    cat = hstack([cat, samples], join_type='exact')

    return cat

def add_data(catalogue, getPSphoto=False, getWSphoto=False, getUKphoto=False,
             getVTphoto=False, getTMphoto=False, extinction=False):

    data_folder = '../data/'
    pstarrs_folder = os.path.join(data_folder, 'pstarrs1')
    wise_folder = os.path.join(data_folder, 'allwise')
    tmass_folder = os.path.join(data_folder, '2mass')
    ukidss_folder = os.path.join(data_folder, 'ukidss')
    vista_folder = os.path.join(data_folder, 'vista')

    # Select only sources in the photoz sample
    catalogue = catalogue[catalogue['sample_photoz']]

    photo_ps_file = os.path.join(pstarrs_folder, 'photometry.fits')
    if getPSphoto:
        photo_ps = pstarrs(catalogue['PSobjID'])
        root, ext = os.path.splitext(photo_ps_file)
        photo_ps.write('{}_all{}'.format(root,ext), overwrite=True)

        #photo_ps = Table.read('{}_all{}'.format(root,ext), memmap=True)
        photo_ps = pstarrs_clean(photo_ps)
        photo_ps.write(photo_ps_file, overwrite=True)

    else:
        photo_ps = Table.read(photo_ps_file, memmap=True)

    photo_ws_file = os.path.join(wise_folder, 'photometry.fits')
    if getWSphoto:
        cols = 'cntr, cc_flags, ph_qual, '
        cols += 'w1mpro, w1sigmpro, w1snr, '
        cols += 'w2mpro, w2sigmpro, w2snr'

        msk = ~catalogue['WSID'].mask
        photo_ws = wise(catalogue[msk]['WSID'], columns=cols)

        photo_ws.write(photo_ws_file, overwrite=True)

    else:
        photo_ws = Table.read(photo_ws_file, memmap=True)

    photo_uk_file = os.path.join(ukidss_folder, 'photometry.fits')
    if getUKphoto:
        cols = 'sourceID, mergedClass, '
        cols += 'j_1AperMag3, j_1AperMag3Err, j_1ppErrBits, '
        cols += 'hAperMag3, hAperMag3Err, hppErrBits, '
        cols += 'kAperMag3, kAperMag3Err, kppErrBits'

        msk = np.logical_and(~catalogue['NIRobjID'].mask,
                             catalogue['NIR_SURVEY'] == 'UKIDSS')
        photo_uk = ukidss(catalogue[msk]['NIRobjID'], columns=cols)

        photo_uk.write(photo_uk_file, overwrite=True)

    else:
        photo_uk = Table.read(photo_uk_file, memmap=True)

    photo_vt_file = os.path.join(vista_folder, 'photometry.fits')
    if getVTphoto:
        cols = 'sourceID, mergedClass, '
        cols += 'jAperMag3, jAperMag3Err, jppErrBits, '
        cols += 'hAperMag3, hAperMag3Err, hppErrBits, '
        cols += 'ksAperMag3, ksAperMag3Err, ksppErrBits'

        msk = np.logical_and(~catalogue['NIRobjID'].mask,
                             catalogue['NIR_SURVEY'] == 'VISTA')
        photo_vt = vista(catalogue[msk]['NIRobjID'], columns=cols)

        photo_vt.write(photo_vt_file, overwrite=True)

    else:
        photo_vt = Table.read(photo_vt_file, memmap=True)

    photo_tm_file = os.path.join(tmass_folder, 'photometry.fits')
    if getTMphoto:
        cols = 'pts_key, cc_flg, ph_qual, '
        cols += 'j_m, j_msigcom, j_snr, '
        cols += 'h_m, h_msigcom, h_snr, '
        cols += 'k_m, k_msigcom, k_snr'

        msk = np.logical_and(~catalogue['NIRobjID'].mask,
                             catalogue['NIR_SURVEY'] == '2MASS')
        photo_tm = tmass(catalogue[msk]['NIRobjID'], columns=cols)

        photo_tm.write(photo_tm_file, overwrite=True)

    else:
        photo_tm = Table.read(photo_tm_file, memmap=True)

    catalogue_photo = merge(catalogue, photo_ps, photo_ws, photo_tm, photo_uk, photo_vt)
    if extinction :
        catalogue_photo = add_extinction(catalogue_photo)
    catalogue_photo = define_samples(catalogue_photo)

    return catalogue_photo

max_stars_repo_name: ruizca/xmmpzcat

max_issues_repo_name: ruizca/xmmpzcat

max_forks_repo_name: ruizca/xmmpzcat

reason: import numpy,from astropy


text:
from util import rich3_lyap, getJacobi, Lagrange
import numpy as np

from scipy import integrate

from tfc.utils import MakePlot

from time import process_time as timer
import tqdm

def PCR3BP(t,x):
    m_E = 5.9724e24
    m_M = 7.346e22
    mu = m_M/(m_M + m_E)

    r1 = np.sqrt( (x[0]+mu)**2    + x[1]**2 )
    r2 = np.sqrt( (x[0]-1.+mu)**2 + x[1]**2 )
    Omega_x = x[0]-(1.-mu)*(x[0]+mu)/r1**3-mu*(x[0]-1.+mu)/r2**3
    Omega_y = x[1]-(1.-mu)*x[1]/r1**3-mu*x[1]/r2**3

    # STM
    Uxx = 1.-(1.-mu)*(1./r1**3-3.*(x[0]+mu)**2/r1**5)-(mu)*(1./r2**3-3.*(x[0]-1.+mu)**2/r2**5);
    Uxy = (1.-mu)*3*x[1]*(x[0]+mu)/r1**5 + mu*3*x[1]*(x[0]-1.+mu)/r2**5;
    Uyy = 1.-(1.-mu)*(1./r1**3-3.*x[1]**2/r1**5)-(mu)*(1./r2**3-3.*x[1]**2/r2**5);

    A11 = np.zeros((2,2))
    A12 = np.eye(2)
    A21 = np.array([[Uxx, Uxy],[Uxy, Uyy]])
    A22 = np.array([[0., 2.],[-2., 0.]])

    F_tilda = np.block([[A11, A12],[A21, A22]])

    phi = np.reshape(x[4:20],(4,4))

    dphi = F_tilda @ phi
    dphi = dphi.flatten()

    dy = np.zeros((21, 1))
    dy[0] = x[2]
    dy[1] = x[3]
    dy[2] =  2.*x[3] + Omega_x
    dy[3] = -2.*x[2] + Omega_y
    dy[4:20] = np.reshape(dphi,(16,1))

    dy = r2*dy      # d/dtau
    dy[20] = r2;    # tracks time

    return dy.flatten()
    ## *******************************************************************************************************

def target_lyap(Lpt, Cstar, mu, iterMax):
    ## DESCRIPTION: ******************************************************************************************
    # This algorithm will find a planar lyapunov orbit at the specified energy level.
    #
    # INPUT:
    #           Lpt         - Lagrange Points (1 or 2)
    #           C           - desired Jacobi constant
    #
    # OUTPUT:
    #
    #           x0          - initial state
    #           period      - period [sec]
    #           period_tau  - period [regularized]
    #           res         - energy residual
    #*******************************************************************************************************
    tol     = 2.220446049250313e-14
    tol_sym = 2.220446049250313e-14
    tol_res = 2.220446049250313e-14

    # find initial guess with richardson
    _, rvi, period, Ax = rich3_lyap(mu, 0.2, Lpt)
    x0 = np.hstack(( rvi[0,0:2], rvi[0,3:5] ))

    # find lyapunov orbit with differential corrector
    rvi = np.reshape(np.hstack((x0[0:2], 0., x0[2:4], 0.)),(1, 6))
    C_desired = getJacobi(rvi,mu)
    dC0 = (Cstar - C_desired)

    x0, period, period_tau, _ = diff_cor_lyap_p(x0, period, Ax, mu, C_desired)

    ## Go through loop
    crit_iters = 5.
    step_up = 3./2.
    step_down = 2./3.
    max_factor = 1.
    min_factor = 1./iterMax
    factor = min_factor*2.

    for i in range(0,iterMax):
        # Target energy
        dC = factor*dC0;
        if np.abs(dC) > np.abs(Cstar-C_desired):
            dC = Cstar - C_desired

        C_desired = C_desired + dC

        # Find lyapunov orbit with diferential corrector
        x0, period, period_tau, count = diff_cor_lyap_p(x0[0:4], period, Ax, mu, C_desired)

        # Adapt step
        if count > crit_iters:
            factor = factor*step_down

        if count < crit_iters:
            factor = factor*step_up

        if factor >= max_factor:
            factor = max_factor

        if factor <= min_factor:
            factor = min_factor

        rvi = np.reshape(np.hstack((x0[0:2], 0., x0[2:4], 0.)),(1, 6))
        C = getJacobi(rvi,mu)

        print(C)

        res =  abs(Cstar - C);
        if res <= tol_res:
            break

    if res > tol_res:
        print('Residual = ' + str(res))
    else:
        print('LYAP FOUND! \n')

    return x0, period, period_tau, res
    ## *******************************************************************************************************

def diff_cor_lyap_p(x0, T, Ax, mu, C_desired):
    ## DESCRIPTION: ******************************************************************************************
    # This algorithm will generate initial conditions lying on periodic orbit
    #
    # INPUT: Data from rich3, mass parameter
    #
    # OUTPUT:  Initial condition and period of Halo orbit (in seconds and tau)
    #
    def halo_e(t,y): return y[1]
    halo_e.direction = -1
    halo_e.terminal  = True

    tol     = 2.220446049250313e-14
    tol_sym = 2.220446049250313e-14
    tol_res = 1e-12

    mu1 = 1. - mu
    mu2 = mu

    phi_0 = np.ndarray.flatten(np.eye(4))
    y_init = np.hstack(( x0, phi_0, 0.))

    # Integrate
    NotPeriodic = True
    d2 = np.sqrt( (x0[0]-1.+mu)**2 + x0[1]**2 ) # time regularization
    tau_max = 2.*T/d2 # guess of maximum time

    count = 0
    while NotPeriodic:
        sol = integrate.solve_ivp(PCR3BP, [0., tau_max], y_init, method='DOP853', events = halo_e, rtol=tol, atol=tol)

        t = sol.t
        y = sol.y

        X  = y[0,-1]
        Y  = y[1,-1]
        Xd = y[2,-1]
        Yd = y[3,-1]

        # check conditons of ''bad'' solution
        if np.abs(Y) >= 1e-3:
            NotPeriodic = False
            print('Integration time exceeded!! \n')

        L = Lagrange(mu)
        d_norm2 = 1. - mu - L[0,0]
        if np.linalg.norm(y_init[0:2] - sol.y[0:2,-1]) >= 100.*Ax*d_norm2: # Stray too far
            NotPeriodic = False
            print('ESCAPED!! \n')

        # Check if symmetry conditons met
        rvi = np.array([X, Y, 0., Xd, Yd, 0.])
        C = getJacobi(rvi,mu)
        if np.abs(Xd) < tol_sym and np.abs(C-C_desired) < tol_sym:
            NotPeriodic = False
            # print('Converged! \n')
        else:
            r1 = np.sqrt( (X+mu)**2 + Y**2 )
            r2 = np.sqrt( (X-1.+mu)**2 + Y**2 )
            Omega_x = X - (1. - mu ) * (X + mu)/r1**3-mu*(X-1.+mu)/r2**3
            xdd = 2.*Yd + Omega_x;

            ## Check energy
            # Get partials (in terms of tau)
            x0  = sol.y[0,0]
            y0  = sol.y[1,0]
            vx0 = sol.y[2,0]
            vy0 = sol.y[3,0]
            den1 = ((x0-mu1)**2 + y0**2)**(1.5)
            den2 = ((x0+mu2)**2 + y0**2)**(1.5)
            C_x = 2.*x0 - 2.*mu2*(x0-mu1)/den1 - 2.*mu1*(x0+mu2)/den2;
            C_ydot = -2.*vy0;

            # Calculate M
            M_2_0_a = np.array([[y[12,-1], y[15,-1]],[C_x, C_ydot]])
            M_2_0_b = 1./Yd * np.array([[xdd],[0.]]) @ np.array([[y[8,-1], y[11,-1]]])
            M_2_0 = M_2_0_a - M_2_0_b

            ## Energy Constrained Approach
            b = -np.array([[y[2,-1]],[-C_desired + C]])

            q =  np.linalg.pinv(M_2_0) @ b
            y_init[0] = y_init[0] + q[0]
            y_init[3] = y_init[3] + q[1]

        x0 = y[:,0]
        period_tau = 2.*t[-1]
        period = 2.*y[-1,-1]

        if np.abs(period) == 0:
            print('FALSE CONVERGENCE LIKELY!!')
            period = np.abs(T);

        count += 1

    return x0, period, period_tau, count
    ## *******************************************************************************************************
## Start files
import pickle

file1 ='Lyap_CP_L1'
file2 ='Lyap_CP_L2'
## TEST PARAMETERS: ***************************************************
sol1 = pickle.load(open('data/' + file1 + '.pickle','rb'))
sol2 = pickle.load(open('data/' + file2 + '.pickle','rb'))
tfc = {'L1':sol1['L1'], 'L2':sol2['L2'], }

m_E = 5.9724e24
m_M = 7.346e22
mu = m_M/(m_M + m_E)
tol     = 2.220446049250313e-14

Lpt = 'L2'
Cstar = tfc[Lpt]['C']

if Lpt == 'L1':
    _, rvi, period, Ax = rich3_lyap(mu, 0.05, 1)
else:
    _, rvi, period, Ax = rich3_lyap(mu, 0.05, 2)

# find initial guess with richardson
x0 = np.hstack(( rvi[0,0:2], rvi[0,3:5] ))

# find lyapunov orbit with differential corrector
rvi = np.reshape(np.hstack((x0[0:2], 0., x0[2:4], 0.)),(1, 6))

dif = {'res':np.zeros(len(Cstar)),\
       'time':np.zeros(len(Cstar)),\
       'C':np.zeros(len(Cstar))}

for i in tqdm.trange(len(Cstar)):

    start = timer()
    x0, period, period_tau, _ = diff_cor_lyap_p(x0[0:4], period, Ax, mu, Cstar[i])
    dif['time'][i] = timer() - start

    sol = integrate.solve_ivp(PCR3BP, [0, period_tau], x0, method='DOP853', rtol=tol, atol=tol) #, t_eval=tspan)

    resState = np.abs(sol.y[0:4,0] - sol.y[0:4,-1])

    X  = sol.y[0,:]
    Y  = sol.y[1,:]
    Xd = sol.y[2,:]
    Yd = sol.y[3,:]

    ctemp = np.zeros(len(sol.t))
    for j in range(len(sol.t)):
        ctemp[j] = getJacobi(np.array([X[j],Y[j],0.,Xd[j],Yd[j],0.]),mu)

    resJc = np.abs(Cstar[i] - ctemp)

    dif['res'][i]   = np.max(np.hstack((resState, resJc)))
    dif['C'][i]     = Cstar[i]

## END ******************************************************************************************************
# import pickle
# with open('data/' + 'diffcor' + str(Lpt) + '.pickle', 'wb') as handle:
#     pickle.dump(dif, handle)

ext: py

max_stars_repo_name: leakec/tfc

max_issues_repo_name: leakec/tfc

max_forks_repo_name: leakec/tfc

reason: import numpy,from scipy


text:
# Copy task for DNC
import os
import torch
from torch import nn
from torch import optim

import numpy as np
import random

from DNC.dnc import DNC_Module              # Importing DNC implementation

class task_copy():

    def __init__(self):
        self.name = "copy_task"
        self.controller_size = 100
        self.controller_layers = 1
        self.num_read_heads = 2
        self.num_write_heads = 1
        self.sequence_width = 8
        self.sequence_min_len = 5
        self.sequence_max_len = 20   # Default: 20
        self.memory_N = 128
        self.memory_M = 20
        self.num_batches = 10000
        self.batch_size = 1
        self.rmsprop_lr = 1e-4
        self.rmsprop_momentum = 0.9
        self.rmsprop_alpha = 0.95
        self.machine = None
        self.loss = None
        self.optimizer = None

    def get_task_name(self):
        return self.name

    def _data_maker(self, num_batches, batch_size, seq_width, min_len, max_len):    # Generates data for copy task
        # The input data vector will be of size (num_data_rows x batch_size x num_data_columns)
        #
        # 1 1 1 0 1  | 1 1 0 1 0 | 1 1 1 0 1 | 1 0 1 1 1
        # 0 0 1 0 1  | 0 1 0 1 1 | 0 1 0 0 1 | 0 0 1 1 0
        # 0 1 1 0 1  | 1 1 0 0 0 | 1 0 1 0 1 | 0 0 1 1 0
        #
        # Above is the example of data. num_data_rows = 3, num_data_columns = 5, batch_size = 4
        #
        # At a time we will give each row strip to the NTM for prediction as shown below. Therefore input size for one interaction will be (batch_size x num_data_columns)
        #
        # 1 1 1 0 1  | 1 1 0 1 0 | 1 1 1 0 1 | 1 0 1 1 1

        for batch_num in range(num_batches+1):
            # All batches have the same sequence length
            seq_len = random.randint(min_len, max_len)
            seq = np.random.binomial(1, 0.5, (seq_len, batch_size, seq_width))  # Here, seq_len = num_data_rows and seq_width = num_data_columns
            seq = torch.from_numpy(seq)

            # The input includes an additional channel used for the delimiter
            inp = torch.zeros(seq_len + 1, batch_size, seq_width + 1)
            inp[:seq_len, :, :seq_width] = seq
            inp[seq_len, :, seq_width] = 1.0 # delimiter in our control channel
            outp = seq.clone()

            yield batch_num+1, inp.float(), outp.float()

    def init_dnc(self):
        self.machine = DNC_Module(self.sequence_width + 1, self.sequence_width, self.controller_size, self.controller_layers, self.num_read_heads, self.num_write_heads, self.memory_N, self.memory_M)

    def init_loss(self):
        self.loss = nn.BCEWithLogitsLoss()  # Binary Cross Entropy Loss -> Sigmoid Activation + Cross Entropy Loss

    def init_optimizer(self):
        self.optimizer = optim.RMSprop(self.machine.parameters(), momentum = self.rmsprop_momentum, alpha = self.rmsprop_alpha, lr = self.rmsprop_lr)

    def calc_loss(self, Y_pred, Y):
        return self.loss(Y_pred, Y)

    def get_sample_data(self):  # Sample data for Testing
        batch_size = 1
        seq_len = random.randint(self.sequence_min_len, self.sequence_max_len)
        seq = np.random.binomial(1, 0.5, (seq_len, batch_size, self.sequence_width))  # Here, seq_len = num_data_rows and seq_width = num_data_columns
        seq = torch.from_numpy(seq)

        # The input includes an additional channel used for the delimiter
        inp = torch.zeros(seq_len + 1, batch_size, self.sequence_width + 1)
        inp[:seq_len, :, :self.sequence_width] = seq
        inp[seq_len, :, self.sequence_width] = 1.0 # delimiter in our control channel
        outp = seq.clone()
        return inp.float(), outp.float()

    def calc_cost(self, Y_out, Y, batch_size):
        y_out_binarized = torch.sigmoid(Y_out).clone().data
        y_out_binarized.apply_(lambda x: 0 if x < 0.5 else 1)

        cost = torch.sum(torch.abs(y_out_binarized - Y.data))
        return cost.item()/batch_size

    def get_training_data(self):
        return self._data_maker(self.num_batches, self.batch_size, self.sequence_width, self.sequence_min_len, self.sequence_max_len)

    def test_model(self):
        self.machine.initialization(self.batch_size)    # Initializing states
        X, Y = self.get_sample_data()
        Y_out = torch.zeros(Y.shape)

        # Feeding the DNC network all the data first and then predicting output
        # by giving zero vector as input and previous read states and hidden vector
        # and thus training vector this way to give outputs matching the labels

        for i in range(X.shape[0]):
            self.machine(X[i])

        for i in range(Y.shape[0]):
            Y_out[i, :, :], _ = self.machine()

        loss = self.calc_loss(Y_out, Y)
        cost = self.calc_cost(Y_out, Y, self.batch_size)    # The cost is the number of error bits per sequence

        print("\n\nTest Data - Loss: " + str(loss.item()) + ", Cost: " + str(cost))

        X.squeeze(1)
        Y.squeeze(1)
        Y_out = torch.sigmoid(Y_out.squeeze(1))

        print("\n------Input---------\n")
        print(X.data)
        print("\n------Labels---------\n")
        print(Y.data)
        print("\n------Output---------")
        print((Y_out.data).apply_(lambda x: 0 if x < 0.5 else 1))
        print("\n")

        return loss.item(), cost, X, Y, Y_out

    def clip_grads(self):       # Clipping gradients for stability
        """Gradient clipping to the range [10, 10]."""
        parameters = list(filter(lambda p: p.grad is not None, self.machine.parameters()))
        for p in parameters:
            p.grad.data.clamp_(-10, 10)

    def train_model(self):
        # Here, the model is optimized using BCE Loss, however, it is evaluated using Number of error bits in predction and actual labels (cost)
        loss_list = []
        cost_list = []
        seq_length = []

        if (self.num_batches/10) > 0:
            model_save_interval = (self.num_batches/10)
        else:
            model_save_interval = 1

        for batch_num, X, Y in self.get_training_data():

            if batch_num>self.num_batches:
                break

            self.optimizer.zero_grad()                      # Making old gradients zero before calculating the fresh ones
            self.machine.initialization(self.batch_size)    # Initializing states
            Y_out = torch.zeros(Y.shape)

            # Feeding the NTM network all the data first and then predicting output
            # by giving zero vector as input and previous read states and hidden vector
            # and thus training vector this way to give outputs matching the labels

            for i in range(X.shape[0]):
                self.machine(X[i])

            for i in range(Y.shape[0]):
                Y_out[i, :, :], _ = self.machine()

            loss = self.calc_loss(Y_out, Y)
            loss.backward()
            self.clip_grads()
            self.optimizer.step()

            # The cost is the number of error bits per sequence
            cost = self.calc_cost(Y_out, Y, self.batch_size)

            loss_list += [loss.item()]
            cost_list += [cost]
            seq_length += [Y.shape[0]]

            if batch_num % model_save_interval == 0:
                self.save_model(batch_num)

            print("Batch: " + str(batch_num) + "/" + str(self.num_batches) + ", Loss: " + str(loss.item()) + ", Cost: " + str(cost) + ", Sequence Length: " + str(Y.shape[0]))

    def save_model(self, curr_epoch):
        if not os.path.exists("Saved_Models/" + self.name):
            os.mkdir("Saved_Models/" + self.name)
        state_dic = {'task_name': self.name, 'start_epoch': curr_epoch + 1, 'state_dict': self.machine.state_dict(), 'optimizer_dic' : self.optimizer.state_dict()}
        filename = "Saved_Models/" + self.name + "/" + self.name + "_" + str(curr_epoch) + "_saved_model.pth.tar"
        torch.save(state_dic, filename)

    def load_model(self, option, epoch):
        path = "Saved_Models/" + self.name + "/" + self.name + "_" + str(epoch) + "_saved_model.pth.tar"
        if option == 1:             # Loading for training
            checkpoint = torch.load(path)
            self.machine.load_state_dict(checkpoint['state_dict'])
            self.optimizer.load_state_dict(checkpoint['optimizer_dic'])
        else:                       # Loading for testing
            checkpoint = torch.load(path)
            self.machine.load_state_dict(checkpoint['state_dict'])
            self.machine.eval()

reason: import numpy


text:
#!/usr/bin/env python

import ctypes
import numpy

from . import libfftwf

FFTW_BACKWARD = 1
FFTW_ESTIMATE = 1 << 6
FFTW_FORWARD = -1
FFTW_MEASURE = 0

class FFTWFComplex(ctypes.Structure):
    '''typedef float fftwf_complex[2];'''
    _fields_ = [('r', ctypes.c_float), ('i', ctypes.c_float)]

class FFTWFPlanStructure(ctypes.Structure):
    '''
    struct fftwf_plan_s {
        plan *pln;
        problem *prb;
        int sign;
    };
    '''
    _fields_ = [
        ('pln', ctypes.c_void_p),
        ('prb', ctypes.c_void_p),
        ('sign', ctypes.c_int),
    ]

'''
typedef struct fftwf_plan_s *fftwf_plan;
'''
FFTWFPlan = ctypes.POINTER(FFTWFPlanStructure)

'''
fftwf_plan fftwf_plan_dft_1d(int, fftwf_complex *, fftwf_complex *,
    int, unsigned);
'''
fftwf_plan_dft_1d = libfftwf.fftwf_plan_dft_1d
fftwf_plan_dft_1d.restype = FFTWFPlan
fftwf_plan_dft_1d.argtypes = [
    ctypes.c_int,
    ctypes.POINTER(FFTWFComplex),
    ctypes.POINTER(FFTWFComplex),
    ctypes.c_int,
    ctypes.c_uint,
]

'''
fftwf_plan fftwf_plan_dft_r2c_1d(int, float *, fftwf_complex *, unsigned);
'''
fftwf_plan_dft_r2c_1d = libfftwf.fftwf_plan_dft_r2c_1d
fftwf_plan_dft_r2c_1d.restype = FFTWFPlan
fftwf_plan_dft_r2c_1d.argtypes = [
    ctypes.c_int,
    ctypes.POINTER(ctypes.c_float),
    ctypes.POINTER(FFTWFComplex),
    ctypes.c_uint,
]

'''
fftwf_plan fftwf_plan_dft_c2r_1d(int, fftwf_complex *, float *, unsigned);
'''
fftwf_plan_dft_c2r_1d = libfftwf.fftwf_plan_dft_c2r_1d
fftwf_plan_dft_c2r_1d.restype = FFTWFPlan
fftwf_plan_dft_c2r_1d.argtypes = [
    ctypes.c_int,
    ctypes.POINTER(FFTWFComplex),
    ctypes.POINTER(ctypes.c_float),
    ctypes.c_uint,
]

'''
void fftwf_destroy_plan(fftwf_plan);
'''
fftwf_destroy_plan = libfftwf.fftwf_destroy_plan
fftwf_destroy_plan.restype = None
fftwf_destroy_plan.argtypes = [FFTWFPlan]

'''
void fftwf_execute(const fftwf_plan);
'''
fftwf_execute = libfftwf.fftwf_execute
fftwf_execute.restype = None
fftwf_execute.argtypes = [FFTWFPlan]

'''
void fftwf_execute_dft(const fftwf_plan, fftwf_complex *, fftwf_complex *);
'''
fftwf_execute_dft = libfftwf.fftwf_execute_dft
fftwf_execute_dft.restype = None
fftwf_execute_dft.argtypes = [
    FFTWFPlan,
    ctypes.POINTER(FFTWFComplex),
    ctypes.POINTER(FFTWFComplex),
]

'''
void fftwf_execute_dft_r2c(const fftwf_plan, float *, fftwf_complex *);
'''
fftwf_execute_dft_r2c = libfftwf.fftwf_execute_dft_r2c
fftwf_execute_dft_r2c.restype = None
fftwf_execute_dft_r2c.argtypes = [
    FFTWFPlan,
    ctypes.POINTER(ctypes.c_float),
    ctypes.POINTER(FFTWFComplex),
]

'''
void fftwf_execute_dft_c2r(const fftwf_plan, fftwf_complex *, float *);
'''
fftwf_execute_dft_c2r = libfftwf.fftwf_execute_dft_c2r
fftwf_execute_dft_c2r.restype = None
fftwf_execute_dft_c2r.argtypes = [
    FFTWFPlan,
    ctypes.POINTER(FFTWFComplex),
    ctypes.POINTER(ctypes.c_float),
]

def data(x):
    if not isinstance(x, numpy.ndarray):
        raise ValueError
    if x.dtype not in [numpy.float32, numpy.complex64]:
        raise ValueError
    if x.dtype == numpy.float32:
        return ctypes.cast(x.ctypes.data, ctypes.POINTER(ctypes.c_float))
    if x.dtype == numpy.complex64:
        return ctypes.cast(x.ctypes.data, ctypes.POINTER(FFTWFComplex))

class FFT(object):

    dtypes = {
        'c2c': (numpy.complex64, numpy.complex64),
        'r2c': (numpy.float32, numpy.complex64),
        'c2r': (numpy.complex64, numpy.float32),
    }

    def __init__(self):
        self.plans = {}
        self.destroy_plan = fftwf_destroy_plan

    def plan(self, x, X):
        if not isinstance(x, numpy.ndarray):
            raise ValueError
        if not isinstance(X, numpy.ndarray):
            raise ValueError
        if (x.dtype, X.dtype) not in self.dtypes.values():
            raise ValueError
        n = x.size
        flags = FFTW_ESTIMATE
        if (x.dtype, X.dtype) == self.dtypes['r2c']:
            return (
                fftwf_plan_dft_r2c_1d(n, data(x), data(X), flags),
                fftwf_plan_dft_c2r_1d(n, data(X), data(x), flags),
            )
        if (x.dtype, X.dtype) == self.dtypes['c2c']:
            return (
                fftwf_plan_dft_1d(n, data(x), data(X), FFTW_FORWARD, flags),
                fftwf_plan_dft_1d(n, data(x), data(X), FFTW_BACKWARD, flags),
            )

    def fft(self, x):
        X = numpy.zeros(x.size, dtype=numpy.complex64)
        if not (x.dtype, X.dtype) == self.dtypes['c2c']:
            raise ValueError
        if (x.dtype, X.dtype, x.size) not in self.plans:
            self.plans[(x.dtype, X.dtype, x.size)] = self.plan(x, X)
        plan = self.plans[(x.dtype, X.dtype, x.size)][0]
        fftwf_execute_dft(plan, data(x), data(X))
        return X

    def ifft(self, X):
        x = numpy.zeros(X.size, dtype=numpy.complex64)
        if not (X.dtype, x.dtype) == self.dtypes['c2c']:
            raise ValueError
        if (x.dtype, X.dtype, x.size) not in self.plans:
            self.plans[(x.dtype, X.dtype, x.size)] = self.plan(x, X)
        plan = self.plans[(x.dtype, X.dtype, x.size)][1]
        fftwf_execute_dft(plan, data(X), data(x))
        x *= 1.0 / x.size
        return x

    def rfft(self, x):
        X = numpy.zeros(x.size, dtype=numpy.complex64)
        if not (x.dtype, X.dtype) == self.dtypes['r2c']:
            raise ValueError
        if (x.dtype, X.dtype, x.size) not in self.plans:
            self.plans[(x.dtype, X.dtype, x.size)] = self.plan(x, X)
        plan = self.plans[(x.dtype, X.dtype, x.size)][0]
        fftwf_execute_dft_r2c(plan, data(x), data(X))
        return X

    def irfft(self, X, n=None):
        x = numpy.zeros(X.size, dtype=numpy.float32)
        if not (X.dtype, x.dtype) == self.dtypes['c2r']:
            raise ValueError
        if (x.dtype, X.dtype, x.size) not in self.plans:
            self.plans[(x.dtype, X.dtype, x.size)] = self.plan(x, X)
        plan = self.plans[(x.dtype, X.dtype, x.size)][1]
        fftwf_execute_dft_c2r(plan, data(X), data(x))
        x *= 1.0 / x.size
        return x

    def __del__(self):
        for plan in self.plans:
            self.destroy_plan(self.plans[plan][0])
            self.destroy_plan(self.plans[plan][1])
        pass

if __name__ == '__main__':
    pass

max_stars_repo_name: eliteraspberries/python-libfftw

max_issues_repo_name: eliteraspberries/python-libfftw

max_forks_repo_name: eliteraspberries/python-libfftw


text:
from tkinter import *
from tkinter import filedialog
from PIL import ImageTk, Image
import numpy as np
import json
import os

root = Tk()
root.title('GUI')
root.geometry('800x800')
root.resizable(0, 0)

picName = StringVar()
files = []
idx = 0
fileName = ""
jsonData = None

def checkPic(fn):
    fn = fn.split('.')
    if len(fn) < 2:
        return False
    if not fn[1] in ('npy'):
        return False
    return True

def checkIdx(x):
    if x < 0:
        return len(files) - 1
    if x >= len(files):
        return 0
    return x

def getImage(file):
    img = np.load(file).astype(np.uint8)
    return ImageTk.PhotoImage(Image.fromarray(img))

def callback():
    global picName, files, idx, fileName, picName, curPic
    fileName = filedialog.askdirectory()
    files = os.listdir(fileName)
    assert files
    file = os.path.join(fileName, files[idx])
    while not checkPic(file):
        idx = checkIdx(idx + 1)
        file = os.path.join(fileName, files[idx])
    picName.set(files[idx])
    canvas.delete('all')
    img = getImage(file)
    image = canvas.create_image(0, 0, anchor='nw', image=img)
    root.mainloop()

def getJson():
    global jsonData
    jsonFile = filedialog.askopenfilename(
        filetypes=[("json文件", [".json"])])
    with open(jsonFile, 'r') as f:
        jsonData = json.load(f)
        if not 'all' in jsonFile:
            jsonData = jsonData['data']

def getPoints(fn):
    fn = fn.split('.')[0]
    for item in jsonData:
        if item['id'] == fn:
            return item['point']
    return None

def nextOne(dx):
    global idx, files, fileName, picName
    idx = checkIdx(idx + dx)
    if not files:
        return
    file = os.path.join(fileName, files[idx])
    while not checkPic(file):
        idx = checkIdx(idx + dx)
        file = os.path.join(fileName, files[idx])
    picName.set(files[idx])
    canvas.delete('all')
    img = getImage(file)
    image = canvas.create_image(0, 0, anchor='nw', image=img)
    if jsonData:
        points = getPoints(files[idx])
        points = sorted(points, key=lambda x: x['coord'][1])
        for i in range(len(points)):
            point = points[i]
            xx, yy = point['coord'][0], point['coord'][1]
            canvas.create_oval(xx - 2, yy - 2, xx + 2, yy + 2, fill='red')
            if i == 0:
                p1 = points[i + 1]
                x1, y1 = p1['coord'][0], p1['coord'][1]
                h = (y1 - yy) * 2
                w = h * 1.4
            if 0 < i < len(points) - 1:
                p1, p2 = points[i + 1], points[i - 1]
                x1, y1 = p1['coord'][0], p1['coord'][1]
                x2, y2 = p2['coord'][0], p2['coord'][1]
                h = max(y1 - yy, yy - y2) * 2
                w = h * 1.4
            if i == len(points) - 1:
                p2 = points[i - 1]
                x2, y2 = p2['coord'][0], p2['coord'][1]
                h = (yy - y2) * 2
                w = h * 1.4

            canvas.create_line(xx - w / 2, yy - h / 2, xx - w / 2, \
                               yy + h / 2, xx + w / 2, yy + h / 2, xx + w / 2,
                               yy - h / 2, xx - w / 2, yy - h / 2, fill='red')

    root.mainloop()

Button(root, text="选择npy文件夹..", command=callback).pack()
Button(root, text="选择json文件..", command=getJson).pack()
Button(root, text="上一张", command=lambda: nextOne(-1)).pack()
Button(root, text="下一张", command=lambda: nextOne(1)).pack()
Label(root, textvar=picName).pack()

W, H = 780, 780
canvas = Canvas(root, width=W, height=H, bg='white')
canvas.pack()

root.mainloop()

max_stars_repo_name: ZhengJun-AI/Spinal-Disease-Detection

max_issues_repo_name: ZhengJun-AI/Spinal-Disease-Detection

max_forks_repo_name: ZhengJun-AI/Spinal-Disease-Detection


text:
import numpy as np
import pylibkriging as lk
import pytest

@pytest.debug_only
@pytest.mark.parametrize("n", [1, 2, 4, 5, 8, 16, 17, 32, 64, 128, 256])
#@pytest.mark.parametrize("n", [17])
def test_two_side_carma_binding_should_be_valid_for_any_size(n):
    a = np.arange(0, n, 1)
    b = np.arange(1, n + 1, 1)

    result = lk.two_side_carma_binding(a, b)
    ref = a + b
    assert np.linalg.norm(result - ref) < 1e-16, f"Bad result\nresult={result}\nref={ref}\ndiff={result - ref}"

max_stars_repo_name: MASCOTNUM/libKriging

max_issues_repo_name: libKriging/libKriging

max_forks_repo_name: MASCOTNUM/libKriging


text:
#####################################################################################
# Credit for the original Navier-Stokes dataset                                     #
#                                                                                #
# Projet: https://github.com/zongyi-li/fourier_neural_operator                      #
# Copyright (c) 2020 Zongyi Li                                                      #
# Licence: https://github.com/zongyi-li/fourier_neural_operator/blob/master/LICENSE #
#####################################################################################

import torch
import numpy as np
from torch.utils.data import Dataset
from scipy.integrate import solve_ivp
from functools import partial
import math
import matplotlib.pyplot as plt
import matplotlib
from timeit import default_timer
import scipy.io
import shelve

class GaussianRF(object):

    def __init__(self, dim, size, alpha=2, tau=3, sigma=None, boundary="periodic"):

        self.dim = dim

        if sigma is None:
            sigma = tau**(0.5*(2*alpha - self.dim))

        k_max = size//2

        if dim == 1:
            k = torch.cat((torch.arange(start=0, end=k_max, step=1), \
                           torch.arange(start=-k_max, end=0, step=1)), 0)

            self.sqrt_eig = size*math.sqrt(2.0)*sigma*((4*(math.pi**2)*(k**2) + tau**2)**(-alpha/2.0))
            self.sqrt_eig[0] = 0.0

        elif dim == 2:
            wavenumers = torch.cat((torch.arange(start=0, end=k_max, step=1), \
                                    torch.arange(start=-k_max, end=0, step=1)), 0).repeat(size,1)

            k_x = wavenumers.transpose(0,1)
            k_y = wavenumers

            self.sqrt_eig = (size**2)*math.sqrt(2.0)*sigma*((4*(math.pi**2)*(k_x**2 + k_y**2) + tau**2)**(-alpha/2.0))
            self.sqrt_eig[0,0] = 0.0

        elif dim == 3:
            wavenumers = torch.cat((torch.arange(start=0, end=k_max, step=1), \
                                    torch.arange(start=-k_max, end=0, step=1)), 0).repeat(size,size,1)

            k_x = wavenumers.transpose(1,2)
            k_y = wavenumers
            k_z = wavenumers.transpose(0,2)

            self.sqrt_eig = (size**3)*math.sqrt(2.0)*sigma*((4*(math.pi**2)*(k_x**2 + k_y**2 + k_z**2) + tau**2)**(-alpha/2.0))
            self.sqrt_eig[0,0,0] = 0.0

        self.size = []
        for j in range(self.dim):
            self.size.append(size)

        self.size = tuple(self.size)

    def sample(self):

        coeff = torch.randn(*self.size, dtype=torch.cfloat)
        coeff = self.sqrt_eig * coeff
        u = torch.fft.ifftn(coeff)
        u = u.real

        return u

class NavierStokesDataset(Dataset):

    def __init__(self, num_traj_per_env, size, time_horizon, dt_eval, params, buffer_filepath=None, group='train'):
        super().__init__()
        self.size = int(size)
        tt = torch.linspace(0, 1, self.size+1)[0:-1]
        X,Y = torch.meshgrid(tt, tt)
        self.params_eq = params
        self.forcing_zero = self.params_eq[0]['f']
        self.num_traj_per_env = num_traj_per_env
        self.num_env = len(params)
        self.len = num_traj_per_env * self.num_env
        self.time_horizon = float(time_horizon)
        self.n = int(time_horizon / dt_eval)

        self.sampler = GaussianRF(2, self.size, alpha=2.5, tau=7)
        self.dt_eval = dt_eval
        self.dt = 1e-3
        self.buffer = shelve.open(buffer_filepath)
        self.test = group == 'test'
        self.max = np.iinfo(np.int32).max
        self.indices = [list(range(env * num_traj_per_env, (env + 1) * num_traj_per_env)) for env in range(self.num_env)]

    def navier_stokes_2d(self, w0, f, visc, T, delta_t, record_steps):

        #Grid size - must be power of 2
        N = w0.size()[-1]

        #Maximum frequency
        k_max = math.floor(N/2.0)

        #Number of steps to final time
        steps = math.ceil(T/delta_t)

        #Initial vorticity to Fourier space
        w_h = torch.fft.fftn(w0, (N,N))

        #Forcing to Fourier space
        f_h = torch.fft.fftn(f, (N,N))

        #If same forcing for the whole batch
        if len(f_h.size()) < len(w_h.size()):
            f_h = torch.unsqueeze(f_h, 0)

        #Record solution every this number of steps
        record_time = math.floor(steps/record_steps)

        #Wavenumbers in y-direction
        k_y = torch.cat((torch.arange(start=0, end=k_max, step=1, device=w0.device), torch.arange(start=-k_max, end=0, step=1, device=w0.device)), 0).repeat(N,1)
        #Wavenumbers in x-direction
        k_x = k_y.transpose(0,1)
        #Negative Laplacian in Fourier space
        lap = 4*(math.pi**2)*(k_x**2 + k_y**2)
        lap[0,0] = 1.0
        #Dealiasing mask
        dealias = torch.unsqueeze(torch.logical_and(torch.abs(k_y) <= (2.0/3.0)*k_max, torch.abs(k_x) <= (2.0/3.0)*k_max).float(), 0)

        #Saving solution and time
        sol = torch.zeros(*w0.size(), record_steps, 1, device=w0.device, dtype=torch.float)
        sol_t = torch.zeros(record_steps, device=w0.device)

        #Record counter
        c = 0
        #Physical time
        t = 0.0
        for j in range(steps):
            if j % record_time == 0:
                #Solution in physical space
                w = torch.fft.ifftn(w_h, (N,N))
                #Record solution and time
                sol[...,c,0] = w.real
                # sol[...,c,1] = w.imag
                sol_t[c] = t

                c += 1

            #Stream function in Fourier space: solve Poisson equation
            psi_h = w_h.clone()
            psi_h = psi_h / lap

            #Velocity field in x-direction = psi_y
            q = psi_h.clone()
            temp = q.real.clone()
            q.real = -2*math.pi * k_y * q.imag
            q.imag =  2*math.pi * k_y * temp
            q = torch.fft.ifftn(q, (N,N))

            #Velocity field in y-direction = -psi_x
            v = psi_h.clone()
            temp = v.real.clone()
            v.real = 2*math.pi*k_x*v.imag
            v.imag = -2*math.pi*k_x*temp
            v = torch.fft.ifftn(v, (N,N))

            #Partial x of vorticity
            w_x = w_h.clone()
            temp = w_x.real.clone()
            w_x.real = -2*math.pi*k_x*w_x.imag
            w_x.imag = 2*math.pi*k_x*temp
            w_x = torch.fft.ifftn(w_x, (N,N))

            #Partial y of vorticity
            w_y = w_h.clone()
            temp = w_y.real.clone()
            w_y.real = -2*math.pi*k_y*w_y.imag
            w_y.imag = 2*math.pi*k_y*temp
            w_y = torch.fft.ifftn(w_y, (N,N))

            #Non-linear term (u.grad(w)): compute in physical space then back to Fourier space
            F_h = torch.fft.fftn(q*w_x + v*w_y, (N,N))

            #Dealias
            F_h = dealias* F_h

            #Cranck-Nicholson update
            w_h = (-delta_t*F_h + delta_t*f_h + (1.0 - 0.5*delta_t*visc*lap)*w_h)/(1.0 + 0.5*delta_t*visc*lap)

            #Update real time (used only for recording)
            t += delta_t

        return sol, sol_t

    def _get_init_cond(self, index):
        torch.manual_seed(index if not self.test else self.max-index)
        if self.buffer.get(f'init_cond_{index}') is None:
            w0 = self.sampler.sample()
            state, _ = self.navier_stokes_2d(w0, f=self.forcing_zero, visc=self.params_eq[0]['visc'], T=30.0, delta_t=self.dt, record_steps=20)
            init_cond = state[:, :, -1, 0]
            self.buffer[f'init_cond_{index}'] = init_cond.numpy()
        else:
            init_cond = torch.from_numpy(self.buffer[f'init_cond_{index}'])

        return init_cond

    def __getitem__(self, index):
        env = index // self.num_traj_per_env
        env_index = index % self.num_traj_per_env
        t = torch.arange(0, self.time_horizon, self.dt_eval).float()
        if self.buffer.get(f'{env},{env_index}') is None:
            print(f'calculating index {env_index} of env {env}')
            w0 = self._get_init_cond(env_index)
            state, _ = self.navier_stokes_2d(w0, f=self.params_eq[env]['f'], visc=self.params_eq[env]['visc'], T=self.time_horizon, delta_t=self.dt, record_steps=self.n)
            state = state.permute(3, 2, 0, 1)[:, :self.n] # nc, t, h, w

            self.buffer[f'{env},{env_index}'] = {
                'state': state.numpy(),
            }
            return {
                'state'  : state,
                't'      : t,
                'env'    : env,
            }
        else:
            buf = self.buffer[f'{env},{env_index}']
            return {
                'state'  : torch.from_numpy(buf['state'][:, :self.n]),
                't'      : t,
                'env'    : env,
            }

    def __len__(self):
        return self.len

max_stars_repo_name: MLIA/LEADS

max_issues_repo_name: MLIA/LEADS

max_forks_repo_name: MLIA/LEADS

reason: import numpy,import scipy,from scipy


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

"""Description

"""

import copy
import numpy as np
from itertools import product

import torch
import torch.nn.functional as F

from ptranking.eval.parameter import ModelParameter
from ptranking.ltr_adversarial.base.ad_player import AdversarialPlayer
from ptranking.ltr_adversarial.base.ad_machine import AdversarialMachine

from ptranking.ltr_global import global_gpu as gpu, global_device as device, torch_zero, torch_one, cpu_torch_one, cpu_torch_zero

class IRGAN_Pair_Generator(AdversarialPlayer):
    def __init__(self, sf_para_dict=None, temperature=None):
        super(IRGAN_Pair_Generator, self).__init__(sf_para_dict=sf_para_dict)
        self.temperature = temperature

    def predict(self, batch_ranking, train=False):
        if train:
            self.train_mode()  # training mode for training
        else:
            self.eval_mode()  # evaluation mode for testing

        batch_pred = self.forward(batch_ranking)

        if self.temperature is not None and 1.0 != self.temperature:
                batch_pred = batch_pred / self.temperature

        return batch_pred

class IRGAN_Pair_Discriminator(AdversarialPlayer):
    def __init__(self, sf_para_dict=None):
        super(IRGAN_Pair_Discriminator, self).__init__(sf_para_dict=sf_para_dict)

    def get_reward(self, pos_docs=None, neg_docs=None, loss_type='svm'):
        ''' used by irgan '''
        batch_pos_preds = self.predict(pos_docs)  # [batch, size_ranking]
        batch_neg_preds = self.predict(neg_docs)  # [batch, size_ranking]

        if 'svm' == loss_type:
            reward = torch.sigmoid(torch.max(torch_zero, torch_one - (batch_pos_preds - batch_neg_preds)))
        elif 'log' == loss_type:
            reward = torch.log(torch.sigmoid(batch_neg_preds - batch_pos_preds))
        else:
            raise NotImplementedError

        return reward

class IRGAN_Pair(AdversarialMachine):
    ''''''
    def __init__(self, eval_dict, data_dict, sf_para_dict=None, temperature=0.2, loss_type='svm', d_epoches=None, g_epoches=None, ad_training_order=None):
        '''

        :param eval_dict:
        :param data_dict:
        :param sf_para_dict:
        :param temperature:   according to the description around Eq-10, temperature is deployed, while it is not used within the provided code
        :param loss_type:
        :param d_epoches:
        :param g_epoches:
        :param ad_training_order:
        '''
        super(IRGAN_Pair, self).__init__(eval_dict=eval_dict, data_dict=data_dict)

        sf_para_dict['ffnns']['apply_tl_af'] = True

        g_sf_para_dict = sf_para_dict

        d_sf_para_dict = copy.deepcopy(g_sf_para_dict)
        d_sf_para_dict['ffnns']['apply_tl_af'] = False
        #d_sf_para_dict['ffnns']['TL_AF'] = 'S'  # as required by the IRGAN model

        self.generator     = IRGAN_Pair_Generator(sf_para_dict=g_sf_para_dict, temperature=temperature)
        self.discriminator = IRGAN_Pair_Discriminator(sf_para_dict=d_sf_para_dict)

        self.loss_type = loss_type
        self.d_epoches = d_epoches
        self.g_epoches = g_epoches
        self.temperature = temperature
        self.ad_training_order = ad_training_order
        #self.samples_per_query = samples_per_query

    def mini_max_train(self, train_data=None, generator=None, discriminator=None, dict_buffer=None):
        if self.ad_training_order == 'DG':
            for d_epoch in range(self.d_epoches):
                if d_epoch % 10 == 0:
                    generated_data = self.generate_data(train_data=train_data, generator=generator)

                self.train_discriminator(train_data=train_data, generated_data=generated_data, discriminator=discriminator)  # train discriminator

            for g_epoch in range(self.g_epoches):
                self.train_generator(train_data=train_data, generator=generator, discriminator=discriminator)  # train generator

        else:
            for g_epoch in range(self.g_epoches):
                self.train_generator(train_data=train_data, generator=generator, discriminator=discriminator)  # train generator

            for d_epoch in range(self.d_epoches):
                if d_epoch % 10 == 0:
                    generated_data = self.generate_data(train_data=train_data, generator=generator)

                self.train_discriminator(train_data=train_data, generated_data=generated_data, discriminator=discriminator)  # train discriminator

        stop_training = False
        return stop_training

    def generate_data(self, train_data=None, generator=None, **kwargs):
        '''
        negative sampling based on current generator for training discriminator
        todo this is a re-implementation as the released irgan-tensorflow, but it seems that this part of irgan-tensorflow is not consistent with the discription of the paper (i.e., the description below Eq. 7)
        '''
        generated_data = dict()
        for entry in train_data:
            qid, batch_ranking, batch_label = entry[0], entry[1], entry[2]

            samples = self.per_query_generation(qid, batch_ranking, batch_label, generator)

            if samples is not None: generated_data[qid] = samples

        return generated_data

    def per_query_generation(self, qid, batch_ranking, batch_label, generator):

        used_batch_label = batch_label

        if gpu: batch_ranking = batch_ranking.to(device)

        # [1, ranking_size] -> [ranking_size]
        # [z, n] If input has n dimensions, then the resulting indices tensor out is of size (z×n), where z is the total number of non-zero elements in the input tensor.
        pos_inds = torch.gt(torch.squeeze(used_batch_label), 0).nonzero()
        num_pos = pos_inds.size()[0]

        if num_pos < 1:
            return None

        if num_pos < used_batch_label.size(1):
            '''
            sampling condition: pos_inds<|tor_batch_stds|
            sample numbers: min(|pos_inds|, |total_neg_inds|)
            '''
            batch_pred = generator.predict(batch_ranking)  # [batch, size_ranking]
            # batch_probs = torch.sigmoid(torch.squeeze(batch_preds))
            batch_prob = F.softmax(torch.squeeze(batch_pred), dim=0)

            # intersection implementation (keeping consistent with the released irgan-tensorflow): remove the positive part, then sample
            reversed_batch_label = torch.where(used_batch_label > 0, cpu_torch_zero, cpu_torch_one)  # negative values
            total_neg_inds = torch.nonzero(torch.squeeze(reversed_batch_label))

            num_samples = min(num_pos, total_neg_inds.size()[0])

            batch_neg_probs = batch_prob[total_neg_inds[:, 0]]
            tmp_neg_inds = torch.multinomial(batch_neg_probs, num_samples, replacement=False)
            sample_neg_inds = total_neg_inds[tmp_neg_inds]  # using the original indices within total_neg_inds rather than batch_neg_probs

            if num_samples < num_pos:
                tmp_pos_inds = torch.multinomial(torch.ones(num_pos), num_samples=num_samples, replacement=False)
                sample_pos_inds = pos_inds[tmp_pos_inds]
            else:
                sample_pos_inds = pos_inds

            # generated_data[qid] = (sample_pos_inds, sample_neg_inds)
            # convert 2-d index matrix into 1-d index matrix
            return (np.squeeze(sample_pos_inds, axis=1), np.squeeze(sample_neg_inds, axis=1))
        else:
            return None

    def train_discriminator(self, train_data=None, generated_data=None, discriminator=None, **kwargs):
        for entry in train_data:
            qid, batch_ranking = entry[0], entry[1]

            if qid in generated_data:
                if gpu: batch_ranking = batch_ranking.to(device)

                pos_inds, neg_inds = generated_data[qid]
                pos_docs = batch_ranking[:, pos_inds, :]
                neg_docs = batch_ranking[:, neg_inds, :]

                batch_pos_preds = discriminator.predict(pos_docs, train=True) # [batch, size_ranking]
                batch_neg_preds = discriminator.predict(neg_docs, train=True) # [batch, size_ranking]

                if 'svm' == self.loss_type:
                    #dis_loss = torch.mean(torch.max(torch.zeros(1), 1.0-(batch_pos_preds-batch_neg_preds)))
                    dis_loss = torch.mean(torch.max(torch_zero, torch_one - (batch_pos_preds - batch_neg_preds)))
                elif 'log' == self.loss_type:
                    dis_loss = -torch.mean(torch.log(torch.sigmoid(batch_pos_preds-batch_neg_preds)))
                else:
                    raise NotImplementedError

                discriminator.optimizer.zero_grad()
                dis_loss.backward()
                discriminator.optimizer.step()

    def train_generator(self, train_data=None, generated_data=None, generator=None, discriminator=None, **kwargs):
        for entry in train_data:
            qid, batch_ranking, batch_label = entry[0], entry[1], entry[2]
            if gpu: batch_ranking = batch_ranking.to(device)

            pos_inds = torch.gt(torch.squeeze(batch_label), 0).nonzero()

            g_preds = generator.predict(batch_ranking, train=True)
            g_probs = torch.sigmoid(torch.squeeze(g_preds))

            neg_inds = torch.multinomial(g_probs, pos_inds.size(0), replacement=True)

            pos_docs = batch_ranking[:, pos_inds[:, 0], :]
            neg_docs = batch_ranking[:, neg_inds, :]

            reward = discriminator.get_reward(pos_docs=pos_docs, neg_docs=neg_docs, loss_type=self.loss_type)

            g_loss = -torch.mean((torch.log(g_probs[neg_inds]) * reward))

            generator.optimizer.zero_grad()
            g_loss.backward()
            generator.optimizer.step()

    def reset_generator(self):
        self.generator.reset_parameters()

    def reset_discriminator(self):
        self.discriminator.reset_parameters()

    def get_generator(self):
        return self.generator

    def get_discriminator(self):
        return self.discriminator

###### Parameter of IRGAN_Pair ######

class IRGAN_PairParameter(ModelParameter):
    ''' Parameter class for Pair_IR_GAN '''
    def __init__(self, debug=False):
        super(IRGAN_PairParameter, self).__init__(model_id='IRGAN_Pair')
        self.debug = debug

    def default_para_dict(self):
        """
        Default parameter setting for IRGAN_Pair
        :return:
        """
        temperature = 0.2
        d_epoches, g_epoches = 1, 1
        ad_training_order = 'DG'
        # ad_training_order = 'GD'

        self.ad_para_dict = dict(model_id=self.model_id, d_epoches=d_epoches, g_epoches=g_epoches,
                                 temperature=temperature, ad_training_order=ad_training_order, loss_type='svm')
        return self.ad_para_dict

    def to_para_string(self, log=False, given_para_dict=None):
        """
        String identifier of parameters
        :param log:
        :param given_para_dict: a given dict, which is used for maximum setting w.r.t. grid-search
        :return:
        """
        # using specified para-dict or inner para-dict
        ad_para_dict = given_para_dict if given_para_dict is not None else self.ad_para_dict

        s1 = ':' if log else '_'
        d_epoches, g_epoches, temperature, ad_training_order, loss_type = \
            ad_para_dict['d_epoches'], ad_para_dict['g_epoches'], ad_para_dict['temperature'],\
            ad_para_dict['ad_training_order'], ad_para_dict['loss_type']

        pair_irgan_paras_str = s1.join([str(d_epoches), str(g_epoches), '{:,g}'.format(temperature),
                                        ad_training_order, loss_type])

        return pair_irgan_paras_str

    def grid_search(self):
        """
        Iterator of parameter settings for IRGAN_Pair
        :param debug:
        :return:
        """
        choice_samples_per_query = [5]
        choice_ad_training_order = ['DG']  # GD for irganlist DG for point/pair
        choice_temperatures = [0.5] if self.debug else [0.5]  # 0.5, 1.0
        choice_d_g_epoches = [(1, 1)] if self.debug else [(1, 1)]  # discriminator-epoches vs. generator-epoches

        choice_losstype_d = ['svm']

        for d_g_epoches, samples_per_query, ad_training_order, temperature, loss_type_d in product(choice_d_g_epoches,
                       choice_samples_per_query, choice_ad_training_order, choice_temperatures, choice_losstype_d):
            d_epoches, g_epoches = d_g_epoches

            self.ad_para_dict = dict(model_id=self.model_id, d_epoches=d_epoches, g_epoches=g_epoches,
                                     samples_per_query=samples_per_query, temperature=temperature,
                                     ad_training_order=ad_training_order, loss_type=loss_type_d)

            yield self.ad_para_dict


text:
import numpy as np
import mnist
from sklearn.model_selection import train_test_split

from RegLib.HelperFunctions import create_frankie_data, create_X, plot_values_with_info,plot_values_with_two_y_axis

from yacs.config import CfgNode as CN

class DataLoader():
    """
    Contains X_train, X_test, y_train, y_test for the dataset given in the config.
    Can also have X_val and y_val for other sets than franke, size depends on the the config
    """

    def __init__(self, cfg: CN, perm_index = [-1], one_hot_encode = True):
        self.data_name = cfg.DATA.NAME

        if self.data_name == "franke":
            self.load_franke_data(cfg, perm_index)
        elif self.data_name == "mnist":
            self.load_mnist_data(cfg, one_hot_encode)
        else:
            raise ValueError(self.data_name, " is not found in DataLoader init")

    def load_franke_data(self, cfg: CN, perm_index):
        x, y, z = create_frankie_data(cfg.SEED, cfg.DATA.FRANKIE.N, cfg.DATA.FRANKIE.NOISE)
        X = create_X(x, y, cfg.DATA.FRANKIE.P)

        self.split_and_scale_train_test(X, z, perm_index, test_size = cfg.TEST_SIZE)
        return self

    def load_mnist_data(self, cfg: CN, one_hot_encode = True):
        val_percent = cfg.DATA.MNIST.VAL_PERCENT
        binary_classes = cfg.DATA.MNIST.BINARY
        num_of_classes = len(binary_classes)
        if(num_of_classes != 0):
            assert num_of_classes == 2, "Cannot have " + str(num_of_classes) + " classes"
            X_train, self.y_train, X_val, self.y_val, X_test, self.y_test = load_binary_dataset(binary_classes[0], binary_classes[1], val_percent)
        else:
            X_train, self.y_train, X_val, self.y_val, X_test, self.y_test = load_full_mnist(val_percent)
            # One hot encode the results
            if one_hot_encode:
                self.y_train = self.one_hot_encode(self.y_train, 10)
                self.y_val = self.one_hot_encode(self.y_val, 10)
                self.y_test = self.one_hot_encode(self.y_test, 10)

        # Pre-process the batch
        X_mean, X_std = (np.mean(X_train), np.std(X_train))
        self.X_train = self.pre_process_images(X_train, X_mean, X_std)
        self.X_val = self.pre_process_images(X_val, X_mean, X_std)
        self.X_test = self.pre_process_images(X_test, X_mean, X_std)

    def one_hot_encode(self, Y: np.ndarray, num_classes: int):
        new_Y = np.zeros((Y.shape[0], num_classes))
        for i in range(len(Y)):
            new_Y[i][Y[i]] = 1
        return new_Y

    def pre_process_images(self, X: np.ndarray, X_mean: float, X_std: float):
        assert X.shape[1] == 784,\
            f"X.shape[1]: {X.shape[1]}, should be 784"
        X = (X - X_mean) / X_std
        X = np.c_[X,  np.ones(X.shape[0])] # Apply bias trick
        return X

    def split_and_scale_train_test(self, X, y, perm_index = [-1], test_size  = 0.2):
        assert X.shape[0] == y.shape[0], ("X.shape[0] and y.shape[0] needs to be the same length, but: " + str(X.shape[0]) + " != " + str(y.shape[0]))
        if(len(perm_index) > 1):
            X = X[perm_index]
            y = y[perm_index]
        self.X_train, self.X_test, self.y_train, self.y_test = train_test_split(X, y, test_size=test_size, shuffle=False)
        self.X_train, self.X_val, self.y_train, self.y_val = train_test_split(self.X_train, self.y_train, test_size=0.25, shuffle = False)
        self.X_train, self.X_test, self.X_val = self.scale_standard(self.X_train, self.X_test, self.X_val)
        # Force the correct shape:
        self.y_test.shape = (self.y_test.shape[0], 1)
        self.y_train.shape = (self.y_train.shape[0], 1)
        self.y_val.shape = (self.y_val.shape[0], 1)
        return self

    def scale_standard(self, train_data, test_data, val_data):
        data_mean = np.mean(train_data[:,1:], axis = 0)
        data_std = np.std(train_data[:,1:], axis = 0)
        train_data_scaled = train_data
        test_data_scaled = test_data
        val_data_scaled = val_data
        train_data_scaled[:,1:] = np.divide((train_data[:,1:] - data_mean), data_std)
        test_data_scaled[:,1:] = np.divide((test_data[:,1:] - data_mean), data_std)
        val_data_scaled[:,1:] = np.divide((val_data[:,1:] - data_mean), data_std)

        return train_data_scaled, test_data_scaled, val_data_scaled

# https://github.com/hukkelas/TDT4265-StarterCode/tree/master/assignment1
def binary_prune_dataset(class1: int, class2: int,
                        X: np.ndarray, Y: np.ndarray):
    """
    Splits the dataset into the class 1 and class2. All other classes are removed.
    Args:
        X: images of shape [batch size, 784] in the range (0, 255)
        Y: labels of shape [batch size]
    """

    mask1 = (Y == class1)
    mask2 = (Y == class2)
    mask_total = np.bitwise_or(mask1, mask2)
    Y_binary = Y.copy()
    Y_binary[mask1] = 1
    Y_binary[mask2] = 0
    return X[mask_total], Y_binary[mask_total]

def train_val_split(X: np.ndarray, Y: np.ndarray, val_percentage: float):
    """
    Randomly splits the training dataset into a training and validation set.
    """
    idx = np.arange(0, X.shape[0])
    np.random.shuffle(idx)
    train_size = int(X.shape[0] * (1 - val_percentage))
    idx_train = idx[:train_size]
    idx_val = idx[train_size:]

    X_train, Y_train = X[idx_train], Y[idx_train]
    X_val, Y_val = X[idx_val], Y[idx_val]

    return X_train, Y_train, X_val, Y_val

def load_binary_dataset(class1: int, class2: int, val_percentage: float):
    """
    Loads, prunes and splits the dataset into train, validation and test.
    """
    train_size = 20000
    test_size = 2000
    X_train, Y_train, X_test, Y_test = mnist.load()

    # First 20000 images from train set
    X_train, Y_train = X_train[:train_size], Y_train[:train_size]
    # Last 2000 images from test set
    X_test, Y_test = X_test[-test_size:], Y_test[-test_size:]
    X_train, Y_train = binary_prune_dataset(
        class1, class2, X_train, Y_train
    )
    X_test, Y_test = binary_prune_dataset(
        class1, class2, X_test, Y_test
    )
    # Reshape to (N, 1)
    Y_train = Y_train.reshape(-1, 1)
    Y_test = Y_test.reshape(-1, 1)
    X_train, Y_train, X_val, Y_val = train_val_split(
        X_train, Y_train, val_percentage
    )
    # print(f"Train shape: X: {X_train.shape}, Y: {Y_train.shape}")
    # print(f"Validation shape: X: {X_val.shape}, Y: {Y_val.shape}")
    # print(f"Test shape: X: {X_test.shape}, Y: {Y_test.shape}")

    return X_train, Y_train, X_val, Y_val, X_test, Y_test

def load_full_mnist(val_percentage: float):
    """
    Loads and splits the dataset into train, validation and test.
    """
    train_size = 20000
    test_size = 2000
    X_train, Y_train, X_test, Y_test = mnist.load()

    # First 20000 images from train set
    X_train, Y_train = X_train[:train_size], Y_train[:train_size]
    # Last 2000 images from test set
    X_test, Y_test = X_test[-test_size:], Y_test[-test_size:]
    # Reshape to (N, 1)
    Y_train = Y_train.reshape(-1, 1)
    Y_test = Y_test.reshape(-1, 1)
    X_train, Y_train, X_val, Y_val = train_val_split(
        X_train, Y_train, val_percentage
    )
    # print(f"Train shape: X: {X_train.shape}, Y: {Y_train.shape}")
    # print(f"Validation shape: X: {X_val.shape}, Y: {Y_val.shape}")
    # print(f"Test shape: X: {X_test.shape}, Y: {Y_test.shape}")

    return X_train, Y_train, X_val, Y_val, X_test, Y_test

max_stars_repo_name: marianylund/fysstkprojects

max_issues_repo_name: marianylund/fysstkprojects

max_forks_repo_name: marianylund/fysstkprojects


text:
"""变化检测数据裁剪"""
import argparse
import glob
import os
import time
from osgeo import gdal
import numpy as np

#  创建文件夹
def mkdir(path):
    if not os.path.exists(path):
        os.mkdir(path)

#  读取栅格数据
def read(fileName, xoff=0, yoff=0, data_width=0, data_height=0):
    dataset = gdal.Open(fileName)
    if dataset == None:
        print(fileName + "文件无法打开")
    #  栅格矩阵的列数
    width = dataset.RasterXSize
    #  栅格矩阵的行数
    height = dataset.RasterYSize
    #  获取数据
    if (data_width == 0 and data_height == 0):
        data_width = width
        data_height = height
    data = dataset.ReadAsArray(xoff, yoff, data_width, data_height)
    #  获取仿射矩阵信息
    geotrans = dataset.GetGeoTransform()
    #  获取投影信息
    proj = dataset.GetProjection()
    return width, height, data, geotrans, proj

#  保存样本函数
def writeTiff(im_data, im_geotrans, im_proj, path):
    if 'int8' in im_data.dtype.name:
        datatype = gdal.GDT_Byte
    elif 'int16' in im_data.dtype.name:
        datatype = gdal.GDT_UInt16
    else:
        datatype = gdal.GDT_Float32
    if len(im_data.shape) == 3:
        im_bands, im_height, im_width = im_data.shape
    elif len(im_data.shape) == 2:
        im_data = np.array([im_data])
        im_bands, im_height, im_width = im_data.shape
    # 创建文件
    driver = gdal.GetDriverByName("GTiff")
    dataset = driver.Create(path, int(im_width), int(im_height), int(im_bands), datatype, options=["TILED=YES", "COMPRESS=LZW"])
    if (dataset != None):
        dataset.SetGeoTransform(im_geotrans)  # 写入仿射变换参数
        dataset.SetProjection(im_proj)  # 写入投影
    for i in range(im_bands):
        dataset.GetRasterBand(i + 1).WriteArray(im_data[i])
    del dataset

#  选取正样本并保存
def select2write(dataA, dataB, label, geotrans, proj, savePath, name, False_sample, mask=255):
    if False_sample:
        if dataA.shape == dataB.shape and label.shape == dataA.shape[1:]:
            if not mask in label:
                p = np.array([0.1, 0.9])  # 10%概率保存负样本
                save = np.random.choice(["yes", "no"], p=p.ravel())
                if save == "yes":
                    writeTiff(dataA, geotrans, proj, os.path.join(savePath, "A/%d.tif" % name))
                    writeTiff(dataB, geotrans, proj, os.path.join(savePath, "B/%d.tif" % name))
                    writeTiff(label, geotrans, proj, os.path.join(savePath, "OUT/%d.tif" % name))
                    name += 1
            else:
                writeTiff(dataA, geotrans, proj, os.path.join(savePath, "A/%d.tif" % name))
                writeTiff(dataB, geotrans, proj, os.path.join(savePath, "B/%d.tif" % name))
                writeTiff(label, geotrans, proj, os.path.join(savePath, "OUT/%d.tif" % name))
                name += 1
    else:
        if mask in label:
            if dataA.shape == dataB.shape and label.shape == dataA.shape[1:]:
                # 判断目标占比
                target = len(label[label == 255])
                num = label.shape[0] * label.shape[1]
                if not target / num <= 0.02:
                    writeTiff(dataA, geotrans, proj, os.path.join(savePath, "A/%d.tif" % name))
                    writeTiff(dataB, geotrans, proj, os.path.join(savePath, "B/%d.tif" % name))
                    writeTiff(label, geotrans, proj, os.path.join(savePath, "OUT/%d.tif" % name))
                    name += 1
    return name

#  裁剪函数
def Crop(width, height, imgA, imgB, label, geotrans, proj, SavePath, CropSize, RepetitionRate, False_sample):
    #  获取当前文件夹的文件个数len,并以len+1命名即将裁剪得到的图像
    new_name = len(os.listdir(os.path.join(SavePath, "A"))) + 1
    #  裁剪图片,重复率为RepetitionRate
    for i in range(int((height - CropSize * RepetitionRate) / (CropSize * (1 - RepetitionRate)))):
        for j in range(int((width - CropSize * RepetitionRate) / (CropSize * (1 - RepetitionRate)))):
            cropped_label = label[
                            int(i * CropSize * (1 - RepetitionRate)): int(
                                i * CropSize * (1 - RepetitionRate)) + CropSize,
                            int(j * CropSize * (1 - RepetitionRate)): int(
                                j * CropSize * (1 - RepetitionRate)) + CropSize]
            croppedA = imgA[:,
                       int(i * CropSize * (1 - RepetitionRate)): int(i * CropSize * (1 - RepetitionRate)) + CropSize,
                       int(j * CropSize * (1 - RepetitionRate)): int(j * CropSize * (1 - RepetitionRate)) + CropSize]
            croppedB = imgB[:,
                       int(i * CropSize * (1 - RepetitionRate)): int(i * CropSize * (1 - RepetitionRate)) + CropSize,
                       int(j * CropSize * (1 - RepetitionRate)): int(j * CropSize * (1 - RepetitionRate)) + CropSize]
            new_name = select2write(croppedA, croppedB, cropped_label, geotrans, proj, SavePath, new_name,
                                    False_sample)
    #  向前裁剪最后一列
    for i in range(int((height - CropSize * RepetitionRate) / (CropSize * (1 - RepetitionRate)))):
        cropped_label = label[
                        int(i * CropSize * (1 - RepetitionRate)): int(i * CropSize * (1 - RepetitionRate)) + CropSize,
                        (width - CropSize): width]
        croppedA = imgA[:,
                   int(i * CropSize * (1 - RepetitionRate)): int(i * CropSize * (1 - RepetitionRate)) + CropSize,
                   (width - CropSize): width]
        croppedB = imgB[:,
                   int(i * CropSize * (1 - RepetitionRate)): int(i * CropSize * (1 - RepetitionRate)) + CropSize,
                   (width - CropSize): width]
        new_name = select2write(croppedA, croppedB, cropped_label, geotrans, proj, SavePath, new_name, False_sample)
    #  向前裁剪最后一行
    for j in range(int((width - CropSize * RepetitionRate) / (CropSize * (1 - RepetitionRate)))):
        cropped_label = label[(height - CropSize): height,
                        int(j * CropSize * (1 - RepetitionRate)): int(j * CropSize * (1 - RepetitionRate)) + CropSize]
        croppedA = imgA[:,
                   (height - CropSize): height,
                   int(j * CropSize * (1 - RepetitionRate)): int(j * CropSize * (1 - RepetitionRate)) + CropSize]
        croppedB = imgB[:,
                   (height - CropSize): height,
                   int(j * CropSize * (1 - RepetitionRate)): int(j * CropSize * (1 - RepetitionRate)) + CropSize]
        new_name = select2write(croppedA, croppedB, cropped_label, geotrans, proj, SavePath, new_name, False_sample)
    #  裁剪右下角
    cropped_label = label[(height - CropSize): height,
                    (width - CropSize): width]
    croppedA = imgA[:,
               (height - CropSize): height,
               (width - CropSize): width]
    croppedB = imgB[:,
               (height - CropSize): height,
               (width - CropSize): width]
    new_name = select2write(croppedA, croppedB, cropped_label, geotrans, proj, SavePath, new_name, False_sample)

#  获取数据
def main(RasterA, RasterB, Label):
    width, height, dataA, geotrans, proj = read(RasterA)
    dataB, label = read(RasterB)[2], read(Label)[2]
    print(os.path.split(RasterA)[1], os.path.split(RasterB)[1],
          os.path.split(Label)[1], "栅格读取完毕,开始裁剪")
    return width, height, dataA, dataB, label, geotrans, proj

if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument('--imageA', type=str, default="", help='ImageA path')
    parser.add_argument('--imageB', type=str, default="", help='ImageB path')
    parser.add_argument('--label', type=str, default="", help='Label path')
    parser.add_argument('--output', type=str, default="", help='output path')
    parser.add_argument('--cropsize', type=int, default=256)
    parser.add_argument('--RepetitionRate', type=float, default=0.)
    parser.add_argument('--FalseSample', action='store_true', help='save false sample')
    opt = parser.parse_args()
    start = time.time()
    RasterA_list = glob.glob(os.path.join(opt.imageA, "*.tif"))
    RasterB_list = glob.glob(os.path.join(opt.imageB, "*.tif"))
    Label_list = glob.glob(os.path.join(opt.label, "*.tif"))
    #  创建文件夹用于保存裁剪结果
    mkdir(os.path.join(opt.output, "A"))
    mkdir(os.path.join(opt.output, "B"))
    mkdir(os.path.join(opt.output, "OUT"))
    for i in range(len(RasterA_list)):
        width, height, dataA, dataB, label, geotrans, proj = main(RasterA_list[i], RasterB_list[i], Label_list[i])
        Crop(width, height, dataA, dataB, label, geotrans, proj, opt.output, opt.cropsize, opt.RepetitionRate,
             opt.FalseSample)
    end = time.time()
    print("用时{}秒".format((end - start)))


text:
"""
    test_cf_hf_gd.py
    ~~~
    This module contains testing function of SVD method to discover hidden
    features in collaborative filtering method

    In this tesing file, we generate rating matrix for 1k user playing history
    of songs in Million Song Dataset. Because there are large amount of miss
    match in two data source, we only generate rate matrix of tracks in MSD
    which are played by 1k user dataset. Then we user Gradient Descent method
    to discover the hidden features in the CF methods.

    :auther: Alexander Z Wang
"""
import numpy
import msd
import echo_nest as en
import cf_hidden_feature as ch

k = 5
lean_rate = 0.00001
lambda_rate = 0.00
max_iter = 5000
GD_method = 1

filename_tracks = "subset_unique_tracks.txt"
filename_echo_nest = "train_triplets.txt"
num = 1000

print "Reading MSD and Echo Nest Data..."
song_index, name_index = msd.get_song_ID_index(filename_tracks)
echo_nest_user_history = en.get_echo_nest_user_history(
                                filename_echo_nest, song_index)
user_rate_dict = en.get_top_user_rating_from_history_echo_nest(
                                echo_nest_user_history, num)

print "Calculating Hidden Features"

user_weight, hidden_feature = ch.get_hidden_feature_matrix_GD(user_rate_dict,
                                k, lean_rate, lambda_rate, max_iter, GD_method)

print "hidden features of 10 songs"
print hidden_feature[0:10,:]
hist, bin_edges = numpy.histogram(hidden_feature, bins=20)
print hist
print bin_edges

max_stars_repo_name: BerryAI/Acai

max_issues_repo_name: BerryAI/Acai

max_forks_repo_name: BerryAI/Acai


text:
import logging
import sys
import random
import torch
import torch.cuda
import numpy as np
import os
import torch
from models.lenet import LeNet5
from models.text_binary_classification import TextBinaryClassificationModel
from models.hate_speech_model import HateSpeechModel
import copy
import yaml
from torch.nn.utils import parameters_to_vector, vector_to_parameters

def getLogger(logFile,stdoutRedirect=False,level=logging.INFO):
    #log_file_name = os.path.basename(args.log_file).split(".")[0]+".log"
    logging.basicConfig(filename=logFile,filemode='w')
    logger = logging.getLogger()
    logger.setLevel(level)
    if(stdoutRedirect):
        handler = logging.StreamHandler(sys.stdout)
        logger.addHandler(handler)
    return logger

def loadConfig(filePath):
    obj = None
    with open(filePath,'r') as f:
        obj = yaml.load(f,Loader=yaml.FullLoader)
    return obj

def seed(seed):
    # seed = 1234
    random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    np.random.seed(seed)
    os.environ['PYTHONHASHSEED'] = str(seed)
    #TODO: Do we need deterministic in cudnn ? Double check
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False
    print ("Seeded everything")

def setParamsToZero(model):
    W = list(model.parameters())
    for i in range(len(W)):
        W[i].data.fill_(0.0)

def saveModel(model,path):
    torch.save(model.state_dict(), path)

def flatModel(mdl):
    W = list(mdl.parameters())
    l = [torch.flatten(w) for w in W]
    return torch.cat(l)

def printStats_(mdl):
    v = flatModel(mdl)
    print(torch.mean(v),torch.std(v))

# Obsolete
def createModel(config):
    model = None
    arch = config['arch']
    lossType = config['loss']
    if(arch=='lenet5'):
        model = LeNet5()

    if(lossType == 'crossEntropy'):
        criterion = torch.nn.CrossEntropyLoss().to(config['device'])

    if(arch=='rnn'):
        model = models.RNNModel(config["model"], config["ntokens"], config["emsize"],
                                       config["nhid"], config["nlayers"], config["dropout"], config["tied"])
    if(lossType == 'crossEntropy'):
        criterion = torch.nn.NLLLoss().to(config['device'])
    if(arch=='textBC'):
        model = TextBinaryClassificationModel(config["modelParams"])
        criterion = None
    if(arch=='hateSpeech'):
        model = HateSpeechModel(config["modelParams"])
        criterion = None

    model.to(config['device'])
    return model,criterion

def loadModel(path,config):
    model,crit = createModel(config)
    model.load_state_dict(torch.load(path))
    return model

#copy from mdl1 to mdl2
def copyParams(mdl1,mdl2):
    W1  = list(mdl1.parameters())
    W2  = list(mdl2.parameters())
    assert len(W1) == len(W2)
    for i in range(len(W1)):
        W2[i].data = copy.deepcopy(W1[i].data)

def addModels(mdl1, mdl2,scale1=1.0,scale2=1.0):
    cls = mdl1.__class__
    mdl3 = cls()
    W3  = list(mdl3.parameters())
    W1  = list(mdl1.parameters())
    W2  = list(mdl2.parameters())
    assert len(W1) == len(W2) == len(W3)
    for i in range(len(W1)):
        W3[i].data = scale1*W1[i].data + scale2*W2[i].data

    return mdl3

def addModelsInPlace(mdl1, mdl2,scale1=1.0,scale2=1.0):

    W1  = list(mdl1.parameters())
    W2  = list(mdl2.parameters())
    assert len(W1) == len(W2)

    for i in range(len(W1)):
        W1[i].data = scale1*W1[i].data + scale2*W2[i].data

def normDiff(mdl1,mdl2):
    W1  = list(mdl1.parameters())
    W2  = list(mdl2.parameters())
    assert len(W1) == len(W2)
    nd = 0
    for i in range(len(W1)):
        nd += torch.norm(W1[i].data - W2[i])**2
    return torch.sqrt(nd).item()

def distModels(mdl1,mdl2,p=2):
    W1 = flatModel(mdl1)
    W2 = flatModel(mdl2)
    return torch.dist(W1,W2,p=p)

def normModel(mdl,p=2):
    v = parameters_to_vector(mdl.parameters())
    return torch.norm(v,p=p)

def constructMNISTBackdoorData(trainData):
    pass


text:
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.colors as colors
from matplotlib.patches import Ellipse
from matplotlib.patches import Rectangle
import warnings
warnings.filterwarnings('ignore')
import os
from os import path
from matplotlib import colors

pitchangle=0.9
d=pitchangle

twilightcmap=plt.cm.twilight

plt.rcParams['mathtext.fontset'] = 'cm'
plt.rcParams['font.family'] = 'cmu serif'
plt.rc('font', size=12)          # controls default text sizes
fig_directory='/Users/emma/OneDrive/PhD/thesis/Figures/'

def main():
	beamnumber=0
	beamcenters=np.empty(shape=(6,6,3))

	#these are the centre coordinates of beam zero for closepack
	x0=3*d - d/(2*np.sqrt(2))
	y0=-1*d/(2*np.sqrt(2)) - np.sqrt(3)*d

	fig,axs=plt.subplots(1,2,sharey=True)
	axs[0].set_xlim(-3.5,3.5)
	axs[0].set_ylim(-3.5,3.5)
	axs[0].set_aspect("equal")
	for j in range(0,6):
		for i in range(0,6):
			y=y0 + j*d*(np.sqrt(3)/2)
			if j==0 or j==2 or j==4:
				x=x0 - i*d
			else:
				x=x0 - (i+0.5)*d

			ellipse=Ellipse(xy=(x,y),width=d,height=d,angle=0,edgecolor=twilightcmap(0.5),fc='none',lw=1)
			axs[0].add_patch(ellipse)
			axs[0].annotate(text=str(beamnumber),xy=(x,y),color=twilightcmap(0.3),ha='center',va='center')
			beamnumber+=1

	#px=x0 - (5.5)*d
	#py=y0+ 5.8*d*(np.sqrt(3)/2)
	#rectangle=Rectangle(xy=(px,py),width=d,height=0.05,color=twilightcmap(0.7))
	#axs[0].add_patch(rectangle)
	#axs[0].annotate(text='Pitch',xy=(px,py+0.12),color=twilightcmap(0.5))
	axs[0].set_xlabel('x offset (degrees)')
	axs[0].set_ylabel('y offset (degrees)')

	x0=+2.5*d
	y0=-2.5*d
	axs[1].set_xlim(-3.5,3.5)
	axs[1].set_ylim(-3.5,3.5)

	square_beams=[31,30,29,28,27,26,32,13,12,11,10,25,33,14,2,3,9,24,34,15,0,1,8,23,35,4,5,6,7,22,16,17,18,19,20,21]
	k=0
	for j in range(0,6):
		for i in range(0,6):
			x=x0-(i*d)
			y=y0+(j*d)
			ellipse=Ellipse(xy=(x,y),width=d,height=d,angle=0,edgecolor=twilightcmap(0.5),fc='none',lw=1)
			axs[1].add_patch(ellipse)
			axs[1].annotate(text=str(square_beams[k]),xy=(x,y),color=twilightcmap(0.3),ha='center',va='center')
			k+=1

	axs[1].set_aspect("equal")
	axs[1].set_xlabel('x offset (degrees)')
	axs[1].tick_params(axis='y',which='both',left=False)

	#px=-2.5*d
	#py=d*3.2
	#rectangle=Rectangle(xy=(px,py),width=d,height=0.05,color=twilightcmap(0.7))
	#axs[1].add_patch(rectangle)
	#axs[1].annotate(text='Pitch',xy=(px,py+0.12),color=twilightcmap(0.5))

	axs[0].set_title('closepack36')
	axs[1].set_title('square_6x6')

	plt.subplots_adjust(left=0.1, bottom=0.01, right=0.99, top=0.99, wspace=0, hspace=0.01)

	plt.savefig(fig_directory+'askapbeams.png',transparent='True',dpi=600)
	#plt.show()

if __name__ == "__main__":
	main()

max_stars_repo_name: EmmaAlexander/possum-tools

max_issues_repo_name: EmmaAlexander/possum-tools

max_forks_repo_name: EmmaAlexander/possum-tools