text:
#! /usr/bin/env python3

# Author: Srinivasa Rao Zinka (srinivas . zinka [at] gmail . com)
# Copyright (c) 2014 Srinivasa Rao Zinka
# License: New BSD License.

import numpy as np
from scipy import integrate, special
import matplotlib.pyplot as plt
from . import planar as pl
from . import Zolotarev as zl
# from mayavi import mlab
import warnings

# adjusting "matplotlib" label fonts
from matplotlib import rc

rc('text', usetex=True)

def ip_format_c(N, radius, A="uniform", starts_at_zero=True, plot_type="2D",
                color='b', linewidth=1, linestyle='-', alpha=1, show=True,
                stem=False, stemline='g--', stemmarker='ro', fgcolor=(1, 1, 1),
                bgcolor=(0.5, 0.5, 0.5), mayavi_app=False):
    r"""
    Function to generate the 'Arraytool' input format for circular ring arrays.

    :param N:          number of elements in the uniformly spaced circular ring array
    :param radius:     radius of the circular ring in wavelengths
    :param A:          a 'column matrix' specifying the excitation values of the
                       circular ring array; by default it will be uniform excitation
    :param plot_type:  can be '2D'/'3D' ... if False, nothing happens
    :param stem:       if True, the array excitation is plotted as 'stem plot'
    :param mayavi_app: if True, the 3D plot will be opened in the MayaVi application
    :param starts_at_zero:   'True' if array starts at beta=0

    All other parameters are nothing but the 'Matplotlib/Mayavi' parameters.
    These should be familiar to 'Matlab' or 'Matplotlib/Mayavi' users.

    :rtype:            array_ip, a Numpy array of size (Number of elements(A),4)
    """

    # Creating Arraytool input form 'array_ip' for the circular ring array
    if (A == "uniform"):
        A = np.ones((N, 1))
    if (starts_at_zero):
        position_beta = (np.linspace(1, N, num=N) - 1) * (2 * np.pi / N)
    else:
        position_beta = (np.linspace(1, N, num=N) - 0.5) * (2 * np.pi / N)
    position_beta = np.reshape(position_beta, (N, -1))
    position_x = radius * np.cos(position_beta)
    position_y = radius * np.sin(position_beta)
    position_z = np.zeros_like(position_x)
    array_ip = np.hstack((position_x, position_y, position_z, A))

    # Plotting 2D/3D plots
    if (plot_type):
        # checking whether 'A' has any imaginary values
        if ((abs(A.imag) > 1e-10).sum()):
            A_plt = abs(A)  # if A.imag are significant, then '|A|' will be plotted
            warnings.warn(
                'Since, the given excitation "A" has significant imaginary parts, stem plot for abs(A) is plotted')
        else:
            A_plt = A.real  # if A.imag are negligible, then 'A'  will be plotted
            warnings.warn(
                'Since, the given excitation "A" has very small imaginary parts, stem plot for "A.real" is plotted')
        if (plot_type == "2D"):  # plot 2D plot in Matplotlib
            plt.plot(position_beta, A_plt, color=color, linewidth=linewidth,
                     linestyle=linestyle, alpha=alpha)
            if (stem): plt.stem(position_beta, A_plt, linefmt=stemline, markerfmt=stemmarker)
            plt.axis('tight');
            plt.grid(True)
            plt.xlabel(r'$y$', fontsize=16);
            plt.ylabel(r'$\left|A_{n}\right|$', fontsize=16)
            if (show): plt.title(r'$\mathrm{Array}\ \mathrm{Excitation}$', fontsize=18); plt.show()
        else:
            if (mayavi_app):  # this option opens the 3D plot in MayaVi Application
                mlab.options.backend = 'envisage'
            mlab.figure(fgcolor=fgcolor, bgcolor=bgcolor)
            s1 = mlab.quiver3d(position_x, position_y, position_z, position_z, position_z,
                               A_plt)  # stem3D representation
            ranges1 = [position_x.min(), position_x.max(), position_y.min(), position_y.max(), A_plt.min(), A_plt.max()]
            mlab.axes(xlabel="x", ylabel="y", zlabel="A", ranges=ranges1, nb_labels=3)
            mlab.colorbar(orientation="vertical", nb_labels=5)
            s1.scene.isometric_view()
            if (show): mlab.show()

    return array_ip

def FS(fun_str_re, fun_str_im='0', T0=2 * np.pi, m_start=-5, m_stop=5, err_lim=1e-8):
    """Function to generate a finite number of Fourier series coefficients of
    a periodic function."""

    N = m_stop - m_start + 1
    FS = np.zeros((N, 1), dtype='complex')
    m_index = list(range(m_start, m_stop + 1))
    w0 = 2 * np.pi / T0

    for m in m_index:
        fun_re = lambda x: (eval(fun_str_re)) * np.cos(m * w0 * x) + (eval(fun_str_im)) * np.sin(m * w0 * x)
        fun_img = lambda x: -(eval(fun_str_re)) * np.sin(m * w0 * x) + (eval(fun_str_im)) * np.cos(m * w0 * x)
        FS_re = integrate.quad(fun_re, 0, 2 * np.pi)
        FS_img = integrate.quad(fun_img, 0, 2 * np.pi)
        if ((FS_re[1] + FS_img[1]) < err_lim):
            FS[m - m_start] = (1 / T0) * (FS_re[0] + 1j * FS_img[0])
        else:
            print("Absolute error of the integration is not less than 1e-10 while calculating Fourier series")
            print("error(FS_re): ", FS_re[1])
            print("error(FS_img): ", FS_img[1])
        m_index = np.array(m_index) * (2 * np.pi / T0)
        m_index = np.reshape(m_index, (m_index.size, -1))

    return m_index, FS

def IFS(FS, T0=2 * np.pi, m_start=-4, m_stop=4, x_min=0, x_max=2 * np.pi, x_num=10):
    """Function to reconstruct (or check) the periodic function from the
    obtained Fourier coefficients"""

    m = np.arange(m_start, m_stop + 1)
    m = np.reshape(m, (-1, m.size))
    M = np.tile(m, (x_num, 1))
    x = np.linspace(x_min, x_max, num=x_num)
    x = np.reshape(x, (x.size, -1))
    X = np.tile(x, (1, m.size))
    FS = np.reshape(FS, (FS.size, -1))
    # Evaluating the inverse of the Fourier series
    IFS = np.dot(np.exp(1j * M * (2 * np.pi / T0) * X), FS)

    return x, IFS

def eval_Taylor(P, R, mbar, alpha_x, x):
    """My function description here."""

    if (P % 2 == 0):
        T0 = 2 * np.pi
    else:
        T0 = 1 * np.pi

    A = pl.dist(1, P + 1, R, dist_type_x='Taylor', mbar=mbar, alpha_x=alpha_x)
    x, result = IFS(A, T0, m_start=-P / 2, m_stop=P / 2, x_min=x, x_max=x, x_num=1)

    #    if(P%2==0):
    #        A = pl.dist(1, P + 1, R, dist_type_x='Taylor', mbar=mbar, alpha_x=alpha_x)
    #        x, result = IFS(A, T0=2*np.pi, m_start= -P / 2, m_stop=P / 2, x_min=x, x_max=x, x_num=1)
    #    else:
    #        A = pl.dist(1, 2*P + 1, R, dist_type_x='Taylor', mbar=mbar, alpha_x=alpha_x)
    #        x, result = IFS(A, T0=1*np.pi, m_start= -P, m_stop=P, x_min=x, x_max=x, x_num=1)

    result = result[0, 0].real

    return result

def eval_Bayliss(P, R, mbar, alpha_x, x):
    """My function description here."""

    if (P % 2 == 0):
        print("Order needs to be an ODD number for null patterns")
    else:
        T0 = 1 * np.pi

    A = pl.dist(1, P + 1, R, dist_type_x='Bayliss', mbar=mbar, alpha_x=alpha_x)
    x, result = IFS(A, T0, m_start=-P / 2, m_stop=P / 2, x_min=x, x_max=x, x_num=1)

    result = result[0, 0].imag

    return result

def FS_Taylor(N, R, mbar, alpha_x, x_min, x_max, x_num, plot_far=False, dB_limit=-40):
    """Function to evaluate Fourier series coefficients of Chebyshev far-field
       pattern"""

    R = str(R)
    mbar = str(mbar)
    alpha_x = str(alpha_x)

    if (N % 2 == 0):
        m_start = int(-N / 2)
        m_stop = int(N / 2)
        N = str(N)
        fun_str_re = 'eval_Taylor(' + N + ',' + R + ',' + mbar + ',' + alpha_x + ',' + 'x' + ')'
        print(fun_str_re)
        m_index, zm = FS(fun_str_re, m_start=m_start, m_stop=m_stop, err_lim=1e-5)
    else:
        m_start = -N
        m_stop = N
        N = str(N)
        fun_str_re = 'eval_Taylor(' + N + ',' + R + ',' + mbar + ',' + alpha_x + ',' + 'x' + ')'
        print(fun_str_re)
        m_index, zm = FS(fun_str_re, m_start=m_start, m_stop=m_stop, err_lim=1e-5)

    if (plot_far):
        x, AF = IFS(zm, 2 * np.pi, m_start, m_stop, x_min, x_max, x_num)
        AF = 20 * np.log10(abs(AF))
        AF = pl.cutoff(AF, dB_limit)
        plt.plot(x * (180 / np.pi), AF);
        plt.axis('tight');
        plt.grid(True)
        plt.title('Far-field Pattern')
        plt.xlabel(r'$\phi$')
        plt.ylabel('AF')
        plt.show()

    return m_index, zm

def FS_Bayliss(N, R, mbar, alpha_x, x_min, x_max, x_num, plot_far=False, dB_limit=-40):
    """Function to evaluate Fourier series coefficients of Chebyshev far-field
       pattern"""

    R = str(R)
    mbar = str(mbar)
    alpha_x = str(alpha_x)

    if (N % 2 == 0):
        print("Order needs to be an ODD number for null patterns")
    else:
        m_start = -N
        m_stop = N
        N = str(N)
        fun_str_re = 'eval_Bayliss(' + N + ',' + R + ',' + mbar + ',' + alpha_x + ',' + 'x' + ')'
        print(fun_str_re)
        m_index, zm = FS(fun_str_re, m_start=m_start, m_stop=m_stop, err_lim=1e-5)

    if (plot_far):
        x, AF = IFS(zm, 2 * np.pi, m_start, m_stop, x_min, x_max, x_num)
        AF = 20 * np.log10(abs(AF))
        AF = pl.cutoff(AF, dB_limit)
        plt.plot(x * (180 / np.pi), AF);
        plt.axis('tight');
        plt.grid(True)
        plt.title('Far-field Pattern')
        plt.xlabel(r'$\phi$')
        plt.ylabel('AF')
        plt.show()

    return m_index, zm

def FS_Chebyshev(N, R, x_min, x_max, x_num, plot_far=False, dB_limit=-40):
    """Function to evaluate Fourier series coefficients of Chebyshev far-field
       pattern"""

    c = np.cosh(np.arccosh(R) / (N))
    c = str(c)

    if (N % 2 == 0):
        m_start = int(-N / 2)
        m_stop = int(N / 2)
        N = str(N)
        fun_str_re = 'special.eval_chebyt(' + N + ',' + c + '*np.cos(x/2))'
        m_index, zm = FS(fun_str_re, m_start=m_start, m_stop=m_stop, err_lim=1e-5)
    else:
        m_start = -N  # make this (2*P+1) ... and take fourier for only half period
        m_stop = N
        N = str(N)
        fun_str_re = 'special.eval_chebyt(' + N + ',' + c + '*np.cos(x))'
        m_index, zm = FS(fun_str_re, m_start=m_start, m_stop=m_stop, err_lim=1e-5)

    if (plot_far):
        x, AF = IFS(zm, 2 * np.pi, m_start, m_stop, x_min, x_max, x_num)
        AF = 20 * np.log10(abs(AF))
        AF = pl.cutoff(AF, dB_limit)
        plt.plot(x * (180 / np.pi), AF);
        plt.axis('tight');
        plt.grid(True)
        plt.title('Far-field Pattern')
        plt.xlabel(r'$\phi$')
        plt.ylabel('AF')
        plt.show()

    return m_index, zm

