text:
import torch
from torch import nn
import gym
from matplotlib import pyplot as plt
from tqdm import tqdm
import numpy as np
import json

from deepq import computations, deepqnetworks, policies, replay_buffers, utils
import experiment_utils
import hardcoded_policies

class QNet(nn.Module):
    """Just a Q Network class to test the components
    """

    def __init__(self, state_shape, act_shape, n_layers, n_units, activation=torch.tanh):
        super().__init__()
        self.state_shape = state_shape
        self.act_shape = act_shape
        self.activation = activation

        self.in_layer = nn.Linear(in_features=state_shape, out_features=n_units)
        self.hidden_layers = [nn.Linear(n_units, n_units) for i in range(n_layers - 2)]
        self.out_layer = nn.Linear(in_features=n_units, out_features=act_shape)

    def forward(self, x):
        x = self.activation(self.in_layer(x))
        for layer in self.hidden_layers:
            x = self.activation(layer(x))
        x = self.out_layer(x)
        return x

def plot_episode_results(episode_lengths, losses, episode_end_steps, argmax_q_values, epsilons,
                         avg_targets):
    print('Training completed in ' + str(len(episode_lengths)) + ' episodes')
    plt.plot(range(len(losses)), losses, label='Loss')
    plt.plot(episode_end_steps, episode_lengths, label='Episode lengths')
    plt.plot(range(len(argmax_q_values)), argmax_q_values,
             label='Q value for argmax action')
    plt.plot(range(len(epsilons)), epsilons, label='Epsilon')
    plt.plot(range(len(avg_targets)), avg_targets, label='Average target value')
    plt.legend()
    plt.show()

def test_agent(q_network, env, num_episodes):
    q_network.eval()
    policy_test = policies.GreedyPolicy()
    episode_rewards = [0]
    state = env.reset()
    for i in range(num_episodes):
        while True:
            with torch.no_grad():
                q_values = q_network.predict_values(state)
                action = policy_test.act(q_values)[0]
            next_state, reward, done, _ = env.step(action)
            episode_rewards[-1] += reward
            if done:
                state = env.reset()
                episode_rewards.append(0)
            else:
                state = next_state
    print('Average total reward over ' + str(num_episodes) + ' episodes: ' +
          str(np.mean(episode_rewards)))

def run_experiment(p_hardcoded_start, p_hardcoded_schedule,
                   p_hardcoded_end, hardcoded_policy=None, training_steps=4000, pretrain_steps=1000,
                   update_net_steps=10, batch_size=128, df=0.99, lr=0.5e-3, buffer_len=10000,
                   epsilon_start=1.0, epsilon_end=0.1, epsilon_schedule='exp', decay_steps=None):
    if decay_steps is None:
        decay_steps = training_steps
    if p_hardcoded_schedule == 'lin_decay':
        assert p_hardcoded_start >= p_hardcoded_end
    if p_hardcoded_schedule == 'lin_upcay':
        assert p_hardcoded_end >= p_hardcoded_start
    p_hardcoded = p_hardcoded_start

    env = gym.make('CartPole-v0')

    # Replay buffer and prefiller
    replay_buffer = replay_buffers.FIFOReplayBuffer(maxlen=buffer_len)
    prefiller = replay_buffers.UniformGymPrefiller()

    # Q Networks
    q_net = QNet(state_shape=4, act_shape=2, n_layers=3, n_units=16)
    dq_nets = deepqnetworks.TargetDQNetworks(q_network=q_net)

    # Target computer and trainer
    target_computer = computations.DoubleQTargetComputer(dq_networks=dq_nets, df=df)
    trainer = computations.DQNTrainer(dq_networks=dq_nets, lr=lr)
    if epsilon_schedule == 'exp_decay':
        policy_train = policies.ExponentialDecayEpsilonGreedy(decay_steps=decay_steps,
                                                              start_epsilon=epsilon_start,
                                                              end_epsilon=epsilon_end)
    elif epsilon_schedule == 'lin_decay':
        policy_train = policies.LinearDecayEpsilonGreedyPolicy(decay_steps=decay_steps,
                                                               start_epsilon=epsilon_start,
                                                               end_epsilon=epsilon_end)
    elif epsilon_schedule == 'const':
        policy_train = policies.FixedEpsilonGreedyPolicy(epsilon=epsilon_start)

    prefiller.fill(replay_buffer, env, pretrain_steps, shuffle=False)

    # Training
    losses = []
    rewards = [0]
    episode_end_steps = []
    argmax_q_values = []
    epsilons = []
    avg_targets = []

    episode_steps = 0
    state = env.reset()
    for step in tqdm(range(training_steps)):
        with torch.no_grad():
            q_values = dq_nets.predict_values(state)
            if np.random.binomial(1, p=p_hardcoded):
                action = hardcoded_policy(state)
            else:
                action = policy_train.act(q_values)[0]
            if p_hardcoded_schedule == 'lin_decay':
                p_hardcoded = p_hardcoded_start - step * \
                              (p_hardcoded_start - p_hardcoded_end) / training_steps
            elif p_hardcoded_schedule == 'lin_upcay':
                p_hardcoded = p_hardcoded_start + \
                              step * (p_hardcoded_end - p_hardcoded_start) / training_steps
        argmax_q_values.append(float(q_values.numpy().max()))
        epsilons.append(policy_train.cur_epsilon)
        next_state, reward, done, _ = env.step(action)
        if done:
            next_state = None
        replay_buffer.remember((state, action, reward, next_state))
        rewards[-1] += reward

        if len(replay_buffer.buffer) >= batch_size:
            batch = replay_buffer.sample(size=batch_size)
            batch = utils.split_replay_batch(batch)
            targets = target_computer.compute_targets(batch)
            avg_targets.append(float(targets.mean()))
            loss, grads = trainer.train(batch, targets)
            losses.append(float(loss))
            if step % update_net_steps == 0:
                dq_nets.update()
        else:
            avg_targets.append(0)
            losses.append(0)

        episode_steps += 1
        if done:
            state = env.reset()
            rewards.append(0)
            episode_end_steps.append(step)
            episode_steps = 0
        else:
            state = next_state
    del rewards[-1]
    return {
        'rewards': rewards,
        'end_steps': episode_end_steps,
        'losses': losses,
        'argmax_q': argmax_q_values,
        'epsilons': epsilons,
        'avg_targets': avg_targets,
        'parameters': {
            'training_steps': training_steps,
            'pretrain_steps': pretrain_steps,
            'buffer_len': buffer_len,
            'update_net_steps': update_net_steps,
            'epsilon_start': epsilon_start,
            'epsilon_end': epsilon_end,
            'epsilon_decay_steps': decay_steps,
            'lr': lr,
            'batch_size': batch_size,
            'df': df,
            'p_hardcoded_start': p_hardcoded_start,
            'p_hardcoded_end': p_hardcoded_end,
            'p_hardcoded_schedule': p_hardcoded_schedule,
        }
    }

if __name__ == '__main__':
    # Only these values should be modified across experiments
    p_hardcoded_starts = [0.1, 0.2, 0.3, 0.4, 0.5]
    p_hardcoded_ends = np.zeros_like(p_hardcoded_starts)

    p_hardcoded_schedule = 'lin_decay'
    NUM_EXPERIMENTS = 20

    results_folder = 'experiment_results/policy_cartpole_medium/'
    file_prefix = 'const_epsilon'
    hardcoded_policy = hardcoded_policies.cartpole_medium
    epsilon_start = 0.2
    epsilon_end = 0.0
    epsilon_schedule = 'const'

    for p_hardcoded_start, p_hardcoded_end in zip(p_hardcoded_starts, p_hardcoded_ends):
        p_hardcoded_start = round(p_hardcoded_start, 1)
        p_hardcoded_end = round(p_hardcoded_end, 1)
        exp_name = file_prefix + p_hardcoded_schedule
        if p_hardcoded_start != p_hardcoded_end:
            exp_name += '_' + str(p_hardcoded_start) + '_' + str(p_hardcoded_end)
        else:
            exp_name += '_' + str(p_hardcoded_start)

        exp_losses = []
        exp_rewards = []
        exp_episode_end_steps = []
        exp_argmax_q_values = []
        exp_epsilons = []
        exp_avg_targets = []

        for e in range(NUM_EXPERIMENTS):
            res = run_experiment(p_hardcoded_start=p_hardcoded_start,
                                 p_hardcoded_end=p_hardcoded_end,
                                 p_hardcoded_schedule=p_hardcoded_schedule,
                                 hardcoded_policy=hardcoded_policy,
                                 epsilon_start=epsilon_start,
                                 epsilon_end=epsilon_end,
                                 epsilon_schedule=epsilon_schedule)
            exp_rewards.append(res['rewards'])
            exp_episode_end_steps.append(res['end_steps'])
            exp_losses.append(res['losses'])
            exp_avg_targets.append(res['avg_targets'])
            exp_argmax_q_values.append(res['argmax_q'])
            exp_epsilons.append(res['epsilons'])

        experiment_data = {
            'n_runs': NUM_EXPERIMENTS,
            'p_hardcoded_schedule': p_hardcoded_schedule,
            'exp_rewards': exp_rewards,
            'exp_losses': exp_losses,
            'exp_end_steps': exp_episode_end_steps,
            'exp_avg_targets': exp_avg_targets,
            'exp_q_values': exp_argmax_q_values,
            'exp_epsilons': exp_epsilons,
            'parameters': res['parameters'],
        }

        reward_plots = [
            np.array([end_steps, lengths])
            for end_steps, lengths in zip(exp_episode_end_steps, exp_rewards)
        ]

        x_avg, y_avg, var, y_arrays = experiment_utils.merge_plots(reward_plots)

        json_content = json.dumps(experiment_data)
        f = open(results_folder + exp_name, 'w')
        f.write(json_content)
        f.close()
        print('Written ' + results_folder + exp_name)

max_stars_repo_name: fedetask/des-rl

max_issues_repo_name: fedetask/des-rl

max_forks_repo_name: fedetask/des-rl


text:
# -*- coding: utf-8 -*-

import numpy as np
import matplotlib.pylab as plt

from dmp_3 import *
from utils.stg.min_jerk_traj import trajectory

# start and goal of the movement (1-dim)
# trajectory is a list of 3-tuples with (pos,vel,acc)

start = 0.5
goal = 1.0
duration = 1.0
delta_t = 0.001

y = trajectory(start, goal, 1.0, delta_t)

plt.figure()
plt.title('Trajectory X - Demo (Td)')
plt.plot(np.matrix(y)[:,0],'r-')
plt.plot(np.matrix(y)[:,1],'g-')
plt.plot(np.matrix(y)[:,2],'b-')
plt.show()

# Learn dmp
dmp1 = DiscreteDMP(OriginalFormulation(),
LWR(activation=0.3, exponentially_spaced=False, n_rfs=8, use_offset=True))

dmp2 = DiscreteDMP(ImprovedFormulation(),
LWR(activation=0.3, exponentially_spaced=False, n_rfs=8, use_offset=True))

# Time is scalar but in the other implementations is a  array of time_steps.
time = 1.0
dmp1.learn(y, time)
dmp2.learn(y, time)

plt.figure()
plt.title('Original From - Forcing Function: Desired (Fd) and predicted (Fp)')
plt.plot(dmp1.Fd,'r--', label='Fd')
plt.plot(dmp1.Fp,'b-', label='Fp')
plt.show()

plt.figure()
plt.title('Improvement From - Forcing Function: Desired (Fd) and predicted (Fp)')
plt.plot(dmp2.Fd,'r--', label='Fd')
plt.plot(dmp2.Fp,'b-', label='Fp')
plt.show()

# Plan a new trajectory
dmp1.use_ft = False
dmp1.setup(start+0.4, goal+0.2, 1.0)
y1, y1d, y1dd = dmp1.plan(time)

