def load_experiments(self):
        """
            Load a specific human experiment.
        """

        fit_mixture_model = self.parameters.get("fit_mixturemodel", False)

        # Load each experimental dataset
        for experiment_id in self.experiment_ids:
            if experiment_id == "bays09":
                self.experimental_datasets[experiment_id] = load_experimental_data.load_data_bays09(
                    data_dir=self.data_dir, fit_mixture_model=fit_mixture_model
                )

                if self.debug:
                    print "Loading Bays09 dataset"

            elif experiment_id == "dualrecall":
                self.experimental_datasets[experiment_id] = load_experimental_data.load_data_dualrecall(
                    data_dir=self.data_dir
                )

                if self.debug:
                    print "Loading double Recall dataset"

            elif experiment_id == "gorgo11":
                self.experimental_datasets[experiment_id] = load_experimental_data.load_data_simult(
                    data_dir=self.data_dir, fit_mixture_model=fit_mixture_model
                )

                if self.debug:
                    print "Loading Gorgo11 simult dataset"
def plots_memory_curves(data_pbs, generator_module=None):
    '''
        Reload and plot memory curve of a feature code.
        Can use Marginal Fisher Information and fitted Mixture Model as well
    '''

    #### SETUP
    #
    savefigs = True
    savedata = True

    plot_pcolor_fit_precision_to_fisherinfo = True
    plot_selected_memory_curves = True

    colormap = None  # or 'cubehelix'
    plt.rcParams['font.size'] = 16
    #
    #### /SETUP

    print "Order parameters: ", generator_module.dict_parameters_range.keys()

    result_all_precisions_mean = utils.nanmean(np.squeeze(data_pbs.dict_arrays['result_all_precisions']['results']), axis=-1)
    result_all_precisions_std = utils.nanstd(np.squeeze(data_pbs.dict_arrays['result_all_precisions']['results']), axis=-1)
    result_em_fits_mean = utils.nanmean(np.squeeze(data_pbs.dict_arrays['result_em_fits']['results']), axis=-1)
    result_em_fits_std = utils.nanstd(np.squeeze(data_pbs.dict_arrays['result_em_fits']['results']), axis=-1)
    result_marginal_inv_fi_mean = utils.nanmean(np.squeeze(data_pbs.dict_arrays['result_marginal_inv_fi']['results']), axis=-1)
    result_marginal_inv_fi_std = utils.nanstd(np.squeeze(data_pbs.dict_arrays['result_marginal_inv_fi']['results']), axis=-1)
    result_marginal_fi_mean = utils.nanmean(1./np.squeeze(data_pbs.dict_arrays['result_marginal_inv_fi']['results']), axis=-1)
    result_marginal_fi_std = utils.nanstd(1./np.squeeze(data_pbs.dict_arrays['result_marginal_inv_fi']['results']), axis=-1)

    M_space = data_pbs.loaded_data['parameters_uniques']['M']
    sigmax_space = data_pbs.loaded_data['parameters_uniques']['sigmax']
    T_space = data_pbs.loaded_data['datasets_list'][0]['T_space']

    print M_space
    print sigmax_space
    print T_space
    print result_all_precisions_mean.shape, result_em_fits_mean.shape, result_marginal_inv_fi_mean.shape

    dataio = DataIO.DataIO(output_folder=generator_module.pbs_submission_infos['simul_out_dir'] + '/outputs/', label='global_' + dataset_infos['save_output_filename'])

    ## Load Experimental data
    data_simult = load_experimental_data.load_data_simult(data_dir=os.path.normpath(os.path.join(os.path.split(load_experimental_data.__file__)[0], '../../experimental_data/')))
    memory_experimental = data_simult['precision_nitems_theo']

    # Compute some landscapes of fit!
    dist_diff_precision_margfi = np.sum(np.abs(result_all_precisions_mean*2. - result_marginal_fi_mean[..., 0])**2., axis=-1)
    dist_diff_precision_margfi_1item = np.abs(result_all_precisions_mean[..., 0]*2. - result_marginal_fi_mean[..., 0, 0])**2.
    dist_diff_all_precision_margfi = np.abs(result_all_precisions_mean*2. - result_marginal_fi_mean[..., 0])**2.

    dist_ratio_precision_margfi = np.sum(np.abs((result_all_precisions_mean*2.)/result_marginal_fi_mean[..., 0] - 1.0)**2., axis=-1)

    dist_diff_emkappa_margfi = np.sum(np.abs(result_em_fits_mean[..., 0]*2. - result_marginal_fi_mean[..., 0])**2., axis=-1)
    dist_ratio_emkappa_margfi = np.sum(np.abs((result_em_fits_mean[..., 0]*2.)/result_marginal_fi_mean[..., 0] - 1.0)**2., axis=-1)

    dist_diff_precision_experim = np.sum(np.abs(result_all_precisions_mean - memory_experimental)**2., axis=-1)
    dist_diff_emkappa_experim = np.sum(np.abs(result_em_fits_mean[..., 0] - memory_experimental)**2., axis=-1)


    if plot_pcolor_fit_precision_to_fisherinfo:
        # Check fit between precision and fisher info
        utils.pcolor_2d_data(dist_diff_precision_margfi, log_scale=True, x=M_space, y=sigmax_space, xlabel='M', ylabel='sigmax')

        if savefigs:
           dataio.save_current_figure('match_precision_margfi_log_pcolor_{label}_{unique_id}.pdf')

        # utils.pcolor_2d_data(dist_diff_precision_margfi, x=M_space, y=sigmax_space[2:], xlabel='M', ylabel='sigmax')
        # if savefigs:
        #    dataio.save_current_figure('match_precision_margfi_pcolor_{label}_{unique_id}.pdf')

        utils.pcolor_2d_data(dist_ratio_precision_margfi[4:], x=M_space[4:], y=sigmax_space, xlabel='M', ylabel='sigmax', log_scale=True)

        utils.pcolor_2d_data(dist_diff_emkappa_margfi, x=M_space, y=sigmax_space, xlabel='M', ylabel='sigmax', log_scale=True)

        utils.pcolor_2d_data(dist_ratio_emkappa_margfi[4:], x=M_space[4:], y=sigmax_space, xlabel='M', ylabel='sigmax', log_scale=True)

        utils.pcolor_2d_data(dist_diff_precision_experim, x=M_space, y=sigmax_space, xlabel='M', ylabel='sigmax', log_scale=True)

        utils.pcolor_2d_data(dist_diff_emkappa_experim, x=M_space, y=sigmax_space, xlabel='M', ylabel='sigmax', log_scale=True)

        utils.pcolor_2d_data(dist_diff_precision_margfi*dist_diff_emkappa_margfi*dist_diff_precision_experim*dist_diff_emkappa_experim, x=M_space, y=sigmax_space, xlabel='M', ylabel='sigmax', log_scale=True)

        utils.pcolor_2d_data(dist_diff_precision_margfi_1item, log_scale=True, x=M_space, y=sigmax_space, xlabel='M', ylabel='sigmax')


    if plot_selected_memory_curves:
        selected_values = [[20, 0.19], [40, 0.15], [100, 0.44], [140, 0.15], [40, 0.17], [60, 0.50]]

        for current_values in selected_values:
            # Find the indices
            M_i         = np.argmin(np.abs(current_values[0] - M_space))
            sigmax_i    = np.argmin(np.abs(current_values[1] - sigmax_space))

            ax = utils.plot_mean_std_area(T_space, memory_experimental, np.zeros(T_space.size), linewidth=3, fmt='o-', markersize=8)

            ax = utils.plot_mean_std_area(T_space, result_all_precisions_mean[M_i, sigmax_i], result_all_precisions_std[M_i, sigmax_i], ax_handle=ax, linewidth=3, fmt='o-', markersize=8)

            ax = utils.plot_mean_std_area(T_space, 0.5*result_marginal_fi_mean[..., 0][M_i, sigmax_i], 0.5*result_marginal_fi_std[..., 0][M_i, sigmax_i], ax_handle=ax, linewidth=3, fmt='o-', markersize=8)

            ax = utils.plot_mean_std_area(T_space, result_em_fits_mean[..., 0][M_i, sigmax_i], result_em_fits_std[..., 0][M_i, sigmax_i], ax_handle=ax, xlabel='Number of items', ylabel='Kappa/Inverse variance', linewidth=3, fmt='o-', markersize=8)

            ax.set_title('M %d, sigmax %.2f' % (M_space[M_i], sigmax_space[sigmax_i]))
            plt.legend(['Experimental data', 'Precision of samples', 'Marginal Fisher Information', 'Fitted kappa'])

            ax.set_xlim([0.9, 5.1])
            ax.set_xticks(range(1, 6))
            ax.set_xticklabels(range(1, 6))

            if savefigs:
                dataio.save_current_figure('memorycurves_M%dsigmax%.2f_{label}_{unique_id}.pdf' % (M_space[M_i], sigmax_space[sigmax_i]))



    all_args = data_pbs.loaded_data['args_list']
    variables_to_save = ['result_all_precisions_mean', 'result_em_fits_mean', 'result_marginal_inv_fi_mean', 'result_all_precisions_std', 'result_em_fits_std', 'result_marginal_inv_fi_std', 'result_marginal_fi_mean', 'result_marginal_fi_std', 'M_space', 'sigmax_space', 'T_space', 'all_args']

    if savedata:
        dataio.save_variables(variables_to_save, locals())


    plt.show()

    return locals()
