text:
"""Provides functionality useful for postporcessing STEM AQOUT files

class aqout_container() provides most of the functionality.  Several
helper functions are also included.
"""

import sys
import numpy as np
import os.path
import os
import pandas as pd
import netCDF4
import warnings
import cPickle
from datetime import datetime, timedelta
import itertools
import subprocess

from stem_pytools import STEM_parsers as sp
from stem_pytools.check_paths import check_path_with_msg
from stem_pytools import NERSC_data_paths as ndp
from stem_pytools import calc_drawdown

class aqout_container(object):
    """class to contain and optionally combine data from one or more
    STEM AQOUT files.

    CONTAINS FIELDS:
    aqout_paths: list or tuple; full paths to the object's AQOUT file(s)
    desc: character string; text describing the COS scenario described
       in the object's data
    key: character string; short string for use a key in a dict of
       aqout scenarios
    aq_keys: list of strings; have have same length as aqout_paths.
       Contains a short key for each individual aqout_file.
    The class also provides a str method to produce a nicely formatted
    printout of the paths it contains.

    """
    def __init__(self,
                 aqout_paths=[],
                 desc='',
                 key='',
                 aq_keys=[]):
        """class constructor"""
        if len(key) is 0:
            key = '-'.join(aq_keys)
        # if the aqout_paths argument contains a single string,
        # convert to a one-element tuple.  Doing that here rather than
        # requiring that the aqout_paths be a tuple already allows for
        # creating arguments for the constructor by, for example,
        # calling itertools.product on two lists of strings.
        if type(aqout_paths) is str:
            aqout_paths = (aqout_paths,)
        self.desc = desc
        self.key = key
        self.aq_keys = aq_keys
        self.aqout_paths = aqout_paths
        self.data = []  # list to hold data arrays
        self.t = []  # list to hold time stamp arrays
        self.cos_total = None  # field to hold cos total

    def __str__(self):
        """ formatted printing of the paths in an aqout_combiner object. """
        return('\n'.join(self.aqout_paths))

    def all_paths_exist(self):
        """Checks whether all files in object's aqout_paths exist.  If so
        returns True.  If not, writes the path(s) that do/does not
        exist to stdout and return False. """
        all_exist = all([check_path_with_msg(p) for p in self.aqout_paths])
        return(all_exist)

    def parse(self, t0, t1, verbose=False):
        """parse the concentrations in the object's AQOUT file(s) for the
        period beginning with t0 and ending with t1.  The
        concentration are placed in a list of numpy ndarray objects;
        that list is placed in the calling object's 'data' field.  The
        corresponding timestamps are also placed in a list.  That list
        is then placed in the calling object's 't' field.

        Assumes that the timestep of all AQOUT files is one hour.
        Results are undefined if this requirement is not met.

        INPUTS
        t0, t1: datetime.datetime objects specifying the starting and
            ending timestamps to combine

        """

        one_hour = 10000  # Models-3 I/O API format for 1 hour
                          # (integer, HHMMSS)

        for p in self.aqout_paths:
            nc = netCDF4.Dataset(p)
            tstep = nc.TSTEP
            nc.close()
            if tstep != one_hour:
                raise ValueError('timestep of {} is not one hour!'
                                 '(it is {})'.format(os.path.basename(p),
                                                     tstep))
            if verbose:
                sys.stdout.write('parsing {}\n'.format(os.path.basename(p)))
                sys.stdout.flush()
            this_cos = sp.parse_STEM_var(p,
                                         varname='CO2_TRACER1',
                                         t0=t0,
                                         t1=t1)
            self.data.append(this_cos['data'])
            this_t = pd.DatetimeIndex(this_cos['t'], freq='1H')
            self.t.append(this_t)

    def sum(self):
        """Add AQOUT concentration fields together.  For now uses the
        (simple-as-can-be) approach of adding together the parsed
        concentration fields element-wise.  I've implemented the
        adding as its own method so that the infrastructure is there
        in the future to handle more complicated cases (such as AQOUT
        files containing different timesteps).  """
        if len(self.data) == 0:
            raise ValueError('object {} contains no '
                             'un-summed AQOUT data'.format(self.key))

        assert all([self.t[0].equals(this_t) for this_t in self.t[1:]])
        self.t = self.t[0]

        self.cos_total = reduce(np.add, self.data)
        del self.data
        self.data = []

    def calc_stats(self):
        """

        """
        if self.cos_total is None:
            raise ValueError('object {} contains no '
                             'summed AQOUT data'.format(self.key))

        self.t_stats, self.cos_mean = daily_window_stats(self.t,
                                                         self.cos_total,
                                                         is_midday,
                                                         np.mean)
        self.t_stats, self.cos_std = daily_window_stats(self.t,
                                                        self.cos_total,
                                                        is_midday,
                                                        np.std)

    def calc_drawdown(self,
                      topo_fname=None, wrfheight_fname=None,
                      lo_height_agl=2000, hi_height_agl=4000):
        """calculate mean COS drawdown from cos_total.  Arguments are
        passed directly to calc_drawdown.calc_JA_midday_drawdown.
        """

        dd = calc_drawdown.calc_STEM_COS_drawdown(self.cos_total,
                                                  topo_fname=None,
                                                  wrfheight_fname=None)
        return dd

    def stats_to_netcdf(self, fname):

        """Write daily COS mean and standard deviation to a netCDF file.
        This is meant to replace the cPickle saving.  My thinking here
        is that these STEM results are useful, hopefully in the
        slightly longer term, and a more portable,
        platform-independent storage and transfer format is useful.

        Note: netCDF file is opened in 'write' mode -- this means an
        existing file named fname will be deleted."""

        # TO DO: check here that mean, std dev have in fact been calculated

        NC_FLOAT = 'f'  # netCDF4 specifier for NC_FLOAT datatype
        NC_INT = 'i4'  # netCDF4 specifier for NC_DOUBLE datatype

        if os.path.exists(fname):
            warnings.warn(("{} already exists. Please delete "
                           "or rename before proceeding. "
                           "Exiting; file not written".format(fname)))
            return
        nc = netCDF4.Dataset(fname, 'w')

        nc.createDimension('T', self.cos_mean.shape[0])
        nc.createDimension('LAY', self.cos_mean.shape[1])
        nc.createDimension('ROW', self.cos_mean.shape[2])
        nc.createDimension('COL', self.cos_mean.shape[3])

        nc.createVariable(varname='cos_mean',
                          datatype=NC_FLOAT,
                          dimensions=(('T', 'LAY', 'ROW', 'COL')))
        nc.createVariable(varname='cos_std',
                          datatype=NC_FLOAT,
                          dimensions=(('T', 'LAY', 'ROW', 'COL')))
        nc.createVariable(varname='time',
                          datatype=NC_INT,
                          dimensions=(('T')))

        # describe units, etc.
        nc.variables['time'].description = 'time stamp'
        nc.variables['time'].units = 'YYYYMMDDhhhmmmss'
        var = nc.variables['cos_mean']
        var.description = 'mid-day COS concentration mean'
        var.units = 'molecules cm-3'
        var = nc.variables['cos_std']
        var.description = ('mid-day COS concentration'
                           'standard deviation')
        var.units = 'molecules cm-3'
        nc.key = self.key
        nc.description = (self.desc +
                          ('July and August North '
                           'American mid-day [COS]'
                           'mean and standard deviation.  Mid-day is '
                           'defined as between 15:00 and 23:00 UTC, '
                           'which is 10:00 EST to 15:00 PST.'))

        nc.variables['cos_mean'][...] = self.cos_mean
        nc.variables['cos_std'][...] = self.cos_std
        nc.variables['time'][...] = map(
            lambda x: np.int(x.strftime('%Y%m%d%H%M%S')), self.t_stats)

        nc.close()

    def align_tstamps(self):
        """This approach will work to combine aqout files at timestamps
        that are present in both.  But it won't work to fill in timestamps
        that are only present in one.  That's OK for now because all our
        AQOUT files are hourly.  I wonder if I should plan ahead now for
        AQOUT files with a different output time step.

        In the case of different timesteps I think combining AQOUT files
        via calls to the Models-3 I/O API routines (currstep, etc) would
        be a better approach.  I wonder how difficult it would be to
        access that fortran via python...  Would certainly be quicker to
        just do the whole thing in Fortran... """

        warnings.warn('this function is a placeholder/work in progress')
        dfs = [pd.DataFrame(data=np.arange(len(t)),
                            index=t)
               for t in self.t]
        # # this works as intended
        # dfs[0].[1:8].merge(dfs[1].iloc[5:],
        #                    how='inner',
        #                    left_index=True, right_index=True)
        dfs[0].merge(dfs[1], how='inner',
                     left_index=True, right_index=True)
        return(dfs)

def is_midday(this_t):
    """
    Return true if this_t is "mid-day" in North America.  "Mid-day" is
    defined as between 15:00 and 23:00 UTC, which is 10:00 EST to
    15:00 PST.

    INPUTS
    this_t: a datetime.datetime object

    OUTPUTS
    True or False
    """
    return ((this_t.hour >= 15) and (this_t.hour <= 23))

def daily_window_stats(t, data, f_wndw, f_stat):
    """
    calculate a statistic within a daily time window.

    INPUT PARAMETERS
    data: np.ndarray of shape [t,z,x,y] with axes time, vertical
       coordinate, horizontal coordinate 1, horizontal coordinate 2.
    t: pandas DatetimeIndex of length t
    f_wndw: a function that accepts a datetime.datetime object and
       returns true if the object is inside the desired daily window and
       false if outside.
    f_stat: function to calculate the statistic.  For example np.mean
       or np.std.  Must accept "axis=0" as an argument.

    OUTPUT PARAMETERS
    data_out: np.ndarray of shape [t1, z, x, y] with the calculated
       statistic
    t_out: np.ndarray of datetime.date objects with the date of the
       data in data_out
    """
    idx_wndw = np.where(t.to_series().apply(f_wndw))[0]
    t = t[idx_wndw]
    data = data[idx_wndw, ...]

    t_date = t.to_series().apply(datetime.date)
    idx_split = np.where(t_date.diff() > timedelta(0))[0]
    lst = np.split(data, idx_split)
    data_lst = list(itertools.imap(lambda x: f_stat(x, axis=0), lst))

    data_out = np.array(data_lst)
    t_out = t_date.unique()

    return(t_out, data_out)

def assemble_data(model_runs=None, pickle_fname=None):
    """calculate and save to a cPickle file:
    (1) daily mid-day [COS] mean
    (2) daily mid-day [COS] standard deviation
    (3) timestamps associated with (1) and (2)

    the data are placed in a dict and saved to a user-specified
    cpickle file.

    INPUT PARAMETERS
    model_runs: dict of STEMRun objects.  If unspecified the default
        is the output of stem_pytools.NERSC_data_paths.get_runs()
    pickle_fname: full path of the cpickle file to create.  If
        unspecified the default is
        /home/thilton/Data/STEM/aq_out_data.cpickle

    OUTPUT PARAMETERS:
    all_data: dict containing (1), (2), and (3) above.

    """

    warnings.warn('assemble_data is deprecated and will be removed in'
                  ' the future.  Please use '
                  'aqout_postprocess.aqout_container instead')

    if model_runs is None:
        model_runs = ndp.get_runs()

    t = []
    cos_mean = []
    cos_std = []

    for k, run in model_runs.items():
        t0 = datetime.now()
        print 'processing {}'.format(k)
        this_cos = sp.parse_STEM_var(run.aqout_path,
                                     varname='CO2_TRACER1',
                                     t0=datetime(2008, 7, 8),
                                     t1=datetime(2008, 8, 31, 23, 59, 59))
        t_data = pd.DatetimeIndex(this_cos['t'], freq='1H')

        this_t, this_mean = daily_window_stats(t_data,
                                               this_cos['data'],
                                               is_midday,
                                               np.mean)
        this_t, this_std = daily_window_stats(t_data,
                                              this_cos['data'],
                                              is_midday,
                                              np.std)
        t.append(this_t)
        cos_mean.append(this_mean)
        cos_std.append(this_std)
        print 'finished {}, time: {}'.format(k, str(datetime.now() - t0))

    t_dict = {key: value for key, value in zip(model_runs.keys(), t)}
    cos_mean_dict = {key: value for key, value in zip(model_runs.keys(),
                                                      cos_mean)}
    cos_std_dict = {key: value for key, value in zip(model_runs.keys(),
                                                     cos_std)}
    all_data_dict = {'t': t_dict,
                     'cos_mean': cos_mean_dict,
                     'cos_std': cos_std_dict}
    outfile = open(pickle_fname, 'wb')
    cPickle.dump(all_data_dict, outfile, protocol=2)
    outfile.close()

    return(all_data_dict)