plt.figure()
plt.title('Original From - Trajectory X - Generated (Tp)')
plt.plot(y1,'r-')
plt.plot(y1d,'g-')
plt.plot(y1dd,'b-')
plt.show()

dmp2.use_ft = False
dmp2.setup(start+0.4, goal+0.2, 1.0)
y2, y2d, y2dd = dmp2.plan(time)

plt.figure()
plt.title('Improved From - Trajectory X - Generated (Tp)')
plt.plot(y2,'r-')
plt.plot(y2d,'g-')
plt.plot(y2dd,'b-')
plt.show()

plt.figure()
plt.title('Trajectory X - Error in position')
err1 = abs(np.squeeze(np.array(np.matrix(y)[:,0])).shape - y1)
err2 = abs(np.squeeze(np.array(np.matrix(y)[:,0])).shape - y2)
plt.plot(err1)
plt.plot(err2)

ext: py

max_stars_repo_name: dgerod/more-dmps

max_issues_repo_name: dgerod/more-dmps

max_forks_repo_name: dgerod/more-dmps

reason: import numpy


text:
import tensorflow as tf
import numpy as np

from bi_tempered_loss import bi_tempered_logistic_loss

t1 = 0.8
t2 = 1.4
label_smoothing = 0.01
y_true = tf.convert_to_tensor([[0, 1, 0], [0, 0, 1]], dtype=tf.float32)
y_pred = tf.convert_to_tensor([[0.05, 0.95, 0], [0.1, 0.8, 0.1]])
loss = bi_tempered_logistic_loss(y_pred, y_true, t1=t1, t2= t2, label_smoothing =label_smoothing)
assert loss.shape == (2,)
print(loss.numpy())

expected = np.array([0.49713564 ,0.92891073])
np.testing.assert_array_almost_equal(loss.numpy(), expected)


text:
#%%
import sys
import os
import os.path as osp

from PIL import Image
import cv2
import numpy as np
import yaml
from colorama import Back, Style # assign color options on your text(ref) https://stackoverflow.com/questions/287871/how-to-print-colored-text-to-the-terminal
import torch
import torchvision.transforms as T

from datasets import make_dataloader
from utils.reranking import re_ranking  # (ref) https://github.com/DoranLyong/person-reid-tiny-baseline/blob/master/utils/reranking.py
from utils.metrics import eval_func
from reid_utils import (ReID_FLAGS,
                        ReID_INFERENCE,
                        )

""" Path checking
"""
python_ver = sys.version
script_path = os.path.abspath(__file__)
cwd = os.getcwd()
os.chdir(cwd) #changing working directory

print(f"Python version: {Back.GREEN}{python_ver}{Style.RESET_ALL}")
print(f"The path of the running script: {Back.MAGENTA}{script_path}{Style.RESET_ALL}")
print(f"CWD is changed to: {Back.RED}{cwd}{Style.RESET_ALL}")

#%% Load configurations in YAML
try:
    with open('config/reid_cfg.yaml', 'r') as cfg_yaml:
        cfg = yaml.load(cfg_yaml, Loader=yaml.FullLoader)
        print("YAML is loaded o_< chu~")

except:
    sys.exit("fail to load YAML...")

reid_flags = ReID_FLAGS(cfg)
#print(vars(reid_flags))  # check the class members(ref) https://www.programiz.com/python-programming/methods/built-in/vars

QUERY_DIR = '/home/kist-ubuntu/workspace_reID/data/Market-1501/query/'

#%%
def vis_tensorImg(img:torch.Tensor):
    """ - tensor img to PIL (ref) https://discuss.pytorch.org/t/pytorch-pil-to-tensor-and-vice-versa/6312/2
        - show normalized tensor image (ref) https://discuss.pytorch.org/t/conversion-from-a-tensor-to-a-pil-image-not-working-well-what-is-going-wrong/26121/2
    """
    img = img.to('cpu')

    # Denormalize the tensor image
    std = [0.229, 0.224, 0.225] # get 'std' from (ref) https://github.com/DoranLyong/person-reid-tiny-baseline/blob/master/datasets/make_dataloader.py
    mean = [0.485, 0.456, 0.406]
    z = img * torch.tensor(std).view(3, 1, 1)
    z = z + torch.tensor(mean).view(3, 1, 1)

    img_pil = T.ToPILImage()(z).convert("RGB")
    img_pil.show()

def visualizer(query_img, img_path, test_img, camid, top_k = 10, img_size=[256,128]):
    figure = np.asarray(query_img.resize((img_size[1],img_size[0])))

    for k in range(top_k):
        name = str(indices[0][k]).zfill(6)
        img = np.asarray(Image.open(img_path[indices[0][k]]).resize((img_size[1],img_size[0])))
        figure = np.hstack((figure, img))
        title=name

    figure = cv2.cvtColor(figure,cv2.COLOR_BGR2RGB)

    cv2.imshow("results", figure)

#    pil_image=Image.fromarray(figure)  # (ref) https://www.delftstack.com/ko/howto/matplotlib/convert-a-numpy-array-to-pil-image-python/
#    pil_image.show()

#%%
#%% Metrics
def euclidean_distance(qf, gf):
    m = qf.shape[0]
    n = gf.shape[0]
    dist_mat = torch.pow(qf, 2).sum(dim=1, keepdim=True).expand(m, n) + \
               torch.pow(gf, 2).sum(dim=1, keepdim=True).expand(n, m).t()
    dist_mat.addmm_(1, -2, qf, gf.t())
    return dist_mat.cpu().numpy()

def cosine_similarity(qf, gf):
    epsilon = 0.00001
    dist_mat = qf.mm(gf.t())
    qf_norm = torch.norm(qf, p=2, dim=1, keepdim=True)  # mx1
    gf_norm = torch.norm(gf, p=2, dim=1, keepdim=True)  # nx1
    qg_normdot = qf_norm.mm(gf_norm.t())

    dist_mat = dist_mat.mul(1 / qg_normdot).cpu().numpy()
    dist_mat = np.clip(dist_mat, -1 + epsilon, 1 - epsilon)
    dist_mat = np.arccos(dist_mat)
    return dist_mat

#%%
if __name__ == '__main__':

    """ Set your device
    """
    gpu_no = 0  # gpu_number
    DEVICE = torch.device( f'cuda:{gpu_no}' if torch.cuda.is_available() else 'cpu')

    print(f"device: { DEVICE }")

    """ init. inference object
    """
    model = ReID_INFERENCE(reid_flags, DEVICE)

    """ set dataloader
    """
    transform = T.Compose([
        T.Resize([256, 128]),
        T.ToTensor(),
        T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
    ])

    """ Test loop
    """
    query_list = sorted(os.listdir(QUERY_DIR))
    for test_img in query_list:

        # load gallery items
        gallery_feats = torch.load('./log/gfeats.pth').to(DEVICE) # gallery features

        gallery_img_path = np.load('./log/imgpath.npy')

        # smaple query
        query_img = Image.open(QUERY_DIR + test_img)
#        query_img.show()

        input = torch.unsqueeze(transform(query_img), 0)  # [3, H, W] -> [1, 3, H, W] for torch tensor
        input = input.to(DEVICE)

        query_feat = model(input)  # get feature in (1, 2048) tensor shape
        norm_query = torch.nn.functional.normalize(query_feat, dim=1, p=2) # gfeats.pth 에 저장된 features도 normalized 된 상태.
        print(f"check if normalized : {(norm_query**2).sum()} ")    # (ref) https://discuss.pytorch.org/t/question-about-functional-normalize-and-torch-norm/27755

        """ feature metrics
        """
#        dist_mat = re_ranking(norm_query , gallery_feats , k1=30, k2=10, lambda_value=0.2) # good & slow
        dist_mat = euclidean_distance(norm_query, gallery_feats) # not bad & fast
#        dist_mat = cosine_similarity(norm_query, gallery_feats)  # bad & fast

        """ Sorting
        """
        indices = np.argsort(dist_mat, axis=1)  # get index order in the best order (short distnace first)

        print(f"Finding ID of {test_img}")
        print(f"Rank-10 gallery paths: {gallery_img_path[indices[0, :10]]}")  # Rank-10 results by slicing 10 items

        """ Visualize
        """
        visualizer(query_img, gallery_img_path , test_img, camid='mixed', top_k = 10, img_size=[256,128])
        cv2.waitKey(32)

    cv2.destroyAllWindows()

# %%

max_stars_repo_name: DoranLyong/person-reid-tiny-baseline

max_issues_repo_name: DoranLyong/person-reid-tiny-baseline

max_forks_repo_name: DoranLyong/person-reid-tiny-baseline


text:
import torch
import torch.nn as nn
from torch.nn.utils.rnn import PackedSequence
from torch import Tensor
from torch import nn, autograd
import torch.nn.functional as F
from torchaudio import functional as F_audio
import math
from torch.nn.utils.rnn import pad_sequence, pack_padded_sequence, pad_packed_sequence
import torchaudio

from torchaudio.compliance.kaldi import mfcc
from speechpy.processing import cmvn, cmvnw

def fast_tanh(x):
    return x / (1 + x.abs())

class FastTanh(nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, x):
        return x / (1 + x.abs())

class TimeReduction(nn.Module):
    def __init__(self, reduction_factor):
        super().__init__()
        self.reduction_factor = reduction_factor

    def forward(self, xs):
        batch_size, xlen, hidden_size = xs.shape
        reduction_factor = xlen % self.reduction_factor
        if reduction_factor != 0:
            reduction_factor =  self.reduction_factor - reduction_factor
        pad_shape = [0, 0, 0, reduction_factor, 0, 0]
        xs = nn.functional.pad(xs, pad_shape)
        xs = xs.view(batch_size, -1, self.reduction_factor, hidden_size)
        xs = xs.mean(dim=2)
        return xs

class MFCC_(torch.nn.Module):
    r"""Create the Mel-frequency cepstrum coefficients from an audio signal.

    By default, this calculates the MFCC on the DB-scaled Mel spectrogram.
    This is not the textbook implementation, but is implemented here to
    give consistency with librosa.

    This output depends on the maximum value in the input spectrogram, and so
    may return different values for an audio clip split into snippets vs. a
    a full clip.

    Args:
        sample_rate (int, optional): Sample rate of audio signal. (Default: ``16000``)
        n_mfcc (int, optional): Number of mfc coefficients to retain. (Default: ``40``)
        dct_type (int, optional): type of DCT (discrete cosine transform) to use. (Default: ``2``)
        norm (str, optional): norm to use. (Default: ``'ortho'``)
        log_mels (bool, optional): whether to use log-mel spectrograms instead of db-scaled. (Default: ``False``)
        melkwargs (dict or None, optional): arguments for MelSpectrogram. (Default: ``None``)
    """
    __constants__ = ['sample_rate', 'n_mfcc', 'dct_type', 'top_db', 'log_mels']

    def __init__(self,
                 sample_rate: int = 16000,
                 n_mfcc: int = 40,
                 dct_type: int = 2,
                 norm: str = 'ortho',
                 log_mels: bool = False,
                 normalize=False,
                 melkwargs= None) -> None:
        super(MFCC_, self).__init__()
        supported_dct_types = [2]
        if dct_type not in supported_dct_types:
            raise ValueError('DCT type not supported'.format(dct_type))
        self.sample_rate = sample_rate
        self.n_mfcc = n_mfcc
        self.dct_type = dct_type
        self.norm = norm
        self.normalize = normalize
        self.top_db = 80.0
        stype = 'power'
        self.multiplier = 10.0 if stype == 'power' else 20.0
        self.amin = 1e-10
        self.ref_value = 1.0
        self.db_multiplier = math.log10(max(self.amin, self.ref_value))
        # self.amplitude_to_DB = AmplitudeToDB('power', self.top_db)

        if melkwargs is not None:
            self.MelSpectrogram = torchaudio.transforms.MelSpectrogram(
                sample_rate=self.sample_rate, **melkwargs)
        else:
            self.MelSpectrogram = torchaudio.transforms.MelSpectrogram(sample_rate=self.sample_rate)

        if self.n_mfcc > self.MelSpectrogram.n_mels:
            raise ValueError(
                'Cannot select more MFCC coefficients than # mel bins')
        dct_mat = F_audio.create_dct(
            self.n_mfcc, self.MelSpectrogram.n_mels, self.norm)
        self.register_buffer('dct_mat', dct_mat)
        self.log_mels = log_mels

    def forward(self, waveform: Tensor) -> Tensor:
        r"""
        Args:
            waveform (Tensor): Tensor of audio of dimension (..., time).

        Returns:
            Tensor: specgram_mel_db of size (..., ``n_mfcc``, time).
        """

        # pack batch
        shape = waveform.size()
        waveform = waveform.view(-1, shape[-1])

        mel_specgram = self.MelSpectrogram(waveform)
        if self.log_mels:
            log_offset = 1e-6
            mel_specgram = torch.log(mel_specgram + log_offset)
        else:
            mel_specgram = F_audio.amplitude_to_DB(mel_specgram, self.multiplier, self.amin, self.db_multiplier, self.top_db)
        # (channel, n_mels, time).tranpose(...) dot (n_mels, n_mfcc)
        # -> (channel, time, n_mfcc).tranpose(...)
        mfcc = torch.matmul(mel_specgram.transpose(1, 2),
                            self.dct_mat).transpose(1, 2)

        # unpack batch
        mfcc = mfcc.view(shape[:-1] + mfcc.shape[-2:])

        if self.normalize:

            mfcc = torch.from_numpy(cmvnw(mfcc.numpy().T, win_size=201)).T
            # mean_vec = mfcc.mean(dim=1)
            # mfcc = torch.sub(mfcc, mean_vec[:, None])

        return mfcc

