<reponame>AMOLFResonantNanophotonics/CPA
'''
Function that simulates photon counting data with the following characteristics
- The data is distributed randomly over two detector channels
- The data is assumed generated by a pulsed laser excitation
- The emitter itself perfectly antibunches,  the two detectors have independent noise
- The data is specified is time tags - detector and reference arrival times
(int 64,  relative to a base unit that is a hypothetical counter card resolution)


The user gets to specify:
-  The number of levels nlevels. The simulated dot switches with no memory.
-  The instantaneous mean count rate and decay rate for levels n = 1...nlevels
-  On times for level n are taken from a powerlaw, exponents alpha_n specifiable
-  Background level
-  Total measurement duration / record length
-  Hypothetical set up info, i.e., laser rep. rate and timing card base unit
'''

import numpy as np
import matplotlib.pyplot as plt
import os

from matplotlib import gridspec as gsp
import pandas as pd
import pyarrow as pa
import pyarrow.parquet as pq

import preamble_for_simulation as pre


def samplepowerlaw(x,alpha,xmin,xmax):
    # %http://up-rs-esp.github.io/bpl/_modules/bpl.html#sample
    beta = 1.-alpha
    a = xmin**beta
    b = xmax**beta-a

    s = (a+b*x)**(1./beta)
    return s