def load_aqout_data(fname='/home/thilton/Data/STEM/aq_out_data.cpickle'):
    """
    load a cpickle data file containing aqout data written by assemble_data
    """
    warnings.warn('load_aqout_data is deprecated and will be removed in'
                  ' the future.  Please use '
                  'aqout_postprocess.aqout_container instead')

    import cPickle
    f = open(fname, 'rb')
    all_data = cPickle.load(f)
    f.close()
    return(all_data)

def combine_aqout_netcdf_files(fnames, fname_out="out.nc"):
    """combine multiple aqout_container netcdf files to a single
    netcdf file.

    Combines multiple aqout_container netcdf files produced by
    aqout_container.stats_to_netcdf() into a single netcdf file.  The
    datasets from each separate netcdf file are placed into netcdf
    groups.

    Depends on ncecat from the `NCO operators <http://nco.sourceforge.net/>`_

    ARGS
    fnames (list): list of strings containing full paths of netcdf
        files to be combined.
    fname_out (string): full path to the combined netcdf file to
        create.  Default is "./out.nc"

    RETURNS:
    0 on success
    """

    for f in fnames:
        if not os.path.exists(f):
            raise IOError("{} not found.".format(f))

    cmd = "ncecat --gag {} {}".format(" ".join(fnames), fname_out)
    subprocess.call(cmd, shell=True)

max_stars_repo_name: Timothy-W-Hilton/STEMPyTools

max_issues_repo_name: Timothy-W-Hilton/STEMPyTools

max_forks_repo_name: Timothy-W-Hilton/STEMPyTools


text:
#!/usr/bin/python

import os
import subprocess
import string

from config import *
import numpy as np
import csv
import locale
import json
from benchmarkParser import *

def main():
    print("Parsing benchmarks")

    for key, value in cmdConfigs.items():
        for mode in layoutMode:
            print("#######")
            print("##%s -- %s####" % (key, mode))
            print("##############")

            filename=results_folder + "/" + key + "_" + mode
            print(filename)

            # Read stdout from file
            stdout_filename = filename + ".out"

    	    # parse stdout
            with open (stdout_filename, "r") as infile:
                stdoutStr=infile.read()

            data = parse_stdout(stdoutStr);

            # Wrtie CSV to file

            # Strip away type information
            outputData = {}
            for k, v in data.items():
                arr = v[1]
                if len(arr) < num_iterations:
                    arr = [None] * (num_iterations - len(arr))
                outputData[k] = arr

            zd = zip(*outputData.values())
            with open(filename + ".csv", 'w') as file:
                writer = csv.writer(file, delimiter=',')
                writer.writerow(outputData.keys())
                writer.writerows(zd)

            # Write JSON to file
            data["scene"] = key
            data["options"] = value
            data["stdout"] = stdout_filename
            data["mode"] = mode

            with open(filename+".json", 'w') as outfile:
                json.dump(data, outfile)

if __name__ == "__main__":
    main()

ext: py

reason: import numpy


text:
#!/usr/bin/env python

import numpy as np
from numpy import *

import os

# pyyaml - https://pyyaml.org/wiki/PyYAMLDocumentation
import yaml

# ROS

import rospy

import rospkg

import std_msgs.msg
from std_msgs.msg import Bool
from std_msgs.msg import Header

import geometry_msgs.msg
from geometry_msgs.msg import PoseStamped

import visualization_msgs.msg
from visualization_msgs.msg import Marker
from visualization_msgs.msg import MarkerArray

import tf_conversions

import tf2_ros

#
import ars_lib_helpers

class ArsSimEnvironmentRos:

  #######

  # World frame
  world_frame = None

  #
  environment_descript = None
  environment_descript_yaml_file_name = None

  #
  flag_dynamic_obstacles = None
  flag_dyn_obst_sub = None

  # Dynam obst loop freq
  # time step
  dynamic_obst_loop_freq = None
  # Timer
  dynamic_obst_loop_timer = None

  # Static obst loop freq
  # time step
  static_obst_loop_freq = None
  # Timer
  static_obst_loop_timer = None

  # Obstacles static pub
  obstacles_static_pub = None
  # Obstacles dynamic pub
  obstacles_dynamic_pub = None

  #
  obstacles_static_msg = None
  #
  obstacles_dynamic_msg = None

  #
  id_first_available = None

  #########

  def __init__(self):

    # World frame
    self.world_frame = 'world'

    #
    self.environment_descript = None
    self.environment_descript_yaml_file_name = ''

    #
    self.flag_dynamic_obstacles = True
    self.flag_dyn_obst_sub = None

    # Dynam obst loop freq
    # time step
    self.dynamic_obst_loop_freq = 10.0
    # Timer
    self.dynamic_obst_loop_timer = None

    # Static obst loop freq
    # time step
    self.static_obst_loop_freq = 10.0
    # Timer
    self.static_obst_loop_timer = None

    #
    self.obstacles_static_msg = MarkerArray()
    #
    self.obstacles_dynamic_msg = MarkerArray()

    #
    self.id_first_available = 0

    # end
    return

  def init(self, node_name='ars_sim_environment_node'):
    #

    # Init ROS
    rospy.init_node(node_name, anonymous=True)

    # Package path
    pkg_path = rospkg.RosPack().get_path('ars_sim_environment')

    #### READING PARAMETERS ###

    # Environment description
    default_environment_descript_yaml_file_name = os.path.join(pkg_path,'config','obstacles_env_01.yaml')
    environment_descript_yaml_file_name_str = rospy.get_param('~environment_description_yaml_file', default_environment_descript_yaml_file_name)
    print(environment_descript_yaml_file_name_str)
    self.environment_descript_yaml_file_name = os.path.abspath(environment_descript_yaml_file_name_str)

    ###

    # Load environment description
    with open(self.environment_descript_yaml_file_name,'r') as file:
        # The FullLoader parameter handles the conversion from YAML
        # scalar values to Python the dictionary format
        self.environment_descript = yaml.load(file)

    # Obstacles static
    self.initObstaclesStatic()

    # Obstacles dynamic
    self.initObstaclesDynamic()

    # End
    return

  def open(self):

    # Publishers

    #
    self.obstacles_static_pub = rospy.Publisher('obstacles_static', MarkerArray, queue_size=1, latch=True)
    #
    self.obstacles_dynamic_pub = rospy.Publisher('obstacles_dynamic', MarkerArray, queue_size=1, latch=True)

    # Subscribers
    self.flag_dyn_obst_sub = rospy.Subscriber('flag_dynamic_obstacles', Bool, self.flagDynamObstCallback)

    # Timers
    #
    self.dynamic_obst_loop_timer = rospy.Timer(rospy.Duration(1.0/self.dynamic_obst_loop_freq), self.dynamicObstaclesLoopTimerCallback)
    #
    self.static_obst_loop_timer = rospy.Timer(rospy.Duration(1.0), self.staticObstaclesLoopTimerCallback, oneshot=True)

    # End
    return

  def run(self):

    rospy.spin()

    return

  def flagDynamObstCallback(self, flag_dynamic_obstacles_msg):

    self.flag_dynamic_obstacles = flag_dynamic_obstacles_msg.data

    if(self.flag_dynamic_obstacles == True):
      self.initObstaclesDynamic()

    else:
      self.emptyObstaclesDynamic()

    # Publish
    self.publishObstacleDynamic()

    return

  def publishObstacleDynamic(self):

    self.obstacles_dynamic_pub.publish(self.obstacles_dynamic_msg)

    return

  def updateDynObstaclesAndPublish(self, time_stamp_current):

    #
    if(self.flag_dynamic_obstacles):

        # Update
        # TODO

        # Publish dynamic obstacles
        self.publishObstacleDynamic()

    # End
    return

  def staticObstaclesLoopTimerCallback(self, timer_msg):

    # Get time
    time_stamp_current = rospy.Time.now()

    # Publish static obstacles
    self.obstacles_static_pub.publish(self.obstacles_static_msg)

    # End
    return

  def dynamicObstaclesLoopTimerCallback(self, timer_msg):

    # Get time
    time_stamp_current = rospy.Time.now()

    #
    self.updateDynObstaclesAndPublish(time_stamp_current)

    # End
    return

  def initObstaclesStatic(self):

    # Set static obstacles

    self.obstacles_static_msg.markers = []

    #
    print('Obstacles static:')

    for object_env in self.environment_descript['static']:

        print('Circles:')

        for circle in object_env['circles']:

            #
            print(circle)

            # obstacle i

            obstacle_i = Marker()

            obstacle_i.header = Header()
            obstacle_i.header.stamp = rospy.Time()
            obstacle_i.header.frame_id = self.world_frame

            obstacle_i.ns = 'static'
            obstacle_i.id = self.id_first_available

            obstacle_i.action = 0

            obstacle_i.type = 3

            obstacle_i.pose.position.x = circle['position'][0]
            obstacle_i.pose.position.y = circle['position'][1]
            obstacle_i.pose.position.z = circle['position'][2]

            obstacle_i.pose.orientation.w = 1.0
            obstacle_i.pose.orientation.x = 0.0
            obstacle_i.pose.orientation.y = 0.0
            obstacle_i.pose.orientation.z = 0.0

            obstacle_i.scale.x = circle['sizes'][0]
            obstacle_i.scale.y = circle['sizes'][1]
            obstacle_i.scale.z = circle['sizes'][2]

            obstacle_i.color.r = 1.0
            obstacle_i.color.g = 0.0
            obstacle_i.color.b = 0.0
            obstacle_i.color.a = 0.3

            obstacle_i.lifetime = rospy.Duration(0.0)

            #
            self.obstacles_static_msg.markers.append(obstacle_i)

            #
            self.id_first_available += 1

    #
    return

  def removeObstaclesStatic(self):

    self.obstacles_static_msg.markers = []

    return

  def addObstaclesStatic(self):

    for obs_i_msg in self.obstacles_static_msg.markers:

      obs_i_msg.action = 0

    return

  def removeObstaclesStatic(self):

    for obs_i_msg in self.obstacles_static_msg.markers:

      obs_i_msg.action = 2

    return

  def initObstaclesDynamic(self):

    # Set dynamic obstacles

    self.obstacles_dynamic_msg.markers = []

    #
    print('Obstacles dynamic:')

    for object_env in self.environment_descript['dynamic']:

        print('Circles:')

        for circle in object_env['circles']:

            #
            print(circle)

            # obstacle i

            obstacle_i = Marker()

            obstacle_i.header = Header()
            obstacle_i.header.stamp = rospy.Time()
            obstacle_i.header.frame_id = self.world_frame

            obstacle_i.ns = 'dynamic'
            obstacle_i.id = self.id_first_available

            obstacle_i.action = 0

            obstacle_i.type = 3

            obstacle_i.pose.position.x = circle['position'][0]
            obstacle_i.pose.position.y = circle['position'][1]
            obstacle_i.pose.position.z = circle['position'][2]

            obstacle_i.pose.orientation.w = 1.0
            obstacle_i.pose.orientation.x = 0.0
            obstacle_i.pose.orientation.y = 0.0
            obstacle_i.pose.orientation.z = 0.0

            obstacle_i.scale.x = circle['sizes'][0]
            obstacle_i.scale.y = circle['sizes'][1]
            obstacle_i.scale.z = circle['sizes'][2]

            obstacle_i.color.r = 0.0
            obstacle_i.color.g = 1.0
            obstacle_i.color.b = 0.0
            obstacle_i.color.a = 0.3

            obstacle_i.lifetime = rospy.Duration(1.0/self.dynamic_obst_loop_freq)

            #
            self.obstacles_dynamic_msg.markers.append(obstacle_i)

            #
            self.id_first_available += 1

    return

  def emptyObstaclesDynamic(self):

    self.obstacles_dynamic_msg.markers = []

    return

