text:
# coding: utf-8
'''
Plot a implied volatility surface by estimating the implied volatiltiy using newton's method

__author__ = "Zhejian Peng"
__copyright__ = "Zhejian Peng"
__license__ = MIT LICENSE
__version__ = "1.0.1"
__maintainer__ = "Zhejian Peng"
__email__ = "zhejianpeng@yahoo.com"
__status__ = "Developing"

__update__ =
'''
import numpy as np
import pandas as pd
import scipy.stats as st
import matplotlib.pyplot as plt
from scipy.optimize import newton
import datetime

def black_scholes(call_or_put, S0, K, r, T, sigma):
    d1 = (1/(sigma*np.sqrt(T)))*(np.log(S0/K)+(r+(sigma**2)/2)*T)
    d2 = d1-sigma*np.sqrt(T)
    if call_or_put.upper() == 'C':
        return S0*st.norm.cdf(d1)-K*np.exp(-r*T)*st.norm.cdf(d2)
    elif call_or_put.upper() == 'P':
        return K*np.exp(-r*T)*st.norm.cdf(-d2)-S0*st.norm.cdf(-d1)
    else:
        print('Please enter C for call option and P for put option')

# print(black_scholes('C', 10, 10, 0.1, 1/12, 0.5))

def find_implied_volatiltiy(call_or_put, mkt_price, S0, K, r, T):
    implied_sigma = newton(lambda x: black_scholes(
        call_or_put, S0, K, r, T, x)-mkt_price, 0.5)
    return implied_sigma

options = pd.read_csv("./Implied Volatility/C_U7P_Comdty.csv")
futures = pd.read_excel("./Implied Volatility/futures_price.xlsx")
futures_price = futures["Last Price"]

date = []
for i in range(len(options["Date"])):
    date.append(datetime.datetime.strptime(
        options["Date"][i], "%m/%d/%Y").date())

expiration = datetime.date(2017, 8, 25)

strike = np.linspace(310, 430, 13)

implied_vol = np.zeros((13, len(date)))
Time_2_maturity = np.zeros(len(date))

for j in range(len(date)):
    for i in range(len(strike)):
        implied_vol[i, j] = find_implied_volatiltiy('p', options[str(
            strike[i])+"0"][j],  futures_price[j], strike[i], 0.0105, (expiration-date[j]).days/360)
    Time_2_maturity[j] = (expiration-date[j]).days/360

plt.title("Implied Volatility Smile (2017-07-13)")
plt.plot(futures_price[0]/strike, implied_vol[:, 0])
plt.scatter(futures_price[0]/strike, implied_vol[:, 0])
plt.grid()
plt.xlabel("Moneyness")
plt.ylabel("Implied Volatility")
plt.show()
# black_scholes('P', 0, 0, 0, 0, 0)

# Volatility surface:
expiration = datetime.date(2017, 8, 25)

strike = np.linspace(310, 430, 13)
surface = np.zeros([len(strike), int(len(date))])
# print(surface.shape)
Y = np.zeros(int(len(date)))

for j in range(int(len(date))):
    for i in range(len(strike)):
        surface[i, j] = implied_vol[i, j]
    Y[j] = Time_2_maturity[j]

from mpl_toolkits.mplot3d import Axes3D
from matplotlib import cm
fig = plt.figure()
ax = fig.gca(projection='3d')

strike, Y = np.meshgrid(strike, Y)
ax.plot_surface(strike, Y, np.transpose(surface), cmap=cm.coolwarm)
# surf = ax.plot_surface(strike, Y, np.transpose(surface))
#     surface), cmap=cm.coolwarm, linewidth=0, antialiased=False)
# #fig.colorbar(surf, shrink=0.5, aspect=5)

ax.set_title("Implied Volatility Surface")
ax.set_xlabel('Strike Price')
ax.set_ylabel('Time to Maturity')
ax.set_zlabel('Implied Volatility')
plt.show()

reason: import numpy,import scipy,from scipy


text:
# -*- coding: utf-8 -*-
"""
Created on Wed Apr 11 14:41:16 2018

@author: JHodges
"""

import geopandas as gpd
import matplotlib.pyplot as plt
import numpy as np
import behavePlus as bp
import glob
from shapely.geometry import Polygon, Point
import datetime
import subprocess
import queryLandFire as qlf
import multiprocessing as mp
import sys
from runFarsiteMp import runFarsite

"""
def runFarsite(commandFile):
    dockerStart = 'docker run -it -v E:\\projects\\wildfire-research\\farsite\\:/commonDir/ farsite'
    dockerCmd = './commonDir/farsite/src/TestFARSITE %s'%(commandFile)

    p = subprocess.Popen('winpty '+dockerStart+' '+dockerCmd,shell=False, creationflags=subprocess.CREATE_NEW_CONSOLE)
    p_status = p.wait()
"""

def getFuelMoistureData(string,params,fuelModels=np.linspace(0,1,2)):
    string = string+'FUEL_MOISTURES_DATA: %.0f\n'%(fuelModels.shape[0])
    m1h = params['m1h']
    m10h = params['m10h']
    m100h = params['m100h']
    lhm = params['lhm']
    lwm = params['lwm']

    for i in range(0,fuelModels.shape[0]):
        #string = string+'%.0f\t%.1f\t%.1f\t%.1f\t%.1f\t%.1f\n'%(fuelModels[i],m1h,m10h,m100h,lhm,lwm)
        string = string+'%.0f %.1f %.1f %.1f %.1f %.1f\n'%(fuelModels[i],m1h,m10h,m100h,lhm,lwm)
    return string

def getMaxDay(Mth):
    maxDay = 31
    if Mth == 4 or Mth == 6 or Mth == 9 or Mth == 11:
        maxDay = 30
    if Mth == 2:
        maxDay = 28
    return maxDay

def incrementDay(Day,Mth):
    Day = Day + 1
    maxDay = getMaxDay(Mth)
    if Day > maxDay:
        Day = 1
        Mth = Mth+1
    if Mth > 12:
        Mth = 1
    return Day, Mth

def getWeatherData(string,params,Elv,totalDays=2):
    string = string+"WEATHER_DATA: %.0f\n"%(totalDays)

    Mth = round(params['Mth'])
    Day = round(params['Day'])
    maxDay = getMaxDay(Mth)
    Day = min(Day,maxDay)

    Pcp = params['Pcp']
    mTH = round(params['mTH'],-2)
    xTH = round(params['xTH'],-2)
    mT = params['mT']
    xT = params['xT']
    xH = params['xH']
    mH = params['mH']
    PST = round(params['PST'],-2)
    PET = round(params['PET'],-2)

    for i in range(0,totalDays):
        # Mth  Day  Pcp  mTH  xTH   mT xT   xH mH   Elv   PST  PET
        Day, Mth = incrementDay(Day,Mth)
        string = string+'%.0f %.0f %.1f %.0f %.0f %.1f %.1f %.1f %.1f %.0f %.0f %.0f\n'%(
                Mth,Day,0,mTH,xTH,mT,xT,xH,xH,Elv,0,0)
    string = string+"WEATHER_DATA_UNITS: Metric\n"
    return string

def getWindData(string,params,totalDays=2):
    string = string+"WIND_DATA: %.0f\n"%(totalDays*24)
    windSpeed = params['windSpeed'] # mph
    windDir = params['windDir']

    windSpeed = windSpeed*5280*12*25.4/(1000*1000) # km/h

    Mth = round(params['Mth'])
    Day = round(params['Day'])
    maxDay = getMaxDay(Mth)
    Day = min(Day,maxDay)

    for i in range(0,totalDays):
        # Mth  Day  Hour   Speed Direction CloudCover
        Day, Mth = incrementDay(Day,Mth)
        for Hour in range(0,2400,100):
            string = string+'%.0f %.0f %.0f %.1f %.1f %.1f\n'%(
                    Mth,Day,Hour,windSpeed,windDir,0)
    string = string+"WIND_DATA_UNITS: Metric\n"
    return string

def getMiscData(string):
    string = string + "FOLIAR_MOISTURE_CONTENT: 100\n"
    string = string + "CROWN_FIRE_METHOD: Finney\n" # Either Finney or ScottRhienhardt
    string = string + "FARSITE_SPOT_PROBABILITY: 0.01\n"
    string = string + "FARSITE_SPOT_IGNITION_DELAY: 15\n"
    string = string + "FARSITE_MINIMUM_SPOT_DISTANCE: 30\n"
    string = string + "FARSITE_ACCELERATION_ON: 1\n"
    return string

def getSimulationData(string,params,totalDays=2):
    Mth = round(params['Mth'])
    Day = round(params['Day'])
    maxDay = getMaxDay(Mth)
    Day = min(Day,maxDay)
    startTime = params['startTime']
    startHour = np.floor(startTime)
    startMin = min(int(round((startTime-startHour)*60,0)),59)
    startHour = int(startHour)

    Day, Mth = incrementDay(Day,Mth)
    Day, Mth = incrementDay(Day,Mth)

    sTime = datetime.datetime(year=2016,month=Mth,day=Day,hour=startHour,minute=startMin)
    eTime = datetime.datetime(year=2016,month=Mth,day=Day)+datetime.timedelta(days=totalDays-2)
    #eTime = sTime+datetime.timedelta(days=totalDays-3)
    sTimeString = sTime.strftime('%m %d %H%M')
    eTimeString = eTime.strftime('%m %d %H%M')

    string = string + "FARSITE_START_TIME: %s\n"%(sTimeString)
    string = string + "FARSITE_END_TIME: %s\n"%(eTimeString)
    string = string + "FARSITE_TIME_STEP: 60\n"
    string = string + "FARSITE_DISTANCE_RES: 30.0\n"
    string = string + "FARSITE_PERIMETER_RES: 60.0\n"
    return string

def getIgnitionData(string,ignitionFile):
    string = string + "FARSITE_IGNITION_FILE: %s\n"%(ignitionFile)
    return string

def saveFarsiteInput(string,file):
    with open(file,'w') as f:
        f.write(string)

def generateFarsiteInput(file,elevation,ignitionFile,totalDays=5):

    paramsInput = bp.getStandardParamsInput()
    params = bp.getRandomConditions(paramsInput,allowDynamicModels=True)
    string = ''
    string = getFuelMoistureData(string,params)
    string = getWeatherData(string,params,elevation,totalDays=totalDays)
    string = getWindData(string,params,totalDays=totalDays)
    string = getSimulationData(string,params,totalDays=totalDays)
    string = getMiscData(string)
    string = getIgnitionData(string,ignitionFile)
    saveFarsiteInput(string,file+'.input')
    #print(string)

    return params

def getLcpElevation(file):
    imgs, names, header = qlf.readLcpFile(file)
    elevation = np.median(imgs[0])
    return elevation

def makeCenterIgnition(file,N=5):
    imgs, names, header = qlf.readLcpFile(file)
    eastUtm = header[1039]; westUtm = header[1040]
    northUtm = header[1041]; southUtm = header[1042]
    xResol = header[1044]; yResol = header[1045];

    centerX = (eastUtm+westUtm)/2
    centerY = (northUtm+southUtm)/2

    corners = [[centerX-N*xResol,centerY-N*yResol],
               [centerX-N*xResol,centerY+N*yResol],
               [centerX+N*xResol,centerY+N*yResol],
               [centerX+N*xResol,centerY-N*yResol]]
    points = [Point(xyz) for xyz in corners]
    geometry = Polygon([[p.x,p.y] for p in points])
    data = gpd.GeoDataFrame([[0,0]],columns=['ENTITY','VALUE'],geometry=[geometry]) # ENTITY=0, VALUE = 0
    file = file.replace('.LCP','_ignite.SHP')
    data.to_file(driver = 'ESRI Shapefile', filename=file)
    return data

def generateCmdFile(lcps,inputs,ignites,outputs,cmdFile):
    cDir = 'commonDir/data/'
    string = ''
    for lcp, Input, ignite, output in zip(lcps,inputs,ignites,outputs):
        string = string+'%s%s %s%s %s%s 0 %s%s 0\n'%(cDir,lcp,cDir,Input,cDir,ignite,cDir,output)
    with open(cmdFile,'w') as f:
        f.write(string)

