def plot_agg()

in automl21/enr.py [0:0]

• 62 lines of code
• 2 McCabe index (conditional complexity)

    def plot_agg(self, n_iter, tag='t', n_sample=50):
        ista_norms, ista_aa_norms, neural_norms = [], [], []
        ista_objs, ista_aa_objs, neural_objs = [], [], []
        for seed in range(n_sample):
            inst = self.sample_single_inst(seed=seed)

            x_cvxpy = self.solve_cvxpy(inst)
            cvxpy_obj = self.obj(inst, x_cvxpy).cpu()

            d_ISTA = self.solve_ISTA(inst, n_iter=n_iter, track_iterates=True)
            d_ISTA_AA = self.solve_ISTA_AA(
                inst, n_iter=n_iter, track_iterates=True)
            d_ISTA_neural = self.solve_ISTA_neural(
                inst, n_iter=n_iter, track_iterates=True)

            ista_norms.append(d_ISTA['rel_residual_norms'])
            ista_aa_norms.append(d_ISTA_AA['rel_residual_norms'])
            neural_norms.append(d_ISTA_neural['rel_residual_norms'])

            ista_objs.append((cvxpy_obj-d_ISTA['objs'].cpu())**2)
            ista_aa_objs.append((cvxpy_obj-d_ISTA_AA['objs'].cpu())**2)
            neural_objs.append((cvxpy_obj-d_ISTA_neural['objs'].cpu())**2)

        ista_norms = torch.tensor(ista_norms)
        ista_aa_norms = torch.tensor(ista_aa_norms)
        neural_norms = torch.tensor(neural_norms)
        ista_objs = torch.stack(ista_objs)
        ista_aa_objs = torch.stack(ista_aa_objs)
        neural_objs = torch.stack(neural_objs)

        nrow, ncol = 1, 2
        fig, axs = plt.subplots(nrow, ncol, figsize=(6*ncol, 4*nrow))

        ax = axs[0]
        mean, std = ista_norms.mean(dim=0), ista_norms.std(dim=0)
        xs = np.arange(len(mean))
        l, = ax.plot(xs, mean, label='ISTA')
        ax.fill_between(xs, mean-std, mean+std, color=l.get_color(), alpha=0.2)

        mean, std = ista_aa_norms.mean(dim=0), ista_aa_norms.std(dim=0)
        l, = ax.plot(xs, mean, label='ISTA+AA')
        ax.fill_between(xs, mean-std, mean+std, color=l.get_color(), alpha=0.2)

        mean, std = neural_norms.mean(dim=0), neural_norms.std(dim=0)
        l, = ax.plot(xs, mean, label='ISTA+Neural')
        ax.fill_between(xs, mean-std, mean+std, color=l.get_color(), alpha=0.2)

        ax.set_yscale('log')
        ax.set_xlabel('Fixed-Point Iteration')
        ax.set_ylabel('Fixed-Point Residual')

        ax = axs[1]
        min_obj_dist = np.inf
        mean, std = ista_objs.mean(dim=0), ista_objs.std(dim=0)
        min_obj_dist = min(min_obj_dist, mean.min())
        l, = ax.plot(xs, mean, label='ISTA')
        ax.fill_between(xs, mean-std, mean+std, color=l.get_color(), alpha=0.2)

        mean, std = ista_aa_objs.mean(dim=0), ista_aa_objs.std(dim=0)
        min_obj_dist = min(min_obj_dist, mean.min())
        l, = ax.plot(xs, mean, label='ISTA+AA')
        ax.fill_between(xs, mean-std, mean+std, color=l.get_color(), alpha=0.2)

        mean, std = neural_objs.mean(dim=0), neural_objs.std(dim=0)
        min_obj_dist = min(min_obj_dist, mean.min())
        l, = ax.plot(xs, mean, label='ISTA+Neural')
        ax.fill_between(xs, mean-std, mean+std, color=l.get_color(), alpha=0.2)
        # ax.set_ylim(min_obj_dist, 1e3)

        ax.set_yscale('log')
        ax.set_xlabel('Fixed-Point Iteration')
        ax.set_ylabel('Objective Distance')

        fname = 'agg_' + tag + '.png'
        fig.tight_layout()
        fig.savefig(fname, transparent=True)
        plt.close(fig)
        os.system(f"convert {fname} -trim {fname}")