reason: import numpy,from numpy


text:
# Copyright 2020 The Magenta Authors.
#
# 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.

# Lint as: python3
"""SketchRNN training."""

import json
import os
import time
import zipfile

import model as sketch_rnn_model
import utils
import numpy as np
import requests
import six
from six.moves.urllib.request import urlretrieve
import tensorflow.compat.v1 as tf

tf.logging.set_verbosity(tf.logging.INFO)

FLAGS = tf.app.flags.FLAGS

tf.app.flags.DEFINE_string(
    'data_dir',
    'https://github.com/hardmaru/sketch-rnn-datasets/raw/master/aaron_sheep',
    'The directory in which to find the dataset specified in model hparams. '
    'If data_dir starts with "http://" or "https://", the file will be fetched '
    'remotely.')
tf.app.flags.DEFINE_string(
    'log_root', '/tmp/sketch_rnn/models/default',
    'Directory to store model checkpoints, tensorboard.')
tf.app.flags.DEFINE_boolean('resume_training', False,
                            'Set to true to load previous checkpoint')
tf.app.flags.DEFINE_string(
    'hparams', '', 'Pass in comma-separated key=value pairs such as '
    '\'save_every=40,decay_rate=0.99\' '
    '(no whitespace) to be read into the HParams object defined in model.py')

PRETRAINED_MODELS_URL = ('http://download.magenta.tensorflow.org/models/'
                         'sketch_rnn.zip')

def reset_graph():
    """Closes the current default session and resets the graph."""
    sess = tf.get_default_session()
    if sess:
        sess.close()
    tf.reset_default_graph()

def load_env(data_dir, model_dir):
    """Loads environment for inference mode, used in jupyter notebook."""
    model_params = sketch_rnn_model.get_default_hparams()
    with tf.gfile.Open(os.path.join(model_dir, 'model_config.json'), 'r') as f:
        model_params.parse_json(f.read())
    return load_dataset(data_dir, model_params, inference_mode=True)

def load_model(model_dir):
    """Loads model for inference mode, used in jupyter notebook."""
    model_params = sketch_rnn_model.get_default_hparams()
    with tf.gfile.Open(os.path.join(model_dir, 'model_config.json'), 'r') as f:
        model_params.parse_json(f.read())

    model_params.batch_size = 1  # only sample one at a time
    eval_model_params = sketch_rnn_model.copy_hparams(model_params)
    eval_model_params.use_input_dropout = 0
    eval_model_params.use_recurrent_dropout = 0
    eval_model_params.use_output_dropout = 0
    eval_model_params.is_training = 0
    sample_model_params = sketch_rnn_model.copy_hparams(eval_model_params)
    sample_model_params.max_seq_len = 1  # sample one point at a time
    return [model_params, eval_model_params, sample_model_params]

def download_pretrained_models(models_root_dir='/tmp/sketch_rnn/models',
                               pretrained_models_url=PRETRAINED_MODELS_URL):
    """Download pretrained models to a temporary directory."""
    tf.gfile.MakeDirs(models_root_dir)
    zip_path = os.path.join(models_root_dir,
                            os.path.basename(pretrained_models_url))
    if os.path.isfile(zip_path):
        tf.logging.info('%s already exists, using cached copy', zip_path)
    else:
        tf.logging.info('Downloading pretrained models from %s...',
                        pretrained_models_url)
        urlretrieve(pretrained_models_url, zip_path)
        tf.logging.info('Download complete.')
    tf.logging.info('Unzipping %s...', zip_path)
    with zipfile.ZipFile(zip_path) as models_zip:
        models_zip.extractall(models_root_dir)
    tf.logging.info('Unzipping complete.')

def load_dataset(data_dir, model_params, inference_mode=False):
    """Loads the .npz file, and splits the set into train/valid/test."""

    # normalizes the x and y columns using the training set.
    # applies same scaling factor to valid and test set.

    if isinstance(model_params.data_set, list):
        datasets = model_params.data_set
    else:
        datasets = [model_params.data_set]

    train_strokes = None
    valid_strokes = None
    test_strokes = None

    for dataset in datasets:
        if data_dir.startswith('http://') or data_dir.startswith('https://'):
            data_filepath = '/'.join([data_dir, dataset])
            tf.logging.info('Downloading %s', data_filepath)
            response = requests.get(data_filepath)
            data = np.load(six.BytesIO(response.content), encoding='latin1')
        else:
            data_filepath = os.path.join(data_dir, dataset)
            data = np.load(data_filepath, encoding='latin1', allow_pickle=True)
        tf.logging.info('Loaded {}/{}/{} from {}'.format(
            len(data['train']), len(data['valid']), len(data['test']),
            dataset))
        if train_strokes is None:
            train_strokes = data['train']
            valid_strokes = data['valid']
            test_strokes = data['test']
        else:
            train_strokes = np.concatenate((train_strokes, data['train']))
            valid_strokes = np.concatenate((valid_strokes, data['valid']))
            test_strokes = np.concatenate((test_strokes, data['test']))

    all_strokes = np.concatenate((train_strokes, valid_strokes, test_strokes))
    num_points = 0
    for stroke in all_strokes:
        num_points += len(stroke)
    avg_len = num_points / len(all_strokes)
    tf.logging.info('Dataset combined: {} ({}/{}/{}), avg len {}'.format(
        len(all_strokes), len(train_strokes), len(valid_strokes),
        len(test_strokes), int(avg_len)))

    # calculate the max strokes we need.
    max_seq_len = utils.get_max_len(all_strokes)
    # overwrite the hps with this calculation.
    model_params.max_seq_len = max_seq_len

    tf.logging.info('model_params.max_seq_len %i.', model_params.max_seq_len)

    eval_model_params = sketch_rnn_model.copy_hparams(model_params)

    eval_model_params.use_input_dropout = 0
    eval_model_params.use_recurrent_dropout = 0
    eval_model_params.use_output_dropout = 0
    eval_model_params.is_training = 1

    if inference_mode:
        eval_model_params.batch_size = 1
        eval_model_params.is_training = 0

    sample_model_params = sketch_rnn_model.copy_hparams(eval_model_params)
    sample_model_params.batch_size = 1  # only sample one at a time
    sample_model_params.max_seq_len = 1  # sample one point at a time

    train_set = utils.DataLoader(
        train_strokes,
        model_params.batch_size,
        max_seq_length=model_params.max_seq_len,
        random_scale_factor=model_params.random_scale_factor,
        augment_stroke_prob=model_params.augment_stroke_prob)

    normalizing_scale_factor = train_set.calculate_normalizing_scale_factor()
    train_set.normalize(normalizing_scale_factor)

    valid_set = utils.DataLoader(valid_strokes,
                                 eval_model_params.batch_size,
                                 max_seq_length=eval_model_params.max_seq_len,
                                 random_scale_factor=0.0,
                                 augment_stroke_prob=0.0)
    valid_set.normalize(normalizing_scale_factor)

    test_set = utils.DataLoader(test_strokes,
                                eval_model_params.batch_size,
                                max_seq_length=eval_model_params.max_seq_len,
                                random_scale_factor=0.0,
                                augment_stroke_prob=0.0)
    test_set.normalize(normalizing_scale_factor)

    tf.logging.info('normalizing_scale_factor %4.4f.',
                    normalizing_scale_factor)

    result = [
        train_set, valid_set, test_set, model_params, eval_model_params,
        sample_model_params
    ]
    return result

def evaluate_model(sess, model, data_set):
    """Returns the average weighted cost, reconstruction cost and KL cost."""
    total_cost = 0.0
    total_r_cost = 0.0
    total_kl_cost = 0.0
    for batch in range(data_set.num_batches):
        unused_orig_x, x, s = data_set.get_batch(batch)
        feed = {model.input_data: x, model.sequence_lengths: s}
        (cost, r_cost,
         kl_cost) = sess.run([model.cost, model.r_cost, model.kl_cost], feed)
        total_cost += cost
        total_r_cost += r_cost
        total_kl_cost += kl_cost

    total_cost /= (data_set.num_batches)
    total_r_cost /= (data_set.num_batches)
    total_kl_cost /= (data_set.num_batches)
    return (total_cost, total_r_cost, total_kl_cost)

def load_checkpoint(sess, checkpoint_path):
    saver = tf.train.Saver(tf.global_variables())
    ckpt = tf.train.get_checkpoint_state(checkpoint_path)
    tf.logging.info('Loading model %s.', ckpt.model_checkpoint_path)
    saver.restore(sess, ckpt.model_checkpoint_path)

def save_model(sess, model_save_path, global_step):
    saver = tf.train.Saver(tf.global_variables())
    checkpoint_path = os.path.join(model_save_path, 'vector')
    tf.logging.info('saving model %s.', checkpoint_path)
    tf.logging.info('global_step %i.', global_step)
    saver.save(sess, checkpoint_path, global_step=global_step)