def FS_Zolotarev(N, R, x_min, x_max, x_num, plot_far=False, dB_limit=-40):
    """Function to evaluate Fourier series coefficients of Chebyshev far-field
       pattern"""

    if (N % 2 == 0):
        print("Order needs to be an ODD number for null patterns")
    else:
        m_start = -N  # make this (2*P+1) ... and take fourier for only half period
        m_stop = N
        m = zl.z_m_frm_R(N, R, a=0.1, b=0.9999999999999)
        m = str(m)
        N = str(N)
        fun_str_re = 'zl.z_Zolotarev(' + N + ',' + 'np.sin(x)' + ',' + m + ')'
        m_index, zm = FS(fun_str_re, m_start=m_start, m_stop=m_stop, err_lim=1e-5)

    if (plot_far):
        x, AF = IFS(zm, 2 * np.pi, m_start, m_stop, x_min, x_max, x_num)
        AF = 20 * np.log10(abs(AF))
        AF = pl.cutoff(AF, dB_limit)
        plt.plot(x * (180 / np.pi), AF);
        plt.axis('tight');
        plt.grid(True)
        plt.title('Far-field Pattern')
        plt.xlabel(r'$\phi$')
        plt.ylabel('AF')
        plt.show()

    return m_index, zm

def dist_c_az(P, N, radius, R, mbar=5, alpha_x=0, starts_at_zero=True,
              dist_type=None, plot_far=False, plot_modes=False, scan=False):
    r"""
    This function gives array excitation coefficients corresponding to various
    'circular' array distribution types (for 'azhimutal' patterns) such as
    Chebyshev, Zolotarev, Taylor, Bayliss, etc.

    :param P:           order of the 'continuous' distribution (be careful! it
                        has nothing to do with the number of elements 'N')
    :param N:           number of elements in the uniformly spaced circular ring
                        array
    :param radius:      radius of the circular ring in wavelengths
    :param R:           side-lobe ratio in linear scale
    :param dist_type:   type of the distribution, e.g., 'Chebyshev', 'Zolotarev', etc
    :param bar:         transition index for dilation
    :param alpha:       Taylor's asymptotic tapering parameter

    :rtype:             A, a Numpy array of size (N, 1)
    """

    if (dist_type == 'Chebyshev'):
        m, zm = FS_Chebyshev(P, R, x_min=0, x_max=2 * np.pi, x_num=500,
                             plot_far=plot_far, dB_limit=-40)
    elif (dist_type == 'Taylor'):
        m, zm = FS_Taylor(P, R, mbar, alpha_x, x_min=0, x_max=2 * np.pi,
                          x_num=500, plot_far=plot_far, dB_limit=-40)
    elif (dist_type == 'Zolotarev'):
        m, zm = FS_Zolotarev(P, R, x_min=0, x_max=2 * np.pi, x_num=500,
                             plot_far=plot_far, dB_limit=-40)
    elif (dist_type == 'Bayliss'):
        m, zm = FS_Bayliss(P, R, mbar, alpha_x, x_min=0, x_max=2 * np.pi,
                           x_num=500, plot_far=plot_far, dB_limit=-40)

    cm = zm / ((1j ** m) * special.jn(m, 2 * np.pi * radius))

    # Plotting the normalized absolute values of zm and cm
    if (plot_modes):
        plt.plot(m, abs(cm) / abs(cm).max(), 'r--', label=r"$\mathrm{a_m}$")
        plt.plot(m, abs(zm) / abs(zm).max(), label=r"$\mathrm{f_m}$")
        plt.axis('tight');
        plt.grid(True)
        plt.title(r'$\mathrm{Far\ \&\  Near-field\ \ Modes \ (f_m\ \& \ a_m)}$')
        plt.xlabel(r'$\mathrm{m\ (Mode\ Number)}$')
        plt.ylabel(r'$\mathrm{Mode\ Amplitude\ (abs)}$')
        plt.legend(bbox_to_anchor=(1, 1), loc=1, borderaxespad=1)
        plt.show()

    # Finally, discretizing the continuous distribution
    if (starts_at_zero):
        beta = (np.linspace(1, N, num=N) - 1) * (2 * np.pi / N)
    else:
        beta = (np.linspace(1, N, num=N) - 0.5) * (2 * np.pi / N)
    beta = np.reshape(beta, (N, -1))
    beta_tile = np.tile(beta, cm.size)
    m_tile = np.tile(m.T, (N, 1))
    A = np.dot(np.exp(1j * m_tile * (beta_tile - scan)), cm)

    return A, cm

if __name__ == '__main__':
    # ==============================================================================
    # Circular Chebyshev related script
    # ==============================================================================

    SLR = 25  # side-lobe ratio in dB
    R = 10 ** (SLR / 20)  # converting SLR from dB scale to linear scale
    P = 21
    N = 150
    radius = 2 / np.pi

    A, cm = dist_c_az(P, N, radius, R, mbar=5, alpha_x=0, starts_at_zero=True,
                      dist_type='Chebyshev', plot_far=False, plot_modes=False, scan=False)
    array_ip = ip_format_c(N, radius, A, starts_at_zero=True, plot_type="3d", stem=False)

    phi, F = pl.pattern_t(array_ip, tht_scan=(0) * np.pi, phi_scan=(0) * np.pi, tht=0.5 * np.pi,
                          phi_min=0, phi_max=2 * np.pi, phi_num=500, scale="dB",
                          dB_limit=-200, factor="NF", plot_type='rect')

# ==============================================================================
# Circular Taylor related script
# ==============================================================================

#    SLR = 25 # side-lobe ratio in dB
#    R = 10 ** (SLR / 20) # converting SLR from dB scale to linear scale
#    P = 9
#    N = 150
#    radius = 2 / np.pi
#    mbar = 3
#    alpha_x =0
#
##    A = pl.dist(1, P + 1, R, dist_type_x='Taylor', mbar=mbar, alpha_x=alpha_x)
##    array_ip = pl.ip_format(1, 0, A, plot=True, stem=True, mayavi_app=False)
##    pl.pattern_u(array_ip, u_scan=0, u_min= 0, u_max=1, u_num=700, scale="dB",
##          dB_limit= -40, factor="AF", plot_type="rect", lattice=True)
#
##    print eval_Bayliss(P, R, mbar, alpha_x, x=0)
##    m_index, zm = FS_Bayliss(P, R, mbar, alpha_x, x_min=0, x_max=2*np.pi, x_num=500, plot_far=True, dB_limit= -40)
#
#    A, cm = dist_c_az(P, N, radius, R, mbar, alpha_x, starts_at_zero=True,
#              dist_type='Zolotarev',plot_far=True, plot_modes=False, scan=False)
#    array_ip = ip_format_c(N, radius, A, starts_at_zero=True, plot_type="3d", stem=False)
#
#    phi, F = pl.pattern_t(array_ip, tht_scan=(0) * np.pi, phi_scan=(0) * np.pi, tht=0.5 * np.pi,
#                 phi_min=0, phi_max=2 * np.pi, phi_num=500, scale="dB",
#                 dB_limit= -120, factor="NF", plot_type='polar')

# ==============================================================================
# Notes to myself
# ==============================================================================

# streamline Fourier coefficient evaluation ... i.e., FS_Chebyshev, etc

max_stars_repo_name: markuskreitzer/arraytool

reason: import numpy,from scipy


text:
import os
from collections import Counter, defaultdict, deque

import numpy as np
import matplotlib.pyplot as plt

sym_splitter = " "
ref_splitter = "_"
alphabet = "abcdefghijklmnopqrstuvwxyzåäö"
alphabet = alphabet + alphabet.upper()

enc_alphabet = "abcdefghijklmnopqrstuvwxyz"
enc_alphabet = enc_alphabet + enc_alphabet.upper()
enc_alphabet = "0123456789" + enc_alphabet
#enc_alphabet = "0123456789"
b = len(enc_alphabet)

root = "0"
symbol_set = set([])

class WZ_Model:
    def __init__(self):
        self.S = {}
        self.G = defaultdict(dict)
        self.s = []
        self.rewind = 0
        self.top = 0
        self.symbol_set = symbol_set
        self.prime_with_symbol_set()

    def prime_with_symbol_set(self):
        for a in self.symbol_set:
            self.update(root, a, sym_splitter)

    def get_s(self):
        return self.s[len(self.symbol_set):]

    def get(self, seq, a):
        return self.G[seq][a]

    '''
    '''
    def update(self, seq, a, splitter):
        self.top += 1
        top = alph(self.top)
        self.S[top] = (seq, a)
        self.G[seq][a] = top

        for _ in range(self.rewind):
            self.s.pop() # remove last elements
        self.rewind = 0

        s = list(self.climb_to_root(seq))
        #print("match:", "".join(s))
        if splitter == ref_splitter:
            self.rewind = self.get_rewind("".join(s))
        #print("rewind:", self.rewind)
        #print()
        self.s.extend(s)
        self.s.append(a)

        return seq, a, splitter

    def update_double(self, seq0, seq1, splitter):
        a = next(self.climb_to_root(seq1))

        self.top += 1
        top = alph(self.top)
        self.S[top] = (seq0, a)
        self.G[seq0][a] = top

        '''
        for _ in range(self.rewind):
            self.s.pop() # remove last elements
        self.rewind = 0

        #print("match:", "".join(s))
        if splitter == ref_splitter:
            self.rewind = self.get_rewind("".join(s))
        #print("rewind:", self.rewind)
        #print()
        '''
        s = self.climb_to_root(top)
        self.s.extend(s)

        #return seq, a, splitter

    def climb_to_root(self, seq):
        stack = []
        while seq != root:
            seq, a = self.S[seq]
            stack.append(a)
        return reversed(stack)

    def get_rewind(self, s):
        if all([ch in alphabet for ch in s]):
            return 0

        m = s.rstrip(alphabet)
        return len(s) - len(m) + 1

    def can_be_stripped(self, seq):
        s = "".join(self.climb_to_root(seq))
        rewind = self.get_rewind(s)
        return rewind

def encode(reader):
    global symbol_set

    seq = root
    seq_old = root
    s = reader.read()
    symbol_set = set(s)
    model = WZ_Model()
    N = len(s)
    if not N:
        return
    idx = 0
    prev_seq = model.get(seq, s[0])
    prev_splitter = sym_splitter
    idx += 1

    while idx < N:
        a = s[idx]
        seq_old = seq

        try:
            seq = model.get(seq, a)
        except KeyError:
            '''
            rewind = model.can_be_stripped(seq)
            if rewind > 0:
                splitter = ref_splitter
                idx -= rewind
            else:
                splitter = sym_splitter
            '''
            yield prev_seq, prev_splitter
            splitter = sym_splitter
            model.update_double(prev_seq, seq, prev_splitter)
            prev_seq = seq
            prev_splitter = splitter
            seq = root

        idx += 1

    if seq != root:
        yield prev_seq, prev_splitter
        model.update_double(prev_seq, seq, prev_splitter)

    print("end of encoding")
    print("G encode:", model.G)
    print("S encode:", model.S)
    print(model.s)

