text:
#!/usr/bin/env python3
"""

Tools to summarise access-om2 runs.

Latest version: https://github.com/aekiss/run_summary
Author: Andrew Kiss https://github.com/aekiss
Apache 2.0 License http://www.apache.org/licenses/LICENSE-2.0.txt
"""

# TODO: use PAYU_N_RUNS - does this tell you whether the run is part of a sequence? if so we can determine queue wait for runs in a sequence - but sometimes it is None

from __future__ import print_function
import sys
try:
    assert sys.version_info >= (3, 3)  # need python >= 3.3 for print flush keyword
except AssertionError:
    print('\nFatal error: Python version too old.')
    print('On NCI, do the following and try again:')
    print('   module use /g/data/hh5/public/modules; module load conda/analysis3\n')
    raise

import os
import glob  # BUG: fails if payu module loaded - some sort of module clash with re
import subprocess
import datetime
import dateutil.parser
from collections import OrderedDict
import csv
import copy

try:
    import numpy as np
    import yaml
    import f90nml  # from https://f90nml.readthedocs.io/en/latest/
except ImportError:  # BUG: don't get this exception if payu module loaded, even if on python 2.6.6
    print('\nFatal error: modules not available.')
    print('On NCI, do the following and try again:')
    print('   module use /g/data/hh5/public/modules; module load conda/analysis3\n')
    raise
import nmltab  # from https://github.com/aekiss/nmltab
import warnings
warnings.simplefilter('ignore', np.RankWarning)

def num(s):
    """
    Return input string as int or float if possible, otherwise return string.
    """
    try:
        return int(s)
    except ValueError:
        try:
            return float(s)
        except ValueError:
            return s

def get_sync_path(fname):
    """
    Return GDATADIR or SYNCDIR path from sync_output_to_gdata.sh or sync_data.sh

    fname: file path to sync_output_to_gdata.sh or sync_data.sh etc

    output: directory string

    """
    dir = None  # default return value
    searchstrings = ['GDATADIR=', 'SYNCDIR=']
    with open(fname, 'r') as infile:
        for line in infile:
            # NB: subsequent matches will replace earlier ones
            for s in searchstrings:
                splitline = line.split(s)
                if len(splitline) == 2 and splitline[0].strip() == '':
                    dir_maybe = splitline[1].strip().rstrip('/')
                    if not(dir_maybe.startswith('/ERROR')):
                        dir = dir_maybe
    return dir

def parse_pbs_log(fname):
    """
    Return dict of items from parsed PBS log file.

    fname: PBS log file path

    output: dict

    example of PBS log file content to parse:
        qsub -q normal -P g40 -l walltime=12600 -l ncpus=2064 -l mem=8256GB -N minimal_01deg_j -l wd -j n -v PAYU_MODULENAME=payu/dev,PYTHONPATH=/projects/access/apps/mnctools/0.1/lib:/projects/access/apps/mnctools/0.1/lib:/projects/access/apps/mnctools/0.1/lib:/projects/v45/apps/payu/dev/lib:/projects/access/apps/mnctools/0.1/lib:/projects/v45/python,PAYU_CURRENT_RUN=137,PAYU_MODULEPATH=/projects/v45/modules,PAYU_N_RUNS=10 -lother=hyperthread -W umask=027 /projects/v45/apps/payu/dev/bin/payu-run
    ...
        git commit -am "2018-10-08 22:32:26: Run 137"
        TODO: Check if commit is unchanged
        ======================================================================================
                          Resource Usage on 2018-10-08 22:32:36:
           Job Id:             949753.r-man2
           Project:            x77
           Exit Status:        0
           Service Units:      20440.40
           NCPUs Requested:    5968                   NCPUs Used: 5968
                                                   CPU Time Used: 20196:31:07
           Memory Requested:   11.66TB               Memory Used: 2.61TB
           Walltime requested: 05:00:00            Walltime Used: 03:25:30
           JobFS requested:    36.43GB                JobFS used: 1.0KB
        ======================================================================================

    """
    def getproject(l):
        return l[1]

    def getpayuversion(l):
        return os.path.dirname(os.path.dirname(l[-1]))
        # return os.path.dirname([s for s in l[0].split(',')[0].split(':')
        #                         if s.find('payu') > -1][0])

    def getpayu(l):
        return l[0].split(',')[0]

    def getpayuint(l):
        return int(l[0].split(',')[0])

    def getrun(l):
        return int(l[4].rstrip('"'))

    def getjob(l):
        return int(l[1].split('.')[0])

    def getint(l):
        return int(l[1])

    def getfloat(l):
        return float(l[1])

    def getsec(l):  # convert hh:mm:ss to sec
        return sum(x * int(t) for x, t in zip([3600, 60, 1], l[1].split(':')))

    def getdatetime(l):  # BUG: doesn't include time zone (can't tell if we're on daylight savings time)
        return l[0]+'T'+l[1].rstrip(':')

    def getbytes(l):  # assumes PBS log info uses binary prefixes - TODO: check
        s = l[1]
        ns = s.strip('BKMGT')  # numerical part
        units = {'B': 1,
                 'KB': 2**10,
                 'MB': 2**20,
                 'GB': 2**30,
                 'TB': 2**40}
        return int(round(float(ns)*units[s[len(ns):]]))

    search_items = {  # keys are strings to search for; items are functions to apply to whitespace-delimited list of strings following key
        'PAYU_CURRENT_RUN': getpayuversion,  # gets path to payu; PAYU_CURRENT_RUN is redundant as this is obtained below from git commit message
        # 'PAYU_CURRENT_RUN=': getpayuint,  # BUG: misses some runs
        'PAYU_MODULENAME=': getpayu,
        'PAYU_MODULEPATH=': getpayu,
        'PAYU_PATH=': getpayu,
        'LD_LIBRARY_PATH=': getpayu,
        'PAYU_N_RUNS=': getpayuint,
        'PYTHONPATH=': getpayu,
# BUG: git commit will be missing if runlog: False in config.yaml - so we won't get run number!
        'git commit': getrun,  # instead of using PAYU_CURRENT_RUN; NB: run with this number might have failed - check Exit Status
        'Resource Usage on': getdatetime,
        'Job Id': getjob,
        'Project': getproject,
        'Exit Status': getint,
        'Service Units': getfloat,
        'NCPUs Requested': getint,
        'NCPUs Used': getint,
        'CPU Time Used': getsec,
        'Memory Requested': getbytes,
        'Memory Used': getbytes,
        'Walltime requested': getsec,
        'Walltime Used': getsec,
        'JobFS requested': getbytes,
        'JobFS used': getbytes}
    parsed_items = search_items.fromkeys(search_items, None)  # set defaults to None

    with open(fname, 'r') as infile:
        for line in infile:
            # NB: subsequent matches will replace earlier ones
            # NB: processes only the first match of each line
            for key, op in search_items.items():
                try:
                    parsed_items[key] = op(line.split(key)[1].split())
                except IndexError:  # key not present in this line
                    continue

    # change to more self-explanatory keys
    rename_keys = {'PAYU_CURRENT_RUN': 'payu version',
                   # 'PAYU_CURRENT_RUN=': 'Run number',
                   'git commit': 'Run number',
                   'Memory Requested': 'Memory Requested (bytes)',
                   'Memory Used': 'Memory Used (bytes)',
                   'Walltime requested': 'Walltime Requested (s)',
                   'Walltime Used': 'Walltime Used (s)',
                   'Resource Usage on': 'Run completion date'}
    for oldkey, newkey in rename_keys.items():
        parsed_items[newkey] = parsed_items.pop(oldkey)

    if parsed_items['Memory Requested (bytes)'] is None:
        parsed_items['Memory Requested (Gb)'] = None
    else:
        parsed_items['Memory Requested (Gb)'] = parsed_items['Memory Requested (bytes)']/2**30

    if parsed_items['Memory Used (bytes)'] is None:
        parsed_items['Memory Used (Gb)'] = None
    else:
        parsed_items['Memory Used (Gb)'] = parsed_items['Memory Used (bytes)']/2**30

    if parsed_items['Walltime Requested (s)'] is None:
        parsed_items['Walltime Requested (hr)'] = None
    else:
        parsed_items['Walltime Requested (hr)'] = parsed_items['Walltime Requested (s)']/3600

    if parsed_items['Walltime Used (s)'] is None:
        parsed_items['Walltime Used (hr)'] = None
    else:
        parsed_items['Walltime Used (hr)'] = parsed_items['Walltime Used (s)']/3600

    try:
        parsed_items['Timeout'] = parsed_items['Walltime Used (s)'] > parsed_items['Walltime Requested (s)']
    except:
        parsed_items['Timeout'] = None

    return parsed_items

def parse_git_log(basepath, datestr):
    """
    Return dict of items from git log from most recent commit before a given date.

    basepath: base directory path string

    datestr: date string

    output: dict
    """
    # possible BUG: what time zone flag should be use? local is problematic if run from overseas....?
    # use Popen for backwards-compatiblity with Python <2.7
    # pretty format is tab-delimited (%x09)
    try:
        p = subprocess.Popen('cd ' + basepath
                             + ' && git log -1 '
                             + '--pretty="format:%H%x09%an%x09%ai%x09%B" '
                             + '`git rev-list -1 --date=local --before="'
                             + datestr + '" HEAD`',  # TODO: add 1 sec to datestr so we don't rely on the delay between git commit and PBS log?
                             stdout=subprocess.PIPE, shell=True)
        log = p.communicate()[0].decode('ascii').split('\t')
        # log = p.communicate()[0].decode('ascii').encode('ascii').split('\t')  # for python 2.6
        log[3] = log[3].strip()  # strip whitespace from message
    except:
        log = [None]*4  # default values in case there's no .git, e.g. if runlog: False in config.yaml
    parsed_items = dict()
    parsed_items['Commit'] = log[0]
    parsed_items['Author'] = log[1]
    parsed_items['Date'] = log[2]
    parsed_items['Message'] = log[3]
    return parsed_items

def parse_mom_time_stamp(paths):
    """
    Return dict of items from parsed MOM time_stamp.out.

    paths: list of base paths

    output: dict parsed from first matching time_stamp.out in paths

    example of MOM time_stamp.out content to parse:
        2001   9   1   0   0   0  Sep
        2001  11   1   0   0   0  Nov

    """
    parsed_items = dict()
    keys = ['Model start time', 'Model end time']
    for path in paths:
        fname = os.path.join(path, 'ocean/time_stamp.out')
        if os.path.isfile(fname):
            parsed_items['Time stamp file'] = fname
            with open(fname, 'r') as infile:
                for key in keys:
                    line = infile.readline()
                    parsed_items[key] = datetime.datetime(
                        *list(map(int, line.split()[0:-1]))).isoformat()
            break
    try:
        d1 = dateutil.parser.parse(parsed_items[keys[0]])
        d2 = dateutil.parser.parse(parsed_items[keys[1]])
        duration = d2-d1  # BUG: presumably assumes Gregorian calendar with leap years and time in UTC
        parsed_items['Model run length (s)'] = duration.total_seconds()
        parsed_items['Model run length (days)'] = duration.total_seconds()/3600/24
    except KeyError:
        pass
    return parsed_items

def parse_yaml(paths, filename):
    """
    Return dict of items from parsed yaml file.

    paths: list of base paths
    filename: yaml filename to attempt to read from base paths

    output: dict parsed from first matching filename in paths
    """
    parsed_items = dict()
    for path in paths:
        fpath = os.path.join(path, filename)
        if os.path.isfile(fpath):
            with open(fpath, 'r') as infile:
                # Need to use load_all to handle manifests. Only return final part.
                parsed_items = list(yaml.load_all(infile, Loader=yaml.FullLoader))[-1]
            break
    return parsed_items