def train(sess, model, eval_model, train_set, valid_set, test_set):
    """Train a sketch-rnn model."""
    # Setup summary writer.
    summary_writer = tf.summary.FileWriter(FLAGS.log_root)

    # Calculate trainable params.
    t_vars = tf.trainable_variables()
    count_t_vars = 0
    for var in t_vars:
        num_param = np.prod(var.get_shape().as_list())
        count_t_vars += num_param
        tf.logging.info('%s %s %i', var.name, str(var.get_shape()), num_param)
    tf.logging.info('Total trainable variables %i.', count_t_vars)
    model_summ = tf.summary.Summary()
    model_summ.value.add(tag='Num_Trainable_Params',
                         simple_value=float(count_t_vars))
    summary_writer.add_summary(model_summ, 0)
    summary_writer.flush()

    # setup eval stats
    best_valid_cost = 100000000.0  # set a large init value
    valid_cost = 0.0

    # main train loop

    hps = model.hps
    start = time.time()

    for _ in range(hps.num_steps):

        step = sess.run(model.global_step)

        curr_learning_rate = ((hps.learning_rate - hps.min_learning_rate) *
                              (hps.decay_rate)**step + hps.min_learning_rate)
        curr_kl_weight = (hps.kl_weight -
                          (hps.kl_weight - hps.kl_weight_start) *
                          (hps.kl_decay_rate)**step)

        _, x, s = train_set.random_batch()
        feed = {
            model.input_data: x,
            model.sequence_lengths: s,
            model.lr: curr_learning_rate,
            model.kl_weight: curr_kl_weight
        }

        (train_cost, r_cost, kl_cost, _, train_step, _) = sess.run([
            model.cost, model.r_cost, model.kl_cost, model.final_state,
            model.global_step, model.train_op
        ], feed)

        if step % 20 == 0 and step > 0:

            end = time.time()
            time_taken = end - start

            cost_summ = tf.summary.Summary()
            cost_summ.value.add(tag='Train_Cost',
                                simple_value=float(train_cost))
            reconstr_summ = tf.summary.Summary()
            reconstr_summ.value.add(tag='Train_Reconstr_Cost',
                                    simple_value=float(r_cost))
            kl_summ = tf.summary.Summary()
            kl_summ.value.add(tag='Train_KL_Cost', simple_value=float(kl_cost))
            lr_summ = tf.summary.Summary()
            lr_summ.value.add(tag='Learning_Rate',
                              simple_value=float(curr_learning_rate))
            kl_weight_summ = tf.summary.Summary()
            kl_weight_summ.value.add(tag='KL_Weight',
                                     simple_value=float(curr_kl_weight))
            time_summ = tf.summary.Summary()
            time_summ.value.add(tag='Time_Taken_Train',
                                simple_value=float(time_taken))

            output_format = ('step: %d, lr: %.6f, klw: %0.4f, cost: %.4f, '
                             'recon: %.4f, kl: %.4f, train_time_taken: %.4f')
            output_values = (step, curr_learning_rate, curr_kl_weight,
                             train_cost, r_cost, kl_cost, time_taken)
            output_log = output_format % output_values

            tf.logging.info(output_log)

            summary_writer.add_summary(cost_summ, train_step)
            summary_writer.add_summary(reconstr_summ, train_step)
            summary_writer.add_summary(kl_summ, train_step)
            summary_writer.add_summary(lr_summ, train_step)
            summary_writer.add_summary(kl_weight_summ, train_step)
            summary_writer.add_summary(time_summ, train_step)
            summary_writer.flush()
            start = time.time()

        if step % hps.save_every == 0 and step > 0:

            (valid_cost, valid_r_cost,
             valid_kl_cost) = evaluate_model(sess, eval_model, valid_set)

            end = time.time()
            time_taken_valid = end - start
            start = time.time()

            valid_cost_summ = tf.summary.Summary()
            valid_cost_summ.value.add(tag='Valid_Cost',
                                      simple_value=float(valid_cost))
            valid_reconstr_summ = tf.summary.Summary()
            valid_reconstr_summ.value.add(tag='Valid_Reconstr_Cost',
                                          simple_value=float(valid_r_cost))
            valid_kl_summ = tf.summary.Summary()
            valid_kl_summ.value.add(tag='Valid_KL_Cost',
                                    simple_value=float(valid_kl_cost))
            valid_time_summ = tf.summary.Summary()
            valid_time_summ.value.add(tag='Time_Taken_Valid',
                                      simple_value=float(time_taken_valid))

            output_format = (
                'best_valid_cost: %0.4f, valid_cost: %.4f, valid_recon: '
                '%.4f, valid_kl: %.4f, valid_time_taken: %.4f')
            output_values = (min(best_valid_cost, valid_cost), valid_cost,
                             valid_r_cost, valid_kl_cost, time_taken_valid)
            output_log = output_format % output_values

            tf.logging.info(output_log)

            summary_writer.add_summary(valid_cost_summ, train_step)
            summary_writer.add_summary(valid_reconstr_summ, train_step)
            summary_writer.add_summary(valid_kl_summ, train_step)
            summary_writer.add_summary(valid_time_summ, train_step)
            summary_writer.flush()

            if valid_cost < best_valid_cost:
                best_valid_cost = valid_cost

                save_model(sess, FLAGS.log_root, step)

                end = time.time()
                time_taken_save = end - start
                start = time.time()

                tf.logging.info('time_taken_save %4.4f.', time_taken_save)

                best_valid_cost_summ = tf.summary.Summary()
                best_valid_cost_summ.value.add(
                    tag='Best_Valid_Cost', simple_value=float(best_valid_cost))

                summary_writer.add_summary(best_valid_cost_summ, train_step)
                summary_writer.flush()

                (eval_cost, eval_r_cost,
                 eval_kl_cost) = evaluate_model(sess, eval_model, test_set)

                end = time.time()
                time_taken_eval = end - start
                start = time.time()

                eval_cost_summ = tf.summary.Summary()
                eval_cost_summ.value.add(tag='Eval_Cost',
                                         simple_value=float(eval_cost))
                eval_reconstr_summ = tf.summary.Summary()
                eval_reconstr_summ.value.add(tag='Eval_Reconstr_Cost',
                                             simple_value=float(eval_r_cost))
                eval_kl_summ = tf.summary.Summary()
                eval_kl_summ.value.add(tag='Eval_KL_Cost',
                                       simple_value=float(eval_kl_cost))
                eval_time_summ = tf.summary.Summary()
                eval_time_summ.value.add(tag='Time_Taken_Eval',
                                         simple_value=float(time_taken_eval))

                output_format = ('eval_cost: %.4f, eval_recon: %.4f, '
                                 'eval_kl: %.4f, eval_time_taken: %.4f')
                output_values = (eval_cost, eval_r_cost, eval_kl_cost,
                                 time_taken_eval)
                output_log = output_format % output_values

                tf.logging.info(output_log)

                summary_writer.add_summary(eval_cost_summ, train_step)
                summary_writer.add_summary(eval_reconstr_summ, train_step)
                summary_writer.add_summary(eval_kl_summ, train_step)
                summary_writer.add_summary(eval_time_summ, train_step)
                summary_writer.flush()

def trainer(model_params):
    """Train a sketch-rnn model."""
    np.set_printoptions(precision=8, edgeitems=6, linewidth=200, suppress=True)

    tf.logging.info('sketch-rnn')
    tf.logging.info('Hyperparams:')
    tf.logging.info('Loading data files.')
    datasets = load_dataset(FLAGS.data_dir, model_params)

    train_set = datasets[0]
    valid_set = datasets[1]
    test_set = datasets[2]
    model_params = datasets[3]
    eval_model_params = datasets[4]

    reset_graph()
    model = sketch_rnn_model.Model(model_params)
    eval_model = sketch_rnn_model.Model(eval_model_params, reuse=True)

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

    if FLAGS.resume_training:
        load_checkpoint(sess, FLAGS.log_root)

    # Write config file to json file.
    tf.gfile.MakeDirs(FLAGS.log_root)
    with tf.gfile.Open(os.path.join(FLAGS.log_root, 'model_config.json'),
                       'w') as f:
        json.dump(list(model_params.values()), f, indent=True)

    train(sess, model, eval_model, train_set, valid_set, test_set)

def main(unused_argv):
    """Load model params, save config file and start trainer."""
    model_params = sketch_rnn_model.get_default_hparams()
    if FLAGS.hparams:
        model_params.parse(FLAGS.hparams)
    trainer(model_params)

def console_entry_point():
    tf.disable_v2_behavior()
    tf.app.run(main)

if __name__ == '__main__':
    console_entry_point()

max_stars_repo_name: JeffreyCA/Skribbl-RNN

max_issues_repo_name: JeffreyCA/Skribbl-RNN

max_forks_repo_name: JeffreyCA/Skribbl-RNN


text:
import torch
import torch.nn as nn
import time, pdb
import numpy as np

from util import  py_softmax,MovingAverage
from multigpu import gpu_mul_Ax, gpu_mul_xA, aggreg_multi_gpu, gpu_mul_AB

def cpu_sk(o1):
    """ Sinkhorn Knopp optimization on CPU
        * stores activations to RAM
        * does matrix-vector multiplies on CPU
        * slower than GPU
    """
    # 1. aggregate inputs:
    N = len(o1.pseudo_loader.dataset)
    if o1.hc == 1:
        o1.PS = np.zeros((N, o1.K), dtype=o1.dtype)
    else:
        o1.PS_pre = np.zeros((N, o1.presize), dtype=o1.dtype)
    now = time.time()
    l_dl = len(o1.pseudo_loader)
    time.time()
    batch_time = MovingAverage(intertia=0.9)
    o1.model.headcount = 1
    for batch_idx, (data, _, _selected) in enumerate(o1.pseudo_loader):
        data = data.to(o1.dev)
        mass = data.size(0)
        if o1.hc == 1:
            p = nn.functional.softmax(o1.model(data), 1)
            o1.PS[_selected, :] = p.detach().cpu().numpy().astype(o1.dtype)
        else:
            p = o1.model(data)
            o1.PS_pre[_selected, :] = p.detach().cpu().numpy().astype(o1.dtype)
        batch_time.update(time.time() - now)
        now = time.time()
        if batch_idx % 50 == 0:
            print(f"Aggregating batch {batch_idx:03}/{l_dl}, speed: {mass / batch_time.avg:04.1f}Hz",
                  end='\r', flush=True)
    o1.model.headcount = o1.hc
    print("Aggreg of outputs  took {0:.2f} min".format((time.time() - now) / 60.), flush=True)

    # 2. solve label assignment via sinkhorn-knopp:
    if o1.hc == 1:
        o1 = optimize_L_sk(o1, nh=0)
    else:
        for nh in range(o1.hc):
            print("computing head %s " % nh, end="\r", flush=True)
            tl = getattr(o1.model, "top_layer%d" % nh)
            time_mat = time.time()

            # clear memory
            try:
                del o1.PS
            except:
                pass

            # apply last FC layer (a matmul and adding of bias)
            o1.PS = (o1.PS_pre @ tl.weight.cpu().numpy().T.astype(o1.dtype)
                       + tl.bias.cpu().numpy().astype(o1.dtype))
            print(f"matmul took {(time.time() - time_mat)/60:.2f}min", flush=True)
            o1.PS = py_softmax(o1.PS, 1)
            o1 = optimize_L_sk(o1, nh=nh)
    return o1

def gpu_sk(o1):
    """ Sinkhorn Knopp optimization on GPU
            * stores activations on multiple GPUs (needed when dataset is large)
            * does matrix-vector multiplies on GPU (extremely fast)
            * recommended variant
            * due to multi-GPU use, it's a bit harder to understand what's happening -> see CPU variant to understand
    """
    # 1. aggregate inputs:
    start_t = time.time()
    if o1.hc == 1:
        o1.PS, indices = aggreg_multi_gpu(o1.model, o1.pseudo_loader,
                                            hc=o1.hc, dim=o1.outs[0], TYPE=o1.dtype)

    else:
        try: # just in case stuff
            del o1.PS_pre
        except:
            pass
        torch.cuda.empty_cache()
        time.sleep(1)
        o1.PS_pre, indices = aggreg_multi_gpu(o1.model, o1.pseudo_loader,
                                                hc=o1.hc, dim=o1.presize, TYPE=torch.float32)
        o1.model.headcount = o1.hc
    print("Aggreg of outputs  took {0:.2f} min".format((time.time() - start_t) / 60.), flush=True)
    # 2. solve label assignment via sinkhorn-knopp:
    if o1.hc == 1:
        o1 = optimize_L_sk_multi(o1, nh=0)
        o1.L[0,indices] = o1.L[0,:]
    else:
        for nh in range(o1.hc):
            tl = getattr(o1.model, "top_layer%d" % nh)
            time_mat = time.time()
            try:
                del o1.PS
                torch.cuda.empty_cache()
            except:
                pass

            # apply last FC layer (a matmul and adding of bias)
            o1.PS = gpu_mul_AB(o1.PS_pre, tl.weight.t(),
                                 c=tl.bias, dim=o1.outs[nh], TYPE=o1.dtype)
            print("matmul took %smin" % ((time.time() - time_mat) / 60.), flush=True)
            o1 = optimize_L_sk_multi(o1, nh=nh)
            o1.L[nh][indices] = o1.L[nh]
    return o1

def optimize_L_sk(o1, nh=0):
    N = max(o1.L.size())
    tt = time.time()
    o1.PS = o1.PS.T # now it is K x N
    r = np.ones((o1.outs[nh], 1), dtype=o1.dtype) / o1.outs[nh]
    c = np.ones((N, 1), dtype=o1.dtype) / N
    o1.PS **= o1.lamb  # K x N
    inv_K = o1.dtype(1./o1.outs[nh])
    inv_N = o1.dtype(1./N)
    err = 1e6
    _counter = 0
    while err > 1e-1:
        r = inv_K / (o1.PS @ c)          # (KxN)@(N,1) = K x 1
        c_new = inv_N / (r.T @ o1.PS).T  # ((1,K)@(KxN)).t() = N x 1
        if _counter % 10 == 0:
            err = np.nansum(np.abs(c / c_new - 1))
        c = c_new
        _counter += 1
    print("error: ", err, 'step ', _counter, flush=True)  # " nonneg: ", sum(I), flush=True)
    # inplace calculations.
    o1.PS *= np.squeeze(c)
    o1.PS = o1.PS.T
    o1.PS *= np.squeeze(r)
    o1.PS = o1.PS.T
    argmaxes = np.nanargmax(o1.PS, 0) # size N
    newL = torch.LongTensor(argmaxes)
    o1.L[nh] = newL.to(o1.dev)
    print('opt took {0:.2f}min, {1:4d}iters'.format(((time.time() - tt) / 60.), _counter), flush=True)
    return o1

