simulate.nf

#!/usr/bin/env nextflow
include { SetupSimulateVisibilities;
          SetupGaussianNoise;
          MakeMachinesFile;
          CopySkyModel;
} from '/home/users/cahofer/code/simple/foo.nf'

include { ApplyPreDIFlag 
} from '/home/users/cahofer/theleap/theleap/leap/main.nf'

params.help = null
params.nodes = "126,127,128,129,130"
//params.nodes = "116,117,118,119,120,121,122,123,124,125,126,127,128,129"
nodes_list = parseNodes()

/*
// Simulate visibilities input parameters with DP3
params.joblist = "my_jobs.list"
params.dp3_parset = "/net/node126/data/users/lofareor/cahofer/simulations/NCP/NCP_fluxcut/chunk5/config_predict.parset"
params.skymodel = "/net/node105/data/users/lofareor/cahofer/simulations/NCP/skymodel/center.txt"

EXAMPLES FOR WORKFLOWS:
-----------------------


=== WORKFLOW TO AVERAGE DATA COLUMN IN TIME AND FREQ ===

workflow {
    filelist = MSList(params.data.msprefix, params.data.path, nodes_list[0], nodes_list[-1])
    ch3 = AverageDataColumn(true, nodes_list[0], nodes_list[-1], "jobs_average_data.list", params.data.path, params.data.msprefix, "DATA", 7191, false, 5, 3, "DATA",  filelist)
}

=== WORKFLOW TO ADD NOISE TO DATA COLUMN ===

workflow {
    filelist = MSList(params.data.msprefix, params.data.path, nodes_list[0], nodes_list[-1])
    noise_parset = SetupGaussianNoise(params.data.path, params.data.msprefix, params.skymodel, 'DATA', 'config_noise.parset')
    ch2 = GaussianNoise(true, nodes_list[0], nodes_list[-1], params.venv, params.joblist_noise, noise_parset, params.data.path, params.outfile)
  
=== WORKFLOW TO PREDICT DIFFUSE EMISSION WITH WSCLEAN FROM IMAGE ===

workflow {
    filelist = MSList(params.data.msprefix, params.data.path, nodes_list[0], nodes_list[-1])
    ch1 = PredictWSClean(true, 126, 130, "wsclean_predict.list", params.data.path, params.predict.nr, filelist)
    ch2 = CreateDataColumn(ch1, 126, 130, "jobs_create_col1.list", params.data.path, params.data.msprefix, params.io.incol, params.io.outcol, filelist)
    ch3 = CreateDataColumn2(ch2, 126, 130, "jobs_create_col2.list", params.data.path, params.data.msprefix, params.io.incol2, params.io.outcol2, filelist)
    ch4 = ApplyPreDIFlag(ch3, params.flag.pssh_hosts_txt_file, params.flag.preprocessing_file, params.data.path, "${params.data.path}/ms_files_002.txt", params.flag.sim)
}

=== WORKFLOW TO APPLY GAINS TO DATA WITH SAGECAL ===

workflow {
    filelist = MSList(params.data.msprefix, params.data.path, nodes_list[0], nodes_list[-1])
    ch1  = ApplyGainsSagecal(nodes_list[0], nodes_list[-1], "jobs_sagecal_apply_gains.list", params.data.path, params.predict.skymodel_file, params.predict.clustering_file, params.z_file, 6, 2, params.soldir, "CORRECTED_DATA2", params.data.msprefix, filelist)
}
*/

process MSList {
    publishDir params.data.outdir, mode: 'copy'

    input:
    val msprefix
    val path
    val ns
    val ne

    output:
    path "file_list_${msprefix}"

    script:
    """
    ns=nodes_list[0]
    ne=nodes_list[-1]
    for ni in {$ns..$ne}; do
        ls -d /net/node\${ni}/${path}/${msprefix}*.MS >> file_list_${msprefix}
    done
    """
}