if __name__ == "__main__":

    distance = 25000

    commandFile = 'commonDir/farsite/example/Panther/runPanther.txt'
    inDir = 'E:/projects/wildfire-research/farsite/data/'
    cDir = 'commonDir/data/'
    #namespace = 'n117-9343_36-5782_3000'
    lcpNames = glob.glob(inDir+"*_25000.LCP")

    totalDays = 5

    cmdFileDockers = []
    for j in range(0,100):
        cmdFile = inDir+'toRun_'+str(j)+'.txt'
        cmdFileDocker = cDir+'toRun_'+str(j)+'.txt'
        lcpFiles = []
        inputFiles = []
        igniteFiles = []
        outputFiles = []
        for i in range(0,50):
            lcpName = np.random.choice(lcpNames).split('\\')[1].split('.LCP')[0]
            namespace = "run_"+str(j)+"_"+(str(i))+"_"+lcpName
            lcpFile = lcpName+'.LCP' # Only one landscape file per location
            igniteFile = lcpName+'_ignite.SHP' # Only one ignition file per location

            inputFile = namespace+'.input' # Different input file per simulation
            outputFile = namespace+'_out' # Different output file per simulation
            elevation = getLcpElevation(inDir+lcpFile)
            if len(glob.glob(igniteFile)) == 0:
                ignitionShape = makeCenterIgnition(inDir+lcpFile)
            fileName = inDir+namespace
            params = generateFarsiteInput(fileName,elevation,cDir+igniteFile,totalDays=totalDays)
            lcpFiles.append(lcpFile)
            inputFiles.append(inputFile)
            igniteFiles.append(igniteFile)
            outputFiles.append(outputFile)
        generateCmdFile(lcpFiles,inputFiles,igniteFiles,outputFiles,cmdFile)
        cmdFileDockers.append(cmdFileDocker)

    totalSamples = 5
    with mp.Pool(processes=6) as pool:
        q = 0
        for i in pool.imap_unordered(runFarsite,cmdFileDockers):
            q += 1
            if q % 2 == 0:
                print("Percent complete: %.2f"%(q/totalSamples*100))
                sys.stdout.flush()

    #runFarsite(cmdFileDocker)

    #dataOut = gpd.GeoDataFrame.from_file(inDir+outputFile+'_Perimeters.shp')
    #dataIn = gpd.GeoDataFrame.from_file(inDir+igniteFile)

    #filename = 'E:/projects/wildfire-research/farsite/data/californiaRaw.LCP'
    #inDir = 'E:/projects/wildfire-research/farsite/data/'

    #print(checkPoint([[lon,lat]]))
    #datas, headers, names = generateLcpFile(lats,lons,distance,inDir)

ext: py

max_stars_repo_name: johodges/datadriven-wildfire-spread

max_issues_repo_name: johodges/datadriven-wildfire-spread

max_forks_repo_name: johodges/datadriven-wildfire-spread

reason: import numpy


text:
"""
Code to generate figures used in our conference paper
"""
import numpy as np
from scipy import sparse
import scipy.io as sio
import matplotlib.pyplot as plt
from ripser import ripser
from persim import plot_diagrams as plot_dgms
from SlidingWindow import *
from Periodicity import *
from Correlations import *
import seaborn as sns

timeticks = [0, 12, 24, 36, 48, 60, 72, 84, 96, 108, 120]

def get0GrayData():
    divs = np.loadtxt("data/div_0Gy.txt")
    allndivs = np.sum(divs, 1)
    X = np.loadtxt("data/p53_0Gy.txt")
    idx = np.arange(X.shape[0])
    idx = idx[np.sum(X, 1) > 0]
    X = X[idx, :]
    allndivs = allndivs[idx]
    return X, allndivs

def getExamples():
    """
    Return an example of p53 time series with 0 divisions
    and an example with 5 divisions.  These examples will be
    used throughout all figures
    Returns
    -------
    {
        'x0': ndarray(N)
            A time series with 0 divisions,
        'x5': ndarray(N)
            A time series with 5 divisions,
        'dividx': ndarray(5)
            Indices of the 5 division events in x5,
        't': ndarray(N),
            Time in hours of each sample in x0 and x5
    }
    """
    X, ndivs = get0GrayData()
    x0 = X[ndivs == 0, :]
    x0 = x0[0, :]
    x5 = X[ndivs == 5, :]
    x5 = x5[0, :]
    divs5 = divs[ndivs == 5, :]
    divs5 = divs5[0, :]
    dividx = np.arange(divs5.size)
    dividx = dividx[divs5 == 1]
    t = 0.5*np.arange(x0.size)
    return {'x0':x0, 'x5':x5, 'dividx':dividx, 't':t}

def PeriodicityScoringExamples():
    """
    Show examples of detrending time series and scoring their
    periodicity
    """
    res = getExamples()
    x0, x5, dividx, t = res['x0'], res['x5'], res['dividx'], res['t']
    ymax = max(np.max(x0), np.max(x5))*1.1

    Win = 11
    for x, s in zip([x0, x5], ["0", "5"]):
        plt.figure(figsize=(8, 6))
        idx, score, y, x_detrended = pitch_detection(x, detrend_win=Win)
        plt.clf()
        plt.subplot(311)
        plt.plot(t, x)
        plt.xlim([0, np.max(t)])
        plt.ylim([0, ymax])
        if s == "5":
            plt.scatter(t[dividx], x[dividx], c='C1')
        plt.xticks(timeticks)
        plt.ylabel("p53 Nuclear\nFluorescence")
        plt.xlabel("Elapsed Time (Hours)")
        plt.title("%s Divisions"%s)

        plt.subplot(312)
        plt.plot(t, x_detrended)
        plt.xlim([0, np.max(t)])
        plt.ylim([-1, 1])
        if s == "5":
            plt.scatter(t[dividx], x_detrended[dividx], c='C1')
        plt.xticks(timeticks)
        plt.ylabel("p53 Nuclear\nFluorescence (Detrended)")
        plt.xlabel("Elapsed Time (Hours)")
        plt.title("%s Divisions Detrended"%s)

        plt.subplot(313)
        plt.plot(t[np.arange(y.size)], y)
        plt.stem([t[idx]], [score])
        plt.xlabel("Shift (hours)")
        plt.xticks(timeticks)
        plt.xlim([0, np.max(t)])
        plt.title("Normalized Autocorrelation (Score %.3g)"%score)
        plt.tight_layout()
        plt.savefig("Figures/Autocorrelation_%s.svg"%s, bbox_inches='tight')

def BlockConcept():
    """
    Make a figure showing what it means to extract blocks
    from a detrended time series and score them with DAPs
    """
    res = getExamples()
    x0, x5, dividx, t = res['x0'], res['x5'], res['dividx'], res['t']
    ymax = max(np.max(x0), np.max(x5))*1.1
    block_len = 48
    Win = 11
    idx, score, y, xd = pitch_detection(x0, detrend_win=Win)

    starts = [0, 40, 77]

    plt.figure(figsize=(8, 4))
    plt.subplot2grid((2, 3), (0, 0), colspan=3)
    plt.plot(t, xd)
    plt.xlim([-1, np.max(t)])
    plt.ylim([-1, 1])
    plt.xticks(timeticks)
    plt.ylabel("p53 Nuclear Fluorescence\n(Detrended)")
    plt.xlabel("Elapsed Time (Hours)")

    Win = int(block_len*0.5) # Block length 24 hours
    c = 'C1'
    AW = 0.1
    AXW = 0.005
    y1, y2 = np.min(xd), np.max(xd)
    pad = 0.3*(y2-y1)
    #c = np.array([1.0, 0.737, 0.667])
    ax = plt.gca()
    for s in starts:
        ax.arrow(s+Win, y2+0.3*pad, 4, 0, head_width = AW, head_length = 2, fc = c, ec = c, width = AXW)
        ax.arrow(s+Win, y1-0.3*pad, 4, 0, head_width = AW, head_length = 2, fc = c, ec = c, width = AXW)
        plt.plot([s, s+Win], [y1-pad, y1-pad], c=c)
        plt.plot([s, s+Win], [y2+pad, y2+pad], c=c)
        plt.plot([s, s], [y1-pad, y2+pad], c=c)
        plt.plot([s+Win, s+Win], [y1-pad, y2+pad], c=c)
    plt.title("0 Divisions Detrended Block Extraction And Scoring")

    starts = [s*2 for s in starts]
    for i, s in enumerate(starts):
        plt.subplot(2, 3, 4+i)
        xdi = xd[s:s+block_len]
        _, score, _, _ = pitch_detection(xdi)
        plt.plot(t[s:s+block_len], xd[s:s+block_len])
        plt.title("Score = %.3g"%score)
        plt.ylim([-1, 1])
        if i == 0:
            plt.ylabel("p53 Nuclear Fluorescence\n(Detrended)")
        plt.xlabel("Elapsed Time (Hours)")
    plt.tight_layout()
    plt.savefig("Figures/BlockConcept.svg", bbox_inches='tight')

def BlockDistributions():
    """
    Show the histograms of the pitch scores for all blocks in the
    0 division and 5 division example next to each other
    """
    res = getExamples()
    x0, x5, dividx, t = res['x0'], res['x5'], res['dividx'], res['t']
    ymax = max(np.max(x0), np.max(x5))*1.1
    block_len = 48
    plt.figure(figsize=(12, 6))
    Win = 11
    allscores = []
    for x, s in zip([x0, x5], ["0", "5"]):
        plt.clf()
        # Get all block scores
        k = 0
        scores = []
        idx, score, y, xd = pitch_detection(x, detrend_win=Win)
        while k+block_len < xd.size:
            xk = xd[k:k+block_len]
            k += 1
            _, score, _, _ = pitch_detection(xk)
            scores.append(score)
        allscores.append(np.array(scores))

    plt.figure(figsize=(6, 3))
    sns.distplot(allscores[0])
    sns.distplot(allscores[1])
    plt.legend(["0 Divisions", "5 Divisions"])
    plt.xlabel("Pitch Score")
    plt.ylabel("Density")
    plt.title("Pitch Score Distributions")
    plt.savefig("Figures/Distributions.svg", bbox_inches='tight')

def ROCFigures():
    plt.figure(figsize=(4, 4))
    path = "tda.mat"
    name = "RawSignal"
    k = 0
    #for k, (path, name) in enumerate(zip(["tda.mat"],
    #                                     ["TDA"])):
    res = sio.loadmat(path)
    #X, grays, allndivs = res["X"], res["grays"], res["allndivs"]
    #X = np.reshape(X, (X.size, 1))
    divs = np.loadtxt("data/div_0Gy.txt")
    allndivs = np.sum(divs, 1)
    X = np.loadtxt("data/p53_0Gy.txt")
    idx = np.arange(X.shape[0])
    idx = idx[np.sum(X, 1) > 0]
    X = X[idx, :]
    allndivs = allndivs[idx]

    #grays = grays.flatten()
    #allndivs = allndivs.flatten()
    #X = X[grays == 0, :]
    #allndivs = allndivs[grays == 0]
    #grays = grays[grays == 0]
    aurocs, results = getEuclideanAUROC(X, allndivs, divgroups = [[0], [1], [2], [3], [4], [5]])

    print("k = ", k)
    plt.subplot(1, 1, k+1)
    sns.heatmap(aurocs, annot=True, fmt='.2f')
    plt.xlabel("Number of Divisions")
    plt.ylabel("Number of Divisions")
    plt.title(name)
    plt.savefig("ROCResults.svg", bbox_inches='tight')

def makeClusterPlot(D, X, allndivs, nclusters):
    """
    Compute and plot clusters for all cells that are 0 gray,
    given a set of distances computed between all cells
    D: ndarray(N, N)
        A distance matrix between all cells
    X: ndarray(N, K)
        A set of K samples in each distribution for each time series
    allndivs: ndarray(N)
        Number of divisions for each time series
    nclusters: int
        Number of clusters to perform
    """
    from scipy.cluster.hierarchy import dendrogram, linkage, fcluster
    N = D.shape[0]
    Z = linkage(D[np.triu_indices(N, 1)], 'ward')
    cluster_idxs = fcluster(Z, nclusters, criterion='maxclust')
    allcounts = []
    labels = []
    allidxs = [] # Indices of time series in each gray/divs group
    allndivs = allndivs.flatten()
    for ndivs in np.unique(allndivs):
        idxs = cluster_idxs[allndivs == ndivs]
        if len(idxs) == 0:
            continue
        allidxs.append(idxs)
        counts = [np.sum(idxs == c) for c in range(1, nclusters+1)]
        allcounts.append(counts)
        labels.append("%g"%ndivs)
    allcounts = np.array(allcounts, dtype=float).T
    idx = np.argmax(allcounts, 0)
    allcounts /= np.sum(allcounts, 0)[None, :]

    # Compute mean periodicity scores for each cluster
    means_per_cluster = np.zeros(nclusters)
    for i in range(nclusters):
        scores = X[allndivs == i, :]
        means_per_cluster[i] = np.mean(scores)

    # Compute mean periodicity scores for each number of divisions
    means_per_type = []
    for ndivs in np.unique(allndivs):
        x = X[allndivs == ndivs, :]
        means_per_type.append(np.mean(x))

    # Sort clusters in ascending order of mean
    idx = np.argsort(means_per_cluster)
    means_per_cluster = means_per_cluster[idx]
    allcounts = allcounts[idx, :]
    # Plot the results
    plt.imshow(allcounts, vmin=0, vmax=1)
    plt.xticks(np.arange(allcounts.shape[1]), ["%s\n%.2g"%(l, m) for (l, m) in zip(labels, means_per_type)])
    plt.yticks(np.arange(len(means_per_cluster)), ["  %.2g  "%m for m in means_per_cluster])
    for x in range(allcounts.shape[1]):
        for y in range(allcounts.shape[0]):
            if allcounts[y][x] < 0.01:
                allcounts[y][x] = 0
            plt.text(x-0.3, y, "%.2g"%allcounts[y][x], color='white')
    plt.xlabel("Number of Divisions")
    plt.ylabel("Cluster")