def parse_accessom2_out(paths):
    """
    Return dict of items from parsed access-om2.out.

    paths: list of base paths

    output: dict of timing names, with dict of statistics

    NB: output may also contain bad data from intermingled CICE output.
    """
    parsed_items = dict()
    for path in paths:
        fname = os.path.join(path, 'access-om2.out')
        if os.path.isfile(fname):
            with open(fname, 'r') as infile:
                for l in infile:
                    if l.startswith('Tabulating mpp_clock statistics'):
                        break
                for l in infile:
                    if l.startswith('                                          tmin'):
                        break
                keys = l.split()
                for l in infile:
                    if l.startswith(' MPP_STACK high water mark='):
                        break
                    name = l[0:32].strip()  # relies on name being cropped at 32 chars
                    vals = [num(n) for n in l[32:].split()]
                    parsed_items[name] = dict(zip(keys, vals))
            break
    return parsed_items

def parse_ice_diag_d(paths):
    """
    Return dict of cice info from ice/ice_diag.d.

    paths: list of base paths

    output: dict
    """
    # this is pretty rough-and-ready, e.g. repeated entries end up containing the final value
    parsed_items = dict()
    for path in paths:
        fname = os.path.join(path, 'ice/ice_diag.d')
        if os.path.isfile(fname):
            with open(fname, 'r') as infile:
                for l in infile:
                    if l.startswith('Timing information:'):
                        break  # handle timing data with parse_cice_timing
                    try:
                        key = l.split('=')[0].strip()
                        val = num(l.split('=')[1].strip())
                        parsed_items[key] = val
                    except:
                        try:
                            key = l.split(':')[0].strip()
                            val = num(l.split(':')[1].strip())
                            parsed_items[key] = val
                        except:
                            pass
            break

    if 'Block size:  nx_block' in parsed_items:
        parsed_items['nx_block'] = parsed_items['Block size:  nx_block']
# NB: in ice_blocks.F90
#       nx_block = block_size_x + 2*nghost,   &!  x,y dir including ghost
#       ny_block = block_size_y + 2*nghost     !  cells
    if 'Number of ghost cells' in parsed_items:
        if 'nx_block' in parsed_items:
            parsed_items['block_size_x'] = parsed_items['nx_block'] - 2*parsed_items['Number of ghost cells']
        if 'ny_block' in parsed_items:
            parsed_items['block_size_y'] = parsed_items['ny_block'] - 2*parsed_items['Number of ghost cells']
    parsed_items['timing'] = parse_cice_timing(paths)
    return parsed_items

def parse_cice_timing(paths):
    """
    Return dict of cice timing info from ice/ice_diag.d.

    paths: list of base paths

    output: dict of timing names, with dict of statistics
    """
#     sample to parse:
# Timing information:
#
# Timer   1:     Total   10894.88 seconds
#   Timer stats (node): min =    10894.69 seconds
#                       max =    10894.88 seconds
#                       mean=    10894.70 seconds
#   Timer stats(block): min =        0.00 seconds
#                       max =        0.00 seconds
#                       mean=        0.00 seconds
# Timer   2:  TimeLoop   10802.50 seconds
#   Timer stats (node): min =    10802.33 seconds
#                       max =    10802.50 seconds
#                       mean=    10802.33 seconds
#   Timer stats(block): min =        0.00 seconds
#                       max =        0.00 seconds
#                       mean=        0.00 seconds

    parsed_items = dict()
    for path in paths:
        fname = os.path.join(path, 'ice/ice_diag.d')
        if os.path.isfile(fname):
            with open(fname, 'r') as infile:
                for l in infile:
                    if l.startswith('Timing information:'):
                        break
                for l in infile:
                    if l.startswith('Timer'):  # ignore time is it it node max
                        timerkey = ' '.join(l[0:21].split()[2:])
                        parsed_items[timerkey] = dict()
                    else:
                        if l.startswith('  Timer'):
                            typekey = l.split('(')[-1].split(')')[0]
                            parsed_items[timerkey][typekey] = dict()
                        try:
                            key = l.split('=')[0].split()[-1]
                            val = num(l.split()[-2])
                            parsed_items[timerkey][typekey][key] = val
                        except:
                            pass
            break
    return parsed_items

def parse_nml(paths):
    """
    Return dict of items from parsed namelists.

    paths: list of base paths to parse for namelists

    output: dict
    """
    parsed_items = dict()
    parsed_items['accessom2.nml'] = None  # default value for non-YATM run
    for path in paths:
        fnames = [os.path.join(path, 'accessom2.nml')]\
                + glob.glob(os.path.join(path, '*/*.nml'))
        for fname in fnames:
            if os.path.isfile(fname):  # no accessom2.nml for non-YATM run
                parsed_items[fname.split(path)[1].strip('/')] \
                        = f90nml.read(fname)
    return parsed_items

def git_diff(basepath, sha1, sha2):
    """
    Return dict of git-tracked differences between two commits.

    basepath: base directory path string

    sha1, sha2: strings; sha1 should be earlier than or same as sha2
    """
    try:
        p = subprocess.Popen('cd ' + basepath
                             + ' && git diff --name-only ' + sha1 + ' ' + sha2,
                             stdout=subprocess.PIPE, shell=True)
        c = ', '.join(
            p.communicate()[0].decode('ascii').split())
        p = subprocess.Popen('cd ' + basepath
                             + ' && git log --ancestry-path --pretty="%B\%x09" '
                             + sha1 + '..' + sha2,
                             stdout=subprocess.PIPE, shell=True)
        m = [s.strip('\n\\')
             for s in p.communicate()[0].decode('ascii').split('\t')][0:-1]
        m.reverse()  # put in chronological order
        if len(m) == 0:
            m = None
    except:
        c = None
        m = None
    parsed_items = dict()
    parsed_items['Changed files'] = c
    parsed_items['Messages'] = m  # NB: will be None if there's no direct ancestry path from sha1 to sha2)
    return parsed_items

def dictget(d, l):
    """
    Lookup item in nested dict using a list of keys, or None if non-existent

    d: nested dict
    l: list of keys, or None
    """
    try:
        dl0 = d[l[0]]
    except (KeyError, TypeError):
        return None
    if len(l) == 1:
        return dl0
    return dictget(dl0, l[1:])

def keylists(d):
    """
    Return list of key lists to every leaf node in a nested dict.
    Each key list can be used as an argument to dictget.

    d: nested dict
    """
    l = []
    for k, v in d.items():
        if isinstance(v, dict):
            sublists = keylists(v)
            for sli in sublists:
                l.append([k]+sli)
        else:
            l.append([k])
    return l

def keylistssuperset(d):
    """
    Apply keylists to every value at the top level of input dict d,
    and return a list containing one instance of every key list found.
    Each key list can be used as an argument to dictget.

    d: nested dict whose values are to be supplied to keylists
    """
    all = set()
    for v in d.values():
        all.update(['\b'.join(l) for l in keylists(v)])
    return [s.split('\b') for s in all]

def tryfunc(func, arg):
    """
    Return func(arg) or None if there's an exception.
    """
    try:
        return func(arg)
    except:
        return None

