text:
#!/usr/bin/env python3.7

# Main script for sequencing analysis
# Input of this script:
#        - JSON file containing input parameters
#    - Fastq files (gzipped)

# what does this script do?
# 1. Read input csv file and convert to json (using csv2Json.R (you need to
# install required packages in R)
# 1. Read paramters from json file
# 2. Read input fastq file, align reads to reference sequence, generate sam files. (bowtie2 required)
# 4. From sam files, count mutations
# 5. Output mutation counts to summary.csv

# modules
import re
import subprocess

import pandas as pd
import numpy as np
import os
import glob
import argparse
import logging
import datetime
import shutil

# pakage modules
from TileSeqMut import help_functions
from TileSeqMut import count_mut
from TileSeqMut import cluster

class fastq2counts(object):
    """Main class to read fastq files and output mutation counts."""

    def __init__(self, param_path, output_dir, main_log, args):
        """
        Initialize mutation counts
        param_path: Full path to param.json file
        output_dir: Directory to save all the output files
        main_log: logging object
        args: user inpu arguments
        """
        # path to this python file
        self._main_path = os.path.abspath(__file__)
        # logger
        self._logging = main_log
        # parameter json file
        self._param_json = param_path
        # user input arguments
        self._args = args
        # convert environment input into upper case
        self._args.environment = self._args.environment.upper()

        # if input arguments provided fastq path, get all fastq files
        if self._args.fastq:
            self._fastq_list = glob.glob(self._args.fastq+"/*.fastq.gz")
        else:
            self._fastq_list = []
        # output dir for this run
        self._output = output_dir

        # parse parameter json file
        self._project, self._seq, self._cds_seq, self._tile_map, self._region_map, \
            self._samples, self._var, self._relations = help_functions.parse_json(param_path)
        self._sample_names = self._samples["Sample ID"].tolist()
        self._project = self._project.replace(" ", "_")  # project name
        # make main log
        self._log = self._logging.getLogger("main.log")
        # self._log.debug(self._args.environment)

    def _init_skip(self, skip, r1=None, r2=None):
        """
        if skip == T: skip alignment
            make suboutput dir with time stamp
        else: not skip alignment
            make main output dir with time stamp
        """
        # initialize read 1/read 2
        self._r1 = ""
        self._r2 = ""
        self._skip = False
        if skip:  # only run mutation counts on the aligned samples
            self._skip = True
            if self._args.r1 and self._args.r2:
                self._r1 = self._args.r1
                self._r2 = self._args.r2
            else:
                self._log.info(f"Sam files are read from {self._output}")

    def _align_sh_(self):
        """
        1. Align each fastq file to reference (submit one job for each alignment)
        Output sam files into output folders (sam_files)
        """
        # VALIDATE: if all samples are present in fastq files
        # also check if any fastq file is empty
        fastq_sample_id = [os.path.basename(i).split("_")[0] for i in self._fastq_list]
        fastq_sample_id = list(set(fastq_sample_id))
        self._log.debug(f"Fastq sample ID in input folder: {fastq_sample_id}")
        # validation
        # check if input fastq files contains all the samples in paramter file
        self._log.debug(f"Sample names in parameter sheet: {set(self._sample_names)}")
        if set(self._sample_names).issubset(fastq_sample_id):
            self._log.info("All samples found")
            self._log.info(f"In total there are {len(list(set(self._sample_names)))} samples in the csv file")
            self._log.info(f"In total there are {len(fastq_sample_id)} fastq files")
        else:
            test= list(np.setdiff1d(self._sample_names,fastq_sample_id))
            join_list = ",".join(test)
            self._log.error("fastq files do not match input samples.")
            self._log.error(f"Fastq files not found for sample {join_list}")
            raise FileNotFoundError()

        # create mapping for R1 and R2 for each sample
        # ONLY samples provided in the parameter files will be analyzed
        fastq_map = []
        for r1 in self._fastq_list:
            ID = os.path.basename(r1).split("_")[0]
            if ("_R1_" in r1) and (ID in self._sample_names):
                    r2 = r1.replace("_R1_", "_R2_")
                    fastq_map.append([r1,r2])

        # make folder to store alignment sam files
        sam_output = os.path.join(self._output, "sam_files/")
        os.system("mkdir " + sam_output)

        # make folder to store ref
        ref_path = os.path.join(self._output, "ref/")
        os.system("mkdir " + ref_path)

        # convert to df
        fastq_map = pd.DataFrame(fastq_map, columns=["R1", "R2"])
        # if phix
        phix_fasta = os.path.join(os.path.dirname(self._main_path), "data/phix.fasta")
        # self._log.info(os.path.dirname(self._main_path))
        # save this to ref path
        # self._log.info(phix_fasta)
        cmd = f"cp {phix_fasta} {ref_path}"
        # self._log.info(cmd)
        os.system(cmd)
        self._phix_fasta = []
        if self._args.calibratePhredPhix:
            # check if Undetermined reads are in the same folder
            self._phix_fastq = glob.glob(self._args.fastq+"/Undetermined_*.fastq.gz")
            if len(self._phix_fastq) < 2:
                raise ValueError("Cannot find Undetermined fastq files for phix alignment")
        #if self._phix_fastq != []:
            # add this to the last row
            fastq_map.loc[len(fastq_map)] = self._phix_fastq

        rc = False
        if self._args.rc:
            rc = True

        # GURU # depreciated
        # if self._args.environment == "GURU":
        #     # make sh files to submit to GURU
        #     sh_output = os.path.join(self._output, "GURU_jobs")
        #     os.system("mkdir "+sh_output)
        #     # for each pair of fastq files, submit jobs to run alignment
        #     # it takes the following arguments: R1, R2, ref, Output sam
        #     # the output log (bowtie log) would be in the same dir
        #     logging.info("Writing sh files for alignment (GURU)")
        #     cluster.alignment_sh_guru(fastq_map, self._project, self._seq.seq.item(), ref_path, sam_output, sh_output, logging)
        #     logging.info("Alignment jobs are submitted to GURU..")
        #     # wait for alignment to finish and call mutations

        if self._args.environment == "GALEN":
            # make sh files to submit to GURU
            sh_output = os.path.join(self._output, "GALEN_jobs")
            os.system("mkdir "+sh_output)
            # for each pair of fastq files, submit jobs to run alignment
            # it takes the following arguments: R1, R2, ref, Output sam
            # the output log (bowtie log) would be in the same dir
            logging.info("Writing sh files for alignment (GALEN)")
            sam_df, job_list = cluster.alignment_sh_galen(fastq_map, self._project, self._seq.seq.values.item(), ref_path,
                                            sam_output, sh_output, self._args.at, self._log, rc)
            logging.info("Alignment jobs are submitted to GALEN..")

        elif self._args.environment == "BC2" or self._args.environment == "BC" or self._args.environment == "DC":
            # make sh files to submit to BC
            sh_output = os.path.join(self._output, "CCBR_sh")
            os.system("mkdir " + sh_output)
            # make sh files to submit to BC
            self._log.info("Submitting alignment jobs...")
            # alignment_sh_bc2(fastq_map, ref_name, ref_seq, ref_path, sam_path, sh_output, at, logging)

            sam_df, job_list = cluster.alignment_sh_ccbr(fastq_map, self._project, self._seq.seq.values.item(),
                                                        ref_path, sam_output, sh_output, self._args.at,
                                                        self._log, rc, self._args.environment)
            self._log.info("Alignment jobs are submitte to BC2. Check pbs-output for STDOUT/STDERR")

        else:
            raise ValueError("Wrong environment name")

        self._log.info(f"Total jobs submitted: {len(job_list)}")
        if self._args.environment == 'GALEN':
            finished = cluster.parse_jobs_galen(job_list, self._logging.getLogger("track.jobs"))
        else:
            finished = cluster.parse_jobs(job_list, self._args.environment, self._logging.getLogger("track.jobs"))  #

        if finished:
            self._log.info(f"Alignment jobs are finished!")
            self._log.info(f"Merging alignment log files...")
            # go through alignment output dir and parse all the log files
            log_files = os.path.join(sam_output, "*.log")
            log_list = glob.glob(log_files)
            alignment_master_log = []
            for f in log_list:
                f_info = []
                sample_id = os.path.basename(f).split(".")[0]
                f_info.append(sample_id)
                with open(f, "r") as fp:
                    for i, line in enumerate(fp):
                        # pick corresponding lines in the alignment log file
                        if i == 0 or i == 5 or i == 6 or i == 11:
                            line = line.split(" ")
                            f_info.append(line[0])
                alignment_master_log.append(f_info)
            df = pd.DataFrame(alignment_master_log)
            df.columns = ["sample_id", "R1_reads", "R1_alignment_rate", "R2_reads", "R2_alignment_rate"]
            alignment_log = os.path.join(sam_output, "aligned_rate.log")
            df.to_csv(alignment_log, index=False)
            self._log.info(f"Alignment log file created")

        return sam_output

    def _mut_count(self):
        """
        Count mutations in input sam files
        1. Log path to sam files and downsampled sam files
        2. If sam df not provided, make df for all the sam files
        3. For each pair of sam files, submit jobs to the cluster
        4. Log to main log
        """
        self._log.info("Counting mutations ...")

        # submit job with main.py -r1 and -r2
        # run main.py with -r1 and -r2
        logger_mut = self._logging.getLogger("count.mut.log")
        self._log.info("Removing existing phred files...")
        #phred_flist = glob.glob(os.path.join(self._output, '*_phred.*'))
        #if phred_flist != []:
        #    os.system(f"rm {os.path.join(self._output, '*_phred.*')}")
        mut_counts = count_mut.readSam(self._r1, self._r2, self._param_json, self._args, self._output, self._args.c, logger_mut)
        self._log.info("Running multi-core analysis ... ")
        if self._args.calibratePhredWT:
            self._log.info("Adjusting phred scores using WT samples...")
            adjusted_phred = mut_counts.adjust_er(wt_override=self._args.wt_override)
        elif self._args.calibratePhredPhix:
            self._log.info("Adjusting phred scores using Phix...")
            adjusted_phred = mut_counts.adjust_er_phix()
        else:
            adjusted_phred = []
        mut_counts.multi_core(adjusted_phred)

    def _makejobs(self, sh_output, sam_dir, samples):
        """
        For each pair of sam files in output/sam_files/
        submit mut count job to the cluster
        """
        if sam_dir == "":
            # get samples in param file
            sam_dir = os.path.join(self._args.output, "sam_files/") # read sam file from sam_file
        # get sam files from parameter file
        if not os.path.isdir(sam_dir):
            self._log.error(f"Directory: ./sam_files/ not found in {self._output}")
            raise ValueError()

        self._log.debug(f"Sam files are read from {sam_dir}")
        job_list = []
        for sample in samples:
            # assume all the sam files have the same name format (id_*.sam)
            sam_f_r1 = glob.glob(f"{sam_dir}/{sample}_*R1_*.sam")
            sam_f_r2 = glob.glob(f"{sam_dir}/{sample}_*R2_*.sam")
            if len(sam_f_r1) == 0 or len(sam_f_r2) == 0:
                self._log.error(f"SAM file for sample {sample} not found. Please check your parameter file")
                raise ValueError()

            else:
                self._log.info(f"Sample {sample}")
                self._log.info(f"Read1: {sam_f_r1[0]}")
                self._log.info(f"Read2: {sam_f_r2[0]}")
                sam_id_1 = os.path.basename(sam_f_r1[0]).split("_")[0]
                sam_id_2 = os.path.basename(sam_f_r2[0]).split("_")[0]
                if (sam_id_1 != sample) or (sam_id_1 != sam_id_2) or (sam_id_2 != sample):
                    self._log.error("IDs in sam files don't match!")
                self._r1 = sam_f_r1[0]
                self._r2 = sam_f_r2[0]

            job_id = self._submit_mutcount_jobs(sample, sh_output)
            job_list.append(job_id)
        self._log.info("all samples submitted")

        # after submitting all jobs, track all submitted jobs on cluster
        jobs = ",".join(job_list)
        self._log.debug(f"All jobs: {jobs}")
        self._log.info(f"Total jobs running: {len(job_list)}")
        if self._args.environment == 'GALEN':
            finished = cluster.parse_jobs_galen(job_list, self._logging.getLogger("track.jobs"))
        else:
            finished = cluster.parse_jobs(job_list, self._args.environment, self._logging.getLogger("track.jobs"))

        return finished

    def _submit_mutcount_jobs(self, sample, sh_output):
        """
        Make mutation counting jobs and submit for this sample
        """
        # submit job with main.py -r1 and -r2
        # run main.py with -r1 and -r2
        cmd = f"tileseq_mut -n {self._args.name} -r1 {self._r1} -r2 {self._r2} -o {self._output} -p" \
                  f" {self._param_json} --skip_alignment -log {self._args.log_level} -env {self._args.environment} -at {self._args.at} -mt {self._args.mt} -c {self._args.c} "
        if self._args.sr_Override:
            cmd = cmd + "-override "

        if self._args.wt_override:
            cmd = cmd + "--wt_override "

        if self._args.calibratePhredPhix:
            cmd = cmd + "--calibratePhredPhix "

        if self._args.calibratePhredWT:
            cmd = cmd + "--calibratePhredWT "

        if self._args.environment == "BC2" or self._args.environment == "BC" or self._args.environment == "DC":
            logging.info("Submitting mutation counts jobs....")
            job_id = cluster.mut_count_sh_ccbr(sample, cmd, self._args.mt, self._args.mm, sh_output, self._log,
                                             self._args.c, self._args.environment)
        # elif :
        #     logging.info("Submitting mutation counts jobs to DC...")
        #     # (sample_name, cmd, mt, sh_output_dir, logger)
        #     job_id = cluster.mut_count_sh_dc(sample, cmd, self._args.mt, self._args.mm, sh_output, self._log,
        #                                      self._args.c)
        elif self._args.environment == "GALEN":
            logging.info("Submitting mutation counts jobs to GALEN...")
            # (sample_name, cmd, mt, sh_output_dir, logger)
            job_id = cluster.mut_count_sh_galen(sample, cmd, self._args.mt, self._args.mm, sh_output, self._log,
                                                self._args.c)  # this
        else:
            raise ValueError("Wrong environment")

        return job_id

    def _checkoutput(self, f):
        """
        Check how many variants are generated in each file
        """
        mut_n = 0
        with open(f, "r") as mut_output:
            for line in mut_output:
                # skip header
                if "c." in line:
                    mut_n += 1
        return mut_n

    def main(self):
        """
        Main for the fastq2counts object
        """
        sam_dir = ""
        if self._skip == False: # submit jobs for alignment
            # get sam_df from self._align_sh()
            sam_dir = self._align_sh_()

        if self._r1 != "" and self._r2 != "":
            # call functions in count_mut.py
            # init(sam_r1, sam_r2, param, args, ouptut_dir, logger)
            self._mut_count()

        # submit jobs for mutation counting
        # if user did not provide r1 and r2 SAM file
        # get r1 and r2 sam files from output dir provided by user
        elif self._r1 == "" and self._r2 == "":
            # resubmit is used when some jobs in the queue did not finish properly
            if self._args.resubmit:
                self._skip = True
                # resubmit failed jobs in existing mut_count dir
                # self.output is the mut_count dir in this case
                # find out which jobs failed by going through all counts files

                # first find if any phred calibration files are empty
                # remove empty files
                self._log.info("Removing existing phred files if they are empty...")
                phred_flist = glob.glob(os.path.join(self._output, '*_phred.*'))
                for ph in phred_flist:
                    if os.stat(ph).st_size == 0:
                        os.system(f"rm {os.path.join(self._output, ph)}")
                # get all the mut count files
                mutcount_list = glob.glob(os.path.join(self._output, "counts_sample_*.csv"))
                failed_samples = []
                all_samples_with_csv = []
                # go through all the mut count files
                for f in mutcount_list:
                    m = re.search('.*counts_sample_(.+?).csv', f)
                    # find sample ID in file name
                    if m:
                        sample_id = m.group(1)
                    else:
                        self._log.error("Sample ID not found")
                        raise ValueError(f"Sample ID not found for {f}")
                    all_samples_with_csv.append(sample_id)
                    # check if the file is empty
                    mut_n = self._checkoutput(f)
                    if mut_n == 0:
                        failed_samples.append(sample_id)
                # if any sample does not have a csv file generated (jobs were killed before started)
                other_missing = [i for i in self._sample_names if i not in all_samples_with_csv]
                failed_samples += other_missing
                self._log.info(f"Failed samples: {failed_samples}")
                self._log.info("Resubmitting samples ...")
                # find the path to jobs dir
                env_jobs_dir = os.path.join(self._output, f"{self._args.environment}_jobs")
                parent_dir = os.path.abspath(os.path.join(self._output, os.pardir))
                sam_dir = os.path.join(parent_dir, "sam_files")
                # remake sh files for these samples and submit them
                finished = self._makejobs(env_jobs_dir, sam_dir, failed_samples)

                self._log.info("ALL DONE!")
                return finished

            # start alignment jobs
            if self._skip == False:
                time_now = datetime.datetime.now().strftime("%Y-%m-%d-%H-%M-%S")
                # make mut count dir alignemnt not skipped
                # otherwise the user input dir should be the output dir generated
                self._output = os.path.join(self._output, self._args.name + "_" + time_now + "_mut_count")
                os.makedirs(self._output)
                # align to phix if specified
                if self._args.calibratePhredPhix:
                    # check if Undetermined reads are in the same folder
                    self._phix_fastq = glob.glob(self._args.fastq+"/Undetermined_*.fastq.gz")
                    if len(self._phix_fastq) < 2:
                        raise ValueError("Cannot find Undetermined fastq files for phix alignment")

            # output directory is the mut_count dir
            # make folder to store all the sh files as well as all the log files
            if self._args.environment == "BC2" or self._args.environment == "BC":
                sh_output = os.path.join(self._output, "BC_jobs")

            elif self._args.environment == "DC":
                sh_output = os.path.join(self._output, "DC_jobs")

            elif self._args.environment == "GALEN":
                sh_output = os.path.join(self._output, "GALEN_jobs")

            else:
                self._log.error("Please provide valid environment: BC/BC2/DC/Galen")
                raise ValueError("Please provide valid environment: BC/BC2/DC/Galen")

            self._log.info(f"Mutation count sh files are made in {sh_output}")

            os.mkdir(sh_output)
            finished = self._makejobs(sh_output, sam_dir, self._sample_names)

            if finished:
                self._log.info(f"Mutation counting jobs are finished!")
                self._log.info("Check mutation counts file ...")
                mutcount_list = glob.glob(os.path.join(self._output, "counts_sample_*.csv"))
                self._log.info(f"{len(mutcount_list)} mutation counts file generated")
                if len(self._sample_names) > len(mutcount_list):
                    self._log.error("Job finished but some variant call files are missing!!")
                    missing = list(set(self._sample_names)-set(mutcount_list))
                    missing = ",".join(missing)
                    self._log.error(f"Missing samples: {missing}")
                # delay here for the system to write to the files
                os.system("sleep 180")
                for f in mutcount_list:
                    mut_n = self._checkoutput(f)
                    if mut_n == 0:
                        self._log.error(f"{f} has 0 variants! Check mut log for this sample.")
                    else:
                        self._log.info(f"{f} has {mut_n} variants")

