sc_count

#!/bin/env python3
"""Create count files for 10X FASTQ files
"""

import argparse
import json
from subprocess import call


def main(args):
    "main method of script"
    with open(args.config) as json_data:
        config = json.load(json_data)
        if args.debug:
            print(config)

        for sample in config["samples"]:
            script_text = """#!/bin/bash
#SBATCH --partition={partition}
#SBATCH --account={account}
#SBATCH --time={time}
#SBATCH -J cellranger_count_{sample}
#SBATCH -o cellranger_count_{sample}_%j_o.txt
#SBATCH --cpus-per-task {cores}
#SBATCH --mem {mem}G
module load biocompute/biocompute-modules
module load cellranger/cellranger-5.0.1

cellranger count \
    --id {sample} \
    --transcriptome {transcriptome_path} \
    --fastqs {fastq_dir} \
    --sample {sample} \
    --jobmode local \
    --localcores $SLURM_CPUS_PER_TASK \
    --localmem {mem}
""".format(
    sample=sample,
    fastq_dir=config["fastq_dir"],
    transcriptome_path=config["transcriptome"],
    mem=args.mem,
    cores=args.cores,
    partition=args.partition,
    account=args.account,
    time=args.time,
)

            script_path = f"{sample}.sbatch"
            with open(script_path, 'w') as file_handle:
                file_handle.write(script_text)

            if args.debug:
                print(['sbatch', script_path])
            else:
                call(['sbatch', script_path])


        script_text = """#!/bin/bash
#SBATCH --partition={partition}
#SBATCH --account={account}
#SBATCH --time={time}
#SBATCH -J cellranger_aggr
#SBATCH -o cellranger_aggr_%j_o.txt
#SBATCH --cpus-per-task {cores}
#SBATCH --mem {mem}G
module load biocompute/biocompute-modules
module load cellranger/cellranger-5.0.1
""".format(
    mem=args.mem,
    cores=args.cores,
    partition=args.partition,
    account=args.account,
    time=args.time,
)
        script_text += f"""cellranger aggr \
    --id=aggregated \
    --csv=molecule_info.csv \
    --localcores=$SLURM_CPUS_PER_TASK \
    --localmem=40 --jobmode=local
"""

        for group in config["groups"]:
            script_text += f"""cellranger aggr \
    --id=aggregated_{group} \
    --csv=molecule_info_{group}.csv \
    --localcores=$SLURM_CPUS_PER_TASK \
    --localmem=40 \
    --jobmode=local
"""

            script_path = f"aggregate.sbatch"
            with open(script_path, 'w') as file_handle:
                file_handle.write(script_text)
            print(f"Run {script_path} when the individual jobs have finished")


if __name__ == '__main__':
    parser = argparse.ArgumentParser(
        description=__doc__
    )
    parser.add_argument(
        'config',
        type=str,
        help='JSON config file',
    )
    parser.add_argument(
        '--debug',
        dest='debug',
        action='store_true',
        default=False,
    )
    parser.add_argument(
        '--cores',
        type=str,
        default='24',
        help='Number of cores(threads) to use to process data',
    )
    parser.add_argument(
        '--mem',
        type=str,
        default='184',
        help='GB of memory to use to process data',
    )
    parser.add_argument(
        '--pre_mRNA',
        type=str,
        default=False,
        help='Whether this is a "pre_mRNA" count job',
    )
    parser.add_argument(
        '--account',
        type=str,
        default="warrenlab",
        help='SLURM account towards which time should be debited',
    )
    parser.add_argument(
        '--time',
        type=str,
        default="2-00:00:00",
        help='Time limit (for SLURM) in DD-HH:MM:SS format',
    )
    parser.add_argument(
        '--partition',
        type=str,
        default='hpc6',
        help='SLURM partition to use',
    )

    main(parser.parse_args())