class ResidualRNNModel(nn.Module):
    def __init__(self, input_size, vocab_size, hidden_size, num_layers, dropout=.2, blank=0, bidirectional=False):
        super(ResidualRNNModel, self).__init__()
        self.hidden_size = hidden_size
        self.num_layers = num_layers
        self.vocab_size = vocab_size
        self.blank = blank
        self.layer_norm = nn.LayerNorm(input_size)
        stacked = StackedRecurrent(residual=True, merge_first=True)
        pre_lstm = nn.LSTM(input_size, hidden_size, 1, batch_first=True)
        stacked.add_module('0', pre_lstm)

        for i in range(num_layers-1):
            layer = nn.LSTM(hidden_size, hidden_size, 1, batch_first=True)
            stacked.add_module(str(i+1), layer)
            stacked.add_module(str(i+1)+'1', nn.LayerNorm(hidden_size))

        # normalize spectrum feature
        # lstm hidden vector: (h_0, c_0) num_layers * num_directions, batch, hidden_size
        # nn.LSTM(input_size, hidden_size, num_layers, batch_first=True, dropout=dropout, bidirectional=False)
        self.lstm = stacked
        # if bidirectional: hidden_size *= 2
        self.linear = None
        if vocab_size == hidden_size:
            self.linear = nn.Linear(hidden_size, vocab_size)

    def forward(self, xs, hid=None):
        xs = self.layer_norm(xs)
        h, hid = self.lstm(xs, hid)
        if self.linear != None:
            h = self.linear(h)
        return h, hid

    def greedy_decode(self, xs):
        xs = self(xs)[0][0]  # only one sequence
        xs = F.log_softmax(xs, dim=1)
        logp, pred = torch.max(xs, dim=1)
        return pred.data.cpu().numpy(), -float(logp.sum())

    def beam_search(self, xs, W):
        ''' CTC '''
        xs = self(xs)[0][0]  # only one sequence
        logp = F.log_softmax(xs, dim=1)
        return ctc_beam(logp.data.cpu().numpy(), W)

class ResidualProjModel(nn.Module):
    def __init__(self, input_size, vocab_size, hidden_size,num_layers, ff_dim=-1, dropout=.2, blank=0, bidirectional=False):
        super(ResidualProjModel, self).__init__()
        self.hidden_size = hidden_size
        self.num_layers = num_layers
        self.vocab_size = vocab_size
        self.blank = blank
        if ff_dim == -1:
            ff_dim = hidden_size//2
        stacked = StackedRecurrent(residual=True, merge_first=True)

        pre_lstm = nn.Sequential(
            nn.LSTM(input_size, hidden_size, 1, batch_first=True),
            nn.Linear(hidden_size, ff_dim),
            FastTanh(),
        )
        stacked.add_module('0', pre_lstm)

        for i in range(num_layers-1):
            layer = nn.Sequential(
                nn.Linear(ff_dim, hidden_size),
                nn.LSTM(hidden_size, hidden_size, 1, batch_first=True),
                nn.Linear(hidden_size, ff_dim),
                FastTanh(),
            )
            stacked.add_module(str(i+1), layer)

        # normalize spectrum feature
        # lstm hidden vector: (h_0, c_0) num_layers * num_directions, batch, hidden_size
        # nn.LSTM(input_size, hidden_size, num_layers, batch_first=True, dropout=dropout, bidirectional=False)
        self.lstm = stacked
        # if bidirectional: hidden_size *= 2
        self.linear = None
        if vocab_size == hidden_size:
            self.linear = nn.Linear(hidden_size, vocab_size)

    def forward(self, xs, hid=None):
        h, hid = self.lstm(xs, hid)
        if self.linear != None:
            h = self.linear(h)
        return h, hid

    def greedy_decode(self, xs):
        xs = self(xs)[0][0]  # only one sequence
        xs = F.log_softmax(xs, dim=1)
        logp, pred = torch.max(xs, dim=1)
        return pred.data.cpu().numpy(), -float(logp.sum())

    def beam_search(self, xs, W):
        ''' CTC '''
        xs = self(xs)[0][0]  # only one sequence
        logp = F.log_softmax(xs, dim=1)
        return ctc_beam(logp.data.cpu().numpy(), W)

class StackedRecurrent(nn.Sequential):

    def __init__(self, dropout=0, residual=False, normalization=False, merge_first=True, reduction_factor=2):
        super(StackedRecurrent, self).__init__()
        self.residual = residual
        self.dropout = dropout
        self.normalization = normalization
        if merge_first:
            self.concat = TimeReduction(reduction_factor=reduction_factor)
        self.merge_first = merge_first

    def forward(self, inputs, hidden=None):
        hidden = hidden or tuple([None] * len(self))
        next_hidden = []
        hidden_idx = 0

        for i, module in enumerate(self._modules.values()):

            if isinstance(module, TimeReduction):
                continue
            elif i == 4 and self.merge_first:
                inputs = self.concat(inputs)
            elif isinstance(module, nn.LayerNorm):
                inputs = module(inputs)
                continue

            output, h = module(inputs, hidden[hidden_idx])
            hidden_idx += 1
            next_hidden.append(h)
            if self.residual and inputs.size(-1) == output.size(-1):
                inputs = output + inputs
            else:
                inputs = output
            if isinstance(inputs, PackedSequence):
                data = nn.functional.dropout(
                    inputs.data, self.dropout, self.training)
                inputs = PackedSequence(data, inputs.batch_sizes)
            else:
                inputs = nn.functional.dropout(
                    inputs, self.dropout, self.training)

        return output, tuple(next_hidden)

class NormalizationLayer(nn.Module):
    def __init__(self, num_features):
        super(NormalizationLayer, self).__init__()
        self.norm = nn.InstanceNorm1d(num_features)

    def forward(self, inputs, hidden=None):
        x = inputs.permute(0, 2, 1)
        x = self.norm(x)
        return x.permute(0, 2, 1)

class ConcatFeature(torch.nn.Module):

    def __init__(self, merge_size=3):
        super(ConcatFeature, self).__init__()
        self.merge_size = merge_size

    def forward(self, waveform: Tensor) -> Tensor:
        batch_size, waveform_len, feat = waveform.shape
        if waveform_len % self.merge_size != 0:
            pad_wave = torch.zeros(
                (batch_size, self.merge_size - (waveform_len % self.merge_size), feat))
            if waveform.is_cuda:
                pad_wave = pad_wave.to(waveform.get_device())
            waveform = torch.cat([waveform, pad_wave], dim=1)

        return waveform.reshape(batch_size, -1, feat*self.merge_size)

if __name__ == "__main__":
    import numpy as np
    import pickle
    # pickle.dump(MFCC_(), open('test.pt', 'wb'))
    import torchaudio

    trans = MFCC_(normalize=True, log_mels=True)
    data, sr = torchaudio.load('bloom.mp3', normalization=True)
    print(data[0].shape)
    mfcc_f = trans(data[0])
    print(mfcc_f.shape)
    print(torch.mean(mfcc_f), torch.var(mfcc_f))
    print('first feature')
    print(torch.mean(mfcc_f[:, 0]), torch.var(mfcc_f))
    print('feature 0')
    print(mfcc_f[0].shape,torch.mean(mfcc_f[0]), torch.var(mfcc_f))

    # stacked = StackedRecurrent(residual=True, merge_first=True)

    # cell = nn.LSTM(40, 64, batch_first=True)
    # stacked.add_module('0',cell)
    # cell = nn.LSTM(64*3, 64*3, batch_first=True)
    # stacked.add_module('1',cell)
    # cell = nn.LSTM(64*3, 64*3, batch_first=True)
    # stacked.add_module('2',cell)

    # inputs = torch.randn((32, 10, 40))
    # outputs, hid = stacked(inputs)
    # inputs = torch.randn((32, 10, 40))
    # outputs, hid = stacked(inputs)
    inputs = torch.randn((32, 10, 40))
    hidden_size = 128
    num_layers = 4
    input_size = 40
    # stacked = StackedRecurrent(residual=True, merge_first=True)

    # pre_lstm = nn.LSTM(input_size, hidden_size, 1, batch_first=True)
    # stacked.add_module('0', pre_lstm)
    # _lstm = nn.LSTM((hidden_size//3)*3, hidden_size, 1, batch_first=True)
    # stacked.add_module('1', _lstm)

    # for i in range(num_layers):
    #     layer = nn.LSTM(hidden_size, hidden_size, 1, batch_first=True)
    #     stacked.add_module(str(i+1), layer)
    # stacked(inputs)

    import torch
    from torch import nn

    a = torch.randn(32, 100, 1)
    m = nn.Conv1d(100, 100, 1, stride=2)
    out = m(a)
    print(out.size())
    print(m)
    # model = ResidualRNNModel(40, 3600, 128, 4)
    # outputs, hid = model(inputs)
    # print(outputs.shape)
    # concat = ConcatFeature()
    # x = torch.from_numpy(np.array([
    #     [
    #         list([1]*4),
    #         list([2]*4),
    #         list([3]*4),
    #         list([4]*4),
    #     ],
    #     [
    #         list([4]*4),
    #         list([5]*4),
    #         list([5]*4),
    #         list([5]*4),
    #     ]
    # ], dtype=float)).float()
    # print(x.shape)
    # x = concat(x)
    # print(x)
    # print(x.shape)

max_stars_repo_name: entn-at/Online-Speech-Recognition

max_issues_repo_name: entn-at/Online-Speech-Recognition

max_forks_repo_name: entn-at/Online-Speech-Recognition


text:
import sys
import os.path
import re
import warnings
from io import StringIO

import matplotlib
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import statsmodels.api as sm
from patsy import dmatrices
import six

class PerfData():
    __DATETIME_HEADER__ = "start-time"
    __PERF_HEADER__ = __DATETIME_HEADER__ + ",response-time"

    def __init__(self, filename):
        self._filename = filename

    def data(self):
        with open(self._filename) as f:
            _data = f.read()
        return _data

    def headers(self):
        return self.__PERF_HEADER__

    def datetime_headers(self):
        return self.__DATETIME_HEADER__