def decode(reader):
    model = WZ_Model()
    reader = iter(reader.read())
    prev_seq, prev_splitter = get_id(reader)

    while True:
        try:
            seq, splitter = get_id(reader)
        except StopIteration:
            break

        model.update_double(prev_seq, seq, prev_splitter)
        prev_seq = seq
        prev_splitter = splitter

    yield from model.get_s()
    print("G decode:", model.G)
    print("S encode:", model.S)
    print(model.s)

def get_id(reader):
    id_num = []
    while True:
        a = next(reader)
        #print("digit", a)
        if a not in enc_alphabet:
            return "".join(id_num), a
        else:
            id_num.append(a)

def alph(seq):
    """
    Convert an int in base-10 to
    base-whatever length the alphabet is
    Assumes 0 is the 0
    """
    if not seq:
        return root
    s = []
    quot = seq
    while quot:
        quot, rest = quot // b, quot % b
        s.append(enc_alphabet[rest])
    return "".join(reversed(s))

def test_alph():
    for seq in range(11):
        print("alph({}): {}".format(seq, alph(seq)))

default_dbg = False
def debug(*s, dbg=False):
    if default_dbg or dbg:
        print(*s)

def main():
    #fn = "/home/jdw/garageofcode/data/compression/nilsholg2.txt"
    #fn = "/home/jdw/garageofcode/data/compression/nilsholg.txt"
    #fn = "/home/jdw/garageofcode/data/compression/medium.txt"
    #fn = "short.txt"
    fn = "veryshort.txt"
    fn_compressed = fn.split(".")[0] + ".wzip"
    fn_reconstructed = fn.split(".")[0] + "_rec.txt"
    # encoding step
    with open(fn, "r") as r:
        with open(fn_compressed, "w") as f:
            for seq, spl in encode(r):
                #if a is None:
                #    f.write("{}{}".format(seq, ref_splitter))
                #else:
                f.write("{}{}".format(seq, spl))

    print("Before ", os.stat(fn).st_size)
    print("After  ", os.stat(fn_compressed).st_size)

    #  exit()

    with open(fn_compressed, "r") as r:
        with open(fn_reconstructed, "w") as f:
            for a in decode(r):
                #if a == "\\n":
                #    f.write("\n")
                #else:
                f.write(a)

    #exit()
    res = os.system("diff {} {}".format(fn, fn_reconstructed))
    if res:
        print("Reconstruction: failed")
    else:
        print("Reconstruction: OK")

if __name__ == '__main__':
    main()

    #test_alph()

ext: py

reason: import numpy


text:
"""Miscellaneous utilities."""
import numpy as np
import logging

class Normalizer():
    """
    An abstract class for normalizing inputs.
    """
    def normalize(self, X):
        raise NotImplementedError()

    def unnormalize(self, X):
        raise NotImplementedError()

class BoundedNormalizer(Normalizer):
    """
    A class for normalizing bounded inputs. Extends the abstract Normalizer
    class.

    :cvar ndarray lb: A matrix of size m-by-1 that contains lower bounds on the
        simulation inputs.
    :cvar ndarray ub: A matrix of size m-by-1 that contains upper bounds on the
        simulation inputs.

    **See Also**

    utils.misc.UnboundedNormalizer
    """
    lb, ub = None, None

    def __init__(self, lb, ub):
        """
        Initialize a BoundedNormalizer.

        :param ndarray lb: A matrix of size m-by-1 that contains lower bounds on
            the simulation inputs.
        :param ndarray ub: A matrix of size m-by-1 that contains upper bounds on
            the simulation inputs.
        """
        m = lb.size
        self.lb = lb.reshape((1, m))
        self.ub = ub.reshape((1, m))

    def normalize(self, X):
        """
        Return corresponding points shifted and scaled to [-1,1]^m.

        :param ndarray X: Contains all input points one wishes to normalize. The
            shape of `X` is M-by-m. The components of each row of `X` should be
            between `lb` and `ub`.

        :return: X_norm, contains the normalized inputs corresponding to `X`.
            The components of each row of `X_norm` should be between -1 and 1.
        :rtype: ndarray
        """
        X, M, m = process_inputs(X)
        X_norm = 2.0 * (X - self.lb) / (self.ub - self.lb) - 1.0
        return X_norm

    def unnormalize(self, X):
        """
        Return corresponding points shifted and scaled to [-1,1]^m.

        :param ndarray X: Contains all input points one wishes to unnormalize.
            The shape of `X` is M-by-m. The components of each row of `X` should
            be between -1 and 1.

        :return: X_unnorm, Contains the unnormalized inputs corresponding to
            `X`. The components of each row of `X_unnorm` should be between `lb`
            and `ub`.
        :rtype: ndarray
        """
        X, M, m = process_inputs(X)
        X_unnorm = (self.ub - self.lb) * (X + 1.0) / 2.0 + self.lb
        return X_unnorm

class UnboundedNormalizer(Normalizer):
    """
    A class for normalizing unbounded, Gaussian inputs to standard normals.
    Extends the abstract Normalizer class.

    :cvar ndarray mu: A matrix of size m-by-1 that contains the mean of the
        Gaussian simulation inputs.
    :cvar ndarray L: A matrix size m-by-m that contains the Cholesky factor of
        the covariance matrix of the Gaussian simulation inputs.

    **See Also**

    utils.misc.BoundedNormalizer

    **Notes**

    A simulation with unbounded inputs is assumed to have a Gaussian weight
    function associated with the inputs. The covariance of the Gaussian weight
    function should be full rank.
    """
    mu, L = None, None

    def __init__(self, mu, C):
        """
        Initialize an UnboundedNormalizer.

        :param ndarray mu: A matrix of size m-by-1 that contains the mean of the
            Gaussian simulation inputs.
        :param ndarray C: A matrix of size m-by-m that contains the covariance
            matrix of the Gaussian simulation inputs.
        """
        self.mu = mu.reshape((1, mu.size))
        self.L = np.linalg.cholesky(C)

    def normalize(self, X):
        """
        Return corresponding points transformed to a standard normal
        distribution.

        :param ndarray X: Contains all input points one wishes to normalize. The
            shape of `X` is M-by-m. The components of each row of `X` should be
            a draw from a Gaussian with mean `mu` and covariance `C`.

        :return: X_norm, Contains the normalized inputs corresponding to `X`.
            The components of each row of `X_norm` should be draws from a
            standard multivariate normal distribution.
        :rtype: ndarray
        """
        X, M, m = process_inputs(X)
        X0 = X - self.mu
        X_norm = np.linalg.solve(self.L,X0.T).T
        return X_norm

    def unnormalize(self, X):
        """
        Return corresponding points transformed to draws from a Gaussian
        distribution with mean `mu` and covariance `C`.

        :param ndarray X: Contains all input points one wishes to unnormalize.
            The shape of `X` is M-by-m. The components of each row of `X` should
            be draws from a standard multivariate normal.

        :return: X_unnorm, Contains the unnormalized inputs corresponding to
            `X`. The components of each row of `X_unnorm` should represent draws
            from a multivariate normal with mean `mu` and covariance `C`.
        :rtype: ndarray
        """
        X, M, m = process_inputs(X)
        X0 = np.dot(X,self.L.T)
        X_unnorm = X0 + self.mu
        return X_unnorm

def process_inputs(X):
    """
    Check a matrix of input values for the right shape.

    :param ndarray X: Contains input points. The shape of `X` should be M-by-m.

    :return: X, The same as the input.
    :rtype: ndarray
    :return: M, Number of rows in `X`.
    :rtype: int
    :return: m, Number of columns in `X`.
    :rtype: int
    """
    if len(X.shape) == 2:
        M, m = X.shape
    else:
        raise ValueError('The inputs X should be a two-d numpy array.')

    X = X.reshape((M, m))
    return X, M, m

def process_inputs_outputs(X, f):
    """
    Check a matrix of input values and a vector of outputs for the right shapes.

    :param ndarray X: Contains input points. The shape of `X` should be M-by-m.
    :param ndarray f: M-by-1 matrix.

    :return: X, The same as the input.
    :rtype: ndarray
    :return: f, The same as the input.
    :rtype: ndarray
    :return: M, Number of rows in `X`.
    :rtype: int
    :return: m, Number of columns in `X`.
    :rtype: int
    """
    X, M, m = process_inputs(X)

    if len(f.shape) == 2:
        Mf, mf = f.shape
    else:
        raise ValueError('The outputs f should be a two-d numpy array.')

    if Mf != M:
        raise Exception('Different number of inputs and outputs.')

    if mf != 1:
        raise Exception('Only scalar-valued functions.')

    f = f.reshape((M, 1))

    return X, f, M, m

def conditional_expectations(f, ind):
    """
    Compute conditional expectations and variances for a set of function values.

    :param ndarray f: An ndarry of function evaluations.
    :param ndarray[int] ind: Index array that tells which values of `f`
        correspond to the same value for the active variable.

    :return: Ef, An ndarray containing the conditional expectations.
    :rtype: ndarray
    :return: Vf, An ndarray containing the conditional variances.
    :rtype: ndarray

    **Notes**

    This function computes the mean and variance for all values in the ndarray
    `f` that have the same index in `ind`. The indices in `ind` correspond to
    values of the active variables.
    """

    n = int(np.amax(ind)) + 1
    NMC = np.sum(ind==0)
    logging.getLogger(__name__).debug('Computing {:d} conditional averages with {:d} MC samples.'.format(n, NMC))

    Ef, Vf = np.zeros((n, 1)), np.zeros((n, 1))
    for i in range(n):
        fi = f[ind == i]
        Ef[i] = np.mean(fi)
        Vf[i] = np.var(fi)
    return Ef, Vf

# thanks to Trent for these functions!!!
def atleast_2d_col(A):
    """
    Return the input `A` as a 2d column array.
    """
    return atleast_2d(A,'col')

def atleast_2d_row(A):
    """
    Return the input `A` as a 2d row array.
    """
    return atleast_2d(A,'row')

def atleast_2d(A, oned_as='row'):
    """
    Ensures the array `A` is at least two dimensions.

    :param ndarray A: matrix
    :param str oned_as: Should be either 'row' or 'col'. It determines whether
        the array `A` should be expanded as a 2d row or 2d column.
    """

    # not an array yet
    if not isinstance(A,(np.ndarray,np.matrixlib.defmatrix.matrix)):
        if not isinstance(A,(list,tuple)):
            A = [A]
        A = np.array(A)

    # check rank
    if np.ndim(A) < 2:
        # expand row or col
        if oned_as == 'row':
            A = A[None,:]
        elif oned_as == 'col':
            A = A[:,None]
        else:
            raise Exception , "oned_as must be 'row' or 'col' "

    return A


text:
#####################################################################
#                                                                   #
# labscript_devices/NewFocusPicoMotorController.py                  #
#                                                                   #
# Copyright 2016, Joint Quantum Institute                           #
#                                                                   #
# This file is part of labscript_devices, in the labscript suite    #
# (see http://labscriptsuite.org), and is licensed under the        #
# Simplified BSD License. See the license.txt file in the root of   #
# the project for the full license.                                 #
#                                                                   #
#####################################################################

from labscript_devices import labscript_device, BLACS_tab, BLACS_worker, runviewer_parser
from labscript import StaticAnalogQuantity, Device, LabscriptError, set_passed_properties
import numpy as np

class NewFocus8742Motor(StaticAnalogQuantity):
    #For Relative and Target Positions
    minval=-2147483648
    maxval=2147483647
    description = 'PicoMotor'