def run_summary(basepath=os.getcwd(), outfile=None, list_available=False,
                dump_all=False, show_fails=False, outfile_syncdir=False,
                no_header=False, no_stats=False):
    """
    Generate run summary
    """
    basepath = os.path.abspath(basepath)
    archive_path = os.path.realpath(os.path.join(basepath, 'archive'))
    print('Generating run summary of ' + basepath, end='')

    # get jobname from config.yaml -- NB: we assume this is the same for all jobs
    with open(os.path.join(basepath, 'config.yaml'), 'r') as infile:
        configyaml = yaml.load(infile, Loader=yaml.FullLoader)
        jobname = configyaml.get('jobname')
        sync_script = configyaml.get('postscript')

    if sync_script:
        sync_path = get_sync_path(os.path.join(basepath, sync_script))
    else:
        sync_path = None

    if outfile is None:
        if outfile_syncdir and sync_path:
            outfile = 'run_summary_' + sync_path.strip(os.sep).replace(os.sep, '_') + '.csv'
        else:
            outfile = 'run_summary_' + basepath.strip(os.sep).replace(os.sep, '_') + '.csv'
        # if show_fails:
        #     outfile = os.path.splitext(outfile)[0]+'_fails.csv'

    try:
        p = subprocess.Popen('cd ' + basepath
                             + ' && git rev-parse --abbrev-ref HEAD',
                             stdout=subprocess.PIPE, shell=True)
        git_branch = p.communicate()[0].decode('ascii').strip()
    except:
        git_branch = None

    # get data from all PBS job logs
    run_data = dict()
    # NB: match jobname[:15] because in some cases the pbs log files use a shortened version of the jobname in config.yaml
    # e.g. see /home/157/amh157/payu/025deg_jra55_ryf8485
    # NB: logs in archive may be duplicated in sync_path, in which case the latter is used
    logfiles = glob.glob(os.path.join(archive_path, 'pbs_logs', jobname[:15] + '*.o*'))\
             + glob.glob(os.path.join(basepath, jobname[:15] + '*.o*'))
    if sync_path:
             logfiles += glob.glob(os.path.join(sync_path, 'pbs_logs', jobname[:15] + '*.o*'))
    logfiles = [f for f in logfiles if '_c.o' not in f]  # exclude collation files *_c.o*
    for f in logfiles:
        print('.', end='', flush=True)
        jobid = int(f.split('.o')[1])
        run_data[jobid] = dict()
        run_data[jobid]['PBS log'] = parse_pbs_log(f)
        run_data[jobid]['PBS log']['PBS log file'] = f
        # fudge: these paths might actually apply only to the latest job
        run_data[jobid]['paths'] = dict()
        run_data[jobid]['paths']['Control path'] = basepath
        run_data[jobid]['paths']['Sync path'] = sync_path
        run_data[jobid]['paths']['Archive path'] = archive_path
        run_data[jobid]['storage'] = dict()

    # get run data for all jobs
    for jobid in run_data:
        print('.', end='', flush=True)
        pbs = run_data[jobid]['PBS log']
        date = pbs['Run completion date']  # BUG: would be better to have time when run began, including time zone
        if date is not None:
            run_data[jobid]['git log'] = parse_git_log(basepath, date)
            # BUG: assumes no commits between run start and end
            # BUG: assumes the time zones match - no timezone specified in date - what does git assume? UTC?
            if pbs['Exit Status'] == 0:  # output dir belongs to this job only if Exit Status = 0
                outdir = 'output' + str(pbs['Run number']).zfill(3)
                restartdir = 'restart' + str(pbs['Run number']).zfill(3)
                paths = []

                if sync_path:
                    sync_output_path = os.path.join(sync_path, outdir)
                    if os.path.isdir(sync_output_path):
                        paths += [sync_output_path]
                        run_data[jobid]['paths']['Sync output path'] =\
                            sync_output_path
                    sync_restart_path = os.path.join(sync_path, restartdir)
                    if os.path.isdir(sync_restart_path):
                        run_data[jobid]['paths']['Sync restart path'] =\
                            sync_restart_path

                archive_output_path = os.path.join(archive_path, outdir)
                if os.path.isdir(archive_output_path):
                    paths += [archive_output_path]
                    run_data[jobid]['paths']['Archive output path'] =\
                        archive_output_path
                archive_restart_path = os.path.join(archive_path, restartdir)
                if os.path.isdir(archive_restart_path):
                    run_data[jobid]['paths']['Archive restart path'] =\
                        archive_restart_path

                # 'Sync output path' if it exists, otherwise 'Archive output path'
                run_data[jobid]['paths']['Output path'] =\
                    run_data[jobid]['paths'].get('Sync output path') or\
                    run_data[jobid]['paths'].get('Archive output path')
                # 'Sync restart path' if it exists, otherwise 'Archive restart path'
                run_data[jobid]['paths']['Restart path'] =\
                    run_data[jobid]['paths'].get('Sync restart path') or\
                    run_data[jobid]['paths'].get('Archive restart path')

                # find GiB for output and restart
                for k in ['Output path', 'Restart path']:
                    path = run_data[jobid]['paths'][k]
                    if path:
                        p = subprocess.Popen('du -bs ' + path,
                                             stdout=subprocess.PIPE, shell=True)
                        ret = p.communicate()[0].decode('ascii')
                        bytes = int(ret.split()[0])
                        run_data[jobid]['storage'][k + ' GiB'] = \
                            round(bytes/1073741824, 3)

                run_data[jobid]['MOM_time_stamp.out'] = parse_mom_time_stamp(paths)
                run_data[jobid]['namelists'] = parse_nml(paths)
                run_data[jobid]['access-om2.out'] = parse_accessom2_out(paths)
                run_data[jobid]['ice_diag.d'] = parse_ice_diag_d(paths)
                run_data[jobid]['metadata.yaml'] = parse_yaml([basepath, sync_path], 'metadata.yaml')
                for fn in ['config.yaml', 'env.yaml', 'job.yaml',
                           'manifests/exe.yaml', 'manifests/input.yaml', 'manifests/restart.yaml']:
                    run_data[jobid][fn] = parse_yaml(paths, fn)

    all_run_data = copy.deepcopy(run_data)  # all_run_data includes failed jobs

    if show_fails:
        # remove all jobs that have no PBS info in log file
        for jobid in all_run_data:
            if all_run_data[jobid]['PBS log']['Run completion date'] is None:
                del run_data[jobid]
        # (jobid, run completion date) tuples sorted by run completion date
        jobid_run_tuples = sorted([(k, v['PBS log']['Run completion date'])
                                   for (k, v) in run_data.items()],
                                  key=lambda t: t[1])
        if len(jobid_run_tuples) == 0:
            print('\nAborting: no jobs?')
            return
        # jobid keys into run_data sorted by run completion date
        sortedjobids = [k[0] for k in jobid_run_tuples]
    else:
        # remove failed jobs from run_data
        for jobid in all_run_data:
            print('.', end='', flush=True)
            pbs = all_run_data[jobid]['PBS log']
            date = pbs['Run completion date']
            if date is None:  # no PBS info in log file
                del run_data[jobid]
            elif pbs['Run number'] is None:  # not a model run log file
                del run_data[jobid]
            elif pbs['Exit Status'] != 0:  # output dir belongs to this job only if Exit Status = 0
                del run_data[jobid]
            elif len(run_data[jobid]['config.yaml']) == 0:  # output dir missing
                del run_data[jobid]

        # (jobid, run number) tuples sorted by run number - re-done below
        jobid_run_tuples = sorted([(k, v['PBS log']['Run number'])
                                   for (k, v) in run_data.items()],
                                  key=lambda t: t[1])
        if len(jobid_run_tuples) == 0:
            print('\nAborting: no successful jobs?')
            return

    # Remove the older jobid if run number is duplicated - assume run was re-done
    # (check by date rather than jobid, since jobid sometimes rolls over)
        prev_jobid_run = jobid_run_tuples[0]
        for jobid_run in jobid_run_tuples[1:]:
            if jobid_run[1] == prev_jobid_run[1]:  # duplicated run number
                if run_data[jobid_run[0]]['PBS log']['Run completion date']\
                 > run_data[prev_jobid_run[0]]['PBS log']['Run completion date']:
                    del run_data[prev_jobid_run[0]]
                    prev_jobid_run = jobid_run
                else:
                    del run_data[jobid_run[0]]
            else:
                prev_jobid_run = jobid_run

        # re-do (jobid, run number) tuples sorted by run number
        jobid_run_tuples = sorted([(k, v['PBS log']['Run number'])
                                   for (k, v) in run_data.items()],
                                  key=lambda t: t[1])
        if len(jobid_run_tuples) == 0:
            print('\nAborting: no successful jobs?')
            return

        # jobid keys into run_data sorted by run number
        sortedjobids = [k[0] for k in jobid_run_tuples]

        # allow referencing by submodel name as well as list index
        for jobid in run_data:
            run_data[jobid]['config.yaml']['submodels-by-name'] = dict()
            for sm in run_data[jobid]['config.yaml']['submodels']:
                run_data[jobid]['config.yaml']['submodels-by-name'][sm['name']] = sm

        # make a 'timing' entry to contain model timestep and run length for both MATM and YATM runs
        # run length is [years, months, days, seconds] to accommodate both MATM and YATM
        prevjobid = -1
        for jobid in sortedjobids:
            r = run_data[jobid]
            timing = dict()
            if r['namelists']['accessom2.nml'] is None:  # non-YATM run
                timing['Timestep'] = r['config.yaml']['submodels'][1]['timestep']  # MOM timestep
                rt = r['config.yaml']['calendar']['runtime']
                timing['Run length'] = [rt['years'], rt['months'], rt['days'], 0]  # insert 0 seconds
            else:
                timing['Timestep'] = r['namelists']['accessom2.nml']['accessom2_nml']['ice_ocean_timestep']
                rp = r['namelists']['accessom2.nml']['date_manager_nml']['restart_period']
                timing['Run length'] = rp[0:2] + [0] + [rp[2]]  # insert 0 days
            yrs = r['MOM_time_stamp.out']['Model run length (days)']/365.25  # FUDGE: assumes 365.25-day year
            timing['SU per model year'] = r['PBS log']['Service Units']/yrs
            timing['Walltime (hr) per model year'] = r['PBS log']['Walltime Used (hr)']/yrs
            storagekeys = list(r['storage'].keys())
            for k in storagekeys:
                timing[k + ' per model year'] = round(r['storage'][k]/yrs, 3)

            if prevjobid >= 0:  # also record time including wait between runs
                d1 = dateutil.parser.parse(run_data[prevjobid]['PBS log']['Run completion date'])
                d2 = dateutil.parser.parse(r['PBS log']['Run completion date'])
                tot_walltime = (d2-d1).total_seconds()/3600
                timing['Walltime (hr) between this completion and previous completion'] = tot_walltime
                timing['Wait (hr) between this run and previous'] = tot_walltime - r['PBS log']['Walltime Used (hr)']
                timing['SU per calendar day'] = r['PBS log']['Service Units']/tot_walltime*24
                timing['Model years per calendar day'] = yrs/tot_walltime*24
                for k in storagekeys:
                    timing[k + ' per calendar day'] = round(r['storage'][k]/tot_walltime*24, 3)

            r['timing'] = timing
            prevjobid = jobid

        # include changes in all git commits since previous run
        for i, jobid in enumerate(sortedjobids):
            print('.', end='', flush=True)
            run_data[jobid]['git diff'] = \
                git_diff(basepath,
                         run_data[sortedjobids[max(i-1, 0)]]['git log']['Commit'],
                         run_data[jobid]['git log']['Commit'])

        # count failed jobs prior to each successful run
        # BUG: always have zero count between two successful runs straddling a jobid rollover
        # BUG: first run also counts all fails after a rollover
        prevjobid = -1
        for jobid in sortedjobids:
            c = [e for e in all_run_data.keys() if e > prevjobid and e < jobid
                 and e not in run_data]
            c.sort()
            run_data[jobid]['PBS log']['Failed previous jobids'] = c
            run_data[jobid]['PBS log']['Failed previous jobs'] = len(c)
            prevjobid = jobid

    if list_available:
        print('\nAvailable data which can be tabulated if added to output_format')
        print('(but you may need to edit some keys to ensure uniqueness):')
        keylist = []
        for k in keylistssuperset(run_data):
            keylist.append((k[-1], "['" + "', '".join(k) + "']"))
        keylist.sort(key = lambda x: x[1])
        maxkeywidth = max([len(k[0]) for k in keylist])
        for k in keylist:
            print("        ('" + k[0] + "', " + " "*(maxkeywidth-len(k[0])) + k[1] + "),")

    if dump_all:
        dumpoutfile = os.path.splitext(outfile)[0]+'.yaml'
        print('\nWriting', dumpoutfile)
        with open(dumpoutfile, 'w') as outf:
            yaml.dump(run_data, outf, default_flow_style=False)

    ###########################################################################
    # Specify the output format here.
    ###########################################################################
    # output_format is a OrderedDict of (key, value) tuples, one for each column.
    # keys are column headers (arbitrary but must be unique)
    # values are lists of keys into run_data (omitting job id)
    # "run_summary.py --list" will list all available data you can add here
    #    (but you may need to edit some keys to ensure uniqueness)
    output_format = OrderedDict([
        ('Run', ['PBS log', 'Run number']),
        ('Run start', ['MOM_time_stamp.out', 'Model start time']),
        ('Run end', ['MOM_time_stamp.out', 'Model end time']),
        ('Run length (y, m, d, s)', ['timing', 'Run length']),
        ('Run length (days)', ['MOM_time_stamp.out', 'Model run length (days)']),
        ('Control directory', ['paths', 'Control path']),
        # ('Archive directory', ['paths', 'Archive path']),
        # ('Sync directory', ['paths', 'Sync path']),
        ('Output directory', ['paths', 'Output path']),
        ('Output GiB', ['storage', 'Output path GiB']),
        ('Restart directory', ['paths', 'Restart path']),
        ('Restart GiB', ['storage', 'Restart path GiB']),
        ('Run by', ['git log', 'Author']),
        ('Run completion date', ['PBS log', 'Run completion date']),
        ('Job Id', ['PBS log', 'Job Id']),
        ('Failed jobs', ['PBS log', 'Failed previous jobs']),
        ('Failed jobids', ['PBS log', 'Failed previous jobids']),
        ('Queue', ['config.yaml', 'queue']),
        ('Service Units', ['PBS log', 'Service Units']),
        ('Walltime Used (hr)', ['PBS log', 'Walltime Used (hr)']),
        ('SU per model year', ['timing', 'SU per model year']),
        ('Walltime (hr) per model year', ['timing', 'Walltime (hr) per model year']),
        ('Wait (hr) between runs', ['timing', 'Wait (hr) between this run and previous']),
        ('SU per calendar day', ['timing', 'SU per calendar day']),
        ('Model years per calendar day', ['timing', 'Model years per calendar day']),
        ('Memory Used (Gb)', ['PBS log', 'Memory Used (Gb)']),
        ('NCPUs Used', ['PBS log', 'NCPUs Used']),
        ('MOM NCPUs', ['config.yaml', 'submodels-by-name', 'ocean', 'ncpus']),
        ('CICE NCPUs', ['config.yaml', 'submodels-by-name', 'ice', 'ncpus']),
        # ('Max Ocean diagnostics (s)', ['access-om2.out', '(Ocean diagnostics)', 'tmax']),
        # ('Max Ocean diagnostics: tracer (s)', ['access-om2.out', '(Ocean diagnostics: tracer)', 'tmax']),
        ('Fraction of MOM runtime in oasis_recv', ['access-om2.out', 'oasis_recv', 'tfrac']),
        ('Max MOM wait for oasis_recv (s)', ['access-om2.out', 'oasis_recv', 'tmax']),
        ('Max CICE wait for coupler (s)', ['ice_diag.d', 'timing', 'waiting_o', 'node', 'max']),
        ('Max CICE I/O time (s)', ['ice_diag.d', 'timing', 'ReadWrite', 'node', 'max']),
        ('MOM tile layout', ['namelists', 'ocean/input.nml', 'ocean_model_nml', 'layout']),
        ('CICE tile distribution', ['namelists', 'ice/cice_in.nml', 'domain_nml', 'distribution_type']),
        ('CICE block_size_x', ['ice_diag.d', 'block_size_x']),
        ('CICE block_size_y', ['ice_diag.d', 'block_size_y']),
        ('Timestep (s)', ['timing', 'Timestep']),
        ('MOM barotropic split', ['namelists', 'ocean/input.nml', 'ocean_model_nml', 'barotropic_split']),
        ('CICE dynamic split (ndtd)', ['namelists', 'ice/cice_in.nml', 'setup_nml', 'ndtd']),
        # ('ktherm', ['namelists', 'ice/cice_in.nml', 'thermo_nml', 'ktherm']),
        # ('Common inputs', ['config.yaml', 'input']),
        # ('Atmosphere executable', ['config.yaml', 'submodels-by-name', 'atmosphere', 'exe']),
        # ('Atmosphere inputs', ['config.yaml', 'submodels-by-name', 'atmosphere', 'input']),
        # ('MOM executable', ['config.yaml', 'submodels-by-name', 'ocean', 'exe']),
        # ('MOM inputs', ['config.yaml', 'submodels-by-name', 'ocean', 'input']),
        # ('CICE executable', ['config.yaml', 'submodels-by-name', 'ice', 'exe']),
        # ('CICE inputs', ['config.yaml', 'submodels-by-name', 'ice', 'input']),
        # ('Payu version', ['PBS log', 'payu version']),
        ('Git hash of run', ['git log', 'Commit']),
        ('Commit date', ['git log', 'Date']),
        ('Git-tracked file changes', ['git diff', 'Changed files']),
        ('Git log messages', ['git diff', 'Messages']),
        ])
    SUdata = [dictget(run_data, [jobid] + ['PBS log', 'Service Units'])
              for jobid in sortedjobids]
    stats = OrderedDict([  # tuples: (label, function)
        ('Total', sum),
        ('Mean', np.mean),
        ('Median', np.median),
        ('Min', min),
        ('Max', max),
        ('Std dev', np.std),
        ('SU correlation', lambda x: np.corrcoef(x, SUdata)[0, 1]),
        ('SU slope', lambda x: np.polyfit(x, SUdata, 1)[0]),
        ('Dimensionless SU slope', lambda x: np.polyfit(x, SUdata, 1)[0]*np.mean(x)/np.mean(SUdata))
        ])
    ###########################################################################
    if no_stats:
        stats = OrderedDict([])
    if show_fails:
        # output crash-related info (redefines order of any keys already in output_format)
        output_format_prefix = OrderedDict([
            ('Job Id', ['PBS log', 'Job Id']),
            ('Run completion date', ['PBS log', 'Run completion date']),
            ('Exit Status', ['PBS log', 'Exit Status']),
            ('Timeout', ['PBS log', 'Timeout']),
            ('Walltime Requested (hr)', ['PBS log', 'Walltime Requested (hr)']),
            ('Walltime Used (hr)', ['PBS log', 'Walltime Used (hr)']),
            ('qsub_flags', ['config.yaml', 'qsub_flags']),
            ])
        output_format_prefix.update(output_format)
        output_format = output_format_prefix
    else:
        # output all namelist changes
        output_format_nmls = OrderedDict()
        nmls_any_runs = set(run_data[list(run_data.keys())[0]]['namelists'].keys())
        nmls_all_runs = nmls_any_runs
        # avoid dict comprehension here to avoid python<2.7 syntax error
        nmls_no_runs = dict([(k, True) for k in nmls_any_runs])  # True for namelists that are None for all runs
        # nmls_no_runs = {k: True for k in nmls_any_runs}  # True for namelists that are None for all runs
        for jobid in run_data:
            run_nmls = run_data[jobid]['namelists']
            nmls_any_runs = set(run_nmls.keys()) | nmls_any_runs
            nmls_all_runs = set(run_nmls.keys()) & nmls_all_runs
            for nml in set(nmls_all_runs):
                if run_nmls[nml] is None:
                    nmls_all_runs.remove(nml)
            for nml in run_nmls:
                newnone = (nml is None)
                if nml in nmls_no_runs:
                    nmls_no_runs[nml] = nmls_no_runs[nml] and newnone
                else:
                    nmls_no_runs.update({nml: newnone})
        for nml in set(nmls_any_runs):
            if nmls_no_runs[nml]:
                nmls_any_runs.remove(nml)

        # add every changed group/variable in nml files that exist in all runs
        for nml in nmls_all_runs:
            # avoid dict comprehension here to avoid python<2.7 syntax error
            nmllistall = dict([(jobid,
                              copy.deepcopy(run_data[jobid]['namelists'][nml]))
                              for jobid in run_data])
            # nmllistall = {jobid: copy.deepcopy(run_data[jobid]['namelists'][nml])
            #               for jobid in run_data}
            groups = nmltab.superset(nmltab.nmldiff(nmllistall))
            for group in groups:
                for var in groups[group]:
                    ngv = [nml, group, var]
                    output_format_nmls.update(OrderedDict([
                        (' -> '.join(ngv), ['namelists'] + ngv)]))

        # add all group/variables in nml files that exist in only some runs
        for nml in nmls_any_runs - nmls_all_runs:
            nmllistall = dict()
            for jobid in run_data:
                if nml in run_data[jobid]['namelists']:
                    if run_data[jobid]['namelists'][nml] is not None:
                        nmllistall.update({jobid:
                              copy.deepcopy(run_data[jobid]['namelists'][nml])})
            groups = nmltab.superset(nmllistall)
            for group in groups:
                for var in groups[group]:
                    ngv = [nml, group, var]
                    output_format_nmls.update(OrderedDict([
                        (' -> '.join(ngv), ['namelists'] + ngv)]))

        # alphabetize
        output_format_nmls = OrderedDict([(k, output_format_nmls[k])
                                 for k in sorted(output_format_nmls.keys())])

        # add output_format entries for every namelist variable that has changed in any run
        output_format.update(output_format_nmls)

    # output csv file according to output_format above
    print('\nWriting', outfile)

    if len(stats) > 0:
        lhcol = [None]
    else:
        lhcol = []  # don't allow space for stats headings if we don't have any

    with open(outfile, 'w', newline='') as csvfile:
        csvw = csv.writer(csvfile, dialect='excel', lineterminator='\n')

        if not no_header:
            csvw.writerow(['Summary report generated by run_summary.py, https://github.com/aekiss/run_summary'])
            csvw.writerow(['report generated:', datetime.datetime.now().replace(microsecond=0).astimezone().isoformat()])
            csvw.writerow(['control directory path:', basepath, 'git branch:', git_branch])
            csvw.writerow(['output path:', sync_path])

        csvw.writerow(lhcol + list(output_format.keys()))  # header
        for jobid in sortedjobids:  # output a row for each jobid
            csvw.writerow(lhcol + [dictget(run_data, [jobid] + keylist) for keylist in output_format.values()])

        if len(stats) > 0:
            # calculate and save summary stats
            statsdata = copy.deepcopy(output_format)
            for k, keylist in output_format.items():  # calculate summary stats
                coldata = [dictget(run_data, [jobid] + keylist) for jobid in sortedjobids]
                coldata = [c for c in coldata if c is not None]
                statsdata[k] = {label:tryfunc(func, coldata) for (label, func) in stats.items()}
            # write summary stats
            csvw.writerow(lhcol + [None]*len(list(output_format.keys())))  # blank row
            csvw.writerow(lhcol + list(output_format.keys()))  # header
            for s in stats:
                csvw.writerow([s] + [statsdata[k][s] for k in statsdata.keys()])

        print('Done.')
    return