def optimize_L_sk_multi(o1, nh=0):
    """ optimizes label assignment via Sinkhorn-Knopp.

         this implementation uses multiple GPUs to store the activations which allow fast matrix multiplies

         Parameters:
             nh (int) number of the head that is being optimized.

    """
    N = max(o1.L.size())
    tt = time.time()
    r = torch.ones((o1.outs[nh], 1), device='cuda:0', dtype=o1.dtype) / o1.outs[nh]
    c = torch.ones((N, 1), device='cuda:0', dtype=o1.dtype) / N
    ones = torch.ones(N, device='cuda:0', dtype=o1.dtype)
    inv_K = 1. / o1.outs[nh]
    inv_N = 1. / N

    # inplace power of softmax activations:
    [qq.pow_(o1.lamb) for qq in o1.PS]  # K x N

    err = 1e6
    _counter = 0
    ngpu = torch.cuda.device_count()
    splits = np.cumsum([0] + [a.size(0) for a in o1.PS])
    while err > 1e-1:
        r = inv_K / (gpu_mul_xA(c.t(), o1.PS,
                                ngpu=ngpu, splits=splits, TYPE=o1.dtype)).t()  # ((1xN)@(NxK)).T = Kx1
        c_new = inv_N / (gpu_mul_Ax(o1.PS, r,
                                    ngpu=ngpu, splits=splits, TYPE=o1.dtype))  # (NxK)@(K,1) = N x 1
        torch.cuda.synchronize()  # just in case
        if _counter % 10 == 0:
            err = torch.sum(torch.abs((c.squeeze() / c_new.squeeze()) - ones)).cpu().item()
        c = c_new
        _counter += 1
    print("error: ", err, 'step ', _counter, flush=True)

    # getting the final tranportation matrix #####################
    for i, qq in enumerate(o1.PS):
        torch.mul(qq, c[splits[i]:splits[i + 1], :].to('cuda:' + str(i + 1)), out=qq)
    [torch.mul(r.to('cuda:' + str(i + 1)).t(), qq, out=qq) for i, qq in enumerate(o1.PS)]
    argmaxes = torch.empty(N, dtype=torch.int64, device='cuda:0')

    start_idx = 0
    for i, qq in enumerate(o1.PS):
        amax = torch.argmax(qq, 1)
        argmaxes[start_idx:start_idx + len(qq)].copy_(amax)
        start_idx += len(qq)
    newL = argmaxes
    print('opt took {0:.2f}min, {1:4d}iters'.format(((time.time() - tt) / 60.), _counter), flush=True)
    # finally, assign the new labels ########################
    o1.L[nh] = newL
    return o1

max_stars_repo_name: BruntonUWBio/cross-modal-ssl-htnet

max_issues_repo_name: BruntonUWBio/cross-modal-ssl-htnet

max_forks_repo_name: BruntonUWBio/cross-modal-ssl-htnet


text:
import numpy as np
import seaborn as sns
import matplotlib.pyplot as plt
# sns.set(style="white", context="talk")
sns.set_style("darkgrid")

rs = np.random.RandomState(8)

# Set up the matplotlib figure
# f, ax3 = plt.subplots(1, 1, figsize=(7, 5), sharex=True)
f, ax3 = plt.subplots()

# Generate some sequential data
x = np.array(list("ABCDEFGHIJ"))
y1 = np.arange(1, 11)
# sns.barplot(x=x, y=y1, palette="rocket", ax=ax1)
# ax1.axhline(0, color="k", clip_on=False)
# ax1.set_ylabel("Sequential")
#
# # Center the data to make it diverging
# y2 = y1 - 5.5
# sns.barplot(x=x, y=y2, palette="vlag", ax=ax2)
# ax2.axhline(0, color="k", clip_on=False)
# ax2.set_ylabel("Diverging")

labels = ['NBC', 'ABC', 'CBS', 'FOX']
sizes = [35, 30, 35, 10]

# Randomly reorder the data to make it qualitative
# y3 = rs.choice(sizes, len(sizes), replace=False)
sns.barplot(x=labels, y=sizes, palette="deep", ax=ax3)
ax3.axhline(0, color="k", clip_on=False)

plt.annotate('actual group', xy=(sizes+0.2,sizes), xytext=(x+0.3, 300),
             arrowprops=dict(facecolor='black', shrink=0.05, headwidth=20, width=7))

# ax3.set_ylabel("Qualitative")

# Finalize the plot
sns.despine(bottom=True)
plt.setp(f.axes, yticks=[])
plt.tight_layout(h_pad=2)

# Configure the y-axis
plt.yticks(np.arange(0, max(sizes)+1, 5))

plt.show()


text:
import os
import argparse
import yaml
import random
import numpy as np
from scipy.special import comb

from Factory import set_problem
from WeightVector import das_dennis, determine_neighbor
from Population import init_pop, eval_pop
from ReferencePoint import init_ref_point, update_ref_point
from Mutation import lf_mutation, poly_mutation, fix_bound
from Decomposition import tchebycheff

###################
#  MOP to solve   #
###################
from pymoo.factory import get_problem
# from platypus.problems import UF9
###################
# parse arguments #
###################

parser = argparse.ArgumentParser()
parser.add_argument('params', type=argparse.FileType('r'))
parser.add_argument('seed', type=int)
args = parser.parse_args()

##################
# set parameters #
##################

params = yaml.safe_load(args.params)                                            # read config file

random.seed(args.seed)                                                          # set the seed for reproducibility purposes
np.random.seed(args.seed)                                                       # set the seed for reproducibility purposes

output = params['output']                                                       # set output of record in this run

n_obj = params['n_obj']                                                         # set number of objectives
n_var = params['n_var']                                                         # set number of variables
xl = params['xl']                                                               # set boundary of variables
xu = params['xu']                                                               # set boundary of variables

prob_name = params['prob_name']
sld_n_part = params['sld_n_part']

problem = set_problem(prob_name, n_var, n_obj, xu, xl)                          # set optimization problem

                                           # set number of partitions

n_eval = params['n_eval']                                           # set maximum number of evaluation
n_pop = int( comb(n_obj + sld_n_part - 1, n_obj - 1) )                  # compute population size

T = params['T']                                                     # set neighbor size
delta = params['delta']                                             # set probability to select parent from neighbor
nr = params['nr']                                                   # set maximum update counts for one offspring

alpha = params['alpha']                                             # set scaling factor of levy flight mutation
beta = params['beta']                                               # set stability parameter of levy flight mutation
etam = params['etam']                                               # set index parameter of polynomial mutation

#################
# start program #
#################

os.makedirs(f'./{output}', exist_ok=True)                           # create a folder to include running results
os.makedirs(f'./{output}/history/', exist_ok=True)
os.makedirs(f'./{output}/history/{prob_name}_{args.seed}', exist_ok=True)
os.makedirs(f'./{output}/final/', exist_ok=True)                                # create a folder to include final results

W = das_dennis(sld_n_part, n_obj)                                   # generate a set of weight vectors
B = determine_neighbor(W, T)                                        # determine neighbor
X = init_pop(n_pop, n_var, xl, xu)                                  # initialize a population

Y = eval_pop(X, problem, prob_name)                                 # evaluate fitness
z = init_ref_point(Y)                                               # determine a reference point

n_fe = n_pop

# for c_gen in range(1, n_gen):                                       # start main loop
c_gen = 1

while n_fe < n_eval:
    result = np.hstack([Y, X])                                                  # record objective values (Y) and decision variables (X)
    np.savetxt(f'./{output}/history/{prob_name}_{args.seed}/{c_gen}_paretos.csv', result)   # record information (Y, X) about the current generation

    info_gen = np.hstack([n_fe, c_gen])                                         # record number of function evaluations and current generation
    np.savetxt(f'./{output}/history/{prob_name}_{args.seed}/{c_gen}_info_gen.csv', info_gen)# record information (n_fe and c_gen) about the current generation

    for i in np.random.permutation(n_pop):                          # traverse the population; randomly permuting solutions in the population set

        n_fe += 1
        xi, fi = X[i, :], Y[i, :]                                   # get current individual and its fitness value

        if random.random() < delta:                                 # determine selection pool by probability
            pool = B[i, :]                                          # neighbor as the pool
        else:
            pool = np.arange(n_pop)                                 # population as the pool

        j = int(np.random.choice(pool))                                  # select a random individual from pool
        xj, fj = X[j, :], Y[j, :]

        xi_ = fix_bound( lf_mutation(xi, xj, alpha, beta), xl, xu ) # levy flight mutation
        xi_ = fix_bound( poly_mutation(xi_, etam, xl, xu), xl, xu ) # polynomial mutation

        fi_ = problem(xi_)                                              # evaluate offspring (new solution)

        z = update_ref_point(z, fi_)                                # update reference point

        nc = 0                                                      # initialize the update counter
        for k in np.random.permutation(len(pool)):                  # traverse the selection pool

            fk = Y[k, :]                                            # get k-th individual fitness
            wk = W[k, :]                                            # get k-th weight vector

            if tchebycheff(fi_, wk, z) <= tchebycheff(fk, wk, z):   # compare tchebycheff cost of offspring and parent

                X[k] = xi_                                          # update parent
                Y[k] = fi_
                nc += 1                                             # cumulate the counter

            if nc >= nr: break                                      # break if counter arrive the upper limit

            # exit()

    c_gen += 1
    # input('next gen normal moea.d')
info_gen = np.hstack([n_fe, c_gen])                                      # record information (n_fe and c_gen) about the current generation
result = np.hstack([Y, X])                                          # record objective values and decision variables

np.savetxt(f'./{output}/history/{prob_name}_{args.seed}/{c_gen}_info_gen.csv', info_gen)
np.savetxt(f'./{output}/history/{prob_name}_{args.seed}/{c_gen}_paretos.csv', result)

np.savetxt(f'./{output}/final/{prob_name}_{args.seed}_paretos.csv', result)
np.savetxt(f'./{output}/final/{prob_name}_{args.seed}_info_gen.csv', info_gen)

reason: import numpy,from scipy


text:
'''
.. _picasso.localize:
    https://picassosr.readthedocs.io/en/latest/localize.html
.. _picasso.render:
    https://picassosr.readthedocs.io/en/latest/render.html
.. _spt:
    https://www.biorxiv.org/content/10.1101/2020.05.17.100354v1
'''

import os
import numpy as np
import numba as numba
import pandas as pd
from tqdm import tqdm
from scipy.spatial import cKDTree

import picasso.localize as localize
import picasso.gausslq as gausslq
import picasso.render as render
import picasso.io as io

import picasso_addon.io as addon_io

PICKED_DTYPE = {'group':np.uint32,
                'frame':np.uint32,
                'x':np.float32,
                'y': np.float32,
                'photons':np.float32,
                'sx':np.float32,
                'sy':np.float32,
                'bg':np.float32,
                'lpx':np.float32,
                'lpy':np.float32,
                'ellipticity':np.float32,
                'net_gradient':np.float32,
                }
LOCS_DTYPE = PICKED_DTYPE.copy()
del LOCS_DTYPE['group']

LOCS_COLS = [field for field in LOCS_DTYPE]
PICKED_COLS = [field for field in PICKED_DTYPE]

#%%
@numba.jit(nopython=True, nogil=True, cache=False)
def check_spots(frame,y,x,box):
    '''
    Check all boxes around x,y for number of localizations, active pixels and center of mass position.

    Args:
        frame (np.array): Subpixel rendered localization image.
        y (int):          Box center y-coordinate (row)
        x (int):          Box center x-coordinate (column)
        box (int):        Box size in subpixels. Must be uneven!

    Returns:
        tuple:
            - [0] (float): Center of mass y-position as measured from y
            - [1] (float): Center of mass x-position as measured from x
            - [2] (int):   Number of localizations within box
            - [3] (int):   Number of subpixels within box with at least one localization
    '''
    box_half = int(box / 2)
    dx = np.zeros(len(x), dtype=np.float32)
    dy = np.zeros(len(x), dtype=np.float32)
    n_locs = np.zeros(len(x), dtype=np.float32)
    n_activepx = np.zeros(len(x), dtype=np.int32)

    for i in range(0,len(x)):
        ### Get spot
        spot=frame[y[i]-box_half:y[i]+box_half+1,
                   x[i]-box_half:x[i]+box_half+1,
                   ]
        ### Get center of mass position, sum and number of active pixels in spot
        n_locs[i],dx[i],dy[i]=gausslq._sum_and_center_of_mass(spot,box)
        dx[i]=dx[i]-box_half # dx is now measured from center position x
        dy[i]=dy[i]-box_half # dy is now measured from center position y
        n_activepx[i]=np.sum(spot>0)

    return dy,dx,n_locs,n_activepx