@labscript_device
class NewFocusPicoMotorController_8742(Device):
    allowed_children = [NewFocus8742Motor]
    generation = 0

    @set_passed_properties(
    )

    def __init__(self, name,
                 host = "", port=23, **kwargs):

        Device.__init__(self, name, None, None, None)
        self.BLACS_connection = '%s' %(host)

    def generate_code(self, hdf5_file):
        data_dict = {}
        for motor in self.child_devices:
            # Call these functions to finalise the motor, they are standard functions of all subclasses of Output:
            ignore = motor.get_change_times()
            motor.make_timeseries([])
            motor.expand_timeseries()
            connection = [int(s) for s in motor.connection.split() if s.isdigit()][0]
            value = motor.raw_output[0]
            if not motor.minval <= value <= motor.maxval:
                # error, out of bounds
                raise LabscriptError('%s %s has value out of bounds. Set value: %s Allowed range: %s to %s.'%(motor.description,motor.name,str(value),motor(motor.minval),str(motor.maxval)))
            if not connection > 0 and not connection < 5:
                # error, invalid connection number
                raise LabscriptError('%s %s has invalid connection number: %s'%(motor.description,motor.name,str(motor.connection)))
            data_dict[str(motor.connection)] = value
        dtypes = [(conn, int) for conn in data_dict]
        data_array = np.zeros(1, dtype=dtypes)
        for conn in data_dict:
            data_array[0][conn] = data_dict[conn]
        grp = hdf5_file.create_group('/devices/'+self.name)
        grp.create_dataset('static_values', data=data_array)

import time

from blacs.tab_base_classes import Worker, define_state
from blacs.tab_base_classes import MODE_MANUAL, MODE_TRANSITION_TO_BUFFERED, MODE_TRANSITION_TO_MANUAL, MODE_BUFFERED

from blacs.device_base_class import DeviceTab

@BLACS_tab
class NewFocusPicoMotorControllerTab(DeviceTab):
    def initialise_GUI(self):
        # Capabilities
        self.base_units = 'steps'
        self.base_min = -2147483648
        self.base_step = 10
        self.base_decimals = 0

        self.device = self.settings['connection_table'].find_by_name(self.device_name)
        self.num_motors = len(self.device.child_list)

        # Create the AO output objects
        ao_prop = {}
        for child_name in self.device.child_list:
            motor_type = self.device.child_list[child_name].device_class
            connection = self.device.child_list[child_name].parent_port
            if motor_type == "NewFocus8742Motor":
                base_max = 2147483647
            else:
                base_max = 2147483647

            ao_prop[connection] = {'base_unit':self.base_units,
                                   'min':self.base_min,
                                   'max':base_max,
                                   'step':self.base_step,
                                   'decimals':self.base_decimals
                                  }

        # Create the output objects
        self.create_analog_outputs(ao_prop)
        # Create widgets for output objects
        dds_widgets,ao_widgets,do_widgets = self.auto_create_widgets()
        # and auto place the widgets in the UI
        self.auto_place_widgets(("Target Position",ao_widgets))

        # Store the address
        self.host = str(self.settings['connection_table'].find_by_name(self.device_name).BLACS_connection)

        # Set the capabilities of this device
        self.supports_remote_value_check(False)
        self.supports_smart_programming(False)

    def initialise_workers(self):
        # Create and set the primary worker
        self.create_worker("main_worker",NewFocusPicoMotorControllerWorker,{'host':self.host})
        self.primary_worker = "main_worker"

@BLACS_worker
class NewFocusPicoMotorControllerWorker(Worker):
    def init(self):
        global socket; import socket
        global zprocess; import zprocess
        global h5py; import labscript_utils.h5_lock, h5py
        global time; import time

        self.port = 23

    def check_connectivity(self, host):
        try:
            response = self.initialise_sockets(self.host, self.port)
        except:
            raise Exception("sockets are not initialised")
        print response
        if '8742' in response:
            return True
        else:
            raise Exception('invalid response from host: ' + str(response))

    def initialise_sockets(self, host, port):
        s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        assert port, 'No port number supplied.'
        assert host, 'No hostname supplied.'
        s.settimeout(5)
        s.connect((host, int(port)))
        time.sleep(0.1)
        s.send("*IDN?"+'\n')
        time.sleep(0.1)
        response = s.recv(1024)
        time.sleep(0.005)
        s.close()
        return response

    def program_manual(self, front_panel_values):
        try:
            self.check_connectivity(self.host)
            # For each motor
            for axis in range(1,5):
                self.program_static(axis, int(front_panel_values["%s" %axis]))
                time.sleep(0.01)
            return {}
        except:
            raise Exception("server not responding, check connectivity failed")

    def program_static(self, axis, value):
        # Target Position ONLY
        s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        time.sleep(0.1)
        assert self.port, 'No port number supplied.'
        assert self.host, 'No hostname supplied.'
        s.settimeout(5)
        s.connect((self.host, int(self.port)))
        command = "xxPAnn"
        axis_command = command.replace("xx", str(axis))
        full_command = axis_command.replace("nn", str(value))
        print full_command
        s.send(full_command +'\n')
        time.sleep(0.1)
        print s.recv(1024)
        time.sleep(0.005)
        s.close()

    def transition_to_buffered(self,device_name,h5file,initial_values,fresh):
        return_data = {}
        with h5py.File(h5file) as hdf5_file:
            group = hdf5_file['/devices/'+device_name]
            if 'static_values' in group:
                data = group['static_values'][:][0]
        self.program_manual(data)

        for motor in data.dtype.names:
            return_data[motor] = data[motor]

        return return_data

    def transition_to_manual(self):
        #self.program_manual(self.initial_values)
        return True

    def abort_buffered(self):
        return True

    def abort_transition_to_buffered(self):
        return True

    def shutdown(self):
        self.s.close()


text:
"""!
@brief A dataset creation which is compatible with pytorch framework
and much faster in loading time depending on the new version of
loading only the appropriate files that might be needed

@author Efthymios Tzinis {etzinis2@illinois.edu}
@copyright University of illinois at Urbana Champaign
"""

import torch
import argparse
import os
import sys
import glob2
import numpy as np
from sklearn.externals import joblib
import scipy.io.wavfile as wavfile
from torch.utils.data import Dataset, DataLoader
from pprint import pprint

root_dir = os.path.join(
           os.path.dirname(os.path.realpath(__file__)),
           '../../../')
sys.path.insert(0, root_dir)
import spatial_two_mics.utils.audio_mixture_constructor as \
    mixture_creator
import spatial_two_mics.config as config
import spatial_two_mics.data_generator.dataset_storage as \
    dataset_storage

class PytorchMixtureDataset(Dataset):
    """
    This is a general compatible class for pytorch datasets.

    @note Each instance of the dataset should be stored using
    joblib.dump() and this is the way that it would be returned.
    After some transformations.

    The path of all datasets should be defined inside config.
    All datasets should be formatted with appropriate subfolders of
    train / test and val and under them there should be all the
    available files.
    """
    def __init__(self,
                 dataset='timit',
                 partition='train',
                 n_samples=[512, 128, 256],
                 n_sources=2,
                 genders=['f', 'm'],
                 n_fft=512,
                 win_len=512,
                 hop_length=128,
                 mixture_duration=2.0,
                 force_delays=[-1, 1],
                 get_top=None,
                 labels_mask='duet',
                 **kwargs):

        self.dataset_params = {
            'dataset': dataset,
            'n_samples': n_samples,
            'n_sources': n_sources,
            'genders': genders,
            'force_delays': force_delays
        }

        if labels_mask == 'duet' or labels_mask == 'ground_truth':
            self.selected_mask = labels_mask
        else:
            raise NotImplementedError("There is no available mask "
                                      "called: {}".format(labels_mask))
        self.partition = partition

        dataset_name = dataset_storage.create_dataset_name(
            self.dataset_params)

        self.dataset_dirpath = os.path.join(
            config.DATASETS_DIR,
            dataset_name,
            partition)

        self.dataset_stats_path = self.dataset_dirpath + '_stats'

        if not os.path.isdir(self.dataset_dirpath):
            raise IOError("Dataset folder {} not found!".format(
                self.dataset_dirpath))
        else:
            print("Loading files from {} ...".format(
                self.dataset_dirpath))

        self.mixture_folders = glob2.glob(os.path.join(
            self.dataset_dirpath, '*'))
        if get_top is not None:
            self.mixture_folders = self.mixture_folders[:get_top]

        self.n_samples = len(self.mixture_folders)

        # preprocess -- store all absolute spectra values for faster
        # loading during run time
        self.store_directly_abs_spectra()

    def __len__(self):
        return self.n_samples

    def __getitem__(self, idx):
        """!
        Depending on the selected partition it returns accordingly
        the following objects:

        if self.partition == 'train':
            (abs_tfs, selected_mask)
        else if partition == 'test' or 'val'
            (abs_tfs, selected_mask, wavs_list, real_tfs, imag_tfs)"""
        mix_folder = self.mixture_folders[idx]
        try:
            abs_tfs = joblib.load(os.path.join(mix_folder, 'abs_tfs'))
        except:
            raise IOError("Failed to load data from path: {} "
                          "for absolute spectra.".format(mix_folder))

        try:
            if self.selected_mask == 'duet':
                mask = joblib.load(os.path.join(mix_folder,
                                                'soft_labeled_mask'))
            else:
                mask = joblib.load(os.path.join(mix_folder,
                                                'ground_truth_mask'))
        except:
            raise IOError("Failed to load data from path: {} "
                          "for tf label masks".format(mix_folder))

        if self.partition == 'train':
            return abs_tfs, mask

        try:
            real_p = os.path.join(mix_folder, 'real_tfs')
            imag_p = os.path.join(mix_folder, 'imag_tfs')
            wavs_p= os.path.join(mix_folder, 'wavs')
            real_tfs = joblib.load(real_p)
            imag_tfs = joblib.load(imag_p)
            wavs_list = joblib.load(wavs_p)
            wavs_list = np.array(wavs_list)
        except:
            raise IOError("Failed to load data from path: {} "
                          "for real, imag tf of the mixture and "
                          "wavs".format(mix_folder))

        return abs_tfs, mask, wavs_list, real_tfs, imag_tfs

    def store_directly_abs_spectra(self):
        for mix_folder in self.mixture_folders:
            abs_p = os.path.join(mix_folder, 'abs_tfs')
            if os.path.lexists(abs_p):
                continue

            try:
                real_p = os.path.join(mix_folder, 'real_tfs')
                imag_p = os.path.join(mix_folder, 'imag_tfs')
                real_tfs = joblib.load(real_p)
                imag_tfs = joblib.load(imag_p)
            except:
                raise IOError("Failed to load data from path: {} "
                              "using joblib.".format(mix_folder))
            abs_tfs = np.abs(real_tfs + 1j * imag_tfs)
            try:
                joblib.dump(abs_tfs, abs_p, compress=0)
            except:
                raise IOError("Failed to save absolute value of "
                              "spectra in path: {}".format(abs_p))

    def extract_stats(self):
        if not os.path.lexists(self.dataset_stats_path):
            mean = 0.
            std = 0.
            for mix_folder in self.mixture_folders:
                try:
                    abs_p = os.path.join(mix_folder, 'abs_tfs')
                    abs_tfs = joblib.load(abs_p)
                except:
                    raise IOError("Failed to load absolute tf "
                                  "representation from path: {} "
                                  "using joblib.".format(abs_p))

                mean += np.mean(np.mean(abs_tfs))
                std += np.std(abs_tfs)
            mean /= self.__len__()
            std /= self.__len__()

            #     store them for later usage
            joblib.dump((mean, std), self.dataset_stats_path)
            print("Saving dataset mean and variance in: {}".format(
                self.dataset_stats_path))
        else:
            mean, std = joblib.load(self.dataset_stats_path)

        return mean, std

def get_data_generator(args,
                       return_stats=False,
                       get_top=None):
    data = PytorchMixtureDataset(**args.__dict__,
                                 get_top=get_top)
    generator_params = {'batch_size': args.batch_size,
                        'shuffle': True,
                        'num_workers': args.num_workers,
                        'drop_last': True}
    data_generator = DataLoader(data,
                                **generator_params,
                                pin_memory=False)
    n_batches = int(len(data) / args.batch_size)
    if return_stats:
        mean, std = data.extract_stats()
        return data_generator, mean, std, n_batches
    else:
        return data_generator, n_batches

