text:
#!/usr/bin/env python

import argparse
import os
import re
import ephem
import pyfits
import numpy as np
from Utils import Constants
from Utils.Catalog import SExtractorDetection
from Utils.TrackableException import TrackableException as STAP_Error
from Utils.TrackableException import ExternalFailure
from Utils.Photometry import calc_seeing, calc_apcor, calc_elong
from STAP_comm import STAP_callexternal, print_cmd_line

def SEx(imname, seeing=5.0, do_apcor=False, do_seeing=False, rerun=False,
        config="STAP_SEx.sex", SEx_switches=None, timeout=None):
    """
    Runs SExtractor on an image.  For input named imname.fits, stores output
    in FITS binary table format in imname.fits.stars.

    imname:     name of SExtractor input FITS image
    seeing:     default expected FWHM of objects and of detection filter
                (will attempt to pull from SEEING keyword in FITS header)
    rerun:      option to run SExtractor twice, once to measure the seeing
                and then once again to re-detect objects and do photometry
    do_seeing:  option to calculate the seeing from detections
    do_apcor:   option to calculate aperture corrections from detections
    timeout:    optional timeout, in seconds (useful in pipeline operations).
    """

    # TODO:  Need to fix display of SEx_switches below.
    print_cmd_line("STAP_SEx.py", imname, seeing=seeing, rerun=rerun,
                   config=config, do_apcor=do_apcor, do_seeing=do_seeing,
                   SEx_switches=SEx_switches, timeout=timeout)
    if rerun:  do_seeing = True

    # RS 2014/07/02:  Verify which pyfits we're using...
    print "Using pyfits", pyfits.__version__

    # RS 2011/08/18:  The convolution kernel, star/galaxy neural net weights,
    # etc. are copied synchronously by the main script into the working path
    # (/ramdisk for pipeline mode running).

    # RS 2012/02/17:  Parse the CCD keyword from the filename for later.
    # Right now this isn't being filled elsewhere (but someday it will).
    iminfo = Constants.Filenames(imname)
    ccd = iminfo.ccd
    #get default
    saturate=Constants.Imager.saturate_adu
    gain=Constants.Imager.gain

    # RS 2011/10/20:  If the SEEING keyword exists in the image header,
    # it probably came from a previous run of SExtractor.  Use it as a more
    # reliable estimate of the appropriate detection filter FWHM.
    with pyfits.open(imname, mode='update') as hdulist:
        imghdr = hdulist[0].header
        if "SEEING" in imghdr:  seeing = float(imghdr["SEEING"])
        # RS 2014/02/05:  Fix the EQUINOX keyword, which needs to be a float
        # and not a string (new TAROS update is filling EQUINOX='2000.0').
        # RS 2014/11/07:  Weirdly, the EQUINOX keyword doesn't seem to exist
        # for old data and pyfits isn't letting us create it via assignment,
        # so this behavior of dictionaries doesn't map exactly onto headers.
        # Make a call to the "update" method instead.
        imghdr.update("EQUINOX", 2000.0, "Equinox of coordinates")
        # RS 2012/06/19:  Extract a few other keywords too...
        exptime = float(imghdr["EXPTIME"])
        airmass = float(imghdr["AIRMASS"])
        filtname = imghdr["FILTNAME"]
        minx, maxx = 0, float(imghdr['NAXIS1'])
        miny, maxy = 0, float(imghdr['NAXIS2'])
        #override defaults
        if 'SATURATE' in imghdr:
            if imghdr['SATURATE'] >0: saturate=imghdr['SATURATE']
        if 'GAIN' in imghdr:
            if imghdr['GAIN']>0: gain=imghdr['GAIN']

    # RS 2011/04/21:  Choose a good convolution kernel.
    if seeing > 0.0:  convfile = "gauss_1.5_3x3.conv"
    if seeing > 2.0:  convfile = "gauss_2.0_3x3.conv"
    if seeing > 2.5:  convfile = "gauss_2.5_5x5.conv"
    if seeing > 3.0:  convfile = "gauss_3.0_5x5.conv"
    if seeing > 3.5:  convfile = "gauss_3.0_7x7.conv"
    if seeing > 4.0:  convfile = "gauss_4.0_7x7.conv"
    if seeing > 4.5:  convfile = "gauss_5.0_9x9.conv"

    # RS 2012/02/27:  Generate the right filename for the star catalog.
    # RS 2012/06/29:  If some other file with this name exists, blow it away.
    starsfile = Constants.sextractor_output(iminfo.basefname)

    # RS 2011/04/21:  List of switches to override sextractor defaults.
    # This maps from the SkyMapper parameters, with whatever names we
    # give them and/or extract from the FITS headers, to the parameter
    # names sextractor uses.
    switches = \
        [
        # Pixel scale (arcsec/pix)
        "-PIXEL_SCALE {0:.2f}".format(Constants.Imager.pixel_scale),
        # FWHM of PSF (seeing, arcsec)
        "-SEEING_FWHM {0:.2f}".format(Constants.Imager.pixel_scale*seeing),
        # Upper saturation limits for pixels
        "-SATUR_LEVEL {0:.1f}".format(saturate),
        # Detector gain
        "-GAIN {0:.3f}".format(gain),
        # Convolution kernel for extraction
        "-FILTER_NAME " + convfile,
        # Output catalog name
        "-CATALOG_NAME " + starsfile,
        ]
    if SEx_switches: switches = switches + SEx_switches.split(" ")

    # RS 2011/04/19:  Now assumes this script is in the $PATH.
    # RS 2011/08/18:  Use new switches.
    print "SExtractor:  Using convolution kernel", convfile
    cmd = "sex {0} -c {1} {2}".format(imname, config, " ".join(switches))

    # RS 2012/02/27:  We'll want to save the output here...
    status, stdoutstr = STAP_callexternal(cmd, combinestderr=True,
                                          getstdout=True, timeout=timeout)
    if status != 0:  raise ExternalFailure(cmd=cmd, exit_code=status)

    # RS 2011/09/14:  Similar to Brian's runsexnewfits.pl, we extract the
    # background mean and sigma from the SExtractor output.
    med, sig, thr = None, None, None
    for line in stdoutstr.split('\n'):
        # Use regular expressions to match the log output.
        mm = re.match(r'.*Background:\s+(\S+).*RMS:\s+(\S+).*Threshold:\s+(\S+)',line)
        if mm:
            med, sig, thr = mm.groups()
            med, sig, thr = float(med), float(sig), float(thr)
            break
    if med is None:
        print ""
        raise STAP_Error(msg="Couldn't find median in SExtractor output")

    # If sextractor didn't find anything, now would be a good time to die.
    stars, starhdr = SExtractorDetection.read_fits_file(starsfile)
    if len(stars) == 0:
        raise STAP_Error(msg="SExtractor detected no objects")
    else:
        print "SExtractor:  Successfully sextracted", len(stars), "objects."
        print "Background:  ", med, "  RMS:  ", sig, "  Thresh:  ", thr

    # RS 2013/10/05:  Use an external call to xy2sky to get the borders of
    # the image in RA and DEC.
    cmd = "xy2sky {0} {1} {2} {3} {4}".format(imname, minx, miny, maxx, maxy)
    status, stdoutstr = STAP_callexternal(cmd, combinestderr=True,
                                          getstdout=True, timeout=timeout)
    if status == 0:
        line1, line2 = stdoutstr.split('\n')[:2]
        coo1 = ephem.Equatorial(*(line1.split()[:2]))
        coo2 = ephem.Equatorial(*(line2.split()[:2]))
        minra, mindec = min(coo1.ra, coo2.ra), min(coo1.dec, coo2.dec)
        maxra, maxdec = max(coo1.ra, coo2.ra), max(coo1.dec, coo2.dec)
        coomin = ephem.Equatorial(minra, mindec)
        coomax = ephem.Equatorial(maxra, maxdec)
        minra_str, mindec_str = str(coomin.ra), str(coomin.dec)
        maxra_str, maxdec_str = str(coomax.ra), str(coomax.dec)
    else:
        minra_str = maxra_str = "00:00:00.00"
        mindec_str = maxdec_str = "+00:00:00.0"

    # If we don't already have a seeing estimate, we should calculate one.
    # RS 2011/10/26:  I'd argue that we use the SExtractor star/galaxy
    # neural net here to select starlike objects, but that gives weird
    # results for normal images.  The RMS (from A_IMAGE*B_IMAGE) seems to
    # fluctuate a lot.  FWHM_IMAGE, on the other hand, is very stable and
    # has a clear mode in the distribution near the (apparent) seeing.
    # So use that for all high-S/N (apsig > 20) detections on the image.
    # Make a cut on elongation (e < 1.2) to weed out really weird sources.
    # Also, don't use any sources flagged as bad (FLAGS = 0).
    tmp_seeing = None
    if do_seeing:  tmp_seeing = calc_seeing(stars, verbose=True)
    if tmp_seeing != None:  seeing = tmp_seeing
    print "Using seeing = {0:.2f} pix ({1:.2f}\") for downstream.".format(
            seeing, seeing*Constants.Imager.pixel_scale)

    # RS 2012/11/27:  Added elongation here.
    elong = calc_elong(stars, verbose=True)

    # RS 2012/06/20:  Include aperture correction here as well.  Right now
    # it's just a crappy weighted sum, but in the future we should use GPs.
    if do_apcor:
        apcor, apcor_err, apcor_chi2 = calc_apcor(stars, verbose=True)

    # Write out the new keyword values to both the SExtractor file header and
    # the main FITS image header (for easy reference).  Some of these keywords
    # will hopefully be filled by the main survey pipeline someday.
    keywords_to_update = \
    [
        ("IMAGEID", ccd, "CCD number in the mosaic"),
        ("EXPTIME", exptime, "Exposure time (sec)"),
        ("FILTNAME", filtname, "Centred Filter"),
        ("BKGMED", med, "SExtractor-calculated background level in ADU"),
        ("BKGSIG", sig, "SExtractor-calculated background sigma in ADU"),
        ("BKGTHR", thr, "SExtractor-calculated pixel threshold in ADU"),
        ("SEEING", seeing, "SExtractor-calculated median seeing in pixels"),
        ("ELONG", elong, "SExtractor-calculated median elongation (a/b)"),
    ]

    # RS 2014/02/19:  Add field boundaries in RA and DEC.
    if status==0:
        keywords_to_update += \
        [
            ('MINRA', minra_str, "Lower boundary of image in RA"),
            ('MINDEC', mindec_str, "Lower boundary of image in RA"),
            ('MAXRA', maxra_str, "Upper boundary of image in RA"),
            ('MAXDEC', maxdec_str, "Upper boundary of image in DEC"),
        ]

    # RS 2014/02/19:  Add aperture correction info, if these were performed.
    if do_apcor and apcor != None:
        for i in range(len(apcor)):
            if apcor[i] in [None, np.nan, np.inf]:  continue
            num = "{0:02d}".format(i+1)
            keywords_to_update += \
            [
                ("APCORR"+num, apcor[i],
                 "Aperture correction from ap #{0} (mag)".format(i+1)),
                ("APCERR"+num, apcor_err[i],
                 "Error on aperture correction #{0}".format(i+1)),
            ]

    # RS 2014/02/19:  Update the image header with this info.
    with pyfits.open(imname, mode="update") as hdulist:
        for kw in keywords_to_update:  hdulist[0].header.update(*kw)
        hdulist.flush()

    # RS 2014/02/19:  Propagate these keywords into the header as well.
    imgcards = imghdr.ascardlist()
    for kw in ['FIELD_ID', 'SUBFIELD', 'DATE-OBS']:
        if kw in imghdr:
            keywords_to_update.append(
                (imgcards[kw].key, imgcards[kw].value, imgcards[kw].comment))

    # RS 2014/02/19:  Update the SExtractor output FITS header with this info.
    with pyfits.open(starsfile, mode="update") as hdulist:
        for kw in keywords_to_update:  hdulist[0].header.update(*kw)
        # RS 2013/10/05:  Need to stick this in there because the .stars file
        # is generated by SExtractor, not by FITSRecord.write_fits_file().
        hdulist[1].header.update("CLASSNAM", "SExtractorDetection",
                                 "Class name of table row objects")
        hdulist.flush()

    # RS 2012/02/27:  If we're going to re-run with the new seeing estimate,
    # go ahead and do that now.  (Assuming it's worthwhile.)
    if rerun and tmp_seeing != None:
        SEx(imname, seeing=seeing, rerun=False, config=config,
            SEx_switches=SEx_switches, timeout=timeout)