def ClusterPlots():
    """
    Plot the results of single linkage clustering using different types
    of distances
    """
    from Wasserstein import getWasserstein
    plt.figure(figsize=(12, 10))
    nclusters=4

    # First make a plot with Euclidean on the raw data
    X, allndivs = get0GrayData()
    plt.subplot(337)
    makeClusterPlot(getSSM(X), X, allndivs, nclusters)
    plt.title("Euclidean Raw Data")

    for row, (filename, score) in enumerate(zip(["pitch.mat", "tda.mat"], ["DAPs", "Sw1PerS"])):
        res = sio.loadmat(filename)
        X, grays, allndivs = res["X"], res["grays"].flatten(), res["allndivs"].flatten()
        X = X[grays == 0, :]
        allndivs = allndivs[grays == 0]
        for col, dist in enumerate(["Euclidean", "Euclidean Sorted", "Wasserstein"]):
            plt.subplot(3, 3, row*3+col+1)
            if col == 0:
                D = getSSM(X)
            elif col == 1:
                D = getSSM(np.sort(X, 1))
            else:
                D = getWasserstein(X)
            makeClusterPlot(D, X, allndivs, nclusters)
            plt.title("%s %s"%(score, dist))
    plt.tight_layout()
    plt.savefig("Clusters.svg", bbox_inches='tight')

def DistributionPlot():
    plt.figure(figsize=(7, 6))
    for row, (filename, score) in enumerate(zip(["pitch.mat", "tda.mat"], ["DAPS", "Sw1PerS"])):
        res = sio.loadmat(filename)
        X, allndivs = res["X"], res["allndivs"].flatten()
        divgroups = [[0], [1], [2], [3], [4], [5]]
        legend = ["0", "1", "2", "3", "4", "5"]
        # Step 1: Plot periodicity distributions
        plt.subplot(2, 2, row*2 + 1)
        xs = []
        for ndivs in divgroups:
            idx = np.zeros(X.shape[0])
            for div in ndivs:
                idx[allndivs == div] = 1
            scores = X[idx == 1, :]
            sns.distplot(scores.flatten())
            xs.append(scores.flatten())
        plt.legend(legend)
        plt.title("%s Block Distributions"%score)
        plt.xlabel("Periodicity Score")
        plt.ylabel("Distribution Density")
        # Step 2: Compute and plot ROC scores
        D = np.zeros((len(divgroups), len(divgroups)))
        for i, divi in enumerate(legend):
            x = xs[i]
            for j in range(i+1, len(divgroups)):
                divj = legend[j]
                y = xs[j]
                D[i, j] = getAUROC(x, y)['auroc']
                print("%s %s divisions vs %s divisions: %.3g"%(score, divi, divj, D[i, j]))
        D += D.T
        np.fill_diagonal(D, 0.5)
        plt.subplot(2, 2, row*2+2)
        plt.imshow(D, cmap='magma', vmin=0.5, vmax=1.0)
        plt.xticks(np.arange(D.shape[0]), legend)
        plt.yticks(np.arange(D.shape[0]), legend)
        plt.xlabel("Number of Divisions")
        plt.ylabel("Number of Divisions")
        plt.title("%s AUROC Scores"%score)
        for i in range(D.shape[0]):
            for j in range(D.shape[1]):
                if i == j:
                    continue
                color = 'black'
                if D[i, j] < 0.75:
                    color = 'white'
                plt.text(j-0.45, i+0.2, "%.2g"%D[i, j], color=color)
    plt.tight_layout()
    plt.savefig("DistributionPlots.svg", bbox_inches='tight')

def ROCExample():
    res = sio.loadmat("pitch.mat")
    X, allndivs = res["X"], res["allndivs"].flatten()
    div1 = 1
    div2 = 4
    legend = ["%i Division"%div1, "%i Divisions"%div2]
    X1 = X[allndivs == div1, :]
    X2 = X[allndivs == div2, :]

    plt.figure(figsize=(10, 3))
    plt.subplot(121)
    sns.distplot(X1.flatten())
    sns.distplot(X2.flatten())
    plt.legend(legend)
    plt.title("DAPS Block Distributions")
    plt.xlabel("Periodicity Score")
    plt.ylabel("Distribution Density")
    # Step 2: Compute and plot ROC scores
    plt.subplot(122)
    auroc = getAUROC(X1.flatten(), X2.flatten(), do_plot=True)['auroc']
    plt.plot([0, 1], [0, 1], linestyle=':', linewidth=1, color='k')
    plt.title("ROC Curve, AUROC = %.3g"%auroc)
    plt.savefig("ROCExample.svg", bbox_inches='tight')

def DetrendingExample():
    """
    To help Mahdi with a block diagram of our procedure
    """
    X, ndivs = get0GrayData()
    X = X[ndivs == 0, :]
    t = np.arange(X.shape[1])*0.5

    i = 6
    Win = 11
    block_hop = 1 # Maximum data augmentation is to take all blocks
    block_len = 48

    plt.clf()
    x = X[i, :]
    y = detrend_timeseries(x, Win)
    score_fn = lambda xk: pitch_detection(xk)[1]

    scores = []
    i1 = 0
    if not block_len:
        block_len = x.size
    while i1+block_len <= y.size: # Loop over all blocks
        print("%i %i"%(i, i1))
        # Pull out block and do sliding window
        yk = np.array(y[i1:i1+block_len])
        score = score_fn(yk)
        scores.append(score)
        i1 += block_hop
    scores = np.array(scores)

    plt.figure(figsize=(12, 8))
    plt.subplot(311)
    plt.plot(t, x)
    plt.xlim([0, 120])
    plt.xlabel("Time (Hours)")
    plt.subplot(312)
    plt.plot(t, y)
    plt.xlim([0, 120])
    plt.title("Detrended")
    plt.xlabel("Time (Hours)")
    plt.subplot(313)
    plt.title("Scores")
    plt.plot(t[0:scores.size], scores)
    plt.xlim([0, 120])
    plt.tight_layout()

    plt.savefig("Detrending%i.svg"%i)

def DAPSConceptFigure():
    """
    An example of a linearly drifting and linearly amplitude modulated
    signal with oscillations hidden inside of it, which allows us
    to show all of the steps of detrending in DAPS
    """
    N = 50
    Win = 5
    x = np.arange(N)*((-1)**(np.arange(N)) + np.arange(N))
    X = getSlidingWindowNoInterp(x, Win)
    XPC = X - np.mean(X, 1)[:, None]
    Norm = np.sqrt(np.sum(XPC**2, 1))
    Norm[Norm == 0] = 1
    XPCSN = XPC/Norm[:, None]
    y = getSlidingWindowL2Inverse(XPCSN)
    Y = getSlidingWindowNoInterp(y, Win)

    plt.figure(figsize=(8, 8))
    plt.subplot(511)
    plt.plot(x)
    plt.title("Original Time Series $x$")
    plt.subplot(512)
    plt.imshow(X.T, aspect='auto', interpolation='nearest')
    plt.axis('off')
    plt.colorbar()
    plt.title("Sliding Window Matrix $X$")
    plt.subplot(513)
    plt.imshow(XPCSN.T, aspect='auto', interpolation='nearest')
    plt.axis('off')
    plt.colorbar()
    plt.title("Point-Centered/Sphere Normalized Matrix $\hat{X}$")
    plt.subplot(514)
    plt.imshow(Y.T, aspect='auto', interpolation='nearest')
    plt.axis('off')
    plt.colorbar()
    plt.title("Projected Hankel Matrix $Y$")
    plt.subplot(515)
    plt.plot(y)
    plt.tight_layout()
    plt.title("Detrended Time Series $y$")
    plt.savefig("DAPSConcept.svg", bbox_inches='tight')

if __name__ == '__main__':
    #PeriodicityScoringExamples()
    #BlockConcept()
    #BlockDistributions()
    #ROCFigures()
    #ClusterPlots()
    #DistributionPlot()
    #ROCExample()
    #DetrendingExample()
    DAPSConceptFigure()


text:
import astropy
import numpy
from astropy.io import fits
import glob
import os
import re
from matplotlib import pylab as plt
import numpy as np
import math
from scipy.signal import find_peaks_cwt
import ccdproc
from astropy.nddata import CCDData
from astropy import units as u
from scipy.interpolate import LSQBivariateSpline, LSQUnivariateSpline
from astropy.stats import sigma_clip
from scipy.optimize import fmin
from scipy.optimize import least_squares
from scipy.interpolate import interp1d
import pickle

poly_arc = {3: np.array([  1.03773471e-05,   5.78487274e-01,   4.45847046e+03]),
            7: np.array([  2.30350858e-05,  -7.64099597e-01,   9.79141140e+03])}

plt.ion()

def get_profile_model(params, ys):
#    a1, cy1, sigma1, a2, cy2, sigma2 = params
    a1, cy1, sigma1 = params

    p1 = np.exp(-(ys - cy1)**2 / 2 / sigma1**2)
    p1 /= p1.max()

#    p2 = np.exp(-(ys - cy2)**2 / 2 / sigma2**2)
#    p2 /= p2.max()
    return a1 * p1 #+ a2 * p2

def get_profile_chisq(params, ys, profile):
    model = get_profile_model(params, ys)
    return np.sum( (profile - model)**2 / np.sqrt(np.abs(profile)) ) / (profile.size - len(params))

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

def get_profile_spl(dls, stamp):
    # need to sort the data (and weights) so that the x values are increasing
    x, y = dls.ravel(), stamp.ravel()
    weights = np.sqrt(np.abs(y)) # not technically optimal for coadded data, but ok
    wsort = x.argsort()
    x, y, weights = x[wsort], y[wsort], weights[wsort]

    # set locations for spline knots
    t = np.linspace(x.min() + 1, x.max() - 1, np.int(x.max() - x.min()))
    spl = LSQUnivariateSpline(x, y, t, weights)
    return x, y, spl

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

def get_dl_model(params, dxs, dys):
    return dxs + params[0] * dys + params[1] * dys ** 2

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

def check_dl_model(params, dxs, dys, stamp):
    dls = get_dl_model(params, dxs, dys)
    x, y, spl = get_profile_spl(dls, stamp)
    chisq = np.sum((stamp - spl(dls)) ** 2 / np.sqrt(np.abs(stamp)))
    return chisq / (stamp.size - len(params))

###############################################################################
# function to fit a 2D spline
def fit_sky(xvals, yvals, image, ky=1, dx=0.5):
    # select knot points in the x (wavelength) direction
    tx = np.arange(xvals.min() + 2, xvals.max() - 2, dx)

    # select knot points in the y (spatial) direction
    ty = [] # if there are no knots, the fit will be a poly nomial of degree ky

    # fit the 2D spline
    return LSQBivariateSpline(xvals.ravel(), yvals.ravel(), image.ravel(), tx, ty, ky=ky)

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