def check(args):
    """
    Validate args and convert csv to JSON
    return path to json
    """

    # two script needed: csv2json.R and calibratePhred.R
    test_csv2json = subprocess.getstatusoutput("csv2json.R -h")
    test_caliPhred = subprocess.getstatusoutput("calibratePhred.R -h")
    if test_csv2json[0] != 0:
        raise OSError(f"csv2json.R failed with error {test_csv2json[1]}")
    if test_caliPhred[0] != 0:
        raise OSError(f"calibratePhred.R failed with error {test_caliPhred[1]}")

    # check bowtie2 and bowtie2-build
    test_bt = subprocess.getstatusoutput("bowtie2 -h")
    test_bt_build = subprocess.getstatusoutput("bowtie2-build -h")
    if test_bt[0] != 0:
        raise OSError(f"bowtie2 failed with error {test_bt[1]}")
    if test_bt_build[0] != 0:
        raise OSError(f"bowtie2-build failed with error {test_bt_build[1]}")

    # check if the correct args are provided
    # if you don't specify --skip_alignment then you cannot provide r1 and r2 sam_files
    if not args.skip_alignment:
        if args.r1 or args.r2:
            raise ValueError(f"Invalid paramters! Please specify --skip_alignment if you want to analyze\
            one pair of sam files")

    # check if output dir exists
    if not os.path.isdir(args.output):
        raise FileNotFoundError(f"Output directory not found: {args.output}")

    # check if fastq dir exists
    if args.fastq and not os.path.isdir(args.fastq):
        raise FileNotFoundError(f"Fastq file path not found: {args.fastq}")

    # try convert csv to json in the same dir as the csv file
    # convert csv file to json
    if args.param.endswith(".csv"):
        print("parameter sheet in csv format, converting to json")
        param_json = args.param.replace(".csv", ".json")
        if os.path.isfile(param_json):
             os.remove(param_json)
        if args.sr_Override:
            convert = f"csv2json.R {args.param} -o {param_json} --srOverride"
        else:
            convert = f"csv2json.R {args.param} -o {param_json}"
        os.system(convert)

    # if the file ends with .json, do nothing
    elif args.param.endswith(".json"):
        param_json = args.param
    else:
        raise ValueError("Please provide valid paramter file format (csv or json)")

    if not os.path.isfile(param_json):
        raise ValueError("Json file does not exist, check conversion!")

    # check environment input
    if args.environment.upper() not in ["BC", "BC2", "DC", "GALEN"]:
        raise ValueError(f"Wrong input for environment: {args.environment}")

    return param_json

def main(args, v):
    """
    Main for fastq2counts
    """
    # get time stamp for this object
    time_now = datetime.datetime.now().strftime("%Y-%m-%d-%H-%M-%S")

    # check input arguments
    param_json = check(args)

    # get basename for the param file
    param_base = os.path.basename(param_json)
    if args.skip_alignment: # if we want to skip the alignment part
        # if user provided R1sam + R2sam
        if args.r1 and args.r2:
            # analyze one pair of r1 and r2 file
            # VALIDATE if both files are sam files
            if not args.r1.endswith(".sam") or not args.r2.endswith(".sam"):
                raise ValueError("Please provide SAM files")
            # copy the parameter file to output dir if the json file is not in the dir
            param_path = os.path.join(args.output, param_base)
            if not os.path.isfile(param_path):
                param_path = shutil.copy(param_json, args.output, follow_symlinks=True)
            # set up loggingthis creates main.log in the mut_count output dir
            main_log_f = os.path.join(args.output, "mutcount_main.log")
            main_log = help_functions.logginginit(args.log_level, main_log_f)

            # initialize mutcount object
            mc = fastq2counts(param_path, args.output, main_log, args)
            mc._init_skip(skip=True, r1=args.r1, r2=args.r2)

        else:
            if not args.resubmit:
                # user provided a csv file and path to sam files
                # for each pair of sam file we would submit one job to the cluster for mutation counting
                # make time stamped output folder for this project
                updated_out = os.path.join(args.output, args.name + "_" + time_now + "_mut_count")
                # sam_dir = os.path.join(args.output, "sam_files/") # read sam file from sam_file
                os.makedirs(updated_out)
            else:
                updated_out = args.output

            args_log_path = os.path.join(updated_out, f"{time_now}_args.log")
            # record input parameters
            write_param(args_log_path, args, v)

            # load json file
            param_path = os.path.join(updated_out, param_base)
            # copy this file to the output folder
            if not os.path.isfile(param_path):
                param_path = shutil.copy(param_json, updated_out, follow_symlinks=True)
            main_log_f = os.path.join(args.output, "mutcount_main.log")
            main_log = help_functions.logginginit(args.log_level, main_log_f)
            # initialize mutcount object
            mc = fastq2counts(param_path, updated_out, main_log, args)
            # skip = skip alignment
            mc._init_skip(skip=True)
    else:
        # alignment
        # make time stamped output folder for this project
        updated_out = os.path.join(args.output, args.name + "_" + time_now)
        os.makedirs(updated_out)  # make directory to save this run
        # write args to file
        args_log_path = os.path.join(updated_out, f"{time_now}_args.log")
        write_param(args_log_path, args, v)

        param_path = os.path.join(updated_out, param_base)
        if not os.path.isfile(param_path):
            param_path = shutil.copy(param_json, updated_out, follow_symlinks=True)

        main_log_f = os.path.join(updated_out, "main.log")
        main_log = help_functions.logginginit(args.log_level, main_log_f)

        # initialize mutcount object
        mc = fastq2counts(param_path, updated_out, main_log, args)
        mc._init_skip(skip=False)

    # start the run
    mc.main()