def main():
    """Wrapper allowing STAP.SEx() to be run from the command line"""

    parser = argparse.ArgumentParser(description='Run SExtractor on an image')
    parser.add_argument('imname',
                        help='filename of input image')
    parser.add_argument('--timeout', type=float, default=30,
                        help='maximum running time allowed for external call (default: %(default)s s)')
    parser.add_argument('--config', default="STAP_SEx.sex",
                        help='SExtractor configuration file (default: %(default)s)')
    parser.add_argument('--seeing', type=float, default=5.0,
                        help='FWHM in pix of detection filter (default: %(default)s)')
    parser.add_argument('--do_seeing', action='store_true', default=False,
                        help='calculate seeing from detections? (default: %(default)s)')
    parser.add_argument('--do_apcor', action='store_true', default=False,
                        help='calculate aperture corrections? (default: %(default)s)')
    parser.add_argument('--rerun', action='store_true', default=False,
                        help='run SExtractor twice to measure seeing? (default: no)')
    parser.add_argument('--SEx_switches',default=None,
                        help='additional SExtractor swithces')
    args = parser.parse_args()
    SEx(args.imname, seeing=args.seeing, rerun=args.rerun, config=args.config,
        do_seeing=args.do_seeing, do_apcor=args.do_apcor,
        timeout=args.timeout, SEx_switches=args.SEx_switches)

if __name__ == "__main__":
    main()

max_stars_repo_name: rscalzo/subpipe

max_issues_repo_name: rscalzo/subpipe

max_forks_repo_name: rscalzo/subpipe


text:
#!/usr/bin/env python

"""
Plot signal heatmaps from TFBS across different bigwigs

@author: Mette Bentsen
@contact: mette.bentsen (at) mpi-bn.mpg.de
@license: MIT
"""

import os
import sys
import argparse
import logging

import numpy as np
import matplotlib as mpl
mpl.use("Agg")	#non-interactive backend
import matplotlib.pyplot as plt
from mpl_toolkits.axes_grid1.inset_locator import inset_axes
import matplotlib.gridspec as gridspec
from datetime import datetime
from sklearn import preprocessing

import pyBigWig
import pysam
import pybedtools as pb

from tobias.parsers import add_heatmap_arguments
from tobias.utils.regions import *
from tobias.utils.utilities import *

#----------------------------------------------------------------------------------------#
def run_heatmap(args):

	#Start logger
	logger = TobiasLogger("PlotHeatmap", args.verbosity)
	logger.begin()

	parser = add_heatmap_arguments(argparse.ArgumentParser())
	logger.arguments_overview(parser, args)
	logger.output_files([args.output])

	check_required(args, ["TFBS", "signals"])

	#Setup TFBS names if not yet
	if args.TFBS_labels == None:
		args.TFBS_labels = [[os.path.basename(fil) for fil in args.TFBS[i]] for i in range(len(args.TFBS))]

	if args.signal_labels == None:
		args.signal_labels = [os.path.basename(fil) for fil in args.signals]

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

	#Check valid input parameters (number of input TFBS vs. bigwig etc.)
	no_signals = len(args.signals)
	no_columns = len(args.show_columns)
	no_TFBS_col = len(args.TFBS)

	if no_TFBS_col > 1 and len(args.show_columns) > 0:
		sys.exit("Error: option --show_columns is not available for multiple --TFBS inputs.")

	if no_TFBS_col > 1 and no_signals != no_TFBS_col:
		sys.exit("Error: Number of --TFBS does not match number of signals")

	elif no_TFBS_col == 1 and no_signals > 1:
		#copy bed_f to other columns

		logger.info("Using bedfiles: {0} across all bigwigs".format(args.TFBS))
		for i in range(no_signals-1):
			args.TFBS.append(args.TFBS[0])
			args.TFBS_labels.append(args.TFBS_labels[0])

	else:
		for i, signal in enumerate(args.signals):
			logger.info("Using {0} with signal from {1}".format(args.TFBS[i], signal))

	#todo: logger overview of bedfiles per column?

	######################################################################################
	##################################### INPUT DATA #####################################
	######################################################################################

	#Setup info dict
	heatmap_info = {col:{row:{"bigwig_f": args.signals[col], "bed_f":args.TFBS[col][row]} for row in range(len(args.TFBS[col]))} for col in range(len(args.signals))}

	#Add extra columns
	for i, bed_column in enumerate(args.show_columns):
		heatmap_info[no_signals+i] = {row:{"column": bed_column, "bed_f":args.TFBS[0][row]} for row in range(len(args.TFBS[0]))}

	#------------------------------------------------------------------------------------#
	#------------------------ Read input files to RegionLists ---------------------------#
	#------------------------------------------------------------------------------------#

	seen_bed = []

	#Read regions per heatmap in grid
	logger.comment("")
	logger.info("Reading bedfiles")
	for col in range(len(heatmap_info)):
		for row in range(len(heatmap_info[col])):

			heatmap_info[col][row]["regions"] = RegionList().from_bed(heatmap_info[col][row]["bed_f"])

			#Estimate region width
			distri = heatmap_info[col][row]["regions"].get_width_distri()
			if len(distri) > 1:
				logger.warning("Input regions have differing lengths: {0}".format(distri))

			heatmap_info[col][row]["width"] = list(distri.keys())[0]

			#Extend to flank
			heatmap_info[col][row]["regions"] = heatmap_info[col][row]["regions"].apply_method(OneRegion.set_width, 2*args.flank)

			#Sort if chosen
			if args.sort_by != None:
				try:
					heatmap_info[col][row]["regions"].sort(key=lambda region: float(region[args.sort_by]), reverse=True)
				except:
					heatmap_info[col][row]["regions"].sort(key=lambda region: region[args.sort_by], reverse=True)

			#Get scores from file
			invalid = []
			for i, bed_column in enumerate(args.show_columns):
				heatmap_info[no_signals+i][row]["column_{0}".format(bed_column)] = [region[bed_column] for region in heatmap_info[col][row]["regions"]]

				try:
					heatmap_info[no_signals+i][row]["column_{0}".format(bed_column)] = [float(element) for element in heatmap_info[no_signals+i][row]["column_{0}".format(bed_column)]]
				except:
					logger.info("Column {0} cannot be converted to float - excluding".format(bed_column))
					del heatmap_info[no_signals+i][row]["column_{0}".format(bed_column)]
					invalid.append(bed_column)

			for bed_column in invalid:
				args.show_columns.remove(bed_column)

			#Logger info about bedfile
			if heatmap_info[col][row]["bed_f"] not in seen_bed:
				logger.info("- Read {1} sites from {0} of width {2}".format(heatmap_info[col][row]["bed_f"], len(heatmap_info[col][row]["regions"]), heatmap_info[col][row]["width"]))
			seen_bed.append(heatmap_info[col][row]["bed_f"])

	#------------------------------------------------------------------------------------#
	#------------------------------ Signals from all sites ------------------------------#
	#------------------------------------------------------------------------------------#

	logger.comment("")
	logger.info("Reading signals from bigwigs")
	for col in range(len(args.TFBS)):

		bigwig_f = heatmap_info[col][0]["bigwig_f"]		#bigwig is the same for all rows, therefore row == 0
		pybw = pyBigWig.open(bigwig_f, "rb")

		for row in heatmap_info[col]:

			logger.info("- Reading {0} from {1}".format(heatmap_info[col][row]["bed_f"], bigwig_f))

			if len(heatmap_info[col][row]["regions"]) > 0:
				heatmap_info[col][row]["signal_mat"] = np.array([region.get_signal(pybw) for region in heatmap_info[col][row]["regions"]])
				heatmap_info[col][row]["aggregate"] = np.mean(heatmap_info[col][row]["signal_mat"], axis=0)
			else:
				heatmap_info[col][row]["signal_mat"] = None
				heatmap_info[col][row]["aggregate"] = None

		pybw.close()

	logger.comment("")

	#------------------------------------------------------------------------------------#
	#---------------------------------- Colorbar min/max --------------------------------#
	#------------------------------------------------------------------------------------#

	#Estimate min/max from all matrices
	if args.share_colorbar == True:

		mats = []
		for col, bigwig in enumerate(args.signals):
			for row in heatmap_info[col]:
				if heatmap_info[col][row]["signal_mat"] is not None:
					mats.append(heatmap_info[col][row]["signal_mat"])

		vmin, vmax = (0,0)
		if len(mats) > 0:
			joined = np.vstack(mats)
			vmin, vmax = np.percentile(joined, [1, 99])

		#Set vmin/vmax for all plots
		for col, bigwig in enumerate(args.signals):
			for row in heatmap_info[col]:
				heatmap_info[col][row].update({"vmin":vmin, "vmax":vmax})

	# Estimate min/max for each bigwig
	else:

		for col, bigwig in enumerate(args.signals):
			mats = [heatmap_info[col][row]["signal_mat"] for row in heatmap_info[col] if heatmap_info[col][row]["signal_mat"] is not None]
			vmin, vmax = (0,0)
			if len(mats) > 0:
				joined = np.vstack(mats)

				vmin, vmax = np.percentile(joined, [1, 99])

			for row in heatmap_info[col]:
				heatmap_info[col][row].update({"vmin":vmin, "vmax":vmax})

	del mats
	del joined

	# Estimate min/max for extra columns
	for i, name in enumerate(args.show_columns):
		col = no_signals + i
		glob_values = []
		for row in range(len(args.TFBS[0])):
			glob_values.extend(heatmap_info[col][row]["column_{0}".format(name)])

		vmin, vmax = np.percentile(glob_values, [1, 99])
		for row in range(len(args.TFBS[0])):
			heatmap_info[col][row]["vmin"] = vmin
			heatmap_info[col][row]["vmax"] = vmax

		del glob_values

	######################################################################################
	##################################### PLOTTING #######################################
	######################################################################################

	#------------------------------------------------------------------------------------#
	#------------------------------------ Set up plots ----------------------------------#
	#------------------------------------------------------------------------------------#

	logger.info("Setting up plotting grid")

	total_columns = no_signals + no_columns
	xvals = np.arange(-args.flank, args.flank)

	fig = plt.figure(figsize = (no_signals*5, 5*5))
	h_ratios = [2,10,0.1]
	w_ratios = [1]*no_signals + [0.1]*no_columns
	gs = gridspec.GridSpec(3, total_columns, height_ratios=h_ratios, width_ratios=w_ratios, hspace=0.1, wspace=0.3)	#aggregate + heatmaps (with sub heatmaps) + colorbar

	#Setup axarr fitting to grid
	axdict = {col:{row:"ax" for row in ["aggregate"] + list(heatmap_info[col]) + ["colorbar"]} for col in range(no_signals)}
	axdict.update({col:{row:"ax" for row in ["aggregate"] + list(heatmap_info[col]) + ["colorbar"]} for col in range(no_signals, no_signals+no_columns)})

	#Per signal column
	xvals = np.arange(-args.flank, args.flank)
	for col in range(no_signals):

		#Aggregates
		axdict[col]["aggregate"] = fig.add_subplot(gs[0,col])
		axdict[col]["aggregate"].set_xlim(left=-args.flank, right=args.flank)
		axdict[col]["aggregate"].set_xlabel('bp from center')
		axdict[col]["aggregate"].set_ylabel('Mean aggregate signal')
		axdict[col]["aggregate"].set_title("{0}".format(args.signal_labels[col]))

		#Heatmaps
		no_beds = len(args.TFBS[col])
		h_ratios = [len(heatmap_info[col][row]["regions"]) for row in heatmap_info[col]]
		h_ratios = [max(num,1) for num in h_ratios] 	#deal with empty beds
		gs_sub = gridspec.GridSpecFromSubplotSpec(no_beds, 1, subplot_spec=gs[1,col], height_ratios=h_ratios, hspace=0.05)

		for row in range(no_beds):
			axdict[col][row] = plt.Subplot(fig, gs_sub[row,0])
			fig.add_subplot(axdict[col][row])

			#Appearance
			plt.setp(axdict[col][row].get_yticklabels(), visible=False)  #Hide y-axis ticks
			plt.setp(axdict[col][row].get_xticklabels(), visible=False)  #Hide x-axis ticks
			axdict[col][row].tick_params(direction="in")
			axdict[col][row].set_ylabel("{0} ({1})".format(args.TFBS_labels[col][row], len(heatmap_info[col][row]["regions"])))

			#Last row
			if row == no_beds-1:
				axdict[col][row].set_xlabel('bp from center')
		#Colorbar
		axdict[col]["colorbar"] = fig.add_subplot(gs[2,col])	#row number 3

	for col in range(no_signals, no_signals + no_columns):
		gs_sub = gridspec.GridSpecFromSubplotSpec(no_beds, 1, subplot_spec=gs[1,col], height_ratios=h_ratios, hspace=0.05)
		for row in range(no_beds):
			axdict[col][row] = plt.Subplot(fig, gs_sub[row,0])

			plt.setp(axdict[col][row].get_yticklabels(), visible=False)  #Hide y-axis ticks
			plt.setp(axdict[col][row].get_xticklabels(), visible=False)  #Hide x-axis ticks
			axdict[col][row].tick_params(direction="in")
			fig.add_subplot(axdict[col][row])

	#------------------------------------------------------------------------------------#
	#--------------------------------- Fill in plots ------------------------------------#
	#------------------------------------------------------------------------------------#

	logger.info("Filling in grid")

	#Colormaps
	for col, bigwig in enumerate(args.signals):
		colors = mpl.cm.jet(np.linspace(0, 1, len(heatmap_info[col]))) 	#colors for aggregate plots

		for row in heatmap_info[col]:
			if heatmap_info[col][row]["signal_mat"] is not None:

				#Aggregate
				axdict[col]["aggregate"].plot(xvals, heatmap_info[col][row]["aggregate"], color=colors[row], linewidth=2, label=args.TFBS_labels[col][row])

				#Heatmap
				lim = np.max([np.abs(heatmap_info[col][row]["vmin"]),np.abs(heatmap_info[col][row]["vmax"])])
				heatmap_info[col][row]["vmin"] = -lim
				heatmap_info[col][row]["vmax"] = lim
				heatmap = axdict[col][row].imshow(heatmap_info[col][row]["signal_mat"], aspect="auto", cmap="seismic", norm=mpl.colors.Normalize(vmin=heatmap_info[col][row]["vmin"], vmax=heatmap_info[col][row]["vmax"]))

				#Insert colorbar (inserted multiple times for each bigwig, but since it is shared for the same bigwig, it doesn't matter)
				fig.colorbar(heatmap, cax=axdict[col]["colorbar"], orientation="horizontal")

	#Extra columns w/ scores from bed
	for i, col in enumerate(range(no_signals, no_signals + no_columns)):
		bed_column = args.show_columns[i]
		for row in heatmap_info[col]:
			values = np.array(heatmap_info[col][row]["column_{0}".format(bed_column)])
			values = values.reshape(-1,1)

			vmin, vmax = np.percentile(values, [1, 99])
			lim = np.max([abs(vmin), abs(vmax)])

			axdict[col][row].imshow(values, aspect="auto", cmap="seismic", norm=mpl.colors.Normalize(vmin=-lim, vmax=lim))

	#------------------------------------------------------------------------------------#
	#-------------------------------- Plot decorations ----------------------------------#
	#------------------------------------------------------------------------------------#

	if args.plot_boundaries:
		for col in heatmap_info:

			motif_len = heatmap_info[col][0]["width"]
			mstart = int(-np.floor(motif_len/2.0))
			mend = int(np.ceil(motif_len/2.0))

			axdict[col]["aggregate"].axvline(mstart, color="black", linestyle="dashed", linewidth=1)
			axdict[col]["aggregate"].axvline(mend, color="black", linestyle="dashed", linewidth=1)

			for row in heatmap_info[col]:

				motif_len = heatmap_info[col][row]["width"]
				mstart = int(-np.floor(motif_len/2.0))
				mend = int(np.ceil(motif_len/2.0))

				axdict[col][row].axvline(mstart+args.flank, color="black", linestyle="dashed", linewidth=1)
				axdict[col][row].axvline(mend+args.flank, color="black", linestyle="dashed", linewidth=1)

	#Add legend to aggregate plots
	for col in range(len(args.signals)):
		axdict[col]["aggregate"].legend(loc=1, prop={"size":6})

	if args.share_colorbar == True:
		ymin = min([axdict[col]["aggregate"].get_ylim()[0] for col in range(no_signals)])
		ymax = max([axdict[col]["aggregate"].get_ylim()[1] for col in range(no_signals)])
		for col in range(no_signals):
			axdict[col]["aggregate"].set_ylim([ymin, ymax])

	#------------------------------------------------------------------------------------#
	#----------------------------- Finish off and output --------------------------------#
	#------------------------------------------------------------------------------------#

	"""
	#For each heatmap
	for row in [1,2]:
		plt.setp(axarr[row].get_yticklabels(), visible=False)  #Hide y-axis ticks
		plt.setp(axarr[row].get_xticklabels(), visible=False)  #Hide x-axis ticks
		axarr[row].tick_params(direction="in")
	"""

	plt.subplots_adjust(top=0.95)
	plt.suptitle(args.title, fontsize=25)

	logger.info("Writing output file")
	plt.savefig(args.output, bbox_inches='tight')
	plt.close()

	logger.end()