def plots_memory_curves(data_pbs, generator_module=None):
    '''
        Reload and plot memory curve of a Mixed code.
        Can use Marginal Fisher Information and fitted Mixture Model as well
    '''

    #### SETUP
    #
    savefigs = True
    savedata = True

    plot_pcolor_fit_precision_to_fisherinfo = True
    plot_selected_memory_curves = False
    plot_best_memory_curves = True

    colormap = None  # or 'cubehelix'
    plt.rcParams['font.size'] = 16
    #
    #### /SETUP

    print "Order parameters: ", generator_module.dict_parameters_range.keys()

    result_all_precisions_mean = utils.nanmean(np.squeeze(data_pbs.dict_arrays['result_all_precisions']['results']), axis=-1)
    result_all_precisions_std = utils.nanstd(np.squeeze(data_pbs.dict_arrays['result_all_precisions']['results']), axis=-1)
    result_em_fits_mean = utils.nanmean(np.squeeze(data_pbs.dict_arrays['result_em_fits']['results']), axis=-1)
    result_em_fits_std = utils.nanstd(np.squeeze(data_pbs.dict_arrays['result_em_fits']['results']), axis=-1)

    ratiohier_space = data_pbs.loaded_data['parameters_uniques']['ratio_hierarchical']
    sigmax_space = data_pbs.loaded_data['parameters_uniques']['sigmax']
    T_space = data_pbs.loaded_data['datasets_list'][0]['T_space']

    print ratiohier_space
    print sigmax_space
    print T_space
    print result_all_precisions_mean.shape, result_em_fits_mean.shape

    dataio = DataIO.DataIO(output_folder=generator_module.pbs_submission_infos['simul_out_dir'] + '/outputs/', label='global_' + dataset_infos['save_output_filename'])

    ## Load Experimental data
    data_simult = load_experimental_data.load_data_simult(data_dir=os.path.normpath(os.path.join(os.path.split(load_experimental_data.__file__)[0], '../../experimental_data/')), fit_mixture_model=True)
    gorgo11_experimental_precision = data_simult['precision_nitems_theo']
    gorgo11_experimental_kappa = np.array([data['kappa'] for _, data in data_simult['em_fits_nitems']['mean'].items()])
    gorgo11_experimental_kappa_std = np.array([data['kappa'] for _, data in data_simult['em_fits_nitems']['std'].items()])
    gorgo11_experimental_emfits_mean = np.array([[data[key] for _, data in data_simult['em_fits_nitems']['mean'].items()] for key in ['kappa', 'mixt_target', 'mixt_nontargets', 'mixt_random']])
    gorgo11_experimental_emfits_std = np.array([[data[key] for _, data in data_simult['em_fits_nitems']['std'].items()] for key in ['kappa', 'mixt_target', 'mixt_nontargets', 'mixt_random']])
    gorgo11_experimental_emfits_sem = gorgo11_experimental_emfits_std/np.sqrt(np.unique(data_simult['subject']).size)
    gorgo11_T_space = data_simult['data_to_fit']['n_items']

    data_bays2009 = load_experimental_data.load_data_bays09(fit_mixture_model=True)
    bays09_experimental_mixtures_mean = data_bays2009['em_fits_nitems_arrays']['mean']
    bays09_experimental_mixtures_std = data_bays2009['em_fits_nitems_arrays']['std']
    # add interpolated points for 3 and 5 items
    emfit_mean_intpfct = spint.interp1d(np.unique(data_bays2009['n_items']), bays09_experimental_mixtures_mean)
    bays09_experimental_mixtures_mean_compatible = emfit_mean_intpfct(np.arange(1, 7))
    emfit_std_intpfct = spint.interp1d(np.unique(data_bays2009['n_items']), bays09_experimental_mixtures_std)
    bays09_experimental_mixtures_std_compatible = emfit_std_intpfct(np.arange(1, 7))
    bays09_T_space_interp = np.arange(1, 7)
    bays09_T_space = data_bays2009['data_to_fit']['n_items']



    # Compute some landscapes of fit!
    dist_diff_precision_experim = np.sum(np.abs(result_all_precisions_mean[..., :gorgo11_experimental_precision.size] - gorgo11_experimental_precision)**2., axis=-1)
    dist_diff_emkappa_experim = np.sum(np.abs(result_em_fits_mean[..., :gorgo11_experimental_precision.size, 0] - gorgo11_experimental_kappa)**2., axis=-1)

    dist_diff_precision_experim_1item = np.abs(result_all_precisions_mean[..., 0] - gorgo11_experimental_precision[0])**2.
    dist_diff_precision_experim_2item = np.abs(result_all_precisions_mean[..., 1] - gorgo11_experimental_precision[1])**2.
    dist_diff_emkappa_experim_1item = np.abs(result_em_fits_mean[..., 0, 0] - gorgo11_experimental_kappa[0])**2.

    dist_diff_em_mixtures_bays09 = np.sum(np.sum((result_em_fits_mean[..., 1:4] - bays09_experimental_mixtures_mean_compatible[1:].T)**2., axis=-1), axis=-1)
    dist_diff_modelfits_experfits_bays09 = np.sum(np.sum((result_em_fits_mean[..., :4] - bays09_experimental_mixtures_mean_compatible.T)**2., axis=-1), axis=-1)

    if plot_pcolor_fit_precision_to_fisherinfo:
        utils.pcolor_2d_data(dist_diff_precision_experim, x=ratiohier_space, y=sigmax_space, xlabel='ratio hier', ylabel='sigmax', log_scale=True)
        if savefigs:
            dataio.save_current_figure('match_diff_precision_experim_log_pcolor_{label}_{unique_id}.pdf')


        utils.pcolor_2d_data(dist_diff_emkappa_experim, x=ratiohier_space, y=sigmax_space, xlabel='ratio hier', ylabel='sigmax', log_scale=True)
        if savefigs:
            dataio.save_current_figure('match_diff_emkappa_experim_log_pcolor_{label}_{unique_id}.pdf')


        utils.pcolor_2d_data(dist_diff_precision_experim*dist_diff_emkappa_experim, x=ratiohier_space, y=sigmax_space, xlabel='ratio hier', ylabel='sigmax', log_scale=True)
        if savefigs:
            dataio.save_current_figure('match_bigmultiplication_log_pcolor_{label}_{unique_id}.pdf')

        utils.pcolor_2d_data(dist_diff_precision_experim_1item, log_scale=True, x=ratiohier_space, y=sigmax_space, xlabel='ratio hier', ylabel='sigmax')
        if savefigs:
            dataio.save_current_figure('match_diff_precision_experim_1item_log_pcolor_{label}_{unique_id}.pdf')


        utils.pcolor_2d_data(dist_diff_emkappa_experim_1item, log_scale=True, x=ratiohier_space, y=sigmax_space, xlabel='ratio hier', ylabel='sigmax')
        if savefigs:
            dataio.save_current_figure('match_diff_emkappa_experim_1item_log_pcolor_{label}_{unique_id}.pdf')

        utils.pcolor_2d_data(dist_diff_precision_experim_2item, log_scale=True, x=ratiohier_space, y=sigmax_space, xlabel='ratio hier', ylabel='sigmax')
        if savefigs:
            dataio.save_current_figure('match_diff_precision_experim_2item_log_pcolor_{label}_{unique_id}.pdf')

        utils.pcolor_2d_data(dist_diff_em_mixtures_bays09, log_scale=True, x=ratiohier_space, y=sigmax_space, xlabel='ratio hier', ylabel='sigmax')
        if savefigs:
            dataio.save_current_figure('match_diff_mixtures_experbays09_pcolor_{label}_{unique_id}.pdf')

        utils.pcolor_2d_data(dist_diff_modelfits_experfits_bays09, log_scale=True, x=ratiohier_space, y=sigmax_space, xlabel='ratio hier', ylabel='sigmax')
        if savefigs:
            dataio.save_current_figure('match_diff_emfits_experbays09_pcolor_{label}_{unique_id}.pdf')


    # Macro plot
    def mem_plot_precision(sigmax_i, ratiohier_i, mem_exp_prec):
        ax = utils.plot_mean_std_area(T_space[:mem_exp_prec.size], mem_exp_prec, np.zeros(mem_exp_prec.size), linewidth=3, fmt='o-', markersize=8, label='Experimental data')

        ax = utils.plot_mean_std_area(T_space[:mem_exp_prec.size], result_all_precisions_mean[ratiohier_i, sigmax_i, :mem_exp_prec.size], result_all_precisions_std[ratiohier_i, sigmax_i, :mem_exp_prec.size], ax_handle=ax, linewidth=3, fmt='o-', markersize=8, label='Precision of samples')

        # ax = utils.plot_mean_std_area(T_space, 0.5*result_marginal_fi_mean[..., 0][ratiohier_i, sigmax_i], 0.5*result_marginal_fi_std[..., 0][ratiohier_i, sigmax_i], ax_handle=ax, linewidth=3, fmt='o-', markersize=8, label='Marginal Fisher Information')

        # ax = utils.plot_mean_std_area(T_space, result_em_fits_mean[..., 0][ratiohier_i, sigmax_i], result_em_fits_std[..., 0][ratiohier_i, sigmax_i], ax_handle=ax, xlabel='Number of items', ylabel="Inverse variance $[rad^{-2}]$", linewidth=3, fmt='o-', markersize=8, label='Fitted kappa')

        ax.set_title('ratio_hier %.2f, sigmax %.2f' % (ratiohier_space[ratiohier_i], sigmax_space[sigmax_i]))
        ax.legend()
        ax.set_xlim([0.9, mem_exp_prec.size + 0.1])
        ax.set_xticks(range(1, mem_exp_prec.size + 1))
        ax.set_xticklabels(range(1, mem_exp_prec.size + 1))

        if savefigs:
            dataio.save_current_figure('memorycurves_precision_ratiohier%.2fsigmax%.2f_{label}_{unique_id}.pdf' % (ratiohier_space[ratiohier_i], sigmax_space[sigmax_i]))

    def mem_plot_kappa(sigmax_i, ratiohier_i, T_space_exp, exp_kappa_mean, exp_kappa_std=None):
        ax = utils.plot_mean_std_area(T_space_exp, exp_kappa_mean, exp_kappa_std, linewidth=3, fmt='o-', markersize=8, label='Experimental data')

        ax = utils.plot_mean_std_area(T_space[:T_space_exp.max()], result_em_fits_mean[..., :T_space_exp.max(), 0][ratiohier_i, sigmax_i], result_em_fits_std[..., :T_space_exp.max(), 0][ratiohier_i, sigmax_i], xlabel='Number of items', ylabel="Memory error $[rad^{-2}]$", linewidth=3, fmt='o-', markersize=8, label='Fitted kappa', ax_handle=ax)

        ax.set_title('ratio_hier %.2f, sigmax %.2f' % (ratiohier_space[ratiohier_i], sigmax_space[sigmax_i]))
        ax.legend()
        ax.set_xlim([0.9, T_space_exp.max()+0.1])
        ax.set_xticks(range(1, T_space_exp.max()+1))
        ax.set_xticklabels(range(1, T_space_exp.max()+1))

        ax.get_figure().canvas.draw()

        if savefigs:
            dataio.save_current_figure('memorycurves_kappa_ratiohier%.2fsigmax%.2f_{label}_{unique_id}.pdf' % (ratiohier_space[ratiohier_i], sigmax_space[sigmax_i]))

    # def em_plot(sigmax_i, ratiohier_i):
    #     # TODO finish checking this up.
    #     f, ax = plt.subplots()
    #     ax2 = ax.twinx()

    #     # left axis, kappa
    #     ax = utils.plot_mean_std_area(T_space, result_em_fits_mean[..., 0][ratiohier_i, sigmax_i], result_em_fits_std[..., 0][ratiohier_i, sigmax_i], xlabel='Number of items', ylabel="Inverse variance $[rad^{-2}]$", ax_handle=ax, linewidth=3, fmt='o-', markersize=8, label='Fitted kappa', color='k')

    #     # Right axis, mixture probabilities
    #     utils.plot_mean_std_area(T_space, result_em_fits_mean[..., 1][ratiohier_i, sigmax_i], result_em_fits_std[..., 1][ratiohier_i, sigmax_i], xlabel='Number of items', ylabel="Mixture probabilities", ax_handle=ax2, linewidth=3, fmt='o-', markersize=8, label='Target')
    #     utils.plot_mean_std_area(T_space, result_em_fits_mean[..., 2][ratiohier_i, sigmax_i], result_em_fits_std[..., 2][ratiohier_i, sigmax_i], xlabel='Number of items', ylabel="Mixture probabilities", ax_handle=ax2, linewidth=3, fmt='o-', markersize=8, label='Nontarget')
    #     utils.plot_mean_std_area(T_space, result_em_fits_mean[..., 3][ratiohier_i, sigmax_i], result_em_fits_std[..., 3][ratiohier_i, sigmax_i], xlabel='Number of items', ylabel="Mixture probabilities", ax_handle=ax2, linewidth=3, fmt='o-', markersize=8, label='Random')

    #     lines, labels = ax.get_legend_handles_labels()
    #     lines2, labels2 = ax2.get_legend_handles_labels()
    #     ax.legend(lines + lines2, labels + labels2)

    #     ax.set_title('ratio_hier %.2f, sigmax %.2f' % (ratiohier_space[ratiohier_i], sigmax_space[sigmax_i]))
    #     ax.set_xlim([0.9, 5.1])
    #     ax.set_xticks(range(1, 6))
    #     ax.set_xticklabels(range(1, 6))

    #     f.canvas.draw()

    #     if savefigs:
    #         dataio.save_current_figure('memorycurves_emfits_ratiohier%.2fsigmax%.2f_{label}_{unique_id}.pdf' % (ratiohier_space[ratiohier_i], sigmax_space[sigmax_i]))

    def em_plot_paper(sigmax_i, ratiohier_i):
        f, ax = plt.subplots()

        # mixture probabilities
        utils.plot_mean_std_area(bays09_T_space_interp, result_em_fits_mean[..., 1][ratiohier_i, sigmax_i][:bays09_T_space_interp.size], result_em_fits_std[..., 1][ratiohier_i, sigmax_i][:bays09_T_space_interp.size], xlabel='Number of items', ylabel="Mixture probabilities", ax_handle=ax, linewidth=3, fmt='o-', markersize=5, label='Target')
        utils.plot_mean_std_area(bays09_T_space_interp, result_em_fits_mean[..., 2][ratiohier_i, sigmax_i][:bays09_T_space_interp.size], result_em_fits_std[..., 2][ratiohier_i, sigmax_i][:bays09_T_space_interp.size], xlabel='Number of items', ylabel="Mixture probabilities", ax_handle=ax, linewidth=3, fmt='o-', markersize=5, label='Nontarget')
        utils.plot_mean_std_area(bays09_T_space_interp, result_em_fits_mean[..., 3][ratiohier_i, sigmax_i][:bays09_T_space_interp.size], result_em_fits_std[..., 3][ratiohier_i, sigmax_i][:bays09_T_space_interp.size], xlabel='Number of items', ylabel="Mixture probabilities", ax_handle=ax, linewidth=3, fmt='o-', markersize=5, label='Random')

        ax.legend(prop={'size':15})

        ax.set_title('ratio_hier %.2f, sigmax %.2f' % (ratiohier_space[ratiohier_i], sigmax_space[sigmax_i]))
        ax.set_xlim([1.0, bays09_T_space_interp.size])
        ax.set_ylim([0.0, 1.1])
        ax.set_xticks(range(1, bays09_T_space_interp.size+1))
        ax.set_xticklabels(range(1, bays09_T_space_interp.size+1))

        f.canvas.draw()

        if savefigs:
            dataio.save_current_figure('memorycurves_emfits_paper_ratiohier%.2fsigmax%.2f_{label}_{unique_id}.pdf' % (ratiohier_space[ratiohier_i], sigmax_space[sigmax_i]))

    if plot_selected_memory_curves:
        selected_values = [[0.84, 0.23], [0.84, 0.19]]

        for current_values in selected_values:
            # Find the indices
            ratiohier_i     = np.argmin(np.abs(current_values[0] - ratiohier_space))
            sigmax_i        = np.argmin(np.abs(current_values[1] - sigmax_space))


            # mem_plot_precision(sigmax_i, ratiohier_i)
            # mem_plot_kappa(sigmax_i, ratiohier_i)

    if plot_best_memory_curves:

        # Best precision fit
        best_axis2_i_all = np.argmin(dist_diff_precision_experim, axis=1)
        for axis1_i, best_axis2_i in enumerate(best_axis2_i_all):
            mem_plot_precision(best_axis2_i, axis1_i, gorgo11_experimental_precision)

        # Best kappa fit
        best_axis2_i_all = np.argmin(dist_diff_emkappa_experim, axis=1)
        for axis1_i, best_axis2_i in enumerate(best_axis2_i_all):
            mem_plot_kappa(best_axis2_i, axis1_i, gorgo11_T_space, gorgo11_experimental_emfits_mean[0], gorgo11_experimental_emfits_std[0])
            # em_plot(best_axis2_i, axis1_i)

        # Best em parameters fit to Bays09
        best_axis2_i_all = np.argmin(dist_diff_modelfits_experfits_bays09, axis=1)
        # best_axis2_i_all = np.argmin(dist_diff_em_mixtures_bays09, axis=1)

        for axis1_i, best_axis2_i in enumerate(best_axis2_i_all):
            mem_plot_kappa(best_axis2_i, axis1_i, bays09_T_space, bays09_experimental_mixtures_mean[0], bays09_experimental_mixtures_std[0])
            # em_plot(best_axis2_i, axis1_i)
            em_plot_paper(best_axis2_i, axis1_i)

    all_args = data_pbs.loaded_data['args_list']
    variables_to_save = ['gorgo11_experimental_precision', 'gorgo11_experimental_kappa', 'bays09_experimental_mixtures_mean_compatible']

    if savedata:
        dataio.save_variables_default(locals(), additional_variables=variables_to_save)

        dataio.make_link_output_to_dropbox(dropbox_current_experiment_folder='memory_curves')

    plt.show()

    return locals()
