<commit_before>#!/usr/bin/env python3.5

import os
import sys
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import matplotlib.colors as colors
import matplotlib.cm as cm
from glob import glob
import musm


plt.style.use('ggplot')


def pad(array, length):
    full = np.zeros((length,) + array.shape[1:], dtype=array.dtype)
    full[:len(array)] = array
    return full


def load(path):
    data = musm.load(path)
    num_users = data['args']['num_users_per_group']
    max_iters = data['args']['max_iters']
    loss1_matrix, time_matrix = [], []
    for trace in data['traces']:
        trace = np.array(trace)
        loss1_matrix.append(pad(trace[:,:num_users], max_iters))
        time_matrix.append(pad(trace[:,-1], max_iters))
    return np.array(loss1_matrix), \
           np.array(time_matrix), \
           data['args']


def prettify(ax, max_iters, title):
    xtick = 5 if max_iters <= 50 else 10
    xticks = np.hstack([[1], np.arange(xtick, max_iters + 1, xtick)])
    ax.set_xticks(xticks)

    ax.xaxis.label.set_fontsize(18)
    ax.yaxis.label.set_fontsize(18)
    ax.grid(True)
    for line in ax.get_xgridlines() + ax.get_ygridlines():
        line.set_linestyle('-.')

    #ax.set_title(title, fontsize=18)
    legend = ax.legend(loc='upper right', fancybox=False, shadow=False)
    for label in legend.get_texts():
        label.set_fontsize('x-large')
    for label in legend.get_lines():
        label.set_linewidth(2)


_WINTER = cm.ScalarMappable(cmap=plt.get_cmap('winter'),
                            norm=colors.Normalize(vmin=0, vmax=1))
_SUMMER = cm.ScalarMappable(cmap=plt.get_cmap('summer'),
                            norm=colors.Normalize(vmin=0, vmax=1))


def get_style(args):
    transform, tau, pick = args['transform'], args['tau'], args['pick']
    lmbda = -1.0
    if transform == 'indep':
        return '#FF0000', 'indep. {}'.format(pick)
    elif transform == 'sumcov':
        return _WINTER.to_rgba(lmbda), 'k only ({}, λ={})'.format(pick, lmbda)
    elif transform == 'varsumvarcov':
        return _SUMMER.to_rgba(lmbda), 'v + k ({}, τ={})'.format(pick, tau)


def draw(args):
    loss1_fig, loss1_ax = plt.subplots(1, 1)<commit_msg>draw: update for lambda again.<commit_after>#!/usr/bin/env python3.5

import os
import sys
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import matplotlib.colors as colors
import matplotlib.cm as cm
from glob import glob
import musm

# TODO: suppress plot at --min-regret


plt.style.use('ggplot')


def pad(array, length):
    full = np.zeros((length,) + array.shape[1:], dtype=array.dtype)
    full[:len(array)] = array
    return full


def load(path):
    data = musm.load(path)
    num_users = data['args']['num_users_per_group']
    max_iters = data['args']['max_iters']
    loss1_matrix, time_matrix = [], []
    for trace in data['traces']:
        trace = np.array(trace)
        loss1_matrix.append(pad(trace[:,:num_users], max_iters))
        time_matrix.append(pad(trace[:,-1], max_iters))
    return np.array(loss1_matrix), \
           np.array(time_matrix), \
           data['args']


def prettify(ax, max_iters, title):
    xtick = 5 if max_iters <= 50 else 10
    xticks = np.hstack([[1], np.arange(xtick, max_iters + 1, xtick)])
    ax.set_xticks(xticks)

    ax.xaxis.label.set_fontsize(18)
    ax.yaxis.label.set_fontsize(18)
    ax.grid(True)
    for line in ax.get_xgridlines() + ax.get_ygridlines():
        line.set_linestyle('-.')

    #ax.set_title(title, fontsize=18)
    legend = ax.legend(loc='upper right', fancybox=False, shadow=False)
    for label in legend.get_texts():
        label.set_fontsize('x-large')
    for label in legend.get_lines():
        label.set_linewidth(2)


_WINTER = cm.ScalarMappable(cmap=plt.get_cmap('winter'),
                            norm=colors.Normalize(vmin=0, vmax=1))
_SUMMER = cm.ScalarMappable(cmap=plt.get_cmap('summer'),
                            norm=colors.Normalize(vmin=0, vmax=1))


def get_style(args):
    transform, tau, pick = args['transform'], args['tau'], args['pick']
    if transform == 'indep':
        return '#FF0000', 'indep. {}'.format(pick)
    elif transform == 'sumcov':
        lmbda = args['lmbda']
        return _WINTER.to_rgba(lmbda), 'k only ({}, λ={})'.format(pick, lmbda)
    elif transform == 'varsumvarcov':
        return _SUMMER.to_rgba(tau), 'v + k ({}, τ={})'.format(pick, tau)


def draw(args):
    loss1_fig, loss1_ax = plt.subplots(1, 1)
{'fraction_non_alphanumeric': 0.1536231884057971, 'fraction_numerical': 0.01516164994425864, 'mean_word_length': 3.2561669829222013, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '384196', 'n_tokens_mistral': 1681, 'n_tokens_neox': 1611, 'n_words': 330}