class PerformanceRunIterator():
    def __init__(self, data, header):
        self._data = data
        self._current_index = 0
        self._perf_header = header

    def __iter__(self):
        self._header_indexes = [m.start() for m in re.finditer(self._perf_header, self._data)]
        self._header_indexes.append(len(self._data))

        return self

    def next(self):
        if self._current_index + 1 >= len(self._header_indexes):
            raise StopIteration
        line = self._line_at_index(self._current_index)
        self._current_index = self._current_index + 1

        return line

    def _line_at_index(self, position):
        start = self._header_indexes[position]
        end = self._header_indexes[position + 1]
        line = self._data[start:end]
        return six.text_type(line)

def read_throughput_data(filename):
    perf_data = PerfData(filename)

    df = pd.DataFrame()
    for run in PerformanceRunIterator(perf_data.data(), perf_data.headers()):
        run_dataframe = pd.read_csv(StringIO(run), parse_dates=[perf_data.datetime_headers()])

        trimmed_section = trim_edges(run_dataframe)

        if len(trimmed_section) > 0:
            df = df.append(trimmed_section)

    # Reset the index because it is a Frankenstein of smaller indexes
    df = df.reset_index().drop('index', axis=1)
    return df

def trim_edges(data):
    indexes = data.set_index('start-time').resample('1S').aggregate(lambda x: 1).index
    test_start_time = indexes[0]
    test_end_time = indexes[-1]
    return data[(data['start-time'] >= test_start_time) & (data['start-time'] <= test_end_time)]

def process_throughput_data(data):
    buckets = data.set_index('start-time')['response-time'].resample('1S')
    throughput_data_set = buckets.aggregate({"throughput": lambda x: 0 if x.count() == 0 else x.count()})

    throughput_data_set = throughput_data_set.reset_index()
    throughput_data_set = throughput_data_set.fillna(method='ffill')
    return buckets, throughput_data_set

def generate_fit_line(data):
    y, x = dmatrices('latency ~ throughput', data=data, return_type='dataframe')
    fit = sm.GLM(y, x, family=sm.families.InverseGaussian(sm.families.links.inverse_squared)).fit()
    max_throughput = data['throughput'].max()
    min_throughput = data['throughput'].min()
    domain = np.arange(min_throughput, max_throughput)
    prediction_inputs = np.ones((len(domain), 2))
    prediction_inputs[:, 1] = domain
    fit_line = fit.predict(prediction_inputs)
    return domain, fit_line, round(max_throughput)

if __name__ == '__main__':
    matplotlib.style.use('ggplot')

    matplotlib.rcParams['figure.figsize'] = 9, 6
    matplotlib.rcParams['legend.loc'] = 'best'
    matplotlib.rcParams['figure.dpi'] = 120

    # We'll need these packages for plotting fit lines
    warnings.filterwarnings('ignore')
    performanceResultsFile = sys.argv[1]
    assert os.path.isfile(performanceResultsFile), 'Missing performance results file'

    compareDatasets = False

    if compareDatasets:
        assert os.path.isfile('old_perfResults.csv'), 'Missing old performance results file "old_perfResults.csv"'

    goData = read_throughput_data(performanceResultsFile)

    throughputBuckets, throughputData = process_throughput_data(goData)

    if compareDatasets:
        oldGoData = read_throughput_data('old_perfResults.csv')
        oldThroughputBuckets, oldThroughputData = process_throughput_data(oldGoData)

    goData['throughput'] = throughputBuckets.transform(len).reset_index()['response-time']
    goData.columns = ['start-time', 'latency', 'throughput']

    if compareDatasets:
        oldGoData['throughput'] = oldThroughputBuckets.transform(len).reset_index()['response-time']
        oldGoData.columns = ['start-time', 'latency', 'throughput']

    domain, goFitLine, xLimit = generate_fit_line(goData)

    if compareDatasets:
        oldDomain, oldGoFitLine, oldXLimit = generate_fit_line(oldGoData)

    fig, ax = plt.subplots()

    # Change the value of `c` to change the color. http://matplotlib.org/api/colors_api.html
    ax = goData.plot(ax=ax, kind='scatter', x='throughput', y='latency', c='b', marker='.', alpha=0.2)
    ax.plot(domain, goFitLine, c='b', lw=2)  # Plot the fit line

    if compareDatasets:
        ax = oldGoData.plot(ax=ax, kind='scatter', x='throughput', y='latency', c='r', marker='.', alpha=0.2)
        ax.plot(oldDomain, oldGoFitLine, c='r', lw=2)  # Plot the fit line
        ax.legend(['after', 'before'])

    # To update x & y axis range change the parameters in function set_(x/y)lim(lower_limit, uppper_limit)

    ax.autoscale(True)
    ax.autoscale_view(True, True, True)
    plt.xlabel('Throughput (requests/sec)')
    plt.ylabel('Latency (sec)')
    plt.title('Headroom plot', y=1.05)
    plt.plot()

    filenameForPlot = performanceResultsFile[:-4] + "Plot.png"
    plt.savefig(filenameForPlot)
    print ("saving graph to " + filenameForPlot)

reason: import numpy,import statsmodels


text:
# -*- coding: utf-8 -*-
# @version : Python3.6
# @Time    : 2017/4/23 9:59
# @Author  : Jianyang-Hu
# @contact : jianyang1993@163.com
# @File    : minlin_04223.py
# @Software: PyCharm

import numpy as np
from numpy import *
import sys
from matplotlib.pyplot import plot
from matplotlib.pyplot import show

#最小二乘法
N = int(sys.argv[-1])
c = loadtxt("stocks.csv", delimiter=",", usecols=(4, ), unpack=True)
b = c[-N:]       #倒着N个数
b = b[::-1]     #颠倒顺序
print("b:",b)

A = np.zeros((N,N),float)
print("A:",A)

for i in range(N):
    A[i, ] = c[-N - 1 - i : -1 -i]
print("A:",A)

(x, residuals,rank,s) = np.linalg.lstsq(A,b)

print(x,residuals,rank,s)
print(np.dot(b,x))

max_stars_repo_name: Jianyang-Hu/numpypractice

max_issues_repo_name: Jianyang-Hu/numpypractice

max_forks_repo_name: Jianyang-Hu/numpypractice

reason: import numpy,from numpy


text:
#! /usr/bin/env python3

import argparse
import math
import os
import sys

import numpy as np

import matplotlib
import matplotlib.pyplot as plt
import matplotlib.image as mpimg

from helpers import *

# Arguments ---

parser = argparse.ArgumentParser(
    description='Displays bounding boxes over SynScapes images.')

parser.add_argument('path',
                    help='Path to SynScapes root directory.')
parser.add_argument('--index', '-i', type=int, required=True,
                    help='Index of image to visualize')
parser.add_argument('--type', choices=['2d', '3d', 'class', 'instance'],
                    default='2d',
                    help='Type of visualization')
parser.add_argument('--threshold', type=float, default=90,
                    help='Percentage of occlusion and truncation at which to cull bounding boxes.')
parser.add_argument('--xkcd', action='store_true',
                    help='Yes.')
parser.add_argument('--save',
                    help='Saves the visualized image to the given filename.')

args = parser.parse_args()

# Constants ---

class_id_to_str = {
    24: 'person',
    25: 'rider',
    26: 'car',
    27: 'truck',
    28: 'bus',
    31: 'train',
    32: 'motorcyclist',
    33: 'bicyclist'
}

colors = [
    (  0,  0,  0),
    (  0,  0,  0),
    (  0,  0,  0),
    (  0,  0,  0),
    (  0,  0,  0),
    (111, 74,  0),
    ( 81,  0, 81),
    (128, 64,128),
    (244, 35,232),
    (250,170,160),
    (230,150,140),
    ( 70, 70, 70),
    (102,102,156),
    (190,153,153),
    (180,165,180),
    (150,100,100),
    (150,120, 90),
    (153,153,153),
    (153,153,153),
    (250,170, 30),
    (220,220,  0),
    (107,142, 35),
    (152,251,152),
    ( 70,130,180),
    (220, 20, 60),
    (255,  0,  0),
    (  0,  0,142),
    (  0,  0, 70),
    (  0, 60,100),
    (  0,  0, 90),
    (  0,  0,110),
    (  0, 80,100),
    (  0,  0,230),
    (119, 11, 32)
]

colormap = matplotlib.colors.ListedColormap(
    [(c[0] / 255.0, c[1] / 255.0, c[2] / 255.0) for c in colors])

# Functions ---

def draw_2d():
    for inst, box in bbox_data.items():
        # Info
        class_idx = metadata['instance']['class'][inst]
        class_name = class_id_to_str[class_idx]
        occluded = metadata['instance']['occluded'][inst]
        if occluded > args.threshold / 100.0:
            continue
        # Box
        xmin = box['xmin']
        xmax = box['xmax']
        ymin = box['ymin']
        ymax = box['ymax']
        c = (xmin * w, ymin * h)
        width = (xmax - xmin) * w
        height = (ymax - ymin) * h
        # Color
        color = tableau20[class_idx % len(tableau20)]
        # Draw
        patch = plt.Rectangle(c, width, height, linewidth=3, edgecolor=color, fill=False)
        ax.add_patch(patch)
        # Text
        plt.text(c[0], c[1] - 2, '{}'.format(class_name),
                 color=text_color, bbox=dict(facecolor=color, alpha=0.5))

def draw_3d():
    vehicle_to_sensor_xform = vehicle_to_sensor_transform(metadata)
    # Projection
    fx = metadata['camera']['intrinsic']['fx']
    fy = metadata['camera']['intrinsic']['fy']
    u0 = metadata['camera']['intrinsic']['u0']
    v0 = metadata['camera']['intrinsic']['v0']
    # Transform corner points to image space
    is_pts = {}
    for inst, box in bbox_data.items():
        # Cull
        occluded = metadata['instance']['occluded'][inst]
        if occluded > args.threshold / 100.0:
            continue
        # Vehicle coordinates
        vc = box['origin']
        # Add zero vector for origin offset
        box['o'] = (0, 0, 0)
        # Image coordinates
        is_pts[inst] = {
            name: vehicle_space_to_image_space(vec_sum(box['origin'], box[name]),
                                               vehicle_to_sensor_xform,
                                               fx, fy, u0, v0) \
            for name in ['o', 'x', 'y', 'z']
        }
    # Create lines
    lines = []
    for inst, box in is_pts.items():
        # Color
        class_idx = metadata['instance']['class'][inst]
        class_name = class_id_to_str[class_idx]
        color = tableau20[class_idx % len(tableau20)]
        # Vertices
        o = box['o']
        x = box['x']
        y = box['y']
        z = box['z']
        # Lines
        for tgt in [x, y, z]:
            verts_x = [o[0], tgt[0]]
            verts_y = [o[1], tgt[1]]
            line = plt.Line2D(verts_x, verts_y, c=color, lw=2)
            ax.add_line(line)
            # plt.scatter(verts_x, verts_y, c=(1, 1, 1))
        # Text
        plt.text(o[0] + 2, o[1] - 2, '{}'.format(class_name),
                 color=text_color, bbox=dict(facecolor=color, alpha=0.5))

# Script ---

root = os.path.abspath(args.path)

# Ensure root exists
if not os.path.exists(root):
    print('Invalid path:', root)
    sys.exit(1)

img_dir, meta_dir = [os.path.join(root, x) for x in ['img', 'meta']]

# Ensure integrity
for d in [img_dir, meta_dir]:
    if not os.path.exists(d):
        print('Missing directory:', d)
        sys.exit(1)

img_path = os.path.join(root, 'img', 'rgb', '{}.png'.format(args.index))
meta_path = os.path.join(root, 'meta')

metadata = read_metadata(meta_path, args.index)[args.index]

img = mpimg.imread(img_path)
w = float(img.shape[1])
h = float(img.shape[0])

text_color = (0.9, 0.9, 0.9)
if args.xkcd:
    text_color = (0, 0, 0)
    plt.xkcd()