def launcher_do_fit_mixturemodels(args):
    '''
        Run the model for 1..T items, computing:
        - Precision of samples
        - EM mixture model fits
        - Theoretical Fisher Information
        - EM Mixture model distances to set of currently working datasets.
    '''

    print "Doing a piece of work for launcher_do_fit_mixturemodels"

    all_parameters = utils.argparse_2_dict(args)
    print all_parameters

    if all_parameters['burn_samples'] + all_parameters['num_samples'] < 200:
        print "WARNING> you do not have enough samples I think!", all_parameters['burn_samples'] + all_parameters['num_samples']

    # Create DataIO
    #  (complete label with current variable state)
    dataio = DataIO.DataIO(output_folder=all_parameters['output_directory'], label=all_parameters['label'].format(**all_parameters))
    save_every = 1
    run_counter = 0

    # Load datasets to compare against
    data_bays2009 = load_experimental_data.load_data_bays09(data_dir=all_parameters['experiment_data_dir'], fit_mixture_model=True)
    bays09_experimental_mixtures_mean = data_bays2009['em_fits_nitems_arrays']['mean']
    # Assume that T_space >= max(T_space_bays09)
    bays09_T_space = np.unique(data_bays2009['n_items'])

    data_gorgo11 = load_experimental_data.load_data_simult(data_dir=all_parameters['experiment_data_dir'], fit_mixture_model=True)
    gorgo11_experimental_emfits_mean = data_gorgo11['em_fits_nitems_arrays']['mean']
    gorgo11_T_space = np.unique(data_gorgo11['n_items'])

    # Parameters to vary
    T_max = all_parameters['T']
    T_space = np.arange(1, T_max+1)
    repetitions_axis = -1

    # Result arrays
    result_all_precisions = np.nan*np.empty((T_space.size, all_parameters['num_repetitions']))
    # result_fi_theo = np.nan*np.empty((T_space.size, all_parameters['num_repetitions']))
    # result_fi_theocov = np.nan*np.empty((T_space.size, all_parameters['num_repetitions']))
    result_em_fits = np.nan*np.empty((T_space.size, 6, all_parameters['num_repetitions']))  # kappa, mixt_target, mixt_nontarget, mixt_random, ll, bic
    # result_em_fits_allnontargets = np.nan*np.empty((T_space.size, 5+(T_max-1), all_parameters['num_repetitions']))  # kappa, mixt_target, mixt_nontarget (T-1), mixt_random, ll, bic
    result_dist_bays09 = np.nan*np.empty((T_space.size, 4, all_parameters['num_repetitions']))  # kappa, mixt_target, mixt_nontarget, mixt_random
    result_dist_gorgo11 = np.nan*np.empty((T_space.size, 4, all_parameters['num_repetitions']))  # kappa, mixt_target, mixt_nontarget, mixt_random
    result_dist_bays09_emmixt_KL = np.nan*np.empty((T_space.size, all_parameters['num_repetitions']))
    result_dist_gorgo11_emmixt_KL = np.nan*np.empty((T_space.size, all_parameters['num_repetitions']))

    # If desired, will automatically save all Model responses.
    if all_parameters['collect_responses']:
        print "--- Collecting all responses..."
        result_responses = np.nan*np.ones((T_space.size, all_parameters['N'], all_parameters['num_repetitions']))
        result_target = np.nan*np.ones((T_space.size, all_parameters['N'], all_parameters['num_repetitions']))
        result_nontargets = np.nan*np.ones((T_space.size, all_parameters['N'], T_max-1, all_parameters['num_repetitions']))

    search_progress = progress.Progress(T_space.size*all_parameters['num_repetitions'])

    for repet_i in xrange(all_parameters['num_repetitions']):
        for T_i, T in enumerate(T_space):
            print "%.2f%%, %s left - %s" % (search_progress.percentage(), search_progress.time_remaining_str(), search_progress.eta_str())

            print "Fit for T=%d, %d/%d" % (T, repet_i+1, all_parameters['num_repetitions'])

            # Update parameter
            all_parameters['T'] = T

            ### WORK WORK WORK work? ###
            # Instantiate
            (_, _, _, sampler) = launchers.init_everything(all_parameters)

            # Sample
            sampler.run_inference(all_parameters)

            # Compute precision
            print "get precision..."
            result_all_precisions[T_i, repet_i] = sampler.get_precision()

            # Fit mixture model
            print "fit mixture model..."
            curr_params_fit = sampler.fit_mixture_model(use_all_targets=False)
            # curr_params_fit['mixt_nontargets_sum'] = np.sum(curr_params_fit['mixt_nontargets'])
            result_em_fits[T_i, :, repet_i] = [curr_params_fit[key] for key in ['kappa', 'mixt_target', 'mixt_nontargets_sum', 'mixt_random', 'train_LL', 'bic']]
            # result_em_fits_allnontargets[T_i, :2, repet_i] = [curr_params_fit['kappa'], curr_params_fit['mixt_target']]
            # result_em_fits_allnontargets[T_i, 2:(2+T-1), repet_i] = curr_params_fit['mixt_nontargets']
            # result_em_fits_allnontargets[T_i, -3:, repet_i] = [curr_params_fit[key] for key in ('mixt_random', 'train_LL', 'bic')]

            # Compute fisher info
            # print "compute fisher info"
            # result_fi_theo[T_i, repet_i] = sampler.estimate_fisher_info_theocov(use_theoretical_cov=False)
            # result_fi_theocov[T_i, repet_i] = sampler.estimate_fisher_info_theocov(use_theoretical_cov=True)

            # Compute distances to datasets
            if T in bays09_T_space:
                result_dist_bays09[T_i, :, repet_i] = (bays09_experimental_mixtures_mean[:, bays09_T_space == T].flatten() - result_em_fits[T_i, :4, repet_i])**2.

                result_dist_bays09_emmixt_KL[T_i, repet_i] = utils.KL_div(result_em_fits[T_i, 1:4, repet_i], bays09_experimental_mixtures_mean[1:, bays09_T_space == T].flatten())

            if T in gorgo11_T_space:
                result_dist_gorgo11[T_i, :, repet_i] = (gorgo11_experimental_emfits_mean[:, gorgo11_T_space == T].flatten() - result_em_fits[T_i, :4, repet_i])**2.

                result_dist_gorgo11_emmixt_KL[T_i, repet_i] = utils.KL_div(result_em_fits[T_i, 1:4, repet_i], gorgo11_experimental_emfits_mean[1:, gorgo11_T_space == T].flatten())


            # If needed, store responses
            if all_parameters['collect_responses']:
                (responses, target, nontarget) = sampler.collect_responses()
                result_responses[T_i, :, repet_i] = responses
                result_target[T_i, :, repet_i] = target
                result_nontargets[T_i, :, :T_i, repet_i] = nontarget

                print "collected responses"


            print "CURRENT RESULTS:\n", result_all_precisions[T_i, repet_i], curr_params_fit, np.sum(result_dist_bays09[T_i, :, repet_i]), np.sum(result_dist_gorgo11[T_i, :, repet_i]), "\n"
            ### /Work ###

            search_progress.increment()
            if run_counter % save_every == 0 or search_progress.done():
                dataio.save_variables_default(locals())
            run_counter += 1

    # Finished
    dataio.save_variables_default(locals())

    print "All finished"
    return locals()