#%%
@numba.jit(nopython=True, nogil=True, cache=False)
def identifyspots_in_image(image, minimum_n_locs, box):
    '''
    Identify boxes above localization threshold in rendered localizations image by using check_spots().

	Args:
		image (np.array): Subpixel rendered localization image.
		minimum_n_locs (int): Number of localization threshold for valid boxes
		box (int): Box size in subpixels. Must be uneven!

	Returns:
		tuple:
    		- [0] (int):   Local box maximum y (i.e. box y center coordinate)
    		- [1] (int):   Local box maximum x (i.e. box x center coordinate)
    		- [2] (float): Center of mass y-position as measured from y
    		- [3] (float): Center of mass x-position as measured from x
    		- [4] (int):   Number of localizations within box
    '''
    ### Get local maxima in each box by covering complete image
    y,x=localize.local_maxima(image, box)
    ### Check spots for number of loclizations and number active pixels
    dy,dx,n_locs, n_activepx = check_spots(image,y,x,box)
    ### Apply filter and return positive spots
    positives = (n_locs >= minimum_n_locs)&(n_activepx>1)
    y = y[positives]
    x = x[positives]
    dy = dy[positives]
    dx = dx[positives]
    n_locs = n_locs[positives]
    n_activepx = n_activepx[positives]
    return y, x, dy, dx, n_locs

#%%
def spotcenters_in_image(image,box,min_nlocs,fit=False):
    '''
    Get pick center coordinates in subpixel rendered localization image by either
    calculating center of mass or fitting 2D Gaussian.

    Args:
  		image (np.array): Subpixel rendered localization image.
  		box (int):        Box size in subpixels. Must be uneven!
  		min_nlocs (int):  Number of localization threshold for valid boxes
  		fit (bool=False): Employ 2D gaussian fitting? Normally not better performing. If False only center of mass position.
    Returns:
        tuple:
            - [0] (pandas.DataFrame): Spot center coordinates (image units) and number of localizations (x,y,n_locs)
            - [1] (bool): True if fit was performed
    '''
    image=image.astype(float)
    ### 1. Cover complete image with boxes and identify local maxima
    ### 2. Reposition boxes around local maxima and check for:
    ###         - Number of localizations (n_locs)
    ###         - Center of mass positions (dx,dy) in spot coordinates with x,y at center position (0,0)
    y,x,dy,dx,n_locs=identifyspots_in_image(image,min_nlocs,box)

    ### Remove spots closer than box_half to image boarder
    box_half=int(box/2)
    im_size_x=np.shape(image)[0]
    im_size_y=np.shape(image)[1]

    positives=(y>box_half)&((y+box_half)<im_size_y)&(x>box_half)&((x+box_half)<im_size_x)
    y=y[positives]
    x=x[positives]
    dy=dy[positives]
    dx=dx[positives]
    n_locs=n_locs[positives]

    ### Try gaussian fitting of boxes to update center position, return center of mass position in case of failure
    if fit==True:
        print('Fitting spots ...')
        ### Prepare spots
        spots = np.zeros((len(x), box, box), dtype=image.dtype) # Initialize spots
        localize._cut_spots_frame(image,0,np.zeros(len(x)),x,y,int(box/2),0,len(x),spots) # Fill spots
        spots=spots.astype(float) # Convert to float for fitting
        try:
            ### Fitting
            theta=gausslq.fit_spots_parallel(spots,asynch=False)
            ### Add box offset
            x_new=theta[:, 0]+x
            y_new=theta[:, 1]+y

            ### Check again for boundaries after fitting
            positives=(y_new>box_half)&((y_new+box_half)<im_size_y)&(x_new>box_half)&((x_new+box_half)<im_size_x)
            y_new=y_new[positives]
            x_new=x_new[positives]
            n_locs=n_locs[positives]
        except:
            print('Fitting failed switching to center of mass position!')
            x_new=x+dx
            y_new=y+dy
            fit=False
    else: # Also return center of mass position
        x_new=x+dx
        y_new=y+dy

    ### Assign
    spot_centers=pd.DataFrame(columns=['x','y','n_locs'])
    spot_centers.x=x_new+0.5 # We have to add 0.5 to go from a pixelated presentation to continuous
    spot_centers.y=y_new+0.5 # We have to add 0.5 to go from a pixelated presentation to continuous
    spot_centers.n_locs=n_locs

    return spot_centers,fit

#%%
def coordinate_convert(spot_centers,viewport_min,oversampling):
    '''
    Convert spot centers as detected in rendered localizations by given viewport minima and oversampling
    back to original localization coordinates.

    Args:
        spot_centers (pandas.DataFrame): Output of function ``spotcenters_in_image()[0]``, i.e. spot center coordinates.
        viewport_min (tuple): ``(min_x,min_y)`` of viewport as in picasso.render.
        oversampling (int): Pixel oversampling as in picasso.render

    Returns:
        pandas.DataFrame: spot_centers converted back to orignal localization coordinates.
    '''
    spot_centers_convert=spot_centers.copy()
    spot_centers_convert.x=spot_centers.x/oversampling+viewport_min[1]
    spot_centers_convert.y=spot_centers.y/oversampling+viewport_min[0]

    return spot_centers_convert
#%%
def query_locs_for_centers(locs,centers,pick_radius=1):
    '''
    Builds up KDtree for locs, queries the tree for localizations within pick_radius (norm of p=2) around centers (xy coordinates).
    Output will be list of indices as result of query.

    Args:
        locs (numpy.recarray or pandas.DataFrame): Localizations as created by picasso.localize with fields ``x`` and ``y``.
        centers (np.array):                        Pick center coordinates (x,y)
        pick_radius (float=1):                     Pick diameter in px.

    Return:
        list:
            List of ``len(centers)``. Single list entries correpond to indices in ``locs`` of localizations
            within ``pick_radius`` around ``centers``. Hence every list indicates one pick (group).
    '''

    #### Prepare data for KDtree
    data=np.vstack([locs.x,locs.y]).transpose()
    #### Build up KDtree
    tree=cKDTree(data,compact_nodes=False,balanced_tree=False)
    #### Query KDtree for indices belonging to picks
    picks_idx=tree.query_ball_point(centers,r=pick_radius,p=2)

    return picks_idx
#%%
def get_picked(locs_in,picks_idx):
    '''
    Get locs corresponding to picks_idx as obtained by ``query_locs_for_centers()`` and assign group ID.

    Args:
        locs (numpy.recarray):      Localizations as loaded by picasso.io (see picasso.localize).
        picks_idx (numpy.ndarray):  Single rows correpond to indices (as list) in locs within pick_radius around centers.
                                    See ``query_locs_for_centers()``.

    Returns:
        pandas.DataFrame: Picked as in `picasso.render`_. Localizations in overlapping picks appear in each pick,
                          i.e. same localization appears multiple times but with different group ID!
                          This means ``_picked`` can have more localizations than original ``_locs``!
    '''

    ######### Prepare locs
    ### Convert to DataFrame for easier access
    locs_df = pd.DataFrame(locs_in,dtype=np.float32)
    ### Rearrange colums according to COLS_ORDER
    locs_df = locs_df[LOCS_COLS]
    ### Convert DataFrame to numpy.array for faster computation
    locs = locs_df.values.astype(np.float32)

    ######### Prepare picks_idx
    ### Convert pick indices np.ndarray to list
    picks_list = picks_idx.tolist()
    ### Get flattened length of pick_list
    n = np.sum([len(l) for l in picks_list])

    ######## This is the important part for speedy group assignment
    ### Get one dimensional array of pick indices in locs
    ids = []
    for l in picks_list: ids.extend(l)
    ### Get one dimensional array of groups ids corresponding to pick indices in locs
    gs = [np.ones(len(l))*group for group,l in enumerate(picks_list)]
    gs = np.concatenate(gs)

    ### Initiate picked
    picked = np.zeros((n,len(PICKED_COLS)), dtype = np.float32)
    ### Assign all localization information to picked
    picked[:,1:] = locs[ids,:]
    picked[:,0] = gs

    picked_df = pd.DataFrame(picked, columns = PICKED_COLS)
    picked_df = picked_df.astype(PICKED_DTYPE)

    return picked_df

#%%
def ck_nlf(trace, M = 10, N = [2, 4, 8, 16], p = 30):
    '''
    Chung-Kennedy nonlinear filter

    Nonlinear local low-pass filter optimized for edge preservation. Used
    quite frequently in step detection in noisy traces.
    For details, see:
    Chung SH & Kennedy RA, J Neurosci Meth 1991,
    DOI: 10.1016/0165-0270(91)90118-J

    Note that the algorithm frequently raises warnings saying "invalid value
    encountered in true_divide". There is no need to worry about that:
    The filter, when used with high p, tends to create extreme numbers in
    calculation intermediates, which may even end up creating zeros. This only
    affects weighting functions of a handful of data points, though, and has
    negligible impact on the overall filter result. It will NOT create any
    extreme values in the returned filtered trace, at least I have never
    encountered such results.

    Args:
        trace (numpyp.array): Time trace to be filtered. Signal only, no time axis required.
        M (int):              Field width for predictor weight determination. Typical values: 5-20.
        N (list):             Field width set for predictor calculation. Typical values: [4, 8, 16, 32] or subsets thereof. Minumum ``len(N) > 1``
        p (int):              Positive single number (usually int, but not strictly required). Nonlinear order of predictor weight calculation.
                              Typical values: 5-100, strongly application-dependent. Higher values
                              mean sharper edge preservation.

    Returns:
        trace_f (numpy.array): 1D filtered time trace of same structure as input.

    '''

    # Initializations
    K = len(N)
    l_trace = len(trace)
    pred_f = np.zeros((l_trace, K))
    pred_b = np.zeros((l_trace, K))
    weight_f = np.zeros((l_trace, K))
    weight_b = np.zeros((l_trace, K))
    weight_f_norm = np.zeros((l_trace, K))
    weight_b_norm = np.zeros((l_trace, K))

    # Raw predictors.
    for k in range(K):
        trace_shift = np.empty((l_trace + N[k], N[k]))
        trace_shift[:, :] = np.nan
        for l in range(N[k]):
            trace_shift[l:l_trace + l, l] = trace
        # Forward predictor
        pred_f[0] = trace[0]
        pred_f[1:, k] = np.nanmean(trace_shift[:l_trace - 1], axis = 1)
        # Backward predictor
        pred_b[0:-1, k] = np.nanmean(trace_shift[N[k]:-1], axis = 1)
        pred_b[-1, k] = trace[-1]

    # Non-normalized weights.
    # First value of forward and last value of backward would lead to div by 0.
    # They are skipped and their weights set to (rather left at) 0.
        # Forward predictor weights
        trace_shift = np.empty((l_trace + M - 1, M))
        trace_shift.fill(np.nan)
        MSE_f = (trace - pred_f[:,k]) ** 2
        for m in range(M):
            trace_shift[m:l_trace + m, m] = MSE_f
        weight_f[1:, k] = np.nansum(trace_shift[1:l_trace, :], axis = 1) ** (-p)
        # Backward predictor weights
        trace_shift = np.empty((l_trace + M - 1, M))
        trace_shift.fill(np.nan)
        MSE_b = (trace - pred_b[:, k]) ** 2
        for m in range(M):
            trace_shift[m:l_trace + m, m] = MSE_b
        weight_b[:-1, k] = np.nansum(trace_shift[-l_trace:-1, :], axis = 1) ** (-p)

    # Normalize weights.
    weightsum = np.nansum(weight_f, axis = 1) + np.nansum(weight_b, axis = 1)
    for k in range(K):
        weight_f_norm[:, k] = weight_f[:, k] / weightsum
        weight_b_norm[:, k] = weight_b[:, k] / weightsum

    # Get filtered trace.
    pred_f_w = pred_f * weight_f_norm
    pred_b_w = pred_b * weight_b_norm
    trace_f = np.nansum(pred_f_w, axis = 1) + np.nansum(pred_b_w, axis = 1)

    return trace_f