fig, ax = plt.subplots(1)
dpi = fig.get_dpi()

# Add image background
ax.imshow(img)

# Add boxes
if args.type == '2d':
    bbox_data = metadata['instance']['bbox{}'.format(args.type)]
    draw_2d()
elif args.type == '3d':
    bbox_data = metadata['instance']['bbox{}'.format(args.type)]
    draw_3d()
elif args.type == 'class':
    class_path = os.path.join(root, 'img', args.type, '{}.png'.format(args.index))
    class_img = mpimg.imread(class_path)
    ax.imshow(class_img, cmap=colormap, alpha=0.66,
              norm=matplotlib.colors.Normalize(vmin=0, vmax=33/255.0), interpolation="none")
elif args.type == 'instance':
    instance_path = os.path.join(root, 'img', args.type, '{}.png'.format(args.index))
    instance_img = mpimg.imread(instance_path)
    ax.imshow(instance_img, alpha=0.5)

fig.set_size_inches(w / dpi, h / dpi, forward=True)
plt.axis('off')
plt.tight_layout()

# Save or Display
if args.save:
    plt.subplots_adjust(left=0, right=1, top=1, bottom=0)
    plt.savefig(args.save, pad_inches=0)
else:
    plt.show()

max_stars_repo_name: TillBeemelmanns/synscapes-utils

max_issues_repo_name: TillBeemelmanns/synscapes-utils

max_forks_repo_name: TillBeemelmanns/synscapes-utils


text:
import os
from distutils.core import setup
from distutils.extension import Extension

import numpy
from Cython.Distutils import build_ext

project_name = 'micromorphic'
cxx_standard = '11'

###############################
# Get the include directories #
###############################

source_files = [os.path.abspath('../cpp/material_python_interface.cpp'),
                os.path.abspath('materials.pyx')]

include_dirs = [numpy.get_include(), os.path.abspath('../cpp/')]
library_dirs = [os.path.abspath('../cpp')]

shared_libraries = ['micromat']

ext_modules = [Extension(project_name,
               sources=source_files,
               language='c++',
               libraries=shared_libraries,
               include_dirs=include_dirs,
               extra_compile_args=[f"-std=c++{cxx_standard}"],
               extra_link_args=[f"-L{d}" for d in library_dirs])]

for e in ext_modules:
    e.cython_directives = {'language_level': "3"}

setup(
  name = project_name,
  cmdclass = {'build_ext': build_ext},
  ext_modules = ext_modules
)

max_stars_repo_name: lanl/tardigrade-micromorphic-element

max_issues_repo_name: lanl/tardigrade-micromorphic-element

max_forks_repo_name: lanl/tardigrade-micromorphic-element


text:
"""
ASAPXYZ class for handing atomic coordinate input and compute/output
"""
import os
import glob
import json
from yaml import dump as ydump
from yaml import Dumper
import numpy as np
from ase import Atoms
from ase.io import read, write
from ase.build import niggli_reduce
from tqdm.auto import tqdm
from joblib import Parallel, delayed
import pandas as pd

from ..io import randomString, NpEncoder
from ..descriptors import Atomic_Descriptors, Global_Descriptors