def checkalldata(directory=False):
    if directory is not False:
        imglist = glob.glob(directory + '*')
    else:
        imglist = glob.glob('*fits')

    listhd = ['GRATENAM','SLMSKNAM','MJD-OBS','DATE-OBS','EXPTIME','AIRMASS','OBJECT',\
              'SLSELNAM','TARGNAME','LAMPS','FLAMPS','OBSTYPE','IMTYPE','RA','DEC']
    dictionary={}
    for img in imglist:
        dictionary[img]={}
        hdu = fits.open(img)
        dictionary[img]['fits']= hdu
        for _header in listhd:
            if hdu[0].header.get(_header) is not None:
                dictionary[img][_header]= hdu[0].header.get(_header)
            else:
                dictionary[img][_header]= None

        dictionary[img]['type']= None
        if dictionary[img]['SLMSKNAM'] =='GOH_X':
            dictionary[img]['type']= 'focus'
        else:
            # select arcs
            if dictionary[img]['OBSTYPE'] =='Line':
                if dictionary[img]['LAMPS']!=['Off']:
                    dictionary[img]['type']= 'arc'
            #   select flats
            if dictionary[img]['OBSTYPE'] =='IntFlat':
                if hdu[0].header.get('FLAMPS')!=['Off']:
                    dictionary[img]['type']= 'flat'
            #   select objecs
            if dictionary[img]['OBSTYPE'] =='Object':
                dictionary[img]['type']= 'object'

    setup_object={}
    setup_flat={}
    setup_arc={}
    for img in dictionary:
        print(img,dictionary[img]['type'],dictionary[img]['OBJECT'],dictionary[img]['OBSTYPE'])
        if dictionary[img]['type'] is None:
            print(dictionary[img])
            answ = input('what is it?')

        _grism = dictionary[img]['GRATENAM']
        _slit = dictionary[img]['SLMSKNAM']

        if dictionary[img]['type']=='object':
            if (_grism,_slit) not in setup_object:
                setup_object[_grism,_slit]=[]
            setup_object[_grism,_slit].append(img)

        if dictionary[img]['type']=='arc':
            if (_grism,_slit) not in setup_arc:
                setup_arc[_grism,_slit]=[]
            setup_arc[_grism,_slit].append(img)

        if dictionary[img]['type']=='flat':
            if (_grism,_slit) not in setup_flat:
                setup_flat[_grism,_slit]=[]
            setup_flat[_grism,_slit].append(img)

    return dictionary, setup_object, setup_arc, setup_flat

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

def trim_rotate_split(setup_object,setup_flat,setup_arc,dictionary,key):
    xmin3,xmax3 = 710,860
    xmin7,xmax7 = 1205,1350
    ############ split files
    for img in setup_object[key] + setup_flat[key]+setup_arc[key]:
#        img3=re.sub('.fits','',img) + '_3.fits'
#        img7=re.sub('.fits','',img) + '_7.fits'
        _header3 = dictionary[img]['fits'][3].header
        _header7 = dictionary[img]['fits'][7].header
        _data3 = np.transpose(dictionary[img]['fits'][3].data)
        _data7 = np.transpose(dictionary[img]['fits'][7].data)

        for ll in ['DATASEC','DETSIZE','DETSEC']:
            del _header3[ll]
            del _header7[ll]

        science3 = CCDData(data=_data3,header=_header3,unit=u.adu)
        science7 = CCDData(data=_data7,header=_header7,unit=u.adu)
        # add header from the
        _header3['exptime'] = dictionary[img]['EXPTIME']
        _header3['MJD-OBS'] = dictionary[img]['MJD-OBS']
        _header3['OBJECT']  = dictionary[img]['OBJECT']
        _header3['OBSTYPE'] = dictionary[img]['OBSTYPE']
        _header3['AIRMASS'] = dictionary[img]['AIRMASS']
        _header3['RA']      = dictionary[img]['RA']
        _header3['DEC']     = dictionary[img]['DEC']

        _header7['exptime'] = dictionary[img]['EXPTIME']
        _header7['MJD-OBS'] = dictionary[img]['MJD-OBS']
        _header7['OBJECT']  = dictionary[img]['OBJECT']
        _header7['OBSTYPE'] = dictionary[img]['OBSTYPE']
        _header7['AIRMASS'] = dictionary[img]['AIRMASS']
        _header7['RA']      = dictionary[img]['RA']
        _header7['DEC']     = dictionary[img]['DEC']

        #  trim images
        trimmed3 = ccdproc.trim_image(science3, fits_section='[:,' + str(xmin3) + ':' + str(xmax3) + ']')
        trimmed7 = ccdproc.trim_image(science7, fits_section='[:,' + str(xmin7) + ':' + str(xmax7) + ']')
        dictionary[img]['trimmed3'] = trimmed3
        dictionary[img]['trimmed7'] = trimmed7
    return dictionary

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

def makeflat(setup_flat,dictionary,key):
    masterflat3 = False
    masterflat7 = False
    ######### make master flat 3
    flatlist = []
    for img in setup_flat[key]:
        if 'trimmed3' in dictionary[img]:
            flatlist.append(dictionary[img]['trimmed3'].data)
            masterflat3 = np.mean(flatlist,axis=0)
        else:
            print('ERROR: Flat not trimmed3')

    ######### make master flat 7
    flatlist = []
    for img in setup_flat[key]:
        if 'trimmed7' in dictionary[img]:
            flatlist.append(dictionary[img]['trimmed7'].data)
            masterflat7 = np.mean(flatlist,axis=0)
        else:
            print('ERROR: Flat not trimmed7')
    return masterflat3, masterflat7

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

def image_plot(image):
    # determine the image pixel distribution (used for displaying below)
    sample = sigma_clip(image)
    vmin = sample.mean() - 1 * sample.std()
    vmax = sample.mean() + 3 * sample.std()
    yvals, xvals = np.indices(image.shape)
    extent = (xvals.min(), xvals.max(), yvals.min(), yvals.max())
    plt.imshow(image, origin='lower', cmap='gray', aspect='auto', vmin=vmin, vmax=vmax, extent=extent)
    plt.xlabel('Column Number')
    plt.ylabel('Row Number');

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

def find_sky_object(image,objectcut=0.2,key=3,interactive=False):
    # get the a pixel coordinate near the image center
    ny, nx = image.shape
    cy, cx = ny//2, nx//2
    # create 1d arays of the possible x and y values
    xs = np.arange(nx)
    ys = np.arange(ny)
    # pixel coordinates for each pixel
    yvals, xvals = np.indices(image.shape)

    # compute the row averages and normalize so that the background is near 0 and the peaks are near 1
    rowaverage = image.mean(axis=1)
#    rowaverage -= np.median(rowaverage)
    rowaverage -= np.percentile(rowaverage,0.1)
    rowaverage /= rowaverage.max()

    image_plot(image)
    fig,ax = plt.subplots()
    ax.plot(ys, rowaverage)
    ax.set_xlabel('Row Number (y-coordinate)'), plt.ylabel('Normalized Row Average')
    plt.grid();
    # find the rows with object light
    objrows = ys[rowaverage > objectcut]

    rangeobj = {3:'70,100', 7:'55,90'}
    if interactive:
        obj = input('give object position [eg ' + str(rangeobj[key]) + ']')
        if not obj:
            obj = rangeobj[key]
        start,end= obj.split(',')
        objrows = np.arange(int(start),int(end))
    else:
        # add some margin to object rows
        ngrow = 5 # number of rows to include above and below object rows
        newobjrows = []
        for row in objrows:
            newobjrows.extend([row + i for i in np.arange(-ngrow, ngrow + 1)])
        objrows = np.unique(newobjrows)

    # mask to mark sky rows
    skymask = np.ones(image.shape, dtype=bool)
    skymask[objrows, :] = False

    # also exclude bad rows
    badrows = ys[rowaverage < -0.05]
    skymask[badrows, :] = False

    # rows with mostly sky background light
    skyrows = ys[skymask.mean(axis=1) == 1]

    # median (unmasked) sky spectrum and standard deviation
    medspec = np.median(image[skyrows, :], axis=0)
    stdspec = np.std(image[skyrows, :], axis=0, ddof=1)

    # exclude deviant pixels from the skymask
    pull = (image - medspec) / stdspec
    w = pull > 5
    skymask[w] = False

    if interactive:
        plt.close()
        plt.clf()
        plt.figure(figsize=(15, 3))
        plt.imshow(skymask, origin='lower', aspect='auto');

    plt.close()
    return objrows, skymask

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

def retify_frame(img0,ext=3, verbose=False):
    image = dictionary[img0]['trimmed'+str(ext)].data
    # this is an arc, we do not ned to mask
    skymask = np.ones(image.shape, dtype=bool)

    # show the mask
    if verbose:
        plt.close()
        plt.figure(figsize=(15, 3))
        plt.imshow(skymask, origin='lower', aspect='auto');
        input('stop')

    # cut out a small image "stamp" near the center of the frame
    ny, nx = image.shape
    cy, cx = ny//2, nx//2
    xs = np.arange(nx)
    ys = np.arange(ny)
    # pixel coordinates for each pixel
    yvals, xvals = np.indices(image.shape)

    row = cy
    col = cx
    hwidth = 50

    stamp = image[:, col - hwidth : col + hwidth]
    yvals, xvals = np.indices(image.shape)
    ys_stamp = yvals[:, col - hwidth : col + hwidth]
    xs_stamp = xvals[:, col - hwidth : col + hwidth]

    if verbose:
        image_plot(stamp)
        plt.close()

        # plot stamp values against column numbers
        plt.close()
        plt.clf()
        plt.plot(xs_stamp.ravel(), stamp.ravel(), 'r.');
        plt.xlabel('Column Number'), plt.ylabel('Counts');
        input('stop')
        plt.close()

    # pixel offsets from the refernece pixel
    dxs = xs_stamp - col
    dys = ys_stamp - row

    # parameter guess
    guess = (0, 0)

    # get the wavelength offsets and plot vs. counts
    dls = get_dl_model(guess, dxs, dys)
    plt.close()
    plt.clf()
    plt.plot(dls.ravel(), stamp.ravel(), 'r.')
    plt.xlabel('Wavelength Offset')
    plt.ylabel('Counts');
    input('stop')

    # fit a spline to the data and plot
    x, y, spl = get_profile_spl(dls, stamp)

    plt.clf()
    fig, axarr = plt.subplots(2, sharex=True)
    axarr[0].plot(x, y, 'r.')
    axarr[0].plot(x, spl(x))
    axarr[1].plot(x, y - spl(x), 'r.')
    plt.ylim(-200, 200)
    plt.xlabel('Wavelength Offset');
    input('stop')
    plt.close()

    # see how good our guess is
    check_dl_model(guess, dxs, dys, stamp)

    # get the best model parameters for this stamp
    params = fmin(check_dl_model, guess, args=(dxs, dys, stamp))
    print("best model parameters are", params)

    hwidth = 300 # width = 2 * hwidth + 1
    cols = np.arange(hwidth, nx, 2 * hwidth)
    cols = cols[1:]

    # reference wavelength offsets to the center row
    row = cy

    # define a 2D array to hold the wavelength offsets for each pixel
    lambdas = np.zeros(image.shape)

    # loop over each central column
    for col in cols:
        print('col = ', col)

        # slice the data
        inds = np.s_[:, col - hwidth : col + hwidth]
        stamp = image[inds]
        mask = skymask[inds]
        dys = yvals[inds] - row
        dxs = xvals[inds] - col

        # initial fit
        params = fmin(check_dl_model, guess, args=(dxs[mask], dys[mask], stamp[mask]))

        # check for outliers
        dls = get_dl_model(guess, dxs, dys)
        x, y, spl = get_profile_spl(dls, stamp)
        model = spl(dls)
        pull = (stamp - model) / np.sqrt(np.abs(model))
        w = (pull < 5) & mask
        params2 = fmin(check_dl_model, params, args=(dxs[w], dys[w], stamp[w]))

        # record
        lambdas[inds] = get_dl_model(params2, dxs, dys) + col
        pickle.dump(lambdas, open('lambdas' + str(ext) + '.dat', 'wb'))

        # just plot offsets for a few of the rows across the image
        plt.clf()
        order = 2
        for y in range(10, ny, 40):
            p = plt.plot(cols, lambdas[y, cols] - xs[cols], 'o')
            c = np.polyfit(cols, lambdas[y, cols] - xs[cols], order)
            plt.plot(xs, np.polyval(c, xs), c=p[0].get_color(), label='row {}'.format(y))

        plt.legend()
        plt.xlabel('Column Number')
        plt.ylabel('Wavelength Offset from Middle Row');
    input('stop')
    return lambdas

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

def model_sky(params, dx, counts):
#    wav0, a, b, c, d, scale, c0 = params
    wav0, a, b, scale, c0 = params

    dtype = []
    dtype.append(('wav', float))
    dtype.append(('flux', float))
    model = np.zeros(counts.shape, dtype=dtype)
#    model['wav'] = wav0 + a * dx + b * dx**2 + c * dx**3 +  d * dx**4
    model['wav'] = wav0 + a * dx + b * dx**2
    model['flux'] = c0 + scale * counts

    return model

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

def get_sky_difference(params, dx, specdata, skyatlas):
    model = model_sky(params, dx, specdata)

    # residual
    res = model['flux'] - skyref_interp(model['wav'])

    return np.sum(res**2 / np.sqrt(np.abs(model['flux'])))

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