def concatenate_for_masks(masks, n_sources, batch_size):
    # create 3d masks for each source
    batch_list = []
    for b in torch.arange(batch_size):
        sources_list = []
        for i in torch.arange(n_sources):
            source_mask = masks[b, :, :] == int(i)
            sources_list.append(source_mask)

        sources_tensor = torch.stack(sources_list,
                                     dim=n_sources)
        batch_list.append(sources_tensor)
    return torch.stack(batch_list, dim=0)

def initialize_and_copy_masks(masks, n_sources, batch_size, device):
    new_masks = torch.empty((batch_size,
                             masks.shape[1],
                             masks.shape[2],
                             n_sources),
                            dtype=torch.uint8)
    new_masks.to(device)
    for i in torch.arange(n_sources):
        new_masks[:, :, :, i] = masks[:, :, :] == int(i)

    return new_masks

def example_of_usage(args):
    import time

    training_data = PytorchMixtureDataset(**args.__dict__)
    mean, std = training_data.extract_stats()
    generator_params = {'batch_size': 128,
                        'shuffle': True,
                        'num_workers': 1,
                        'drop_last': True}
    training_generator = DataLoader(training_data, **generator_params)
    device = torch.device("cuda")

    timing_dic = {}
    n_sources = 2

    batch_now = time.time()
    # just iterate over the data
    for batch_data in training_generator:
        timing_dic['Loading batch'] = time.time() - batch_now
        batch_now = time.time()

        before = time.time()
        (abs_tfs, masks) = batch_data
        now = time.time()
        timing_dic['Loading from disk'] = now-before

        before = time.time()
        input_tf, masks_tf = abs_tfs.to(device), masks.to(device)
        now = time.time()
        timing_dic['Loading to GPU'] = now - before

        before = time.time()
        duet_stack = concatenate_for_masks(masks,
                                           n_sources,
                                           generator_params['batch_size'])
        now = time.time()
        timing_dic['Stacking in appropriate dimensions the masks'] = \
            now - before

        before = time.time()
        duet_copy = initialize_and_copy_masks(masks,
                                              n_sources,
                                              generator_params[
                                                  'batch_size'],
                                              device)
        now = time.time()
        timing_dic['Initializing and copying for masks'] = now - before

        pprint(timing_dic)
        batch_now = time.time()

def get_args():
    """! Command line parser """
    parser = argparse.ArgumentParser(description='Pytorch Fast Dataset '
                                                 'Loader')
    parser.add_argument("--dataset", type=str,
                        help="Dataset name", default="timit")
    parser.add_argument("--n_sources", type=int,
                        help="How many sources in each mix", default=2)
    parser.add_argument("--n_samples", type=int, nargs='+',
                        help="How many samples do u want to be "
                             "created for train test val",
                        required=True)
    parser.add_argument("--genders", type=str, nargs='+',
                        help="Genders that will correspond to the "
                             "genders in the mixtures",
                        default=['m', 'f'])
    parser.add_argument("-f", "--force_delays", nargs='+', type=int,
                        help="""Whether you want to force integer
                        delays of +- 1 in the sources e.g.""",
                        default=[-1,1])
    return parser.parse_args()

if __name__ == "__main__":
    args = get_args()
    example_of_usage(args)

reason: import numpy,import scipy


text:
import math

from PIL import Image
import numpy as np
import os
from skimage import transform, data
def resizePicture(path, window_size):
    """
    重新改变图片的大小
    :param path:源图片的路径
    :return:处理后的图片、图片的宽、图片的高
    """
    img = Image.open(path).convert('RGB')
    wid, hig = img.size[0], img.size[1]  # 宽、高
    img = np.asarray(img)
    if wid < window_size:
        img = transform.rescale(img, [math.ceil(window_size / wid), 1])
    if hig < window_size:
        img = transform.rescale(img, [1, math.ceil(window_size / wid)])

    return img, img.shape[0], img.shape[1]

def cutPicture(path, windows_size=300):
    """
    将图片按照窗口的大小进行切分
    :param path:
    :param windows_size:
    :return:
    """
    # 对图片的大小进行resize
    img, wide, high =resizePicture (path, windows_size)
    hig_num = int(wide/windows_size)
    wid_num = int(high/windows_size)

    """
    图片切割
    """
    for wid_ in range(wid_num):
        for len_ in range(hig_num):
            img_new = img[len_*windows_size:(len_+1)*windows_size, wid_*windows_size:(wid_+1)*windows_size]
            image = Image.fromarray(np.uint8(img_new))

            path = path.replace('C:\CR_\data\Train','C:\CR_\data\cut_picture')
            path_new = path.split('.')[0]+'_'+str(wid_)+'_'+str(len_)+'.png'
            image.save(path_new)

if __name__=="__main__":
    root = r'C:\CR_\data\Train'
    list_name = os.listdir(root)
    for num, file_name in enumerate(list_name):
        print('正在处理图片{}'.format(file_name))
        cutPicture(os.path.join(root, file_name))

    # transform.rescale(img, 0.1)
    # skimage.transform.resize(image, [300,300])

max_stars_repo_name: Mini-tact/CR

max_issues_repo_name: Mini-tact/CR

max_forks_repo_name: Mini-tact/CR


text:
# coding: utf-8
"""
Name: japanmap.py

Plot japan map using cartopy.

Usage: This is the module.

Author: Ryosuke Tomita
Date: 2021/09/08
"""
import matplotlib.pyplot as plt
import matplotlib.ticker as mticker
import cartopy.crs as ccrs
import cartopy.feature as cfea
from cartopy.mpl.ticker import LatitudeFormatter, LongitudeFormatter
import numpy as np

class JpMap:
    """Make map near japan."""
    def __init__(self):
        self.fig = plt.figure(figsize=(36, 24), facecolor='w')
        self.ax = self.fig.add_subplot(1, 1, 1,
                                       projection=ccrs.PlateCarree(central_longitude=0.0))

        self.ax.set_global()
        self.ax.coastlines()

        # using axis lat,lon format
        dlon, dlat = 10, 10  # grid line interval.
        xticks = np.arange(-180, 180.1, dlon)
        yticks = np.arange(-90, 90.1, dlat)
        self.ax.set_xticks(xticks, crs=ccrs.PlateCarree())
        self.ax.set_yticks(yticks, crs=ccrs.PlateCarree())
        latfmt = LatitudeFormatter()  # axis = degree
        lonfmt = LongitudeFormatter(
            zero_direction_label=True)  # No NS mark 0 degree
        self.ax.xaxis.set_major_formatter(lonfmt)
        self.ax.yaxis.set_major_formatter(latfmt)
        self.ax.axes.tick_params(labelsize=12)

        # land color,ocean color
        self.ax.add_feature(cfea.LAND, color='#98fb98')
        self.ax.add_feature(cfea.OCEAN, color='#87cefa')

        # grid setting
        grid = self.ax.gridlines(crs=ccrs.PlateCarree(),
                                 draw_labels=False,
                                 linewidth=1,
                                 alpha=0.8,
                                 color='k',)
        grid.xlocator = mticker.FixedLocator(xticks)
        grid.ylocator = mticker.FixedLocator(yticks)

        # chose near japan from the world map
        self.ax.set_xlim(110, 180)
        self.ax.set_ylim(20, 60)

        self.color_list = ['#b22222', '#ffa500', '#ffd700', '#006400',
                      '#3cb371', '#00bfff', '#0000ff', '#9932cc', '#c71585']

        self.plt_cnt = 0

    def plot_data(self, x, y, label):
        """plot location data. plot type is dot."""
        self.ax.plot(y, x, 'bo',
                     label           = label,
                     markersize      = 15,
                     color           =self.color_list[self.plt_cnt],
                     markeredgewidth =1.6,
                     markeredgecolor ='#000000',
                     linestyle       ="solid",)
        self.plt_cnt += 1

    def plot_value(self, x, y, z):
        """plot z value on japan map. x, y are location data. """
        for i in range(len(x)):
            self.ax.plot(y[i], x[i], marker=str(f'${z[i]:.2f}$'), markersize=90, color='k')

    def contour_plot(self, x, y, z):
        """make contour map."""
        plt.rcParams["font.size"] = 30
        X, Y = np.meshgrid(x, y)
        contour = self.ax.contour(X, Y, z, colors="black")
        #contour = self.ax.contour(X, Y, z, 1000,colors="black",vmin=0.8, vmax=1.0)
        self.ax.clabel(contour)

    def shade_plot(self, x, y, z, label):
        """make color map."""
        X, Y = np.meshgrid(x, y)
        shade = self.ax.pcolormesh(X, Y, z, cmap = "rainbow")
        color_bar = plt.colorbar(shade, orientation="vertical")
        color_bar.set_label(label, fontsize=24)

    def vector_plot(self, x, y, u, v, vector_interval, vector_scale):
        """make vector map."""
        def skip_value(array, len_x, len_y, vector_interval):
            skiped_array = [array[i][j] for i in range(len_y) for j in range(len_x) if i%vector_interval == 0 and j % vector_interval == 0]
            return skiped_array

        u_skipped = skip_value(u, len(x), len(y), vector_interval)
        v_skipped = skip_value(v, len(x), len(y), vector_interval)
        x_2d, y_2d = np.meshgrid(x, y)
        x_2d_skipped = skip_value(x_2d, len(x), len(y), vector_interval)
        y_2d_skipped = skip_value(y_2d, len(x), len(y), vector_interval)

        self.ax.quiver(x_2d_skipped, y_2d_skipped, u_skipped, v_skipped, angles='uv',
                       scale_units='xy', scale=vector_scale, headlength=3.5)

    def save_fig(self, outname, title):
        """save figure."""
        self.ax.set_title(title, fontsize=24)
        bar, label = self.ax.get_legend_handles_labels()
        self.ax.legend(bar, label, loc='upper left',borderaxespad=0,
                       bbox_to_anchor=(1.05, 1), fontsize=30)
        self.fig.savefig(outname, bbox_inches="tight", pad_inches=0.5)
        plt.close()

max_stars_repo_name: RyosukeDTomita/gcmPlot

max_issues_repo_name: RyosukeDTomita/gcmPlot

max_forks_repo_name: RyosukeDTomita/gcmPlot


text:
#srun -c12 --mem-per-cpu=7G --pty -p gpu --gres=gpu:ampere:1 --pty bash
#sbatch -c12 --mem-per-cpu=7G -p gpu --wrap="python train_ssl_minimax2.py --run_id="" --kin_pretrained" --gres=gpu:ampere:1 --job-name="ours" --output=".out"

import torch
from torch.autograd import Variable
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.optim.lr_scheduler import *
import time, random, os, cv2, math, argparse, itertools, warnings, sys, traceback
import numpy as np

import parameters_SSL as params
import config as cfg
from keras.utils import to_categorical

from sklearn.metrics import precision_recall_fscore_support, average_precision_score
from tensorboardX import SummaryWriter
from torch.utils.data import DataLoader

from dl_vispr_ssl import vispr_ssl, collate_fn_train
from dl_vispr import vispr_dataset_generator
from dl_vispr import collate_fn_train as collate_fn_train_regular

from dl_ft_nov import baseline_dataloader_train_strong, multi_baseline_dataloader_val_strong, collate_fn_train_ucf101, collate_fn_val_ucf101
from model import build_r3d_classifier, load_r3d_classifier, ssl_pretrained_load, load_unet, build_privacy_classifier, load_privacy_classifier, ssl_privacy_load
from nt_xent_original import NTXentLoss
from torchvision.utils import save_image

warnings.filterwarnings("ignore", category=DeprecationWarning)

torch.backends.cudnn.benchmark = True