def write_param(args_log_path, args, v):
    """
    Write input arguments to param.log
    """
    with open(args_log_path, "w") as args_log:
        args_log.write(f"{v}\n")
        for arg in vars(args):
            args_log.write(arg+",")
            args_log.write(f"{getattr(args, arg, 'N/A')}\n")

if __name__ == "__main__":

    parser = argparse.ArgumentParser(description='TileSeq mutation counts')
    # user input arguments
    parser.add_argument("-f", "--fastq", help="Path to all fastq files you want to analyze", type=str)
    parser.add_argument("-o", "--output", help="Output folder", type=str, required=True)
    parser.add_argument("-p", "--param", help="csv paramter file", type=str, required=True)
    parser.add_argument("-n", "--name", help="Name for this run", type=str, required=True)
    parser.add_argument("--skip_alignment", action="store_true", help="skip alignment for this analysis, "
                                                                      "ONLY submit jobs for counting mutations in existing output folder")
    parser.add_argument("-r1", help="r1 SAM file", type=str)
    parser.add_argument("-r2", help="r2 SAM file", type=str)

    # user input arguments with default values set
    parser.add_argument("-log", "--log_level", help="set log level: debug, \
        info, warning, error, critical. (default = info)", type=str, default="info")
    parser.add_argument("-env", "--environment", help= "The cluster used to \
        run this script (default = DC)",type=str, default="DC")
    parser.add_argument("-at", type = int, help="Alignment time \
        (default = 8h)", default=8)
    parser.add_argument("-mt", type = int, help="Mutation call time \
        (default = 36h)", default=36)
    parser.add_argument("-mm", type=int, help="Mutation call request memory \
            (default = 15GB)", default=15)
    parser.add_argument("-c", type=int, help="Number of cores", default=16)
    parser.add_argument("-b", "--base", help="ASCII code base", default=33)
    parser.add_argument("-test", action="store_true", help="Turn on testing mode")
    parser.add_argument("-rc", action="store_true", help="Turn on rc mode, both direction of the reads will be "
                                                         "aligned to the reference. Variant calling will be "
                                                         "performed on all the reads that are aligned (BE CAREFUL!)")

    parser.add_argument("-override", "--sr_Override", action="store_true", help="Provide this argument when there is only one replicate")
    parser.add_argument("--posteriorQC", action="store_true", help="Turn on posterior QC mode, this requires more "
                                                                   "memory and runtime, please change the variable "
                                                                   "accordingly")
    parser.add_argument("--wt_override", action="store_true", help="When no wt conditions defined in the parameter sheet, turn on this option will treat EVERYTHING as wt. Phred scores will be adjusted based on the first replicate")

    parser.add_argument("--calibratePhredWT", action="store_true", help="When this parameter is provided, "
                                                                      "use wt to calibrate phred scores")
    parser.add_argument("--calibratePhredPhix", action="store_true", help="When this parameter is provided, "
                                                                      "use phix alignments to calibrate phred scores")
    parser.add_argument("--resubmit", action="store_true", help="For a finished run, batch resubmit failed scripts ("
                                                                "if any)")
    args = parser.parse_args()
    args.environment = args.environment.upper()

    # main(args)


text:
#!/usr/bin/env python

import rospy
import rosnode
import numpy as np
from std_msgs.msg import String
from std_msgs.msg import Float64
from dynamixel_msgs.msg import JointState
from math import sin
import time

RA1 = rospy.Publisher('/RA1_controller/command', Float64, queue_size=10)
RA2 = rospy.Publisher('/RA2_controller/command', Float64, queue_size=10)
RA3 = rospy.Publisher('/RA3_controller/command', Float64, queue_size=10)
RA4 = rospy.Publisher('/RA4_controller/command', Float64, queue_size=10)
LA1 = rospy.Publisher('/LA1_controller/command', Float64, queue_size=10)
LA2 = rospy.Publisher('/LA2_controller/command', Float64, queue_size=10)
LA3 = rospy.Publisher('/LA3_controller/command', Float64, queue_size=10)
LA4 = rospy.Publisher('/LA4_controller/command', Float64, queue_size=10)
Goal_pos = 0;
Start_pos = -.50;
Tf = 50 ; #no of sec * 50

# 4 angle values to be sent for each joint 2 shoulder 1 elbow 1 wrist
def right_arm(Angles):
    Curr_pos = 0 ;

    rate = rospy.Rate(50) # 50hz
    time.sleep(1);
    flag = 1;
    while (flag == 1 and rospy.is_shutdown() == 0):
	rospy.loginfo(rospy.get_caller_id() + " Publishing %s to right arm 1" ,Curr_pos );
	RA1.publish(Curr_pos );
	time.sleep(1);

	rospy.loginfo(rospy.get_caller_id() + " Publishing %s to right arm 2" ,Curr_pos );
	RA2.publish(1.57);
	time.sleep(1);

	rospy.loginfo(rospy.get_caller_id() + " Publishing %s to right arm 3" ,Curr_pos );
	RA3.publish(Curr_pos );
	time.sleep(1);

	rospy.loginfo(rospy.get_caller_id() + " Publishing %s to right arm 4" ,Curr_pos );
	RA4.publish(Curr_pos );
	flag = 0;
	#rate.sleep();

def left_arm():
    Curr_pos = 0 ;

    rate = rospy.Rate(50) # 50hz
    timer = 0;
    forw = 1;
    time.sleep(1);
    flag = 1;
    while (flag == 1 and rospy.is_shutdown() == 0):
	rospy.loginfo(rospy.get_caller_id() + " Publishing %s to left arm 1" ,Curr_pos );
	RA1.publish(1.0 );
	time.sleep(1);

	rospy.loginfo(rospy.get_caller_id() + " Publishing %s to left arm 2" ,Curr_pos );
	LA2.publish(-1.57 );
	time.sleep(1);

	rospy.loginfo(rospy.get_caller_id() + " Publishing %s to left arm 3" ,Curr_pos );
	LA3.publish(Curr_pos );
	time.sleep(1);

	rospy.loginfo(rospy.get_caller_id() + " Publishing %s to left arm 4" ,Curr_pos );
	LA4.publish(Curr_pos );
	time.sleep(1);
	flag = 0;
	#rate.sleep();

if __name__ == '__main__':
    try:

    	rospy.init_node('main_client', anonymous=True)
        right_arm(['1','1',0,0]);
        left_arm();

    except rospy.ROSInterruptException:
        pass

ext: py

max_stars_repo_name: DIvyanshu-Goel/Biped

max_issues_repo_name: DIvyanshu-Goel/Biped

max_forks_repo_name: DIvyanshu-Goel/Biped

reason: import numpy


text:
import numpy as np
import glob
import os
from scipy.optimize import curve_fit
from scipy.interpolate import interp1d
import matplotlib.pyplot as plt

def sigmoid(x, scale, steep, loc):
    """ sigmoid function """
    return scale/(1+np.exp(-steep * (x - loc)))

def sigmoid_mixture_model(scales, steeps, locs):
    """ Returns a sigmoid mixture model """
    def model(x):
        y = np.zeros(x.shape)
        for scale, steep, loc in zip(scales, steeps, locs):
            y += sigmoid(x, scale, steep, loc)
        return y
    return model

def sigmoid_mixture_model_wrap(x, *pars):
    n = int(len(pars) / 3)
    scales = [pars[i] for i in range(n)]
    steeps = [pars[n+i] for i in range(n)]
    locs = [pars[2*n+i] for i in range(n)]
    return sigmoid_mixture_model(scales, steeps, locs)(x)

def fit_sigmoid_mixture_model(x, y, n=4):
    """ Fits and returns the sigmoid mixture model """
    locs_init = np.min(x) + (np.max(x) - np.min(x)) * (np.arange(n) + 0.5) / n
    scales_init = (np.max(y) - np.min(y)) * np.ones(n) / n
    steeps_init = 10 * np.ones(n) / ((np.max(x) - np.min(x)) / n)
    pars = np.append(np.append(locs_init, scales_init), steeps_init)
    p_opt, _ = curve_fit(sigmoid_mixture_model_wrap, x, y, p0=list(pars), method='trf')
    scales = [p_opt[i] for i in range(n)]
    steeps = [p_opt[n+i] for i in range(n)]
    locs = [p_opt[2*n+i] for i in range(n)]
    return sigmoid_mixture_model(scales, steeps, locs)

# data can be found at:
# https://hitran.org/hitemp/

# cross section for wavelength between 625/cm and 750/cm
data_625_750 = np.genfromtxt('02_625-750_HITEMP2010.par', usecols=(1, 2))
wavelength = data_625_750[:, 0]    # in 1/cm
cross_section = data_625_750[:, 1]  # in cm^2

plt.plot(wavelength, cross_section, color='k')
plt.xlabel('wave number [1 / cm]', fontsize=16)
plt.ylabel('cross section [cm^2]', fontsize=16)
plt.savefig('c02_wavelength_625_750.png', bbox_inches='tight')
plt.close()

# cross section for wavelength between 2250/cm and 2500/cm
data_2250_2500 = np.genfromtxt('02_2250-2500_HITEMP2010.par', usecols=(1, 2))
wavelength = data_2250_2500[:, 0]    # in 1/cm
cross_section = data_2250_2500[:, 1]  # in cm^2

plt.plot(wavelength, cross_section, color='k')
plt.xlabel('wave number [1 / cm]', fontsize=16)
plt.ylabel('cross section [cm^2]', fontsize=16)
plt.savefig('c02_wavelength_2250_2500.png', bbox_inches='tight')
plt.close()

# plot all cross sections in the relevant regime
wave_numbers_all = np.array([])
cross_section_all = np.array([])
for _file in glob.glob('02_*.par'):
    data = np.genfromtxt(_file, usecols=(1, 2))
    wave_numbers_all = np.append(wave_numbers_all, data[:, 0])
    cross_section_all = np.append(cross_section_all, data[:, 1])

# sort wave numbers
_sort = np.argsort(wave_numbers_all)
wave_numbers_all = wave_numbers_all[_sort]
# remove double occuring wavelengths
wave_numbers_all, indices = np.unique(wave_numbers_all, return_index=True)
cross_section_all = cross_section_all[_sort][indices]

plt.step(wave_numbers_all, cross_section_all, where='mid', color='k', linewidth=0.1)
plt.xlabel('wave number [1 / cm]', fontsize=16)
plt.ylabel('cross section [cm^2]', fontsize=16)
plt.savefig('c02_wavelength_all.pdf', bbox_inches='tight')
plt.close()