#--------------------------------------------------------------------------------------------------------#
if __name__ == '__main__':

	parser = argparse.ArgumentParser()
	parser = add_heatmap_arguments(parser)
	args = parser.parse_args()

	if len(sys.argv[1:]) == 0:
		parser.print_help()
		sys.exit()

	run_heatmap(args)

ext: py

max_stars_repo_name: chris-cheshire/TOBIAS

max_issues_repo_name: chris-cheshire/TOBIAS

max_forks_repo_name: chris-cheshire/TOBIAS

reason: import numpy


text:
#! /usr/bin/env python
# Copyright 2021 Intel Corporation
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from __future__ import print_function
from timeit import default_timer as time

import sys
import numpy as np
import numba_dppy, numba_dppy as dppy
import dpctl

@dppy.kernel
def dppy_gemm(a, b, c):
    i = dppy.get_global_id(0)
    j = dppy.get_global_id(1)
    if i >= c.shape[0] or j >= c.shape[1]:
        return
    c[i, j] = 0
    for k in range(c.shape[0]):
        c[i, j] += a[i, k] * b[k, j]

# Array dimesnions
X = 1024
Y = 16
global_size = X, X

griddim = X, X
blockdim = Y, Y

def driver(a, b, c):
    # Invoke the kernel
    dppy_gemm[griddim, blockdim](a, b, c)

def main():
    a = np.arange(X * X, dtype=np.float32).reshape(X, X)
    b = np.array(np.random.random(X * X), dtype=np.float32).reshape(X, X)
    c = np.ones_like(a).reshape(X, X)

    if dpctl.has_gpu_queues():
        with dpctl.device_context("opencl:gpu") as gpu_queue:
            driver(a, b, c)
    elif dpctl.has_cpu_queues():
        with dpctl.device_context("opencl:cpu") as cpu_queue:
            driver(a, b, c)
    else:
        print("No device found")
        exit()

    # Host compute using standard Numpy
    Amat = np.matrix(a)
    Bmat = np.matrix(b)
    Cans = Amat * Bmat

    # Check result
    assert np.allclose(c, Cans)
    print("Done...")

if __name__ == "__main__":
    main()

max_stars_repo_name: vlad-perevezentsev/numba-dppy

max_issues_repo_name: vlad-perevezentsev/numba-dppy

max_forks_repo_name: vlad-perevezentsev/numba-dppy

reason: import numpy,import numba


text:
import context
import numpy as np
from tensorflow.keras.models import Sequential

from diffraction_image_rcg.nn_model import NNModel

def test_generate_combination_low_high():
    depth = 3
    low_no = 5
    high_no = 7
    values = NNModel.generate_combination_low_high(depth, low_no, high_no)

    expected_values = np.asarray(
        [
            [5, 5, 5],
            [5, 5, 6],
            [5, 5, 7],
            [5, 6, 5],
            [5, 6, 6],
            [5, 6, 7],
            [5, 7, 5],
            [5, 7, 6],
            [5, 7, 7],
            [6, 5, 5],
            [6, 5, 6],
            [6, 5, 7],
            [6, 6, 5],
            [6, 6, 6],
            [6, 6, 7],
            [6, 7, 5],
            [6, 7, 6],
            [6, 7, 7],
            [7, 5, 5],
            [7, 5, 6],
            [7, 5, 7],
            [7, 6, 5],
            [7, 6, 6],
            [7, 6, 7],
            [7, 7, 5],
            [7, 7, 6],
            [7, 7, 7],
        ]
    )

    assert np.array_equal(values, expected_values)

def test_generate_combination_low_high_different():
    low_high_arr = np.asarray([[2, 5], [7, 8]])

    values = NNModel.generate_combination_low_high_different(low_high_arr)

    expected_values = np.asarray(
        [[2, 7], [2, 8], [3, 7], [3, 8], [4, 7], [4, 8], [5, 7], [5, 8]]
    )

    assert np.array_equal(values, expected_values)

def test_nn_model_build():
    # Arrange
    hidden_units = np.asarray([10, 10])
    input_dimension = 16
    learning_rate = 0.005
    activation_fun = "relu"

    # Act
    model = NNModel.build(
        hidden_units=hidden_units,
        input_dimension=input_dimension,
        learning_rate=learning_rate,
        activation_func=activation_fun,
    )

    assert isinstance(model, Sequential)

def test_nn_model_build_dropout():
    # Arrange
    hidden_units = np.asarray([10, 10])
    input_dimension = 16
    learning_rate = 0.005
    activation_fun = "relu"
    dropout_rate = np.asarray([0.5, 0.5])

    # Act
    model = NNModel.build_dropout(
        hidden_units=hidden_units,
        input_dimension=input_dimension,
        learning_rate=learning_rate,
        activation_func=activation_fun,
        dropout_rate=dropout_rate,
    )

    assert isinstance(model, Sequential)


text:
#!/usr/bin/python
# -*- coding: utf-8 -*-
""" Load dataset and preprocess data before model training """

from sklearn.model_selection import train_test_split
from sklearn import preprocessing
import pandas as pd
import numpy as np
import shutil
import os

from joblib import Memory

__all__ = ['load_data']

def _osrm(datadir):
    features = ['total_distance', 'total_travel_time', 'number_of_steps']
    fr1 = pd.read_csv(os.path.join(datadir, 'osrm/fastest_routes_train_part_1.csv'), usecols=['id', *features])
    fr2 = pd.read_csv(os.path.join(datadir, 'osrm/fastest_routes_train_part_2.csv'), usecols=['id', *features])
    test_street_info = pd.read_csv(os.path.join(datadir, 'osrm/fastest_routes_test.csv'), usecols=['id', *features])
    train_street_info = pd.concat((fr1, fr2))
    return train_street_info, test_street_info, features

def _bucketize(train_targets, valid_targets, bucket_count):
    """ Process buckets from train targets and deduce labels of trainset and testset """
    sorted_targets = np.sort(train_targets)
    bucket_size = len(sorted_targets) // bucket_count
    buckets = [sorted_targets[i * bucket_size: (1 + i) * bucket_size] for i in range(bucket_count)]
    bucket_maxs = [np.max(b) for b in buckets]
    bucket_maxs[-1] = float('inf')

    # Bucketize targets (labels are bucket indices)
    def _find_indice(value): return np.searchsorted(bucket_maxs, value)
    train_labels = np.vectorize(_find_indice)(train_targets)
    valid_labels = np.vectorize(_find_indice)(valid_targets)
    # Process buckets means
    buckets_means = [np.mean(bucket) for bucket in buckets]
    return train_labels, valid_labels, buckets_means