def plot_sky_model(skyref, model):
    plt.plot(model['wav'], model['flux'], label='Extracted Sky Background');
    plt.plot(skyref['wav'], skyref['flux'], label='Reference Sky Spectrum');
    plt.xlim(model['wav'].min(), model['wav'].max());
    plt.ylim(model['flux'].min() - 0.25, model['flux'].max() + 0.25)

    plt.xlabel('Wavelength ($\AA$)')
    plt.ylabel('Scaled Counts or Flux')
    plt.legend();

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

def trace(img,dictionary, g0=10, g1=85, g2=3, key=3):
    if 'sky' + str(key) in dictionary[img]:
        sky = dictionary[img]['sky' + str(key)].data
        image = dictionary[img]['trimmed' + str(key)].data
        nosky = image - sky

        # get the a pixel coordinate near the image center
        ny, nx = image.shape
        cy, cx = ny//2, nx//2
        # create 1d arays of the possible x and y values
        xs = np.arange(nx)
        ys = np.arange(ny)
        # pixel coordinates for each pixel
        yvals, xvals = np.indices(image.shape)

        # get the median for each row
        profile = np.median(image - sky, axis=1)

        # starting guess for the profile model
        guess = (g0, g1, g2)
        model = get_profile_model(guess, ys)

        # fit for the best model
        params = fmin(get_profile_chisq, guess, args=(ys, profile))
        print("best fit parameters are", params)

        model = get_profile_model(params, ys)

        # fit the profile centered at these columns
        hwidth = 50
        cols = np.arange(hwidth, nx + 1, 2 * hwidth)

        ycenter = np.zeros( (len(cols), 2) )
        for icol, col in enumerate(cols):
             stamp = (image - sky)[:, col - hwidth : col + hwidth]
             profile = np.mean(stamp, axis=1)
             params = fmin(get_profile_chisq, guess, args=(ys, profile))
             ycenter[icol, :] = params[[1]]

        # fit the relation with a polynomial
        ind = 0 # which trace 0 or 1?
        t_order = 3
        trace_c = np.polyfit(cols, ycenter[:, ind], t_order)
        fig, axarr = plt.subplots(2, sharex=True)
        axarr[0].plot(cols, ycenter[:, ind], 'ro')
        axarr[0].plot(xs, np.polyval(trace_c, xs), 'r')
        axarr[0].axes.set_ylabel('y-coordinate')
        axarr[1].plot(cols, ycenter[:, ind] - np.polyval(trace_c, cols), 'ro')
        axarr[1].axes.set_ylim(-0.5, 0.5)
        axarr[1].axes.set_ylabel('Fit Residual (pixels)')
        plt.xlabel('Column Number');
        input('trace completed')

        slitpos = yvals - np.polyval(trace_c, yvals)

        dictionary[img]['trace' + str(key)] = slitpos
        dictionary[img]['tracefit' + str(key)] = trace_c
    return dictionary

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