def train_epoch(run_id, epoch, data_loader_vispr, dataloader_ucf101, ft_model, fa_model, fb_model,  criterion_ft, criterion_fb, optimizer_fa, optimizer_fb, optimizer_ft, writer, use_cuda,learning_rate_fa, learning_rate_fb, learning_rate_ft):
    for param_group in optimizer_fa.param_groups:
        param_group['lr']=learning_rate_fa
    for param_group_fb in optimizer_fb.param_groups:
        param_group_fb['lr']=learning_rate_fb
    for param_group_ft in optimizer_ft.param_groups:
        param_group_ft['lr']=learning_rate_ft
    writer.add_scalar('Learning Rate', learning_rate_fa, epoch)
    writer.add_scalar('Learning Rate', learning_rate_fb, epoch)
    writer.add_scalar('Learning Rate', learning_rate_ft, epoch)

    print("Learning rate of fa is: {}".format(param_group['lr']))
    print("Learning rate of fb is: {}".format(param_group_fb['lr']))
    print("Learning rate of ft is: {}".format(param_group_ft['lr']))

    losses_fa, losses_fb, losses_ft = [], [], []

    step = 1
    for i, (data1, data2) in enumerate(zip(data_loader_vispr, dataloader_ucf101)):
        if use_cuda:
            inputs_vispr = [data1[ii].cuda() for ii in range(2)]
            labels_vispr = torch.from_numpy(np.asarray(data1[2])).float().cuda()
            inputs_ucf101 = data2[0].cuda()
            labels_ucf101 = torch.from_numpy(np.asarray(data2[1])).cuda()
        # print(labels_vispr)
        # print(labels_ucf101)
        #step-1 update fa
        optimizer_fa.zero_grad()
        optimizer_fb.zero_grad()
        optimizer_ft.zero_grad()

        if step ==1:
            fa_model.train()
            ft_model.eval()
            fb_model.eval()

            output1 = [fb_model(fa_model(inputs_vispr[ii])) for ii in range(2)]

            ori_bs, ori_t, ori_c, ori_h, ori_w = inputs_ucf101.shape
            inputs_ucf101 = inputs_ucf101.view(-1, inputs_ucf101.shape[2], inputs_ucf101.shape[3], inputs_ucf101.shape[4])

            output2 = ft_model(fa_model(inputs_ucf101).reshape(ori_bs, ori_t, ori_c, ori_h, ori_w).permute(0,2,1,3,4))

            con_loss_criterion = NTXentLoss(device = 'cuda', batch_size= output1[0].shape[0], temperature= 0.1, use_cosine_similarity= False)

            loss_fb = con_loss_criterion(output1[0], output1[1])
            loss_ft = criterion_ft(output2,labels_ucf101)

            loss_fa = -loss_fb + params.ft_loss_weight*loss_ft
            losses_fa.append(loss_fa.item())
            # print('loss computed finished')

            loss_fa.backward()
            # print('backward finished')

            optimizer_fa.step()
            step = 2
            # print('step-1 finished')
            if i % 100 == 0:
                print(f'Training Epoch {epoch}, Batch {i}, loss_fa: {np.mean(losses_fa) :.5f}', flush = True)

            continue

        if step == 2:
            fa_model.eval()
            ft_model.train()
            fb_model.train()

            output1 = [fb_model(fa_model(inputs_vispr[ii])) for ii in range(2)]

            ori_bs, ori_t, ori_c, ori_h, ori_w = inputs_ucf101.shape
            inputs_ucf101 = inputs_ucf101.view(-1, inputs_ucf101.shape[2], inputs_ucf101.shape[3], inputs_ucf101.shape[4])

            output2 = ft_model(fa_model(inputs_ucf101).reshape(ori_bs, ori_t, ori_c, ori_h, ori_w).permute(0,2,1,3,4))

            con_loss_criterion = NTXentLoss(device = 'cuda', batch_size= output1[0].shape[0], temperature= 0.1, use_cosine_similarity= False)

            loss_fb = con_loss_criterion(output1[0], output1[1])

            loss_ft = criterion_ft(output2,labels_ucf101)

            losses_ft.append(loss_ft.item())
            losses_fb.append(loss_fb.item())

            loss_fb.backward()
            optimizer_fb.step()

            loss_ft.backward()
            optimizer_ft.step()
            step = 1
            if (i-1) % 100 == 0:
                print(f'Training Epoch {epoch}, Batch {i}, loss_fb: {np.mean(losses_fb) :.5f}, loss_ft: {np.mean(losses_ft) :.5f}', flush = True)
            continue

    print('Training Epoch: %d, loss_fa: %.4f, loss_fb: %.4f, loss_ft: %.4f' % (epoch,  np.mean(losses_fa), np.mean(losses_fb), np.mean(losses_ft)))
    writer.add_scalar('Training loss_fa', np.mean(losses_fa), epoch)
    writer.add_scalar('Training loss_fb', np.mean(losses_fb), epoch)
    writer.add_scalar('Training loss_ft', np.mean(losses_ft), epoch)

    del loss_fb, loss_ft, loss_fa,  inputs_vispr, inputs_ucf101,  output1, output2

    return fa_model, fb_model, ft_model

def val_visualization_fa_vispr(save_dir, epoch, validation_dataloader, fa_model):
    for i, (inputs, label, vid_path) in enumerate(validation_dataloader):
    # inputs = inputs.permute(0,4,1,2,3)
        if len(inputs.shape) == 1:
            continue

            # inputs = inputs.permute(0, 2, 1, 3, 4)

        inputs = inputs.cuda()
        # label1 = torch.from_numpy(np.asarray(label))[0].item()
        # label = classes_dict_rev[label1]
        # print(label1, label)
        # label = label.replace(' ', '_')
        # ori_bs, ori_t, ori_c, ori_h, ori_w = inputs.shape
        # inputs = inputs.view(-1, inputs.shape[2], inputs.shape[3], inputs.shape[4])

        # print(inputs.shape)
        # exit()
        # save_image(torch.flip(inputs[::2,:],dims = [1]) , 'input_' + label + '.png', padding=1, nrow=inputs.shape[0])
        image_full_name = save_dir + 'combined_e'+ str(epoch) + '.png'
        # print(inputs.shape)
        with torch.no_grad():

            outputs = fa_model(inputs)
            vis_image = torch.cat([inputs, outputs], dim=0)
            save_image(vis_image, image_full_name, padding=5, nrow=int(inputs.shape[0]))

            # if i ==0:
            break

def val_epoch_vispr(run_id, epoch, validation_dataloader, fa_model, fb_model, criterion, writer, use_cuda):

    fa_model.eval()
    fb_model.eval()
    losses = []
    predictions, ground_truth = [], []
    vid_paths = []
    label_dict, pred_dict = {}, {}

    for i, (inputs, label, vid_path) in enumerate(validation_dataloader):
        # inputs = inputs.permute(0,4,1,2,3)
        if len(inputs.shape) != 1:

            # inputs = inputs.permute(0, 2, 1, 3, 4)

            if use_cuda:
                inputs = inputs.cuda()
                label = torch.from_numpy(np.asarray(label)).float().cuda()

            # print(inputs.shape)

            # print(f'fa_model output shape is {output1.shape}')
            # output1 = output1.reshape(ori_bs, ori_t, ori_c, ori_h, ori_w)
            # print(f'fa_model reshaped output shape is {output1.shape}')
            # exit()

            # random.sample(range(output1.shape[1]), 2)
            # print(output1[:,0,].shape, flush = True)

            # print(inputs.shape)
            with torch.no_grad():

                # output1 = fa_model(inputs)
                output1 = fa_model(inputs)
                output = fb_model(output1)

                # print(f'fa_model output shape is {output1.shape}')

                '''# print(label)
                if use_cuda:
                    inputs = inputs.cuda()
                    label = torch.from_numpy(np.asarray(label)).cuda()
                # print(label)

                with torch.no_grad():

                    output = model(inputs)
                    loss = criterion(output,label)'''
                loss = criterion(output,label)
                losses.append(loss.item())

            predictions.extend(output.cpu().data.numpy())
            vid_paths.extend(vid_path)
            ground_truth.extend(label.cpu().data.numpy())

            # print(len(predictions))

            if i+1 % 99 == 0:
                print("Validation Epoch ", epoch, " Batch ", i, "- Loss : ", np.mean(losses))

    del inputs, output, label, loss
    print("Validation Epoch ", epoch, "- Loss : ", np.mean(losses))

    ground_truth = np.asarray(ground_truth)
    # predictions = np.asarray(predictions)

    prec, recall, f1, _ = precision_recall_fscore_support(ground_truth, (np.array(predictions) > 0.5).astype(int))
    predictions = np.asarray(predictions)
    # try:
    #     print(f'GT shape before putting in ap: {ground_truth.shape}')
    #     print(f'pred shape before putting in ap: {predictions.shape}')
    # except:
    #     print(f'GT len before putting in ap: {len(ground_truth)}')
    #     print(f'pred len before putting in ap: {len(predictions)}')

    ap = average_precision_score(ground_truth, predictions, average=None)

    print(f'Macro f1 is {np.mean(f1)}')
    print(f'Macro prec is {np.mean(prec)}')
    print(f'Macro recall is {np.mean(recall)}')
    print(f'Classwise AP is {ap}')
    print(f'Macro AP is {np.mean(ap)}')

    for entry in range(len(vid_paths)):
        if str(vid_paths[entry].split('/')[-1]) not in pred_dict.keys():
            pred_dict[str(vid_paths[entry].split('/')[-1])] = []
            pred_dict[str(vid_paths[entry].split('/')[-1])].append(predictions[entry])

        else:
            # print('yes')
            pred_dict[str(vid_paths[entry].split('/')[-1])].append(predictions[entry])

    for entry in range(len(vid_paths)):
        if str(vid_paths[entry].split('/')[-1]) not in label_dict.keys():
            label_dict[str(vid_paths[entry].split('/')[-1])]= ground_truth[entry]

    return pred_dict, label_dict, np.mean(ap)