def load_data(datadir, trainset, testset, valid_size, output_size, embed_discrete_features=False, max_distinct_values=None, cache_read_only=False):
    if cache_read_only:
        dest = '/output/cache'
        shutil.copytree(datadir, dest)
        datadir = dest
    memory = Memory(cachedir=os.path.join(datadir, 'cache'))

    @memory.cache(ignore=['datadir'])
    def _cached_light_load_data(datadir, trainset, testset, valid_size, embed_discrete_features, max_distinct_values):
        features = ['vendor_id', 'passenger_count', 'pickup_latitude', 'pickup_longitude', 'dropoff_latitude', 'dropoff_longitude']
        df_all = pd.concat((pd.read_csv(os.path.join(datadir, trainset)), pd.read_csv(os.path.join(datadir, testset))))
        df_all['pickup_datetime'] = df_all['pickup_datetime'].apply(pd.Timestamp)
        df_all['dropoff_datetime'] = df_all['dropoff_datetime'].apply(pd.Timestamp)
        df_all['trip_duration_log'] = np.log(df_all['trip_duration'] + 1)

        def _add_feature(name, value):
            features.append(name)
            df_all[name] = value

        _, indices = np.unique(df_all['id'], return_inverse=True)
        _add_feature('id_idx', indices)
        # Date time features
        pickup_dt = df_all['pickup_datetime'].dt
        _add_feature('week_delta', pickup_dt.weekday + ((pickup_dt.hour + (pickup_dt.minute / 60.0)) / 24.0))
        _add_feature('weekofyear', pickup_dt.weekofyear)
        _add_feature('weekday', pickup_dt.weekday)
        _add_feature('seconds', pickup_dt.second + pickup_dt.minute * 60.)
        _add_feature('month', pickup_dt.month)
        _add_feature('hour', pickup_dt.hour)

        # Get test set and train set from df_all
        X_train = df_all[df_all['trip_duration'].notnull()]
        y_train = df_all[df_all['trip_duration'].notnull()]['trip_duration_log'].values.flatten()
        X_test = df_all[df_all['trip_duration'].isnull()]
        test_ids = df_all[df_all['trip_duration'].isnull()]['id'].values

        # Add OSRM data
        train_street_info, test_street_info, osrm_features = _osrm(datadir)
        features.extend(osrm_features)
        X_train = X_train.merge(train_street_info, how='left', on='id')[features]
        X_test = X_test.merge(test_street_info, how='left', on='id')[features]

        # Fill missing osrm data
        mean_distance, mean_travel_time, mean_steps = X_train.total_distance.mean(), X_train.total_travel_time.mean(), X_train.number_of_steps.mean()

        def _fillnan(df):
            df.total_distance = df.total_distance.fillna(mean_distance)
            df.total_travel_time = df.total_travel_time.fillna(mean_travel_time)
            df.number_of_steps = df.number_of_steps.fillna(round(mean_steps))
        _fillnan(X_train)
        _fillnan(X_test)

        # Split dataset into trainset and testset
        train_data, valid_data, train_targets, valid_targets = train_test_split(X_train.values, y_train, test_size=valid_size, random_state=459)

        # Normalize feature columns
        standardizer = preprocessing.StandardScaler()
        train_data = standardizer.fit_transform(train_data)
        valid_data = standardizer.transform(valid_data)
        test_data = standardizer.transform(X_test.values)

        # Encode discrete features with less than 'max_distinct_values' unique values
        discrete_features = []
        if embed_discrete_features:
            for feature_idx in range(len(features)):
                train_feature = train_data[:, feature_idx]
                # TODO: apply a threshold on values count in order to filter out outlier values
                # (values, counts = np.unique(train_data[:, feature_idx], return_counts=True))
                unique_values = np.unique(train_feature)
                if max_distinct_values is None or len(unique_values) < max_distinct_values:
                    discrete_features.append((feature_idx, np.append(unique_values, [len(unique_values)])))

                    def _encode(array, encoded_values):
                        rslt = np.empty_like(array)
                        for val in np.unique(array):
                            encode_idx = np.argwhere(encoded_values == val)
                            arr_idx = np.argwhere(array == val)
                            rslt[arr_idx] = len(encoded_values) if len(encode_idx) == 0 else encode_idx[0, 0]
                        return rslt
                    unique_values, train_data[:, feature_idx] = np.unique(train_feature, return_inverse=True)
                    valid_data[:, feature_idx] = _encode(valid_data[:, feature_idx], unique_values)
                    test_data[:, feature_idx] = _encode(test_data[:, feature_idx], unique_values)

        return len(features), discrete_features, (test_ids, test_data), (train_data, valid_data, train_targets, valid_targets)

    # Parse and preprocess data
    features_len, discrete_features, (test_ids, testset), dataset = _cached_light_load_data(
        datadir, trainset, testset, valid_size, embed_discrete_features, max_distinct_values)
    (train_data, valid_data, train_targets, valid_targets) = dataset

    # Get buckets from train targets
    train_labels, valid_labels, bucket_means = _bucketize(train_targets, valid_targets, output_size)

    return features_len, discrete_features, (test_ids, testset), (train_data, valid_data, train_targets, valid_targets, train_labels, valid_labels), bucket_means


text:
"""
  Unit tests for metropolis sampling of discrete distributions.
  -- kvysyara@andrew.cmu.edu
"""
from __future__ import absolute_import

# pylint: disable=invalid-name
# pylint: disable=no-self-use

import numpy as np

# Local imports
from .discrete import DiscreteUniform, Bernoulli
from ..utils.base_test_class import BaseTestClass, execute_tests

class MetropolisDiscreteTestCases(BaseTestClass):
  """ Unit tests for metropolis sampling of discrete distributions """

  def setUp(self):
    """ Sets up unit tests. """
    self.size = 100000
    self.threshold = 0.1

  def _check_sample_sizes(self, samples):
    """ Compares the sample sizes with the size parameter"""
    assert self.size == len(samples)

  def _compute_mean(self, samples):
    """ Computes Mean """
    return np.mean(samples)

  def _compute_variance(self, samples):
    """ Computes Variance """
    return np.var(samples)

  def test_metropolis_discreteuniform(self):
    """ Tests metropolis sampling from Discrete Uniform distribution """
    self.report('Test metropolis sampling of Discrete Uniform Distribution.')
    lower = -5
    upper = 5
    dist = DiscreteUniform(lower, upper)
    mean = dist.get_mean()
    var = dist.get_variance()
    samples = dist.draw_samples('metropolis', self.size, np.array([0.1]))
    mean_r = self._compute_mean(samples)
    var_r = self._compute_variance(samples)
    self._check_sample_sizes(samples)
    assert (samples >= lower).all() and (samples <= upper).all()
    assert abs(mean - self._compute_mean(samples)) <= self.threshold
    assert abs(var - self._compute_variance(samples)) <= self.threshold
    self.report('%s, mean=%0.3f, variance=%0.3f :: test result: mean=%0.3f, '
                'variance=%0.3f'%(str(dist), mean, var, mean_r, var_r), 'test_result')

  def test_metropolis_bernoulli(self):
    """ Tests metropolis sampling from Bernoulli distribution """
    self.report('Test metropolis sampling of Bernoulli Distribution.')
    p = 0.6
    dist = Bernoulli(p)
    mean = dist.get_mean()
    var = dist.get_variance()
    samples = dist.draw_samples('binarymetropolis', self.size, np.array([0]))
    mean_r = self._compute_mean(samples)
    var_r = self._compute_variance(samples)
    self._check_sample_sizes(samples)
    assert (samples >= 0).all() and (samples <= 1).all()
    assert abs(mean - self._compute_mean(samples)) <= self.threshold
    assert abs(var - self._compute_variance(samples)) <= self.threshold
    self.report('%s, mean=%0.3f, variance=%0.3f :: test result: mean=%0.3f, '
                'variance=%0.3f'%(str(dist), mean, var, mean_r, var_r), 'test_result')

if __name__ == '__main__':
  execute_tests()

max_stars_repo_name: qboticslabs/dragonfly

max_issues_repo_name: qboticslabs/dragonfly

max_forks_repo_name: qboticslabs/dragonfly


text:
import gym
import pybulletgym
import tensorflow as tf
import numpy as np
import math
import utils.functions as functions

EPS = 1e-8

class DiscreteRewardToGoAgent:
    '''
    Agent with environment
    '''
    def __init__(self, env = None, learning_rate = 1e-2, seed = 0, sizes = [32], activation = tf.tanh, save = False):
        '''Initializing all specific variables for bot'''

        assert env, "Value env is required!"

        self.env = gym.make(env)
        self.seed = 0
        self.obs = self.env.reset()
        self.activation = activation
        self.sizes = sizes
        self.obs_dim = self.env.observation_space.shape[0]
        self.n_acts = self.env.action_space.n
        self.learning_rate = learning_rate
        self.obs_ph, self.logits = functions.create_model(obs_dim = self.obs_dim, n_acts = self.n_acts, sizes = self.sizes, activation = self.activation)
        self.action = tf.squeeze(tf.random.categorical(logits = self.logits,\
                                 num_samples = 1),axis = 1)
        self.weights_ph, self.act_ph, self.loss = functions.categorical_loss_func_weights(self.logits, self.n_acts)
        #optimizer
        self.train_op = tf.train.AdamOptimizer(learning_rate = self.learning_rate).minimize(self.loss)

        self.sess = tf.InteractiveSession()
        self.sess.run(tf.global_variables_initializer())

    def act(self, obs):
        '''Choosing action for self.obs '''
        action = self.sess.run(self.action, feed_dict = {
            self.obs_ph : [obs]
        })#[0] or [1]
        return action[0]

    def play_episode(self, rendering = False):
        '''
        Playing only one episode, collecting trajectory and rewards
        '''
        episode_obs_batch = list()
        episode_action_batch = list()
        episode_trajectory = list()
        rewards = list()

        self.obs = self.env.reset()
        while True:
            episode_obs_batch.append(self.obs.copy())
            if rendering == True:
                self.env.render()
            action = self.act(self.obs)
            episode_trajectory.append((self.obs, action))
            self.obs, reward, is_done, info = self.env.step(action)

            episode_action_batch.append(action)

            rewards.append(reward)
            if is_done:
                break
        if rendering == True:
            return sum(rewards)
        else:
            return rewards, episode_obs_batch, episode_action_batch, episode_trajectory

    def play_epoch(self, n):
        '''
        Collecting samples with current policy
        n = number of samples from 1 epoch
        '''
        epoch_obs_batch = list()
        epoch_action_batch = list()
        epoch_rewards = list()
        epoch_trajectory = list()
        epoch_reward_to_go = list()

        while len(epoch_obs_batch)<n:
            reward, obs_batch, action_batch, trajectory_batch = self.play_episode()

            epoch_obs_batch+=(obs_batch)

            epoch_action_batch+=(action_batch)
            epoch_rewards.append(sum(reward))
            epoch_reward_to_go += list(functions.reward_to_go(reward))

            epoch_trajectory.append(trajectory_batch)
        return epoch_obs_batch, epoch_action_batch, epoch_trajectory, epoch_reward_to_go, sum(epoch_rewards)/len(epoch_rewards)

    def train_epoch(self):
        '''
        Training agent for n epochs
        '''
        obs_batches, action_batches, trajectory_batches, reward_to_go, mean_reward = self.play_epoch(5000)
        self.play_episode(True)
        #print("obs: {}, \nactions: {},\n rewards: {}".format(obs_batches[0:8], action_batches[0:8], reward_to_go[0:8]))
        loss, _ = self.optimize(obs_batches, action_batches, reward_to_go)

        return loss, mean_reward

    def can_solve(self):
        for i in range(100):
            reward, _, _, _= self.play_episode()
            if sum(reward) != 200:
                print("{} Cannot solve yet {}".format(i, sum(reward)))
                return False
        return True

    def train_n_epochs(self, n):
        for i in range(n):
            loss, reward = self.train_epoch()
            if reward == 200:
                if self.can_solve():
                    print("SOLVED")
                    return True

            print("#i:{} loss: {}, mean reward: {}".format(i, loss, reward))
        return False

    def loss_func(self):
        '''Calculating loss, returning specific placeholders'''
        weights_ph = tf.placeholder(shape=(None,), dtype = tf.float32)
        act_ph = tf.placeholder(shape=(None,), dtype = tf.int32)

        action_masks = tf.one_hot(act_ph, self.n_acts)#array([[0., 1.]...[1., 0.], dtype=float32)

        log_probs = tf.reduce_sum(action_masks * tf.nn.log_softmax(self.logits), axis=1)
        #array([-0.6952652,...], dtype=float32)
        loss = -tf.reduce_mean(weights_ph* log_probs)#9.812662
        return weights_ph, act_ph, loss

    def optimize(self, obs, act, weights):
        '''Optimizing logits with loss, using self.train_op(optimizer)'''
        return self.sess.run([self.loss, self.train_op], feed_dict = {
            self.obs_ph : obs,
            self.act_ph: act,
            self.weights_ph: weights
        })

