plot_result.py

# Copyright 2020 Google LLC
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Functions for plotting the results."""

from absl import app
from absl import flags

import numpy as np
import os
import matplotlib.pyplot as plt

from privateFM.utils import EasyDict, read_from_file
from privateFM.utils import get_default_delta
from privateFM.FM_simulate import make_estimate, set_k_p_eta, get_debias

FLAGS = flags.FLAGS

flags.DEFINE_float('epsilon', -1.0, 'DP epsilon. -1 for non-private.')
flags.DEFINE_float('gamma', 1.0, 'Accuracy parameter.')
flags.DEFINE_string('res_dir', './res', 'Directory to write results.')


def plot_figure(ax1, x, ys, xlabel=None, ylabel=None, title=None, xscale='log', yscale=None, ylim=None):
    """
    Args:
        x: x axis value
        ys: a dict of y values to plot. If y is 2d, plot error bar
    """
    params = {'legend.handlelength': 1,
              'grid.alpha': 1,
              'grid.color': "#cccccc",
              'grid.linestyle': '--',
              'legend.fontsize': 14,
              }
    plt.rcParams.update(params)
    markersize = 5

    x = np.array(x)
    plt.setp(ax1, xticks=x, xticklabels=['2^{}'.format(int(np.log2(xx))) for xx in x])

    shapes = 'o^+*x'
    line_styles = ['--', '-', '-.', ':']
    labels = list(ys.keys())

    for idx, label in enumerate(labels):
        y = np.array(ys[label])
        style = shapes[idx % len(shapes)] + line_styles[idx % len(line_styles)]
        if y.ndim == 1:
            ax1.plot(x, y, style, markersize=markersize, label=labels[idx])
        elif y.ndim == 2:
            ax1.plot(x, y[:,0], style, markersize=markersize, label=labels[idx])
            ax1.fill_between(x, y[:,0] - y[:,1], y[:,0] + y[:,1], alpha=0.2, linewidth=3, capstyle='round', linestyle='-')
        else:
            raise ValueError('y has {} dimensions'.format(y.ndim))

    ax1.set_title(title)
    ax1.set_xlabel(xlabel)
    ax1.set_ylabel(ylabel)
    ax1.grid(linestyle='--')

    if xscale:
        assert xscale in ['log', 'linear']
        ax1.set_xscale(xscale)
    if yscale:
        ax1.set_yscale(yscale)
    if ylim:
        ax1.set_ylim(ylim)
    return ax1


def process_estimation(estimation, k):
    """Get some statistics of the estimations."""
    estimation = np.array(estimation)
    est_mean = np.mean(estimation)
    est_std = np.std(estimation)

    error = np.abs(estimation - k)
    mre_mean = np.mean(error / k)
    mre_std = np.std(error / k)
    return ((mre_mean, mre_std), (est_mean, est_std))


def process_data(data, k, config, factor=0):
    """Read data, make estimation and return statistics of estimations."""
    Is = data['I'].values()
    assert len(Is) == 100
    param = EasyDict({'config': config, 'k_p': set_k_p_eta(config)[0], 'factor':factor})
    est_quantile = [make_estimate(I, 'quantile', param) for I in Is]
    est_geom = [make_estimate(I, 'mean_geom', param) for I in Is]
    est_harmo = [make_estimate(I, 'mean_harmo', param) for I in Is]
    return (process_estimation(est, k) for est in [est_quantile, est_geom, est_harmo])


def main(unused_argv):
    epsilon = FLAGS.epsilon
    gamma = FLAGS.gamma
    ks = [2**i for i in range(12, 21)]
    ms = [1024, 4096, 32768]
    quantile_factor = 1/12
    morePhantom = False

    est_qs, est_gs, est_hs = {}, {}, {}
    mre_qs, mre_gs, mre_hs = {}, {}, {}
    for m in ms:
        label = 'm={}'.format(m)
        est_qs[label], est_gs[label], est_hs[label] = [], [], []
        mre_qs[label], mre_gs[label], mre_hs[label] = [], [], []
        for k in ks:
            config = EasyDict(m=m, gamma=gamma, epsilon=epsilon, delta=get_default_delta(k, epsilon), morePhantom=morePhantom)
            data = read_from_file(k, config, FLAGS.res_dir)

            (mre_q, est_q), (mre_g, est_g), (mre_h, est_h) = process_data(data, k, config, quantile_factor)
            est_qs[label].append(est_q)
            est_gs[label].append(est_g)
            est_hs[label].append(est_h)
            mre_qs[label].append(mre_q)
            mre_gs[label].append(mre_g)
            mre_hs[label].append(mre_h)

    ms = np.array(ms)
    figsize=(12, 3)
    ylim = (0, 0.07)

    plt.style.use('seaborn-colorblind')
    fig, axs = plt.subplots(nrows=1, ncols=3, constrained_layout=True, figsize=figsize)
    for i, (ests, mres) in enumerate([(est_qs, mre_qs), (est_gs, mre_gs), (est_hs, mre_hs)]):
        ylabel = r'mean relative error' if i == 0 else None
        axs[i] = plot_figure(axs[i], ks, mres, xlabel=r'$F_0$', ylabel=ylabel, ylim=ylim)
        plt.setp(axs[i], xticks=ks, xticklabels=['$2^{'+str(int(np.log2(xx)))+'}$' for xx in ks])

    handles, labels = axs[-1].get_legend_handles_labels()
    fig.legend(handles, labels, loc='upper center', ncol=ms.size, framealpha=1)  # put legend in the top center

    fig_name = os.path.join(FLAGS.res_dir, 'mre_m_eps{}_gamma{}{}'.format(float(epsilon), gamma, ('_morePhantom' if morePhantom else ''))) + '.pdf'
    print(fig_name)
    fig.savefig(fig_name, bbox_inches='tight')


if __name__ == '__main__':
    app.run(main)