#%%
def main(locs,info,path,**params):
    '''
    Cluster detection (pick) in localizations by thresholding in number of localizations per cluster.
    Cluster centers are determined by creating images of localization list with set oversampling using picasso.render.

    Args:
        locs(numpy.recarray):           (Undrifted) localization list as created by `picasso.render`_.
        info(list(dict)):               Info to localization list when loaded with picasso.io.load_locs().

    Keyword Arguments:
        pick_diameter(float=1.2):       Pick diameter in original pixels.
        oversampling(int=5):            Oversampling for rendering of localization list, i.e. sub-pixels per pixel of original image.
        min_n_locs(float=0.2*NoFrames): Detection threshold for number of localizations in cluster.
                                        Standard value is set for `spt`_.
                                        Set to lower value for usual DNA-PAINT signal (see ``lbfcs``).
        fit_center(bool=False):         False = Center of mass. True = 2D Gaussian fitting of center.
        lbfcs(bool=False):              If set to True will overrun min_n_locs and sets it to 0.02*NoFrames.

    Returns:
        list:
            - [0] (dict):             kwargs passed to function.
            - [1] (pandas.DataFrame): Picked localizations, saved with extension _picked.hdf5.
            - [2] (pandas.DataFrame): Center positions and number of localizations per pick, saved with extension _autopick.yaml.
            - [3] (str):              Full save path.

    '''
    ### Path of file that is processed
    path=os.path.splitext(path)[0]

    ### Set standard conditions if not set as input
    NoFrames=info[0]['Frames']

    standard_params={'pick_diameter':1.2,
                     'oversampling':5,
                     'min_n_locs':0.2*NoFrames,
                     'fit_center':False,
                     }
    ### If lbFCS given overrun min_n_locs
    try:
        if params['lbfcs']==True:
            standard_params['min_n_locs']=0.02*NoFrames
            standard_params['lbfcs']=True
    except KeyError:
        pass

    ### Remove keys in params that are not needed
    for key, value in standard_params.items():
        try:
            params[key]
            if params[key]==None: params[key]=standard_params[key]
        except:
            params[key]=standard_params[key]

    ### Remove keys in params that are not needed
    delete_key=[]
    for key, value in params.items():
        if key not in standard_params.keys():
            delete_key.extend([key])
    for key in delete_key:
        del params[key]

    ### Procsessing marks: generatedby
    params['generatedby']='picasso_addon.autopick.main()'

    print('Minimum number of localizations in pick set to %i'%(params['min_n_locs']))

    ### Check if locs is given as numpy.recarray as created by picasso.localize
    if type(locs) is not np.recarray:
        raise SystemExit('locs must be given as numpy.recarray')

    ### Render locs
    print('Rendering locs for pick detection ...(oversampling = %i)'%params['oversampling'])
    image=render.render(locs,
                        info,
                        oversampling=params['oversampling'],
                        )[1]

    ### Get pick centers in image coordinates ...
    ### ... i.e. first define correct box size for given pick_diameter
    pick_box_half = params['pick_diameter'] * 0.5 * params['oversampling'] # Half the box in oversampled size
    pick_box_half = np.floor(pick_box_half - 0.5)                          #
    pick_box =  int(2*pick_box_half + 1)                                   # Mutliply half by 2x and add 1 to get uneven integer for box
    pick_box = max(3,pick_box)                                             # Ensure that box >= 3
    print('Identifiying valid picks ...(box = %i)'%pick_box)

    centers_image,fit=spotcenters_in_image(image,
                                           pick_box,
                                           params['min_n_locs'],
                                           fit=params['fit_center'])
    params['fit_center']=fit # Update if it was not successful!

    ### Convert image coordinate centers to original values as in locs
    centers=coordinate_convert(centers_image,
                               (0,0),
                               params['oversampling'],
                               )
    ### Save converted centers as picks.yaml for later usage in picasso.render
    addon_io.save_picks(centers,
                        params['pick_diameter'],
                        path+'_autopick.yaml',
                        )

    ### Query locs for centers
    print('Build up and query KDtree ...(pick_diameter = %.1f)'%params['pick_diameter'])
    picks_idx=query_locs_for_centers(locs,
                                     centers[['x','y']].values,
                                     pick_radius=params['pick_diameter']/2,
                                     )
    ### Assign group ID to locs
    print('Assigning %i groups ...'%len(picks_idx))
    locs_picked=get_picked(locs,
                           picks_idx,
                           )

    ### Apply Chung-Kennedy local mean filter to photons of each group
    print('Applying Chung-Kennedy filter ...')
    tqdm.pandas()
    locs_picked = locs_picked.groupby('group').progress_apply(lambda df: df.assign( photons_ck = ck_nlf(df.photons.values).astype(np.float32) ) )

    ### Save locs_picked as .hdf5 and info+params as .yaml
    print('Saving _picked ...')
    info_picked=info.copy()+[params]      # Update info
    next_path=path+'_picked.hdf5'         # Update path
    io.save_locs(next_path,
                 locs_picked.to_records(index=False),
                 info_picked,
                 )

    return [params,locs_picked,centers,next_path]

reason: import numpy,from scipy,import numba


text:
"""
Generates random sample of tax filing units with attributes such that
generated file can be directly uploaded to Internet TAXSIM version 27.
"""
# CODING-STYLE CHECKS:
# pycodestyle taxsim_input.py
# pylint --disable=locally-disabled taxsim_input.py

import argparse
import sys
import numpy as np
import pandas as pd

VALID_LETTERS = ['a', 'b', 'c']

def main():
    """
    High-level logic.
    """
    # parse command-line arguments:
    usage_str = 'python taxsim_input.py YEAR LETTER [OFFSET] [--help]'
    parser = argparse.ArgumentParser(
        prog='',
        usage=usage_str,
        description=('Generates random sample of tax filing units with '
                     'attributes and format such that the file can be '
                     'directly uploaded to Internet TAXSIM version 27. '
                     'For details on Internet TAXSIM version 27 INPUT '
                     'format, go to '
                     'https://users.nber.org/~taxsim/taxsim27/'))
    parser.add_argument('YEAR', nargs='?', type=int, default=0,
                        help=('YEAR specifies calendar year assumed in '
                              'generated input data.'))
    parser.add_argument('LETTER', nargs='?', default='',
                        help=('LETTER specifies assumption set '
                              'used to generate input data.'))
    parser.add_argument('OFFSET', nargs='?', type=int, default=0,
                        help=('optional OFFSET alters the '
                              'random-number seed used to generate '
                              'sample of filing units.  Default OFFSET '
                              'value is zero.'))
    args = parser.parse_args()
    # check YEAR value
    if args.YEAR < 2013 or args.YEAR > 2023:
        sys.stderr.write('ERROR: YEAR not in [2013,2023] range\n')
        sys.stderr.write('USAGE: {}\n'.format(usage_str))
        return 1
    # check LETTER value
    if args.LETTER == '':
        sys.stderr.write('ERROR: must specify LETTER\n')
        sys.stderr.write('USAGE: {}\n'.format(usage_str))
        return 1
    if args.LETTER not in VALID_LETTERS:
        sys.stderr.write('ERROR: LETTER not in VALID_LETTERS, where\n')
        sys.stderr.write('       VALID_LETTERS={}\n'.format(VALID_LETTERS))
        sys.stderr.write('USAGE: {}\n'.format(usage_str))
    # check OFFSET value
    if args.OFFSET < 0 or args.OFFSET > 999:
        sys.stderr.write('ERROR: OFFSET not in [0,999] range\n')
        sys.stderr.write('USAGE: {}\n'.format(usage_str))
        return 1
    # get dictionary containing assumption set
    assump = assumption_set(args.YEAR, args.LETTER)
    # generate sample as pandas DataFrame
    sample = sample_dataframe(assump, args.YEAR, args.OFFSET)
    # write sample to input file
    filename = '{}{}.in'.format(args.LETTER, args.YEAR % 100)
    sample.to_csv(filename, sep=' ', header=False, index=False)
    # return no-error exit code
    return 0
# end of main function code

def assumption_set(year, letter):
    """
    Return dictionary containing assumption parameters.
    """
    adict = dict()
    if letter in VALID_LETTERS:  # <===========================================
        # basic assumption parameters for all ?YY.in samples:
        adict['sample_size'] = 100000
        adict['year'] = year  # TAXSIM ivar 2
        # demographic attributes:
        adict['joint_frac'] = 0.60  # fraction of sample with joint MARS
        adict['min_age'] = 17  # TAXSIM ivar 5 (primary taxpayer age)
        adict['max_age'] = 77  # TAXSIM ivar 5 (primary taxpayer age)
        adict['min_age_diff'] = -10  # min spouse age difference
        adict['max_age_diff'] = 10  # max spouse age difference
        adict['max_depx'] = 5  # TAXSIM ivar 7 (total number of dependents)
        adict['max_dep13'] = 4  # TAXSIM ivar 8 (Child/Dependent Care Credit)
        adict['max_dep17'] = 4  # TAXSIM ivar 9 (Child Credit)
        adict['max_dep18'] = 4  # TAXSIM ivar 10 (EITC)
        # labor income:
        adict['max_pwages_yng'] = 500  # TAXSIM ivar 11
        adict['max_pwages_old'] = 30  # TAXSIM ivar 11 (65+ ==> old)
        adict['max_swages_yng'] = 500  # TAXSIM ivar 12
        adict['max_swages_old'] = 30  # TAXSIM ivar 12 (65+ ==> old)
        # non-labor income (all zeros):
        adict['max_divinc'] = 0  # TAXSIM ivar 13
        adict['max_intinc'] = 0  # TAXSIM ivar 14
        adict['min_stcg'] = 0  # TAXSIM ivar 15
        adict['max_stcg'] = 0  # TAXSIM ivar 15
        adict['min_ltcg'] = 0  # TAXSIM ivar 16
        adict['max_ltcg'] = 0  # TAXSIM ivar 16
        adict['max_other_prop_inc'] = 0  # TAXSIM ivar 17
        adict['max_other_nonprop_inc'] = 0  # TAXSIM ivar 18
        adict['max_pnben'] = 0  # TAXSIM ivar 19
        adict['max_ssben'] = 0  # TAXSIM ivar 20
        adict['max_uiben'] = 0  # TAXSIM ivar 21
        # childcare expense amount (all zero):
        adict['max_ccexp'] = 0  # TAXSIM ivar 26
        # itemized expense amounts (all zero):
        adict['max_ided_proptax'] = 0  # TAXSIM ivar 24
        adict['max_ided_nopref'] = 0  # TAXSIM ivar 25
        adict['max_ided_mortgage'] = 0  # TAXSIM ivar 27
        # end if letter in VALID_LETTERS
    if letter in ['b', 'c']:  # <==============================================
        # non-labor income:
        adict['max_divinc'] = 20  # TAXSIM ivar 13
        adict['max_intinc'] = 20  # TAXSIM ivar 14
        adict['min_stcg'] = -10  # TAXSIM ivar 15
        adict['max_stcg'] = 10  # TAXSIM ivar 15
        adict['min_ltcg'] = -10  # TAXSIM ivar 16
        adict['max_ltcg'] = 10  # TAXSIM ivar 16
        adict['max_other_prop_inc'] = 30  # TAXSIM ivar 17
        adict['max_other_nonprop_inc'] = 30  # TAXSIM ivar 18
        adict['max_pnben'] = 60  # TAXSIM ivar 19
        adict['max_ssben'] = 60  # TAXSIM ivar 20
        adict['max_uiben'] = 10  # TAXSIM ivar 21
    if letter == 'c':  # <=====================================================
        # childcare expense amount:
        adict['max_ccexp'] = 10  # TAXSIM ivar 26
        # itemized expense amounts:
        adict['max_ided_proptax'] = 30  # TAXSIM ivar 24
        adict['max_ided_nopref'] = 10  # TAXSIM ivar 25
        adict['max_ided_mortgage'] = 40  # TAXSIM ivar 27
    return adict