def plots_memory_curves(data_pbs, generator_module=None):
    '''
        Reload and plot memory curve of a Mixed code.
        Can use Marginal Fisher Information and fitted Mixture Model as well
    '''

    #### SETUP
    #
    savefigs = True
    savedata = True

    plot_pcolor_fit_precision_to_fisherinfo = False
    plot_selected_memory_curves = False
    plot_best_memory_curves = False
    plot_subplots_persigmax = True

    colormap = None  # or 'cubehelix'
    plt.rcParams['font.size'] = 16
    #
    #### /SETUP

    print "Order parameters: ", generator_module.dict_parameters_range.keys()

    result_all_precisions_mean = utils.nanmean(np.squeeze(data_pbs.dict_arrays['result_all_precisions']['results']), axis=-1)
    result_all_precisions_std = utils.nanstd(np.squeeze(data_pbs.dict_arrays['result_all_precisions']['results']), axis=-1)
    # ratio_space, sigmax_space, T, 5
    result_em_fits_mean = utils.nanmean(np.squeeze(data_pbs.dict_arrays['result_em_fits']['results']), axis=-1)
    result_em_fits_std = utils.nanstd(np.squeeze(data_pbs.dict_arrays['result_em_fits']['results']), axis=-1)
    result_marginal_inv_fi_mean = utils.nanmean(np.squeeze(data_pbs.dict_arrays['result_marginal_inv_fi']['results']), axis=-1)
    result_marginal_inv_fi_std = utils.nanstd(np.squeeze(data_pbs.dict_arrays['result_marginal_inv_fi']['results']), axis=-1)
    result_marginal_fi_mean = utils.nanmean(np.squeeze(1./data_pbs.dict_arrays['result_marginal_inv_fi']['results']), axis=-1)
    result_marginal_fi_std = utils.nanstd(np.squeeze(1./data_pbs.dict_arrays['result_marginal_inv_fi']['results']), axis=-1)

    ratioconj_space = data_pbs.loaded_data['parameters_uniques']['ratio_conj']
    sigmax_space = data_pbs.loaded_data['parameters_uniques']['sigmax']
    T_space = data_pbs.loaded_data['datasets_list'][0]['T_space']

    print ratioconj_space
    print sigmax_space
    print T_space
    print result_all_precisions_mean.shape, result_em_fits_mean.shape, result_marginal_inv_fi_mean.shape

    dataio = DataIO.DataIO(output_folder=generator_module.pbs_submission_infos['simul_out_dir'] + '/outputs/', label='global_' + dataset_infos['save_output_filename'])

    ## Load Experimental data
    data_simult = load_experimental_data.load_data_simult(data_dir=os.path.normpath(os.path.join(os.path.split(load_experimental_data.__file__)[0], '../../experimental_data/')), fit_mixture_model=True)
    memory_experimental_precision = data_simult['precision_nitems_theo']
    memory_experimental_kappa = np.array([data['kappa'] for _, data in data_simult['em_fits_nitems']['mean'].items()])
    memory_experimental_kappa_std = np.array([data['kappa'] for _, data in data_simult['em_fits_nitems']['std'].items()])

    experim_datadir = os.environ.get('WORKDIR_DROP', os.path.split(load_experimental_data.__file__)[0])
    data_bays2009 = load_experimental_data.load_data_bays09(data_dir=os.path.normpath(os.path.join(experim_datadir, '../../experimental_data/')), fit_mixture_model=True)
    bays09_experimental_mixtures_mean = data_bays2009['em_fits_nitems_arrays']['mean']
    # add interpolated points for 3 and 5 items
    bays3 = (bays09_experimental_mixtures_mean[:, 2] + bays09_experimental_mixtures_mean[:, 1])/2.
    bays5 = (bays09_experimental_mixtures_mean[:, -1] + bays09_experimental_mixtures_mean[:, -2])/2.
    bays09_experimental_mixtures_mean_compatible = np.c_[bays09_experimental_mixtures_mean[:,:2], bays3, bays09_experimental_mixtures_mean[:, 2], bays5]

    # Boost non-targets
    # bays09_experimental_mixtures_mean_compatible[1] *= 1.5
    # bays09_experimental_mixtures_mean_compatible[2] /= 1.5
    # bays09_experimental_mixtures_mean_compatible /= np.sum(bays09_experimental_mixtures_mean_compatible, axis=0)

    # Compute some landscapes of fit!
    dist_diff_precision_margfi = np.sum(np.abs(result_all_precisions_mean*2. - result_marginal_fi_mean[..., 0])**2., axis=-1)
    dist_ratio_precision_margfi = np.sum(np.abs((result_all_precisions_mean*2.)/result_marginal_fi_mean[..., 0] - 1.0)**2., axis=-1)
    dist_diff_emkappa_margfi = np.sum(np.abs(result_em_fits_mean[..., 0]*2. - result_marginal_fi_mean[..., 0])**2., axis=-1)
    dist_ratio_emkappa_margfi = np.sum(np.abs((result_em_fits_mean[..., 0]*2.)/result_marginal_fi_mean[..., 0] - 1.0)**2., axis=-1)

    dist_diff_precision_experim = np.sum(np.abs(result_all_precisions_mean - memory_experimental_precision)**2., axis=-1)
    dist_diff_emkappa_experim = np.sum(np.abs(result_em_fits_mean[..., 0] - memory_experimental_kappa)**2., axis=-1)

    dist_diff_precision_experim_1item = np.abs(result_all_precisions_mean[..., 0] - memory_experimental_precision[0])**2.
    dist_diff_precision_experim_2item = np.abs(result_all_precisions_mean[..., 1] - memory_experimental_precision[1])**2.
    dist_diff_precision_margfi_1item = np.abs(result_all_precisions_mean[..., 0]*2. - result_marginal_fi_mean[..., 0, 0])**2.
    dist_diff_emkappa_experim_1item = np.abs(result_em_fits_mean[..., 0, 0] - memory_experimental_kappa[0])**2.
    dist_diff_margfi_experim_1item = np.abs(result_marginal_fi_mean[..., 0, 0] - memory_experimental_precision[0])**2.

    dist_diff_emkappa_mixtures_bays09 = np.sum(np.sum((result_em_fits_mean[..., 1:4] - bays09_experimental_mixtures_mean_compatible[1:].T)**2., axis=-1), axis=-1)
    dist_diff_modelfits_experfits_bays09 = np.sum(np.sum((result_em_fits_mean[..., :4] - bays09_experimental_mixtures_mean_compatible.T)**2., axis=-1), axis=-1)


    if plot_pcolor_fit_precision_to_fisherinfo:
        # Check fit between precision and fisher info
        utils.pcolor_2d_data(dist_diff_precision_margfi, log_scale=True, x=ratioconj_space, y=sigmax_space, xlabel='ratio conj', ylabel='sigmax')

        if savefigs:
            dataio.save_current_figure('match_precision_margfi_log_pcolor_{label}_{unique_id}.pdf')

        # utils.pcolor_2d_data(dist_diff_precision_margfi, x=ratioconj_space, y=sigmax_space[2:], xlabel='ratio conj', ylabel='sigmax')
        # if savefigs:
        #    dataio.save_current_figure('match_precision_margfi_pcolor_{label}_{unique_id}.pdf')

        utils.pcolor_2d_data(dist_ratio_precision_margfi, x=ratioconj_space, y=sigmax_space, xlabel='ratio conj', ylabel='sigmax', log_scale=True)
        if savefigs:
            dataio.save_current_figure('match_ratio_precision_margfi_log_pcolor_{label}_{unique_id}.pdf')

        utils.pcolor_2d_data(dist_diff_emkappa_margfi, x=ratioconj_space, y=sigmax_space, xlabel='ratio conj', ylabel='sigmax', log_scale=True)
        if savefigs:
            dataio.save_current_figure('match_diff_emkappa_margfi_log_pcolor_{label}_{unique_id}.pdf')

        utils.pcolor_2d_data(dist_ratio_emkappa_margfi, x=ratioconj_space, y=sigmax_space, xlabel='ratio conj', ylabel='sigmax', log_scale=True)
        if savefigs:
            dataio.save_current_figure('match_ratio_emkappa_margfi_log_pcolor_{label}_{unique_id}.pdf')


        utils.pcolor_2d_data(dist_diff_precision_experim, x=ratioconj_space, y=sigmax_space, xlabel='ratio conj', ylabel='sigmax', log_scale=True)
        if savefigs:
            dataio.save_current_figure('match_diff_precision_experim_log_pcolor_{label}_{unique_id}.pdf')


        utils.pcolor_2d_data(dist_diff_emkappa_experim, x=ratioconj_space, y=sigmax_space, xlabel='ratio conj', ylabel='sigmax', log_scale=True)
        if savefigs:
            dataio.save_current_figure('match_diff_emkappa_experim_log_pcolor_{label}_{unique_id}.pdf')


        utils.pcolor_2d_data(dist_diff_precision_margfi*dist_diff_emkappa_margfi*dist_diff_precision_experim*dist_diff_emkappa_experim, x=ratioconj_space, y=sigmax_space, xlabel='ratio conj', ylabel='sigmax', log_scale=True)
        if savefigs:
            dataio.save_current_figure('match_bigmultiplication_log_pcolor_{label}_{unique_id}.pdf')

        utils.pcolor_2d_data(dist_diff_precision_margfi_1item, log_scale=True, x=ratioconj_space, y=sigmax_space, xlabel='ratio conj', ylabel='sigmax')
        if savefigs:
            dataio.save_current_figure('match_diff_precision_margfi_1item_log_pcolor_{label}_{unique_id}.pdf')


        utils.pcolor_2d_data(dist_diff_precision_experim_1item, log_scale=True, x=ratioconj_space, y=sigmax_space, xlabel='ratio conj', ylabel='sigmax')
        if savefigs:
            dataio.save_current_figure('match_diff_precision_experim_1item_log_pcolor_{label}_{unique_id}.pdf')


        utils.pcolor_2d_data(dist_diff_emkappa_experim_1item, log_scale=True, x=ratioconj_space, y=sigmax_space, xlabel='ratio conj', ylabel='sigmax')
        if savefigs:
            dataio.save_current_figure('match_diff_emkappa_experim_1item_log_pcolor_{label}_{unique_id}.pdf')


        utils.pcolor_2d_data(dist_diff_margfi_experim_1item, log_scale=True, x=ratioconj_space, y=sigmax_space, xlabel='ratio conj', ylabel='sigmax')
        if savefigs:
            dataio.save_current_figure('match_diff_margfi_experim_1item_log_pcolor_{label}_{unique_id}.pdf')

        utils.pcolor_2d_data(dist_diff_precision_experim_2item, log_scale=True, x=ratioconj_space, y=sigmax_space, xlabel='ratio conj', ylabel='sigmax')
        if savefigs:
            dataio.save_current_figure('match_diff_precision_experim_2item_log_pcolor_{label}_{unique_id}.pdf')

        utils.pcolor_2d_data(dist_diff_emkappa_mixtures_bays09, log_scale=True, x=ratioconj_space, y=sigmax_space, xlabel='ratio conj', ylabel='sigmax')
        if savefigs:
            dataio.save_current_figure('match_diff_mixtures_experbays09_pcolor_{label}_{unique_id}.pdf')

        utils.pcolor_2d_data(dist_diff_modelfits_experfits_bays09, log_scale=True, x=ratioconj_space, y=sigmax_space, xlabel='ratio conj', ylabel='sigmax')
        if savefigs:
            dataio.save_current_figure('match_diff_emfits_experbays09_pcolor_{label}_{unique_id}.pdf')

    # Macro plot
    def mem_plot_precision(sigmax_i, ratioconj_i):
        ax = utils.plot_mean_std_area(T_space, memory_experimental_precision, np.zeros(T_space.size), linewidth=3, fmt='o-', markersize=8, label='Experimental data')

        ax = utils.plot_mean_std_area(T_space, result_all_precisions_mean[ratioconj_i, sigmax_i], result_all_precisions_std[ratioconj_i, sigmax_i], ax_handle=ax, linewidth=3, fmt='o-', markersize=8, label='Precision of samples')

        # ax = utils.plot_mean_std_area(T_space, 0.5*result_marginal_fi_mean[..., 0][ratioconj_i, sigmax_i], 0.5*result_marginal_fi_std[..., 0][ratioconj_i, sigmax_i], ax_handle=ax, linewidth=3, fmt='o-', markersize=8, label='Marginal Fisher Information')

        # ax = utils.plot_mean_std_area(T_space, result_em_fits_mean[..., 0][ratioconj_i, sigmax_i], result_em_fits_std[..., 0][ratioconj_i, sigmax_i], ax_handle=ax, xlabel='Number of items', ylabel="Inverse variance $[rad^{-2}]$", linewidth=3, fmt='o-', markersize=8, label='Fitted kappa')

        ax.set_title('ratio_conj %.2f, sigmax %.2f' % (ratioconj_space[ratioconj_i], sigmax_space[sigmax_i]))
        ax.legend()
        ax.set_xlim([0.9, 5.1])
        ax.set_xticks(range(1, 6))
        ax.set_xticklabels(range(1, 6))

        if savefigs:
            dataio.save_current_figure('memorycurves_precision_ratioconj%.2fsigmax%.2f_{label}_{unique_id}.pdf' % (ratioconj_space[ratioconj_i], sigmax_space[sigmax_i]))

    def mem_plot_kappa(sigmax_i, ratioconj_i):
        ax = utils.plot_mean_std_area(T_space, memory_experimental_kappa, memory_experimental_kappa_std, linewidth=3, fmt='o-', markersize=8, label='Experimental data')

        ax = utils.plot_mean_std_area(T_space, result_em_fits_mean[..., 0][ratioconj_i, sigmax_i], result_em_fits_std[..., 0][ratioconj_i, sigmax_i], xlabel='Number of items', ylabel="Memory error $[rad^{-2}]$", ax_handle=ax, linewidth=3, fmt='o-', markersize=8, label='Fitted kappa')

        # ax = utils.plot_mean_std_area(T_space, 0.5*result_marginal_fi_mean[..., 0][ratioconj_i, sigmax_i], 0.5*result_marginal_fi_std[..., 0][ratioconj_i, sigmax_i], ax_handle=ax, linewidth=3, fmt='o-', markersize=8, label='Marginal Fisher Information')

        ax.set_title('ratio_conj %.2f, sigmax %.2f' % (ratioconj_space[ratioconj_i], sigmax_space[sigmax_i]))
        ax.legend()
        ax.set_xlim([0.9, 5.1])
        ax.set_xticks(range(1, 6))
        ax.set_xticklabels(range(1, 6))

        if savefigs:
            dataio.save_current_figure('memorycurves_kappa_ratioconj%.2fsigmax%.2f_{label}_{unique_id}.pdf' % (ratioconj_space[ratioconj_i], sigmax_space[sigmax_i]))

    def em_plot(sigmax_i, ratioconj_i):
        # TODO finish checking this up.
        f, ax = plt.subplots()
        ax2 = ax.twinx()

        # left axis, kappa
        ax = utils.plot_mean_std_area(T_space, result_em_fits_mean[..., 0][ratioconj_i, sigmax_i], result_em_fits_std[..., 0][ratioconj_i, sigmax_i], xlabel='Number of items', ylabel="Inverse variance $[rad^{-2}]$", ax_handle=ax, linewidth=3, fmt='o-', markersize=8, label='Fitted kappa', color='k')

        # Right axis, mixture probabilities
        utils.plot_mean_std_area(T_space, result_em_fits_mean[..., 1][ratioconj_i, sigmax_i], result_em_fits_std[..., 1][ratioconj_i, sigmax_i], xlabel='Number of items', ylabel="Mixture probabilities", ax_handle=ax2, linewidth=3, fmt='o-', markersize=8, label='Target')
        utils.plot_mean_std_area(T_space, result_em_fits_mean[..., 2][ratioconj_i, sigmax_i], result_em_fits_std[..., 2][ratioconj_i, sigmax_i], xlabel='Number of items', ylabel="Mixture probabilities", ax_handle=ax2, linewidth=3, fmt='o-', markersize=8, label='Nontarget')
        utils.plot_mean_std_area(T_space, result_em_fits_mean[..., 3][ratioconj_i, sigmax_i], result_em_fits_std[..., 3][ratioconj_i, sigmax_i], xlabel='Number of items', ylabel="Mixture probabilities", ax_handle=ax2, linewidth=3, fmt='o-', markersize=8, label='Random')

        lines, labels = ax.get_legend_handles_labels()
        lines2, labels2 = ax2.get_legend_handles_labels()
        ax.legend(lines + lines2, labels + labels2)

        ax.set_title('ratio_conj %.2f, sigmax %.2f' % (ratioconj_space[ratioconj_i], sigmax_space[sigmax_i]))
        ax.set_xlim([0.9, 5.1])
        ax.set_xticks(range(1, 6))
        ax.set_xticklabels(range(1, 6))

        f.canvas.draw()

        if savefigs:
            dataio.save_current_figure('memorycurves_emfits_ratioconj%.2fsigmax%.2f_{label}_{unique_id}.pdf' % (ratioconj_space[ratioconj_i], sigmax_space[sigmax_i]))

    def em_plot_paper(sigmax_i, ratioconj_i):
        f, ax = plt.subplots()

        # Right axis, mixture probabilities
        utils.plot_mean_std_area(T_space, result_em_fits_mean[..., 1][ratioconj_i, sigmax_i], result_em_fits_std[..., 1][ratioconj_i, sigmax_i], xlabel='Number of items', ylabel="Mixture probabilities", ax_handle=ax, linewidth=3, fmt='o-', markersize=5, label='Target')
        utils.plot_mean_std_area(T_space, result_em_fits_mean[..., 2][ratioconj_i, sigmax_i], result_em_fits_std[..., 2][ratioconj_i, sigmax_i], xlabel='Number of items', ylabel="Mixture probabilities", ax_handle=ax, linewidth=3, fmt='o-', markersize=5, label='Nontarget')
        utils.plot_mean_std_area(T_space, result_em_fits_mean[..., 3][ratioconj_i, sigmax_i], result_em_fits_std[..., 3][ratioconj_i, sigmax_i], xlabel='Number of items', ylabel="Mixture probabilities", ax_handle=ax, linewidth=3, fmt='o-', markersize=5, label='Random')

        ax.legend(prop={'size':15})

        ax.set_title('ratio_conj %.2f, sigmax %.2f' % (ratioconj_space[ratioconj_i], sigmax_space[sigmax_i]))
        ax.set_xlim([1.0, 5.0])
        ax.set_ylim([0.0, 1.1])
        ax.set_xticks(range(1, 6))
        ax.set_xticklabels(range(1, 6))

        f.canvas.draw()

        if savefigs:
            dataio.save_current_figure('memorycurves_emfits_paper_ratioconj%.2fsigmax%.2f_{label}_{unique_id}.pdf' % (ratioconj_space[ratioconj_i], sigmax_space[sigmax_i]))


    if plot_selected_memory_curves:
        selected_values = [[0.84, 0.23], [0.84, 0.19]]

        for current_values in selected_values:
            # Find the indices
            ratioconj_i     = np.argmin(np.abs(current_values[0] - ratioconj_space))
            sigmax_i        = np.argmin(np.abs(current_values[1] - sigmax_space))


            mem_plot_precision(sigmax_i, ratioconj_i)
            mem_plot_kappa(sigmax_i, ratioconj_i)

    if plot_best_memory_curves:
        # Best precision fit
        best_axis2_i_all = np.argmin(dist_diff_precision_experim, axis=1)

        for axis1_i, best_axis2_i in enumerate(best_axis2_i_all):
            mem_plot_precision(best_axis2_i, axis1_i)

        # Best kappa fit
        best_axis2_i_all = np.argmin(dist_diff_emkappa_experim, axis=1)

        for axis1_i, best_axis2_i in enumerate(best_axis2_i_all):
            mem_plot_kappa(best_axis2_i, axis1_i)
            em_plot(best_axis2_i, axis1_i)

        # Best em parameters fit to Bays09
        best_axis2_i_all = np.argmin(dist_diff_modelfits_experfits_bays09, axis=1)

        for axis1_i, best_axis2_i in enumerate(best_axis2_i_all):
            mem_plot_kappa(best_axis2_i, axis1_i)
            # em_plot(best_axis2_i, axis1_i)
            em_plot_paper(best_axis2_i, axis1_i)

    if plot_subplots_persigmax:
        # Do subplots with ratio_conj on x, one plot per T and different figures per sigmax. Looks a bit like reloader_hierarchical_constantMMlower_maxlik_allresponses_221213.py

        sigmax_selected_indices = np.array([np.argmin((sigmax_space - sigmax_select)**2.) for sigmax_select in [0.1, 0.2, 0.3, 0.4, 0.5]])

        for sigmax_select_i in sigmax_selected_indices:
            # two plots per sigmax.

            f, axes = plt.subplots(nrows=T_space[-1], ncols=1, sharex=True, figsize=(10, 12))
            # all_lines_bis = []

            for T_i, T in enumerate(T_space):
                # Plot EM mixtures
                utils.plot_mean_std_area(ratioconj_space, result_em_fits_mean[:, sigmax_select_i, T_i, 1], result_em_fits_std[:, sigmax_select_i, T_i, 1], ax_handle=axes[T_i], linewidth=3, fmt='o-', markersize=5)
                utils.plot_mean_std_area(ratioconj_space, result_em_fits_mean[:, sigmax_select_i, T_i, 2], result_em_fits_std[:, sigmax_select_i, T_i, 2], ax_handle=axes[T_i], linewidth=3, fmt='o-', markersize=5)
                utils.plot_mean_std_area(ratioconj_space, result_em_fits_mean[:, sigmax_select_i, T_i, 3], result_em_fits_std[:, sigmax_select_i, T_i, 3], ax_handle=axes[T_i], linewidth=3, fmt='o-', markersize=5)

                # ratio_MMlower_space, result_emfits_filtered[:, i, 1:4], 0*result_emfits_filtered[:, i, 1:4], ax_handle=axes[T_i], linewidth=2) #, color=all_lines[T_i].get_color())
                # curr_lines = axes[T_i].plot(ratio_MMlower_space, results_precision_filtered[:, i], linewidth=2, color=all_lines[T_i].get_color())
                axes[T_i].grid()
                axes[T_i].set_xticks(np.linspace(0., 1.0, 5))
                axes[T_i].set_xlim((0.0, 1.0))
                # axes[T_i].set_yticks([])
                # axes[T_i].set_ylim((np.min(result_emfits_filtered[:, i, 0]), result_emfits_filtered[max_ind, i, 0]*1.1))
                axes[T_i].set_ylim((0.0, 1.05))
                axes[T_i].locator_params(axis='y', tight=True, nbins=4)
                # all_lines_bis.extend(curr_lines)

            axes[0].set_title('Sigmax: %.3f' % sigmax_space[sigmax_select_i])
            axes[-1].set_xlabel('Proportion of conjunctive units')

            if savefigs:
                dataio.save_current_figure('results_subplots_emmixtures_sigmax%.2f_{label}_global_{unique_id}.pdf' % sigmax_space[sigmax_select_i])

            f, axes = plt.subplots(nrows=T_space[-1], ncols=1, sharex=True, figsize=(10, 12))

            # Kappa kappa
            for T_i, T in enumerate(T_space):

                # Plot kappa mixture
                utils.plot_mean_std_area(ratioconj_space, result_em_fits_mean[:, sigmax_select_i, T_i, 0], result_em_fits_std[:, sigmax_select_i, T_i, 0], ax_handle=axes[T_i], linewidth=3, fmt='o-', markersize=5)
                # utils.plot_mean_std_area(ratio_MMlower_space, result_emfits_filtered[:, i, 0], 0*result_emfits_filtered[:, i, 0], ax_handle=axes[T_i], linewidth=2) #, color=all_lines[T_i].get_color())
                # curr_lines = axes[T_i].plot(ratio_MMlower_space, results_precision_filtered[:, i], linewidth=2, color=all_lines[T_i].get_color())
                axes[T_i].grid()
                axes[T_i].set_xticks(np.linspace(0., 1.0, 5))
                axes[T_i].set_xlim((0.0, 1.0))
                # axes[T_i].set_yticks([])
                # axes[T_i].set_ylim((np.min(result_emfits_filtered[:, i, 0]), result_emfits_filtered[max_ind, i, 0]*1.1))
                # axes[T_i].set_ylim((0.0, 1.0))
                axes[T_i].locator_params(axis='y', tight=True, nbins=4)
                # all_lines_bis.extend(curr_lines)

            axes[0].set_title('Sigmax: %.3f' % sigmax_space[sigmax_select_i])
            axes[-1].set_xlabel('Proportion of conjunctive units')

            # f.subplots_adjust(right=0.75)
            # plt.figlegend(all_lines_bis, ['%d item' % i + 's'*(i>1) for i in xrange(1, T+1)], loc='right', bbox_to_anchor=(1.0, 0.5))

            if savefigs:
                dataio.save_current_figure('results_subplots_emkappa_sigmax%.2f_{label}_global_{unique_id}.pdf' % sigmax_space[sigmax_select_i])



    all_args = data_pbs.loaded_data['args_list']
    variables_to_save = ['memory_experimental_precision', 'memory_experimental_kappa', 'bays09_experimental_mixtures_mean_compatible']

    if savedata:
        dataio.save_variables_default(locals(), variables_to_save)

        dataio.make_link_output_to_dropbox(dropbox_current_experiment_folder='memory_curves')

    plt.show()

    return locals()