if __name__ == "__main__":
    agent = DiscreteRewardToGoAgent("CartPole-v0")#gym.make("InvertedPendulumPyBulletEnv-v0")
    agent.train_n_epochs(500)
    agent.play_episode(True)
    agent.env.close()


text:
# -*- coding: utf-8 -*-
#
# @Author: Florian Briegel (briegel@mpia.de)
# @Date: 2021-08-18
# @Filename: __main__.py
# @License: BSD 3-clause (http://www.opensource.org/licenses/BSD-3-Clause)

import numpy as np

import matplotlib.pyplot as plt
from matplotlib.colors import LogNorm
from matplotlib.patches import Ellipse, Rectangle
from matplotlib import use as plt_use

from mpl_toolkits.axes_grid1 import make_axes_locatable
from scipy.ndimage.filters import gaussian_filter, gaussian_gradient_magnitude
import scipy.ndimage
import sep

from sdsstools.logger import get_logger

class PlotIt:
    def __init__(self, data, rectw=60, logger=get_logger("plotit")):
        self.data = data
        self.rectw = rectw
        self.log = logger

#        plt_use('TkAgg')
        plt_use('Qt5Agg')

        self.fig = plt.figure(figsize=(17, 8))

        self.ax_img = plt.subplot2grid((3,4), (0 ,1), rowspan=3, colspan=3)
        divider = make_axes_locatable(self.ax_img)
        cax_r = divider.append_axes('right', size='3%', pad=0.05)

        self.mean, self.sigma = np.mean(data), np.std(data)
        self.lperc, self.uperc = np.percentile(data, 0.5), np.percentile(data, 99.8)

        self.pl_img = self.ax_img.imshow(data, interpolation='nearest', vmin=self.mean-self.sigma, vmax=self.mean+self.sigma, origin='lower')
        self.ax_img.set_title('skymaker cam', fontweight ="bold")
        self.fig.colorbar(self.pl_img, cax=cax_r, orientation='vertical')
        # plt.colorbar(ticks=[1.0,3.0,6.0])
        # cbar.ax.set_yticklabels(['1','3','6'])
        self.pl_ax = [plt.subplot2grid((3,4), (0,0)),
                      plt.subplot2grid((3,4), (1,0)),
                      plt.subplot2grid((3,4), (2,0))]

        self.pl_xs, self.pl_ys = [None] * len(self.pl_ax), [None] * len(self.pl_ax)

        self.sep_objects()
        self.plot_objects()
        self.plot_star_cuts(update=False)

        self.fig.tight_layout()
        plt.pause(.5)

    def sep_objects(self):
        bkg = sep.Background(self.data)
        data_sub = self.data - bkg
        self.objects = sep.extract(data_sub, 1.5, err=bkg.globalrms)
        self.object_index_sorted_by_peak=list({k: v for k, v in sorted({idx: self.objects['peak'][idx] for idx in range(len(self.objects))}.items(), key=lambda item: item[1], reverse=True)}.keys())

    def plot_objects(self):
        while len(self.ax_img.artists):
            for a in self.ax_img.artists:
                a.remove()

        for i in range(len(self.objects)):
            e = Ellipse(xy=(self.objects['x'][i], self.objects['y'][i]),
                        width=6*self.objects['a'][i],
                        height=6*self.objects['b'][i],
                        angle=self.objects['theta'][i] * 180. / np.pi)
            e.set_facecolor('none')
            e.set_edgecolor('red')
            self.ax_img.add_artist(e)

    def plot_star_cuts(self, update=True):
        self.pl_objects = [self.object_index_sorted_by_peak[0],
                           self.object_index_sorted_by_peak[round(len(self.object_index_sorted_by_peak)/2)],
                           self.object_index_sorted_by_peak[-1]]

        i = 0
        for idx in self.pl_objects:
            xc, yc, p = round(self.objects['x'][idx]), round(self.objects['y'][idx]), self.objects['peak'][idx]

            x0, y0 = xc - self.rectw/2, yc - self.rectw/2
            x1, y1 = xc + self.rectw/2, yc + self.rectw/2
            t = np.arange(-self.rectw/2, +self.rectw/2)

            xx, yx = np.linspace(x0, x1, self.rectw), np.linspace(yc, yc, self.rectw)
            xy, yy = np.linspace(xc, xc, self.rectw), np.linspace(y0, y1, self.rectw)

            if not update:
                #print(f"object {idx}: xy = [{xc}, {yc}]: {p}")
                r = Rectangle(xy=(xc-self.rectw/2, yc-self.rectw/2), width=self.rectw, height=self.rectw)
                r.set_facecolor('none')
                r.set_edgecolor('blue')
                self.ax_img.add_artist(r)

                self.pl_xs[i] = scipy.ndimage.map_coordinates(self.data, np.vstack((yx,xx)))
                self.pl_ys[i] = scipy.ndimage.map_coordinates(self.data, np.vstack((yy,xy)))
                self.pl_ax[i].plot(t, self.pl_xs[i], t, self.pl_ys[i])
            else:
                self.pl_ax[i].clear()
                xi = scipy.ndimage.map_coordinates(self.data, np.vstack((yx,xx)))
                yi = scipy.ndimage.map_coordinates(self.data, np.vstack((yy,xy)))
                self.pl_ax[i].plot(t, self.pl_xs[i], t, self.pl_ys[i], dashes=[4, 2])
                self.pl_ax[i].plot(t, xi, t, yi, linewidth=2)

            self.pl_ax[i].grid(True)
            i+=1

    def update(self, data, new_objects=False):
        self.data = data

        self.pl_img.set_data(data)
        if(new_objects):
            self.sep_objects()
            self.plot_objects()
            self.plot_star_cuts(False)
        else:
            self.plot_star_cuts()

        self.fig.tight_layout()
        self.fig.canvas.draw_idle()
        self.fig.canvas.start_event_loop(0.001)

max_stars_repo_name: mgjeon/skymakercam

max_issues_repo_name: mgjeon/skymakercam

max_forks_repo_name: mgjeon/skymakercam

reason: import numpy,import scipy,from scipy


text:
#!/bin/python
## author: Ethan D. Evans
## date: 10/10/2020
import argparse, os
import pandas as pd
import numpy as np
import mdtraj as md

def parse_args():
	parser = argparse.ArgumentParser()
	parser.add_argument('-p', help='path to the folder containing replica folders and state csv file')
	parser.add_argument('--equil', help='length of the preproduction equilibration to subtract off from times', default=0, type=int)
	parser.add_argument('--remove', help='frames to start at...ie remove the first N frame if remove is N', default=0, type=int)
	parser.add_argument('--restart_num', help='specific restart number to run on, default is to run on all', default=None)
	return parser.parse_args()

def main():
	# python sort_replicas.py -p /media/ee/Seagate\ Portable\ Drive/MD/ala/ala_rest_1/ --equil 50000 --remove 1
	### normally you would not need the --remove above.
	args = parse_args()
	### remove is basically only used for initial ala REST analysis, i changed code s.t. it should normally be 0
	remove = args.remove
	### 1) read in files listing replica and which state it was in at a given snapshot
	work_dir = os.listdir(args.p)
	state_files = [f for f in work_dir if 'csv' in f]
	restart_range = [f[-5] for f in state_files] ### assumes it ends with '#.csv', we want the #
	if args.restart_num != None:
		state_files = [f for f in state_files if 'restart_{}.csv'.format(args.restart_num) in f]
	state_files = [os.path.join(args.p, sf) for sf in state_files]
	### restart_states maps from the restart number to the csv file that hold the states + replica
	restart_states = {restart:pd.read_csv(state_file) for restart,state_file in zip(restart_range,state_files)}
	### replica_n will be 0...N for the different dirs, also the headers for the replicas in the csvs
	replica_n = [f for f in work_dir if 'csv' not in f]
	### get the temps and map to replicas with 0 mapping to lowest real / effective temp (ie highest REST fraction=1)
	temps_obtained = False
	for _,file in restart_states.items():
		### fix the start time, maybe later fixed the restart times if its important
		file['step'] = file['step'] - args.equil
		if not temps_obtained:
			headers = [h for h in list(file) if h != 'step']
			temps = sorted(file[headers].iloc[0,:], reverse=True)
			replica_n_to_temp = {r:i for r,i in zip(replica_n, temps)}
			temps_obtained = True

	### 2) read in DCD files one at a time and split frames into different new dcd files for each state
	for restart in restart_states:
		state_corrected_reps = {s:[] for s in replica_n}
		sorted_f_names = {}
		for r in replica_n:
			rep_dir = os.path.join(args.p,r)
			rep_dir_files = os.listdir(rep_dir)
			f = [f for f in rep_dir_files if restart+'.dcd' in f][0]
			restart_dcd = os.path.join(rep_dir,f)
			sorted_f_names[r] = os.path.join(rep_dir, f[:-4]+'_sorted.dcd')
			top = os.path.join(rep_dir,[f for f in rep_dir_files if restart+'.pdb' in f][0])
			traj = md.load_dcd(restart_dcd, top=top)
			traj = traj[remove:]
			# print(traj.n_frames, traj.xyz.shape, traj.time, traj[0])
			### currently the traj has frames from different states
			rep_i_states_and_restart_n = restart_states[restart][r]

			#### used the following for the trpcage sim to get only the 300k data
			# rep = '0'
			# temp = 1
			# temp_data = rep_i_states_and_restart_n[rep_i_states_and_restart_n==temp]
			# try:
			# 	state_corrected_reps[rep].append(traj[temp_data.index])
			# except:
			# 	indicies = temp_data.index
			# 	indicies = [i for i in indicies if i in range(traj.n_frames)]
			# 	state_corrected_reps[rep].append(traj[indicies])

			for rep,temp in replica_n_to_temp.items():
				temp_data = rep_i_states_and_restart_n[rep_i_states_and_restart_n==temp]
				try:
					state_corrected_reps[rep].append(traj[temp_data.index])
				except:
					indicies = temp_data.index
					indicies = [i for i in indicies if i in range(traj.n_frames)]
					state_corrected_reps[rep].append(traj[indicies])
		### 3) sort the new DCD files to give proper time ordering
		for s, data in state_corrected_reps.items():
			if data == []:
				continue
			d_ = data[0]
			for d in data[1:]:
				d_ = d_.join(d)
			order = np.argsort(d_.time)
			d_ = d_[order]
			d_.save_dcd(sorted_f_names[s])

if __name__ == '__main__':
	main()


text:
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import os

from six.moves import range  # pylint: disable=redefined-builtin
import tensorflow as tf
import numpy as np
import random
from scripts.generate_dataset import GO_ID, EOS_ID
flags = tf.app.flags

# Basic model parameters.
DATA_DIR = os.path.join('..', 'data', 'instagram', 'caption_dataset')
FEATURE_DIR = os.path.join('features_incomplete')
root_path = "/"
train_fpath = 'train.txt'
val_fpath = 'trial.txt'

flags.DEFINE_string("data_dir",
    DATA_DIR,
    "data directory [data]"
)
flags.DEFINE_string("img_data_dir",
    FEATURE_DIR,
    "data directory [data]"
)

flags.DEFINE_integer("max_context_length", 60,
    "User contex max length default [60]"
)
flags.DEFINE_integer("max_output_length", 16,
    "User contex max length default [60]"
)
flags.DEFINE_integer('batch_size', 200,
    """Number of examples to process in a batch."""
)
flags.DEFINE_integer('vocab_size', 40000,
    """Number of vocab."""
)
flags.DEFINE_integer('word_emb_dim', 512,
    """Dimensions of word embeddings."""
)
flags.DEFINE_integer('mem_dim', 1024,
    """Dimensions of memories."""
)
flags.DEFINE_integer('num_channels', 300,
    """Number of channels of memory cnn."""
)

FLAGS = flags.FLAGS