class ASAPXYZ:
    """extended xyz class

    Parameters
    ----------
    fxyz: string_like
         the path to the extended xyz file
    fmat: string_like
           the name of the descriptors in the extended xyz file
    use_atomic_desc: bool
            return the descriptors for each atom, read from the xyz file
    stride: int
           the stride when reading the xyz file
    """

    def __init__(self, fxyz=None, stride=1, periodic=True, fileformat=None):
        # compile a list of matching xyz files
        # in fact they don't strictly need to be xyz format, anything that can be read by ASE is fine
        # a list of possible file formats: https://wiki.fysik.dtu.dk/ase/ase/io/io.html
        if '*' in fxyz:
            self.fxyz = glob.glob(fxyz)
            print("Find matching input files with coordinates: ", self.fxyz)
        else:
            self.fxyz = fxyz
        # essential
        self.stride = stride
        self.periodic = periodic
        if fileformat is not None:
            import ast
            self.fileformat = ast.literal_eval(fileformat)
        else:
            self.fileformat = {}

        # store the xyz file
        self.frames = None
        self.nframes = 0
        self.natom_list = []  # number of atoms for each frame
        self.total_natoms = 0  # total number of atoms for all frames
        self.global_species = []  # list of elements contains in all frames

        # record the state of the computation, e.g. which descriptors have been computed
        self.computed_desc_dict = {'data': {'fxyz': fxyz}}
        self.computed_desc_dict = {'descriptors': {}}
        # the conversion between tag of the descriptors and their acronyms
        self.tag_to_acronym = {'global': {}, 'atomic': {}}

        # we make a dictionary to store the computed descriptors
        self.global_desc = {}
        # this is for the atomic ones
        self.atomic_desc = {}

        # try to read the xyz file
        try:
            if isinstance(self.fxyz, (tuple, list)):
                self.frames = []
                for f in self.fxyz:
                    self.frames += read(f, slice(0, None, self.stride), **self.fileformat)
            else:
                self.frames = read(self.fxyz, slice(0, None, self.stride), **self.fileformat)
        except:
            raise ValueError('Exception occurred when loading the input file')

        self.nframes = len(self.frames)
        all_species = []
        for i, frame in enumerate(self.frames):
            # record the total number of atoms
            self.natom_list.append(len(frame.get_positions()))
            all_species.extend(frame.get_atomic_numbers())
            if frame.get_pbc()[0] and frame.get_pbc()[1] and frame.get_pbc()[2]:
                # niggli_reduce
                niggli_reduce(frame)
            if not self.periodic or not np.sum(frame.get_cell()) > 0:
                frame.set_pbc([False, False, False])
            # we also initialize the descriptor dictionary for each frame
            self.global_desc[i] = {}
            self.atomic_desc[i] = {}

        self.total_natoms = np.sum(self.natom_list)
        self.max_atoms = max(self.natom_list)
        # Keep things in plain python for serialisation
        self.global_species = np.unique(all_species).tolist()
        print('load xyz file: ', self.fxyz,
              ', a total of ', str(self.nframes), 'frames',
              ', a total of ', str(self.total_natoms), 'atoms',
              ', with elements: ', self.global_species, '.')

    def get_xyz(self):
        return self.frames

    def get_num_frames(self):
        return self.nframes

    def get_total_natoms(self):
        return self.total_natoms

    def get_natom_list(self):
        return self.natom_list

    def get_global_species(self):
        return self.global_species

    def get_natom_list_by_species(self, species_name=None):
        if species_name is None:
            return self.natom_list
        elif species_name in self.global_species:
            return [a.get_atomic_numbers().tolist().count(species_name) for a in self.frames]
        else:
            raise ValueError("Cannot find the specified chemical species in the data set.")

    def save_state(self, filename, mode='yaml'):
        if mode == 'yaml':
            with open(filename + '-state.yaml', 'w') as yd:
                ydump(self.computed_desc_dict, yd, sort_keys=True, Dumper=Dumper)
        else:
            with open(filename + '-state.json', 'w') as jd:
                json.dump(self.computed_desc_dict, jd, sort_keys=True, cls=NpEncoder)

    def save_descriptor_acronym_state(self, filename, mode='yaml'):

        if mode == 'yaml':
            with open(filename + '-descriptor-acronyms.yaml', 'w') as yd:
                ydump(self.tag_to_acronym, yd, sort_keys=True)
        else:
            with open(filename + '-descriptor-acronyms.json', 'w') as jd:
                json.dump(self.tag_to_acronym, jd, sort_keys=True, cls=NpEncoder)

    def symmetrise(self, sbs=[], symprec=1e-2):
        import spglib
        if len(sbs) == 0:
            sbs = range(self.nframes)
        for i in sbs:
            frame = self.frames[i]
            frame.info['space_group'] = spglib.get_spacegroup(frame, symprec=symprec)

    def standardize(self, sbs=[], symprec=1e-2):
        """
        reduce to primitive cell
        """
        import spglib
        if len(sbs) == 0:
            sbs = range(self.nframes)
        for i in sbs:
            frame = self.frames[i]
            lattice, scaled_positions, numbers = spglib.standardize_cell(frame,
                                                                         to_primitive=1,
                                                                         no_idealize=1,
                                                                         symprec=symprec)
            self.frames[i] = Atoms(numbers=numbers, cell=lattice, scaled_positions=scaled_positions,
                                   pbc=frame.get_pbc())

    def _add_info_to_desc_spec(self, desc_spec_dict):
        """
        add some system specific information to the list to descriptor specifications
        Parameters
        ----------
        desc_spec_dict: dictionaries
                        dict that specify which global descriptor to use.
        """
        for element in desc_spec_dict.keys():
            desc_spec_dict[element]['species'] = self.global_species
            desc_spec_dict[element]['periodic'] = self.periodic
            desc_spec_dict[element]['max_atoms'] = self.max_atoms

    def compute_atomic_descriptors(self, desc_spec_dict={}, sbs=[], tag=None, n_process=1):
        """
        compute the atomic descriptors for selected frames
        Parameters
        ----------
        desc_spec: a list of dictionaries
                 contrains infos on the descriptors to use
        e.g.
        .. code-block:: python

            atomic_desc_dict = {
            "firstsoap":
            {"type": 'SOAP',"species": [1, 6, 7, 8], "cutoff": 2.0, "atom_gaussian_width": 0.2, "n": 4, "l": 4}}

        sbs: array, integer
             the index of the subset of structures to compute
        """

        if len(sbs) == 0:
            sbs = range(self.nframes)
        if tag is None: tag = randomString(6)

        # add some system specific information to the list to descriptor specifications
        self._add_info_to_desc_spec(desc_spec_dict)

        # business!
        atomic_desc = Atomic_Descriptors(desc_spec_dict)

        # serial computation
        if n_process == 1:
            for i in tqdm(sbs):
                frame = self.frames[i]
                # compute atomic descriptor
                self.atomic_desc[i].update(atomic_desc.compute(frame))
        # parallel computation
        elif n_process >= 2:
            results = Parallel(n_jobs=n_process, verbose=1)(delayed(atomic_desc.compute)(self.frames[i]) for i in sbs)
            for i, atomic_desc_dict_now in enumerate(results):
                self.atomic_desc[i].update(atomic_desc_dict_now)
        else:
            raise ValueError("Please set the number of processes to be a positive integer.")

        # we mark down that this descriptor has been computed
        self.computed_desc_dict[tag] = atomic_desc.desc_spec_dict

    def compute_global_descriptors(self, desc_spec_dict={}, sbs=[], keep_atomic=False, tag=None, n_process=1):
        """
        compute the atomic descriptors for selected frames
        Parameters
        ----------
        desc_spec_dict: dictionaries that specify which global descriptor to use.

        e.g.
        .. code-block:: python

            {'global_desc1':
                          {"type": 'CM'}}

            # or

            {'global_desc2':
                          {'atomic_descriptor':
                                       atomic_desc_dict,
                           'reducer_function':
                                       reducer_dict
                          }}

            atomic_desc_dict = {
                              "firstsoap":
                                       {"type": 'SOAP',
                                       "species": [1, 6, 7, 8],
                                       "cutoff": 2.0,
                                       "atom_gaussian_width": 0.2,
                                       "n": 4,
                                       "l": 4}}

            reducer_dict = {'first_reducer':
                                     {'reducer_type': reducer_type,
                                     'zeta': zeta,
                                     'species': species,
                                     'element_wise': element_wise}}

        sbs: array, integer
             list of the indexes of the subset
        """

        if len(sbs) == 0:
            sbs = range(self.nframes)
        if tag is None: tag = randomString(6)

        # add some system specific information to the list to descriptor specifications
        self._add_info_to_desc_spec(desc_spec_dict)

        # business! Intialize a Global_Descriptors object
        global_desc = Global_Descriptors(desc_spec_dict)

        # serial computation
        if n_process == 1:
            for i in tqdm(sbs):
                frame = self.frames[i]
                # compute atomic descriptor
                desc_dict_now, atomic_desc_dict_now = global_desc.compute(frame)
                self.global_desc[i].update(desc_dict_now)
                if keep_atomic:
                    for _, v in atomic_desc_dict_now.items():
                        self.atomic_desc[i].update(v)
        # parallel computation
        elif n_process >= 2:
            results = Parallel(n_jobs=n_process, verbose=1)(delayed(global_desc.compute)(self.frames[i]) for i in sbs)
            for i, (desc_dict_now, atomic_desc_dict_now) in enumerate(results):
                self.global_desc[i].update(desc_dict_now)
                if keep_atomic:
                    for _, v in atomic_desc_dict_now.items():
                        self.atomic_desc[i].update(v)
        else:
            raise ValueError("Please set the number of processes to be a positive integer.")

        # we mark down that this descriptor has been computed
        self.computed_desc_dict['descriptors'][tag] = global_desc.desc_spec_dict

    def fetch_computed_descriptors(self, desc_dict_keys=[], sbs=[]):
        """
        Fetch the computed descriptors for selected frames
        Parameters
        ----------
        desc_spec_keys: a list (str-like) of keys
                    for which computed descriptors to fetch.
        sbs: array, integer

        Returns
        -------
        desc: np.matrix [n_frame, n_desc]
        """
        if len(sbs) == 0:
            sbs = list(range(self.nframes))
        return np.vstack([self._parse_computed_descriptors_singleframe(desc_dict_keys, i) for i in sbs])

    def fetch_computed_atomic_descriptors(self, desc_dict_keys=[], sbs=[]):
        """
        Fetch the computed atomic descriptors for selected frames
        Parameters
        ----------
        desc_spec_keys: a list (str-like) of keys
                    for which computed descriptors to fetch.
        sbs: array, integer

        Returns
        -------
        desc: np.matrix [n_atoms, n_desc]
        """
        if len(sbs) == 0:
            sbs = list(range(self.nframes))
        return np.vstack([self._parse_computed_atomic_descriptors_singleframe(desc_dict_keys, i) for i in sbs])

    def _parse_computed_descriptors_singleframe(self, desc_dict_keys=[], i=0):
        """return the global descriptor computed for frame i"""
        # TODO: use the nested dictionary search `extract_from_nested_dict` in ..io
        desc_array = np.array([])
        for e in desc_dict_keys:
            try:
                desc_array = np.append(desc_array, self.global_desc[i][e]['descriptors'])
            except:
                # if we use atomic to global descriptor, this is a nested dictionary
                for e2 in self.global_desc[i][e].keys():
                    for e3 in self.global_desc[i][e][e2].keys():
                        desc_array = np.append(desc_array, self.global_desc[i][e][e2][e3]['descriptors'])
        return desc_array

    def _parse_computed_atomic_descriptors_singleframe(self, desc_dict_keys=[], i=0):
        """return the atomic descriptor computed for frame i"""
        return np.hstack( self.atomic_desc[i][e]['atomic_descriptors'] for e in desc_dict_keys)

    def _write_computed_descriptors_to_xyz(self, desc_dict_now, frame):
        """
        we recorded the computed descriptors to the xyz object
        we use acronym to record the entry in the extended xyz file, so it's much easier to ready by human
        """
        for e in desc_dict_now.keys():
            try:
                frame.info[desc_dict_now[e]['acronym']] = desc_dict_now[e]['descriptors']
                self.tag_to_acronym['global'][e] = desc_dict_now[e]['acronym']
            except:
                # if we use atomic to global descriptor, this is a nested dictionary
                self.tag_to_acronym['global'][e] = {}
                for e2 in desc_dict_now[e].keys():
                    self.tag_to_acronym['global'][e][e2] = {}
                    for e3 in desc_dict_now[e][e2].keys():
                        frame.info[desc_dict_now[e][e2][e3]['acronym']] = desc_dict_now[e][e2][e3]['descriptors']
                        self.tag_to_acronym['global'][e][e2][e3] = desc_dict_now[e][e2][e3]['acronym']

    def _write_computed_atomic_descriptors_to_xyz(self, atomic_desc_dict_now, frame):
        """
        we recorded the computed descriptors to the xyz object
        we use acronym to record the entry in the extended xyz file, so it's much easier to ready by human
        """
        for e in atomic_desc_dict_now.keys():
            frame.new_array(atomic_desc_dict_now[e]['acronym'],
                            atomic_desc_dict_now[e]['atomic_descriptors'])
            self.tag_to_acronym['atomic'][e] = atomic_desc_dict_now[e]['acronym']

    def _desc_name_with_wild_card(self, desc_name_list, atomic_desc=False):
        """
        Use a wildcard when specifying the name of the descriptors
        """
        new_desc_name = []
        for desc_name in desc_name_list:
            # print("desc_name", desc_name)
            if desc_name == '*':
                import re
                possible_desc_prefix = ['SOAP', 'ACSF', 'LMBTR', 'FCHL19', 'CM', 'pca', 'skpca', 'umap', 'tsne']
                for pre in possible_desc_prefix:
                    if atomic_desc:
                        for key in self.frames[0].arrays.keys():
                            if re.search(pre + '.+', key):
                                new_desc_name.append(key)
                    else:
                        for key in self.frames[0].info.keys():
                            if re.search(pre + '.+', key):
                                new_desc_name.append(key)
            elif '*' in desc_name:
                import re
                if atomic_desc:
                    for key in self.frames[0].arrays.keys():
                        if re.search(desc_name.replace('*', '.+'), key):
                            new_desc_name.append(key)
                else:
                    for key in self.frames[0].info.keys():
                        if re.search(desc_name.replace('*', '.+'), key):
                            new_desc_name.append(key)
            else:
                new_desc_name.append(desc_name)
        return new_desc_name

    def get_descriptors(self, desc_name_list=[], use_atomic_desc=False, species_name=None):
        """ extract the descriptor array from each frame

        Parameters
        ----------
        desc_name_list: a list of strings
                        the name of the .info[] in the extended xyz file
        use_atomic_desc: bool
                         return the descriptors for each atom, read from the xyz file
        species_name: int
                      the atomic number of the species selected.
                      Only the desciptors of atoms of the specified specied will be returned.
                      species_name=None means all atoms are selected.

        Returns
        -------
        desc: np.matrix
        atomic_desc: np.matrix
        """
        desc = []
        atomic_desc = []

        if isinstance(desc_name_list, str):
            desc_name_list = [desc_name_list]

        desc_name_list = self._desc_name_with_wild_card(desc_name_list)
        print("Find the following descriptor names that match the specifications: ", desc_name_list)

        # load from xyz file
        try:
            # retrieve the descriptor vectors --- both of these throw a ValueError if any are missing or are of wrong shape
            desc = np.hstack(
                np.vstack([a.info[desc_name] for a in self.frames]) for desc_name in desc_name_list)
            print("Use global descriptor matrix with shape: ", np.shape(desc))
        except:
            print(desc[-1])
            print("Cannot find the specified descriptors from xyz")

        try:
            # get the atomic descriptors with the same name
            if use_atomic_desc:
                atomic_desc = self.get_atomic_descriptors(desc_name_list, species_name)
        except:
            print("Cannot find the specified atomic descriptors from xyz")
        return desc, atomic_desc

    def get_atomic_descriptors(self, desc_name_list=[], species_name=None):
        """ extract the descriptor array from each frame

        Parameters
        ----------
        desc_name_list: a list of strings
                     the name of the .info[] in the extended xyz file
        species_name: int
                        the atomic number of the species selected.
                        Only the desciptors of atoms of the specified specied will be returned.
                    species_name=None means all atoms are selected.

        Returns
        -------
        atomic_desc: np.matrix
        """

        atomic_desc = []

        if isinstance(desc_name_list, str):
            desc_name_list = [desc_name_list]

        desc_name_list = self._desc_name_with_wild_card(desc_name_list, True)
        print("Find the following atomic descriptor names that match the specifications: ", desc_name_list)

        # load from xyz file
        try:
            if species_name is None:
                atomic_desc = np.vstack(
                    np.concatenate([a.get_array(desc_name) for a in self.frames]) for desc_name in desc_name_list)
            elif species_name in self.global_species:
                atomic_desc = np.hstack(np.concatenate(
                    [self._get_atomic_descriptors_by_species(a, desc_name, species_name) for a in self.frames]) for
                                              desc_name in desc_name_list)
            else:
                raise ValueError("Cannot find the specified chemical species in the data set.")
            print("Use atomic descriptor matrix with shape: ", np.shape(atomic_desc))
        except:
            print("Cannot find the specified atomic descriptors from xyz")

        return atomic_desc

    def _get_atomic_descriptors_by_species(self, frame, desc_name, species_name=None):
        species_index = [i for i, s in enumerate(frame.get_atomic_numbers()) if s == species_name]
        return frame.get_array(desc_name)[species_index]

    def get_property(self, y_key=None, extensive=False, sbs=[]):
        """ extract specified property from selected frames

        Parameters
        ----------
        y_key: string_like
               the name of the property in the extended xyz file
        sbs: array, integer

        Returns
        -------
        y_all: array [N_samples]
        """

        if len(sbs) == 0:
            sbs = range(self.nframes)

        y_all = []
        try:
            for i in sbs:
                frame = self.frames[i]
                if y_key == 'volume' or y_key == 'Volume':
                    y_all.append(frame.get_volume() / len(frame.get_positions()))
                elif y_key == 'size' or y_key == 'Size':
                    y_all.append(len(frame.get_positions()))
                elif y_key == 'index' or y_key == 'Index' or y_key == None:
                    y_all.append(i)
                elif isinstance(y_key, int) and int(y_key) in self.global_species:
                    # count the number of atoms of this specified chemical element
                    y_all.append(frame.get_atomic_numbers().tolist().count(int(y_key)))
                elif extensive:
                    y_all.append(frame.info[y_key] / len(frame.get_positions()))
                else:
                    y_all.append(frame.info[y_key])
        except:
            try:
                for i in sbs:
                    frame = self.frames[i]
                    if extensive:
                        # use the sum of atomic properties
                        y_all.append(np.sum(frame.get_array(y_key)))
                    else:
                        # use the average of atomic properties
                        y_all.append(np.mean(frame.get_array(y_key)))
            except:
                raise ValueError('Cannot load the property vector')
        if len(np.shape(y_all)) > 1:
            raise ValueError('The property from the xyz file has more than one column')
        return np.array(y_all)

    def get_atomic_property(self, y_key=None, extensive=False, sbs=[], species_name=None):
        """ extract the property array from each atom

        Parameters
        ----------
        y_key: string_like
             the name of the property in the extended xyz file
        sbs: array, integer
        specie: int
                 the atomic number of the species selected.
                      Only the properties of atoms of the specified specied will be returned.
                      species_name=None means all atoms are selected.
        Returns
        -------
        y_all: array [N_atoms]
        """

        if len(sbs) == 0:
            sbs = range(self.nframes)

        y_all = []
        try:
            # y_all = np.concatenate([a.get_array(y_key) for a in self.frames[sbs]]) # this doesn't work ?!
            for i in sbs:
                frame = self.frames[i]
                if species_name is None:
                    y_all = np.append(y_all, frame.get_array(y_key))
                elif species_name in self.global_species:
                    y_all = np.append(y_all, self._get_atomic_descriptors_by_species(frame, y_key, species_name))
                else:
                    raise ValueError("Cannot find the specified chemical species in the data set.")
        except:
            try:
                natom_used = self.get_natom_list_by_species(species_name)
                for index, y in enumerate(self.get_property(y_key, extensive, sbs)):
                    y_all = np.append(y_all, y * np.ones(natom_used[index]))
                print("Cannot find the atomic properties, use the per-frame property instead")
            except:
                raise ValueError('Cannot load the property vector')
        if len(np.shape(y_all)) > 1:
            raise ValueError('The property from the xyz file has more than one column')
        return y_all

    def set_descriptors(self, desc=None, desc_name=None):
        """ write the descriptor array to the atom object

        Parameters
        ----------
        desc: np.matrix, shape=[n_descriptors, n_frames]

        Returns
        -------
        """
        # TODO: (maybe?) rename this into set_info. Need to change pca.py/kpca.py/... that uses this function
        # check if the length of the descriptor matrix is the same as the number of frames
        if len(desc) != self.nframes and self.nframes > 1:
            raise ValueError('The length of the descriptor matrix is not the same as the number of frames.')
        if len(desc) != self.total_natoms and self.nframes == 1:
            raise ValueError('The length of the descriptor matrix is not the same as the number of atoms.')

        if self.nframes > 1:
            for i, frame in enumerate(self.frames):
                frame.info[desc_name] = desc[i]
        else:
            self.frames[0].new_array(desc_name, desc)

    def set_atomic_descriptors(self, atomic_desc=None, atomic_desc_name=None, species_name=None):
        """ write the descriptor array to the atom object

        Parameters
        ----------
        desc: np.matrix, shape=[n_descriptors, n_atoms]

        Returns
        -------
        """
        # check if the length of the descriptor matrix is the same as the total number of atoms
        if species_name is None and len(atomic_desc) != self.total_natoms:
            raise ValueError('The length of the atomic descriptor matrix is not the same as the total number of atoms.')

        atom_index = 0

        if species_name is None:
            for i, frame in enumerate(self.frames):
                natomnow = self.natom_list[i]
                frame.new_array(atomic_desc_name, np.array(atomic_desc)[atom_index:atom_index + natomnow])
                atom_index += natomnow
        else:
            if atomic_desc.ndim == 1:
                n_desc = 1
            else:
                n_desc = np.shape(atomic_desc)[1]
            for i, frame in enumerate(self.frames):
                array_now = np.zeros((self.natom_list[i], n_desc), dtype=float)
                for j, s in enumerate(frame.get_atomic_numbers()):
                    if s == species_name:
                        array_now[j] = atomic_desc[atom_index]
                        atom_index += 1
                    else:
                        array_now[j] = np.nan
                frame.new_array(atomic_desc_name, np.array(array_now))

    def remove_descriptors(self, desc_name_list=[]):
        """
        remove the desciptors
        """
        if isinstance(desc_name_list, str):
            desc_name_list = [desc_name_list]

        desc_name_list = self._desc_name_with_wild_card(desc_name_list)
        print("removing the global descriptors from output xyz with the names: ", desc_name_list)

        for dn in desc_name_list:
            for frame in self.frames:
                if dn in frame.info:
                    del frame.info[dn]
                else:
                    pass
                    # print("Warning: Cannot parse desc_name "+str(dn)+" when remove_descriptors.")

    def remove_atomic_descriptors(self, desc_name_list=[]):
        """
        remove the desciptors
        """
        if isinstance(desc_name_list, str):
            desc_name_list = [desc_name_list]
        desc_name_list = self._desc_name_with_wild_card(desc_name_list, True)
        print("removing the atomic descriptors from output xyz with the names: ", desc_name_list)

        for dn in desc_name_list:
            for frame in self.frames:
                if dn in frame.arrays:
                    del frame.arrays[dn]
                else:
                    pass
                    # print("Warning: Cannot parse desc_name "+str(dn)+" when remove_descriptors.")

    def load_properties(self, filename, header='infer', prefix='X', **kwargs):
        """
        Load properties from a CSV file

        Read in the CSV file and save the columns to the `info` dictionary of the
        frames.

        Parameters
        ----------
        filename: str
                Name of the CSV file.
        header: int
            Row number of the header. Defaults to use the first row unless explicit
          names for the columns are given
        """
        data = pd.read_csv(filename, header=header, prefix=prefix, **kwargs)

        for frame, (_, row) in zip(self.frames, data.iterrows()):
            frame.info.update(row.to_dict())

    def write(self, filename, sbs=[], save_acronym=False, wrap_output=True):
        """
        write the selected frames or all the frames to a xyz file

        Parameters
        ----------
        filename: str
        sbs: array, integer
        """

        if len(sbs) == 0:
            sbs = range(self.nframes)

        # prepare for the output
        if os.path.isfile(str(filename) + ".xyz"):
            os.rename(str(filename) + ".xyz", "bck." + str(filename) + ".xyz")

        for i in sbs:
            if wrap_output: self.frames[i].wrap()
            self._write_computed_descriptors_to_xyz(self.global_desc[i], self.frames[i])
            self._write_computed_atomic_descriptors_to_xyz(self.atomic_desc[i], self.frames[i])
            write(str(filename) + ".xyz", self.frames[i], append=True)

        # this acronym state file lets us know how the descriptors correspond to the outputs in the xyz file
        if save_acronym:
            self.save_descriptor_acronym_state(filename)

    def write_chemiscope(self, filename, sbs=None, save_acronym=False, cutoff=None, wrap_output=True):
        """
        write the selected frames or all the frames to ChemiScope JSON

        Parameters
        ----------
        filename: str
        sbs: array, integer
        cutoff: float
                generate cutoff for atomic environments, set to None to disable atomic environments
        """

        from asaplib.io.cscope import write_chemiscope_input

        if sbs is None:
            sbs = range(self.nframes)

        # prepare for the output
        if os.path.isfile(str(filename) + ".xyz"):
            os.rename(str(filename) + ".xyz", "bck." + str(filename) + ".xyz")

        for i in sbs:
            if wrap_output: self.frames[i].wrap()
            self._write_computed_descriptors_to_xyz(self.global_desc[i], self.frames[i])
            self._write_computed_atomic_descriptors_to_xyz(self.atomic_desc[i], self.frames[i])

        # this acronym state file lets us know how the descriptors correspond to the outputs in the xyz file
        if save_acronym:
            self.save_descriptor_acronym_state(filename)

        # Disable atomic environments if there isn't any data
        write_chemiscope_input(filename + '.json.gz', self.frames, cutoff=cutoff)

    def write_descriptor_matrix(self, filename, desc_name_list, sbs=[], comment=''):
        """
        write the selected descriptor matrix in a matrix format to file

        Parameters
        ----------
        filename: str
        desc_name_list: a list of str.
                Name of the properties/descriptors to write
        sbs: array, integer
        comment: str
        """
        if len(sbs) == 0:
            sbs = range(self.nframes)

        desc, _ = self.get_descriptors(desc_name_list, False, sbs)

        if os.path.isfile(str(filename) + ".desc"):
            os.rename(str(filename) + ".desc", "bck." + str(filename) + ".desc")
        np.savetxt(str(filename) + ".desc", desc, fmt='%4.8f', header=comment)

    def write_atomic_descriptor_matrix(self, filename, desc_name, sbs=[], comment=''):
        """
        write the selected descriptor matrix in a matrix format to file

        Parameters
        ----------
        filename: str
        desc_name: str
                Name of the properties/descriptors to write
        sbs: array, integer
        comment: str
        """

        if len(sbs) == 0:
            sbs = range(self.nframes)

        _, atomic_desc = self.get_descriptors(desc_name, True, sbs)

        if os.path.isfile(str(filename) + ".atomic-desc"):
            os.rename(str(filename) + ".atomic-desc", "bck." + str(filename) + "atomic-desc")
        np.savetxt(str(filename) + ".atomic-desc", desc, fmt='%4.8f', header=comment)

    def write_computed_descriptors(self, filename, desc_dict_keys=[], sbs=[], comment=''):
        """
        write the computed descriptors for selected frames
        Parameters
        ----------
        desc_spec_keys: list
              a list (str-like) of keys for which computed descriptors to fetch.
        sbs: array, integer

        Returns
        -------
        desc: np.matrix [n_frame, n_desc]
        """
        if len(sbs) == 0:
            sbs = range(self.nframes)
        np.savetxt(str(filename) + ".desc", self.fetch_computed_descriptors(desc_dict_keys, sbs), fmt='%4.8f',
                   header=comment)