# absorption height for vertical upward going rays
h0 = 1e4
co2_fraction = 0.0004
n_density = 2.504e25
attenuation = cross_section_all * 1e-4 * co2_fraction * n_density
block = 1 - np.exp(-attenuation*h0)
# cut out only high cross sections
cut = block > 1e-3
wn_mid = (wave_numbers_all[1:] + wave_numbers_all[:-1]) / 2.
wn_left = np.append(wn_mid[0]-2*(wn_mid[0]-wave_numbers_all[0]), wn_mid)[cut]
wn_right = np.append(wn_mid, wn_mid[-1]+2*(wave_numbers_all[-1]-wn_mid[-1]))[cut]
print('total wavelength in database: ', len(wave_numbers_all))
print('subset with cut 1/(100*h0): ', len(wave_numbers_all[cut]))
block_cut = np.copy(block)
block_cut[~cut] = 0
os.makedirs('c02_absorped', exist_ok=True)
for xmin, xmax in zip([500, 500, 800, 1150, 1850, 2000, 2200], [2500, 800, 1150, 1850, 2000, 2200, 2500]):
    plt.step(wave_numbers_all, block, where='mid', color='grey', linewidth=0.1)
    plt.step(wave_numbers_all, block_cut, where='mid', color='red', linewidth=0.1)
    plt.xlabel('wave number [1 / cm]', fontsize=16)
    plt.ylabel('fraction absorped', fontsize=16)
    plt.xlim([xmin, xmax])
    plt.savefig('c02_absorped/min%i_max%i.pdf' % (xmin, xmax) if (xmax-xmin != 2000) else 'c02_absorped.pdf', bbox_inches='tight')
    plt.close()

interp_indices = interp1d(wave_numbers_all[cut], np.cumsum(cut)[cut]-1, kind='nearest',
                          assume_sorted=True, bounds_error=False)
plt.plot(wave_numbers_all, np.cumsum(cut)-1, color='k')
_wn = np.linspace(np.min(wave_numbers_all), np.max(wave_numbers_all), 100000)
plt.plot(_wn, interp_indices(_wn), color='red', alpha=0.3)
plt.xlabel('wave number [1 / cm]', fontsize=16)
plt.ylabel('indice to take', fontsize=16)
plt.savefig('c02_wave_number_indice.pdf', bbox_inches='tight')
plt.close()
np.savez('co2_cross_section.npz', wave_numbers_left=wn_left, wave_numbers_right=wn_right,
         cross_sections=cross_section_all[cut], interp_indices=interp_indices)

# _model = fit_sigmoid_mixture_model(wn_interpolate, np.cumsum(cut)[indices]-1, n=3)
# _wn = np.linspace(np.min(wavelength_all), np.max(wavelength_all), 100000)
# print(_model)
# print(_model(600 * np.ones(1)))
# print(_model(_wn))
#
# plt.plot(wavelength_all[_sort], np.cumsum(cut)-1, color='k')
# plt.plot(_wn, _model(_wn), color='red', alpha=0.5)
# plt.xlabel('wave number [1 / cm]', fontsize=16)
# plt.ylabel('indice to take', fontsize=16)
# plt.savefig('c02_wave_number_indice_sigmoid_model.pdf', bbox_inches='tight')
# plt.close()

# relevance at 270 K temperature
c = 3e8

def lam2nu(lam):
    """ Wavelength to frequency converter """
    return c / lam

def planck(nu, T, norm=False):
    """
    Planck distribution for black body radiation.

    :param nu: frequency (in Hz)
    :param T: temperature (in K)
    :return: energy density, in J / m^2
    """
    nu = np.atleast_1d(nu)
    a = 2.0 * 6.62607e-34 * nu**3 / (3e8)**2
    b = 6.62607e-34 * nu / 1.38065e-23 / T
    intensity = a / (np.exp(b) - 1.)
    if norm:
        intensity /= np.sum(intensity)
    return intensity

nu = lam2nu(1 / wave_numbers_all * 1e-2)  # frequency in Hz
plt.plot(wave_numbers_all, cross_section_all * planck(nu, 270) * nu)
plt.xlabel('wave number [1 / cm]', fontsize=16)
plt.ylabel('relevance', fontsize=16)
plt.savefig('c02_relevance_all.pdf', bbox_inches='tight')
plt.close()

max_stars_repo_name: MRCWirtz/ghsim

max_issues_repo_name: MRCWirtz/ghsim

max_forks_repo_name: MRCWirtz/ghsim

reason: import numpy,from scipy


text:
# -*- coding: utf-8 -*-
"""
Created on Tue May 24 18:05:04 2016

@author: bing
"""

#!python

import numpy as np
from fastkde import fastKDE
import pylab as PP

#Generate two random variables dataset (representing 100000 pairs of datapoints)
N = 2e5
var1 = 50*np.random.normal(size=N) + 0.1
var2 = 0.01*np.random.normal(size=N) - 300

#Do the self-consistent density estimate
myPDF,axes = fastKDE.pdf(var1,var2)

#Extract the axes from the axis list
v1,v2 = axes

#Plot contours of the PDF should be a set of concentric ellipsoids centered on
#(0.1, -300) Comparitively, the y axis range should be tiny and the x axis range
#should be large
PP.contour(v1,v2,myPDF)
PP.show()


text:
import numpy as np
import pandas as pd
import plotly.graph_objs as go
import numpy as np
import dash
import dash_core_components as dcc
import dash_html_components as html
from dash.dependencies import Input,Output,State

df = pd.read_csv('OldFaithful.csv')

# Css Styles
tabs_styles = {
    'height': '44px'
}
tab_style = {
    'borderBottom': '1px solid #d6d6d6',
    'padding': '6px',
    'fontWeight': 'bold'
}

tab_selected_style = {
    'borderTop': '1px solid #d6d6d6',
    'borderBottom': '1px solid #d6d6d6',
    'backgroundColor': '#119DFF',
    'color': 'white',
    'padding': '6px'
}

app = dash.Dash()

app.layout = html.Div([
                        dcc.Tabs(id="tabs", value='tab-1',
                        children=[dcc.Tab(label='Tab one', value='tab-1', style=tab_style,selected_style=tab_selected_style),
                                    dcc.Tab(label='Tab two', value='tab-2', style=tab_style,selected_style=tab_selected_style)],
                        style=tabs_styles),

                        html.Div(id='show-tabs-content'),

                                            ])

@app.callback(Output('show-tabs-content', 'children'),
              [Input('tabs', 'value')])

def render_content(tab):
    value = 'OldFaithful'
    if tab == 'tab-1':
        return html.Div([dcc.Graph(id='OldFaithful',
                                figure = {'data':[go.Scatter(x=df['X'],y=df['Y'],mode='markers',
                                marker=dict(symbol='pentagon',size=12,color='rgb(50,200,50')
                                )],
                                'layout':go.Layout(title='Old Faithful',
                                                    xaxis=dict(title='Duration'),yaxis=dict(title='Time until Next')),
                                                    })])

    elif tab == 'tab-2':
        return html.Div(html.H3('Tab content 2'))

if __name__ == '__main__':
    app.run_server()


text:
import matplotlib.pyplot as plt
import matplotlib.pylab as pylab

import requests
from io import BytesIO
from PIL import Image
import numpy as np
import os

# this makes our figures bigger
pylab.rcParams['figure.figsize'] = 20, 12
os.environ['CUDA_VISIBLE_DEVICES']='2'

from maskrcnn_benchmark.config import cfg
from predictor import COCODemo

"""
RPN
    POST_NMS_TOP_N_TRAIN    2000
    POST_NMS_TOP_N_TEST     1000
    NMS_THRESH              0.7

ROI_HEAD
    ROI_HEADS.SCORE_THRESH  0.01
    ROI_HEADS.NMS           0.3
    DETECTIONS_PER_IMG      100

"""
config_file = "../configs/e2e_faster_rcnn_R_50_C4_1x.yaml"

# update the config options with the config file
cfg.merge_from_file(config_file)
# manual override some options
cfg.merge_from_list(["MODEL.DEVICE", "cuda"])

def load(file_name):
    """
    Gvien the file name, open the image and return PIL format image
    """
    image_dir = "../flickr_datasets/flickr30k_images"
    pil_image = Image.open(os.path.join(image_dir,file_name))
    # convert to BGR format
    image = np.array(pil_image)[:, :, [2, 1, 0]]
    return image

def imshow(img):
    plt.imshow(img[:, :, [2, 1, 0]])
    plt.axis("off")
    plt.show()
#! create Faster RCNN with the pre-trained weight
coco_demo = COCODemo(
    cfg,
    min_image_size=800,
    confidence_threshold=0.7,
    weight_loading="../pretrained_models/e2e_faster_rcnn_R_50_C4_1x.pth"
)

#! extract the training image's feature and proposal
with open('/home/zhangjiacheng/Code/maskrcnn-benchmark/flickr_datasets/split/train_backpack.txt','r') as f:
    image_name_list = f.readlines()
    for image_name in image_name_list:
        image_file = image_name.strip()+'.jpg'

        image = load(image_file)
        predictions = coco_demo.run_on_opencv_image(image, image_file)


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

import pathlib
from collections import OrderedDict
import tensorflow as tf
from tensorflow.keras import layers
import numpy as np
import pandas as pd
import scipy.stats
import math
import matplotlib.pyplot as plt
from baselines.common.schedules import LinearSchedule
from baselines.deepq.replay_buffer import ReplayBuffer, PrioritizedReplayBuffer

tf.enable_eager_execution()
tf.keras.backend.set_floatx('float32')

# Gaussian Noise Class (Diagonal Covariance Matrix)
class GaussActionNoise:
    def __init__(self, mean, std_deviation, dim = 2):
        self.mean = mean
        self.std_dev = std_deviation
        self.dim = dim

    def __call__(self):
        x = np.random.normal(self.mean, self.std_dev, self.dim)
        return x

# Parent Buffer class
class Buffer:
    def __init__(self, buffer_capacity=100000, batch_size=64, prioritized_replay_eps=1e-6):
        # Number of "experiences" to store at max
        self.buffer_capacity = buffer_capacity
        # Num of tuples to train on.
        self.batch_size = batch_size
        self.prioritized_replay_eps = prioritized_replay_eps

    def __len__(self):
        return len(self.buffer)

    # Takes (s,a,r,s') obervation tuple as input
    def record(self, obs, action, rew, new_obs, done):
        self.buffer.add(obs, action, rew, new_obs, float(done))

    def learn(self, beta):

        experience = self.buffer.sample(self.batch_size, beta = beta)
        (obses_t, actions, rewards, obses_tp1, dones, weights, batch_idxes) = experience

        rewards = tf.expand_dims(rewards, 1)

        update_metrics = self.update(obses_t, actions, rewards, obses_tp1, dones, weights.astype(np.float32))

        td_errors = update_metrics[0]

        # update priorities
        new_priorities = np.abs(td_errors) + self.prioritized_replay_eps

        self.buffer.update_priorities(batch_idxes, new_priorities)

        return update_metrics

# Q(s,a) Buffer
class QsaBuffer(Buffer):

    def __init__(self, buffer_capacity=100000, batch_size=64, alpha = 0.6):

        super(QsaBuffer, self).__init__(buffer_capacity, batch_size, prioritized_replay_eps = 1e-6)

        self.buffer = PrioritizedReplayBuffer(self.buffer_capacity, alpha = alpha)

    # Eager execution is turned on by default in TensorFlow 2. Decorating with tf.function allows
    # TensorFlow to build a static graph out of the logic and computations in our function.
    # This provides a large speed up for blocks of code that contain many small TensorFlow operations such as this one.
    @tf.function
    def update(
        self, state_batch, action_batch, reward_batch, next_state_batch, dones_batch, impt_weights_batch
    ):
        with tf.GradientTape() as tape:
            target_actions = mm_target_actor(next_state_batch, training=True)
            y = reward_batch + gamma * mm_target_qsa(
                [next_state_batch, target_actions], training=True
            )
            qsa_value = mm_qsa([state_batch, action_batch], training=True)
            td_errors = y - qsa_value
            qsa_loss = tf.math.reduce_mean(impt_weights_batch * tf.math.square(td_errors))

        qsa_grad = tape.gradient(qsa_loss, mm_qsa.trainable_variables)
        mm_qsa_optimizer.apply_gradients(
            zip(qsa_grad, mm_qsa.trainable_variables)
        )

        qsa_grad_list = []
        for grad in qsa_grad:
            qsa_grad_list.append(tf.math.reduce_mean(tf.abs(grad)))

        return td_errors, qsa_loss, tf.math.reduce_mean(qsa_grad_list)