def _generate_data_and_label_batch(inputs, min_queue_examples,
                                    batch_size, shuffle):
  """Construct a queued batch

  """
  # Create a queue that shuffles the examples, and then
  # read 'batch_size' data from the example queue.
  num_preprocess_threads = 8
  if shuffle:
    outputs = tf.train.shuffle_batch(
        inputs,
        batch_size=1,
        num_threads=num_preprocess_threads,
        enqueue_many=False,
        capacity=min_queue_examples + 3 ,
        min_after_dequeue=min_queue_examples
    )
  else:
    outputs = tf.train.batch(
        inputs,
        batch_size=1,
        num_threads=num_preprocess_threads,
        enqueue_many=False,
        capacity=min_queue_examples + 3
    )

  return outputs

def numpy_read_func(batch_path):
  np_list = []
  for i in range(len(batch_path)):
    b_path = batch_path[i].decode('utf-8')
    np_list.append(
        np.load(os.path.join(FLAGS.img_data_dir, b_path))
    )
  return np.array(np_list)

def token_split_func(batch_token, max_length, cap=None):
  all_ids = np.zeros([FLAGS.batch_size, max_length], dtype=np.int32)
  for i in range(FLAGS.batch_size):
    b_token = batch_token[i].decode('utf-8')
    if "_" in b_token:
      valid_ids = list(map(int, b_token.split('_')))[:max_length]
    elif len(b_token) == 0:
      valid_ids = []
    else:
      valid_ids = [int(b_token)]
    if cap:
      valid_ids = valid_ids[:max_length-1]
      if cap == 'caption' :
        all_ids[i, :len(valid_ids)+1]= [GO_ID] + valid_ids[:]
      else:
        all_ids[i, :len(valid_ids)+1]= valid_ids[:] + [EOS_ID]
    else:
      all_ids[i, :len(valid_ids)] = valid_ids[:]
  return all_ids

def mask_build_func(batch_length, max_length, cap=None):
  mask = np.zeros([FLAGS.batch_size, max_length], dtype=np.bool)
  for i in range(FLAGS.batch_size):
    if cap:
      length = int(batch_length[i]) + 1
    else:
      length = int(batch_length[i])
    length = max_length if length > max_length else length
    mask[i, :length] = True
  return mask

def read_numpy_format_and_label(filename_queue):
  # Tricks to make the batch have
  # almost same caption length
  # We should group it by using dequeue_many and enqueue many
  # from sorted list.
  filename_and_label_tensor = filename_queue.dequeue_many(
      FLAGS.batch_size
  )
  batch_filename, batch_context_length, batch_caption_length, \
      batch_context, batch_caption = tf.decode_csv(
          filename_and_label_tensor,
          [[""], [""], [""], [""], [""]]
      )
  batch_context_length = tf.minimum(
      tf.cast(
          tf.string_to_number(batch_context_length),
          tf.int32
      ),
      FLAGS.max_context_length
  )
  batch_caption_length = tf.minimum(
      tf.cast(
          tf.string_to_number(batch_caption_length),
          tf.int32
      )+1,
      FLAGS.max_output_length
  )
  batch_img_embedding = tf.py_func(
      numpy_read_func,
      [batch_filename],
      tf.float32
  )
  batch_context_id = tf.py_func(
      token_split_func,
      [batch_context, FLAGS.max_context_length],
      tf.int32
  )
  batch_caption_id = tf.py_func(
      token_split_func,
      [batch_caption, FLAGS.max_output_length, 'caption'],
      tf.int32
  )
  batch_answer_id = tf.py_func(
      token_split_func,
      [batch_caption, FLAGS.max_output_length, 'answer'],
      tf.int32
  )
  batch_context_mask = tf.py_func(
      mask_build_func,
      [batch_context_length, FLAGS.max_context_length],
      tf.bool
  )
  batch_caption_mask = tf.py_func(
      mask_build_func,
      [batch_caption_length, FLAGS.max_output_length, 'caption'],
      tf.bool
  )

  return batch_img_embedding, batch_context_length, \
      batch_caption_length, batch_context_id, batch_caption_id, \
      batch_answer_id, batch_context_mask, batch_caption_mask
def chunks(data, batch_size):
  chunk_l = []
  N = int(len(data)/batch_size)
  for i in range(N):
    chunk_l.append(data[i*batch_size:(i+1)*batch_size])
  chunk_l.append(data[(i+1)*batch_size:])
  return chunk_l

def mergers(chunk_l):
  newdata = []
  for datas in chunk_l:
    for data in datas:
      newdata.append(data)
  return newdata

def enqueue(eval_data):
  # string input format
  # numpyfname,contextlength,captionlength,contexttoken1_contexttoken2,wordtoken1_wordtoken2
  # e.g. 12345.npy,4,3,445_24_445_232,134_466_234
  if not eval_data:
    filenames = [l.strip() for l in open(os.path.join(FLAGS.data_dir, train_fpath)).readlines()]
    chunk_list = chunks(filenames, FLAGS.batch_size)
    random.shuffle(chunk_list)
    filenames = mergers(chunk_list)
  else:
    filenames = [l.strip() for l in open(os.path.join(FLAGS.data_dir, val_fpath)).readlines()]
  num_examples_per_epoch = len(filenames)

  # Create a queue that produces the filenames to read.
  # Don't shuffle.. to maintain the order
  # Data will be shuffled in the process of the tf.train.shuffle_batch
  filename_queue = tf.train.string_input_producer(
    filenames,
    shuffle=False,
    capacity=num_examples_per_epoch
  )

  img_embedding_list = []
  context_length_list = []
  caption_length_list = []
  context_id_list = []
  caption_id_list = []
  answer_id_list = []
  context_mask_list = []
  caption_mask_list = []

  for it in range(FLAGS.num_gpus):
    img_embedding, context_length, caption_length, context_id, caption_id, \
        answer_id, context_mask, caption_mask = read_numpy_format_and_label( filename_queue)

    # Set the shapes of tensors
    img_embedding.set_shape([FLAGS.batch_size, 2048,1,1])
    img_embedding = tf.transpose(
        tf.reshape(img_embedding, [FLAGS.batch_size, 2048,1]),
        perm=[0, 2, 1]
    )
    context_length.set_shape([FLAGS.batch_size])
    caption_length.set_shape([FLAGS.batch_size])
    context_id.set_shape([FLAGS.batch_size, FLAGS.max_context_length])
    caption_id.set_shape([FLAGS.batch_size, FLAGS.max_output_length])
    answer_id.set_shape([FLAGS.batch_size, FLAGS.max_output_length])
    context_mask.set_shape([FLAGS.batch_size, FLAGS.max_context_length])
    caption_mask.set_shape([FLAGS.batch_size, FLAGS.max_output_length])

    # Ensure that the random shuffling has good mixing properties.
    min_fraction_of_examples_in_queue = 0.01
    min_queue_examples = int(
        num_examples_per_epoch * min_fraction_of_examples_in_queue / FLAGS.batch_size
    )

    # Generate a batch of images and labels by building up a queue of examples.
    img_embeddings, context_lengths, caption_lengths, \
        context_ids, caption_ids, answer_ids, context_masks, \
        caption_masks = _generate_data_and_label_batch(
            [img_embedding, context_length, caption_length, \
            context_id, caption_id, answer_id, context_mask, caption_mask],
            min_queue_examples, FLAGS.batch_size,
            shuffle=(not eval_data)
        )
    img_embeddings = tf.squeeze(img_embeddings, axis = [0])
    context_lengths = tf.squeeze(context_lengths, axis = [0])
    caption_lengths = tf.squeeze(caption_lengths, axis = [0])
    context_ids = tf.squeeze(context_ids, axis = [0])
    caption_ids = tf.squeeze(caption_ids, axis = [0])
    answer_ids = tf.squeeze(answer_ids, axis = [0])
    context_masks = tf.squeeze(context_masks, axis = [0])
    caption_masks = tf.squeeze(caption_masks, axis = [0])

    # Save tensors
    img_embedding_list.append(img_embeddings)
    context_length_list.append(context_lengths)
    caption_length_list.append(caption_lengths)
    context_id_list.append(context_ids)
    caption_id_list.append(caption_ids)
    answer_id_list.append(answer_ids)
    context_mask_list.append(context_masks)
    caption_mask_list.append(caption_masks)

  return num_examples_per_epoch, img_embedding_list, context_length_list, \
      caption_length_list, context_id_list, caption_id_list, answer_id_list, \
      context_mask_list, caption_mask_list


text:
import numpy as np
from pathlib import Path
import pandas as pd
import h5py
import matplotlib.pyplot as plt
import obspy
import bisect
from pyproj import Geod
import sys
sys.path.append('./src/S_EqT_codes/src/')
sys.path.append('./src/PickNet')