def extract(img,dictionary, key=3):
    if 'trace' + str(key) in dictionary[img]:
        sky = dictionary[img]['sky' + str(key)].data
        image = dictionary[img]['trimmed' + str(key) ].data
        nosky = image - sky

        # get the a pixel coordinate near the image center
        ny, nx = image.shape
        cy, cx = ny//2, nx//2
        # create 1d arays of the possible x and y values
        xs = np.arange(nx)
        ys = np.arange(ny)
        # pixel coordinates for each pixel
        yvals, xvals = np.indices(image.shape)

        slitpos = dictionary[img]['trace' + str(key)]
        trace_c = dictionary[img]['tracefit' + str(key)]

        # normalize to the pixel brightness at the trace center
        yinds = (np.round(np.polyval(trace_c, xs))).astype(int)
        normed = nosky / nosky[yinds, xs]

        # get 1D arrays with the positions along the slit and the normalized counts
        pos = slitpos.flatten()
        counts = normed.flatten()

        # sort by slit position
        sort_inds = pos.argsort()
        pos, counts = pos[sort_inds], counts[sort_inds]

        # fit a spline to model the spatial profile
        t = np.linspace(pos.min() + 2, pos.max() - 2, ny // 2) # spline knot points
        profile_spl = LSQUnivariateSpline(pos, counts, t)

        # remove outliers and re-fit
        diff = counts - profile_spl(pos)
        sample = sigma_clip(diff)
        w = ((np.abs(diff) / sample.std()) < 5) & np.isfinite(diff)
        profile_spl = LSQUnivariateSpline(pos[w], counts[w], t)

        # create the profile image
        profile_image = profile_spl(slitpos)

        # de-weight negative values in provile_image
        profile_image[profile_image < 0] = 0

        # select which rows to sum
        w = (slitpos > -10) & (slitpos < 10)
        ymin, ymax = yvals[w].min(), yvals[w].max()

        # calculate the sum
        spec_basic = nosky[ymin:ymax, :].sum(axis=0)

        # sky background
        skybg_basic = np.array(sky)[ymin:ymax, :].sum(axis=0)

        # calculate the weighted average (for each column)
        spec_opt = (nosky * profile_image)[ymin:ymax, :].sum(axis=0) / profile_image.sum(axis=0)

        # calculate the bias factor needed to scale the average to a sum
        bias_factor = np.median(spec_basic / spec_opt)
        spec_opt *= bias_factor

        # same for the sky background
        skybg_opt = (sky * profile_image)[ymin:ymax, :].sum(axis=0) / profile_image.sum(axis=0)
        bias_factor_sky = np.median(skybg_basic / skybg_opt)
        skybg_opt *= bias_factor_sky
        skybg_opt -= np.min(skybg_opt)

        # plot the extracted spectrum
        plt.clf()
        plt.plot(xs, spec_basic, label='basic extraction')
        plt.plot(xs, spec_opt, label='optimal extraction')
        plt.legend()
        dictionary[img]['spec_basic' + str(key)]= spec_basic
        dictionary[img]['spec_opt' + str(key)]= spec_opt
        dictionary[img]['skybg_opt' + str(key)]= skybg_opt
        input('extraction completed')
    return dictionary

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

def fit_another(xs, cols, hwidths, fitparams, i, plot=True):
    col = cols[i]
    hwidth = hwidths[i]

    # find the closest matching index with a valid fit
    inds = np.arange(cols.size)
    w = np.isfinite(fitparams[:, 0]) & (inds != i)
    iref = inds[w][np.argmin(np.abs(inds[w] - i))]
    print("reference index is ", iref)

    # get the wavelength guess
    if w.sum() == 0:
        print("not enough data to make guess...bye!")
        return
    if w.sum() < 4:
        guess = fitparams[iref, :].copy()
        guess[0] = guess[0] + guess[1] * (cols[i] - cols[iref])
    else:
        if w.sum() > 9:
            order = 3
        elif w.sum() > 6:
            order = 2
        else:
            order = 1
        print("order is", order)
        cwav = np.polyfit(cols[w], fitparams[w, 0], order)
        cdwav = np.polyfit(cols[w], fitparams[w, 1], order)
        cscale = np.polyfit(cols[w], fitparams[w, 3], order)
        guess = (np.polyval(cwav, col), np.polyval(cdwav, col), 0.0, np.polyval(cscale, col), 0.0)
    print("guess is", guess)

    w = (xs > (cols[i] - hwidths[i])) & (xs < (cols[i] + hwidths[i]))
    output = fmin(get_sky_difference, guess, args=(xs[w] - cols[i], bgspec[w], skyref)
                  , full_output=True, disp=False, maxiter=10000)
    if output[4] == 0:
        fitparams[i, :] = output[0]
        print(output[0])

        if plot:
            model = model_sky(output[0], xs[w] - col, bgspec[w])
            plot_sky_model(skyref, model)
            plt.title('Fit to Section {}'.format(i))
    else:
        print("fit failed!")

#######################################################
#            # apply flat
#            if masterflat3 is not False:
#                # get the a pixel coordinate near the image center
#                ny, nx = masterflat3.shape
#                cy, cx = ny//2, nx//2
#                # create 1d arays of the possible x and y values
#                xs = np.arange(nx)
#                ys = np.arange(ny)
#                # pixel coordinates for each pixel
#                yvals, xvals = np.indices(masterflat3.shape)
#                from scipy.interpolate import LSQBivariateSpline, LSQUnivariateSpline

#                lambdafit = np.zeros(masterflat3.shape)
#                x = lambdafit.ravel()
#                y = masterflat3.ravel()
#                weights = np.sqrt(np.abs(y))
#                wsort = x.argsort()
#                x, y, weights = x[wsort], y[wsort], weights[wsort]
#                t = np.linspace(x.min() + 1, x.max() - 1, nx)
#                spl = LSQUnivariateSpline(x, y, t, weights)
#                plt.plot(x, spl(x))
#                plt.xlabel('Wavelength ($\AA$)')
#                plt.ylabel('Counts');
#                input('stop')
###############################################################

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

max_stars_repo_name: svalenti/deimos

max_issues_repo_name: svalenti/deimos

max_forks_repo_name: svalenti/deimos

reason: import numpy,from scipy,import astropy,from astropy


text:
from __future__ import print_function

import glob
import os
import os.path as osp
import warnings

import librosa
import numpy as np
import pandas as pd
import soundfile
import torch
from desed.utils import create_folder
from torch import nn

def read_audio(path, target_fs=None):
    """Read a wav file
    Args:
        path: str, path of the audio file
        target_fs: int, (Default value = None) sampling rate of the returned audio file, if not specified, the sampling
            rate of the audio file is taken

    Returns:
        tuple
        (numpy.array, sampling rate), array containing the audio at the sampling rate given

    """
    (audio, fs) = soundfile.read(path)
    if audio.ndim > 1:
        audio = np.mean(audio, axis=1)
    if target_fs is not None and fs != target_fs:
        audio = librosa.resample(audio, orig_sr=fs, target_sr=target_fs)
        fs = target_fs
    return audio, fs

def weights_init(m):
    """Initialize the weights of some layers of neural networks, here Conv2D, BatchNorm, GRU, Linear
        Based on the work of Xavier Glorot
    Args:
        m: the model to initialize
    """
    classname = m.__class__.__name__
    if classname.find("Conv2d") != -1:
        nn.init.xavier_uniform_(m.weight, gain=np.sqrt(2))
        m.bias.data.fill_(0)
    elif classname.find("BatchNorm") != -1:
        m.weight.data.normal_(1.0, 0.02)
        m.bias.data.fill_(0)
    elif classname.find("GRU") != -1:
        for weight in m.parameters():
            if len(weight.size()) > 1:
                nn.init.orthogonal_(weight.data)
    elif classname.find("Linear") != -1:
        m.weight.data.normal_(0, 0.01)
        m.bias.data.zero_()

def to_cuda_if_available(*args):
    """Transfer object (Module, Tensor) to GPU if GPU available
    Args:
        args: torch object to put on cuda if available (needs to have object.cuda() defined)

    Returns:
        Objects on GPU if GPUs available
    """
    res = list(args)
    if torch.cuda.is_available():
        for i, torch_obj in enumerate(args):
            res[i] = torch_obj.cuda()
    if len(res) == 1:
        return res[0]
    return res

class SaveBest:
    """Callback to get the best value and epoch
    Args:
        val_comp: str, (Default value = "inf") "inf" or "sup", inf when we store the lowest model, sup when we
            store the highest model
    Attributes:
        val_comp: str, "inf" or "sup", inf when we store the lowest model, sup when we
            store the highest model
        best_val: float, the best values of the model based on the criterion chosen
        best_epoch: int, the epoch when the model was the best
        current_epoch: int, the current epoch of the model
    """

    def __init__(self, val_comp="inf"):
        self.comp = val_comp
        if val_comp in ["inf", "lt", "desc"]:
            self.best_val = np.inf
        elif val_comp in ["sup", "gt", "asc"]:
            self.best_val = 0
        else:
            raise NotImplementedError("value comparison is only 'inf' or 'sup'")
        self.best_epoch = 0
        self.current_epoch = 0

    def apply(self, value):
        """Apply the callback
        Args:
            value: float, the value of the metric followed
        """
        decision = False
        if self.current_epoch == 0:
            decision = True
        if (self.comp == "inf" and value < self.best_val) or (self.comp == "sup" and value > self.best_val):
            self.best_epoch = self.current_epoch
            self.best_val = value
            decision = True
        self.current_epoch += 1
        return decision

class EarlyStopping:
    """Callback to stop training if the metric have not improved during multiple epochs.
    Args:
        patience: int, number of epochs with no improvement before stopping the model
        val_comp: str, (Default value = "inf") "inf" or "sup", inf when we store the lowest model, sup when we
            store the highest model
    Attributes:
        patience: int, number of epochs with no improvement before stopping the model
        val_comp: str, "inf" or "sup", inf when we store the lowest model, sup when we
            store the highest model
        best_val: float, the best values of the model based on the criterion chosen
        best_epoch: int, the epoch when the model was the best
        current_epoch: int, the current epoch of the model
    """

    def __init__(self, patience, val_comp="inf", init_patience=0):
        self.patience = patience
        self.first_early_wait = init_patience
        self.val_comp = val_comp
        if val_comp == "inf":
            self.best_val = np.inf
        elif val_comp == "sup":
            self.best_val = 0
        else:
            raise NotImplementedError("value comparison is only 'inf' or 'sup'")
        self.current_epoch = 0
        self.best_epoch = 0

    def apply(self, value):
        """Apply the callback

        Args:
            value: the value of the metric followed
        """
        current = False
        if self.val_comp == "inf":
            if value < self.best_val:
                current = True
        if self.val_comp == "sup":
            if value > self.best_val:
                current = True
        if current:
            self.best_val = value
            self.best_epoch = self.current_epoch
        elif self.current_epoch - self.best_epoch > self.patience and self.current_epoch > self.first_early_wait:
            self.current_epoch = 0
            return True
        self.current_epoch += 1
        return False

class AverageMeterSet:
    def __init__(self):
        self.meters = {}

    def __getitem__(self, key):
        return self.meters[key]

    def update(self, name, value, n=1):
        if name not in self.meters:
            self.meters[name] = AverageMeter()
        self.meters[name].update(value, n)

    def reset(self):
        for meter in self.meters.values():
            meter.reset()

    def values(self, postfix=""):
        return {name + postfix: meter.val for name, meter in self.meters.items()}

    def averages(self, postfix="/avg"):
        return {name + postfix: meter.avg for name, meter in self.meters.items()}

    def sums(self, postfix="/sum"):
        return {name + postfix: meter.sum for name, meter in self.meters.items()}

    def counts(self, postfix="/count"):
        return {name + postfix: meter.count for name, meter in self.meters.items()}

    def __str__(self):
        string = ""
        for name, meter in self.meters.items():
            fmat = ".4f"
            if meter.val < 0.01:
                fmat = ".2E"
            string += "{} {:{format}} \t".format(name, meter.val, format=fmat)
        return string

class AverageMeter:
    """Computes and stores the average and current value"""

    def __init__(self):
        self.reset()

    def reset(self):
        self.val = 0
        self.avg = 0
        self.sum = 0
        self.count = 0

    def update(self, val, n=1):
        self.val = val
        self.sum += val * n
        self.count += n
        self.avg = self.sum / self.count

    def __format__(self, format):
        return "{self.avg:{format}}".format(self=self, format=format)

def generate_tsv_wav_durations(audio_dir, out_tsv):
    """Generate a dataframe with filename and duration of the file
    Args:
        audio_dir: str, the path of the folder where audio files are (used by glob.glob)
        out_tsv: str, the path of the output tsv file

    Returns:
        pd.DataFrame, the dataframe containing filenames and durations
    """
    meta_list = []
    for file in glob.glob(os.path.join(audio_dir, "*.wav")):
        d = soundfile.info(file).duration
        meta_list.append([os.path.basename(file), d])
    meta_df = pd.DataFrame(meta_list, columns=["filename", "duration"])
    if out_tsv is not None:
        meta_df.to_csv(out_tsv, sep="\t", index=False, float_format="%.1f")
    return meta_df

def generate_tsv_from_isolated_events(wav_folder, out_tsv=None):
    """Generate list of separated wav files in a folder and export them in a tsv file
    Separated audio files considered are all wav files in 'subdirectories' of the 'wav_folder'
    Args:
        wav_folder: str, path of the folder containing subdirectories (one for each mixture separated)
        out_tsv: str, path of the csv in which to save the list of files
    Returns:
        pd.DataFrame, having only one column with the filename considered
    """
    if out_tsv is not None and os.path.exists(out_tsv):
        source_sep_df = pd.read_csv(out_tsv, sep="\t")
    else:
        source_sep_df = pd.DataFrame()
        list_dirs = [d for d in os.listdir(wav_folder) if osp.isdir(osp.join(wav_folder, d))]
        for dirname in list_dirs:
            list_isolated_files = []
            for directory, subdir, fnames in os.walk(osp.join(wav_folder, dirname)):
                for fname in fnames:
                    if osp.splitext(fname)[1] in [".wav"]:
                        # Get the level folders and keep it in the tsv
                        subfolder = directory.split(dirname + os.sep)[1:]
                        if len(subfolder) > 0:
                            subdirs = osp.join(*subfolder)
                        else:
                            subdirs = ""
                        # Append the subfolders and name in the list of files
                        list_isolated_files.append(osp.join(dirname, subdirs, fname))
                    else:
                        warnings.warn(
                            f"Not only wav audio files in the separated source folder,"
                            f"{fname} not added to the .tsv file"
                        )
            source_sep_df = source_sep_df.append(pd.DataFrame(list_isolated_files, columns=["filename"]))
        if out_tsv is not None:
            create_folder(os.path.dirname(out_tsv))
            source_sep_df.to_csv(out_tsv, sep="\t", index=False, float_format="%.3f")
    return source_sep_df

def meta_path_to_audio_dir(tsv_path):
    return os.path.splitext(tsv_path.replace("metadata", "audio"))[0]

def audio_dir_to_meta_path(audio_dir):
    return audio_dir.replace("audio", "metadata") + ".tsv"

def get_durations_df(gtruth_path, audio_dir):
    path, ext = os.path.splitext(gtruth_path)
    path_durations_synth = path + "_durations" + ext
    if not os.path.exists(path_durations_synth):
        durations_df = generate_tsv_wav_durations(audio_dir, path_durations_synth)
    else:
        durations_df = pd.read_csv(path_durations_synth, sep="\t")
    return durations_df

max_stars_repo_name: m-koichi/ConformerSED

max_issues_repo_name: m-koichi/ConformerSED

max_forks_repo_name: m-koichi/ConformerSED


text:
"""
Module containing utility functions needed for PRMSE simulations.

:author: Nitin Madnani
:author: Anastassia Loukina
:organization: ETS
:date: March 2020
"""

import itertools

import numpy as np
import pandas as pd
from rsmtool.analyzer import Analyzer
from rsmtool.utils.prmse import prmse_true
from scipy.stats import pearsonr

def get_rater_pairs(rater_ids, num_pairs, seed=1234567890):
    """
    Randomly sample given number of rater pairs from given rater IDs.

    Parameters
    ----------
    rater_ids : list of str
        A list of rater IDs from which we have to randomly sample pairs.
    num_pairs : int
        Number of rater pairs we want to sample.
    seed : int, optional
        The seed for the random number generator that will be
        used to sample the rater pairs.
        Defaults to 1234567890.

    Returns
    -------
    rater_pairs : list of str
        A list containing the required number of randomly sampled
        rater pairs. Each pair is of the form "h_X+h_Y".
    """
    # first we generate all possible rater pair combinations
    all_pairs = [f"{rater1}+{rater2}" for (rater1, rater2)
                 in itertools.combinations(rater_ids, 2)]

    # next we randomly sample as many rater pairs as we need
    prng = np.random.RandomState(seed)
    chosen_pairs = prng.choice(all_pairs, size=num_pairs, replace=False)
    return [pair.split('+') for pair in chosen_pairs]

def compute_agreement_one_system_one_rater_pair(df_scores,
                                                system_id,
                                                rater_id1,
                                                rater_id2,
                                                include_mean=False):
    """
    Evaluate the given system against the given pair of raters.

    This function computes the agreement metrics between the scores
    assigned by the given simulated system (``system_id``) against the scores
    assigned by the two simulated raters ``rater_id1`` and ``rater_id2``.

    The agreement metrics computed are: Pearson's correlation, adjusted R^2,
    quadratically-weighted kappa, and the difference between the human-human
    Pearson correlation and the human-machine Pearson correlation (commonly
    known as "degradation"). All 4 metrics are computed against the scores
    of the first rater in the pair and, if ``include_mean`` is ``True``, also
    against the average of the scores assigned by both raters in the pair.

    Parameters
    ----------
    df_scores : pandas.DataFrame
        The data frame containing the simulated scores.
        This is usually one of the data frames returned
        by the ``simulation.dataset.Dataset.to_frame()``
        method.
    system_id : str
        The ID for the simulated system to be evaluated.
        This must be a column in ``df_scores``.
    rater_id1 : str
        The ID for the first rater in the rater pair
        being used to evaluate the given system.
        This must be a column in ``df_scores``.
    rater_id2 : str
        The ID for the second rater in the rater pair
        being used to evaluate the given system.
    include_mean : bool, optional
        If set to ``True``, also include the metric values
        computed against the average of the scores assigned
        by both raters in the given pair.

    Returns
    -------
    metrics_series : list of pandas.Series
        A list containing 1 or 2 pandas series depending on the value
        of ``include_mean``. If it is ``True``, this list contains
        two series:  the first containing the values of the metrics
        against the average of the two rater scores and the second
        containing the value of the metrics against the scores of
        the first rater. If ``include_mean`` is ``False``, this list
        only contains a single series - that containing the metric
        values against the scores of the first rater. Any series
        returned will contains the following columns:
        1. "r" - the Pearson's correlation between the system score
           and the average and the first rater scores.
        2. "QWK" - the quadratically-weighted kapp between the system score
           and the average and the first rater scores.
        3. "R2" - the R^2 score between the system score
           and the average and the first rater scores.
        4. "degradation" - the difference between the human-human correlation
           score and the human-machine correlation score. Note that this column
           may not be included in the output if any of the scores for either of
           the two simulated raters are null, e.g., if some of the responses are
           single scored.
        5. "reference" - a column containing whether the metric values were
           computed against the average of the two rater scores (``h1-h2 mean``)
           or the first rater scores (``h1``).
    """
    # compute the inter-rater correlation that we need for degradation
    try:
        rater1_rater2_correlation = pearsonr(df_scores[rater_id1], df_scores[rater_id2])[0]
    except ValueError:
        rater1_rater2_correlation = None
        pass

    # we only want these 3 metrics to start with
    chosen_metrics = ['wtkappa', 'corr', 'R2']

    # compute the metrics against the first rater as a series
    h1_metric_values = Analyzer.metrics_helper(df_scores[rater_id1], df_scores[system_id])
    h1_metric_values = h1_metric_values[chosen_metrics]

    # compute the degradation values
    if rater1_rater2_correlation:
        h1_metric_values['degradation'] = rater1_rater2_correlation - h1_metric_values['corr']

    # add a new column called "reference" indicating whether we used
    # the h1-h2 average score for just the h1 score
    h1_metric_values['reference'] = 'h1'

    # rename some of the metrics to have more recognizable names
    h1_metric_values.rename({'wtkappa': 'QWK', 'corr': 'r'}, inplace=True)

    # compute the metrics against the average ot the two rater scores
    # as a series if it was requested
    if include_mean:
        mean_metric_values = Analyzer.metrics_helper(df_scores[[rater_id1, rater_id2]].mean(axis=1),
                                                     df_scores[system_id])
        mean_metric_values = mean_metric_values[chosen_metrics]

        if rater1_rater2_correlation:
            mean_metric_values['degradation'] = rater1_rater2_correlation - mean_metric_values['corr']
        mean_metric_values['reference'] = 'h1-h2 mean'
        mean_metric_values.rename({'wtkappa': 'QWK', 'corr': 'r'}, inplace=True)

    # return the right number of metric series
    ans = [mean_metric_values, h1_metric_values] if include_mean else [h1_metric_values]
    return ans

def compute_agreement_one_system_multiple_rater_pairs(df_scores,
                                                      system_id,
                                                      rater_pairs,
                                                      include_mean=False):
    """
    Compute agreement for system against all given rater pairs.

    This function computes the values of conventional metrics of agreement
    between the scores of the given system (``system_id``) against the scores
    assigned by the two simulated raters ``rater_id1`` and ``rater_id2``.

    This function simply calls the ``compute_agreement_one_system_one_rater_pair()``
    function for each rater pair and combines the output. Refer
    to ``compute_agreements_for_rater_pair()`` for more details.

    Parameters
    ----------
    df_scores : pandas.DataFrame
        The data frame containing the simulated scores.
        This is usually one of the data frames returned
        by the ``simulation.dataset.Dataset.to_frame()``
        method.
    system_id : str
        The ID for the simulated system to be evaluated.
        This must be a column in ``df_scores``.
        Description
    rater_pairs : list of lists of str
        A list containing rater pairs against which
        the system is to be evaluated. Each rater
        pair is a list of rater ID, e.g.,
        ``[h_1, h_33]``.
    include_mean : bool, optional
        If set to ``True``, also include the metric values
        computed against the average of the scores assigned
        by both raters in the given pair.

    Returns
    -------
    df_metrics : pandas.DataFrame
        A pandas data frame which has the same columns as the series
        returned by ``compute_agreements_for_rater_pair()``.
    """
    # initialize a list that will hold the series
    metrics_for_all_pairs = []

    # iterate over each given rater pair
    for rater_id1, rater_id2 in rater_pairs:

        # call the per-pair function
        metrics_for_this_pair = compute_agreement_one_system_one_rater_pair(df_scores,
                                                                            system_id,
                                                                            rater_id1,
                                                                            rater_id2,
                                                                            include_mean=include_mean)
        # add the returned series to the list of all metrics
        metrics_for_all_pairs.extend(metrics_for_this_pair)

    # create a data frame from the lists of series
    df_metrics = pd.DataFrame(metrics_for_all_pairs)

    return df_metrics

def compute_agreement_multiple_systems_one_rater_pair(df_scores,
                                                      system_ids,
                                                      rater_id1,
                                                      rater_id2,
                                                      include_mean=False):
    """
    Compute agreement for given systems against the given rater pair.

    This function computes the values of conventional metrics of agreement
    between the scores of the given list of systems (with IDs ``system_ids``)
    against the scores assigned by the two simulated raters ``rater_id1``
    and ``rater_id2``.

    This function simply calls the ``compute_agreement_one_system_one_rater_pair()``
    function for each system and combines the output. Refer to
    ``compute_agreements_for_rater_pair()`` for more details.

    Parameters
    ----------
    df_scores : pandas.DataFrame
        The data frame containing the simulated scores.
        This is usually one of the data frames returned
        by the ``simulation.dataset.Dataset.to_frame()``
        method.
    system_ids : list of str
        The list of IDs of the simulated systems to be evaluated.
        Each ID must be a column in ``df_scores``.
    rater_id1 : str
        The ID for the first rater in the rater pair
        being used to evaluate the given system.
        This must be a column in ``df_scores``.
    rater_id2 : str
        The ID for the second rater in the rater pair
        being used to evaluate the given system.
        This must be a column in ``df_scores``.

    Returns
    -------
    data frame
        Description
    """
    # initialize an empty list we will use to save each system ID's results
    metrics = []

    # iterate over each system ID
    for system_id in system_ids:

        # compute the metric series for this system ID against the rater pair
        metric_series = compute_agreement_one_system_one_rater_pair(df_scores,
                                                                    system_id,
                                                                    rater_id1,
                                                                    rater_id2,
                                                                    include_mean=include_mean)
        # save the current system ID in the same series
        for series in metric_series:
            series['system_id'] = system_id

        # save the series in the list
        metrics.extend(metric_series)

    # convert the list of series into a data frame and return
    return pd.DataFrame(metrics)

def compute_prmse_one_system_multiple_rater_pairs(df_scores, system_id, rater_pairs):
    """
    Compute the PRMSE score for the system against all given rater pairs.

    This function computes the value of the PRMSE metric between
    the scores of the given system (``system_id``) against the scores
    assigned by the two simulated raters ``rater_id1`` and ``rater_id2``.

    Parameters
    ----------
    df_scores : pandas.DataFrame
        The data frame containing the simulated scores.
        This is usually one of the data frames returned
        by the ``simulation.dataset.Dataset.to_frame()``
        method.
    system_id : str
        The ID for the simulated system to be evaluated.
        This must be a column in ``df_scores``.
        Description
    rater_pairs : list of lists of str
        A list containing rater pairs against which
        the system is to be evaluated. Each rater
        pair is a list of rater ID, e.g.,
        ``[h_1, h_33]``.

    Returns
    -------
    prmse_values : list of float
        A list containing the values for the PRMSE metric
        for each of the given rater pairs.
    """
    # initialize a list that will hold the series
    prmse_for_all_pairs = []

    # iterate over each given rater pair
    for rater_id1, rater_id2 in rater_pairs:

        # call the per-pair function
        prmse_for_this_pair = prmse_true(df_scores[system_id],
                                         df_scores[[rater_id1, rater_id2]])
        # save the returned lists of series
        prmse_for_all_pairs.append(prmse_for_this_pair)

    return prmse_for_all_pairs

def compute_cumulative_mean_for_raters(df_scores, rater_ids):
    """
    Compute cumulative average of given rater's scores.

    This function computes the cumulative average of scores
    assigned by the raters identified by the given list of
    ``rater_ids``. The cumulative average is computed by
    adding one rater at a time.

    Parameters
    ----------
    df_scores : pandas.DataFrame
        The data frame containing the simulated scores.
        This is usually one of the data frames returned
        by the ``simulation.dataset.Dataset.to_frame()``
        method.
    rater_ids : numpy.ndarray
        An array of simulated rater IDs whose scores want
        to compute the cumulative average over.

    Returns
    -------
    df_cumulative_mean_scores : pandas.DataFrame
        A data frame containing the cumulative average rater scores.
        It has ``len(df_scores)`` rows and ``len(rater_ids`` columns.
        Each of the columns is named ``N=n``, where n is the number of
        raters included in computing the cumulative average.
    """
    df_rater_scores = df_scores[rater_ids]
    df_cumulative_average_scores = df_rater_scores.expanding(1, axis=1).mean()
    df_cumulative_average_scores.columns = [f"N={num_raters+1}" for num_raters in range(len(rater_ids))]
    return df_cumulative_average_scores

def compute_ranks_from_metrics(df_metrics):
    """
    Compute ranks given metric values for systems.

    This function computes ranks for a list of systems
    according to different metrics present in the given
    data frame.

    Parameters
    ----------
    df_metrics : pandas.DataFrame
        A data frame with one row for each system to be ranked and
        the following columns:
        1. "system_id" : the ID of the system
        2. "system_category" : the performance category that the system belongs to.
        3. At least one other column containing values for a metric with the name
       of the metric being the column name. For example, "QWK" or "PRMSE" etc.

    Returns
    -------
    df_metric_ranks : pandas.DataFrame
        A data frame with the same number of rows and columns as the input
        data frame except that the values in each "metric" column are now the
        ranks of the systems rather than the metric values themselves.
    """
    # first we set our indices to be the system IDs and categories so that they are
    # retained when we compute the ranks for the systems
    df_metrics_for_ranks = df_metrics.set_index(['system_category', 'system_id'])

    # if degradation is one of the metrics, multiply it by -1 to make it behave
    # like other metrics for ranking purposes
    if "degradation" in df_metrics_for_ranks.columns:
        df_metrics_for_ranks['degradation'] = -1 * df_metrics_for_ranks['degradation']

    # compute the ranks in descending order since lower ranks are better;
    # also reset the indices so that we get the IDs and categories back
    df_metric_ranks = df_metrics_for_ranks.rank(ascending=False).reset_index()

    # return the data frame
    return df_metric_ranks

def simulate_percent_double_scored(df_scores, rater_id1, rater_id2, percentage):
    """
    Simulate a dataset with only given percentage of double-scored responses.

    This function takes two given rater IDs and a percentage value and returns
    a data frame that contains the scores from only those two simulated raters
    with the additional constraint that only a certain percentage of the
    hypothetical responses in the dataset have a score from the second rater
    and the remaining responses have the scores from the second rater set
    to ``NaN``.

    Parameters
    ----------
    df_scores : pandas.DataFrame
        The data frame containing the simulated scores.
        This is usually one of the data frames returned
        by the ``simulation.dataset.Dataset.to_frame()``
        method.
    rater_id1 : str
        The ID for the first rater in the rater pair
        being used to evaluate the given system.
        This must be a column in ``df_scores``.
    rater_id2 : str
        The ID for the second rater in the rater pair
        being used to evaluate the given system.
        This must be a column in ``df_scores``.
    percentage : float
        The percentage of randomly chosen responses in the dataset
        that should still have scores from both raters.

    Returns
    -------
    df_rater_pair_scores : pandas.DataFrame
        A data frame containing the same number of rows as ``df_scores`` but
        only two columns : ``rater_id1`` and ``rater_id2``. Only ``percentage``
        percent of responses in the ``rater_id2`` column have valid scores
        with the rest set to ``Nan``.
    num_double_scored : int
        The number of double-scored responses in the returned data frame.
    """
    # initialize a random number generator
    prng = np.random.RandomState(987654321)

    # create a data frame that only has the scores from the raters in the given pair
    df_rater_pair_scores = df_scores[[rater_id1, rater_id2]].copy()

    # stochastically mask the second rater's scores unless we want 100% double-scored
    if percentage < 100:

        # calculate the number of responses we want to simulate as double scored
        num_double_scored = int(percentage / 100 * len(df_rater_pair_scores))

        # randomly choose this many number of responses
        chosen = prng.choice(df_rater_pair_scores.index.values,
                             size=num_double_scored,
                             replace=False)

        # for all the other responses in the dataset, set the second rater's score to null
        # that is, make them single scored
        df_rater_pair_scores.loc[~df_rater_pair_scores.index.isin(chosen), rater_id2] = np.nan

    else:

        # if the percentge is 100%, then the whole dataset is double-scored
        # and we don't need to do anything else
        num_double_scored = len(df_rater_pair_scores)

    return df_rater_pair_scores, num_double_scored

reason: import numpy,from scipy


text:
import os

import numpy as np
import oct2py

from pilco.gaussian_process.gaussian_process import GaussianProcess
from pilco.gaussian_process.multivariate_gp import MultivariateGP
from pilco.gaussian_process.sparse_multivariate_gp import SparseMultivariateGP

octave = oct2py.Oct2Py()
dir_path = "pilco/test/matlab_code"
octave.addpath(dir_path)

def test_mgpr():
    np.random.seed(1)

    state_dim = 3
    n_targets = 2

    n_samples = 100

    # Training Dataset
    X0 = np.random.rand(n_samples, state_dim)
    A = np.random.rand(state_dim, n_targets)
    Y0 = np.sin(X0).dot(A) + 1e-3 * (np.random.rand(n_samples, n_targets) - 0.5)
    length_scales = np.random.rand(n_targets, state_dim)
    sigma_f = np.ones(state_dim)
    sigma_eps = np.ones(state_dim)

    mgpr = MultivariateGP(X0, Y0, container=GaussianProcess, length_scales=length_scales, n_targets=n_targets,
                          sigma_f=sigma_f, sigma_eps=sigma_eps)

    mgpr.optimize()

    # Generate input
    mu = np.random.rand(1, state_dim)
    sigma = np.random.rand(state_dim, state_dim)
    sigma = sigma.dot(sigma.T)

    _ = mgpr.predict_from_dist(mu, sigma)

    # Change the dataset to avoid any caching issues for K and beta
    X0 = 5 * np.random.rand(100, state_dim)
    mgpr.fit(X0, Y0)

    M, S, V = mgpr.predict_from_dist(mu, sigma)

    # convert data to the struct expected by the MATLAB implementation
    length_scales = mgpr.length_scales()
    sigma_f = mgpr.sigma_f()
    sigma_eps = mgpr.sigma_eps()

    hyp = np.hstack(
        (length_scales,
         sigma_f,
         sigma_eps)
    ).T

    gpmodel = oct2py.io.Struct()
    gpmodel.hyp = hyp
    gpmodel.inputs = X0
    gpmodel.targets = Y0

    # Call function in octave
    M_mat, S_mat, V_mat = octave.gp0(gpmodel, mu.T, sigma, nout=3)
    M_mat = np.asarray(M_mat).flatten()
    S_mat = np.atleast_2d(S_mat)
    V_mat = np.asarray(V_mat)

    assert M.shape == M_mat.T.shape
    assert S.shape == S_mat.shape
    assert V.shape == V_mat.shape
    np.testing.assert_allclose(M, M_mat.T, rtol=1e-5)
    np.testing.assert_allclose(S, S_mat, rtol=1e-5)
    np.testing.assert_allclose(V, V_mat, rtol=1e-5)

def test_smgpr():
    np.random.seed(1)

    state_dim = 3
    n_targets = 1

    n_samples = 5
    n_inducing_points = 2

    # Training Dataset
    X0 = np.random.rand(n_samples, state_dim)
    A = np.random.rand(state_dim, n_targets)
    Y0 = np.sin(X0).dot(A) + 1e-3 * (np.random.rand(n_samples, n_targets) - 0.5)
    length_scales = np.random.rand(n_targets, state_dim)
    sigma_eps = np.ones(state_dim)
    sigma_f = np.ones(state_dim)

    smgpr = SparseMultivariateGP(length_scales=length_scales, n_targets=n_targets, x=X0, y=Y0,
                                 n_inducing_points=n_inducing_points, sigma_eps=sigma_eps, sigma_f=sigma_f)
    smgpr.optimize()

    # Generate input
    mu = np.random.rand(1, state_dim)
    sigma = np.random.rand(state_dim, state_dim)
    sigma = sigma.dot(sigma.T)

    _ = smgpr.predict_from_dist(mu, sigma)

    # Change the dataset to avoid any caching issues for K and beta
    X0 = 5 * np.random.rand(n_samples, state_dim)
    smgpr.fit(X0, Y0)

    M, S, V = smgpr.predict_from_dist(mu, sigma)

    # convert data to the struct expected by the MATLAB implementation
    length_scales = smgpr.length_scales()
    sigma_f = smgpr.sigma_f()
    sigma_eps = smgpr.sigma_eps()

    hyp = np.hstack(
        (length_scales,
         sigma_f,
         sigma_eps)
    ).T

    gpmodel = oct2py.io.Struct()
    gpmodel.hyp = hyp
    gpmodel.inputs = X0
    gpmodel.targets = Y0
    gpmodel.induce = np.stack([gp.Z.T for gp in smgpr.models]).T
    # gpmodel.induce = smgpr.gp_container[0].Z

    # Call function in octave
    M_mat, S_mat, V_mat = octave.gp1(gpmodel, mu.T, sigma, nout=3)
    M_mat = np.asarray(M_mat).flatten()
    S_mat = np.atleast_2d(S_mat)
    V_mat = np.asarray(V_mat)

    assert M.shape == M_mat.T.shape
    assert S.shape == S_mat.shape
    assert V.shape == V_mat.shape
    np.testing.assert_allclose(M, M_mat.T, rtol=1e-3)
    np.testing.assert_allclose(S, S_mat, rtol=1e-3)
    np.testing.assert_allclose(V, V_mat, rtol=1e-3)

if __name__ == '__main__':
    test_mgpr()
    test_smgpr()

max_stars_repo_name: BoboDance/RL-Project

max_issues_repo_name: ottofabian/RL-Project

max_forks_repo_name: ottofabian/RL-Project


text:
import pandas as pd
import numpy as np
import warnings
warnings.filterwarnings('ignore')
import os
os.environ['OMP_NUM_THREADS'] = '4'
import gc
import time
os.environ["CUDA_VISIBLE_DEVICES"] = '3'
import setproctitle
setproctitle.setproctitle('Kaggle@shihongzhi')

t0 = time.time()
path = '/data/mas/shihongzhi/Kaggle/'
dtypes = {
        'ip'            : 'uint32',
        'app'           : 'uint16',
        'device'        : 'uint16',
        'os'            : 'uint16',
        'channel'       : 'uint16',
        'is_attributed' : 'uint8',
        'click_id'      : 'uint32'
        }

print('load train....')
is_sample = 0
if(is_sample):
	train_file = "train_sample.csv"; test_file = "test_sample.csv"
else:
	train_file = "train.csv"; test_file = "test.csv"
train_df = pd.read_csv(path+train_file, dtype=dtypes, skiprows = range(1, 131886954), usecols=['ip','app','device','os', 'channel', 'click_time', 'is_attributed'], parse_dates=['click_time'])
print('load test....')
test_df = pd.read_csv(path+test_file, dtype=dtypes, usecols=['ip','app','device','os', 'channel', 'click_time', 'click_id'], parse_dates=['click_time'])
len_train = len(train_df)
train_df=train_df.append(test_df)
del test_df; gc.collect()

print('click time....')
train_df['click_time'] = (train_df['click_time'].astype(np.int64) // 10 ** 9).astype(np.int32)
train_df['next_click'] = (train_df.groupby(['ip', 'app', 'device', 'os']).click_time.shift(-1) - train_df.click_time).astype(np.float32)
train_df['next_click'].fillna((train_df['next_click'].mean()), inplace=True)

print('hour, day, wday....')
train_df['hour'] = pd.to_datetime(train_df.click_time).dt.hour.astype('uint8')
train_df['day'] = pd.to_datetime(train_df.click_time).dt.day.astype('uint8')
train_df['wday']  = pd.to_datetime(train_df.click_time).dt.dayofweek.astype('uint8')
print('grouping by ip-day-hour combination....')
gp = train_df[['ip','day','hour','channel']].groupby(by=['ip','day','hour'])[['channel']].count().reset_index().rename(index=str, columns={'channel': 'qty'})
train_df = train_df.merge(gp, on=['ip','day','hour'], how='left')
del gp; gc.collect()
print('group by ip-app combination....')
gp = train_df[['ip','app', 'channel']].groupby(by=['ip', 'app'])[['channel']].count().reset_index().rename(index=str, columns={'channel': 'ip_app_count'})
train_df = train_df.merge(gp, on=['ip','app'], how='left')
del gp; gc.collect()
print('group by ip-app-os combination....')
gp = train_df[['ip','app', 'os', 'channel']].groupby(by=['ip', 'app', 'os'])[['channel']].count().reset_index().rename(index=str, columns={'channel': 'ip_app_os_count'})
train_df = train_df.merge(gp, on=['ip','app', 'os'], how='left')
del gp; gc.collect()
print("vars and data type....")
train_df['qty'] = train_df['qty'].astype('uint16')
train_df['ip_app_count'] = train_df['ip_app_count'].astype('uint16')
train_df['ip_app_os_count'] = train_df['ip_app_os_count'].astype('uint16')
print("label encoding....")
from sklearn.preprocessing import LabelEncoder
train_df[['app','device','os', 'channel', 'hour', 'day', 'wday']].apply(LabelEncoder().fit_transform)
print ('final part of preparation....')
test_df = train_df[len_train:]
train_df = train_df[:len_train]
y_train = train_df['is_attributed'].values
train_df.drop(['click_id', 'click_time','ip','is_attributed'],1,inplace=True)
output = open('data.pkl', 'wb')

print ('neural network....')
from keras.layers import Input, Embedding, Dense, Flatten, Dropout, concatenate
from keras.layers import BatchNormalization, SpatialDropout1D, Conv1D
from keras.callbacks import Callback
from keras.models import Model
from keras.optimizers import Adam

max_app = np.max([train_df['app'].max(), test_df['app'].max()])+1
max_ch = np.max([train_df['channel'].max(), test_df['channel'].max()])+1
max_dev = np.max([train_df['device'].max(), test_df['device'].max()])+1
max_os = np.max([train_df['os'].max(), test_df['os'].max()])+1
max_h = np.max([train_df['hour'].max(), test_df['hour'].max()])+1
max_d = np.max([train_df['day'].max(), test_df['day'].max()])+1
max_wd = np.max([train_df['wday'].max(), test_df['wday'].max()])+1
max_qty = np.max([train_df['qty'].max(), test_df['qty'].max()])+1
max_c1 = np.max([train_df['ip_app_count'].max(), test_df['ip_app_count'].max()])+1
max_c2 = np.max([train_df['ip_app_os_count'].max(), test_df['ip_app_os_count'].max()])+1
def get_keras_data(dataset):
    X = {
        'app': np.array(dataset.app),
        'ch': np.array(dataset.channel),
        'dev': np.array(dataset.device),
        'os': np.array(dataset.os),
        'h': np.array(dataset.hour),
        'd': np.array(dataset.day),
        'wd': np.array(dataset.wday),
        'qty': np.array(dataset.qty),
        'c1': np.array(dataset.ip_app_count),
        'c2': np.array(dataset.ip_app_os_count),
        'nc': np.array(dataset.next_click)
    }
    return X
train_df = get_keras_data(train_df)

emb_n = 50
dense_n = 1000
in_app = Input(shape=[1], name = 'app')
emb_app = Embedding(max_app, emb_n)(in_app)
in_ch = Input(shape=[1], name = 'ch')
emb_ch = Embedding(max_ch, emb_n)(in_ch)
in_dev = Input(shape=[1], name = 'dev')
emb_dev = Embedding(max_dev, emb_n)(in_dev)
in_os = Input(shape=[1], name = 'os')
emb_os = Embedding(max_os, emb_n)(in_os)
in_h = Input(shape=[1], name = 'h')
emb_h = Embedding(max_h, emb_n)(in_h)
in_d = Input(shape=[1], name = 'd')
emb_d = Embedding(max_d, emb_n)(in_d)
in_wd = Input(shape=[1], name = 'wd')
emb_wd = Embedding(max_wd, emb_n)(in_wd)
in_qty = Input(shape=[1], name = 'qty')
emb_qty = Embedding(max_qty, emb_n)(in_qty)
in_c1 = Input(shape=[1], name = 'c1')
emb_c1 = Embedding(max_c1, emb_n)(in_c1)
in_c2 = Input(shape=[1], name = 'c2')
emb_c2 = Embedding(max_c2, emb_n)(in_c2)
in_nc = Input(shape=[1], name = 'nc')
fe = concatenate([(emb_app), (emb_ch), (emb_dev), (emb_os), (emb_h),
                 (emb_d), (emb_wd), (emb_qty), (emb_c1), (emb_c2)])
s_dout = SpatialDropout1D(0.2)(fe)
conv = Conv1D(100, kernel_size=4, strides=1, padding='same')(s_dout)
concat = concatenate([Flatten()(s_dout), Flatten()(conv), (in_nc)])
x = Dropout(0.2)(Dense(dense_n,activation='relu')(concat))
x = Dropout(0.2)(Dense(dense_n,activation='relu')(x))
outp = Dense(1,activation='sigmoid')(x)
model = Model(inputs=[in_app,in_ch,in_dev,in_os,in_h,in_d,in_wd,in_qty,in_c1,in_c2,in_nc], outputs=outp)

batch_size = 50000
epochs = 2
exp_decay = lambda init, fin, steps: (init/fin)**(1/(steps-1)) - 1
steps = int(len(list(train_df)[0]) / batch_size) * epochs
lr_init, lr_fin = 0.002, 0.0002
lr_decay = exp_decay(lr_init, lr_fin, steps)
optimizer_adam = Adam(lr=0.002, decay=lr_decay)
model.compile(loss='binary_crossentropy',optimizer=optimizer_adam,metrics=['accuracy'])

model.summary()

class_weight = {0:.01,1:.99} # magic
model.fit(train_df, y_train, batch_size=batch_size, epochs=2, class_weight=class_weight, shuffle=True, verbose=2)
del train_df, y_train; gc.collect()
model.save_weights('submission1.h5')

sub = pd.DataFrame()
sub['click_id'] = test_df['click_id'].astype('int')
test_df.drop(['click_id', 'click_time','ip','is_attributed'],1,inplace=True)
test_df = get_keras_data(test_df)

print("predicting....")
sub['is_attributed'] = model.predict(test_df, batch_size=batch_size, verbose=2)
del test_df; gc.collect()
print("writing....")
sub.to_csv('submission1.csv',index=False)
print('Time passed %fs'%(time.time() - t0))


text:
"""

"""
import math
from collections import defaultdict

import numpy as np
import numpy.random

from sentence_transformers import InputExample
from . import RoundRobinRankingDataset

class QueryExclusiveSimilarityDataset(RoundRobinRankingDataset):
    def __init__(self, model, queries, rel_queries, rel_corpus, batch_size=32, n_positives=2, temperature=1, shuffle=True, n_negatives=0, neg_rel_corpus=None, replace=True):
        super().__init__(model, queries, None, rel_queries, rel_corpus, None, batch_size, n_positives, temperature, shuffle, n_negatives, neg_rel_corpus)
        self.replace = replace
        self.counts = self.rel_queries.map(len)

    def __iter__(self):
        for _ in range(math.ceil(self.__len__() / self.batch_size)):
            current_batch_formation = defaultdict(int)
            available_d_ids = set(self.rel_queries.keys())
            current_batch_size = 0

            while current_batch_size < self.batch_size:
                # sampling n=batch_size covers worst case with 1 positive and no repetitions
                d_sample = self.rel_queries.sample(self.batch_size, weights=self.weights, replace=self.replace).keys()
                d_sample = d_sample[d_sample.isin(available_d_ids)]
                for d_id in d_sample:
                    if current_batch_size >= self.batch_size:
                        break

                    if d_id in available_d_ids:
                        if current_batch_formation[d_id] < self.counts[d_id]:
                            n_samples = min(self.n_positives, self.batch_size - current_batch_size, self.counts[d_id] - current_batch_formation[d_id])
                            current_batch_formation[d_id] += n_samples
                            current_batch_size += n_samples
                        else:
                            available_d_ids.remove(d_id)

            for d_id, sample_size in current_batch_formation.items():
                q_ids = [*np.random.choice(self.rel_queries[d_id], sample_size, replace=False)]
                for p_id in q_ids:
                    yield InputExample(texts=([self.queries[p_id]]), label=d_id)

max_stars_repo_name: t-mesq/sentence-transformers

max_issues_repo_name: t-mesq/sentence-transformers

max_forks_repo_name: t-mesq/sentence-transformers