if __name__ == '__main__':
    import argparse
    parser = argparse.ArgumentParser(description=
        'Summarise ACCESS-OM2 runs.\
        Latest version and help: https://github.com/aekiss/run_summary')
    parser.add_argument('-f', '--show_fails',
                        action='store_true', default=False,
                        help='include failed runs (disables some output columns)')
    parser.add_argument('-l', '--list',
                        action='store_true', default=False,
                        help='list all data that could be tabulated by adding it to output_format')
    parser.add_argument('-d', '--dump_all',
                        action='store_true', default=False,
                        help='also dump all data to <outfile>.yaml')
    parser.add_argument('-o', '--outfile', type=str,
                        metavar='file',
                        default=None,
                        help="output file path; default is\
                        'run_summary_<path>.csv';\
                        overrides --outfile_syncdir if set.\
                        WARNING: output file will be overwritten")
    parser.add_argument('--outfile_syncdir',
                        action='store_true', default=False,
                        help="set output file path to\
                        'run_summary_<sync dir path>.csv'\
                        or 'run_summary_<path>.csv' if sync dir\
                        path is invalid;\
                        ignored if '-o', '--outfile' is set.\
                        WARNING: output file will be overwritten")
    parser.add_argument('--no_header',
                        action='store_true', default=False,
                        help="don't write header rows in output .csv")
    parser.add_argument('--no_stats',
                        action='store_true', default=False,
                        help="don't output summary statistics")
    parser.add_argument('path', metavar='path', type=str, nargs='*',
                        help='zero or more ACCESS-OM2 control directory paths; default is current working directory')
    args = parser.parse_args()
    show_fails = vars(args)['show_fails']
    lst = vars(args)['list']
    dump_all = vars(args)['dump_all']
    outfile = vars(args)['outfile']
    outfile_syncdir = vars(args)['outfile_syncdir']
    no_header = vars(args)['no_header']
    no_stats = vars(args)['no_stats']
    basepaths = vars(args)['path']  # a list of length >=0 since nargs='*'
    if not basepaths:
        basepaths = [os.getcwd()]
    for bp in basepaths:
        try:
            run_summary(show_fails=show_fails, basepath=bp,
                        outfile=outfile, list_available=lst,
                        dump_all=dump_all,
                        outfile_syncdir=outfile_syncdir,
                        no_header=no_header,
                        no_stats=no_stats)
        except:
            print('\nFailed. Error:', sys.exc_info())


text:
import torch
import argparse
from pandas import json_normalize
import torch.optim as optim
from torch.optim import *
from kme.data.utils import get_loaders, DATASETS, REVERSE_TRANSFORM, TRANSFORMS
from kme.models.utils import build_kme_net, count_parameters
from torch.utils.tensorboard import SummaryWriter
from kme.tools.checkpoint import load_checkpoint, save_checkpoint
from kme.tools.training import train_routine, test_routine
from kme.tools.config import load_config, save_config, DATASET_KEY, TRAINING_KEY, MODEL_KEY, CKPT_KEY
from kme.data.text_classification import DATASETS as TXT_DS, MAX_SENTENCE_LEN
from torch.nn import DataParallel
from kme.extern.radam import RAdam
import tqdm
import os
import numpy as np
from sklearn.neighbors import KDTree
import pickle
from captum.attr import IntegratedGradients, Saliency, InputXGradient
from captum.attr._utils.visualization import visualize_image_attr
from torchvision.transforms import Resize
from matplotlib import pyplot as plt
from kme.tools.interpreting import match_single_patches
from matplotlib.cm import get_cmap
import pandas as pd

IMAGE_DATASETS = {'cifar10', 'cifar100', 'cub', 'svhn', 'svhn_no_augment',
                  'mnist', 'fashionmnist', 'imagenet', 'cub_no_augment'}
N_CHANNELS = {
    'mnist': 1,
    'fashionmnist': 1,
    'cifar10': 3,
    'cifar100': 3,
    'cub': 3,
    'cub_no_augment': 3,
    'svhn': 3,
    'imagenet': 3,
    'svhn_no_augment': 3
}
IMG_SZ = {
    'mnist': 28,
    'fashionmnist': 28,
    'cifar10': 32,
    'cifar100': 32,
    'cub': 224,
    'cub_no_augment': 224,
    'svhn': 32,
    'imagenet': 224,
    'svhn_no_augment': 32
}
TEXT_DATASETS = set(TXT_DS.keys())
TABULAR_DATASETS = set(DATASETS.keys()).difference(IMAGE_DATASETS)

X_fname = 'X.npy'
Y_fname = 'Y.npy'
probs_fname = 'probs.npy'
latent_fname = 'latent.npy'
latent_norms_fname = 'latent_norms.npy'
latent_probs_fname = 'latent_probs.npy'
sample_tree_fname = 'sample_tree.pkl'
per_feature_tree_fname = 'feat_{}_tree.pkl'
joint_tree_fname = 'joint_tree.pkl'
feature_attribution_fname = 'feat_attr_{}.pkl'

def check_latent_files(root):
    for fname in [X_fname, Y_fname, probs_fname, latent_fname]:
        path = os.path.join(root, fname)
        if not os.path.exists(path):
            return False
    return True