import tensorflow as tf
from data_preprocessing import build_phase_dict
import keras.backend as K
import keras
keras.backend.set_floatx('float32')
import yaml
from random import shuffle
import os
import argparse
from fcn.tester import FCNTester
from scipy import signal

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='03_run_PickNet')
    parser.add_argument('--config-file', dest='config_file',
                        type=str, help='Configuration file path',default='./default_pipline_config.yaml')
    args = parser.parse_args()
    cfgs = yaml.load(open(args.config_file,'r'),Loader=yaml.SafeLoader)
    task_dir = './' + cfgs['TASKID'] + '/'
    os.chdir(task_dir)
    import tensorflow as tf
    import keras.backend.tensorflow_backend as KTF
    def get_session(gpu_fraction=0.4):
        gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=gpu_fraction)
        return tf.Session(config=tf.ConfigProto(gpu_options=gpu_options))
    session = get_session()
    KTF.set_session(session)
    #os.environ['CUDA_VISIBLE_DEVICES']  = cfgs['EqT']['gpuid']

    #session = tf.Session(config=tf.ConfigProto())
    tester = FCNTester('../src/PickNet/fcn/configs/test_config_example_P_pipeline.yaml')
    tester.setup(session)
    tester.setup_testing_for_pipeline(session)

    wavetype = 'P'
    print('On P-- {}'.format(cfgs['TASKID']))
    reftime = obspy.UTCDateTime(cfgs['REAL']['ref_time'])
    sta_file_name = cfgs['REAL']['save_sta']
    raw_folder_name = './' + cfgs['EqT']['mseed_dir']
    res_folder_name = cfgs['S_EqT']['txt_folder']

    if os.path.exists(cfgs['PickNet']['txt_folder']):
        pass
    else:
        os.makedirs(cfgs['PickNet']['txt_folder'])

    phase_dict, station_list = build_phase_dict(sta_file_name, res_folder_name, wavetype='P')
    for sta_info in station_list:
        print('On {}'.format(sta_info))
        sta_name = sta_info[1][3:]
        sta_path = Path(raw_folder_name + '/{}/'.format(sta_name))
        t_waveform = None
        t_wavekey = None

        f_picknet = open('./{}/{}.{}.txt'.format(cfgs['PickNet']['txt_folder'],sta_info[1],wavetype),'w')

        for phase_dx, phase_time in enumerate(phase_dict[sta_info[1]]['{}'.format(wavetype)]):
            seek_time = reftime + phase_time
            #print(seek_time)
            y_key = '{}'.format(seek_time.year)
            m_key = '{}'.format(seek_time.month).rjust(2,'0')
            d_key = '{}'.format(seek_time.day).rjust(2,'0')
            seek_key = '*{}*T*T*.mseed'.format(y_key+m_key+d_key)
            if t_wavekey == seek_key:
                pass
            else:
                t_wavekey = seek_key
                st = obspy.Stream()
                for waveform in sta_path.glob(seek_key):
                    st += obspy.read(str(waveform))
                st.merge()
                # st.filter('bandpass',freqmin=0.01,freqmax=20,zerophase=True)
                st.sort(keys=['channel'])
            #print(st)
            # sort order ENZ 012
            if len(st.traces) == 3:
                t_st_z = st[2].slice(seek_time - 6, seek_time + 6)
                t_st_z.resample(100.0)
                # perform PickNet - P wave
                if len(t_st_z.data[:]) < 1200:
                    refine_time = seek_time
                else:
                    input_mat = np.zeros([1,1,1200,1])
                    input_mat[0,0,:,0] = t_st_z.data[:1200]
                    input_mat -= np.mean(input_mat)
                    input_mat /= np.max(input_mat)
                    res = tester.run_on_one(input_mat, session)
                    shift_t = np.argmax(res[-1][0,0,540:660,0]) + 540

                    if shift_t <= 550 or shift_t >= 650:
                        refine_time = seek_time
                    else:
                        refine_time = seek_time - 6 + 0.01*shift_t
            else:
                refine_time = seek_time
                # calculate max amp
            write_time = refine_time - reftime
            write_line = '{:.3f} {:.5f} {:.8f}\n'.format(write_time, 1.0, 0.0)

            f_picknet.write(write_line)
        f_picknet.close()

    K.clear_session()
    session = get_session()
    KTF.set_session(session)
    #session = tf.Session(config=tf.ConfigProto())
    tester = FCNTester('../src/PickNet/fcn/configs/test_config_example_S_pipeline.yaml')
    tester.setup(session)
    tester.setup_testing_for_pipeline(session)
    wavetype = 'S'

    print('On S-- {}'.format(cfgs['TASKID']))
    reftime = obspy.UTCDateTime(cfgs['REAL']['ref_time'])
    sta_file_name = cfgs['REAL']['save_sta']
    raw_folder_name = './' + cfgs['EqT']['mseed_dir']
    res_folder_name = cfgs['S_EqT']['txt_folder']

    if os.path.exists(cfgs['PickNet']['txt_folder']):
        pass
    else:
        os.makedirs(cfgs['PickNet']['txt_folder'])

    phase_dict, station_list = build_phase_dict(sta_file_name, res_folder_name, wavetype)
    for sta_info in station_list:
        print('On {}'.format(sta_info))
        sta_name = sta_info[1][3:]
        sta_path = Path(raw_folder_name + '/{}/'.format(sta_name))
        t_waveform = None
        t_wavekey = None

        f_picknet = open('./{}/{}.{}.txt'.format(cfgs['PickNet']['txt_folder'],sta_info[1],wavetype),'w')

        for phase_dx, phase_time in enumerate(phase_dict[sta_info[1]]['{}'.format(wavetype)]):
            seek_time = reftime + phase_time
            #print(seek_time)
            y_key = '{}'.format(seek_time.year)
            m_key = '{}'.format(seek_time.month).rjust(2,'0')
            d_key = '{}'.format(seek_time.day).rjust(2,'0')
            seek_key = '*{}*T*T*.mseed'.format(y_key+m_key+d_key)
            if t_wavekey == seek_key:
                pass
            else:
                t_wavekey = seek_key
                st = obspy.Stream()
                for waveform in sta_path.glob(seek_key):
                    st += obspy.read(str(waveform))
                st.merge()
                # st.filter('bandpass',freqmin=0.01,freqmax=20,zerophase=True)
                st.sort(keys=['channel'])
            #print(st)
            # sort order ENZ 012
            if len(st.traces) == 3:
                t_st_e = st[0].slice(seek_time - 8, seek_time + 8)
                t_st_e.resample(100.0)
                t_st_n = st[1].slice(seek_time - 8, seek_time + 8)
                t_st_n.resample(100.0)
                # perform PickNet - S wave
                # print(len(t_st_e.data[:]))
                if len(t_st_e.data[:]) < 1600 or len(t_st_n.data[:]) < 1600:
                    refine_time = seek_time
                else:
                    input_mat = np.zeros([1,1,1600,2])

                    input_mat[0,0,:,0] = t_st_e.data[:1600]
                    input_mat[0,0,:,0] -= np.mean(input_mat[0,0,:,0])
                    input_mat[0,0,:,0] /= np.max(input_mat[0,0,:,0])

                    input_mat[0,0,:,1] = t_st_n.data[:1600]
                    input_mat[0,0,:,1] -= np.mean(input_mat[0,0,:,1])
                    input_mat[0,0,:,1] /= np.max(input_mat[0,0,:,1])

                    res = tester.run_on_one(input_mat, session)
                    shift_t = np.argmax(res[-1][0,0,690:911,0]) + 690
                    #print(shift_t)
                    if shift_t <= 700 or shift_t >= 900:
                        refine_time = seek_time
                    else:
                        refine_time = seek_time - 8 + 0.01*shift_t
            else:
                refine_time = seek_time
            write_time = refine_time - reftime
            write_line = '{:.3f} {:.5f} {:.8f}\n'.format(write_time, 1.0, 0.0)
            f_picknet.write(write_line)
        f_picknet.close()

max_stars_repo_name: MrXiaoXiao/ESPRH

reason: import numpy,from scipy


text:
import os
import numpy as np
import torch
import torch.nn.functional as F
from tqdm import tqdm
from test import test
from utils import LabelSmoothingLoss, GradualWarmupScheduler