# MM Actor Buffer Class
class ActorBuffer(Buffer):

    def __init__(self, buffer_capacity=100000, batch_size=64):

        super(ActorBuffer, self).__init__(buffer_capacity, batch_size)

        self.buffer = ReplayBuffer(self.buffer_capacity)

    # Eager execution is turned on by default in TensorFlow 2. Decorating with tf.function allows
    # TensorFlow to build a static graph out of the logic and computations in our function.
    # This provides a large speed up for blocks of code that contain many small TensorFlow operations such as this one.
    @tf.function
    def update(
        self, state_batch, action_batch, reward_batch, next_state_batch,
    ):
        with tf.GradientTape() as tape:

            actions = mm_actor(state_batch, training = True)

            nextState_Mean = state_batch \
                                + time_delta / cap * (actions * cap_air * (temp_air - state_batch)) * is_room \
                                + (tf.linalg.matmul(state_batch, tf.cast(adj_matrix / r_wall_matrix, tf.float32)) \
                                   - tf.math.multiply(state_batch, tf.cast(tf.math.reduce_sum(adj_matrix / r_wall_matrix, 1), tf.float32))) \
                                + adj_out * temp_out_mean / r_outside \
                                - adj_out * state_batch / r_outside \
                                + adj_hall * temp_hall_mean / r_hall \
                                - adj_hall * state_batch / r_hall

            ndim = state_batch.get_shape().as_list()[1]
            diffMatrix = next_state_batch - nextState_Mean

            prob_nextState = tf.math.exp(-0.5 * tf.reduce_sum(tf.matmul(diffMatrix, tf.cast(tf.linalg.inv(gaussian_cov), tf.float32)) * diffMatrix, 1)) / tf.math.sqrt((2 * math.pi)**ndim * tf.cast(tf.linalg.det(gaussian_cov), tf.float32))

            prob_nextState /= max_pdf # Divide by maximum possible pdf (normalize to below 1)
            prob_nextState += 1e-12 # Small value tolerance (avoid pdf being 0 due to limited accuracy)

            #next_state_actions = tf.dtypes.cast(actor_model(next_state_batch, training = True), tf.float64)
            #next_state_actions = tf.dtypes.cast(target_actor(next_state_batch, training = True), tf.float64)
            V_currState = mm_qsa([state_batch, actions], training=True)
            next_state_actions = mm_lag_actor(next_state_batch, training = True)
            V_nextState = mm_qsa([next_state_batch, next_state_actions], training=True)

            out_of_range_bool = (state_batch < temp_low_vec) | (state_batch > temp_up_vec)
            reward_sa = -tf.math.reduce_sum(is_room * (actions * cost_air_var \
                                                       + tf.cast(out_of_range_bool, tf.float32) * (penalty_var) \
                                                       + 10.0 * tf.math.abs((temp_up_vec + temp_low_vec) / 2.0 - state_batch)),
                                            axis = 1)

            #actor_loss = -tf.math.reduce_mean(gamma * V_nextState * tf.math.log(prob_nextState))

            actor_loss = -tf.math.reduce_mean(reward_sa + gamma * (V_nextState - V_currState) * tf.math.log(prob_nextState))

        actor_grad = tape.gradient(actor_loss, mm_actor.trainable_variables)
        mm_actor_optimizer.apply_gradients(
            zip(actor_grad, mm_actor.trainable_variables)
        )

        actor_grad_list = []
        for grad in actor_grad:
            actor_grad_list.append(tf.math.reduce_mean(tf.abs(grad)))

        return actor_loss, tf.math.reduce_mean(actor_grad_list)

    # For Actor buffer
    def learn(self):

        obses_t, actions, rewards, obses_tp1, dones = self.buffer.sample(self.batch_size)

        rewards = tf.expand_dims(rewards, 1)

        update_metrics = self.update(obses_t, actions, rewards, obses_tp1)

        return update_metrics

# Actor Network Architecture
def get_actor():
    # Initialize weights between -3e-3 and 3-e3
    last_init = tf.random_uniform_initializer(minval=-3e-3, maxval=3e-3)

    inputs = layers.Input(shape=(num_states,))
    out = layers.Dense(2048, activation = "relu")(inputs)
    out = layers.LayerNormalization()(out)
    outputs = layers.Dense(num_actions, activation="sigmoid",
                           kernel_initializer=last_init)(out)

    outputs = outputs * air_max_vec
    model = tf.keras.Model(inputs, outputs)
    return model

# Qsa Network Architecture
def get_qsa():

    last_init = tf.random_uniform_initializer(minval=-3e-3, maxval=3e-3)

    # State as input
    state_input = layers.Input(shape=(num_states))
    state_out = layers.Dense(32, activation = "relu")(state_input)
    state_out = layers.LayerNormalization()(state_out)

    # Action as input
    action_input = layers.Input(shape=(num_actions))
    action_out = layers.Dense(32, activation = "relu")(action_input)
    action_out = layers.LayerNormalization()(action_out)

    # Both are passed through seperate layer before concatenating
    concat = layers.Concatenate()([state_out, action_out])

    out = layers.Dense(2048, activation = "relu")(concat)
    out = layers.LayerNormalization()(out)
    outputs = layers.Dense(1,
                           activation="relu",
                           kernel_initializer=last_init)(out)

    # Try RELU output to make it +ve (QMix)
    outputs = outputs * -1.0
    # Outputs single value for give state-action
    model = tf.keras.Model([state_input, action_input], outputs)

    return model

# Policy Function
def policy(actor_model, state, noise_object, t):
    sampled_actions = tf.squeeze(actor_model(state))
    noise = noise_object()
    # Adding noise to action
    sampled_actions = sampled_actions.numpy() + (noise * air_max_vec)

    # We make sure action is within bounds
    legal_action = np.clip(sampled_actions, 0, air_max_vec)

    return np.squeeze(legal_action)

# Calculate Cumulative Discounted Rewards
def calcDiscRewards(traj_rewards, gamma):
    i, total_reward = 0, 0
    for r in traj_rewards:
        total_reward += ((gamma**i) * r)
        i += 1
    return total_reward

# This updates target parameters slowly
# Based on rate `tau`, which is much less than one.
@tf.function
def update_target(target_weights, weights, tau):
    for (a, b) in zip(target_weights, weights):
        a.assign(b * tau + a * (1 - tau))

@tf.function
def update_lag_actor(lag_actor_weights, weights):
    for (a, b) in zip(lag_actor_weights, weights):
        a.assign(b)

def set_start_state(rddl_env, start_state):

    # Initialize Environment
    state, t = rddl_env.reset()

    # Set Start State
    state['temp/1'] = np.array(start_state)
    env._state['temp/1'] = state['temp/1']

    return rddl_env, state, t

def simulate_test(rddl_env, start_state, total_episodes, seed):
    # list of trajectories
    listTrajectory = []
    listTestTrajectories = []
    disc_rewards_arr, mean_qsa_loss_arr, mean_actor_loss_arr, mean_qsa_grad_arr, mean_actor_grad_arr = [], [], [], [], []

    beta_schedule = LinearSchedule(total_episodes * horizon,
                                   initial_p = 1.0,
                                   final_p = 1.0)

    noise_schedule = LinearSchedule(total_episodes * horizon,
                                    initial_p = 0.2,
                                    final_p = 0.0)

    t_iter, best_mean_undisc_reward = 0, float('-inf')

    for ep in range(total_episodes):

        # initialize environament
        if ep == 0:
            rddl_env, state, t = set_start_state(rddl_env, start_state)
        else:
            rddl_env, state, t = set_start_state(rddl_env, np.random.uniform([0] * num_states, [30] * num_states))
        done = False

        # create a trajectory container
        trajectory = rddlgym.Trajectory(rddl_env)
        qsa_loss_arr, actor_loss_arr, qsa_grad_arr, actor_grad_arr = [], [], [], []

        while not done:

            curr_state = state['temp/1'].astype(np.float32)

            tf_state = tf.expand_dims(tf.convert_to_tensor(curr_state), 0)

            action = OrderedDict({'air/1': policy(mm_actor,
                                                  tf_state,
                                                  GaussActionNoise(mean = 0,
                                                                   std_deviation = noise_schedule.value(t_iter),
                                                                   dim = num_actions),
                                                  t)})

            next_state, reward, done, info = rddl_env.step(action)
            reward, nextState = reward.astype(np.float32), next_state['temp/1'].astype(np.float32)

            # Reward scaling for HVAC-6, for training only
            scaled_reward = reward / np.abs(penalty_var) * 10.0
            scaled_reward = scaled_reward.astype(np.float32)

            q_buffer.record(curr_state, action['air/1'].astype(np.float32), scaled_reward, nextState, done)
            actor_buffer.record(curr_state, action['air/1'].astype(np.float32), scaled_reward, nextState, done)
            if len(q_buffer) > q_buffer.batch_size:
                td_errors, qsa_loss, ave_qsa_grad = q_buffer.learn(beta = beta_schedule.value(t_iter))
                qsa_loss_arr.append(qsa_loss)
                qsa_grad_arr.append(ave_qsa_grad)
                update_target(mm_target_qsa.variables, mm_qsa.variables, tau)
            if (len(actor_buffer) > actor_buffer.batch_size):
                actor_loss, ave_actor_grad = actor_buffer.learn()
                actor_loss_arr.append(actor_loss)
                actor_grad_arr.append(ave_actor_grad)
                update_target(mm_target_actor.variables, mm_actor.variables, tau)
                update_lag_actor(mm_lag_actor.variables, mm_actor.variables)

            trajectory.add_transition(t, state, action, reward, next_state, info, done)

            state = next_state
            t = rddl_env.timestep

            t_iter += 1

        disc_rewards = calcDiscRewards(trajectory.as_dataframe().reward, gamma)
        disc_rewards_arr.append(disc_rewards)

        if len(qsa_loss_arr) == 0:
            mean_qsa_loss = None
            mean_qsa_loss_arr.append(float('nan'))
        else:
            mean_qsa_loss = np.mean(qsa_loss_arr)
            mean_qsa_loss_arr.append(mean_qsa_loss)

        if len(actor_loss_arr) == 0:
            mean_actor_loss = None
            mean_actor_loss_arr.append(float('nan'))
        else:
            mean_actor_loss = np.mean(actor_loss_arr)
            mean_actor_loss_arr.append(mean_actor_loss)

        if len(qsa_grad_arr) == 0:
            mean_qsa_grad = None
            mean_qsa_grad_arr.append(float('nan'))
        else:
            mean_qsa_grad = np.mean(qsa_grad_arr)
            mean_qsa_grad_arr.append(mean_qsa_grad)

        if len(actor_grad_arr) == 0:
            mean_actor_grad = None
            mean_actor_grad_arr.append(float('nan'))
        else:
            mean_actor_grad = np.mean(actor_grad_arr)
            mean_actor_grad_arr.append(mean_actor_grad)

        print("Episode * {} * Total Reward is ==> {}".format(ep, disc_rewards))
        print("Qsa loss: {}".format(mean_qsa_loss))
        print("Actor loss: {}".format(mean_actor_loss))
        print("Average Qsa gradient: {}".format(mean_qsa_grad))
        print("Average actor gradient: {}".format(mean_actor_grad))
        print()

        listTrajectory.append(trajectory.as_dataframe())

        if (ep + 1) % test_interval == 0:

            l_test_trajs, mean_disc_r, mean_undisc_r = test_actor_loop(folderName + '/' + 'mm_test_log_' + str(ep + 1) + '.csv', env, start_state)
            listTestTrajectories.append(l_test_trajs)

            if mean_undisc_r > best_mean_undisc_reward:

                best_mm_actor.set_weights(mm_actor.get_weights())
                best_mm_qsa.set_weights(mm_qsa.get_weights())

                best_mean_undisc_reward = mean_undisc_r

    return disc_rewards_arr, mean_qsa_loss_arr, mean_actor_loss_arr, mean_qsa_grad_arr, mean_actor_grad_arr, listTrajectory, listTestTrajectories