def plots_memory_curves(data_pbs, generator_module=None):
    """
        Reload and plot memory curve of a Mixed code.
        Can use Marginal Fisher Information and fitted Mixture Model as well
    """

    #### SETUP
    #
    savefigs = True
    savedata = True

    do_error_distrib_fits = True

    plot_pcolor_fit_precision_to_fisherinfo = True
    plot_selected_memory_curves = False
    plot_best_memory_curves = False
    plot_best_error_distrib = True

    colormap = None  # or 'cubehelix'
    plt.rcParams["font.size"] = 16
    #
    #### /SETUP

    print "Order parameters: ", generator_module.dict_parameters_range.keys()

    result_all_precisions_mean = utils.nanmean(data_pbs.dict_arrays["result_all_precisions"]["results"], axis=-1)
    result_all_precisions_std = utils.nanstd(data_pbs.dict_arrays["result_all_precisions"]["results"], axis=-1)
    result_em_fits_mean = utils.nanmean(data_pbs.dict_arrays["result_em_fits"]["results"], axis=-1)
    result_em_fits_std = utils.nanstd(data_pbs.dict_arrays["result_em_fits"]["results"], axis=-1)
    result_marginal_inv_fi_mean = utils.nanmean(data_pbs.dict_arrays["result_marginal_inv_fi"]["results"], axis=-1)
    result_marginal_inv_fi_std = utils.nanstd(data_pbs.dict_arrays["result_marginal_inv_fi"]["results"], axis=-1)
    result_marginal_fi_mean = utils.nanmean(1.0 / data_pbs.dict_arrays["result_marginal_inv_fi"]["results"], axis=-1)
    result_marginal_fi_std = utils.nanstd(1.0 / data_pbs.dict_arrays["result_marginal_inv_fi"]["results"], axis=-1)
    result_responses_all = data_pbs.dict_arrays["result_responses"]["results"]
    result_target_all = data_pbs.dict_arrays["result_target"]["results"]
    result_nontargets_all = data_pbs.dict_arrays["result_nontargets"]["results"]

    M_space = data_pbs.loaded_data["parameters_uniques"]["M"].astype(int)
    sigmax_space = data_pbs.loaded_data["parameters_uniques"]["sigmax"]
    T_space = data_pbs.loaded_data["datasets_list"][0]["T_space"]
    nb_repetitions = result_responses_all.shape[-1]

    print M_space
    print sigmax_space
    print T_space
    print result_all_precisions_mean.shape, result_em_fits_mean.shape, result_marginal_inv_fi_mean.shape

    dataio = DataIO.DataIO(
        output_folder=generator_module.pbs_submission_infos["simul_out_dir"] + "/outputs/",
        label="global_" + dataset_infos["save_output_filename"],
    )

    ## Load Experimental data
    experim_datadir = os.environ.get("WORKDIR_DROP", os.path.split(load_experimental_data.__file__)[0])
    data_simult = load_experimental_data.load_data_simult(
        data_dir=os.path.normpath(os.path.join(experim_datadir, "../../experimental_data/")), fit_mixture_model=True
    )
    gorgo11_experimental_precision = data_simult["precision_nitems_theo"]
    gorgo11_experimental_kappa = np.array([data["kappa"] for _, data in data_simult["em_fits_nitems"]["mean"].items()])
    gorgo11_experimental_emfits_mean = np.array(
        [
            [data[key] for _, data in data_simult["em_fits_nitems"]["mean"].items()]
            for key in ["kappa", "mixt_target", "mixt_nontargets", "mixt_random"]
        ]
    )
    gorgo11_experimental_emfits_std = np.array(
        [
            [data[key] for _, data in data_simult["em_fits_nitems"]["std"].items()]
            for key in ["kappa", "mixt_target", "mixt_nontargets", "mixt_random"]
        ]
    )
    gorgo11_experimental_emfits_sem = gorgo11_experimental_emfits_std / np.sqrt(np.unique(data_simult["subject"]).size)

    experim_datadir = os.environ.get("WORKDIR_DROP", os.path.split(load_experimental_data.__file__)[0])
    data_bays2009 = load_experimental_data.load_data_bays09(
        data_dir=os.path.normpath(os.path.join(experim_datadir, "../../experimental_data/")), fit_mixture_model=True
    )
    bays09_experimental_mixtures_mean = data_bays2009["em_fits_nitems_arrays"]["mean"]
    bays09_experimental_mixtures_std = data_bays2009["em_fits_nitems_arrays"]["std"]
    # add interpolated points for 3 and 5 items
    emfit_mean_intpfct = spint.interp1d(np.unique(data_bays2009["n_items"]), bays09_experimental_mixtures_mean)
    bays09_experimental_mixtures_mean_compatible = emfit_mean_intpfct(np.arange(1, 7))
    emfit_std_intpfct = spint.interp1d(np.unique(data_bays2009["n_items"]), bays09_experimental_mixtures_std)
    bays09_experimental_mixtures_std_compatible = emfit_std_intpfct(np.arange(1, 7))
    T_space_bays09 = np.arange(1, 6)

    # Boost non-targets
    # bays09_experimental_mixtures_mean_compatible[1] *= 1.5
    # bays09_experimental_mixtures_mean_compatible[2] /= 1.5
    # bays09_experimental_mixtures_mean_compatible /= np.sum(bays09_experimental_mixtures_mean_compatible, axis=0)

    # Compute some landscapes of fit!
    # dist_diff_precision_margfi = np.sum(np.abs(result_all_precisions_mean*2. - result_marginal_fi_mean[..., 0])**2., axis=-1)
    # dist_ratio_precision_margfi = np.sum(np.abs((result_all_precisions_mean*2.)/result_marginal_fi_mean[..., 0] - 1.0)**2., axis=-1)
    # dist_diff_emkappa_margfi = np.sum(np.abs(result_em_fits_mean[..., 0]*2. - result_marginal_fi_mean[..., 0])**2., axis=-1)
    # dist_ratio_emkappa_margfi = np.sum(np.abs((result_em_fits_mean[..., 0]*2.)/result_marginal_fi_mean[..., 0] - 1.0)**2., axis=-1)

    dist_diff_precision_experim = np.sum(
        np.abs(result_all_precisions_mean[..., : gorgo11_experimental_kappa.size] - gorgo11_experimental_precision)
        ** 2.0,
        axis=-1,
    )
    dist_diff_emkappa_experim = np.sum(
        np.abs(result_em_fits_mean[..., 0, : gorgo11_experimental_kappa.size] - gorgo11_experimental_kappa) ** 2.0,
        axis=-1,
    )

    dist_diff_em_mixtures_bays09 = np.sum(
        np.sum((result_em_fits_mean[..., 1:4] - bays09_experimental_mixtures_mean_compatible[1:].T) ** 2.0, axis=-1),
        axis=-1,
    )
    dist_diff_modelfits_experfits_bays09 = np.sum(
        np.sum((result_em_fits_mean[..., :4] - bays09_experimental_mixtures_mean_compatible.T) ** 2.0, axis=-1), axis=-1
    )

    if do_error_distrib_fits:
        print "computing error distribution histograms fits"
        # Now try to fit histograms of errors to target/nontargets
        bays09_hist_target_mean = data_bays2009["hist_cnts_target_nitems_stats"]["mean"]
        bays09_hist_target_std = data_bays2009["hist_cnts_target_nitems_stats"]["std"]
        bays09_hist_nontarget_mean = data_bays2009["hist_cnts_nontarget_nitems_stats"]["mean"]
        bays09_hist_nontarget_std = data_bays2009["hist_cnts_nontarget_nitems_stats"]["std"]
        T_space_bays09_filt = np.unique(data_bays2009["n_items"])

        angle_space = np.linspace(-np.pi, np.pi, bays09_hist_target_mean.shape[-1] + 1)
        bins_center = angle_space[:-1] + np.diff(angle_space)[0] / 2

        errors_targets = utils.wrap_angles(result_responses_all - result_target_all)
        hist_targets_all = np.empty(
            (M_space.size, sigmax_space.size, T_space_bays09_filt.size, angle_space.size - 1, nb_repetitions)
        )

        errors_nontargets = np.nan * np.empty(result_nontargets_all.shape)
        hist_nontargets_all = np.empty(
            (M_space.size, sigmax_space.size, T_space_bays09_filt.size, angle_space.size - 1, nb_repetitions)
        )
        for M_i, M in enumerate(M_space):
            for sigmax_i, sigmax in enumerate(sigmax_space):
                for T_bays_i, T_bays in enumerate(T_space_bays09_filt):
                    for repet_i in xrange(nb_repetitions):
                        # Could do a nicer indexing but f**k it

                        # Histogram errors to targets
                        hist_targets_all[M_i, sigmax_i, T_bays_i, :, repet_i], x, bins = utils.histogram_binspace(
                            utils.dropnan(errors_targets[M_i, sigmax_i, T_bays - 1, ..., repet_i]),
                            bins=angle_space,
                            norm="density",
                        )

                        # Compute the error between the responses and nontargets.
                        errors_nontargets[M_i, sigmax_i, T_bays - 1, :, :, repet_i] = utils.wrap_angles(
                            (
                                result_responses_all[M_i, sigmax_i, T_bays - 1, :, repet_i, np.newaxis]
                                - result_nontargets_all[M_i, sigmax_i, T_bays - 1, :, :, repet_i]
                            )
                        )
                        # Histogram it
                        hist_nontargets_all[M_i, sigmax_i, T_bays_i, :, repet_i], x, bins = utils.histogram_binspace(
                            utils.dropnan(errors_nontargets[M_i, sigmax_i, T_bays - 1, ..., repet_i]),
                            bins=angle_space,
                            norm="density",
                        )

        hist_targets_mean = utils.nanmean(hist_targets_all, axis=-1).filled(np.nan)
        hist_targets_std = utils.nanstd(hist_targets_all, axis=-1).filled(np.nan)
        hist_nontargets_mean = utils.nanmean(hist_nontargets_all, axis=-1).filled(np.nan)
        hist_nontargets_std = utils.nanstd(hist_nontargets_all, axis=-1).filled(np.nan)

        # Compute distances to experimental histograms
        dist_diff_hist_target_bays09 = np.nansum(
            np.nansum((hist_targets_mean - bays09_hist_target_mean) ** 2.0, axis=-1), axis=-1
        )
        dist_diff_hist_nontargets_bays09 = np.nansum(
            np.nansum((hist_nontargets_mean - bays09_hist_nontarget_mean) ** 2.0, axis=-1), axis=-1
        )
        dist_diff_hist_nontargets_5_6items_bays09 = np.nansum(
            np.nansum((hist_nontargets_mean[:, :, -2:] - bays09_hist_nontarget_mean[-2:]) ** 2.0, axis=-1), axis=-1
        )

    if plot_pcolor_fit_precision_to_fisherinfo:
        # Check fit between precision and Experiments
        utils.pcolor_2d_data(
            dist_diff_precision_experim,
            log_scale=True,
            x=M_space,
            y=sigmax_space,
            xlabel="M",
            ylabel="sigmax",
            xlabel_format="%d",
        )
        if savefigs:
            dataio.save_current_figure("match_precision_exper_log_pcolor_{label}_{unique_id}.pdf")

        utils.pcolor_2d_data(
            dist_diff_emkappa_experim, x=M_space, y=sigmax_space, xlabel="M", ylabel="sigmax", xlabel_format="%d"
        )
        if savefigs:
            dataio.save_current_figure("match_emkappa_model_exper_pcolor_{label}_{unique_id}.pdf")

        utils.pcolor_2d_data(
            dist_diff_em_mixtures_bays09,
            x=M_space,
            y=sigmax_space,
            xlabel="M",
            ylabel="sigmax",
            log_scale=True,
            xlabel_format="%d",
        )
        if savefigs:
            dataio.save_current_figure("match_emmixtures_experbays09_log_pcolor_{label}_{unique_id}.pdf")

        utils.pcolor_2d_data(
            dist_diff_modelfits_experfits_bays09,
            log_scale=True,
            x=M_space,
            y=sigmax_space,
            xlabel="M",
            ylabel="sigmax",
            xlabel_format="%d",
        )
        if savefigs:
            dataio.save_current_figure("match_diff_emfits_experbays09_pcolor_{label}_{unique_id}.pdf")

        if do_error_distrib_fits:
            utils.pcolor_2d_data(
                dist_diff_hist_target_bays09,
                x=M_space,
                y=sigmax_space,
                xlabel="M",
                ylabel="sigmax",
                log_scale=True,
                xlabel_format="%d",
            )
            if savefigs:
                dataio.save_current_figure("match_hist_targets_experbays09_log_pcolor_{label}_{unique_id}.pdf")

            utils.pcolor_2d_data(
                dist_diff_hist_nontargets_bays09,
                x=M_space,
                y=sigmax_space,
                xlabel="M",
                ylabel="sigmax",
                log_scale=True,
                xlabel_format="%d",
            )
            if savefigs:
                dataio.save_current_figure("match_hist_nontargets_experbays09_log_pcolor_{label}_{unique_id}.pdf")

            utils.pcolor_2d_data(
                dist_diff_hist_nontargets_5_6items_bays09,
                x=M_space,
                y=sigmax_space,
                xlabel="M",
                ylabel="sigmax",
                log_scale=True,
                xlabel_format="%d",
            )
            if savefigs:
                dataio.save_current_figure(
                    "match_hist_nontargets_6items_experbays09_log_pcolor_{label}_{unique_id}.pdf"
                )

    # Macro plot
    def mem_plot_precision(sigmax_i, M_i, mem_exp_prec):
        ax = utils.plot_mean_std_area(
            T_space[: mem_exp_prec.size],
            mem_exp_prec,
            np.zeros(mem_exp_prec.size),
            linewidth=3,
            fmt="o-",
            markersize=8,
            label="Experimental data",
        )

        ax = utils.plot_mean_std_area(
            T_space[: mem_exp_prec.size],
            result_all_precisions_mean[M_i, sigmax_i, : mem_exp_prec.size],
            result_all_precisions_std[M_i, sigmax_i, : mem_exp_prec.size],
            ax_handle=ax,
            linewidth=3,
            fmt="o-",
            markersize=8,
            label="Precision of samples",
        )

        # ax = utils.plot_mean_std_area(T_space, 0.5*result_marginal_fi_mean[..., 0][M_i, sigmax_i], 0.5*result_marginal_fi_std[..., 0][M_i, sigmax_i], ax_handle=ax, linewidth=3, fmt='o-', markersize=8, label='Marginal Fisher Information')

        # ax = utils.plot_mean_std_area(T_space, result_em_fits_mean[..., 0][M_i, sigmax_i], result_em_fits_std[..., 0][M_i, sigmax_i], ax_handle=ax, xlabel='Number of items', ylabel="Inverse variance $[rad^{-2}]$", linewidth=3, fmt='o-', markersize=8, label='Fitted kappa')

        ax.set_title("M %d, sigmax %.2f" % (M_space[M_i], sigmax_space[sigmax_i]))
        ax.legend()
        ax.set_xlim([0.9, mem_exp_prec.size + 0.1])
        ax.set_xticks(range(1, mem_exp_prec.size + 1))
        ax.set_xticklabels(range(1, mem_exp_prec.size + 1))

        if savefigs:
            dataio.save_current_figure(
                "memorycurves_precision_M%dsigmax%.2f_{label}_{unique_id}.pdf" % (M_space[M_i], sigmax_space[sigmax_i])
            )

    def mem_plot_kappa(sigmax_i, M_i, exp_kappa_mean, exp_kappa_std=None):
        ax = utils.plot_mean_std_area(
            T_space[: exp_kappa_mean.size],
            exp_kappa_mean,
            exp_kappa_std,
            linewidth=3,
            fmt="o-",
            markersize=8,
            label="Experimental data",
        )

        ax = utils.plot_mean_std_area(
            T_space[: exp_kappa_mean.size],
            result_em_fits_mean[..., : exp_kappa_mean.size, 0][M_i, sigmax_i],
            result_em_fits_std[..., : exp_kappa_mean.size, 0][M_i, sigmax_i],
            xlabel="Number of items",
            ylabel="Memory error $[rad^{-2}]$",
            linewidth=3,
            fmt="o-",
            markersize=8,
            label="Fitted kappa",
            ax_handle=ax,
        )

        # ax = utils.plot_mean_std_area(T_space, 0.5*result_marginal_fi_mean[..., 0][M_i, sigmax_i], 0.5*result_marginal_fi_std[..., 0][M_i, sigmax_i], ax_handle=ax, linewidth=3, fmt='o-', markersize=8, label='Marginal Fisher Information')

        ax.set_title("M %d, sigmax %.2f" % (M_space[M_i], sigmax_space[sigmax_i]))
        ax.legend()
        ax.set_xlim([0.9, exp_kappa_mean.size + 0.1])
        ax.set_xticks(range(1, exp_kappa_mean.size + 1))
        ax.set_xticklabels(range(1, exp_kappa_mean.size + 1))

        ax.get_figure().canvas.draw()

        if savefigs:
            dataio.save_current_figure(
                "memorycurves_kappa_M%dsigmax%.2f_{label}_{unique_id}.pdf" % (M_space[M_i], sigmax_space[sigmax_i])
            )

    def em_plot(sigmax_i, M_i):
        f, ax = plt.subplots()
        ax2 = ax.twinx()

        # left axis, kappa
        ax = utils.plot_mean_std_area(
            T_space,
            result_em_fits_mean[..., 0][M_i, sigmax_i],
            result_em_fits_std[..., 0][M_i, sigmax_i],
            xlabel="Number of items",
            ylabel="Inverse variance $[rad^{-2}]$",
            ax_handle=ax,
            linewidth=3,
            fmt="o-",
            markersize=8,
            label="Fitted kappa",
            color="k",
        )

        # Right axis, mixture probabilities
        utils.plot_mean_std_area(
            T_space,
            result_em_fits_mean[..., 1][M_i, sigmax_i],
            result_em_fits_std[..., 1][M_i, sigmax_i],
            xlabel="Number of items",
            ylabel="Mixture probabilities",
            ax_handle=ax2,
            linewidth=3,
            fmt="o-",
            markersize=8,
            label="Target",
        )
        utils.plot_mean_std_area(
            T_space,
            result_em_fits_mean[..., 2][M_i, sigmax_i],
            result_em_fits_std[..., 2][M_i, sigmax_i],
            xlabel="Number of items",
            ylabel="Mixture probabilities",
            ax_handle=ax2,
            linewidth=3,
            fmt="o-",
            markersize=8,
            label="Nontarget",
        )
        utils.plot_mean_std_area(
            T_space,
            result_em_fits_mean[..., 3][M_i, sigmax_i],
            result_em_fits_std[..., 3][M_i, sigmax_i],
            xlabel="Number of items",
            ylabel="Mixture probabilities",
            ax_handle=ax2,
            linewidth=3,
            fmt="o-",
            markersize=8,
            label="Random",
        )

        lines, labels = ax.get_legend_handles_labels()
        lines2, labels2 = ax2.get_legend_handles_labels()
        ax.legend(lines + lines2, labels + labels2)

        ax.set_title("M %d, sigmax %.2f" % (M_space[M_i], sigmax_space[sigmax_i]))
        ax.set_xlim([0.9, T_space.size])
        ax.set_xticks(range(1, T_space.size + 1))
        ax.set_xticklabels(range(1, T_space.size + 1))

        f.canvas.draw()

        if savefigs:
            dataio.save_current_figure(
                "memorycurves_emfits_M%dsigmax%.2f_{label}_{unique_id}.pdf" % (M_space[M_i], sigmax_space[sigmax_i])
            )

    def em_plot_paper(sigmax_i, M_i):
        f, ax = plt.subplots()

        # Right axis, mixture probabilities
        utils.plot_mean_std_area(
            T_space_bays09,
            result_em_fits_mean[..., 1][M_i, sigmax_i][: T_space_bays09.size],
            result_em_fits_std[..., 1][M_i, sigmax_i][: T_space_bays09.size],
            xlabel="Number of items",
            ylabel="Mixture probabilities",
            ax_handle=ax,
            linewidth=3,
            fmt="o-",
            markersize=5,
            label="Target",
        )
        utils.plot_mean_std_area(
            T_space_bays09,
            result_em_fits_mean[..., 2][M_i, sigmax_i][: T_space_bays09.size],
            result_em_fits_std[..., 2][M_i, sigmax_i][: T_space_bays09.size],
            xlabel="Number of items",
            ylabel="Mixture probabilities",
            ax_handle=ax,
            linewidth=3,
            fmt="o-",
            markersize=5,
            label="Nontarget",
        )
        utils.plot_mean_std_area(
            T_space_bays09,
            result_em_fits_mean[..., 3][M_i, sigmax_i][: T_space_bays09.size],
            result_em_fits_std[..., 3][M_i, sigmax_i][: T_space_bays09.size],
            xlabel="Number of items",
            ylabel="Mixture probabilities",
            ax_handle=ax,
            linewidth=3,
            fmt="o-",
            markersize=5,
            label="Random",
        )

        ax.legend(prop={"size": 15})

        ax.set_title("M %d, sigmax %.2f" % (M_space[M_i], sigmax_space[sigmax_i]))
        ax.set_xlim([1.0, T_space_bays09.size])
        ax.set_ylim([0.0, 1.1])
        ax.set_xticks(range(1, T_space_bays09.size + 1))
        ax.set_xticklabels(range(1, T_space_bays09.size + 1))

        f.canvas.draw()

        if savefigs:
            dataio.save_current_figure(
                "memorycurves_emfits_paper_M%dsigmax%.2f_{label}_{unique_id}.pdf"
                % (M_space[M_i], sigmax_space[sigmax_i])
            )

    def hist_errors_targets_nontargets(hists_toplot_mean, hists_toplot_std, title="", M=0, sigmax=0, yaxis_lim="auto"):

        f1, axes1 = plt.subplots(
            ncols=hists_toplot_mean.shape[-2], figsize=(hists_toplot_mean.shape[-2] * 6, 6), sharey=True
        )

        for T_bays_i, T_bays in enumerate(T_space_bays09_filt):
            if not np.all(np.isnan(hists_toplot_mean[T_bays_i])):
                axes1[T_bays_i].bar(
                    bins_center,
                    hists_toplot_mean[T_bays_i],
                    width=2.0 * np.pi / (angle_space.size - 1),
                    align="center",
                    yerr=hists_toplot_std[T_bays_i],
                )
                axes1[T_bays_i].set_title("N=%d" % T_bays)

                # axes1[T_{}]

                axes1[T_bays_i].set_xlim(
                    [bins_center[0] - np.pi / (angle_space.size - 1), bins_center[-1] + np.pi / (angle_space.size - 1)]
                )
                if yaxis_lim == "target":
                    axes1[T_bays_i].set_ylim([0.0, 2.0])
                elif yaxis_lim == "nontarget":
                    axes1[T_bays_i].set_ylim([0.0, 0.3])
                else:
                    axes1[T_bays_i].set_ylim([0.0, np.nanmax(hists_toplot_mean + hists_toplot_std) * 1.1])
                axes1[T_bays_i].set_xticks((-np.pi, -np.pi / 2, 0, np.pi / 2.0, np.pi))
                axes1[T_bays_i].set_xticklabels(
                    (r"$-\pi$", r"$-\frac{\pi}{2}$", r"$0$", r"$\frac{\pi}{2}$", r"$\pi$"), fontsize=16
                )

        f1.canvas.draw()

        if savefigs:
            dataio.save_current_figure(
                "memorycurves_hist_%s_paper_M%dsigmax%.2f_{label}_{unique_id}.pdf" % (title, M, sigmax)
            )

    #################################

    if plot_selected_memory_curves:
        selected_values = [[0.84, 0.23], [0.84, 0.19]]

        for current_values in selected_values:
            # Find the indices
            M_i = np.argmin(np.abs(current_values[0] - M_space))
            sigmax_i = np.argmin(np.abs(current_values[1] - sigmax_space))

            mem_plot_precision(sigmax_i, M_i)
            mem_plot_kappa(sigmax_i, M_i)

    if plot_best_memory_curves:
        # Best precision fit
        best_axis2_i_all = np.argmin(dist_diff_precision_experim, axis=1)

        for axis1_i, best_axis2_i in enumerate(best_axis2_i_all):
            mem_plot_precision(best_axis2_i, axis1_i, gorgo11_experimental_precision)

        # Best kappa fit
        best_axis2_i_all = np.argmin(dist_diff_emkappa_experim, axis=1)

        for axis1_i, best_axis2_i in enumerate(best_axis2_i_all):
            mem_plot_kappa(
                best_axis2_i, axis1_i, gorgo11_experimental_emfits_mean[0], gorgo11_experimental_emfits_std[0]
            )
            # em_plot(best_axis2_i, axis1_i)

        # Best em parameters fit to Bays09
        best_axis2_i_all = np.argmin(dist_diff_modelfits_experfits_bays09, axis=1)
        # best_axis2_i_all = np.argmin(dist_diff_em_mixtures_bays09, axis=1)

        for axis1_i, best_axis2_i in enumerate(best_axis2_i_all):
            mem_plot_kappa(
                best_axis2_i,
                axis1_i,
                bays09_experimental_mixtures_mean_compatible[0, : T_space_bays09.size],
                bays09_experimental_mixtures_std_compatible[0, : T_space_bays09.size],
            )
            # em_plot(best_axis2_i, axis1_i)
            em_plot_paper(best_axis2_i, axis1_i)

    if plot_best_error_distrib and do_error_distrib_fits:

        # Best target histograms
        best_axis2_i_all = np.argmin(dist_diff_hist_target_bays09, axis=1)

        for axis1_i, best_axis2_i in enumerate(best_axis2_i_all):
            hist_errors_targets_nontargets(
                hist_targets_mean[axis1_i, best_axis2_i],
                hist_targets_std[axis1_i, best_axis2_i],
                "target",
                M=M_space[axis1_i],
                sigmax=sigmax_space[best_axis2_i],
                yaxis_lim="target",
            )

        # Best nontarget histograms
        best_axis2_i_all = np.argmin(dist_diff_hist_nontargets_bays09, axis=1)

        for axis1_i, best_axis2_i in enumerate(best_axis2_i_all):
            hist_errors_targets_nontargets(
                hist_nontargets_mean[axis1_i, best_axis2_i],
                hist_nontargets_std[axis1_i, best_axis2_i],
                "nontarget",
                M=M_space[axis1_i],
                sigmax=sigmax_space[best_axis2_i],
                yaxis_lim="nontarget",
            )

    all_args = data_pbs.loaded_data["args_list"]
    variables_to_save = [
        "gorgo11_experimental_precision",
        "gorgo11_experimental_kappa",
        "bays09_experimental_mixtures_mean_compatible",
        "T_space",
    ]

    if savedata:
        dataio.save_variables_default(locals(), variables_to_save)

        dataio.make_link_output_to_dropbox(dropbox_current_experiment_folder="memory_curves")

    plt.show()

    return locals()