def val_epoch_ucf101(run_id, epoch,mode, skip, hflip, cropping_fac, pred_dict,label_dict, data_loader, fa_model, model, criterion, writer, use_cuda):
    print(f'validation at epoch {epoch} - mode {mode} - skip {skip} - hflip {hflip} - cropping_fac {cropping_fac}')

    model.eval()
    fa_model.eval()
    losses = []
    predictions, ground_truth = [], []
    vid_paths = []

    for i, (inputs, label, vid_path, _) in enumerate(data_loader):
        vid_paths.extend(vid_path)
        ground_truth.extend(label)
        # inputs = inputs.permute(0,4,1,2,3)
        if len(inputs.shape) != 1:

            # inputs = inputs.permute(0, 2, 1, 3, 4)

            if use_cuda:
                inputs = inputs.cuda()
                label = torch.from_numpy(np.asarray(label)).cuda()
            ori_bs, ori_t, ori_c, ori_h, ori_w = inputs.shape
            inputs = inputs.view(-1, inputs.shape[2], inputs.shape[3], inputs.shape[4])
            # print(inputs.shape)
            with torch.no_grad():

                output1 = fa_model(inputs)
                # print(f'fa_model output shape is {output1.shape}')
                output1 = output1.reshape(ori_bs, ori_t, ori_c, ori_h, ori_w)
                # print(f'fa_model reshaped output shape is {output1.shape}')
                # exit()
                output1 = output1.permute(0,2,1,3,4)
                output = model(output1)

                '''# print(label)
                if use_cuda:
                    inputs = inputs.cuda()
                    label = torch.from_numpy(np.asarray(label)).cuda()
                # print(label)

                with torch.no_grad():

                    output = model(inputs)
                    loss = criterion(output,label)'''
                loss = criterion(output,label)
                losses.append(loss.item())

            predictions.extend(nn.functional.softmax(output, dim = 1).cpu().data.numpy())
            # print(len(predictions))

            if i+1 % 99 == 0:
                print("Validation Epoch ", epoch , "mode", mode, "skip", skip, "hflip", hflip, " Batch ", i, "- Loss : ", np.mean(losses))

    del inputs, output, label, loss

    ground_truth = np.asarray(ground_truth)
    pred_array = np.flip(np.argsort(predictions,axis=1),axis=1) # Prediction with the most confidence is the first element here
    c_pred = pred_array[:,0] #np.argmax(predictions,axis=1).reshape(len(predictions))

    for entry in range(len(vid_paths)):
        if str(vid_paths[entry].split('/')[-1]) not in pred_dict.keys():
            pred_dict[str(vid_paths[entry].split('/')[-1])] = []
            pred_dict[str(vid_paths[entry].split('/')[-1])].append(predictions[entry])

        else:
            # print('yes')
            pred_dict[str(vid_paths[entry].split('/')[-1])].append(predictions[entry])

    for entry in range(len(vid_paths)):
        if str(vid_paths[entry].split('/')[-1]) not in label_dict.keys():
            label_dict[str(vid_paths[entry].split('/')[-1])]= ground_truth[entry]

    print_pred_array = []

    # for entry in range(pred_array.shape[0]):
    #     temp = ''
    #     for i in range(5):
    #         temp += str(int(pred_array[entry][i]))+' '
    #     print_pred_array.append(temp)
    # print(f'check {print_pred_array[0]}')
    # results = open('Submission1.txt','w')
    # for entry in range(len(vid_paths)):
    #     content = str(vid_paths[entry].split('/')[-1] + ' ' + print_pred_array[entry])[:-1]+'\n'
    #     results.write(content)
    # results.close()

    correct_count = np.sum(c_pred==ground_truth)
    accuracy = float(correct_count)/len(c_pred)

    # print(f'Correct Count is {correct_count}')
    print(f'Epoch {epoch}, mode {mode}, skip {skip}, hflip {hflip}, cropping_fac {cropping_fac}, Accuracy: {accuracy*100 :.3f}')
    return pred_dict, label_dict, accuracy, np.mean(losses)

def train_classifier(run_id, restart, saved_model= None, kin_pretrained= False):
    use_cuda = True
    best_score = 0
    writer = SummaryWriter(os.path.join(cfg.logs, str(run_id)))

    save_dir = os.path.join(cfg.saved_models_dir, run_id)
    if not os.path.exists(save_dir):
        os.makedirs(save_dir)
    fa_model = load_unet(saved_model_file = '/home/c3-0/ishan/privacy_preserving1/unet_reconstruction_training/saved_models/recon1/model_20_bestAcc_0.0100.pth') #UNet(n_channels = 3, n_classes=3)
    if kin_pretrained:
        ft_model = build_r3d_classifier(kin_pretrained= True, self_pretrained = False)
    else:
        ft_model =  load_r3d_classifier(saved_model_file= '/home/c3-0/ishan/privacy_preserving1/ucf101_with_unetrecon/saved_models_and_logs/model_50_bestAcc_0.6262.pth')

    fb_model = ssl_privacy_load(saved_model_file = '/home/c3-0/ishan/privacy_preserving1/vispr_with_unetreocon/saved_models/SSL_unet_recon/model_160.pth')

    # fb_model = load_privacy_classifier(saved_model_file = '/home/c3-0/ishan/privacy_preserving1/vispr_with_unetreocon/saved_models/vispr_unetrecon1/model_57_bestAP_0.6271.pth', arch = 'mv1')

    learning_rate_fa = params.learning_rate_fa
    learning_rate_fb = params.learning_rate_fb
    learning_rate_ft = params.learning_rate_ft

    # if kin_pretrained:
    #     learning_rate_fa = params.learning_rate_fa/10
    #     learning_rate_fb = params.learning_rate_fb/10
    #     learning_rate_ft = params.learning_rate_ft/10

    criterion_ft = nn.CrossEntropyLoss()
    criterion_fb = nn.BCEWithLogitsLoss()

    if torch.cuda.device_count()>1:
        print(f'Multiple GPUS found!')
        ft_model=nn.DataParallel(ft_model)
        fa_model=nn.DataParallel(fa_model)
        fb_model=nn.DataParallel(fb_model)
        ft_model.cuda()
        fa_model.cuda()
        fb_model.cuda()

    else:
        print('Only 1 GPU is available')
        ft_model.cuda()
        fa_model.cuda()
        fb_model.cuda()

    optimizer_fa = optim.Adam(fa_model.parameters(),lr=params.learning_rate_fa)
    optimizer_fb = optim.Adam(fb_model.parameters(),lr=params.learning_rate_fb)
    optimizer_ft = optim.Adam(ft_model.parameters(),lr=params.learning_rate_ft)

    train_dataset_ucf101 = baseline_dataloader_train_strong(shuffle = False, data_percentage = params.data_percentage_ucf101)
    # train_dataset = baseline_dataloader(shuffle = False, data_percentage = 0.1)

    train_dataloader_ucf101 = DataLoader(train_dataset_ucf101, batch_size=params.batch_size_ucf101, shuffle=True, num_workers=params.num_workers, collate_fn=collate_fn_train_ucf101)
    print(f'UCF101 Train dataset length: {len(train_dataset_ucf101)}')
    print(f'UCF101 Train dataset steps per epoch: {len(train_dataset_ucf101)/params.batch_size_ucf101}')

    train_dataset_vispr = vispr_ssl(datasplit = 'train', shuffle = False, data_percentage = params.data_percentage_vispr)
    # train_dataset = baseline_dataloader(shuffle = False, data_percentage = 0.1)

    train_dataloader_vispr = DataLoader(train_dataset_vispr, batch_size=params.batch_size_vispr, shuffle=True, num_workers=params.num_workers, collate_fn=collate_fn_train)
    print(f'VISPR Train dataset length: {len(train_dataset_vispr)}')
    print(f'VISPR Train dataset steps per epoch: {len(train_dataset_vispr)/params.batch_size_vispr}')
    epoch0  = 0
    val_array = [0, 3, 5,10,15,20, 25, 30, 35, 40, 45] + [50+ x for x in range(100)]

    modes = list(range(params.num_modes))
    skip = list(range(1,params.num_skips+1))
    hflip = params.hflip#list(range(2))
    cropping_fac1 = params.cropping_fac1#[0.7,0.85,0.8,0.75]
    print(f'Num modes {len(modes)}')
    print(f'Num skips {skip}')
    print(f'Cropping fac {cropping_fac1}')
    print(f'Base learning rate {params.learning_rate}')
    print(f'Scheduler patient {params.lr_patience}')
    print(f'Scheduler drop {params.scheduled_drop}')

    modes, skip,hflip, cropping_fac =  list(zip(*itertools.product(modes,skip,hflip,cropping_fac1)))
    accuracy = 0

    for epoch in range(epoch0, params.num_epochs):
        # if epoch < params.warmup and lr_flag1 ==0:
        #     learning_rate2 = params.warmup_array[epoch]*params.learning_rate

        print(f'Epoch {epoch} started')
        start=time.time()

        fa_model, fb_model, ft_model = train_epoch(run_id, epoch, train_dataloader_vispr, train_dataloader_ucf101, ft_model, fa_model, fb_model,  criterion_ft, criterion_fb, optimizer_fa, optimizer_fb, optimizer_ft, writer, use_cuda,learning_rate_fa, learning_rate_fb, learning_rate_ft)

        if epoch in val_array:
            pred_dict = {}
            label_dict = {}
            val_losses =[]

            validation_dataset = vispr_dataset_generator(datasplit = 'test', shuffle = False, data_percentage = params.data_percentage)

            validation_dataloader = DataLoader(validation_dataset, batch_size=10, shuffle=True, num_workers=params.num_workers, collate_fn=collate_fn_train_regular)
            '''
            pred_dict, label_dict, macro_ap = val_epoch_vispr(run_id, epoch, validation_dataloader, fa_model, fb_model, criterion_fb, writer, use_cuda)
            '''
            val_visualization_fa_vispr(save_dir, epoch, validation_dataloader, fa_model)

            pred_dict = {}
            label_dict = {}
            val_losses =[]

            for val_iter in range(len(modes)):
                try:
                    validation_dataset = multi_baseline_dataloader_val_strong(shuffle = True, data_percentage = params.data_percentage,\
                        mode = modes[val_iter], skip = skip[val_iter], hflip= hflip[val_iter], \
                        cropping_factor= cropping_fac[val_iter])
                    validation_dataloader = DataLoader(validation_dataset, batch_size=params.v_batch_size, shuffle=True, num_workers=params.num_workers, collate_fn=collate_fn_val_ucf101)
                    if val_iter ==0:
                        print(f'Validation dataset length: {len(validation_dataset)}')
                        print(f'Validation dataset steps per epoch: {len(validation_dataset)/params.v_batch_size}')

                    pred_dict, label_dict, accuracy, loss = val_epoch_ucf101(run_id, epoch,modes[val_iter],skip[val_iter],hflip[val_iter],cropping_fac[val_iter], \
                        pred_dict, label_dict, validation_dataloader, fa_model, ft_model, criterion_ft, writer, use_cuda)
                    val_losses.append(loss)

                    predictions1 = np.zeros((len(list(pred_dict.keys())),params.num_classes))
                    ground_truth1 = []
                    entry = 0
                    for key in pred_dict.keys():
                        predictions1[entry] = np.mean(pred_dict[key], axis =0)
                        entry+=1

                    for key in label_dict.keys():
                        ground_truth1.append(label_dict[key])

                    pred_array1 = np.flip(np.argsort(predictions1,axis=1),axis=1) # Prediction with the most confidence is the first element here
                    c_pred1 = pred_array1[:,0]

                    correct_count1 = np.sum(c_pred1==ground_truth1)
                    accuracy11 = float(correct_count1)/len(c_pred1)

                    print(f'Running Avg Accuracy is for epoch {epoch}, mode {modes[val_iter]}, skip {skip[val_iter]}, hflip {hflip[val_iter]}, cropping_fac {cropping_fac[val_iter]} is {accuracy11*100 :.3f}% ')
                except:
                    print(f'Failed epoch {epoch}, mode {modes[val_iter]}, skip {skip[val_iter]}, hflip {hflip[val_iter]}, cropping_fac {cropping_fac[val_iter]} is {accuracy11*100 :.3f}% ')

            val_loss = np.mean(val_losses)
            predictions = np.zeros((len(list(pred_dict.keys())),params.num_classes))
            ground_truth = []
            entry = 0
            for key in pred_dict.keys():
                predictions[entry] = np.mean(pred_dict[key], axis =0)
                entry+=1

            for key in label_dict.keys():
                ground_truth.append(label_dict[key])

            pred_array = np.flip(np.argsort(predictions,axis=1),axis=1) # Prediction with the most confidence is the first element here
            c_pred = pred_array[:,0]

            correct_count = np.sum(c_pred==ground_truth)
            accuracy1 = float(correct_count)/len(c_pred)
            print(f'Correct Count is {correct_count} out of {len(c_pred)}')
            writer.add_scalar('Validation Loss', np.mean(val_loss), epoch)
            writer.add_scalar('Validation Accuracy', np.mean(accuracy1), epoch)

            print(f'Overall Ft Accuracy is for epoch {epoch} is {accuracy1*100 :.3f}% ')
            # file_name = f'RunID_{run_id}_Acc_{accuracy1*100 :.3f}_cf_{len(cropping_fac1)}_m_{params.num_modes}_s_{params.num_skips}.pkl'
            # pickle.dump(pred_dict, open(file_name,'wb'))
            accuracy = accuracy1

        if epoch % 3 == 0:
            save_file_path = os.path.join(save_dir, 'model_{}.pth'.format(epoch))
            states = {
                'epoch': epoch + 1,
                # 'arch': params.arch,
                'fa_model_state_dict': fa_model.state_dict(),
                'fb_model_state_dict': fb_model.state_dict(),
                'ft_model_state_dict': ft_model.state_dict(),
            }
            torch.save(states, save_file_path)
        # We will save optimizer weights for each temp model, not all saved models to reduce the storage
        save_file_path = os.path.join(save_dir, 'model_temp.pth')
        states = {
                'epoch': epoch + 1,
                # 'arch': params.arch,
                'fa_model_state_dict': fa_model.state_dict(),
                'fb_model_state_dict': fb_model.state_dict(),
                'ft_model_state_dict': ft_model.state_dict(),
                'optimizer_fa': optimizer_fa.state_dict(),
                'optimizer_fb': optimizer_fb.state_dict(),
                'optimizer_ft': optimizer_ft.state_dict(),
            }
        torch.save(states, save_file_path)
        taken = time.time()-start
        print(f'Time taken for Epoch-{epoch} is {taken}')
        print()
        train_dataset_ucf101 = baseline_dataloader_train_strong(shuffle = False, data_percentage = params.data_percentage_ucf101)
        # train_dataset = baseline_dataloader(shuffle = False, data_percentage = 0.1)

        train_dataloader_ucf101 = DataLoader(train_dataset_ucf101, batch_size=params.batch_size_ucf101, shuffle=True, num_workers=params.num_workers, collate_fn=collate_fn_train_ucf101)
        print(f'UCF101 Train dataset length: {len(train_dataset_ucf101)}')
        print(f'UCF101 Train dataset steps per epoch: {len(train_dataset_ucf101)/params.batch_size_ucf101}')

        train_dataset_vispr = vispr_ssl(datasplit = 'train', shuffle = False, data_percentage = params.data_percentage_vispr)
        # train_dataset = baseline_dataloader(shuffle = False, data_percentage = 0.1)

        train_dataloader_vispr = DataLoader(train_dataset_vispr, batch_size=params.batch_size_vispr, shuffle=True, num_workers=params.num_workers, collate_fn=collate_fn_train)
        print(f'VISPR Train dataset length: {len(train_dataset_vispr)}')
        print(f'VISPR Train dataset steps per epoch: {len(train_dataset_vispr)/params.batch_size_vispr}')

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='Script to train baseline')

    parser.add_argument("--run_id", dest='run_id', type=str, required=False, default= "dummy_ssl_minmax",
                        help='run_id')
    parser.add_argument("--restart", action='store_true')
    parser.add_argument("--kin_pretrained", action='store_true')

    args = parser.parse_args()
    print(f'Restart {args.restart}')

    print(f'kin_pretrained {args.kin_pretrained}')
    run_id = args.run_id

    train_classifier(str(run_id), args.restart, None, args.kin_pretrained)