def MakeSimulatedData():

    if not (len(pre.alpha_lst) == len(pre.I_lst) == len(pre.g_lst) == pre.nlevels):
        print('all input lists must have length '+str(pre.nlevels)+'!')

    # =============================================================================
    # draw from a multi-powerlaw distribution
    # output is a sequence of nominal jumptimes, and level assignments
    # =============================================================================

    print('Performing random draw of jumptimes and level assignments')

    n=1
    success = 0
    while success==0:

        # sample nlevels powerlaws
        N_powerlawelements = int(n*1E4) # the number of powerlaw elements to sample. This is the nr of jumps we simulate
        # Note that this will be cropped to fit in the measurement time

        which_level = (np.random.random(N_powerlawelements)*pre.nlevels).astype(int) # which intensity level are you in
        which_level = which_level[:-1][np.diff(which_level) != 0] # ensure you are not switching to an identical level

        # create the switching durations and timestamps
        time_durations_bin = np.zeros(N_powerlawelements)
        for i, j in enumerate(which_level):
            time_durations_bin[i] = samplepowerlaw(np.random.rand(1),pre.alpha_lst[j], pre.Tshortest_ns/pre.minibinsize_ns, pre.Tlongest_ns/pre.minibinsize_ns)[0]
        time_durations_bin = time_durations_bin.astype(int)

        # trim the jump times to the measurement time
        seg_times_mbin = np.cumsum(np.int64(time_durations_bin))
        seg_times_mbin = seg_times_mbin[seg_times_mbin<int(pre.NN)]

        # did we succesfully fill the entire measurement range with jumps
        if len(seg_times_mbin) < len(time_durations_bin):
            success = 1
        else:
            n = n+1
            print('Another round of sampling powerlaws')

    # adding the start and the end of the measurement time to the jump times, and giving it a more convenient name
    jumptimes = np.append(np.insert(seg_times_mbin, 0, 0), pre.NN-1)


    # =============================================================================
    # Create discrete photon events [does a give laser pulse [minibinsize]  result in a click
    # ....  and delay times, i.e., how long after the laser pulse does the click arise
    # =============================================================================

    trace_I = np.zeros(len(jumptimes)-1)
    trace_g = np.zeros(len(jumptimes)-1)


    # in principal we want to use the uniform probability distribution to draw zeroes and ones for empty, and full,
    # laser pulse cycles. However, this is memory wise quite prohibitive, as we are talking about maybe 1E9 or more bins, of which perhaps one promille is filled
    # hence we just try to allocate space for the photon events
    #
    # As we cycle over segments, the logical approach is to append photon events to an already existing list to build the time trace
    # However, append in Python is slow, as it reallocates the array
    # to avoid appending, we over allocate space for the photon counting array and insert values. We trim out overallocated space later
    maximumallocatedlength= int(max(pre.I_lst)*pre.NN*pre.minibinsize_ns*1E-9*2)  # this should be so overdimensioned that no allocation error could occur later on
    maximumallocatednoise= int(pre.I_noise*pre.NN*pre.minibinsize_ns*1E-9*10)


    timestamps_bin = np.zeros(maximumallocatedlength).astype(np.int64)
    delaylist_bin  = np.zeros(maximumallocatedlength).astype(np.int64)
    noiseclicks_chA = np.zeros(maximumallocatednoise).astype(np.int64)
    noiseclicks_chB = np.zeros(maximumallocatednoise).astype(np.int64)





    idxclick=0
    idxnoiseA=0
    idxnoiseB=0

    print('Performing random draw of photon detection events')

    for q in range(len(jumptimes)-1):
        pval = pre.I_lst[which_level[q]]*1E-9*pre.minibinsize_ns
        trace_I[q] = pre.I_lst[which_level[q]]  #nominal instantaneous intensity
        trace_g[q] = pre.g_lst[which_level[q]]  #nominal instantaneous decay rate


        # for given segment:
        #clicks due to the emitter with time granularity equal to the laser rep rate
        timestamps = ((jumptimes[q]+np.nonzero((np.random.rand(jumptimes[q+1]-jumptimes[q])<=pval).astype(int))[0]) * pre.minibinsize_ns / pre.dtau_ns).astype(np.int64)
        numclicks=len(timestamps)

        # arrival times relative to laser
        delaylist_ns = np.random.exponential(1/trace_g[q], size=numclicks) #random draws, expontl prob distrib



        #bookkeeping to insert segment clicks in bigger list
        if(numclicks>0):
            timestamps_bin[idxclick: idxclick+numclicks] = timestamps[:]
            delaylist_bin[ idxclick: idxclick+numclicks] =((delaylist_ns[:] / pre.dtau_ns).astype(np.int64))


        # same for noise for  both detectors (granularity -  per laser repeition rate here, fine time shifts added below)
        # doing this per segment avoids the memory over allocation problem.
        pval_noise = pre.I_noise*1E-9*pre.minibinsize_ns

        noiseA = ((jumptimes[q]+np.nonzero((np.random.rand(jumptimes[q+1]-jumptimes[q])<=pval_noise).astype(int))[0]) * pre.minibinsize_ns / pre.dtau_ns).astype(np.int64)
        noiseB = ((jumptimes[q]+np.nonzero((np.random.rand(jumptimes[q+1]-jumptimes[q])<=pval_noise).astype(int))[0]) * pre.minibinsize_ns / pre.dtau_ns).astype(np.int64)


        numnoiseA=len(noiseA)
        if(numnoiseA>0):
            noiseclicks_chA[idxnoiseA:idxnoiseA+numnoiseA]=noiseA[:]
        numnoiseB=len(noiseB)
        if(numnoiseB>0):
            noiseclicks_chB[idxnoiseB:idxnoiseB+numnoiseB]=noiseB[:]

        idxclick=idxclick+numclicks
        idxnoiseA=idxnoiseA+numnoiseA
        idxnoiseB=idxnoiseB+numnoiseB



    # =============================================================================
    # clip overdimensioned arrays
    # =============================================================================

    timestamps_bin = timestamps_bin[0:idxclick-1]
    delaylist_bin  =  delaylist_bin[0:idxclick-1]
    noiseclicks_chA=noiseclicks_chA[0:idxnoiseA]
    noiseclicks_chB=noiseclicks_chB[0:idxnoiseB]

    # =============================================================================
    # put the emitter event times and delays together. First in absence of background
    # =============================================================================

    timestamps_chR_bin = timestamps_bin  # spoofs laser reference channel - for every detector click, the laser pulse clock
    timestamps_bin = timestamps_bin - int(pre.minibinsize_ns/pre.dtau_ns) + delaylist_bin  # actual photon clicks, cares about fluorescence decay dynamics


    # =============================================================================
    # antibunch data into separate detectors as if for Hanbury-Brown Twiss
    # =============================================================================
    print('Antibunching the simulated emission events over detectors, subsequently including uncorrelated detector noise counts')

    whichdet = np.random.random(len(timestamps_bin))>0.5  # 50:50 routing ratio
    timestamps_chA_bin = timestamps_bin[whichdet]
    timestamps_chB_bin = timestamps_bin[np.logical_not(whichdet)]


    # =============================================================================
    # adding background counts a posteriori  so the detector noise doesn't antibunch
    # =============================================================================

    #Generate random time moments relative to the 100 ns trigger windows that we already determined
    delay_noise_chA_bin = (np.random.rand(len(noiseclicks_chA))*(pre.minibinsize_ns/pre.dtau_ns)).astype(np.int64)
    delay_noise_chB_bin = (np.random.rand(len(noiseclicks_chB))*(pre.minibinsize_ns/pre.dtau_ns)).astype(np.int64)

    #full finegrained random click arrival times, intersperse in data
    timestamps_chA_bin = np.sort(np.concatenate((timestamps_chA_bin,noiseclicks_chA-delay_noise_chA_bin)))
    timestamps_chB_bin = np.sort(np.concatenate((timestamps_chB_bin,noiseclicks_chB-delay_noise_chB_bin)))

    #interspersing the related reference trigger data for the noise count
    noisereferencepulses=np.sort(np.concatenate((noiseclicks_chA,noiseclicks_chB)))
    timestamps_chR_bin = np.sort(np.concatenate((timestamps_chR_bin,noisereferencepulses)))


    # =============================================================================
    # Save the data
    # =============================================================================

    # create the folder to save the data

    savepath = pre.simulated_parquetpath + pre.dot_foldername
    if not os.path.exists(savepath):
        os.makedirs(savepath)
    
    # save the photon events
    print('\nSaving generated photon and reference events in parquet files:')
    print('Folder:  '+savepath)
    for label in ['A', 'B', 'R']:
        datalabel = 'timestamps_ch'+label+'_bin'
        df = pd.DataFrame({'ch'+label: eval(datalabel)})
        table = pa.Table.from_pandas(df, preserve_index=False)
        print('File:  '+datalabel+'.parquet')
        pq.write_table(table, savepath+datalabel+'.parquet')

    # save the generated jump times
    print('\nFor reference, also saving a .csv file with')
    print('the nominal jump time,intensity and decay rate sequence')
    print('according to which the photon events were drawn')

    filepath_seg= pre.simulated_parquetpath + pre.dot_foldername + pre.seg_filename
    print('File:  '+filepath_seg)

    with open(filepath_seg,"w") as f:
        print("jumptime [bin, end]   intensity  decayrate [ns-1] ", file=f)
        np.savetxt(f, np.stack((jumptimes[1:], trace_I, trace_g),1), delimiter=', ')