def plots_memory_curves(data_pbs, generator_module=None):
    """
        Reload and plot memory curve of a feature code.
        Can use Marginal Fisher Information and fitted Mixture Model as well
    """

    #### SETUP
    #
    savefigs = True
    savedata = True

    plot_pcolor_fit_precision_to_fisherinfo = True
    plot_selected_memory_curves = True
    plot_best_memory_curves = True

    colormap = None  # or 'cubehelix'
    plt.rcParams["font.size"] = 16
    #
    #### /SETUP

    print "Order parameters: ", generator_module.dict_parameters_range.keys()

    result_all_precisions_mean = utils.nanmean(
        np.squeeze(data_pbs.dict_arrays["result_all_precisions"]["results"]), axis=-1
    )
    result_all_precisions_std = utils.nanstd(
        np.squeeze(data_pbs.dict_arrays["result_all_precisions"]["results"]), axis=-1
    )
    result_em_fits_mean = utils.nanmean(np.squeeze(data_pbs.dict_arrays["result_em_fits"]["results"]), axis=-1)
    result_em_fits_std = utils.nanstd(np.squeeze(data_pbs.dict_arrays["result_em_fits"]["results"]), axis=-1)
    result_marginal_inv_fi_mean = utils.nanmean(
        np.squeeze(data_pbs.dict_arrays["result_marginal_inv_fi"]["results"]), axis=-1
    )
    result_marginal_inv_fi_std = utils.nanstd(
        np.squeeze(data_pbs.dict_arrays["result_marginal_inv_fi"]["results"]), axis=-1
    )
    result_marginal_fi_mean = utils.nanmean(
        1.0 / np.squeeze(data_pbs.dict_arrays["result_marginal_inv_fi"]["results"]), axis=-1
    )
    result_marginal_fi_std = utils.nanstd(
        1.0 / np.squeeze(data_pbs.dict_arrays["result_marginal_inv_fi"]["results"]), axis=-1
    )

    M_space = data_pbs.loaded_data["parameters_uniques"]["M"]
    sigmax_space = data_pbs.loaded_data["parameters_uniques"]["sigmax"]
    T_space = data_pbs.loaded_data["datasets_list"][0]["T_space"]

    print M_space
    print sigmax_space
    print T_space
    print result_all_precisions_mean.shape, result_em_fits_mean.shape, result_marginal_inv_fi_mean.shape

    dataio = DataIO.DataIO(
        output_folder=generator_module.pbs_submission_infos["simul_out_dir"] + "/outputs/",
        label="global_" + dataset_infos["save_output_filename"],
    )

    ## Load Experimental data
    data_simult = load_experimental_data.load_data_simult(
        data_dir=os.path.normpath(
            os.path.join(os.path.split(load_experimental_data.__file__)[0], "../../experimental_data/")
        )
    )
    memory_experimental = data_simult["precision_nitems_theo"]

    data_bays2009 = load_experimental_data.load_data_bays09(
        data_dir=os.path.normpath(
            os.path.join(os.path.split(load_experimental_data.__file__)[0], "../../experimental_data/")
        ),
        fit_mixture_model=True,
    )
    bays09_experimental_mixtures_mean = data_bays2009["em_fits_nitems_arrays"]["mean"][1:]
    # add interpolated points for 3 and 5 items
    bays3 = (bays09_experimental_mixtures_mean[:, 2] + bays09_experimental_mixtures_mean[:, 1]) / 2.0
    bays5 = (bays09_experimental_mixtures_mean[:, -1] + bays09_experimental_mixtures_mean[:, -2]) / 2.0
    bays09_experimental_mixtures_mean_compatible = c_[
        bays09_experimental_mixtures_mean[:, :2], bays3, bays09_experimental_mixtures_mean[:, 2], bays5
    ]

    # Boost non-targets
    bays09_experimental_mixtures_mean_compatible[1] *= 1.5
    bays09_experimental_mixtures_mean_compatible[2] /= 1.5
    bays09_experimental_mixtures_mean_compatible /= np.sum(bays09_experimental_mixtures_mean_compatible, axis=0)

    # Force non target em fit mixture to be zero and not nan
    result_em_fits_mean[..., 0, 2] = 0
    result_em_fits_std[..., 0, 2] = 0

    # Compute some landscapes of fit!
    dist_diff_precision_margfi = np.sum(
        np.abs(result_all_precisions_mean * 2.0 - result_marginal_fi_mean[..., 0]) ** 2.0, axis=-1
    )
    dist_diff_precision_margfi_1item = (
        np.abs(result_all_precisions_mean[..., 0] * 2.0 - result_marginal_fi_mean[..., 0, 0]) ** 2.0
    )

    dist_diff_emkappa_margfi = np.sum(
        np.abs(result_em_fits_mean[..., 0] * 2.0 - result_marginal_fi_mean[..., 0]) ** 2.0, axis=-1
    )
    dist_ratio_emkappa_margfi = np.sum(
        np.abs((result_em_fits_mean[..., 0] * 2.0) / result_marginal_fi_mean[..., 0] - 1.0) ** 2.0, axis=-1
    )

    dist_diff_precision_experim = np.sum(np.abs(result_all_precisions_mean - memory_experimental) ** 2.0, axis=-1)
    dist_diff_precision_experim_1item = np.abs(result_all_precisions_mean[..., 0] - memory_experimental[0]) ** 2.0

    dist_diff_emkappa_experim = np.sum(np.abs(result_em_fits_mean[..., 0] - memory_experimental) ** 2.0, axis=-1)
    dist_diff_emkappa_experim_1item = np.abs(result_em_fits_mean[..., 0, 0] - memory_experimental[0]) ** 2.0

    dist_diff_margfi_experim_1item = np.abs(result_marginal_fi_mean[..., 0, 0] - memory_experimental[0]) ** 2.0

    dist_diff_emkappa_mixtures_bays09 = np.sum(
        np.sum((result_em_fits_mean[..., 1:4] - bays09_experimental_mixtures_mean_compatible.T) ** 2.0, axis=-1),
        axis=-1,
    )

    if plot_pcolor_fit_precision_to_fisherinfo:
        # Check fit between precision and fisher info
        utils.pcolor_2d_data(
            dist_diff_precision_margfi, log_scale=True, x=M_space, y=sigmax_space, xlabel="M", ylabel="sigmax"
        )

        if savefigs:
            dataio.save_current_figure("match_precision_margfi_log_pcolor_{label}_{unique_id}.pdf")

        # utils.pcolor_2d_data(dist_diff_precision_margfi, x=M_space, y=sigmax_space[2:], xlabel='M', ylabel='sigmax')
        # if savefigs:
        #    dataio.save_current_figure('match_precision_margfi_pcolor_{label}_{unique_id}.pdf')

        utils.pcolor_2d_data(
            dist_diff_precision_experim, x=M_space, y=sigmax_space, xlabel="M", ylabel="sigmax", log_scale=True
        )

        utils.pcolor_2d_data(
            dist_diff_emkappa_experim, x=M_space, y=sigmax_space, xlabel="M", ylabel="sigmax", log_scale=True
        )

        utils.pcolor_2d_data(
            dist_diff_precision_margfi
            * dist_diff_emkappa_margfi
            * dist_diff_precision_experim
            * dist_diff_emkappa_experim,
            x=M_space,
            y=sigmax_space,
            xlabel="M",
            ylabel="sigmax",
            log_scale=True,
        )

        utils.pcolor_2d_data(
            dist_diff_precision_margfi_1item, log_scale=True, x=M_space, y=sigmax_space, xlabel="M", ylabel="sigmax"
        )

        utils.pcolor_2d_data(
            dist_diff_precision_experim_1item, log_scale=True, x=M_space, y=sigmax_space, xlabel="M", ylabel="sigmax"
        )

        utils.pcolor_2d_data(
            dist_diff_emkappa_experim_1item, log_scale=True, x=M_space, y=sigmax_space, xlabel="M", ylabel="sigmax"
        )

        utils.pcolor_2d_data(
            dist_diff_margfi_experim_1item, log_scale=True, x=M_space, y=sigmax_space, xlabel="M", ylabel="sigmax"
        )

        utils.pcolor_2d_data(
            dist_diff_emkappa_mixtures_bays09, log_scale=False, x=M_space, y=sigmax_space, xlabel="M", ylabel="sigmax"
        )

    if plot_selected_memory_curves:
        selected_values = [[100, 0.8], [200, 0.27], [100, 0.1], [200, 0.8], [100, 0.17], [100, 0.08], [50, 0.21]]

        for current_values in selected_values:
            # Find the indices
            M_i = np.argmin(np.abs(current_values[0] - M_space))
            sigmax_i = np.argmin(np.abs(current_values[1] - sigmax_space))

            ax = utils.plot_mean_std_area(
                T_space, memory_experimental, np.zeros(T_space.size), linewidth=3, fmt="o-", markersize=8
            )

            ax = utils.plot_mean_std_area(
                T_space,
                result_all_precisions_mean[M_i, sigmax_i],
                result_all_precisions_std[M_i, sigmax_i],
                ax_handle=ax,
                linewidth=3,
                fmt="o-",
                markersize=8,
            )

            ax = utils.plot_mean_std_area(
                T_space,
                0.5 * result_marginal_fi_mean[..., 0][M_i, sigmax_i],
                0.5 * result_marginal_fi_std[..., 0][M_i, sigmax_i],
                ax_handle=ax,
                linewidth=3,
                fmt="o-",
                markersize=8,
            )

            ax = utils.plot_mean_std_area(
                T_space,
                result_em_fits_mean[..., 0][M_i, sigmax_i],
                result_em_fits_std[..., 0][M_i, sigmax_i],
                ax_handle=ax,
                xlabel="Number of items",
                ylabel="Inverse variance $[rad^{-2}]$",
                linewidth=3,
                fmt="o-",
                markersize=8,
            )

            # ax.set_title('M %d, sigmax %.2f' % (M_space[M_i], sigmax_space[sigmax_i]))
            plt.legend(["Experimental data", "Precision of samples", "Marginal Fisher Information", "Fitted kappa"])

            ax.set_xlim([0.9, 5.1])
            ax.set_xticks(range(1, 6))
            ax.set_xticklabels(range(1, 6))

            if savefigs:
                dataio.save_current_figure(
                    "memorycurves_M%dsigmax%.2f_{label}_{unique_id}.pdf" % (M_space[M_i], sigmax_space[sigmax_i])
                )

    def em_plot_paper(sigmax_i, M_i):
        f, ax = plt.subplots()

        # Right axis, mixture probabilities
        utils.plot_mean_std_area(
            T_space,
            result_em_fits_mean[..., 1][M_i, sigmax_i],
            result_em_fits_std[..., 1][M_i, sigmax_i],
            xlabel="Number of items",
            ylabel="Mixture probabilities",
            ax_handle=ax,
            linewidth=3,
            fmt="o-",
            markersize=5,
            label="Target",
        )
        utils.plot_mean_std_area(
            T_space,
            result_em_fits_mean[..., 2][M_i, sigmax_i],
            result_em_fits_std[..., 2][M_i, sigmax_i],
            xlabel="Number of items",
            ylabel="Mixture probabilities",
            ax_handle=ax,
            linewidth=3,
            fmt="o-",
            markersize=5,
            label="Nontarget",
        )
        utils.plot_mean_std_area(
            T_space,
            result_em_fits_mean[..., 3][M_i, sigmax_i],
            result_em_fits_std[..., 3][M_i, sigmax_i],
            xlabel="Number of items",
            ylabel="Mixture probabilities",
            ax_handle=ax,
            linewidth=3,
            fmt="o-",
            markersize=5,
            label="Random",
        )

        ax.legend(prop={"size": 15})

        ax.set_title("M %d, sigmax %.2f" % (M_space[M_i], sigmax_space[sigmax_i]))
        ax.set_xlim([1.0, 5.0])
        ax.set_ylim([0.0, 1.1])
        ax.set_xticks(range(1, 6))
        ax.set_xticklabels(range(1, 6))

        f.canvas.draw()

        if savefigs:
            dataio.save_current_figure(
                "memorycurves_emfits_paper_M%.2fsigmax%.2f_{label}_{unique_id}.pdf"
                % (M_space[M_i], sigmax_space[sigmax_i])
            )

    if plot_best_memory_curves:

        # Best mixtures fit
        best_axis2_i_all = np.argmin(dist_diff_emkappa_mixtures_bays09, axis=1)
        for axis1_i, best_axis2_i in enumerate(best_axis2_i_all):
            em_plot_paper(best_axis2_i, axis1_i)

    all_args = data_pbs.loaded_data["args_list"]
    variables_to_save = [
        "result_all_precisions_mean",
        "result_em_fits_mean",
        "result_marginal_inv_fi_mean",
        "result_all_precisions_std",
        "result_em_fits_std",
        "result_marginal_inv_fi_std",
        "result_marginal_fi_mean",
        "result_marginal_fi_std",
        "M_space",
        "sigmax_space",
        "T_space",
        "all_args",
    ]

    if savedata:
        dataio.save_variables(variables_to_save, locals())

    # Make link to Dropbox
    dataio.make_link_output_to_dropbox(dropbox_current_experiment_folder="memory_curves")

    plt.show()

    return locals()