class TrainLoop(object):

	def __init__(self, models_dict, optimizer_task, source_loader, test_source_loader, target_loader, nadir_slack, alpha, patience, factor, label_smoothing, warmup_its, lr_threshold, verbose=-1, cp_name=None, save_cp=False, checkpoint_path=None, checkpoint_epoch=None, cuda=True, logging=False, ablation='no', train_mode='hv'):
		if checkpoint_path is None:
			# Save to current directory
			self.checkpoint_path = os.getcwd()
		else:
			self.checkpoint_path = checkpoint_path
			if not os.path.isdir(self.checkpoint_path):
				os.mkdir(self.checkpoint_path)

		self.save_epoch_fmt_task = os.path.join(self.checkpoint_path, 'task'+cp_name) if cp_name else os.path.join(self.checkpoint_path, 'task_checkpoint_{}ep.pt')
		self.save_epoch_fmt_domain = os.path.join(self.checkpoint_path, 'Domain_{}'+cp_name) if cp_name else os.path.join(self.checkpoint_path, 'Domain_{}.pt')

		self.cuda_mode = cuda
		self.feature_extractor = models_dict['feature_extractor']
		self.task_classifier = models_dict['task_classifier']
		self.domain_discriminator_list = models_dict['domain_discriminator_list']
		self.optimizer_task = optimizer_task
		self.source_loader = source_loader
		self.test_source_loader = test_source_loader
		self.target_loader = target_loader
		self.history = {'loss_task': [], 'hypervolume': [], 'loss_domain': [], 'accuracy_source':[], 'accuracy_target':[]}
		self.cur_epoch = 0
		self.total_iter = 0
		self.nadir_slack = nadir_slack
		self.alpha = alpha
		self.ablation = ablation
		self.train_mode = train_mode
		self.device = next(self.feature_extractor.parameters()).device

		its_per_epoch = len(source_loader.dataset)//(source_loader.batch_size) + 1 if len(source_loader.dataset)%(source_loader.batch_size)>0 else len(source_loader.dataset)//(source_loader.batch_size)
		patience = patience * (1+its_per_epoch)
		self.after_scheduler_task = torch.optim.lr_scheduler.ReduceLROnPlateau(self.optimizer_task, factor=factor, patience=patience, verbose=True if verbose>0 else False, threshold=lr_threshold, min_lr=1e-7)
		self.after_scheduler_disc_list = [torch.optim.lr_scheduler.ReduceLROnPlateau(disc.optimizer, factor=factor, patience=patience, verbose=True if verbose>0 else False, threshold=lr_threshold, min_lr=1e-7) for disc in self.domain_discriminator_list]
		self.verbose = verbose
		self.save_cp = save_cp

		self.scheduler_task = GradualWarmupScheduler(self.optimizer_task, total_epoch=warmup_its, after_scheduler=self.after_scheduler_task)
		self.scheduler_disc_list = [GradualWarmupScheduler(self.domain_discriminator_list[i].optimizer, total_epoch=warmup_its, after_scheduler=sch_disc) for i, sch_disc in enumerate(self.after_scheduler_disc_list)]

		if checkpoint_epoch is not None:
			self.load_checkpoint(checkpoint_epoch)

		self.logging = logging
		if self.logging:
			from torch.utils.tensorboard import SummaryWriter
			self.writer = SummaryWriter()

		if label_smoothing > 0.0:
			self.ce_criterion = LabelSmoothingLoss(label_smoothing, lbl_set_size = 7)
		else:
			self.ce_criterion = torch.nn.CrossEntropyLoss()	#torch.nn.NLLLoss()#

		#loss_domain_discriminator = F.binary_cross_entropy_with_logits(y_predict, curr_y_domain)
		weight = torch.tensor([2.0/3.0, 1.0/3.0]).to(self.device)
			#d_cr=torch.nn.CrossEntropyLoss(weight=weight)
		self.d_cr=torch.nn.NLLLoss(weight=weight)
		#self.d_cr=  F.binary_cross_entropy_with_logits()
	#### Edit####
	def adjust_learning_rate(self,optimizer, epoch=1,every_n=700,In_lr=0.01):
		"""Sets the learning rate to the initial LR decayed by 10 every n epoch epochs"""
		every_n_epoch=every_n#n_epoch/n_step
		lr = In_lr * (0.1** (epoch // every_n_epoch))
		for param_group in optimizer.param_groups:
			param_group['lr'] = lr

		return lr
	######
	def train(self, n_epochs=1, save_every=1):

		'''
		##### Edit ####
		# init necessary objects
		num_steps = n_epochs #* (len(train_loader.dataset) / train_loader.batch_size)
		#yd = Variable(torch.from_numpy(np.hstack([np.repeat(1, int(batch_size / 2)), np.repeat(0, int(batch_size / 2))]).reshape(50, 1)))
		j = 0
		lambd=0
		max_lambd= 1 - self.alpha

		max_lr_t = list(self.optimizer_task.param_groups)[-1]['initial_lr']

		max_lr_d = list(self.domain_discriminator_list[0].optimizer.param_groups)[-1]['initial_lr']
		#0.005
		#lr_2=0
		#lr=In_lr
		#Domain_Classifier.set_lambda(lambd)
		#print(num_steps)

		'''
		max_lr_t = list(self.optimizer_task.param_groups)[-1]['initial_lr']

		max_lr_d = list(self.domain_discriminator_list[0].optimizer.param_groups)[-1]['initial_lr']

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

		while self.cur_epoch < n_epochs:

			cur_loss_task = 0
			cur_hypervolume = 0
			cur_loss_total = 0

			source_iter = tqdm(enumerate(self.source_loader), disable=False)

			'''
			## Edit ##
			## update lambda and learning rate as suggested in the paper
			p = float(j) / num_steps
			lambd = max_lambd * round(2. / (1. + np.exp(-10. * p)) - 1, 3) #I changed 10 to 7
			#lambd=lambd

			lr_t = max_lr_t / (1. + 10 * p)**0.75 #
			lr_d = max_lr_d / (1. + 10 * p)**0.75 #

			lr_t= self.adjust_learning_rate(self.optimizer_task,In_lr=lr_t)
			for _, disc in enumerate(self.domain_discriminator_list):
				lr_d= self.adjust_learning_rate(disc.optimizer,In_lr=lr_d)

			self.alpha= 1 - lambd
			#print('Alpha = {}' .format(self.alpha))
			j += 1
			'''
			lr_t = max_lr_t #/ (1. + 10 * p)**0.75 #
			lr_d = max_lr_d #/ (1. + 10 * p)**0.75 #
			'''
			lr_t= self.adjust_learning_rate(self.optimizer_task, epoch= self.cur_epoch, every_n=0.8 * n_epochs, In_lr=lr_t)
			for _, disc in enumerate(self.domain_discriminator_list):
				lr_d= self.adjust_learning_rate(disc.optimizer, epoch= self.cur_epoch, every_n=0.8 * n_epochs, In_lr=lr_d)
			'''
			##########

			print('Epoch {}/{} | Alpha = {:1.3} | Lr_task = {:1.4} | Lr_dis = {:1.4} '.format(self.cur_epoch + 1, n_epochs, self.alpha, lr_t , lr_d))

			for t, batch in source_iter:
				if self.ablation == 'all':
					cur_losses = self.train_step_ablation_all(batch)
				else:
					cur_losses = self.train_step(batch)

				self.scheduler_task.step(epoch = self.total_iter, metrics = 1.- self.history['accuracy_source'][-1] if self.cur_epoch>0 else np.inf)
				for sched in self.scheduler_disc_list:
					sched.step(epoch = self.total_iter, metrics = 1.- self.history['accuracy_source'][-1] if self.cur_epoch>0 else np.inf)

				cur_loss_task += cur_losses[0]
				cur_hypervolume += cur_losses[1]
				cur_loss_total += cur_losses[2]
				self.total_iter += 1

				if self.logging:
					self.writer.add_scalar('train/task_loss', cur_losses[0], self.total_iter)
					self.writer.add_scalar('train/hypervolume_loss', cur_losses[1], self.total_iter)
					self.writer.add_scalar('train/total_loss', cur_losses[2], self.total_iter)

			self.history['loss_task'].append(cur_loss_task/(t+1))
			self.history['hypervolume'].append(cur_hypervolume/(t+1))

			print('Current task loss: {}.'.format(cur_loss_task/(t+1)))
			print('Current hypervolume: {}.'.format(cur_hypervolume/(t+1)))

			self.history['accuracy_source'].append(test(self.test_source_loader, self.feature_extractor, self.task_classifier, self.domain_discriminator_list, self.device, source_target = 'source', epoch = self.cur_epoch, tb_writer = self.writer if self.logging else None))
			self.history['accuracy_target'].append(test(self.target_loader, self.feature_extractor, self.task_classifier, self.domain_discriminator_list, self.device, source_target = 'target', epoch = self.cur_epoch, tb_writer = self.writer if self.logging else None))

			idx = np.argmax(self.history['accuracy_source'])

			print('Valid. on SOURCE data - Current acc., best acc., best acc target, and epoch: {:0.4f}, {:0.4f}, {:0.4f}, {}'.format(self.history['accuracy_source'][-1], np.max(self.history['accuracy_source']), self.history['accuracy_target'][idx], 1+np.argmax(self.history['accuracy_source'])))
			print('Valid. on TARGET data - Current acc., best acc., and epoch: {:0.4f}, {:0.4f}, {}'.format(self.history['accuracy_target'][-1], np.max(self.history['accuracy_target']), 1+np.argmax(self.history['accuracy_target'])))

			if self.logging:
				self.writer.add_scalar('misc/LR-task', self.optimizer_task.param_groups[0]['lr'], self.total_iter)
				for i, disc in enumerate(self.domain_discriminator_list):
					self.writer.add_scalar('misc/LR-disc{}'.format(i), disc.optimizer.param_groups[0]['lr'], self.total_iter)

			self.cur_epoch += 1

			if self.save_cp and (self.cur_epoch % save_every == 0 or self.history['accuracy_target'][-1] > np.max([-np.inf]+self.history['accuracy_target'][:-1])):
				self.checkpointing()

		if self.logging:
			self.writer.close()

		idx_final = np.argmax(self.history['accuracy_source'])

		return np.max(self.history['accuracy_target']), self.history['accuracy_target'][-1]

	def train_step(self, batch):
		self.feature_extractor.train()
		self.task_classifier.train()
		for disc in self.domain_discriminator_list:
			disc = disc.train()

		x_1, x_2, x_3, y_task_1, y_task_2, y_task_3, y_domain_1, y_domain_2, y_domain_3 = batch

		x = torch.cat((x_1, x_2, x_3), dim=0)
		y_task = torch.cat((y_task_1, y_task_2, y_task_3), dim=0)
		y_domain = torch.cat((y_domain_1, y_domain_2, y_domain_3), dim=0)

		if self.cuda_mode:
			x = x.to(self.device)
			y_task = y_task.to(self.device)

		# COMPUTING FEATURES
		features = self.feature_extractor.forward(x)
		features_ = features.detach()

		# DOMAIN DISCRIMINATORS (First)
		for i, disc in enumerate(self.domain_discriminator_list):
			y_predict = disc.forward(features_).squeeze()

			curr_y_domain = torch.where(y_domain == i, torch.ones(y_domain.size(0)), torch.zeros(y_domain.size(0)))
			curr_y_domain.type_as(y_domain)
			#print(y_domain.shape,curr_y_domain.shape,y_predict.shape)
			#print(sum(curr_y_domain))
			curr_y_domain = curr_y_domain.long()
			if self.cuda_mode:
				curr_y_domain = curr_y_domain.long().to(self.device)

			#loss_domain_discriminator = F.binary_cross_entropy_with_logits(y_predict, curr_y_domain)
			#weight = torch.tensor([2.0/3.0, 1.0/3.0]).to(self.device)
			#d_cr=torch.nn.CrossEntropyLoss(weight=weight)
			#d_cr=torch.nn.NLLLoss(weight=weight)
			loss_domain_discriminator = self.d_cr (y_predict, curr_y_domain)
			#print(loss_domain_discriminator)

			if self.logging:
				self.writer.add_scalar('train/D{}_loss'.format(i), loss_domain_discriminator, self.total_iter)

			disc.optimizer.zero_grad()
			loss_domain_discriminator.backward()
			disc.optimizer.step()

		# UPDATE TASK CLASSIFIER AND FEATURE EXTRACTOR
		task_out = self.task_classifier.forward(features)

		loss_domain_disc_list = []
		loss_domain_disc_list_float = []
		for i, disc in enumerate(self.domain_discriminator_list):
			y_predict = disc.forward(features).squeeze()
			curr_y_domain = torch.where(y_domain == i, torch.zeros(y_domain.size(0)), torch.ones(y_domain.size(0)))

			curr_y_domain = curr_y_domain.long()
			if self.cuda_mode:
				curr_y_domain = curr_y_domain.long().to(self.device)

			#loss_domain_disc_list.append(F.binary_cross_entropy_with_logits(y_predict, curr_y_domain))
			#y_predict= y_predict.long().to(self.device)
			loss_domain_disc_list.append(self.d_cr(y_predict, curr_y_domain))

			loss_domain_disc_list_float.append(loss_domain_disc_list[-1].detach().item())

		if self.train_mode == 'hv':
			self.update_nadir_point(loss_domain_disc_list_float)

		hypervolume = 0
		for loss in loss_domain_disc_list:
			if self.train_mode == 'hv':
				hypervolume -= torch.log(self.nadir - loss + 1e-6)
				#hypervolume -= torch.log(loss)

			elif self.train_mode == 'avg':
				hypervolume -= loss

		task_loss = self.ce_criterion(task_out, y_task)
		loss_total = self.alpha*task_loss + (1-self.alpha)*hypervolume/len(loss_domain_disc_list)
		#loss_total = task_loss + self.alpha *hypervolume/len(loss_domain_disc_list)

		self.optimizer_task.zero_grad()
		loss_total.backward()
		self.optimizer_task.step()

		'''
		# DOMAIN DISCRIMINATORS (Second-After)
		for i, disc in enumerate(self.domain_discriminator_list):
			y_predict = disc.forward(features_).squeeze()

			curr_y_domain = torch.where(y_domain == i, torch.ones(y_domain.size(0)), torch.zeros(y_domain.size(0)))
			curr_y_domain.type_as(y_domain)
			#print(y_domain.shape,curr_y_domain.shape,y_predict.shape)
			#print(sum(curr_y_domain))
			if self.cuda_mode:
				curr_y_domain = curr_y_domain.long().to(self.device)

			loss_domain_discriminator = self.d_cr (y_predict, curr_y_domain)
			#print(loss_domain_discriminator)

			if self.logging:
				self.writer.add_scalar('train/D{}_loss'.format(i), loss_domain_discriminator, self.total_iter)

			disc.optimizer.zero_grad()
			loss_domain_discriminator.backward()
			disc.optimizer.step()

			####

		'''
		losses_return = task_loss.item(), hypervolume.item(), loss_total.item()
		return losses_return

	def train_step_ablation_all(self, batch):

		self.feature_extractor.train()
		self.task_classifier.train()
		for disc in self.domain_discriminator_list:
			disc = disc.train()

		x_1, x_2, x_3, y_task_1, y_task_2, y_task_3, _, _, _ = batch

		x = torch.cat((x_1, x_2, x_3), dim=0)
		y_task = torch.cat((y_task_1, y_task_2, y_task_3), dim=0)

		if self.cuda_mode:
			x = x.to(self.device)
			y_task = y_task.to(self.device)

		# COMPUTING FEATURES
		features = self.feature_extractor.forward(x)
		task_out = self.task_classifier.forward(features)
		task_loss = torch.nn.CrossEntropyLoss()(task_out, y_task)

		self.optimizer_task.zero_grad()
		task_loss.backward()
		self.optimizer_task.step()

		losses_return = task_loss.item(), 0

		return losses_return

	def checkpointing(self):
		if self.verbose>0:
			print(' ')
			print('Checkpointing...')

		ckpt = {'feature_extractor_state': self.feature_extractor.state_dict(),
				'task_classifier_state': self.task_classifier.state_dict(),
				'optimizer_task_state': self.optimizer_task.state_dict(),
				'scheduler_task_state': self.scheduler_task.state_dict(),
				'history': self.history,
				'cur_epoch': self.cur_epoch}
		torch.save(ckpt, self.save_epoch_fmt_task.format(self.cur_epoch))

		for i, disc in enumerate(self.domain_discriminator_list):
			ckpt = {'model_state': disc.state_dict(),
					'optimizer_disc_state': disc.optimizer.state_dict(),
					'scheduler_disc_state': self.scheduler_disc_list[i].state_dict()}
			torch.save(ckpt, self.save_epoch_fmt_domain.format(i + 1))

	def load_checkpoint(self, epoch):
		ckpt = self.save_epoch_fmt_task.format(epoch)

		if os.path.isfile(ckpt):

			ckpt = torch.load(ckpt)
			# Load model state
			self.feature_extractor.load_state_dict(ckpt['feature_extractor_state'])
			self.task_classifier.load_state_dict(ckpt['task_classifier_state'])
			self.domain_classifier.load_state_dict(ckpt['domain_classifier_state'])
			# Load optimizer state
			self.optimizer.load_state_dict(ckpt['optimizer_task_state'])
			# Load scheduler state
			self.scheduler_task.load_state_dict(ckpt['scheduler_task_state'])
			# Load history
			self.history = ckpt['history']
			self.cur_epoch = ckpt['cur_epoch']

			for i, disc in enumerate(self.domain_discriminator_list):
				ckpt = torch.load(self.save_epoch_fmt_domain.format(i + 1))
				disc.load_state_dict(ckpt['model_state'])
				disc.optimizer.load_state_dict(ckpt['optimizer_disc_state'])
				self.scheduler_disc_list[i].load_state_dict(ckpt['scheduler_disc_state'])

		else:
			print('No checkpoint found at: {}'.format(ckpt))

	def print_grad_norms(self, model):
		norm = 0.0
		for params in list(filter(lambda p: p.grad is not None, model.parameters())):
			norm += params.grad.norm(2).item()
		print('Sum of grads norms: {}'.format(norm))

	def update_nadir_point(self, losses_list):
		self.nadir = float(np.max(losses_list) * self.nadir_slack + 1e-8)


text:
import numpy as np

from perform.rom.projection_rom.linear_proj_rom.linear_proj_rom import LinearProjROM

class LinearLSPGProj(LinearProjROM):
    """Class for projection-based ROM with linear decoder and least-squares Petrov-Galerkin projection.

    Inherits from LinearProjROM.

    Trial basis is assumed to represent the conservative variables. Allows implicit time integration only.

    Args:
        model_idx: Zero-indexed ID of a given RomModel instance within a RomDomain's model_list.
        rom_domain: RomDomain within which this RomModel is contained.
        sol_domain: SolutionDomain with which this RomModel's RomDomain is associated.

    Attributes:
        trial_basis_scaled: 2D NumPy array of trial basis scaled by norm_fac_prof_cons. Precomputed for cost savings.
        hyper_reduc_operator: 2D NumPy array of gappy POD projection operator. Precomputed for cost savings.
    """

    def __init__(self, model_idx, rom_domain, sol_domain):

        # I'm not going to code LSPG with an explicit time integrator, it's a pointless exercise.
        if rom_domain.time_integrator.time_type == "explicit":
            raise ValueError(
                "LSPG with an explicit time integrator deteriorates to Galerkin, please use"
                + " Galerkin or select an implicit time integrator."
            )

        if rom_domain.time_integrator.dual_time:
            raise ValueError("LSPG is intended for conservative variable evolution, please set dual_time = False")

        super().__init__(model_idx, rom_domain, sol_domain)

        self.trial_basis_scaled = self.trial_basis * self.norm_fac_prof_cons.ravel(order="C")[:, None]

        # Precompute hyper-reduction projector, [S^T * U]^+
        # TODO: may want to have an option to compute U * [S^T * U]^+
        #   The matrix can be big, but it's necessary for conservative LSPG, different least-square problem
        if self.hyper_reduc:
            self.hyper_reduc_operator = np.linalg.pinv(self.hyper_reduc_basis[self.direct_samp_idxs_flat, :])