def test_actor_loop(filename, rddl_env, start_state):

    list_traj_df, list_disc_reward, list_undisc_reward = [], [], []

    for i in range(test_loops):

        # initialize environament
        rddl_env, state, t = set_start_state(rddl_env, start_state)
        done = False

        test_trajectory = rddlgym.Trajectory(rddl_env)

        while not done:

            curr_state = state['temp/1'].astype(np.float32)

            tf_state = tf.expand_dims(tf.convert_to_tensor(curr_state), 0)

            action = OrderedDict({'air/1': policy(mm_actor, tf_state, lambda : np.array([0] * num_actions), t)})
            next_state, reward, done, info = rddl_env.step(action)

            test_trajectory.add_transition(t, state, action, reward, next_state, info, done)

            state = next_state
            t = rddl_env.timestep

        test_log_df = test_trajectory.as_dataframe()
        disc_reward = calcDiscRewards(test_log_df.reward, gamma)
        test_log_df['Total Discounted Rewards'] = [disc_reward for i in range(test_log_df.shape[0])]
        undisc_reward = calcDiscRewards(test_log_df.reward, 1.0)
        test_log_df['Total Undiscounted Rewards'] = [undisc_reward for i in range(test_log_df.shape[0])]

        list_traj_df.append(test_log_df)
        list_disc_reward.append(disc_reward)
        list_undisc_reward.append(undisc_reward)

    return list_traj_df, np.mean(list_disc_reward), np.mean(list_undisc_reward)

def log_learn(folderName, lDiscRewards, lQsaLoss, lActorLoss, lQsaGrad, lActorGrad):

    learn_log_df = pd.DataFrame({'Episode': [i for i in range(len(lDiscRewards))],
                                 'Discounted Rewards': lDiscRewards,
                                 'Qsa Loss': lQsaLoss,
                                 'Actor Loss': lActorLoss,
                                 'Qsa Gradient': lQsaGrad,
                                 'Actor Gradient': lActorGrad})

    learn_log_df.to_csv(folderName + 'learn_log.csv', index = False)

def log_trajectories(folderName, lTrainTraj, lListTestTraj):

    for i in range(len(lTrainTraj)):
        lTrainTraj[i].to_csv(folderName + 'E' + str(i + 1) + '.csv', index = False)

    testTrajFolder = folderName + 'test_trajs/'
    for i in range(len(lListTestTraj)):
        testTraj_subFolder = testTrajFolder + 'E' + str((i + 1) * test_interval) + '/'
        pathlib.Path(testTraj_subFolder).mkdir(parents = True, exist_ok = True)
        for j in range(len(lListTestTraj[i])):
            lListTestTraj[i][j].to_csv(testTraj_subFolder + str(j + 1) + '.csv', index = False)

def plot_graphs(dirName, lQsaLoss, lActorLoss, lQsaGrad, lActorGrad, lTrajList):

    numEpisodes = len(lQsaLoss)
    plt.figure()
    plt.plot(range(numEpisodes), lQsaLoss)
    plt.xlabel("Episode")
    plt.ylabel("Average Qsa Loss Across Minibatches")
    plt.savefig(dirName + 'qsa_loss.png')
    plt.close()

    # Plot mean_actor_loss_arr
    plt.figure()
    plt.plot(range(numEpisodes), lActorLoss)
    plt.xlabel("Episode")
    plt.ylabel("Average Actor Loss Across Minibatches")
    plt.savefig(dirName + 'actor_loss.png')
    plt.close()

    # Plot mean_qsa_grad_arr
    plt.figure()
    plt.plot(range(numEpisodes), lQsaGrad)
    plt.xlabel("Episode")
    plt.ylabel("Average Qsa Gradient Across Minibatches")
    plt.savefig(dirName + 'qsa_grad.png')
    plt.close()

    # Plot mean_qsa_loss_arr, mean_actor_loss_arr
    plt.figure()
    plt.plot(range(numEpisodes), lActorGrad)
    plt.xlabel("Episode")
    plt.ylabel("Average Actor Gradient Across Minibatches")
    plt.savefig(dirName + 'actor_grad.png')
    plt.close()

    ave_disc_rewards = []
    ave_undisc_rewards = []
    for TrajList in lTrajList:
        ave_disc_rewards.append(np.mean([df['Total Discounted Rewards'][0] for df in TrajList]))
        ave_undisc_rewards.append(np.mean([df['Total Undiscounted Rewards'][0] for df in TrajList]))

    testEpisode_range = range(test_interval, len(ave_undisc_rewards) * test_interval + 1, test_interval)
    # Plot ave_disc_rewards, ave_undisc_rewards
    plt.figure()
    plt.plot(testEpisode_range, ave_disc_rewards)
    plt.xlabel("Sample Episodes Collected")
    plt.ylabel("Accumulated Discounted Rewards")
    plt.savefig(dirName + 'test_disc_rewards.png')
    #plt.show()
    plt.close()

    # Plot ave_disc_rewards, ave_undisc_rewards
    plt.figure()
    plt.plot(testEpisode_range, ave_undisc_rewards)
    plt.xlabel("Sample Episodes Collected")
    plt.ylabel("Accumulated Undiscounted Rewards")
    plt.savefig(dirName + 'test_undisc_rewards.png')
    #plt.show()
    plt.close()

    ave_reward_df = pd.DataFrame(data = {'Episode': testEpisode_range,
                                         'Average Discounted Reward': ave_disc_rewards,
                                         'Average Undiscounted Reward': ave_undisc_rewards})

    ave_reward_df.to_csv(dirName + 'average_test_reward.csv', index = False)

def init_networks(qsa_lr, actor_lr):

    actor_model = get_actor()
    qsa_model = get_qsa()
    target_actor = get_actor()
    target_qsa = get_qsa()

    # Making the weights equal initially
    target_actor.set_weights(actor_model.get_weights())
    target_qsa.set_weights(qsa_model.get_weights())

    qsa_optimizer = tf.keras.optimizers.RMSprop(qsa_lr)
    actor_optimizer = tf.keras.optimizers.RMSprop(actor_lr)

    lag_actor = get_actor()
    lag_actor.set_weights(actor_model.get_weights())
    return actor_model, qsa_model, target_actor, target_qsa, lag_actor, qsa_optimizer, actor_optimizer

# Start States
with open('./start_state.txt') as f:
    startState_list = [tuple(map(float, i.split(','))) for i in f]

# Seeds
with open('./seed.txt') as f:
    seed_list = [int(i) for i in f]

import rddlgym

# create RDDLGYM environment
rddl_id = "HVAC-6" # see available RDDL domains/instances with `rddlgym ls` command
env = rddlgym.make(rddl_id, mode = rddlgym.GYM)

# you can also wrap your own RDDL files (domain + instance)
# env = rddlgym.make("/path/to/your/domain_instance.rddl", mode=rddlgym.GYM)

num_states = env.observation_space['temp/1'].shape[0]
print("Size of State Space ->  {}".format(num_states))
num_actions = env.action_space['air/1'].shape[0]
print("Size of Action Space ->  {}".format(num_actions))

upper_bound = env.action_space['air/1'].high[0]
lower_bound = env.action_space['air/1'].low[0]

#upper_bound = 1.0
#lower_bound = -1.0

print("Max Value of Action ->  {}".format(upper_bound))
print("Min Value of Action ->  {}".format(lower_bound))

adj_matrix = env.non_fluents['ADJ/2']

for i in range(adj_matrix.shape[0]):
    for j in range(i + 1, adj_matrix.shape[0]):
        adj_matrix[j, i] = adj_matrix[i, j]

#adj_out = [1, 0, 1, 1, 0, 1]
#adj_hall = [1, 1, 1, 1, 1, 1]
adj_out = env.non_fluents['ADJ_OUTSIDE/1']
adj_hall = env.non_fluents['ADJ_HALL/1']
time_delta = env.non_fluents['TIME_DELTA/0']
cap = env.non_fluents['CAP/1']
cap_air = env.non_fluents['CAP_AIR/0']
is_room = env.non_fluents['IS_ROOM/1']
r_wall_matrix = env.non_fluents['R_WALL/2']
r_outside = env.non_fluents['R_OUTSIDE/1']
r_hall = env.non_fluents['R_HALL/1']
temp_air = env.non_fluents['TEMP_AIR/0']
temp_out_mean = env.non_fluents['TEMP_OUTSIDE_MEAN/1']
temp_out_var = env.non_fluents['TEMP_OUTSIDE_VARIANCE/1']
temp_hall_mean = env.non_fluents['TEMP_HALL_MEAN/1']
temp_hall_var = env.non_fluents['TEMP_HALL_VARIANCE/1']
cost_air_var = env.non_fluents['COST_AIR/0']
temp_up_vec = env.non_fluents['TEMP_UP/1']
temp_low_vec = env.non_fluents['TEMP_LOW/1']
penalty_var = env.non_fluents['PENALTY/0']
air_max_vec = env.non_fluents['AIR_MAX/1']
horizon = env.horizon

#gaussian_sd = 1e-1
#gaussian_sd = 0.05**0.5
gaussian_cov = np.diag(adj_out * temp_out_var / r_outside**2 + adj_hall * temp_hall_var / r_hall**2)
gaussian_zero_mean = np.array([0] * num_states)
max_pdf = scipy.stats.multivariate_normal.pdf(gaussian_zero_mean, gaussian_zero_mean, gaussian_cov)

# Learning rate for Qsa and DRP
mm_qsa_lr = 0.001
mm_actor_lr = 0.0001

total_episodes = 5000
# Discount factor for future rewards
gamma = 0.99
# Used to update target networks
tau = 0.005

test_interval = 5
test_loops = 64

for startState_idx in range(len(startState_list)):

    for seed in seed_list:

        folderName = './' + rddl_id + '_Results/MM/SIdx' + str(startState_idx) + '/Seed' + str(seed) + '/'
        pathlib.Path(folderName).mkdir(parents = True, exist_ok = True)

        # Initialize for MM
        tf.set_random_seed(seed)
        mm_actor, mm_qsa, mm_target_actor, mm_target_qsa, mm_lag_actor, mm_qsa_optimizer, mm_actor_optimizer = init_networks(mm_qsa_lr, mm_actor_lr)
        best_mm_qsa = get_qsa()
        best_mm_qsa.set_weights(mm_qsa.get_weights())
        best_mm_actor = get_actor()
        best_mm_actor.set_weights(mm_actor.get_weights())
        q_buffer = QsaBuffer(1000000, 64, alpha = 0.2)
        actor_buffer = ActorBuffer(1000, 64)

        print("MM")
        # MM Second
        discRewards_l, qsaLoss_l, actorLoss_l, qsaGrad_l, actorGrad_l, MMTrainTraj_l, listMMTestTraj_l = simulate_test(env, startState_list[startState_idx], total_episodes, seed)

        mm_qsa.save_weights(folderName + 'checkpoints/mm_qsa')
        mm_actor.save_weights(folderName + 'checkpoints/mm_actor')
        best_mm_qsa.save_weights(folderName + 'checkpoints/best_mm_qsa')
        best_mm_actor.save_weights(folderName + 'checkpoints/best_mm_actor')

        log_learn(folderName, discRewards_l, qsaLoss_l, actorLoss_l, qsaGrad_l, actorGrad_l)
        log_trajectories(folderName, MMTrainTraj_l, listMMTestTraj_l)

        plot_graphs(folderName, qsaLoss_l, actorLoss_l, qsaGrad_l, actorGrad_l, listMMTestTraj_l)

max_stars_repo_name: siowmeng/ilbo

max_issues_repo_name: siowmeng/ilbo

max_forks_repo_name: siowmeng/ilbo

reason: import numpy,import scipy


text:
"""
Vectorized plotting of multiple trajectories at once

2013 (BSD-3) California Institute of Technology
"""
from __future__ import division
import numpy as np

from multidim_plot import text