def load(config_path, device):
    config = load_config(config_path)

    dataset_params = config[DATASET_KEY]
    if dataset_params['dataset'] == 'cub':
        dataset_params['dataset'] = 'cub_no_augment'
    training_params = config[TRAINING_KEY]
    model_params = config[MODEL_KEY]
    ckpt_root = config[CKPT_KEY]
    dataset_params['shuffle'] = False

    # training params
    N_EPOCHS = training_params['n_epochs']
    BATCH_SIZE = 10
    OPTIMIZER = training_params['optimizer']
    OPTIMIZER_PARAMS = training_params['optimizer_params']
    SCHEDULER = training_params['scheduler']
    if 'norm_reg' in training_params:
        NORM_REG = training_params['norm_reg']
    else:
        NORM_REG = 0.

    # dataset
    train_loader, valid_loader, test_loader = get_loaders(
        **dataset_params, batch_size=BATCH_SIZE)

    # create model
    if 'use_mean' not in model_params:
        model_params['use_mean'] = False
    if dataset_params['dataset'] in TEXT_DATASETS:
        model_params['feature_net_args']['vocabulary'] = train_loader.dataset.dataset.get_vocab()
        model_params['feature_net_args']['max_sentence_len'] = MAX_SENTENCE_LEN[dataset_params['dataset']]
    kme_net = build_kme_net(model_params, device=device)
    kme_net._feat_net.use_mean(model_params['use_mean'])
    # optimization
    optimizer = globals()[OPTIMIZER](kme_net.parameters(), **OPTIMIZER_PARAMS)
    scheduler = None
    if SCHEDULER is not None:
        SCHEDULER_PARAMS = training_params['scheduler_params']
        scheduler = getattr(optim.lr_scheduler, SCHEDULER)(
            optimizer, **SCHEDULER_PARAMS)
    # checkpoint and summary writer
    kme_net, optimizer, scheduler, epoch = load_checkpoint(
        ckpt_root, kme_net, optimizer, scheduler, device=device)
    print(count_parameters(kme_net))

    if torch.cuda.device_count() > 1:
        kme_net = DataParallel(kme_net)
    kme_net = kme_net.to(device)

    return kme_net, train_loader, test_loader, ckpt_root, dataset_params['dataset']

def compute_latent(model, loader, root, device):
    X = []
    Y = []
    probs = []
    latent_repr_pre_agg = []

    for x, y in tqdm.tqdm(loader):
        x = x.to(device)
        p = model(x)
        latent = model._feat_net.process_samples(x)

        X.append(x.detach().cpu().numpy())
        probs.append(p.detach().cpu().numpy())
        Y.append(y.detach().cpu().numpy())
        latent_repr_pre_agg.append(latent.detach().cpu().numpy())

    X = np.concatenate(X, axis=0)
    Y = np.concatenate(Y, axis=0)
    probs = np.concatenate(probs, axis=0)
    latent_repr_pre_agg = np.concatenate(latent_repr_pre_agg, axis=0)

    np.save(os.path.join(root, X_fname), X)
    np.save(os.path.join(root, Y_fname), Y)
    np.save(os.path.join(root, probs_fname), probs)
    np.save(os.path.join(root, latent_fname), latent_repr_pre_agg)
    return

def create_trees(root):
    fpath = os.path.join(root, latent_fname)
    if not os.path.exists(fpath):
        raise RuntimeError(
            "File '{}' not detected. Please create it using this program with flag '--compute_latent'".format(fpath))
    latent = np.load(fpath)

    # create per feature tree
    for i in tqdm.tqdm(range(latent.shape[1])):
        feat_repr = latent[:, i, :]
        tree = KDTree(feat_repr, leaf_size=40)
        with open(os.path.join(root, per_feature_tree_fname.format(i)), 'wb') as out_file:
            pickle.dump(tree, out_file)

    # create sample tree
    sample_repr = np.sum(latent, axis=1)
    tree = KDTree(sample_repr, leaf_size=40)
    with open(os.path.join(root, sample_tree_fname), 'wb') as out_file:
        pickle.dump(tree, out_file)

    # create joint tree
    joint_repr = np.reshape(latent, (-1, latent.shape[2]))
    tree = KDTree(joint_repr, leaf_size=40)
    with open(os.path.join(root, joint_tree_fname), 'wb') as out_file:
        pickle.dump(tree, out_file)

    return

def compute_native_attributions(model, loader, dataset, device):
    attributions = []
    if dataset in IMAGE_DATASETS:
        fold = torch.nn.Fold(
            output_size=IMG_SZ[dataset], kernel_size=model._feat_net._patch_sz, stride=model._feat_net._patch_sz)
    for x, _ in tqdm.tqdm(loader):
        x = x.to(device)
        latent_repr = model._feat_net.process_samples(x)
        curr_attr = torch.norm(latent_repr, dim=2, p=2).unsqueeze(dim=1)

        if dataset in IMAGE_DATASETS:
            curr_attr = curr_attr.repeat(
                1, N_CHANNELS[dataset]*(model._feat_net._patch_sz**2), 1)
            curr_attr = fold(curr_attr)

        curr_attr = curr_attr.transpose(1, 2).transpose(2, 3)
        attributions.append(curr_attr.detach().cpu().numpy())
    return np.concatenate(attributions, axis=0)

def compute_captum_attributions(model, loader, dataset, method, device):
    attributions = []

    attr_algo = globals()[method](model)

    for x, _ in tqdm.tqdm(loader):
        x = x.to(device)
        preds = torch.argmax(model(x), dim=1)

        if method == 'IntegratedGradients':
            curr_attr = attr_algo.attribute(x, target=preds, baselines=0.)
        else:
            curr_attr = attr_algo.attribute(x, target=preds)

        if dataset in IMAGE_DATASETS:
            curr_attr = curr_attr.transpose(1, 2).transpose(2, 3)

        attributions.append(curr_attr.detach().cpu().numpy())

    return np.concatenate(attributions, axis=0)

def compute_feature_attributions(model, loader, results_root, methods, dataset, device):
    for m in methods:
        if m == 'native':
            attributions = compute_native_attributions(
                model, loader, dataset, device)
        else:
            attributions = compute_captum_attributions(
                model, loader, dataset, m, device)
        np.save(os.path.join(results_root,
                feature_attribution_fname.format(m)), attributions)
    return

def attribute_image(model, loader, id, device, root):
    folder = os.path.join(root, 'attributions')
    os.makedirs(folder, exist_ok=True)
    path = os.path.join(folder, 'flan_{}.png'.format(id))
    img = loader.dataset[id][0].to(device).unsqueeze(0)
    resize = Resize(loader.dataset.get_size(id))
    unnormalize = REVERSE_TRANSFORM['cub']

    fold = torch.nn.Fold(
        output_size=IMG_SZ[dataset], kernel_size=model._feat_net._patch_sz, stride=model._feat_net._patch_sz)
    latent_repr = model._feat_net.process_samples(img)

    curr_attr = torch.norm(latent_repr, dim=2, p=2).unsqueeze(dim=1)
    curr_attr = curr_attr.repeat(
        1, N_CHANNELS[dataset]*(model._feat_net._patch_sz**2), 1)
    curr_attr = fold(curr_attr)
    curr_attr = resize(unnormalize(curr_attr)).transpose(
        1, 2).transpose(2, 3)[0].detach().cpu().numpy()

    img = resize(unnormalize(img)).transpose(
        1, 2).transpose(2, 3)[0].detach().cpu().numpy()

    fig, ax = visualize_image_attr(
        curr_attr, img, cmap='Reds', use_pyplot=False, method='blended_heat_map')
    fig.savefig(path, dpi=300)
    print('Saved to {}'.format(path))

def attribute_image_comparison(model, loader, id, device, root):
    folder = os.path.join(root, 'attributions')
    os.makedirs(folder, exist_ok=True)
    path = os.path.join(folder, 'comparison_{}.png'.format(id))
    img = loader.dataset[id][0].to(device).unsqueeze(0)
    size = loader.dataset.get_size(id)
    resize = Resize(size)
    unnormalize = REVERSE_TRANSFORM['cub']
    target = torch.argmax(model(img), dim=1)

    fold = torch.nn.Fold(
        output_size=IMG_SZ[dataset], kernel_size=model._feat_net._patch_sz, stride=model._feat_net._patch_sz)
    latent_repr = model._feat_net.process_samples(img)

    # Compute native attributions
    attr_native = torch.norm(latent_repr, dim=2, p=2).unsqueeze(dim=1)
    attr_native = attr_native.repeat(
        1, N_CHANNELS[dataset]*(model._feat_net._patch_sz**2), 1)
    attr_native = fold(attr_native)
    attr_native = resize(attr_native).transpose(
        1, 2).transpose(2, 3)[0].detach().cpu().numpy()

    # Compute IG attributions
    algo = IntegratedGradients(model)
    attr_ig = algo.attribute(img, target=target)
    attr_ig = resize(attr_ig).transpose(
        1, 2).transpose(2, 3)[0].detach().cpu().numpy()

    # Compute Saliency
    algo = Saliency(model)
    attr_saliency = algo.attribute(img, target=target)
    attr_saliency = resize(attr_saliency).transpose(
        1, 2).transpose(2, 3)[0].detach().cpu().numpy()

    # Compute Saliency
    algo = InputXGradient(model)
    attr_ixg = algo.attribute(img, target=target)
    attr_ixg = resize(attr_ixg).transpose(
        1, 2).transpose(2, 3)[0].detach().cpu().numpy()

    img_numpy = resize(unnormalize(img)).transpose(
        1, 2).transpose(2, 3)[0].detach().cpu().numpy()

    fig, ax = plt.subplots(1, 4)

    visualize_image_attr(attr_native, img_numpy, cmap='Reds', use_pyplot=False,
                         method='blended_heat_map', plt_fig_axis=(fig, ax[0]))
    ax[0].set_title('Latent Norms')
    ax[0].title.set_size(8)
    visualize_image_attr(attr_ig, img_numpy, cmap='Reds', use_pyplot=False,
                         method='blended_heat_map', plt_fig_axis=(fig, ax[1]))
    ax[1].set_title('Integrated Gradients')
    ax[1].title.set_size(8)
    visualize_image_attr(attr_saliency, img_numpy, cmap='Reds', use_pyplot=False,
                         method='blended_heat_map', plt_fig_axis=(fig, ax[2]))
    ax[2].set_title('Saliency')
    ax[2].title.set_size(8)
    visualize_image_attr(attr_ixg, img_numpy, cmap='Reds', use_pyplot=False,
                         method='blended_heat_map', plt_fig_axis=(fig, ax[3]))
    ax[3].set_title('InputXGradient')
    ax[3].title.set_size(8)

    fig.savefig(path, dpi=300, bbox_inches='tight')
    print('Saved to {}'.format(path))