process MSListSolutions {
    publishDir params.data.outdir, mode: 'copy'

    input:
    val msprefix
    val path
    val ns
    val ne

    output:
    path "file_list_${msprefix}"

    script:
    """
    for ni in {$ns..$ne}; do
        ls -d /net/node\${ni}/${path}/${msprefix}*.MS.solutions >> file_list_${msprefix}
    done
    """
}


/*
A task to simulate visibilites.
A bash script that produces a DP3 job list on the specified nodes using nodetool from https://gitlab.com/flomertens/libpipe.

Parameters
----------
ns : int
    Start node in list of nodes to use
ne : int
    End node in list of nodes to use
joblist : string
    Filname of joblist file
path : string
    The path/to/file where measurements sets and parset file for prediction are located
msprefix : string
    The name of the measurement set file without the subband number

Returns
-------
joblist : path
    The path to the joblist file

*/

process SimulateVisibilitiesDP3 {
    publishDir params.data.outdir

    input:
    val ns
    val ne
    val joblist
    val path
    val dp3_parset
    path mslist

    output:
    path joblist

    shell:
    '''
    input=!{mslist}

    while read -r line
    do
        nodetool add_job !{joblist} "DP3 !{dp3_parset} msin=$line"
    
    done < "$input"

    nodetool run_joblist -n node[!{ns}-!{ne}] !{joblist}
    '''
}

/*
A task to simulate visibilites with SAGECAL.
A bash script that produces a sagecal_gpu job list on the specified nodes using nodetool from https://gitlab.com/flomertens/libpipe.

Parameters
----------
ns : int
    Start node in list of nodes to use
ne : int
    End node in list of nodes to use
joblist : string
    Filname of joblist file
path : string
    The path/to/file where measurements sets and parset file for prediction are located
skymodel_file : string
    The relative /path/to/file where skymodel file is located (i.e. without /net/nodeXXX/).
    Requires: CopySkyModel to be completed, such that skymodel_file is on every node that sagecal will run.
clustering_file : string
    The relative /path/to/file where clustering file is located (i.e. without /net/nodeXXX/).
    Requires: CopySkyModel to be completed, such that clustering_file is on every node that sagecal will run.
np : int
    Number of worker threads (CPUs)
beammode : int
    According to sagecal documentation, 1: array beam, 2: array+element beam, 3: element beam, default 0
outcol : string
    The name of the output data column to create.
msprefix : string
    The name of the measurement set file without the subband number
mslist : path
    Path to text file containing the list of measurement set files.
    Requires: MSList to be completed such that text file exsists.

Returns
-------
true : bool
    Task is completed.
*/

process SimulateVisibilitiesSagecal {
    debug true
    publishDir params.data.outdir

    input:
    val ready
    val ns
    val ne
    val joblist
    val path
    val skymodel_file
    val clustering_file
    val np
    val beammode 
    val outcol
    val msprefix
    path mslist

    output:
    val true

    shell:
    '''
    input=!{mslist}

    while read -r line
    do
        nodetool add_job !{joblist} "sagecal_gpu -d $line -s !{skymodel_file} -c !{clustering_file} -n !{np} -B !{beammode} -E 100 -a 1 -t 1 -b 1 -I DATA -O !{outcol}"
    
    done < "$input"

    nodetool run_joblist -n node[!{ns}-!{ne}] !{joblist}
    '''
}

/* 
Apply gains with SAGECAL. This is not done yet. Need to tell sagecal the matching solution file for each file. 
*/