def plot_trajectory(
    ax, xtraj, x0=np.array([]), xd, x0str, xdstr,
    xtraj_style, x0_style, xd_style, n_subsample
):
    """Plot trajectory, initial condition and desired destination.

     input
       ax = axes object handle
       xtraj = intermediate trajectory points
             = [#dim x #pnts] |
               {1 x #traj} = {[#dim x #pnt1], [#dim x #pnt2], ... }

     optional input
       x0 = initial point(s) [xtraj
          = [#dim x #traj]
       xd = destination point(s) [none]
          = [#dim x #traj]
       x0str = initial condition text annotation ['$x_0$']
             = string
       xdstr = destination text annotation ['$x_0$']
             = string
       xtraj_style = trajectory style for plot function, for example:
                   = {'g-o'}
                   = {'Color', 'g', 'Marker', 'o', 'LineStyle', '-'}
       x0_style = initial point style for plot function, for example:
                = {'rs'}
                = {'Color', 'r', 'Marker', 's'}
       xd_style = destination point style for plot function, for example:
                = {'go'}
                = {'Color', 'g', 'Marker', 'o'}
       n_subsample = use at most this number of points per trajectory
                     (economizes on file size)
                   >0 (default=100) | =0 (disable)

     caution
       if subsampling, all trajectories should have same number of points
    """
    #depends
    #   plot, text, takehold, restorehold, subsample
    # todo
    #   update code of: nf_spline_plot_results, nf_spline_plot_results_md, plotq0qsqd
    #   to use this for multiple trajectories

    # input
    # single traj ?
    if not  iscell(xtraj):
        xtraj = [xtraj]
    if nargin < 3:
        x0 = np.array([])
        # will be extracted later based on this "flag"
    else:
        if iscell(x0):
            # temporary compatibility check
            error('update your code, plot_trajectory changed order of args 2,3')
    if (x0.shape[1] > 1) and (max(xtraj.shape) == 1):
        error('Multiple x0, one xtraj, check order of args 2,3.')
    if nargin < 4:
        xd = np.array([])
    if nargin < 5:
        x0str = '$x_0$'
    if nargin < 6:
        xdstr = '$x_d$'
    if nargin < 7:
        xtraj_style = ['Color', 'b', 'LineWidth', 2]
    if nargin < 8:
        x0_style = ['Color', 'r', 'Marker', 's', 'LineStyle', 'none']
    if not  ismember('LineStyle', x0_style):
        x0_style = np.array([x0_style, ['LineStyle', 'none']]).reshape(1, -1)
    if nargin < 9:
        xd_style = ['Color', 'g', 'Marker', 'o', 'LineStyle', 'none']
    if not  ismember('LineStyle', xd_style):
        xd_style = np.array([xd_style, ['LineStyle', 'none']]).reshape(1, -1)
    if nargin < 10:
        ddisp('No subsampling for plot_trajectory, using: n = 100.')
        n_subsample = 100
    held = takehold(ax)
    #
    # data
    # subsample
    npnt = xtraj[0, 0].shape[1]
    if (npnt > n_subsample) and (n_subsample != 0):
        xtraj = cellfun(subsample, xtraj, 'UniformOutput', false())
    # concatenate multiple lines, separated by NaN columns
    ndim = xtraj[0, 0].shape[0]
    sep = nan(ndim, 1)
    xt = cellfun(lambda x: np.array([x, sep]).reshape(1, -1), xtraj, 'UniformOutput', false())
    xt = cell2mat(xt)
    # get actual start-points (if none provided)
    if (0 in x0.shape):
        x0 = cellfun(lambda x: x[:, 0], xtraj, 'UniformOutput', false())
        x0 = cell2mat(x0)
    # get actual end-points
    xfinal = cellfun(lambda x: x[:, (end -1)], xtraj, 'UniformOutput', false())
    xfinal = cell2mat(xfinal)
    #
    # plot
    plot(ax, xt, xtraj_style[:])
    # trajectory
    # note
    #   no HandleVisibility off needed, since concatenated into single line object
    plot(ax, x0, x0_style[:])
    # initial condition
    plot(ax, xfinal, 'Color', 'b', 'Marker', 'o', 'LineStyle', 'None', 'HandleVisibility', 'off')
    # actual final point
    # no destination ?
    if not  (0 in xd.shape):
        plot(ax, xd, xd_style[:])
        # desired destination
    # annotate only single initial condition and desired destination
    if not  (0 in xd.shape):
        annot_x0_xd(ax, x0[:, 0], x0str, xd[:, 0], xdstr)
    else:
        annot_x0_xd(ax, x0[:, 0], x0str, np.array([]), xdstr)
    # not nice result to annotate all initial conditions
    #annot_x0_xd(ax, x0, x0str, xd, xdstr)
    restorehold(ax, held)

def annot_x0_xd(ax, x0, x0str, xd, xdstr):
    # annotate initial conditions
    n0 = x0.shape[1]
    for i in range(1, (n0 +1)):
        curx0 = x0[:, (i -1)]
        text(1.1 * curx0, x0str, 'Interpreter', 'Latex', 'FontSize', 15, 'Parent', ax)
    # no destination ?
    if (0 in xd.shape):
        return
    # annotate desired destinations
    nd = xd.shape[1]
    for i in range(1, (nd +1)):
        curxd = xd[:, (i -1)]
        text(1.1 * curxd, xdstr, 'Interpreter', 'Latex', 'FontSize', 15, 'Parent', ax)
    nnotate('desired', 'destinations')
    nd = xd.shape[1]
    for i in range(1, (nd +1)):
        curxd = xd[:, (i -1)]
        text(1.1 * curxd, xdstr, 'Interpreter', 'Latex', 'FontSize', 15, 'Parent', ax)
    nnotate('desired', 'destinations')
    nd = xd.shape[1]
    for i in range(1, (nd +1)):
        curxd = xd[:, (i -1)]
        text(1.1 * curxd, xdstr, 'Interpreter', 'Latex', 'FontSize', 15, 'Parent', ax)
    nd
    nnotate('desired', 'destinations')
    nd = xd.shape[1]
    for i in range(1, (nd +1)):
        curxd = xd[:, (i -1)]
        text(1.1 * curxd, xdstr, 'Interpreter', 'Latex', 'FontSize', 15, 'Parent', ax)
    return

def plot_path(
    ax, x0, xtraj, xd, x0str='',
    xdstr='', zoff=None, zonsurf=None, Rs=None,
    xtraj_style=None, x0_style=None,
    xd_style=None,
    xtraj_onsurf_style=None
):
    """Plot trajectory on function surface.

     input
       ax = axes object handle
       x0 = initial condition
       xtraj = path subsequent points
       xd = agent destination
       x0str = initial condition text annotation
             = string
       xdstr = destination text annotation
             = string
       zoff = the contour plot level z offset
            = real | []
       zonsurf = navigation function values at
                   corresponding points
               = [1 x size(xtraj, 2) ] | []
       Rs = sensing radius
          > 0 | []
    """
    # depends
    #   plot, zoffset, plot_trajectory,
    #   plot_circle, plotSphere

    if xtraj_style is None:
        xtraj_style = ['Linewidth', 2, 'Color', 'r', 'Linestyle', '-']
    elif (0 in xtraj_style.shape):
            xtraj_style = ['Linewidth', 2, 'Color', 'r', 'Linestyle', '-']

    if x0_style is None:
        x0_style = ['Color', 'r', 'Marker', 's']
    elif (0 in x0_style.shape):
            x0_style = ['Color', 'r', 'Marker', 's']

    if nargin < 12:
        xd_style = ['Color', 'g', 'Marker', 'o']
    elif (0 in xd_style.shape):
            xd_style = ['Color', 'g', 'Marker', 'o']

    if xtraj_onsurf_style is None:
        xtraj_onsurf_style = ['Linewidth', 3, 'Color', 'm', 'Linestyle', '-']
    elif (0 in xtraj_onsurf_style.shape):
            xtraj_onsurf_style = ['Linewidth', 3, 'Color', 'm', 'Linestyle', '-']

    ndim = x0.shape[0]

    # plot
    held = takehold(ax)
    if ndim == 2:
        # vertical offset ?
        if not  (0 in zoff.shape):
            x0 = zoffset(x0, zoff)
            xtraj = zoffset(xtraj, zoff)
            xd = zoffset(xd, zoff)
        # 3D path (imagined on field surface)
        if not  (0 in zonsurf.shape):
            xonsurf = zoffset(xtraj, zonsurf)
            plot(ax, xonsurf, xtraj_onsurf_style[:])
    plot_trajectory(ax, xtraj, x0, xd, x0str, xdstr, xtraj_style, x0_style, xd_style)

    # sensing on ?
    if not  (0 in Rs.shape):
        if ndim == 2:
            plot_circle(ax, xtraj[:, (end -1)], Rs, 'r', 'Color', 'm', 'LineStyle', '--', 'CenterStyle', 'mo')
        else:
            if ndim == 3:
                plotSphere(ax, xtraj[:, (end -1)], Rs, 'Color', 'r', 'Opacity', 0)
            else:
                error('Path plot only available for 2 and 3 dimensions.')
    restorehold(ax, held)

def qtraj_mat2cell(qtraj):
    """Convert 3d matrix of trajectories to row cell array of 2d matrices

     input
       qtraj = 3d matrix of trajectories
             = [#dim x #time_iterations x #trajectories]

     output
       qtraj = cell row array of 2d matrices of column vectors of points in
               each trajectory
             = {1 x #trajectories} = {[#dim x #time_iterations], ... }

     note
       input qtraj is a matrix, usually produced by a parallel integration of
       the flow, so all trajectories stop at the same number of iterations.
       Otherwise nan values could be used to signify the (different) end of
       each trajectory.
    """
    #todo
    #   dev faster alternative for plotting/post-processing directly with 3d
    #   matrices of trajectories
    ndim, niter, ntraj = qtraj.shape # nargout=3
    qtraj = mat2cell(qtraj, ndim, niter, ones(1, ntraj))
    qtraj = squeeze(qtraj)
    qtraj = qtraj.T
    return qtraj

def test_plot_traj():
    """test trajectory plotting code for multiple trajectories.
    """
    ax = plt.gca()
    ndim = 3
    ntraj = 10
    npnt = 1000
    x0 = rand(ndim, ntraj)
    xd = rand(ndim, ntraj)
    xtraj = cell(1, ntraj)
    for i in range(1, (ntraj +1)):
        xtraj[0, (i -1)] = rand(ndim, npnt)
    x0str = '$x_0$'
    xdstr = '$x_d$'
    xtraj_style = ['Color', 'b', 'LineStyle', '--']
    x0_style = ['Color', 'r', 'Marker', 's']
    xd_style = ['Color', 'g', 'Marker', 'o']
    plot_trajectory(ax, x0, xtraj, xd, x0str, xdstr, xtraj_style, x0_style, xd_style)
    plot_scalings(ax, 0)
    grid(ax, 'on')

max_stars_repo_name: johnyf/pyvectorized

max_issues_repo_name: johnyf/pyvectorized

max_forks_repo_name: johnyf/pyvectorized


text:
#!/usr/bin/env python

"""Pseudo Python code for calculating the Kalman filter for learning."""

import matplotlib.pyplot as plt
import numpy.random

class KalmanFilter:
    """
    Kalman filter for finding the true value of data comming iteratively.

    Attributes
    ----------
    estimate : list of floats
    estimate_error : float
    """

    def __init__(self, estimate, estimate_error):
        assert estimate_error >= 0.0  # it can be bigger than 1!
        self.estimate = estimate
        self.estimate_error = estimate_error

    def add_measurement(self, measurement, m_error):
        estimate, error = self.calculate_new_estimate(measurement, m_error)
        self.estimate = estimate
        self.estimate_error = error
        return (self.estimate, self.estimate_error)

    def calculate_new_estimate(self, measurement, m_error):
        # Calculate Kalman gain
        kg = self.get_kalman_gain(m_error)
        new_estimate = self.estimate + kg * (measurement - self.estimate)
        new_error = (1.0 - kg) * self.estimate_error
        assert new_error >= 0.0, "Error was %0.4f" % new_error
        return new_estimate, new_error

    def get_kalman_gain(self, m_error):
        """
        Calculate the Kalman gain.

        The higher the Kalman gain, the more important is the measurement
        compared to the current estimate

        Returns
        -------
        float
            In [0, 1]
        """
        return self.estimate_error / (self.estimate_error + m_error)