def sample_dataframe(assump, year, offset):
    """
    Construct DataFrame containing sample specified by assump and year+offset.
    """
    # pylint: disable=too-many-locals
    np.random.seed(123456789 + year + offset)
    size = assump['sample_size']
    zero = np.zeros(size, dtype=np.int64)
    sdict = dict()
    # (01) RECID
    sdict[1] = range(1, size + 1)
    # (02) YEAR
    sdict[2] = np.full_like(zero, assump['year'], dtype=np.int64)
    # (03) STATE
    sdict[3] = zero
    # (04) MSTAT
    urn = np.random.random(size)
    mstat = np.where(urn < assump['joint_frac'], 2, 1)
    sdict[4] = mstat
    # (05) PAGE
    sdict[5] = np.random.randint(assump['min_age'],
                                 assump['max_age']+1,
                                 size)
    # (06) SAGE
    age_diff = np.random.randint(assump['min_age_diff'],
                                 assump['max_age_diff']+1,
                                 size)
    sage = sdict[5] + age_diff
    sdict[6] = np.where(mstat == 2, np.maximum(sage, assump['min_age']), zero)
    # (07-10) DEPX, DEP13, DEP17, DEP18
    depx = np.random.randint(0, assump['max_depx']+1, size)
    d18 = np.random.randint(0, assump['max_dep18']+1, size)
    dep18 = np.where(d18 <= depx, d18, depx)
    d17 = np.random.randint(0, assump['max_dep17']+1, size)
    dep17 = np.where(d17 <= dep18, d17, dep18)
    d13 = np.random.randint(0, assump['max_dep13']+1, size)
    dep13 = np.where(d13 <= dep17, d13, dep17)
    sdict[7] = depx
    sdict[8] = dep13
    sdict[9] = dep17
    sdict[10] = dep18
    # (11) PWAGES
    pwages_yng = np.random.randint(0, assump['max_pwages_yng']+1, size)
    pwages_old = np.random.randint(0, assump['max_pwages_old']+1, size)
    sdict[11] = np.where(sdict[5] >= 65, pwages_old, pwages_yng) * 1000
    # (12) SWAGES
    swages_yng = np.random.randint(0, assump['max_swages_yng']+1, size)
    swages_old = np.random.randint(0, assump['max_swages_old']+1, size)
    swages = np.where(sdict[6] >= 65, swages_old, swages_yng) * 1000
    sdict[12] = np.where(mstat == 2, swages, zero)
    # (13) DIVIDENDS
    sdict[13] = np.random.randint(0, assump['max_divinc']+1, size) * 1000
    # (14) INTREC
    sdict[14] = np.random.randint(0, assump['max_intinc']+1, size) * 1000
    # (15) STCG
    sdict[15] = np.random.randint(assump['min_stcg'],
                                  assump['max_stcg']+1,
                                  size) * 1000
    # (16) LTCG
    sdict[16] = np.random.randint(assump['min_ltcg'],
                                  assump['max_ltcg']+1,
                                  size) * 1000
    # (17) OTHERPROP
    sdict[17] = np.random.randint(0,
                                  assump['max_other_prop_inc']+1,
                                  size) * 1000
    # (18) NONPROP
    sdict[18] = np.random.randint(0,
                                  assump['max_other_nonprop_inc']+1,
                                  size) * 1000
    # (19) PENSIONS
    sdict[19] = np.random.randint(0, assump['max_pnben']+1, size) * 1000
    # (20) GSSI
    sdict[20] = np.random.randint(0, assump['max_ssben']+1, size) * 1000
    # (21) UI
    sdict[21] = np.random.randint(0, assump['max_uiben']+1, size) * 1000
    # (22) TRANSFERS (non-taxable in federal income tax)
    sdict[22] = zero
    # (23) RENTPAID (used only in some state income tax laws)
    sdict[23] = zero
    # (24) PROPTAX
    sdict[24] = np.random.randint(0,
                                  assump['max_ided_proptax']+1,
                                  size) * 1000
    # (25) OTHERITEM
    sdict[25] = np.random.randint(0,
                                  assump['max_ided_nopref']+1,
                                  size) * 1000
    # (26) CHILDCARE (TAXSIM-27 EXPECTS ZERO IF NO QUALIFYING CHILDRED)
    ccexp = np.random.randint(0, assump['max_ccexp']+1, size) * 1000
    sdict[26] = np.where(dep13 > 0, ccexp, zero)
    # (27) MORTGAGE
    sdict[27] = np.random.randint(0,
                                  assump['max_ided_mortgage']+1,
                                  size) * 1000
    smpl = pd.DataFrame(sdict)
    return smpl

if __name__ == '__main__':
    sys.exit(main())

max_stars_repo_name: grantseiter/Tax-Benefits-Of-Parenthood

max_issues_repo_name: grantseiter/Tax-Benefits-Of-Parenthood

max_forks_repo_name: grantseiter/Tax-Benefits-Of-Parenthood


text:
import matplotlib.pyplot as plt
import numpy as np
from landlab.plot import graph, plot_network_and_parcels
from landlab.plot.imshow import imshow_grid, imshow_grid_at_node
from matplotlib.colors import LightSource
import imageio
import glob
import matplotlib.colors as mcolors
colors = [(0,0.2,1,i) for i in np.linspace(0,1,3)]
mycmap = mcolors.LinearSegmentedColormap.from_list('mycmap', colors, N=10)

def plot_discharge(hydrograph_time, discharge):
    fig=plt.figure()
    plt.plot(hydrograph_time, discharge, "b-", label="at outlet")
    plt.ylabel("Discharge (cms)")
    plt.xlabel("Time (seconds)")
    plt.legend(loc="upper right")
    return fig
def plot_depth(hydrograph_time, height_at_outlet):
    fig=plt.figure()
    plt.plot(hydrograph_time, height_at_outlet, "b-", label="at outlet")
    plt.ylabel("Water depth (m)")
    plt.xlabel("Time (seconds)")
    plt.legend(loc="upper right")
    return fig

def plot_sed_volume(parcels):
    parcel_vol_on_grid = parcels.dataset["volume"].values
    #if s
    parcel_vol_on_grid[parcels.dataset["element_id"].values==-2]=0
    sum_parcel_vol_on_grid = np.sum(parcel_vol_on_grid, axis=0)
    fig=plt.figure(4)
    plt.plot(np.asarray(parcels.time_coordinates),
             sum_parcel_vol_on_grid[0]-sum_parcel_vol_on_grid,'b-'
            )
    plt.ylabel('Total volume of parcels that left catchment $[m^3]$')
    plt.xlabel('Time (seconds)')
    return fig

def plot_outlet_conditions(hydrograph_time, discharge, height_at_outlet, parcels, filepath = ""):
    fig1 = plot_discharge(hydrograph_time, discharge)
    fig2 = plot_depth(hydrograph_time, height_at_outlet)
    fig3 = plot_sed_volume(parcels)
    if len(filepath) > 0:
        fig1.savefig(filepath+"runoff_discharge.jpeg", dpi = 1000)
        fig2.savefig(filepath+"runoff_height.jpeg", dpi = 1000)
        fig3.savefig(filepath+"sediment_volume.jpeg", dpi = 1000)
        print("outlet figures saved")
    else:
        plt.show()

def plot_network(network_grid):
    ## Plot nodes

    plt.figure(figsize=(18,5))
    plt.subplot(1,2,1)
    graph.plot_nodes(network_grid)
    plt.title("Nodes")

    ## Plot nodes + links
    plt.subplot(1,2,2)
    graph.plot_nodes(network_grid,with_id=False,markersize=4)
    graph.plot_links(network_grid)
    plt.title("Links")
    plt.show()
def plot_topography(rmg, network_grid=None):
# Plot topography
    plt.figure(figsize=(12,8))

    imshow_grid(rmg, 'topographic__elevation',\
                plot_name="Basin topography",\
                color_for_closed=None,\
                colorbar_label="$z$ [m]")
    if network_grid is not None:
        graph.plot_nodes(network_grid,with_id=False,markersize=4)
        graph.plot_links(network_grid,with_id=False)
    plt.show()

def plot_storm(rmg,h):
    plt.figure(figsize=(14,8))

    imshow_grid(rmg, 'topographic__elevation',\
                plot_name="Basin topography",\
                color_for_closed=None,\
                colorbar_label="$z$ [m]",\
                output=None,shrink=0.5)

    imshow_grid(rmg, h,\
                plot_name="Rain distribution",\
                color_for_closed=None,\
                cmap = mycmap,\
                colorbar_label="Initial rain $h$ [m]",\
                output=None,shrink=0.5)
    plt.show()
def plot_initial_floodplain(zFP, nX, nY, grid):
    # Plot 2D domain topography

    fig = plt.figure(figsize=(14,8))
    plt.subplot(1,2,1)
    ls = LightSource(azdeg=315, altdeg=45)
    plt.imshow(ls.hillshade(np.reshape(zFP,[nX,nY]), vert_exag=10), cmap='gray')
    plt.title('initial floodplain topography')

    #Plot 2D domain hydraulic cond.
    plt.subplot(1,2,2)
    imshow_grid(grid,'hydraulic_conductivity', plot_name="Hydraulic Conductivity", cmap="winter")
    fig.show()
def plot_overland_flow(rmg,outlet_nearest_raster_cell,elapsed_time,filepath = ""):
    #Plot overland flow
    fig=plt.figure()
    imshow_grid(rmg,'topographic__elevation',colorbar_label='Elevation (m)')
    imshow_grid(rmg,'surface_water__depth',limits=(0,2),cmap=mycmap,colorbar_label='Water depth (m)')
    plt.title(f'Time = {round(elapsed_time,1)} s')
    plt.plot(rmg.node_x[outlet_nearest_raster_cell], rmg.node_y[outlet_nearest_raster_cell], "yo")
    if len(filepath) > 0:
        fig.savefig(filepath + "flow/" + str(int(elapsed_time)).zfill(5) + ".png", dpi=150)
        plt.close(fig)
    else:
        plt.show()
def plot_floodplain(grid, zFP, nX, nY, elapsed_time, filepath = ""):
    fig = plt.figure(figsize=(14, 8))
    ls = LightSource(azdeg=315, altdeg=45)
    plt.imshow(ls.hillshade(np.reshape(zFP,[nX,nY]), vert_exag=10), cmap='gist_earth',origin="lower")
    imshow_grid(grid,'surface_water__depth',
          limits=(0,1),
          colorbar_label="Water depth (m)",
          cmap = mycmap,
          plot_name="Time = %i" %elapsed_time)
    if len(filepath) > 0:
        fig.savefig(filepath + "floodplain/" + str(int(elapsed_time)).zfill(5) + ".png", dpi = 150)
        plt.close(fig)
    else:
        plt.show()
def plot_parcels(new_grid, parcels, elapsed_time, outlet_nearest_raster_cell, filepath = ""):
    #Plot sediment parcels locationss

   # parcel_filter = np.zeros((parcels.dataset.dims["item_id"]), dtype=bool)

    #parcel_filter[[parcels.dataset["active_layer"].values==1][0][:,-1]==True] = True

    #pc_opts= {"parcel_filter": parcel_filter}
    #fig = plt.figure(figsize = (14,8))

    #plt.subplot(plt.subplot(1,2,1))
    #plot_network_and_parcels(
    #        new_grid, parcels,
    #        parcel_time_index=len(parcels.time_coordinates)-1)#, parcel_filter = parcel_filter)
    #plt.plot(rmg.node_x[outlet_nearest_raster_cell], rmg.node_y[outlet_nearest_raster_cell], "yo")
    #plt.title(f'Time = {round(elapsed_time,1)} s')

    fig = plt.figure()

        #grain size
    parcel_D = parcels.dataset.D.values
    parcel_D_off_grid=parcel_D[parcels.dataset["element_id"].values==-2]

    # the histogram of the data
    plt.hist(parcel_D_off_grid*1000, histtype='bar')
    plt.xlabel('grain size (mm)')
    plt.ylabel('Count')
    plt.title('Histogram of grain sizes that left grid')
    plt.text(0.011, 700, r'original distribution $\mu=2 mm$')
    plt.xlim(0, 20)
    plt.ylim(0, 4000)
    plt.grid(True)
    plt.show()
    if len(filepath) > 0:
        fig.savefig(filepath + "sed/" +str(int(elapsed_time)).zfill(5) + ".png", dpi = 150)
        plt.close(fig)
    else:
        plt.show()
def save_animations(source):

    source_list=np.asarray(["./output/flow/*", "./output/sed/*", "./output/floodplain/*"])
    animation_list = np.asarray(['./output/animationupland.gif', './output/animationsed.gif', './output/animationlowland.gif'])
    for i in range(0,3):
        images = []
        original_files=list(glob.glob(source_list[i]))
        original_files.sort(reverse=False)
        print(source_list[i])
        for file_ in original_files:
            images.append(imageio.imread(file_))
        imageio.mimsave(animation_list[i], images, duration=1/5, subrectangles=True)
    print("animations saved")