def plots_memory_curves(data_pbs, generator_module=None):
    '''
        Reload and plot memory curve of a conjunctive code.
        Can use Marginal Fisher Information and fitted Mixture Model as well
    '''

    #### SETUP
    #
    savefigs = True
    savedata = True

    plot_pcolor_fit_precision_to_fisherinfo = True
    plot_selected_memory_curves = False
    plot_best_memory_curves = True

    colormap = None  # or 'cubehelix'
    plt.rcParams['font.size'] = 16
    #
    #### /SETUP

    print "Order parameters: ", generator_module.dict_parameters_range.keys()

    result_all_precisions_mean = utils.nanmean(np.squeeze(data_pbs.dict_arrays['result_all_precisions']['results']), axis=-1)
    result_all_precisions_std = utils.nanstd(np.squeeze(data_pbs.dict_arrays['result_all_precisions']['results']), axis=-1)
    result_em_fits_mean = utils.nanmean(np.squeeze(data_pbs.dict_arrays['result_em_fits']['results']), axis=-1)
    result_em_fits_std = utils.nanstd(np.squeeze(data_pbs.dict_arrays['result_em_fits']['results']), axis=-1)
    result_marginal_inv_fi_mean = utils.nanmean(np.squeeze(data_pbs.dict_arrays['result_marginal_inv_fi']['results']), axis=-1)
    result_marginal_inv_fi_std = utils.nanstd(np.squeeze(data_pbs.dict_arrays['result_marginal_inv_fi']['results']), axis=-1)
    result_marginal_fi_mean = utils.nanmean(1./np.squeeze(data_pbs.dict_arrays['result_marginal_inv_fi']['results']), axis=-1)
    result_marginal_fi_std = utils.nanstd(1./np.squeeze(data_pbs.dict_arrays['result_marginal_inv_fi']['results']), axis=-1)

    M_space = data_pbs.loaded_data['parameters_uniques']['M']
    sigmax_space = data_pbs.loaded_data['parameters_uniques']['sigmax']
    T_space = data_pbs.loaded_data['datasets_list'][0]['T_space']

    print M_space
    print sigmax_space
    print T_space
    print result_all_precisions_mean.shape, result_em_fits_mean.shape, result_marginal_inv_fi_mean.shape

    dataio = DataIO.DataIO(output_folder=generator_module.pbs_submission_infos['simul_out_dir'] + '/outputs/', label='global_' + dataset_infos['save_output_filename'])

    ## Load Experimental data
    experim_datadir = os.environ.get('WORKDIR_DROP', os.path.split(load_experimental_data.__file__)[0])

    data_simult = load_experimental_data.load_data_simult(data_dir=os.path.normpath(os.path.join(experim_datadir, '../../experimental_data/')), fit_mixture_model=True)
    gorgo11_experimental_precision = data_simult['precision_nitems_theo']
    gorgo11_experimental_kappa = np.array([data_s['kappa'] for _, data_s in data_simult['em_fits_nitems']['mean'].items()])
    gorgo11_experimental_emfits_mean = np.array([[data[key] for _, data in data_simult['em_fits_nitems']['mean'].items()] for key in ['kappa', 'mixt_target', 'mixt_nontargets', 'mixt_random']])
    gorgo11_experimental_emfits_std = np.array([[data[key] for _, data in data_simult['em_fits_nitems']['std'].items()] for key in ['kappa', 'mixt_target', 'mixt_nontargets', 'mixt_random']])
    gorgo11_experimental_emfits_sem = gorgo11_experimental_emfits_std/np.sqrt(np.unique(data_simult['subject']).size)


    data_bays2009 = load_experimental_data.load_data_bays09(data_dir=os.path.normpath(os.path.join(experim_datadir, '../../experimental_data/')), fit_mixture_model=True)
    # add interpolated points for 3 and 5 items
    emfit_mean_intpfct = spint.interp1d(np.unique(data_bays2009['n_items']), data_bays2009['em_fits_nitems_arrays']['mean'])
    bays09_experimental_mixtures_mean_compatible = emfit_mean_intpfct(np.arange(1, 6))
    emfit_std_intpfct = spint.interp1d(np.unique(data_bays2009['n_items']), data_bays2009['em_fits_nitems_arrays']['std'])
    bays09_experimental_mixtures_std_compatible = emfit_std_intpfct(np.arange(1, 6))

    # Boost non-targets
    # bays09_experimental_mixtures_mean_compatible[1] *= 1.5
    # bays09_experimental_mixtures_mean_compatible[2] /= 1.5
    # bays09_experimental_mixtures_mean_compatible /= np.sum(bays09_experimental_mixtures_mean_compatible, axis=0)

    # Force non target em fit mixture to be zero and not nan
    result_em_fits_mean[..., 0, 2] = 0
    result_em_fits_std[..., 0, 2] = 0

    # Compute some landscapes of fit!
    dist_diff_precision_margfi = np.sum(np.abs(result_all_precisions_mean*2. - result_marginal_fi_mean[..., 0])**2., axis=-1)
    dist_ratio_precision_margfi = np.sum(np.abs((result_all_precisions_mean*2.)/result_marginal_fi_mean[..., 0] - 1.0)**2., axis=-1)
    dist_diff_emkappa_margfi = np.sum(np.abs(result_em_fits_mean[..., 0]*2. - result_marginal_fi_mean[..., 0])**2., axis=-1)
    dist_ratio_emkappa_margfi = np.sum(np.abs((result_em_fits_mean[..., 0]*2.)/result_marginal_fi_mean[..., 0] - 1.0)**2., axis=-1)

    dist_diff_precision_experim = np.sum(np.abs(result_all_precisions_mean - gorgo11_experimental_precision)**2., axis=-1)
    dist_diff_emkappa_experim = np.sum(np.abs(result_em_fits_mean[..., 0] - gorgo11_experimental_kappa)**2., axis=-1)

    dist_diff_precision_experim_1item = np.abs(result_all_precisions_mean[..., 0] - gorgo11_experimental_precision[0])**2.
    dist_diff_precision_experim_2item = np.abs(result_all_precisions_mean[..., 1] - gorgo11_experimental_precision[1])**2.
    dist_diff_precision_margfi_1item = np.abs(result_all_precisions_mean[..., 0]*2. - result_marginal_fi_mean[..., 0, 0])**2.
    dist_diff_emkappa_experim_1item = np.abs(result_em_fits_mean[..., 0, 0] - gorgo11_experimental_kappa[0])**2.
    dist_diff_margfi_experim_1item = np.abs(result_marginal_fi_mean[..., 0, 0] - gorgo11_experimental_precision[0])**2.

    dist_diff_emkappa_mixtures_bays09 = np.sum(np.sum((result_em_fits_mean[..., 1:4] - bays09_experimental_mixtures_mean_compatible[1:].T)**2., axis=-1), axis=-1)

    if plot_pcolor_fit_precision_to_fisherinfo:
        # Check fit between precision and fisher info
        utils.pcolor_2d_data(dist_diff_precision_margfi, log_scale=True, x=M_space, y=sigmax_space, xlabel='M', ylabel='sigmax')

        if savefigs:
            dataio.save_current_figure('match_precision_margfi_log_pcolor_{label}_{unique_id}.pdf')

        # utils.pcolor_2d_data(dist_diff_precision_margfi, x=M_space, y=sigmax_space[2:], xlabel='M', ylabel='sigmax')
        # if savefigs:
        #    dataio.save_current_figure('match_precision_margfi_pcolor_{label}_{unique_id}.pdf')

        utils.pcolor_2d_data(dist_ratio_precision_margfi, x=M_space, y=sigmax_space, xlabel='M', ylabel='sigmax', log_scale=True)
        if savefigs:
            dataio.save_current_figure('match_ratio_precision_margfi_log_pcolor_{label}_{unique_id}.pdf')

        utils.pcolor_2d_data(dist_diff_emkappa_margfi, x=M_space, y=sigmax_space, xlabel='M', ylabel='sigmax', log_scale=True)
        if savefigs:
            dataio.save_current_figure('match_diff_emkappa_margfi_log_pcolor_{label}_{unique_id}.pdf')

        utils.pcolor_2d_data(dist_ratio_emkappa_margfi, x=M_space, y=sigmax_space, xlabel='M', ylabel='sigmax', log_scale=True)
        if savefigs:
            dataio.save_current_figure('match_ratio_emkappa_margfi_log_pcolor_{label}_{unique_id}.pdf')


        utils.pcolor_2d_data(dist_diff_precision_experim, x=M_space, y=sigmax_space, xlabel='M', ylabel='sigmax', log_scale=True)
        if savefigs:
            dataio.save_current_figure('match_diff_precision_experim_log_pcolor_{label}_{unique_id}.pdf')


        utils.pcolor_2d_data(dist_diff_emkappa_experim, x=M_space, y=sigmax_space, xlabel='M', ylabel='sigmax', log_scale=True)
        if savefigs:
            dataio.save_current_figure('match_diff_emkappa_experim_log_pcolor_{label}_{unique_id}.pdf')

        utils.pcolor_2d_data(dist_diff_precision_margfi*dist_diff_emkappa_margfi*dist_diff_precision_experim*dist_diff_emkappa_experim, x=M_space, y=sigmax_space, xlabel='M', ylabel='sigmax', log_scale=True)
        if savefigs:
            dataio.save_current_figure('match_bigmultiplication_log_pcolor_{label}_{unique_id}.pdf')

        utils.pcolor_2d_data(dist_diff_precision_margfi_1item, log_scale=True, x=M_space, y=sigmax_space, xlabel='M', ylabel='sigmax')
        if savefigs:
            dataio.save_current_figure('match_diff_precision_margfi_1item_log_pcolor_{label}_{unique_id}.pdf')


        utils.pcolor_2d_data(dist_diff_precision_experim_1item, log_scale=True, x=M_space, y=sigmax_space, xlabel='M', ylabel='sigmax')
        if savefigs:
            dataio.save_current_figure('match_diff_precision_experim_1item_log_pcolor_{label}_{unique_id}.pdf')


        utils.pcolor_2d_data(dist_diff_emkappa_experim_1item, log_scale=True, x=M_space, y=sigmax_space, xlabel='M', ylabel='sigmax')
        if savefigs:
            dataio.save_current_figure('match_diff_emkappa_experim_1item_log_pcolor_{label}_{unique_id}.pdf')


        utils.pcolor_2d_data(dist_diff_margfi_experim_1item, log_scale=True, x=M_space, y=sigmax_space, xlabel='M', ylabel='sigmax')
        if savefigs:
            dataio.save_current_figure('match_diff_margfi_experim_1item_log_pcolor_{label}_{unique_id}.pdf')

        utils.pcolor_2d_data(dist_diff_emkappa_mixtures_bays09, log_scale=True, x=M_space, y=sigmax_space, xlabel='M', ylabel='sigmax')
        if savefigs:
            dataio.save_current_figure('match_diff_mixtures_experbays09_pcolor_{label}_{unique_id}.pdf')


    # Macro plot
    def mem_plot_precision(sigmax_i, M_i):
        ax = utils.plot_mean_std_area(T_space, gorgo11_experimental_precision, np.zeros(T_space.size), linewidth=3, fmt='o-', markersize=8, label='Experimental data')

        ax = utils.plot_mean_std_area(T_space, result_all_precisions_mean[M_i, sigmax_i], result_all_precisions_std[M_i, sigmax_i], ax_handle=ax, linewidth=3, fmt='o-', markersize=8, label='Precision of samples')

        ax = utils.plot_mean_std_area(T_space, 0.5*result_marginal_fi_mean[..., 0][M_i, sigmax_i], 0.5*result_marginal_fi_std[..., 0][M_i, sigmax_i], ax_handle=ax, linewidth=3, fmt='o-', markersize=8, label='Marginal Fisher Information')

        # ax = utils.plot_mean_std_area(T_space, result_em_fits_mean[..., 0][M_i, sigmax_i], result_em_fits_std[..., 0][M_i, sigmax_i], ax_handle=ax, xlabel='Number of items', ylabel="Inverse variance $[rad^{-2}]$", linewidth=3, fmt='o-', markersize=8, label='Fitted kappa')

        ax.set_title('M %d, sigmax %.2f' % (M_space[M_i], sigmax_space[sigmax_i]))
        ax.legend()
        ax.set_xlim([0.9, 5.1])
        ax.set_xticks(range(1, 6))
        ax.set_xticklabels(range(1, 6))

        ax.get_figure().canvas.draw()

        if savefigs:
            dataio.save_current_figure('memorycurves_precision_M%dsigmax%.2f_{label}_{unique_id}.pdf' % (M_space[M_i], sigmax_space[sigmax_i]))


    def mem_plot_kappa(sigmax_i, M_i, experim_data_mean, experim_data_std=None):
        ax = utils.plot_mean_std_area(T_space, experim_data_mean, experim_data_std, linewidth=3, fmt='o-', markersize=8, label='Experimental data')

        ax = utils.plot_mean_std_area(T_space, result_em_fits_mean[..., 0][M_i, sigmax_i], result_em_fits_std[..., 0][M_i, sigmax_i], xlabel='Number of items', ylabel="Inverse variance $[rad^{-2}]$", ax_handle=ax, linewidth=3, fmt='o-', markersize=8, label='Fitted kappa')

        # ax = utils.plot_mean_std_area(T_space, 0.5*result_marginal_fi_mean[..., 0][M_i, sigmax_i], 0.5*result_marginal_fi_std[..., 0][M_i, sigmax_i], ax_handle=ax, linewidth=3, fmt='o-', markersize=8, label='Marginal Fisher Information')

        ax.set_title('M %d, sigmax %.2f' % (M_space[M_i], sigmax_space[sigmax_i]))
        ax.legend()
        ax.set_xlim([0.9, 5.1])
        ax.set_xticks(range(1, 6))
        ax.set_xticklabels(range(1, 6))

        ax.get_figure().canvas.draw()

        if savefigs:
            dataio.save_current_figure('memorycurves_kappa_M%dsigmax%.2f_{label}_{unique_id}.pdf' % (M_space[M_i], sigmax_space[sigmax_i]))

    def em_plot_paper(sigmax_i, M_i):
        f, ax = plt.subplots()

        # Right axis, mixture probabilities
        utils.plot_mean_std_area(T_space, result_em_fits_mean[..., 1][M_i, sigmax_i], result_em_fits_std[..., 1][M_i, sigmax_i], xlabel='Number of items', ylabel="Mixture probabilities", ax_handle=ax, linewidth=3, fmt='o-', markersize=5, label='Target')
        utils.plot_mean_std_area(T_space, result_em_fits_mean[..., 2][M_i, sigmax_i], result_em_fits_std[..., 2][M_i, sigmax_i], xlabel='Number of items', ylabel="Mixture probabilities", ax_handle=ax, linewidth=3, fmt='o-', markersize=5, label='Nontarget')
        utils.plot_mean_std_area(T_space, result_em_fits_mean[..., 3][M_i, sigmax_i], result_em_fits_std[..., 3][M_i, sigmax_i], xlabel='Number of items', ylabel="Mixture probabilities", ax_handle=ax, linewidth=3, fmt='o-', markersize=5, label='Random')

        ax.legend(prop={'size':15})

        ax.set_title('M %d, sigmax %.2f' % (M_space[M_i], sigmax_space[sigmax_i]))
        ax.set_xlim([1.0, 5.0])
        ax.set_ylim([0.0, 1.1])
        ax.set_xticks(range(1, 6))
        ax.set_xticklabels(range(1, 6))

        f.canvas.draw()

        if savefigs:
            dataio.save_current_figure('memorycurves_emfits_paper_M%.2fsigmax%.2f_{label}_{unique_id}.pdf' % (M_space[M_i], sigmax_space[sigmax_i]))


    if plot_selected_memory_curves:
        selected_values = [[100, 0.1], [196, 0.28], [64, 0.05]]

        for current_values in selected_values:
            # Find the indices
            M_i         = np.argmin(np.abs(current_values[0] - M_space))
            sigmax_i    = np.argmin(np.abs(current_values[1] - sigmax_space))

            mem_plot_precision(sigmax_i, M_i)

            mem_plot_kappa(sigmax_i, M_i)


    if plot_best_memory_curves:
        # Best precision fit
        best_axis2_i_all = np.argmin(dist_diff_precision_experim, axis=1)

        for axis1_i, best_axis2_i in enumerate(best_axis2_i_all):
            mem_plot_precision(best_axis2_i, axis1_i)

        # Best kappa fit
        best_axis2_i_all = np.argmin(dist_diff_emkappa_experim, axis=1)

        for axis1_i, best_axis2_i in enumerate(best_axis2_i_all):
            mem_plot_kappa(best_axis2_i, axis1_i, gorgo11_experimental_emfits_mean[0], gorgo11_experimental_emfits_std[0])

            mem_plot_precision(best_axis2_i, axis1_i)

        # Best mixtures fit for Bays09
        best_axis2_i_all = np.argmin(dist_diff_emkappa_mixtures_bays09, axis=1)

        for axis1_i, best_axis2_i in enumerate(best_axis2_i_all):
            em_plot_paper(best_axis2_i, axis1_i)



    all_args = data_pbs.loaded_data['args_list']
    variables_to_save = ['gorgo11_experimental_precision', 'gorgo11_experimental_kappa']

    if savedata:
        dataio.save_variables_default(locals(), variables_to_save)

        dataio.make_link_output_to_dropbox(dropbox_current_experiment_folder='memory_curves')

    plt.show()

    return locals()