if __name__ == '__main__':
    # Generate data
    n = 500
    variance = 4
    measurements = [numpy.random.normal(loc=0.0, scale=variance)
                    for i in range(n)]
    errors = [variance * 3 for i in range(n)]
    estimates = []
    estimates_error = []

    # Run Kalman filter
    kf = KalmanFilter(measurements[0], errors[0])
    estimates.append(measurements[0])
    estimates_error.append(errors[0])
    for measurement, error in zip(measurements[1:], errors[1:]):
        new_est, err = kf.add_measurement(measurement, error)
        estimates.append(new_est)
        estimates_error.append(err)

    # Visualize
    plt.plot([0, len(measurements)], [0, 0], 'k-', lw=2)  # 0-line
    me, = plt.plot(range(len(measurements)), measurements, 'ro')
    es, = plt.plot(range(len(measurements)), estimates, '.b-')
    er, = plt.plot(range(len(measurements)), estimates_error, '.g-')
    plt.axis([0, len(measurements), -15, 15])
    plt.legend([me, es, er], ['Measurement', 'Estimate', 'Error estimate'])
    plt.show()

max_issues_repo_name: Kerwin-Xie/algorithms

max_forks_repo_name: Kerwin-Xie/algorithms


text:
import pyasdf
import numpy as np
import scipy.fftpack
import matplotlib.pyplot as plt

'''
this script takes a chunk of noise spectrum for a station pair and
compare their cross-correlation functions computed using two schemes:
one is averaging the frequency domain and the other is in the time
domain
'''

def cross_correlation1(fft1,fft2,maxlag,dt,Nfft):
    #------convert all 2D arrays into 1D to speed up--------
    corr = np.zeros(fft1.shape,dtype=np.complex64)
    corr = np.conj(fft1) * fft2

    ncorr = np.zeros((fft1.shape[0],Nfft),dtype=np.complex64)
    ncorr[:,:Nfft//2] = corr[:,:]
    ncorr[:,-(Nfft//2)+1:]=np.flip(np.conj(ncorr[:,1:(Nfft//2)]),axis=1)
    ncorr[:,0]=complex(0,0)
    ncorr = np.real(np.fft.ifftshift(scipy.fftpack.ifft(ncorr, Nfft, axis=1)))

    tcorr = np.arange(-Nfft//2 + 1, Nfft//2)*dt
    ind   = np.where(np.abs(tcorr) <= maxlag)[0]
    ncorr = ncorr[:,ind]
    ncorr = np.mean(ncorr,axis=0)
    return ncorr

def cross_correlation2(fft1,fft2,maxlag,dt,Nfft):
    #------convert all 2D arrays into 1D to speed up--------
    corr = np.zeros(fft1.shape,dtype=np.complex64)
    corr = np.conj(fft1) * fft2

    ncorr = np.zeros(shape=Nfft,dtype=np.complex64)
    ncorr[:Nfft//2] = np.mean(corr,axis=0)
    ncorr[-(Nfft//2)+1:]=np.flip(np.conj(ncorr[1:(Nfft//2)]),axis=0)
    ncorr[0]=complex(0,0)
    ncorr = np.fft.ifftshift(scipy.fftpack.ifft(ncorr, Nfft, axis=0))
    print(ncorr.real,ncorr.imag)

    tcorr = np.arange(-Nfft//2 + 1, Nfft//2)*dt
    ind   = np.where(np.abs(tcorr) <= maxlag)[0]
    ncorr = ncorr[ind]
    return ncorr

#-----common parameters------
iday   = '2010_01_10'
icomp  = 'EHZ'
dt     = 0.05
maxlag = 800

sfile1 = '/Users/chengxin/Documents/Harvard/Kanto_basin/code/KANTO/FFT/N.AC2H.h5'
sfile2 = '/Users/chengxin/Documents/Harvard/Kanto_basin/code/KANTO/FFT/N.CHHH.h5'

#-----------reading the data------------
ds1 = pyasdf.ASDFDataSet(sfile1,mode='r')
ds2 = pyasdf.ASDFDataSet(sfile2,mode='r')

spect1 = ds1.auxiliary_data[icomp][iday].data[:]
spect2 = ds2.auxiliary_data[icomp][iday].data[:]
std1 = ds1.auxiliary_data[icomp][iday].parameters['std']
std2 = ds2.auxiliary_data[icomp][iday].parameters['std']
nwin = spect1.shape[0]
nfft = spect1.shape[1]*2

print('data dimension for spect1 and spect2 are %d and %d' % (spect1.ndim,spect2.ndim))

#------select the sections-------
indx1 = np.where(std1<10)[0]
indx2 = np.where(std2<10)[0]
bb=np.intersect1d(indx1,indx2)
print(spect1[bb,:],spect2[bb,:])

corr1=cross_correlation1(spect1[bb,:],spect2[bb,:],np.round(maxlag),dt,nfft)
corr2=cross_correlation2(spect1[bb,:],spect2[bb,:],np.round(maxlag),dt,nfft)

#---plotting----
plt.subplot(311)
plt.plot(corr1)
plt.subplot(312)
plt.plot(corr2)
plt.subplot(313)
plt.plot(corr2)
plt.plot(corr1)
plt.show()

max_stars_repo_name: kjdavidson/NoisePy

max_issues_repo_name: kjdavidson/NoisePy

max_forks_repo_name: kjdavidson/NoisePy


text:
import networkx as nx
import matplotlib.pyplot as plt

#Creating a networkX digraph
G=nx.DiGraph()
#Adding nodes of int, str and tuple type
G.add_node(1)
G.add_node("string1")
G.add_node((6,2))
G.add_node('string2')
G.add_nodes_from([2,3,5])
G.add_nodes_from([(4,{'color':'red'})])
#Adding edges between the nodes
G.add_edge(4,1)
G.add_edges_from([(1,2),(1,3),(1,'string1'),(2,3),('string2',1),(1,'string2'),(1,(6,2)),(2,5),((6,2),'string2'),('string2',4)])
#Calculating shortest path between all nodes in the graph
p=nx.shortest_path(G)
print('Printing shortest path')
#Storing existing nodes in the graph
G_node_list = G.nodes
for index,node in enumerate(G_node_list):
    print('From',node)
    for index2,node2 in enumerate(G_node_list):
        if index!=index2 and nx.has_path(G, node, node2):
            print('To', node2, ':', p[node][node2])
#Plotting the graph
nx.draw(G,pos=nx.circular_layout(G),with_labels=True)
plt.show()

reason: import networkx


text:
from phase_diagram.phase_diagram import PhaseDiagram
from src.helpers import compound_index, state_index
import pint
import numpy as np

def test_water():
    water_data = PhaseDiagram('water')
    assert isinstance(water_data, PhaseDiagram)
    assert water_data.idx == 1
    assert water_data.name == 'water'
    assert water_data.cas == '7732-18-5'
    assert water_data.formula == 'H2O'
    assert water_data.molar_mass == pint.Quantity(18.0153, 'gram/mole')
    assert water_data.alternative_names == (None, None, None)
    assert isinstance(water_data.density_solid, pint.Quantity)
    assert water_data.density_solid == pint.Quantity(0.9167, 'gram/cm**3')
    assert water_data.density_liquid == pint.Quantity(0.9970474, 'gram/cm**3')
    assert water_data.antoine == (0.01, 373.98, 8.05573, 1723.6425, 233.08)
    assert water_data.boiling_point == (pint.Quantity(373.124, 'kelvin'), pint.Quantity(101325.0, 'pascal'))
    assert water_data.melting_point == (pint.Quantity(273.15, 'kelvin'), pint.Quantity(101325.0, 'pascal'))
    assert water_data.triple_point == (pint.Quantity(273.16, 'kelvin'), pint.Quantity(611.657, 'pascal'))
    assert water_data.critical_point == (pint.Quantity(647.1, 'kelvin'), pint.Quantity(2.206E7, 'pascal'))
    assert water_data.enthalpy_fusion == pint.Quantity(6.009, 'kJ/mole')
    assert water_data.enthalpy_sublimation == pint.Quantity(44.0, 'kJ/mole')
    assert water_data.enthalpy_vaporization == pint.Quantity(40.66, 'kJ/mole')
    assert np.isclose(water_data.volume_change_fusion, pint.Quantity(-1.583, 'cm**3/mole'), atol=1E-3)

def test_water_cas():
    water_data = PhaseDiagram('7732-18-5')
    assert isinstance(water_data, PhaseDiagram)
    assert water_data.idx == 1
    assert water_data.name == 'water'
    assert water_data.cas == '7732-18-5'
    assert water_data.formula == 'H2O'

def test_water_formula():
    water_data = PhaseDiagram('H2O')
    assert isinstance(water_data, PhaseDiagram)
    assert water_data.idx == 1
    assert water_data.name == 'water'
    assert water_data.cas == '7732-18-5'
    assert water_data.formula == 'H2O'

def test_carbon_dioxide():
    carbon_dioxide_data = PhaseDiagram('carbon dioxide')
    assert isinstance(carbon_dioxide_data, PhaseDiagram)
    assert carbon_dioxide_data.idx == 2
    assert carbon_dioxide_data.name == 'carbon dioxide'
    assert carbon_dioxide_data.cas == '124-38-9'
    assert carbon_dioxide_data.formula == 'CO2'
    assert carbon_dioxide_data.molar_mass == pint.Quantity(44.0095, 'gram/mole')
    assert carbon_dioxide_data.alternative_names == ('carbonic anhydride', None, None)
    assert isinstance(carbon_dioxide_data.density_solid, pint.Quantity)
    assert carbon_dioxide_data.density_solid == pint.Quantity(1.562, 'gram/cm**3')
    assert carbon_dioxide_data.density_liquid == pint.Quantity(1.179, 'gram/cm**3')
    assert carbon_dioxide_data.antoine == (-56.57, 31.04, 7.58828, 861.8205, 271.88)
    assert carbon_dioxide_data.boiling_point == (pint.Quantity(194.69, 'kelvin'), pint.Quantity(101325.0, 'pascal'))
    assert carbon_dioxide_data.melting_point == (pint.Quantity(216.59, 'kelvin'), pint.Quantity(101325.0, 'pascal'))
    assert carbon_dioxide_data.triple_point == (pint.Quantity(216.58, 'kelvin'), pint.Quantity(518500.0, 'pascal'))
    assert carbon_dioxide_data.critical_point == (pint.Quantity(304.13, 'kelvin'), pint.Quantity(7.375E6, 'pascal'))
    assert carbon_dioxide_data.enthalpy_fusion == pint.Quantity(9.02, 'kJ/mole')
    assert carbon_dioxide_data.enthalpy_sublimation == pint.Quantity(25.2, 'kJ/mole')
    assert carbon_dioxide_data.enthalpy_vaporization == pint.Quantity(15.8, 'kJ/mole')

def test_carbon_dioxide_alt_name():
    carbon_dioxide_data = PhaseDiagram('carbonic anhydride')
    assert isinstance(carbon_dioxide_data, PhaseDiagram)
    assert carbon_dioxide_data.idx == 2
    assert carbon_dioxide_data.name == 'carbon dioxide'
    assert carbon_dioxide_data.cas == '124-38-9'
    assert carbon_dioxide_data.formula == 'CO2'

def test_state_index():
    assert state_index('solid') == 1
    assert state_index('liquid') == 2

def test_compound_index():
    assert compound_index('water') == 1
    assert compound_index('CO2') == 2
    assert compound_index('iodine') == 3

max_stars_repo_name: fhfraga/PhaseDiagram

max_issues_repo_name: fhfraga/PhaseDiagram

max_forks_repo_name: fhfraga/PhaseDiagram