max_stars_repo_name: BingqingCheng/ASAP

max_issues_repo_name: FelixWodaczek/ASAP

max_forks_repo_name: FelixWodaczek/ASAP


text:
import numpy as np

def knn(X_train, y_train, X_test, k, dist):
    def classify_single(x):
        dists = [dist(x, i) for i in X_train]
        indices = np.argpartition(dists, k)[:k]
        return np.argmax(np.bincount(y_train[indices]))

    return [classify_single(x) for x in X_test]

max_forks_repo_path: Neighbors and wine/k-nearest neighbors/metric_classification.py


text:
# coding = utf-8

import numpy as np


text:
import matplotlib
matplotlib.use('Agg')
import os
import numpy as np
from scipy.io import wavfile as wf
import librosa
from matplotlib import pyplot as plt
import soundfile as sf

filename = 'arctic_a0001'
folder = 'data_new'
files = sorted(os.listdir(folder))
wave_array = []
for file in files:
   if filename in file and 'audio' in file:
      A = np.load(folder + '/' + file)
      for a in A:
          wave_array.append(a)

w = np.array(wave_array) * 1.0 #/ 256.0 # * 32000
print len(w), w[0:10], np.max(w)
librosa.output.write_wav('00_noquantization.wav', np.asarray(w), 16000)
sf.write('00_noquantization.wav', np.asarray(w), 16000,format='wav',subtype="PCM_16")

plt.plot(w)
plt.savefig('00_noquantization.png')

fs,A = wf.read('../voices/cmu_us_slt/wav/arctic_a0001.wav')
plt.close()
plt.plot(A)
plt.savefig('orig_wav.png')

reason: import numpy,from scipy


text:
import numpy as np

def slice_perc(lst, start_perc, end_perc):
    new_lst = lst[int(len(lst) * (start_perc / 100)): int(len(lst) * (end_perc / 100))]

    return new_lst

# Creates a float range
def frange(start, stop, step):
    i = start
    while i < stop:
        yield i
        i += step

# Originally from https://github.com/saurabhnagrecha/Pandas-to-ARFF
def pandas2arff(df, filename, wekaname="pandasdata", cleanstringdata=True, cleannan=True):
    """
    converts the pandas dataframe to a weka compatible file
    df: dataframe in pandas format
    filename: the filename you want the weka compatible file to be in
    wekaname: the name you want to give to the weka dataset (this will be visible to you when you open it in Weka)
    cleanstringdata: clean up data which may have spaces and replace with "_", special characters etc which seem to annoy Weka.
                     To suppress this, set this to False
    cleannan: replaces all nan values with "?" which is Weka's standard for missing values.
              To suppress this, set this to False
    """
    import re

    def cleanstring(s):
        if s != "?":
            return re.sub('[^A-Za-z0-9]+', "_", str(s))
        else:
            return "?"

    dfcopy = df  # all cleaning operations get done on this copy

    if cleannan != False:
        dfcopy = dfcopy.fillna(-999999999)  # this is so that we can swap this out for "?"
        # this makes sure that certain numerical columns with missing values don't get stuck with "object" type

    f = open(filename, "w")
    arffList = []
    arffList.append("@relation " + wekaname + "\n")
    # look at each column's dtype. If it's an "object", make it "nominal" under Weka for now (can be changed in source for dates.. etc)
    for i in range(df.shape[1]):
        if dfcopy.dtypes[i] == 'O' or (df.columns[i] in ["Class", "CLASS", "class"]):
            if cleannan != False:
                dfcopy.iloc[:, i] = dfcopy.iloc[:, i].replace(to_replace=-999999999, value="?")
            if cleanstringdata != False:
                dfcopy.iloc[:, i] = dfcopy.iloc[:, i].apply(cleanstring)
            _uniqueNominalVals = [str(_i) for _i in np.unique(dfcopy.iloc[:, i])]
            _uniqueNominalVals = ",".join(_uniqueNominalVals)
            _uniqueNominalVals = _uniqueNominalVals.replace("[", "")
            _uniqueNominalVals = _uniqueNominalVals.replace("]", "")
            _uniqueValuesString = "{" + _uniqueNominalVals + "}"
            arffList.append("@attribute " + df.columns[i] + _uniqueValuesString + "\n")
        else:
            arffList.append("@attribute attribute%d real\n" % i)
            # even if it is an integer, let's just deal with it as a real number for now
    arffList.append("@data\n")
    for i in dfcopy.values.tolist():

        _instanceString = ','.join(map(str, i))
        _instanceString += "\n"

        if cleannan != False:
            _instanceString = _instanceString.replace("-999999999.0", "?")  # for numeric missing values
            _instanceString = _instanceString.replace("\"?\"", "?")  # for categorical missing values
        arffList.append(_instanceString)
    f.writelines(arffList)
    f.close()
    del dfcopy
    return True