def analyze_top_patches(model, loader, id, device, root, top_k, top_k_vis=3):
    folder = os.path.join(root, 'masking')
    os.makedirs(folder, exist_ok=True)
    path = os.path.join(folder, 'top{}_{}.png'.format(top_k, id))
    img = loader.dataset[id][0].to(device).unsqueeze(0)
    true_label = loader.dataset[id][1]
    size = loader.dataset.get_size(id)
    resize = Resize(size)
    unnormalize = REVERSE_TRANSFORM['cub']

    probs = torch.softmax(model(img), dim=1)
    target = torch.argmax(probs, dim=1).detach(
    ).cpu().numpy().astype(np.int)[0]
    probs = probs.detach().cpu().numpy()[0]

    fold = torch.nn.Fold(
        output_size=IMG_SZ[dataset], kernel_size=model._feat_net._patch_sz, stride=model._feat_net._patch_sz)
    unfold = model._feat_net._unfold
    img_orig_color = unnormalize(img)
    patches = unfold(img_orig_color).transpose(
        1, 2).reshape(-1, 3, model._feat_net._patch_sz, model._feat_net._patch_sz)
    latent_repr = model._feat_net.process_samples(img)

    # Compute native attributions and retrieve top_k
    attr_native = torch.norm(latent_repr, dim=2, p=2).squeeze()
    top_k_norms = torch.topk(attr_native, k=top_k)[1]

    fig, ax = plt.subplots(1, top_k_vis + 1)

    for i in range(top_k_vis):
        idx = top_k_norms[i]
        patch_vis = patches[idx:idx+1]
        patch_vis = resize(patch_vis).transpose(
            1, 2).transpose(2, 3)[0].detach().cpu().numpy()

        patch_pred = torch.softmax(model._classifier(
            latent_repr[:, idx, :]), dim=1).squeeze()
        top_3_patch_pred = torch.topk(patch_pred, k=3)
        patch_pred_idx = top_3_patch_pred[1].detach(
        ).cpu().numpy().astype(np.int)
        patch_pred_probs = top_3_patch_pred[0].detach().cpu().numpy()

        curr_title = '{}: {:.2f}\n{}: {:.2f}\n{}: {:.2f}'.format(
            loader.dataset.get_label(patch_pred_idx[0]), patch_pred_probs[0],
            loader.dataset.get_label(patch_pred_idx[1]), patch_pred_probs[1],
            loader.dataset.get_label(patch_pred_idx[2]), patch_pred_probs[2],
        )

        ax[i+1].imshow(patch_vis)
        ax[i+1].xaxis.set_ticks_position("none")
        ax[i+1].yaxis.set_ticks_position("none")
        ax[i+1].set_yticklabels([])
        ax[i+1].set_xticklabels([])
        ax[i+1].grid(b=False)
        ax[i+1].set_title(curr_title)
        ax[i+1].title.set_size(4)

    masked_img = torch.zeros(patches.shape, device=device)
    masked_img[top_k_norms] = patches[top_k_norms]
    masked_img = fold(masked_img.reshape(-1, model._feat_net._n_patches,
                      3*model._feat_net._patch_sz*model._feat_net._patch_sz).transpose(1, 2))
    masked_img = resize(masked_img).transpose(
        1, 2).transpose(2, 3)[0].detach().cpu().numpy()
    partial_repr = torch.sum(latent_repr[:, top_k_norms, :], dim=1)
    partial_probs = torch.softmax(model._classifier(partial_repr), dim=1)
    partial_target = torch.argmax(
        partial_probs, dim=1).detach().cpu().numpy().astype(np.int)[0]
    partial_probs = partial_probs.detach().cpu().numpy()[0]
    probs_ord_idx = np.argsort(partial_probs)
    last_title = '(Full) {}: {:.2f}\n(Partial-1) {}: {:.2f}\n(Partial-2) {}: {:.2f}\n(Partial-3) {}: {:.2f}\n(True) {}'.format(
        loader.dataset.get_label(target), probs[target],
        loader.dataset.get_label(
            partial_target), partial_probs[partial_target],
        loader.dataset.get_label(
            probs_ord_idx[-2]), partial_probs[probs_ord_idx[-2]],
        loader.dataset.get_label(
            probs_ord_idx[-3]), partial_probs[probs_ord_idx[-3]],
        loader.dataset.get_label(true_label),
    )
    ax[0].imshow(masked_img)
    ax[0].xaxis.set_ticks_position("none")
    ax[0].yaxis.set_ticks_position("none")
    ax[0].set_yticklabels([])
    ax[0].set_xticklabels([])
    ax[0].grid(b=False)
    ax[0].set_title(last_title)
    ax[0].title.set_size(4)

    fig.savefig(path, dpi=300, bbox_inches='tight')
    print('Saved to {}'.format(path))

def correspondence_images(model, loader, train_loader, id, device, root, top_k=10):
    if not check_latent_files(root):
        print('Pre-Computing latent representations')
        compute_latent(model, train_loader, root, device)

    folder = os.path.join(root, 'correspondences')
    os.makedirs(folder, exist_ok=True)

    # useful operations
    fold = torch.nn.Fold(
        output_size=IMG_SZ[dataset], kernel_size=model._feat_net._patch_sz, stride=model._feat_net._patch_sz)
    unfold = model._feat_net._unfold

    # get the image of interest
    img = loader.dataset[id][0].to(device).unsqueeze(0)
    true_label = loader.dataset[id][1]

    # compute its representation
    patches_rep = model._feat_net.process_samples(img).detach().cpu().numpy()
    sample_repr = np.sum(patches_rep, axis=1).squeeze()

    #
    size = loader.dataset.get_size(id)
    resize = Resize(size)
    unnormalize = REVERSE_TRANSFORM['cub']

    fig, ax = plt.subplots(1, 3)

    def plot_img_at_ax(_img, ax_id, _hl, title='', **kwargs):
        img_numpy = resize(unnormalize(_img))
        img_numpy = img_numpy.transpose(1, 2).transpose(2, 3)[
            0].detach().cpu().numpy()
        ax[ax_id].imshow(img_numpy, alpha=0.5)
        ax[ax_id].imshow(_hl, alpha=0.5, cmap='Reds')
        ax[ax_id].xaxis.set_ticks_position("none")
        ax[ax_id].yaxis.set_ticks_position("none")
        ax[ax_id].set_yticklabels([])
        ax[ax_id].set_xticklabels([])
        ax[ax_id].grid(b=False)
        ax[ax_id].set_title(title)
        ax[ax_id].title.set_size(4)

    # get training latent representation
    train_patches_rep = np.load(os.path.join(root, latent_fname))
    train_samples_rep = np.sum(train_patches_rep, axis=1)

    # retrieve closest samples and farthest sample
    distances = np.linalg.norm(train_samples_rep - sample_repr, axis=1)
    sorted_ids = np.argsort(distances)
    closest_idx = sorted_ids[0]
    farthest_idx = sorted_ids[-1]

    closest_img = train_loader.dataset[closest_idx][0].to(device).unsqueeze(0)
    farthest_img = train_loader.dataset[farthest_idx][0].to(
        device).unsqueeze(0)

    # compute top predictions
    def compute_top_3_probs(_model, _img):
        _probs = torch.softmax(_model(_img), dim=1)
        _top_3_patch_pred = torch.topk(_probs, k=3)
        _pred_idx = _top_3_patch_pred[1].squeeze(
        ).detach().cpu().numpy().astype(np.int)
        _pred_probs = _top_3_patch_pred[0].squeeze().detach().cpu().numpy()

        text = '{}: {:.2f}\n{}: {:.2f}\n{}: {:.2f}'.format(
            loader.dataset.get_label(_pred_idx[0]), _pred_probs[0],
            loader.dataset.get_label(_pred_idx[1]), _pred_probs[1],
            loader.dataset.get_label(_pred_idx[2]), _pred_probs[2],
        )
        return text

    def compute_correspondences(_model, _ref_img, _other_img, minimize=True):
        # get patches representations
        _ref_patches_repr = _model._feat_net.process_samples(_ref_img)
        attr_native = torch.norm(_ref_patches_repr, dim=2, p=2).squeeze()
        top_k_norms = torch.topk(attr_native, k=top_k)[
            1].squeeze().detach().cpu().numpy().astype(np.int)
        _ref_patches_repr = _ref_patches_repr.detach().cpu().numpy()
        _other_patches_repr = _model._feat_net.process_samples(
            _other_img).detach().cpu().numpy()

        correspondence = match_single_patches(
            _ref_patches_repr[0][top_k_norms], _other_patches_repr[0], minimize)

        # get colormap
        m = get_cmap('tab10')
        _ref_alpha = np.zeros((_ref_patches_repr.shape[1], 1, 1, 1))
        _other_alpha = np.zeros((_ref_patches_repr.shape[1], 1, 1, 1))

        _ref_hl = np.zeros((_ref_patches_repr.shape[1], 3, 1, 1))
        _other_hl = np.zeros((_ref_patches_repr.shape[1], 3, 1, 1))

        def fold_patch(_patches):
            C = _patches.shape[1]
            P_SZ = _patches.shape[2]

            # _patches = N_P x C x P_SZ x P_SZ
            _patches = torch.FloatTensor(_patches)
            _patches = _patches.reshape(1, -1, C*P_SZ*P_SZ).transpose(1, 2)
            folded_img = resize(fold(_patches)[0]).transpose(
                0, 1).transpose(1, 2).detach().cpu().numpy()
            return folded_img

        color = np.array([1.0, 0., 0.])
        for k, i in enumerate(top_k_norms):
            _other_idx = correspondence[1][k]

            _ref_alpha[i, 0, 0, 0] = 1.
            _other_alpha[_other_idx, 0, 0, 0] = 1.

            _ref_hl[i, :, 0, 0] = color
            _other_hl[_other_idx, :, 0, 0] = color

        _ref_alpha = np.repeat(_ref_alpha, _model._feat_net._patch_sz, axis=2).repeat(
            _model._feat_net._patch_sz, axis=3)
        _other_alpha = np.repeat(_other_alpha, _model._feat_net._patch_sz, axis=2).repeat(
            _model._feat_net._patch_sz, axis=3)
        _ref_hl = np.repeat(_ref_hl, _model._feat_net._patch_sz, axis=2).repeat(
            _model._feat_net._patch_sz, axis=3)
        _other_hl = np.repeat(_other_hl, _model._feat_net._patch_sz, axis=2).repeat(
            _model._feat_net._patch_sz, axis=3)

        _ref_alpha = fold_patch(_ref_alpha)
        _other_alpha = fold_patch(_other_alpha)
        _ref_hl = fold_patch(_ref_hl)
        _other_hl = fold_patch(_other_hl)

        _ref_hl = np.concatenate([_ref_hl, _ref_alpha], axis=2)
        _other_hl = np.concatenate([_other_hl, _other_alpha], axis=2)

        return _ref_hl, _other_hl

    _, closest_hl = compute_correspondences(
        model, img, closest_img, minimize=True)
    ref_hl, farthest_hl = compute_correspondences(
        model, img, farthest_img, minimize=False)
    plot_img_at_ax(img, 0, ref_hl, compute_top_3_probs(model, img))
    plot_img_at_ax(closest_img, 1, closest_hl,
                   compute_top_3_probs(model, closest_img))
    plot_img_at_ax(farthest_img, 2, farthest_hl,
                   compute_top_3_probs(model, farthest_img))

    path = os.path.join(folder, '{}_{}_{}.png'.format(
        id, closest_idx, farthest_idx))
    fig.savefig(path, dpi=300, bbox_inches='tight')
    print('Saved to {}'.format(path))

def compute_latent_norms(root):
    latent_reprs = np.load(os.path.join(root, latent_fname))
    latent_norms = np.linalg.norm(latent_reprs, axis=2)
    np.save(os.path.join(root, latent_norms_fname), latent_norms)
    return latent_norms

def compute_latent_probs(model, root, device):
    latent_reprs = np.load(os.path.join(root, latent_fname))
    n_patches = latent_reprs.shape[1]
    latent_probs = []

    latent_reprs = torch.FloatTensor(
        latent_reprs).reshape(-1, latent_reprs.shape[-1])

    for i in tqdm.tqdm(range(0, len(latent_reprs), 30)):
        curr_repr = latent_reprs[i:i+30].to(device)

        probs = torch.softmax(model._classifier(
            curr_repr), dim=1).detach().cpu().numpy()
        latent_probs.append(probs)

    latent_probs = np.concatenate(latent_probs, axis=0).reshape(
        (-1, n_patches, probs.shape[-1]))

    np.save(os.path.join(root, latent_probs_fname), latent_probs)
    return latent_probs