process ApplyGainsSagecal {
    debug true
    publishDir params.data.outdir

    input:
    val ns
    val ne
    val joblist
    val path
    val skymodel_file
    val clustering_file
    val z_file
    val np
    val beammode
    val soldir
    val outcol
    val msprefix
    path mslist

    output:
    val true

    shell:
    '''
    input=!{mslist}

    while read -r line
    do
        nodetool add_job !{joblist} "sagecal_gpu -d ${line:13} -s !{skymodel_file} -c !{clustering_file} -B !{beammode} -p !{soldir}/!{msprefix}${line:124:3}_uv_002_simulation_BP.MS.solutions -E 1 -a 3 -z !{z_file} -k -1 -t 91 -b 1 -I MODEL_DATA -O !{outcol} > !{path}/logfile_${line:122:5}_applygain.output-init" -t ${line:5:7}
    
    done < "$input"

    nodetool run_joblist -n node[!{ns}-!{ne}] !{joblist}
    '''
}

/*
Takes the output of SetupGaussianNoise, which creates a noise parset file according to the measurement sets that are processed.
A task to simulate thermal noise in the visibilities.
A bash script that uses the python package losito from https://github.com/darafferty/losito/

Parameters
----------
ready : bool
    Previous task is completed.
ns : int
    Start node in list of nodes to use
ne : int
    End node in list of nodes to use
venv : string
    The path/to/file where the virtual environment of losito is located.
joblist : string
    Filname of joblist file
noise_parset : string
    The path/to/file of a default noise parset file.
outfile : string
    Name of logfile of this task. 
Returns
-------
true : bool
    Task is completed.
    
    nodetool add_job !{joblist} "source !{venv}; losito /net/node!{ne}/!{path}/!{noise_parset} > !{path}/!{outfile}" -t node$ni
*/

//losito /path/to/noise_parset noise_parset
// or try path! 

process GaussianNoise {
    debug true
    publishDir params.data.outdir

    input:
    val ready
    val ns
    val ne
    val joblist
    val noise_parset
    val path
    val outfile

    output:
    val true

    shell:
    
    '''
    for ni in {!{ns}..!{ne}}; do
        nodetool add_job !{joblist} "source activate lofar_pipeline; losito /net/node129/!{noise_parset} > !{path}/!{outfile}" -t node$ni
    done

    nodetool run_joblist -n node[!{ns}-!{ne}] !{joblist}
    '''
}

/*
A task to create an additional data column in the measurement set in parallel using nodetool. 

Parameters
----------
ready : bool
    Ready to process.
ns : int
    Start node in list of nodes to use
ne : int
    End node in list of nodes to use
joblist : string
    Filname of joblist file
path : string
    The path/to/file where measurements sets and parset file for prediction are located
msprefix : string
    The name of the measurement set file without the subband number
incol : string
    The name of the input data column in the measrument set to copy from.
outcol : string
    The name of the output data column to create.
mslist : path
    Path to text file containing the list of measurement set files.
  
Returns
-------
bool : True
    The process is done.

*/

process CreateDataColumn {
    debug true 
    publishDir params.data.outdir

    input:
    val ready
    val ns
    val ne
    val joblist
    val path
    val msprefix
    val incol
    val outcol
    val mslist

    output:
    val true

    shell:
    '''
    input=!{mslist}

    while read -r line
    do
        nodetool add_job !{joblist} "DP3 msin=$line msin.datacolumn=!{incol} msout=. msout.datacolumn=!{outcol} steps=[]"
    
    done < "$input"

    nodetool run_joblist -n node[!{ns}-!{ne}] !{joblist}
    '''
}

process CreateDataColumn2 {
    debug true 
    publishDir params.data.outdir

    input:
    val ready
    val ns
    val ne
    val joblist
    val path
    val msprefix
    val incol
    val outcol
    path mslist

    output:
    val true

    shell:
    '''
    input=!{mslist}

    while read -r line
    do
        nodetool add_job !{joblist} "DP3 msin=$line msin.datacolumn=!{incol} msout=. msout.datacolumn=!{outcol} steps=[]"
    
    done < "$input"

    nodetool run_joblist -n node[!{ns}-!{ne}] !{joblist}
    '''
}