max_stars_repo_name: jorkro/wirecaml

max_issues_repo_name: pacraswill/wirecaml

max_forks_repo_name: pacraswill/wirecaml


text:
import os
from collections import deque, namedtuple

from tqdm import tqdm
import numpy as np
import torch
from torch import nn
from torch.utils.data import Dataset, DataLoader
from torch.optim import Adam
from torch.distributions import Categorical
import torch.nn.functional as F

import gym

# constants

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

# data

Memory = namedtuple('Memory', ['state', 'action', 'action_log_prob', 'reward', 'done', 'value'])
AuxMemory = namedtuple('Memory', ['state', 'target_value', 'old_values'])

class ExperienceDataset(Dataset):
    def __init__(self, data):
        super().__init__()
        self.data = data

    def __len__(self):
        return len(self.data[0])

    def __getitem__(self, ind):
        return tuple(map(lambda t: t[ind], self.data))

def create_shuffled_dataloader(data, batch_size):
    ds = ExperienceDataset(data)
    return DataLoader(ds, batch_size = batch_size, shuffle = True)

# helpers

def exists(val):
    return val is not None

def normalize(t, eps = 1e-5):
    return (t - t.mean()) / (t.std() + eps)

def update_network_(loss, optimizer):
    optimizer.zero_grad()
    loss.mean().backward()
    optimizer.step()

def init_(m):
    if isinstance(m, nn.Linear):
        gain = torch.nn.init.calculate_gain('tanh')
        torch.nn.init.orthogonal_(m.weight, gain)
        if m.bias is not None:
            torch.nn.init.zeros_(m.bias)

# networks

class Actor(nn.Module):
    def __init__(self, state_dim, hidden_dim, num_actions):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(state_dim, hidden_dim),
            nn.Tanh(),
            nn.Linear(hidden_dim, hidden_dim),
            nn.Tanh(),
            nn.Linear(hidden_dim, hidden_dim),
            nn.Tanh()
        )

        self.action_head = nn.Sequential(
            nn.Linear(hidden_dim, num_actions),
            nn.Softmax(dim=-1)
        )

        self.value_head = nn.Linear(hidden_dim, 1)
        self.apply(init_)

    def forward(self, x):
        hidden = self.net(x)
        return self.action_head(hidden), self.value_head(hidden)

class Critic(nn.Module):
    def __init__(self, state_dim, hidden_dim):
        super().__init__()
        self.net = nn.Sequential(
            nn.Linear(state_dim, hidden_dim),
            nn.Tanh(),
            nn.Linear(hidden_dim, hidden_dim),
            nn.Tanh(),
            nn.Linear(hidden_dim, 1),
        )
        self.apply(init_)

    def forward(self, x):
        return self.net(x)

# agent

def clipped_value_loss(values, rewards, old_values, clip):
    value_clipped = old_values + (values - old_values).clamp(-clip, clip)
    value_loss_1 = (value_clipped.flatten() - rewards) ** 2
    value_loss_2 = (values.flatten() - rewards) ** 2
    return torch.mean(torch.max(value_loss_1, value_loss_2))

class PPG:
    def __init__(
        self,
        state_dim,
        num_actions,
        actor_hidden_dim,
        critic_hidden_dim,
        epochs,
        epochs_aux,
        minibatch_size,
        lr,
        betas,
        lam,
        gamma,
        beta_s,
        eps_clip,
        value_clip
    ):
        self.actor = Actor(state_dim, actor_hidden_dim, num_actions).to(device)
        self.critic = Critic(state_dim, critic_hidden_dim).to(device)
        self.opt_actor = Adam(self.actor.parameters(), lr=lr, betas=betas)
        self.opt_critic = Adam(self.critic.parameters(), lr=lr, betas=betas)

        self.minibatch_size = minibatch_size

        self.epochs = epochs
        self.epochs_aux = epochs_aux

        self.lam = lam
        self.gamma = gamma
        self.beta_s = beta_s

        self.eps_clip = eps_clip
        self.value_clip = value_clip

    def save(self):
        torch.save({
            'actor': self.actor.state_dict(),
            'critic': self.critic.state_dict()
        }, f'./ppg.pt')

    def load(self):
        if not os.path.exists('./ppg.pt'):
            return

        data = torch.load(f'./ppg.pt')
        self.actor.load_state_dict(data['actor'])
        self.critic.load_state_dict(data['critic'])

    def learn(self, memories, aux_memories, next_state):
        # retrieve and prepare data from memory for training
        states = []
        actions = []
        old_log_probs = []
        rewards = []
        masks = []
        values = []

        for mem in memories:
            states.append(mem.state)
            actions.append(torch.tensor(mem.action))
            old_log_probs.append(mem.action_log_prob)
            rewards.append(mem.reward)
            masks.append(1 - float(mem.done))
            values.append(mem.value)

        # calculate generalized advantage estimate
        next_state = torch.from_numpy(next_state).to(device)
        next_value = self.critic(next_state).detach()
        values = values + [next_value]

        returns = []
        gae = 0
        for i in reversed(range(len(rewards))):
            delta = rewards[i] + self.gamma * values[i + 1] * masks[i] - values[i]
            gae = delta + self.gamma * self.lam * masks[i] * gae
            returns.insert(0, gae + values[i])

        # convert values to torch tensors
        to_torch_tensor = lambda t: torch.stack(t).to(device).detach()

        states = to_torch_tensor(states)
        actions = to_torch_tensor(actions)
        old_values = to_torch_tensor(values[:-1])
        old_log_probs = to_torch_tensor(old_log_probs)

        rewards = torch.tensor(returns).float().to(device)

        # store state and target values to auxiliary memory buffer for later training
        aux_memory = AuxMemory(states, rewards, old_values)
        aux_memories.append(aux_memory)

        # prepare dataloader for policy phase training
        dl = create_shuffled_dataloader([states, actions, old_log_probs, rewards, old_values], self.minibatch_size)

        # policy phase training, similar to original PPO
        for _ in range(self.epochs):
            for states, actions, old_log_probs, rewards, old_values in dl:
                action_probs, _ = self.actor(states)
                values = self.critic(states)
                dist = Categorical(action_probs)
                action_log_probs = dist.log_prob(actions)
                entropy = dist.entropy()

                # calculate clipped surrogate objective, classic PPO loss
                ratios = (action_log_probs - old_log_probs).exp()
                advantages = normalize(rewards - old_values.detach())
                surr1 = ratios * advantages
                surr2 = ratios.clamp(1 - self.eps_clip, 1 + self.eps_clip) * advantages
                policy_loss = - torch.min(surr1, surr2) - self.beta_s * entropy

                update_network_(policy_loss, self.opt_actor)

                # calculate value loss and update value network separate from policy network
                value_loss = clipped_value_loss(values, rewards, old_values, self.value_clip)

                update_network_(value_loss, self.opt_critic)

    def learn_aux(self, aux_memories):
        # gather states and target values into one tensor
        states = []
        rewards = []
        old_values = []
        for state, reward, old_value in aux_memories:
            states.append(state)
            rewards.append(reward)
            old_values.append(old_value)

        states = torch.cat(states)
        rewards = torch.cat(rewards)
        old_values = torch.cat(old_values)

        # get old action predictions for minimizing kl divergence and clipping respectively
        old_action_probs, _ = self.actor(states)
        old_action_probs.detach_()

        # prepared dataloader for auxiliary phase training
        dl = create_shuffled_dataloader([states, old_action_probs, rewards, old_values], self.minibatch_size)

        # the proposed auxiliary phase training
        # where the value is distilled into the policy network, while making sure the policy network does not change the action predictions (kl div loss)
        for epoch in range(self.epochs_aux):
            for states, old_action_probs, rewards, old_values in tqdm(dl, desc=f'auxiliary epoch {epoch}'):
                action_probs, policy_values = self.actor(states)
                action_logprobs = action_probs.log()

                # policy network loss copmoses of both the kl div loss as well as the auxiliary loss
                aux_loss = clipped_value_loss(policy_values, rewards, old_values, self.value_clip)
                loss_kl = F.kl_div(action_logprobs, old_action_probs, reduction='batchmean')
                policy_loss = aux_loss + loss_kl

                update_network_(policy_loss, self.opt_actor)

                # paper says it is important to train the value network extra during the auxiliary phase
                values = self.critic(states)
                value_loss = clipped_value_loss(values, rewards, old_values, self.value_clip)

                update_network_(value_loss, self.opt_critic)

# main

def main(
    env_name = 'procgen:procgen-fruitbot-v0',
    num_episodes = 50000,
    max_timesteps = 500,
    actor_hidden_dim = 32,
    critic_hidden_dim = 256,
    minibatch_size = 64,
    lr = 0.0005,
    betas = (0.9, 0.999),
    lam = 0.95,
    gamma = 0.99,
    eps_clip = 0.2,
    value_clip = 0.4,
    beta_s = .01,
    update_timesteps = 5000,
    num_policy_updates_per_aux = 32,
    epochs = 1,
    epochs_aux = 6,
    seed = None,
    render = False,
    render_every_eps = 250,
    save_every = 1000,
    load = False,
    monitor = False
):
    env = gym.make(env_name)

    if monitor:
        env = gym.wrappers.Monitor(env, './tmp/', force=True)

    state_dim = env.observation_space.shape[0]
    num_actions = env.action_space.n

    memories = deque([])
    aux_memories = deque([])

    agent = PPG(
        state_dim,
        num_actions,
        actor_hidden_dim,
        critic_hidden_dim,
        epochs,
        epochs_aux,
        minibatch_size,
        lr,
        betas,
        lam,
        gamma,
        beta_s,
        eps_clip,
        value_clip
    )

    if load:
        agent.load()

    if exists(seed):
        torch.manual_seed(seed)
        np.random.seed(seed)

    time = 0
    updated = False
    num_policy_updates = 0

    for eps in tqdm(range(num_episodes), desc='episodes'):
        render_eps = render and eps % render_every_eps == 0
        state = env.reset()
        for timestep in range(max_timesteps):
            time += 1

            if updated and render_eps:
                env.render()

            state = torch.from_numpy(state).to(device)
            action_probs, _ = agent.actor(state)
            value = agent.critic(state)

            dist = Categorical(action_probs)
            action = dist.sample()
            action_log_prob = dist.log_prob(action)
            action = action.item()

            next_state, reward, done, _ = env.step(action)

            memory = Memory(state, action, action_log_prob, reward, done, value)
            memories.append(memory)

            state = next_state

            if time % update_timesteps == 0:
                agent.learn(memories, aux_memories, next_state)
                num_policy_updates += 1
                memories.clear()

                if num_policy_updates % num_policy_updates_per_aux == 0:
                    agent.learn_aux(aux_memories)
                    aux_memories.clear()

                updated = True

            if done:
                if render_eps:
                    updated = False
                break

        if render_eps:
            env.close()

        if eps % save_every == 0:
            agent.save()

if __name__ == '__main__':
    main()