'''
Function that simulates photon counting data with the following characteristics
- The data is distributed randomly over two detector channels
- The data is assumed generated by a pulsed laser excitation
- The emitter itself perfectly antibunches,  the two detectors have independent noise
- The data is specified is time tags - detector and reference arrival times
(int 64,  relative to a base unit that is a hypothetical counter card resolution)


The user gets to specify:
-  The number of levels nlevels. The simulated dot switches with no memory.
-  The instantaneous mean count rate and decay rate for levels n = 1...nlevels
-  On times for level n are taken from a powerlaw, exponents alpha_n specifiable
-  Background level
-  Total measurement duration / record length
-  Hypothetical set up info, i.e., laser rep. rate and timing card base unit
'''

import numpy as np
import matplotlib.pyplot as plt
import os

from matplotlib import gridspec as gsp
import pandas as pd
import pyarrow as pa
import pyarrow.parquet as pq

import preamble_for_simulation as pre


def samplepowerlaw(x,alpha,xmin,xmax):
    # %http://up-rs-esp.github.io/bpl/_modules/bpl.html#sample
    beta = 1.-alpha
    a = xmin**beta
    b = xmax**beta-a

    s = (a+b*x)**(1./beta)
    return s



def MakeSimulatedData():

    if not (len(pre.alpha_lst) == len(pre.I_lst) == len(pre.g_lst) == pre.nlevels):
        print('all input lists must have length '+str(pre.nlevels)+'!')

    # =============================================================================
    # draw from a multi-powerlaw distribution
    # output is a sequence of nominal jumptimes, and level assignments
    # =============================================================================

    print('Performing random draw of jumptimes and level assignments')

    n=1
    success = 0
    while success==0:

        # sample nlevels powerlaws
        N_powerlawelements = int(n*1E4) # the number of powerlaw elements to sample. This is the nr of jumps we simulate
        # Note that this will be cropped to fit in the measurement time

        which_level = (np.random.random(N_powerlawelements)*pre.nlevels).astype(int) # which intensity level are you in
        which_level = which_level[:-1][np.diff(which_level) != 0] # ensure you are not switching to an identical level

        # create the switching durations and timestamps
        time_durations_bin = np.zeros(N_powerlawelements)
        for i, j in enumerate(which_level):
            time_durations_bin[i] = samplepowerlaw(np.random.rand(1),pre.alpha_lst[j], pre.Tshortest_ns/pre.minibinsize_ns, pre.Tlongest_ns/pre.minibinsize_ns)[0]
        time_durations_bin = time_durations_bin.astype(int)

        # trim the jump times to the measurement time
        seg_times_mbin = np.cumsum(np.int64(time_durations_bin))
        seg_times_mbin = seg_times_mbin[seg_times_mbin<int(pre.NN)]

        # did we succesfully fill the entire measurement range with jumps
        if len(seg_times_mbin) < len(time_durations_bin):
            success = 1
        else:
            n = n+1
            print('Another round of sampling powerlaws')

    # adding the start and the end of the measurement time to the jump times, and giving it a more convenient name
    jumptimes = np.append(np.insert(seg_times_mbin, 0, 0), pre.NN-1)


    # =============================================================================
    # Create discrete photon events [does a give laser pulse [minibinsize]  result in a click
    # ....  and delay times, i.e., how long after the laser pulse does the click arise
    # =============================================================================

    trace_I = np.zeros(len(jumptimes)-1)
    trace_g = np.zeros(len(jumptimes)-1)


    # in principal we want to use the uniform probability distribution to draw zeroes and ones for empty, and full,
    # laser pulse cycles. However, this is memory wise quite prohibitive, as we are talking about maybe 1E9 or more bins, of which perhaps one promille is filled
    # hence we just try to allocate space for the photon events
    #
    # As we cycle over segments, the logical approach is to append photon events to an already existing list to build the time trace
    # However, append in Python is slow, as it reallocates the array
    # to avoid appending, we over allocate space for the photon counting array and insert values. We trim out overallocated space later
    maximumallocatedlength= int(max(pre.I_lst)*pre.NN*pre.minibinsize_ns*1E-9*2)  # this should be so overdimensioned that no allocation error could occur later on
    maximumallocatednoise= int(pre.I_noise*pre.NN*pre.minibinsize_ns*1E-9*10)


    timestamps_bin = np.zeros(maximumallocatedlength).astype(np.int64)
    delaylist_bin  = np.zeros(maximumallocatedlength).astype(np.int64)
    noiseclicks_chA = np.zeros(maximumallocatednoise).astype(np.int64)
    noiseclicks_chB = np.zeros(maximumallocatednoise).astype(np.int64)





    idxclick=0
    idxnoiseA=0
    idxnoiseB=0

    print('Performing random draw of photon detection events')

    for q in range(len(jumptimes)-1):
        pval = pre.I_lst[which_level[q]]*1E-9*pre.minibinsize_ns
        trace_I[q] = pre.I_lst[which_level[q]]  #nominal instantaneous intensity
        trace_g[q] = pre.g_lst[which_level[q]]  #nominal instantaneous decay rate


        # for given segment:
        #clicks due to the emitter with time granularity equal to the laser rep rate
        timestamps = ((jumptimes[q]+np.nonzero((np.random.rand(jumptimes[q+1]-jumptimes[q])<=pval).astype(int))[0]) * pre.minibinsize_ns / pre.dtau_ns).astype(np.int64)
        numclicks=len(timestamps)

        # arrival times relative to laser
        delaylist_ns = np.random.exponential(1/trace_g[q], size=numclicks) #random draws, expontl prob distrib



        #bookkeeping to insert segment clicks in bigger list
        if(numclicks>0):
            timestamps_bin[idxclick: idxclick+numclicks] = timestamps[:]
            delaylist_bin[ idxclick: idxclick+numclicks] =((delaylist_ns[:] / pre.dtau_ns).astype(np.int64))


        # same for noise for  both detectors (granularity -  per laser repeition rate here, fine time shifts added below)
        # doing this per segment avoids the memory over allocation problem.
        pval_noise = pre.I_noise*1E-9*pre.minibinsize_ns

        noiseA = ((jumptimes[q]+np.nonzero((np.random.rand(jumptimes[q+1]-jumptimes[q])<=pval_noise).astype(int))[0]) * pre.minibinsize_ns / pre.dtau_ns).astype(np.int64)
        noiseB = ((jumptimes[q]+np.nonzero((np.random.rand(jumptimes[q+1]-jumptimes[q])<=pval_noise).astype(int))[0]) * pre.minibinsize_ns / pre.dtau_ns).astype(np.int64)


        numnoiseA=len(noiseA)
        if(numnoiseA>0):
            noiseclicks_chA[idxnoiseA:idxnoiseA+numnoiseA]=noiseA[:]
        numnoiseB=len(noiseB)
        if(numnoiseB>0):
            noiseclicks_chB[idxnoiseB:idxnoiseB+numnoiseB]=noiseB[:]

        idxclick=idxclick+numclicks
        idxnoiseA=idxnoiseA+numnoiseA
        idxnoiseB=idxnoiseB+numnoiseB



    # =============================================================================
    # clip overdimensioned arrays
    # =============================================================================

    timestamps_bin = timestamps_bin[0:idxclick-1]
    delaylist_bin  =  delaylist_bin[0:idxclick-1]
    noiseclicks_chA=noiseclicks_chA[0:idxnoiseA]
    noiseclicks_chB=noiseclicks_chB[0:idxnoiseB]

    # =============================================================================
    # put the emitter event times and delays together. First in absence of background
    # =============================================================================

    timestamps_chR_bin = timestamps_bin  # spoofs laser reference channel - for every detector click, the laser pulse clock
    timestamps_bin = timestamps_bin - int(pre.minibinsize_ns/pre.dtau_ns) + delaylist_bin  # actual photon clicks, cares about fluorescence decay dynamics


    # =============================================================================
    # antibunch data into separate detectors as if for Hanbury-Brown Twiss
    # =============================================================================
    print('Antibunching the simulated emission events over detectors, subsequently including uncorrelated detector noise counts')

    whichdet = np.random.random(len(timestamps_bin))>0.5  # 50:50 routing ratio
    timestamps_chA_bin = timestamps_bin[whichdet]
    timestamps_chB_bin = timestamps_bin[np.logical_not(whichdet)]


    # =============================================================================
    # adding background counts a posteriori  so the detector noise doesn't antibunch
    # =============================================================================

    #Generate random time moments relative to the 100 ns trigger windows that we already determined
    delay_noise_chA_bin = (np.random.rand(len(noiseclicks_chA))*(pre.minibinsize_ns/pre.dtau_ns)).astype(np.int64)
    delay_noise_chB_bin = (np.random.rand(len(noiseclicks_chB))*(pre.minibinsize_ns/pre.dtau_ns)).astype(np.int64)

    #full finegrained random click arrival times, intersperse in data
    timestamps_chA_bin = np.sort(np.concatenate((timestamps_chA_bin,noiseclicks_chA-delay_noise_chA_bin)))
    timestamps_chB_bin = np.sort(np.concatenate((timestamps_chB_bin,noiseclicks_chB-delay_noise_chB_bin)))

    #interspersing the related reference trigger data for the noise count
    noisereferencepulses=np.sort(np.concatenate((noiseclicks_chA,noiseclicks_chB)))
    timestamps_chR_bin = np.sort(np.concatenate((timestamps_chR_bin,noisereferencepulses)))


    # =============================================================================
    # Save the data
    # =============================================================================

    # create the folder to save the data

    savepath = pre.simulated_parquetpath + pre.dot_foldername
    if not os.path.exists(savepath):
        os.makedirs(savepath)
    
    # save the photon events
    print('\nSaving generated photon and reference events in parquet files:')
    print('Folder:  '+savepath)
    for label in ['A', 'B', 'R']:
        datalabel = 'timestamps_ch'+label+'_bin'
        df = pd.DataFrame({'ch'+label: eval(datalabel)})
        table = pa.Table.from_pandas(df, preserve_index=False)
        print('File:  '+datalabel+'.parquet')
        pq.write_table(table, savepath+datalabel+'.parquet')

    # save the generated jump times
    print('\nFor reference, also saving a .csv file with')
    print('the nominal jump time,intensity and decay rate sequence')
    print('according to which the photon events were drawn')

    filepath_seg= pre.simulated_parquetpath + pre.dot_foldername + pre.seg_filename
    print('File:  '+filepath_seg)

    with open(filepath_seg,"w") as f:
        print("jumptime [bin, end]   intensity  decayrate [ns-1] ", file=f)
        np.savetxt(f, np.stack((jumptimes[1:], trace_I, trace_g),1), delimiter=', ')