def create_counterfactual(model, loader, train_loader, id, device, root, top_k, target_class, consistent):
    if not check_latent_files(root):
        print('Pre-Computing latent representations')
        compute_latent(model, train_loader, root, device)

    latent_norms_path = os.path.join(root, latent_norms_fname)
    if not os.path.exists(latent_norms_path):
        latent_norms = compute_latent_norms(root)
    else:
        latent_norms = np.load(latent_norms_path)

    latent_probs_path = os.path.join(root, latent_probs_fname)
    if not os.path.exists(latent_probs_path):
        latent_probs = compute_latent_probs(model, root, device)
    else:
        latent_probs = np.load(latent_probs_path)

    folder = os.path.join(root, 'counterfactuals')
    os.makedirs(folder, exist_ok=True)

    # useful operations
    fold = torch.nn.Fold(
        output_size=IMG_SZ[dataset], kernel_size=model._feat_net._patch_sz, stride=model._feat_net._patch_sz)
    unfold = model._feat_net._unfold

    # get the image of interest
    img = loader.dataset[id][0].to(device).unsqueeze(0)
    true_label = loader.dataset[id][1]

    # compute its representation
    patches_rep = model._feat_net.process_samples(img).detach().cpu().numpy()
    sample_repr = np.sum(patches_rep, axis=1).squeeze()

    #
    size = loader.dataset.get_size(id)
    resize = Resize(size)
    unnormalize = REVERSE_TRANSFORM['cub']

    if size[0] > size[1]:
        fig, ax = plt.subplots(3, 1)
    else:
        fig, ax = plt.subplots(1, 3)

    def img2patches(_img):
        patches = unfold(_img).transpose(1, 2).reshape(-1, 3,
                                                       model._feat_net._patch_sz, model._feat_net._patch_sz)
        return patches

    def patches2img(_patches):
        C = _patches.shape[1]
        P_SZ = _patches.shape[2]

        # _patches = N_P x C x P_SZ x P_SZ
        _patches = _patches.reshape(1, -1, C*P_SZ*P_SZ).transpose(1, 2)
        folded_img = fold(_patches)
        return folded_img

    def compute_counterfactual(_model, _ref_img, _change_top):
        # get patches representations
        _ref_patches_repr = _model._feat_net.process_samples(_ref_img)
        attr_native = torch.norm(_ref_patches_repr, dim=2, p=2).squeeze()
        top_k_norms = torch.topk(attr_native, k=top_k, largest=_change_top)[
            1].squeeze().detach().cpu().numpy().astype(np.int)
        _ref_patches_repr = _ref_patches_repr.detach().cpu().numpy()

        def get_best_candidate(_idx):
            curr_other_norms = latent_norms[:, _idx:_idx+1]
            curr_other_probs = latent_probs[:, _idx:_idx+1, target_class]

            df = pd.DataFrame(np.concatenate(
                [curr_other_norms, curr_other_probs], axis=1), columns=['norms', 'probs'])
            df['rank'] = df.sort_values(by=['norms', 'probs'], ascending=False)[
                'probs'].index + 1
            return df.index[df['rank'] == 1].tolist()[0]

        _ref_patches = img2patches(img)

        if consistent:
            best_candidate_idx = get_best_candidate(top_k_norms[0])
            candidate_img = train_loader.dataset[best_candidate_idx][0].to(
                device).unsqueeze(0)
            candidate_patches = img2patches(candidate_img)
            for i in top_k_norms:
                _ref_patches[i] = candidate_patches[i]
        else:
            for i in top_k_norms:
                best_candidate_idx = get_best_candidate(i)
                candidate_img = train_loader.dataset[best_candidate_idx][0].to(
                    device).unsqueeze(0)
                candidate_patches = img2patches(candidate_img)
                _ref_patches[i] = candidate_patches[i]

        return patches2img(_ref_patches)

    def plot_img_at_ax(_img, ax_id, title='', **kwargs):
        img_numpy = resize(unnormalize(_img))
        img_numpy = img_numpy.transpose(1, 2).transpose(2, 3)[
            0].detach().cpu().numpy()
        ax[ax_id].imshow(img_numpy)
        ax[ax_id].xaxis.set_ticks_position("none")
        ax[ax_id].yaxis.set_ticks_position("none")
        ax[ax_id].set_yticklabels([])
        ax[ax_id].set_xticklabels([])
        ax[ax_id].grid(b=False)
        ax[ax_id].set_title(title)
        ax[ax_id].title.set_size(4)

    def compute_top_3_probs(_model, _img):
        _probs = torch.softmax(_model(_img), dim=1)
        _top_3_patch_pred = torch.topk(_probs, k=3)
        _pred_idx = _top_3_patch_pred[1].squeeze(
        ).detach().cpu().numpy().astype(np.int)
        _pred_probs = _top_3_patch_pred[0].squeeze().detach().cpu().numpy()

        text = '{}: {:.2f}\n{}: {:.2f}\n{}: {:.2f}'.format(
            loader.dataset.get_label(_pred_idx[0]), _pred_probs[0],
            loader.dataset.get_label(_pred_idx[1]), _pred_probs[1],
            loader.dataset.get_label(_pred_idx[2]), _pred_probs[2],
        )
        return text

    counter_top = compute_counterfactual(model, img, _change_top=True)
    counter_bottom = compute_counterfactual(model, img, _change_top=False)
    plot_img_at_ax(img, 0, compute_top_3_probs(model, img))
    plot_img_at_ax(counter_top, 1, compute_top_3_probs(model, counter_top))
    plot_img_at_ax(counter_bottom, 2,
                   compute_top_3_probs(model, counter_bottom))

    path = os.path.join(folder, '{}_{}_{}_{}.png'.format(
        id, target_class, loader.dataset.get_label(target_class), consistent))
    fig.savefig(path, dpi=300, bbox_inches='tight')
    print('Saved to {}'.format(path))

def transfer_explanation(model, loader, id, device, root, top_k=10):
    mnist_config_path = model._classifier._config_path
    mnist_model, mnist_train_loader, mnist_test_loader, mnist_results_root, mnist_dataset \
        = load(mnist_config_path, device)

    if not check_latent_files(mnist_results_root):
        print('Pre-Computing latent representations')
        compute_latent(mnist_model, mnist_train_loader,
                       mnist_results_root, device)

    folder = os.path.join(root, 'correspondences')
    os.makedirs(folder, exist_ok=True)

    # useful operations
    fold = torch.nn.Fold(
        output_size=IMG_SZ[dataset], kernel_size=model._feat_net._patch_sz, stride=model._feat_net._patch_sz)
    mnist_fold = torch.nn.Fold(
        output_size=IMG_SZ[mnist_dataset], kernel_size=mnist_model._feat_net._patch_sz, stride=mnist_model._feat_net._patch_sz)

    # get the image of interest
    img = loader.dataset[id][0].to(device).unsqueeze(0)

    # compute its representation
    patches_rep = model._feat_net.process_samples(img).detach().cpu().numpy()
    sample_repr = np.sum(patches_rep, axis=1).squeeze()

    #
    size = [32, 32]
    resize = Resize(size)
    unnormalize = REVERSE_TRANSFORM['svhn']
    mnist_unnormalize = REVERSE_TRANSFORM['mnist']

    # prepare plotting
    fig, ax = plt.subplots(1, 2)

    def plot_img_at_ax(_img, ax_id, _unnormalize, _hl, title='', **kwargs):
        img_numpy = resize(_img)
        img_numpy = img_numpy.transpose(1, 2).transpose(2, 3)[
            0].detach().cpu().numpy()
        if img_numpy.shape[2] == 1:
            ax[ax_id].imshow(img_numpy, alpha=0.8, cmap='gray')
        else:
            ax[ax_id].imshow(img_numpy, alpha=0.8)
        ax[ax_id].imshow(_hl, alpha=0.2, cmap='Reds')
        ax[ax_id].xaxis.set_ticks_position("none")
        ax[ax_id].yaxis.set_ticks_position("none")
        ax[ax_id].set_yticklabels([])
        ax[ax_id].set_xticklabels([])
        ax[ax_id].grid(b=False)
        ax[ax_id].set_title(title)
        ax[ax_id].title.set_size(4)

    # get training latent representation
    train_patches_rep = np.load(os.path.join(mnist_results_root, latent_fname))
    train_samples_rep = np.sum(train_patches_rep, axis=1)

    # retrieve closest samples and farthest sample
    distances = np.linalg.norm(train_samples_rep - sample_repr, axis=1)
    sorted_ids = np.argsort(distances)
    closest_idx = sorted_ids[0]

    closest_img = mnist_train_loader.dataset[closest_idx][0].to(
        device).unsqueeze(0)

    # compute top predictions
    def compute_top_3_probs(_model, _img):
        _probs = torch.softmax(_model(_img), dim=1)
        _top_3_patch_pred = torch.topk(_probs, k=3)
        _pred_idx = _top_3_patch_pred[1].squeeze(
        ).detach().cpu().numpy().astype(np.int)
        _pred_probs = _top_3_patch_pred[0].squeeze().detach().cpu().numpy()

        text = '{}: {:.2f}\n{}: {:.2f}\n{}: {:.2f}'.format(
            _pred_idx[0], _pred_probs[0],
            _pred_idx[1], _pred_probs[1],
            _pred_idx[2], _pred_probs[2],
        )
        return text

    def compute_correspondences(_model, _ref_img, _other_model, _other_img):
        # get patches representations
        _ref_patches_repr = _model._feat_net.process_samples(_ref_img)
        attr_native = torch.norm(_ref_patches_repr, dim=2, p=2).squeeze()
        top_k_norms = torch.topk(attr_native, k=top_k)[
            1].squeeze().detach().cpu().numpy().astype(np.int)
        _ref_patches_repr = _ref_patches_repr.detach().cpu().numpy()
        _other_patches_repr = _other_model._feat_net.process_samples(
            _other_img).detach().cpu().numpy()

        correspondence = match_single_patches(
            _ref_patches_repr[0][top_k_norms], _other_patches_repr[0])

        # get colormap
        _ref_alpha = np.zeros((_ref_patches_repr.shape[1], 1, 1, 1))
        _other_alpha = np.zeros((_other_patches_repr.shape[1], 1, 1, 1))

        _ref_hl = np.zeros((_ref_patches_repr.shape[1], 3, 1, 1))
        _other_hl = np.zeros((_other_patches_repr.shape[1], 3, 1, 1))

        def fold_patch(_patches, _fold):
            C = _patches.shape[1]
            P_SZ = _patches.shape[2]

            # _patches = N_P x C x P_SZ x P_SZ
            _patches = torch.FloatTensor(_patches)
            _patches = _patches.reshape(1, -1, C*P_SZ*P_SZ).transpose(1, 2)
            folded_img = resize(_fold(_patches)[0]).transpose(
                0, 1).transpose(1, 2).detach().cpu().numpy()
            return folded_img

        color = np.array([1.0, 0., 0.])
        for k, i in enumerate(top_k_norms):
            _other_idx = correspondence[1][k]

            _ref_alpha[i, 0, 0, 0] = 1.
            _other_alpha[_other_idx, 0, 0, 0] = 1.

            _ref_hl[i, :, 0, 0] = color
            _other_hl[_other_idx, :, 0, 0] = color

        _ref_alpha = np.repeat(_ref_alpha, _model._feat_net._patch_sz, axis=2).repeat(
            _model._feat_net._patch_sz, axis=3)
        _other_alpha = np.repeat(_other_alpha, _other_model._feat_net._patch_sz, axis=2).repeat(
            _other_model._feat_net._patch_sz, axis=3)
        _ref_hl = np.repeat(_ref_hl, _model._feat_net._patch_sz, axis=2).repeat(
            _model._feat_net._patch_sz, axis=3)
        _other_hl = np.repeat(_other_hl, _other_model._feat_net._patch_sz, axis=2).repeat(
            _other_model._feat_net._patch_sz, axis=3)

        _ref_alpha = fold_patch(_ref_alpha, fold)
        _other_alpha = fold_patch(_other_alpha, mnist_fold)
        _ref_hl = fold_patch(_ref_hl, fold)
        _other_hl = fold_patch(_other_hl, mnist_fold)

        _ref_hl = np.concatenate([_ref_hl, _ref_alpha], axis=2)
        _other_hl = np.concatenate([_other_hl, _other_alpha], axis=2)

        return _ref_hl, _other_hl

    ref_hl, closest_hl = compute_correspondences(
        model, img, mnist_model, closest_img)
    plot_img_at_ax(img, 0, unnormalize, ref_hl,
                   compute_top_3_probs(model, img))
    plot_img_at_ax(closest_img, 1, mnist_unnormalize, closest_hl,
                   compute_top_3_probs(mnist_model, closest_img))

    path = os.path.join(folder, '{}_{}.png'.format(id, closest_idx))
    fig.savefig(path, dpi=300, bbox_inches='tight')
    print('Saved to {}'.format(path))