/*
A task to average data in time and frequency for a selected number of time samples in a measurement set in parallel using nodetool. 

Parameters
----------
ready : bool
    Ready to process.
ns : int
    Start node in list of nodes to use
ne : int
    End node in list of nodes to use
joblist : string
    Filname of joblist file , input paramters for nodetool
path : string
    The path/to/file where measurements are located
msprefix : string
    The name of the measurement set file without the subband number
incol : string
    The name of the input data column in the measurument set to copy from.
outcol : string
    The name of the output data column to create.
outname : string
    The name of the averaged measurement set file.
mslist : path
    Path to text file containing the list of measurement set files.
  
Returns
-------
bool : True
    The process is done.

*/

process AverageDataColumn {
    debug true 
    publishDir params.data.outdir

    input:
    val ready
    val ns
    val ne
    val joblist
    val path
    val msprefix
    val incol
    val ntimes
    val flagbool
    val timestep
    val freqstep
    val outcol
    path mslist

    output:
    val true

    shell:
    '''
    input=!{mslist}

    while read -r line
    do
        msoutname=$(echo "${line##*/}" | sed "s/R/simulation/")
        nodetool add_job !{joblist} "DP3 msin=$line msin.datacolumn=!{incol} msin.ntimes=!{ntimes} msin.useflag=!{flagbool} steps=[average] average.timestep=!{timestep} average.freqstep=!{freqstep} msout=!{path}/$msoutname msout.datacolumn=!{outcol}" -t ${line:5:7}
    
    done < "$input"

    nodetool run_joblist -n node[!{ns}-!{ne}] !{joblist}
    '''
}



/*
A task to simulate visibilities with wgridder of WSCLEAN in parallel using nodetool.

Parameters
----------
ready : bool
    Ready to process.
*/


process PredictWSClean {
    debug true 
    publishDir params.data.outdir

    input:
    val ready
    val ns
    val ne
    val joblist
    val path
    val nr
    val mslist

    output:
    val true

    shell:
    '''
    input=!{mslist}

    while read -r line
    do
        msoutname=$(echo "${line##*/}" | sed "s/_uv_002_simulation.MS/-!{nr}/")
        echo $msoutname
        nodetool add_job !{joblist} "wsclean -j 64 -data-column MODEL_DATA -predict -gridder idg -grid-with-beam -no-small-inversion -pol I -channels-out 1 -name !{path}/$msoutname $line" -t ${line:5:7}

    done < "$input"

    nodetool run_joblist -n node[!{ns}-!{ne}] !{joblist}
    '''
}
// Command for WGRIDDER - but WGRIDDER doesn't use the primary beam.
//        nodetool add_job !{joblist} "wsclean -j 64 -data-column MODEL_DATA -predict -gridder wgridder -wgridder-accuracy 1e-6 -no-small-inversion -pol I -channels-out 1 -name !{path}/$msoutname $line" -t ${line:5:7}
// -n node[!{ns}-!{ne}]


def write_nodes_for_pssh(nodes_list, pssh_hosts_txt_file) {
    pssh_hosts_txt_file = new File(pssh_hosts_txt_file)
    if (pssh_hosts_txt_file.exists()){
        pssh_hosts_txt_file.delete()
    }
    pssh_hosts_txt_file.createNewFile()
    pssh_hosts_txt_file.withWriter { out ->
    nodes_list.each {
      out.println("node${it}") // node129
    }
  }
}

//return a list of the nodes with the 'node' prefix given an input string e.g. [node129]
def parseNodes(){
    if (params.nodes instanceof String){
        nodes_list = params.nodes.split(',').collect{"${it}"} as List
    }
    //when a single node is given..
    else if (params.nodes instanceof Integer) {
        nodes_list = params.nodes.collect{"${it}"} as List
    }
    //Define 2 more required params
    // params.masternode = nodes_list[-1] //The last node in the list is used as the masternode
    // params.number_of_processes = params.slots * nodes_list.size

    return nodes_list
}