Function that simulates photon counting data with the following characteristics - The data is distributed randomly over two detector channels - The data is assumed generated by a pulsed laser excitation - The emitter itself perfectly antibunches,  the two detectors have independent noise - The data is specified is time tags - detector and reference arrival times (int 64,  relative to a base unit that is a hypothetical counter card resolution)   The user gets to specify: -  The number of levels nlevels. The simulated dot switches with no memory. -  The instantaneous mean count rate and decay rate for levels n = 1...nlevels -  On times for level n are taken from a powerlaw, exponents alpha_n specifiable -  Background level -  Total measurement duration / record length -  Hypothetical set up info, i.e., laser rep. rate and timing card base unit # %http://up-rs-esp.github.io/bpl/_modules/bpl.html#sample # ============================================================================= # draw from a multi-powerlaw distribution # output is a sequence of nominal jumptimes, and level assignments # ============================================================================= # sample nlevels powerlaws # the number of powerlaw elements to sample. This is the nr of jumps we simulate # Note that this will be cropped to fit in the measurement time # which intensity level are you in # ensure you are not switching to an identical level # create the switching durations and timestamps # trim the jump times to the measurement time # did we succesfully fill the entire measurement range with jumps # adding the start and the end of the measurement time to the jump times, and giving it a more convenient name # ============================================================================= # Create discrete photon events [does a give laser pulse [minibinsize]  result in a click # ....  and delay times, i.e., how long after the laser pulse does the click arise # ============================================================================= # in principal we want to use the uniform probability distribution to draw zeroes and ones for empty, and full, # laser pulse cycles. However, this is memory wise quite prohibitive, as we are talking about maybe 1E9 or more bins, of which perhaps one promille is filled # hence we just try to allocate space for the photon events # # As we cycle over segments, the logical approach is to append photon events to an already existing list to build the time trace # However, append in Python is slow, as it reallocates the array # to avoid appending, we over allocate space for the photon counting array and insert values. We trim out overallocated space later # this should be so overdimensioned that no allocation error could occur later on #nominal instantaneous intensity #nominal instantaneous decay rate # for given segment: #clicks due to the emitter with time granularity equal to the laser rep rate # arrival times relative to laser #random draws, expontl prob distrib #bookkeeping to insert segment clicks in bigger list # same for noise for  both detectors (granularity -  per laser repeition rate here, fine time shifts added below) # doing this per segment avoids the memory over allocation problem. # ============================================================================= # clip overdimensioned arrays # ============================================================================= # ============================================================================= # put the emitter event times and delays together. First in absence of background # ============================================================================= # spoofs laser reference channel - for every detector click, the laser pulse clock # actual photon clicks, cares about fluorescence decay dynamics # ============================================================================= # antibunch data into separate detectors as if for Hanbury-Brown Twiss # ============================================================================= # 50:50 routing ratio # ============================================================================= # adding background counts a posteriori  so the detector noise doesn't antibunch # ============================================================================= #Generate random time moments relative to the 100 ns trigger windows that we already determined #full finegrained random click arrival times, intersperse in data #interspersing the related reference trigger data for the noise count # ============================================================================= # Save the data # ============================================================================= # create the folder to save the data # save the photon events # save the generated jump times