max_stars_repo_name: DAVEISHAN/SPAct

max_issues_repo_name: DAVEISHAN/SPAct

max_forks_repo_name: DAVEISHAN/SPAct


text:
from numpy.testing import *
from pymc.gp import *
from pymc.gp.cov_funs import *
from numpy import *
from copy import copy

from pymc import six
xrange = six.moves.xrange

def zero_fun(x):
    return zeros(x.shape[:-1], dtype=float)
M = Mean(zero_fun)

lin = linspace(-1., 1., 50)
x, y = meshgrid(lin, lin)
xstack = zeros(x.shape+(2,), dtype=float)
xstack[:,:, 0] = x
xstack[:,:, 1] = y

lin_obs = linspace(-.5, .5, 2)
xobs, yobs = meshgrid(lin_obs, lin_obs)
ostack = zeros(xobs.shape+(2,), dtype=float)
ostack[:,:, 0] = xobs
ostack[:,:, 1] = yobs

V = .02*ones((2, 2), dtype=float)
data = ones((2, 2), dtype=float)
N = 10

coef_cov = ones((2*N+1, 2*N+1), dtype=float)
for i in xrange(1, 2*N+1):
    int_i = int(((i+1)/2))
    for j in xrange(1, 2*N+1):
        int_j = int(((j+1)/2))
        coef_cov[i, j] = 1./(sqrt((int_i + int_j)**2))**2.5

basis = fourier_basis([N, N])
C = SeparableBasisCovariance(basis, coef_cov, xmin = [-4., -.4], xmax = [4., 4.])

class test_basiscov(TestCase):
    def test_sep(self):

        observe(M, C, obs_mesh = ostack, obs_V = V, obs_vals = data)

    def test_nonsep(self):

        basis_array = zeros((2*N+1, 2*N+1), dtype=object)
        for i in arange(2*N+1):

            if i%2 == 0:
                def funi(x, xmin, xmax, i=i):
                    T = xmax[0] - xmin[0]
                    return cos(i*pi*(x - xmin[0]) / T)

            else:
                def funi(x, xmin, xmax, i=i):
                    T = xmax[0] - xmin[0]
                    return sin((i+1)*pi*(x - xmin[0]) / T)

            for j in arange(2*N+1):

                if j%2 == 0:
                    def funj(x, xmin, xmax,j=j):
                        T = xmax[1] - xmin[1]
                        return cos(j*pi*(x - xmin[1]) / T)
                else:
                    def funj(x, xmin, xmax,j=j):
                        T = xmax[1] - xmin[1]
                        return sin((j+1)*pi*(x - xmin[1]) / T)

                def fun_now(x, xmin, xmax, funi=funi, funj=funj):
                    return funi(x[:, 0], xmin, xmax) * funj(x[:, 1], xmin, xmax)

                basis_array[i, j] = fun_now

        C2 = BasisCovariance(basis_array, coef_cov, xmin = [-4., -.4], xmax = [4., 4.])
        assert_almost_equal(C2(xstack), C(xstack))

        observe(M, C, obs_mesh = ostack, obs_V = V, obs_vals = data)

max_stars_repo_name: CamDavidsonPilon/pymc

max_issues_repo_name: kforeman/pymc

max_forks_repo_name: kforeman/pymc

reason: from numpy


text:
#!/usr/bin/env python3
# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.

import logging

import numpy as np
import torch
from caffe2.python import workspace
from ml.rl.test.gym.open_ai_gym_environment import ModelType
from ml.rl.training.training_data_page import TrainingDataPage

logger = logging.getLogger(__name__)

class OpenAIGymMemoryPool:
    def __init__(self, max_replay_memory_size):
        """
        Creates an OpenAIGymMemoryPool object.

        :param max_replay_memory_size: Upper bound on the number of transitions
            to store in replay memory.
        """
        self.replay_memory = []
        self.max_replay_memory_size = max_replay_memory_size
        self.memory_num = 0
        self.skip_insert_until = self.max_replay_memory_size

    def sample_memories(self, batch_size, model_type):
        """
        Samples transitions from replay memory uniformly at random.

        :param batch_size: Number of sampled transitions to return.
        :param model_type: Model type (discrete, parametric).
        """
        cols = [[], [], [], [], [], [], [], [], [], [], []]
        indices = np.random.permutation(len(self.replay_memory))[:batch_size]
        for idx in indices:
            memory = self.replay_memory[idx]
            for col, value in zip(cols, memory):
                col.append(value)

        states = torch.tensor(cols[0], dtype=torch.float32)
        next_states = torch.tensor(cols[3], dtype=torch.float32)

        if model_type == ModelType.PYTORCH_PARAMETRIC_DQN.value:
            num_possible_actions = len(cols[7][0])

            actions = torch.tensor(cols[1], dtype=torch.float32)
            possible_actions = []
            for pa_matrix in cols[8]:
                logger.info("PA" + str(pa_matrix))
                for row in pa_matrix:
                    possible_actions.append(row)

            tiled_states = states.repeat(1, num_possible_actions).reshape(
                -1, states.shape[1]
            )
            possible_actions = torch.tensor(possible_actions, dtype=torch.float32)
            possible_actions_state_concat = torch.cat(
                (tiled_states, possible_actions), dim=1
            )
            possible_actions_mask = torch.tensor(cols[9], dtype=torch.float32)

            next_actions = torch.tensor(cols[4], dtype=torch.float32)
            possible_next_actions = []
            for pna_matrix in cols[6]:
                for row in pna_matrix:
                    logger.info("PNA" + str(row))
                    possible_next_actions.append(row)

            tiled_next_states = next_states.repeat(1, num_possible_actions).reshape(
                -1, next_states.shape[1]
            )

            possible_next_actions = torch.tensor(
                possible_next_actions, dtype=torch.float32
            )
            possible_next_actions_state_concat = torch.cat(
                (tiled_next_states, possible_next_actions), dim=1
            )
            possible_next_actions_mask = torch.tensor(cols[7], dtype=torch.float32)
        else:
            possible_actions = None
            possible_actions_state_concat = None
            possible_next_actions = None
            possible_next_actions_state_concat = None
            if cols[7] is None or cols[7][0] is None:
                possible_next_actions_mask = None
            else:
                possible_next_actions_mask = torch.tensor(cols[7], dtype=torch.float32)
            if cols[9] is None or cols[9][0] is None:
                possible_actions_mask = None
            else:
                possible_actions_mask = torch.tensor(cols[9], dtype=torch.float32)

            actions = torch.tensor(cols[1], dtype=torch.float32)
            next_actions = torch.tensor(cols[4], dtype=torch.float32)

        return TrainingDataPage(
            states=states,
            actions=actions,
            propensities=None,
            rewards=torch.tensor(cols[2], dtype=torch.float32).reshape(-1, 1),
            next_states=next_states,
            next_actions=next_actions,
            possible_next_actions=possible_next_actions,
            not_terminal=torch.from_numpy(
                np.logical_not(np.array(cols[5]), dtype=np.bool).astype(np.int32)
            ).reshape(-1, 1),
            time_diffs=torch.tensor(cols[10], dtype=torch.int32).reshape(-1, 1),
            possible_actions_mask=possible_actions_mask,
            possible_actions_state_concat=possible_actions_state_concat,
            possible_next_actions_mask=possible_next_actions_mask,
            possible_next_actions_state_concat=possible_next_actions_state_concat,
        )

    def insert_into_memory(
        self,
        state,
        action,
        reward,
        next_state,
        next_action,
        terminal,
        possible_next_actions,
        possible_next_actions_mask,
        time_diff,
        possible_actions,
        possible_actions_mask,
    ):
        """
        Inserts transition into replay memory in such a way that retrieving
        transitions uniformly at random will be equivalent to reservoir sampling.
        """
        item = (
            state,
            action,
            reward,
            next_state,
            next_action,
            terminal,
            possible_next_actions,
            possible_next_actions_mask,
            possible_actions,
            possible_actions_mask,
            time_diff,
        )

        if self.memory_num < self.max_replay_memory_size:
            self.replay_memory.append(item)
        elif self.memory_num >= self.skip_insert_until:
            p = float(self.max_replay_memory_size) / self.memory_num
            self.skip_insert_until += np.random.geometric(p)
            rand_index = np.random.randint(self.max_replay_memory_size)
            self.replay_memory[rand_index] = item
        self.memory_num += 1

max_stars_repo_name: johncliu/Horizon

max_issues_repo_name: johncliu/Horizon

max_forks_repo_name: johncliu/Horizon


text:
import numpy as np
import matplotlib.pyplot as plt

data = np.loadtxt('../data/populations.txt')
year, hares, lynxes, carrots = data.T  # trick: columns to variables

plt.axes([0.2, 0.1, 0.5, 0.8])
plt.plot(year, hares, year, lynxes, year, carrots)
plt.legend(('Hare', 'Lynx', 'Carrot'), loc=(1.05, 0.5))
plt.show()