if __name__ == '__main__':
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

    parser = argparse.ArgumentParser(
        prog='Tool to analyze KME-net'
    )
    parser.add_argument('--config_file', type=str,
                        help='Path to the config to load', required=True)
    parser.add_argument('--print_test_accuracy', action='store_true')
    parser.add_argument('--compute_latent', action='store_true',
                        help='Preprocess the samples to get the latent representations of both samples and features.')
    parser.add_argument('--create_nn_tree', action='store_true',
                        help='If samples have been preprocessed, compute trees.')
    parser.add_argument('--top_k', type=int, default=3,
                        help='How many patches to consider')

    feat_attr_parser_group = parser.add_argument_group('feat_attr')

    feat_attr_parser_group.add_argument('--methods', nargs='*', default=['native', 'IntegratedGradients', 'Saliency',
                                                                         'InputXGradient'],
                                        help='Methods to use to compute the feature attributions.')
    feat_attr_parser_group.add_argument('--compute_on_train', action='store_true',
                                        help='If true, use the training dataset to compute the attributions.')
    feat_attr_parser_group.add_argument('--compute_feature_attribution', action='store_true',
                                        help='Whether to perform feature attribution analysis.')
    feat_attr_parser_group.add_argument('--visualize_feature_attribution', action='store_true',
                                        help='Whether to visualize feature attributions, i.e. '
                                        'generate an image showing the attribution.')
    feat_attr_parser_group.add_argument('--perturbation_analysis', action='store_true',
                                        help='Test attributions.')

    cub_parser_group = parser.add_argument_group('cub dataset')
    cub_parser_group.add_argument('--cub', action='store_true',
                                  help='Perform CUB analysis')
    cub_parser_group.add_argument('--sample_id', type=int, default=-1,
                                  help='ID of the test sample to analyze')
    cub_parser_group.add_argument('--create_attribution_image', action='store_true',
                                  help='Takes a test sample and computes attributions')
    cub_parser_group.add_argument('--create_attribution_comparison', action='store_true',
                                  help='Takes a test sample and compares attributions')
    cub_parser_group.add_argument('--top_k_analysis', action='store_true',
                                  help='Takes a test sample and analyzes top k patches')
    cub_parser_group.add_argument('--target_class', type=int, default=-1,
                                  help='Target Class for counterfactual')
    cub_parser_group.add_argument('--top_k_vis', type=int, default=3,
                                  help='How many patches to visualize')
    cub_parser_group.add_argument('--create_correspondence', action='store_true',
                                  help='Takes a test sample and computes correspondence')
    cub_parser_group.add_argument('--consistent', action='store_true',
                                  help='If create a consistent counterfactual')
    cub_parser_group.add_argument('--create_counterfactual', action='store_true',
                                  help='Takes a test sample and computes counterfactual')

    svhn_parser_group = parser.add_argument_group('svhn dataset')
    svhn_parser_group.add_argument('--svhn', action='store_true',
                                   help='Perform SVHN analysis')
    svhn_parser_group.add_argument('--transfer_explanation', action='store_true',
                                   help='Produce transfer explanation with MNIST')

    args = parser.parse_args()
    model, train_loader, test_loader, results_root, dataset = load(
        args.config_file, device)
    model.eval()

    if args.print_test_accuracy:
        test_routine(-1, model, test_loader, device)

    if args.compute_latent:
        compute_latent(model, train_loader, results_root, device)

    if args.create_nn_tree:
        create_trees(results_root)

    if args.compute_feature_attribution:
        loader = test_loader
        if args.compute_on_train:
            loader = train_loader
        compute_feature_attributions(
            model, loader, results_root, args.methods, dataset, device)

    if args.cub:
        if args.sample_id == -1:
            args.sample_id = np.random.randint(0, len(test_loader.dataset))

        if args.create_attribution_image:
            attribute_image(model, test_loader, args.sample_id,
                            device, results_root)

        if args.create_attribution_comparison:
            attribute_image_comparison(
                model, test_loader, args.sample_id, device, results_root)

        if args.top_k_analysis:
            analyze_top_patches(model, test_loader, args.sample_id, device, results_root,
                                args.top_k, args.top_k_vis)

        if args.create_correspondence:
            correspondence_images(model, test_loader, train_loader,
                                  args.sample_id, device, results_root, top_k=args.top_k)

        if args.create_counterfactual:
            if args.target_class == -1:
                args.target_class = np.random.randint(
                    0, len(test_loader.dataset.label_map))
            create_counterfactual(model, test_loader, train_loader, args.sample_id, device, results_root, top_k=args.top_k,
                                  target_class=args.target_class, consistent=args.consistent)

    if args.svhn:
        if args.sample_id == -1:
            args.sample_id = np.random.randint(0, len(test_loader.dataset))
        if args.transfer_explanation:
            transfer_explanation(
                model, test_loader, args.sample_id, device, results_root, top_k=args.top_k)

ext: py

max_stars_repo_name: phineasng/FLAN

max_issues_repo_name: phineasng/FLAN

max_forks_repo_name: phineasng/FLAN

reason: import numpy


text:
# Import all packages
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import scipy.io
import datetime
import re

# Create dict d
# d = {'a': [1], 'b': [2]}
# print(d)

# # Convert dict 'data' into DataFrame
# df = pd.DataFrame(data=d)
# print(df)

# Create DataFrame using zip of lists
# names = ['One', 'Two', 'Three', 'Four']

# dates = [0, 1, 2, 3]

# ages = [20, 30, 40, 50]

# places = ['City', 'Country', 'Suburb', 'Desert']

# list_labels = ['Name', 'Date', 'Age', 'Location']

# list_cols = [names, dates, ages, places]

# zipped = list(zip(list_labels, list_cols))
# # print(zipped)

# data = dict(zipped)

# df = pd.DataFrame(data)
# print(df.columns)
# print(df.index)
# print(df)

# # Create DataFrame from list using key/value pairs
# cities = ['Manheim', 'Preston park', 'Biglerville', 'Indiana', 'Curwensville', 'Crown', 'Harveys lake', 'Mineral springs', 'Cassville', 'Hannastown', 'Saltsburg', 'Tunkhannock', 'Pittsburgh', 'Lemasters', 'Great bend']

# state = 'PA'

# data = {'state': state, 'city': cities}

# df = pd.DataFrame(data)
# print(df.head())

# Import dataset from url
from urllib.request import urlretrieve
# url = 'https://assets.datacamp.com/production/course_1639/datasets/weather_data_austin_2010.csv'
url = 'https://assets.datacamp.com/production/course_1639/datasets/austin_airport_departure_data_2015_july.csv'
# url = 'https://assets.datacamp.com/production/course_1639/datasets/NOAA_QCLCD_2011_hourly_13904.txt'
# urlretrieve(url, 'dob.csv')
# file = 'dob.csv'

# Inspect file
# df = pd.read_csv(url, sep=',', delimiter=',')
df = pd.read_csv(url, sep=',', skiprows=12, delimiter=',')
# df.index.rename('Date', inplace=True)
date_time = df['Date (MM/DD/YYYY)']
# print(date_time)
date_time = pd.to_datetime(date_time)
df.index = date_time.rename('Date')
print(df.head(2))
print(df.info())
# print(df.columns)

# df_temp = df.resample('D').mean()
# print(df_temp.head(2))
# df1 = df_temp.loc['2010-01-01':'2010-01-08'].plot(style='b.-', subplots=True)
# plt.show()

# # Inspect 'Scheduled Elapsed Time(Minutes)'
# print(df['Scheduled Elapsed Time(Minutes)'][0:5])
# df.plot(kind='hist', y='Scheduled Elapsed Time(Minutes)', bins=20)
# plt.show()

# Setup plots for/from DataFrame
# temp = df['Temperature']
# graph = plt.plot(temp)
# plt.show()

# Plot DataFrame
# plot_list = df.iloc[:, 0:2]
# plot_list.plot(subplots=True)
# plt.show()

reason: import numpy,import scipy


text:
import numpy as np
import MFIRAP.d00_utils.project as pr
DTYPE = pr.FLOATX

def read_txt_header(filepath: str):
    """
    Read Heimann HTPA .txt header.
    Parameters
    ----------
    filepath : str
    Returns
    -------
    str
        TPA file header
    """
    with open(filepath) as f:
        header = f.readline().rstrip()
    return header

def txt2np(filepath: str, array_size: int = 32):
    """
    Convert Heimann HTPA .txt to NumPy array shaped [frames, height, width].
    Parameters
    ----------
    filepath : str
    array_size : int, optional
    Returns
    -------
    np.array
        3D array of temperature distribution sequence, shaped [frames, height, width].
    list
        list of timestamps
    """
    with open(filepath) as f:
        # discard the first line
        _ = f.readline()
        # read line by line now
        line = "dummy line"
        frames = []
        timestamps = []
        while line:
            line = f.readline()
            if line:
                split = line.split(" ")
                frame = split[0: array_size ** 2]
                timestamp = split[-1]
                frame = np.array([int(T) for T in frame], dtype=DTYPE)
                frame = frame.reshape([array_size, array_size], order="F")
                frame *= 1e-2
                frames.append(frame)
                timestamps.append(float(timestamp))
        frames = np.array(frames)
        # the array needs rotating 90 CW
        frames = np.rot90(frames, k=-1, axes=(1, 2))
    return frames, timestamps

max_stars_repo_name: igor-morawski/MFIR-AP

max_issues_repo_name: igor-morawski/MFIR-AP

max_forks_repo_name: igor-morawski/MFIR-AP


text:
import random

import numpy as np
from nose import with_setup

from pypair.contingency import BinaryTable, CategoricalTable, ConfusionMatrix

def setup():
    """
    Setup.
    :return: None.
    """
    np.random.seed(37)
    random.seed(37)

def teardown():
    """
    Teardown.
    :return: None.
    """
    pass

@with_setup(setup, teardown)
def test_confusion_matrix_creation():
    """
    Tests creating ConfusionMatrix.

    :return: None.
    """
    a = [0, 0, 0, 0, 0, 1, 1, 1, 1, 1]
    b = [0, 0, 0, 1, 1, 0, 0, 1, 1, 1]

    table = ConfusionMatrix(a, b)
    for measure in ConfusionMatrix.measures():
        stats = table.get(measure)
        if isinstance(stats, tuple):
            print(f'{stats[0]:.8f}, {stats[1]:.8f}: {measure}')
        else:
            print(f'{stats:.8f}: {measure}')

@with_setup(setup, teardown)
def test_binary_table_creation():
    """
    Tests creating BinaryTable. The data is simulated from this `site <https://www.mathsisfun.com/data/chi-square-test.html>`_.

    :return: None.
    """
    get_data = lambda x, y, n: [(x, y) for _ in range(n)]
    data = get_data(1, 1, 207) + get_data(1, 0, 282) + get_data(0, 1, 231) + get_data(0, 0, 242)
    a = [a for a, _ in data]
    b = [b for _, b in data]

    table = BinaryTable(a, b)
    for measure in BinaryTable.measures():
        stats = table.get(measure)
        if isinstance(stats, tuple):
            print(f'{stats[0]:.8f}, {stats[1]:.8f}: {measure}')
        else:
            print(f'{stats:.8f}: {measure}')

@with_setup(setup, teardown)
def test_categorical_table_creation():
    """
    Tests creating CategoricalTable.

    :return: None.
    """
    a = [0, 0, 0, 0, 0, 1, 1, 1, 1, 1]
    b = [0, 0, 0, 1, 1, 0, 0, 1, 1, 1]

    table = CategoricalTable(a, b)
    chisq = table.get('chisq')
    phi = table.get('phi')
    print(chisq, phi)

max_stars_repo_name: oneoffcoder/py-pair

max_issues_repo_name: oneoffcoder/py-pair

max_forks_repo_name: oneoffcoder/py-pair