Comparisons.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "01b74424",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "import numpy as np\n",
    "import copy\n",
    "import matplotlib.pyplot as plt\n",
    "from src.group_model import EarlyStopping, GaussianSimulation, Simulation, groupModel, groupTrainer, summary_plot"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0743f189",
   "metadata": {},
   "outputs": [],
   "source": [
    "from yaglm.toy_data import sample_sparse_lin_reg\n",
    "\n",
    "from yaglm.GlmTuned import GlmCV, GlmTrainMetric\n",
    "\n",
    "from yaglm.config.loss import Huber\n",
    "from yaglm.config.penalty import Lasso, GroupLasso\n",
    "from yaglm.config.flavor import Adaptive, NonConvex\n",
    "\n",
    "from yaglm.metrics.info_criteria import InfoCriteria\n",
    "from yaglm.infer.Inferencer import Inferencer\n",
    "from yaglm.infer.lin_reg_noise_var import ViaRidge\n",
    "\n",
    "from yaglm.Glm import Glm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "020ae7dd",
   "metadata": {},
   "outputs": [],
   "source": [
    "from iterreg.sparse import dual_primal\n",
    "from iterreg.utils import datadriven_ratio"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "99a8542c",
   "metadata": {},
   "outputs": [],
   "source": [
    "def one_fit_group_lasso(sim, penalty, group_size):\n",
    "    res = Glm(loss='lin_reg', penalty=GroupLasso(groups=[range(i*group_size,i*group_size+group_size) for i in range((sim.p)//group_size)],pen_val=penalty)).fit(sim.X, sim.y)\n",
    "    val_error = ((sim.y_val - res.decision_function(sim.X_val))**2).mean().item()\n",
    "    return (val_error, penalty)\n",
    "\n",
    "def val_group_lasso(sim, penalties, group_size=4):\n",
    "    val_errors = []\n",
    "    for penalty in penalties:\n",
    "        val_error = one_fit_group_lasso(sim, penalty, group_size)\n",
    "        val_errors.append(val_error)\n",
    "    val_errors.sort()\n",
    "    \n",
    "    res = Glm(loss='lin_reg', penalty=GroupLasso(groups=[range(i*group_size,i*group_size+group_size) for i in range((sim.p)//group_size)],pen_val=val_errors[0][1])).fit(sim.X, sim.y)\n",
    "    val_error = ((sim.y_val - res.decision_function(sim.X_val))**2).mean().item()\n",
    "    est_err = ((res.coef_ - sim.w_star.numpy())**2).sum()\n",
    "    return est_err, val_error,res, val_errors"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f57a4647",
   "metadata": {},
   "outputs": [],
   "source": [
    "class groupTrainer:\n",
    "    def __init__(self, model, sim, is_two_lr=False, varepsilon = 1e-12, tol_on_u = 5e-3, \n",
    "                 lr=0.01, is_small_train=False, is_monitor_u_diff=False, verbose=False):\n",
    "        self.model = model\n",
    "        self.sim = sim\n",
    "        self.lr = lr\n",
    "        self.optimizer = torch.optim.SGD(self.model.u.parameters(), lr=lr)\n",
    "        self.all_optimizer = torch.optim.SGD(self.model.parameters(), lr=lr)\n",
    "\n",
    "        self.loss_criterion = torch.nn.MSELoss()\n",
    "        self.monitor = {'w':[], 'u':[], 'v':[]}\n",
    "        self.loss = []\n",
    "        self.early_stopping = EarlyStopping(patience=500)\n",
    "        self.flag = True\n",
    "        self.early_stopped_epoch = None\n",
    "        self.early_stopped_model = None\n",
    "        \n",
    "        self.val_err = []\n",
    "        self.params_est_err = []\n",
    "        self.dir_monitor = []\n",
    "        \n",
    "        self.change_epoch = 0\n",
    "        self.is_two_lr = is_two_lr\n",
    "        self.tol_on_u = tol_on_u\n",
    "        self.is_small_train = is_small_train\n",
    "        self.is_monitor_u_diff = is_monitor_u_diff\n",
    "\n",
    "        self.num_groups = self.sim.p//self.model.group_size\n",
    "        self.v_optimizers = []\n",
    "        for i in range(self.num_groups):\n",
    "            tmp_optim = torch.optim.SGD(self.model.vs[i].parameters(), lr=lr)\n",
    "            self.v_optimizers.append(tmp_optim)\n",
    "        self.varepsilon = varepsilon\n",
    "        self.verbose = verbose\n",
    "    def _one_epoch(self):\n",
    "        \n",
    "        y_pred = self.model(self.sim.X)\n",
    "        loss = self.loss_criterion(y_pred, self.sim.y)\n",
    "        self.optimizer.zero_grad()\n",
    "        for i in range(self.num_groups):\n",
    "            self.v_optimizers[i].zero_grad()\n",
    "        loss.backward()\n",
    "        for i in range(self.num_groups):\n",
    "            for g in self.v_optimizers[i].param_groups:\n",
    "                g['lr'] = 1./self.model.u.weight.data.detach().clone()[0,i]**(self.model.depth*2)\n",
    "            self.v_optimizers[i].step()\n",
    "            self.model.vs[i].weight.data = self.model.vs[i].weight.data.clone() / (self.model.vs[i].weight.data.clone() ** 2).sum().sqrt()\n",
    "\n",
    "        self.optimizer.step()\n",
    "\n",
    "        est_err = self._monitor()\n",
    "        self.loss.append(loss.item())\n",
    "        return loss.item(), est_err, self.model.u.weight.data.detach().clone()\n",
    "    \n",
    "    def _small_lr_one_epoch(self):\n",
    "        y_pred = self.model(self.sim.X)\n",
    "        loss = self.loss_criterion(y_pred, self.sim.y)\n",
    "        self.all_optimizer.zero_grad()\n",
    "        loss.backward()\n",
    "        self.all_optimizer.step()\n",
    "        \n",
    "        for i in range(self.sim.p//self.model.group_size):\n",
    "            self.model.vs[i].weight.data = self.model.vs[i].weight.data / (self.model.vs[i].weight.data ** 2).sum().sqrt()\n",
    "\n",
    "        est_err = self._monitor()\n",
    "        self.loss.append(loss.item())\n",
    "        return loss.item(), est_err\n",
    "    \n",
    "    def _monitor(self):\n",
    "        k = self.sim.k\n",
    "        num = k//self.model.group_size\n",
    "        params = self.model.get_params().squeeze(0)\n",
    "        est_err = ((self.sim.w_star - self.model.get_params().detach())**2).sum().item() # self.loss_criterion(params, self.sim.w_star).item()\n",
    "        \n",
    "        val_err = ((self.sim.y_val - np.matmul(self.sim.X_val, self.model.get_params().detach().numpy()))**2).sum().item()/sim.y_val.shape[0]\n",
    "\n",
    "        self.val_err.append(val_err)\n",
    "        self.params_est_err.append(est_err)\n",
    "        params = params.numpy().tolist()\n",
    "        \n",
    "        u = self.model.u.weight.detach().squeeze(0).numpy().tolist()\n",
    "        v = [self.model.vs[i].weight.detach().squeeze(0).numpy().tolist() for i in range(self.model.num_groups)]\n",
    "        v = [item for sublist in v for item in sublist]\n",
    "        \n",
    "        self.monitor['w'].append([*params[:k],max([abs(x) for x in params[k:]])])\n",
    "        self.monitor['u'].append([*u[:num],max([abs(x) for x in u[num:]])])\n",
    "        self.monitor['v'].append([*v[:k],max([abs(x) for x in v[k:]])])\n",
    "        \n",
    "        return est_err\n",
    "    \n",
    "    def _small_train(self, epochs):\n",
    "        for epoch in range(epochs):\n",
    "            loss, est_err = self._small_lr_one_epoch()\n",
    "            if self.verbose:\n",
    "                print(f'{epoch}/{epochs}, loss: {loss:.4f}, est error: {est_err:.4f}')\n",
    "\n",
    "            \n",
    "    def _optimal_train(self, epochs):\n",
    "        for epoch in range(epochs):\n",
    "            loss, est_err, _ = self._one_epoch()\n",
    "            if self.verbose:\n",
    "                print(f'{epoch}/{epochs}, loss: {loss:.4f}, est error: {est_err:.4f}')\n",
    "            \n",
    "    def _two_train(self, epochs):\n",
    "        prev_u = 1\n",
    "        flag = True\n",
    "        for epoch in range(epochs):\n",
    "            if epoch < 200:\n",
    "                loss, est_err, u = self._one_epoch()\n",
    "                u_diff = ((u-prev_u).abs()/np.abs(prev_u + self.varepsilon)).max()\n",
    "                prev_u = u\n",
    "                if self.is_monitor_u_diff:\n",
    "                    print(f'{u_diff.item()}')\n",
    "                if u_diff < self.tol_on_u:\n",
    "                    flag = False\n",
    "                    print(f'Change epoch: {epoch} with {u_diff.item()}')\n",
    "                    self.change_epoch = epoch\n",
    "            else:\n",
    "                loss, est_err = self._small_lr_one_epoch()\n",
    "\n",
    "            if self.verbose:\n",
    "                print(f'{epoch}/{epochs}, loss: {loss:.4f}, est error: {est_err:.4f}')  \n",
    "            \n",
    "    def train(self, epochs=500):\n",
    "        if self.is_two_lr:\n",
    "            self._two_train(epochs=epochs)\n",
    "        elif self.is_small_train:\n",
    "            self._small_train(epochs=epochs)\n",
    "        else:\n",
    "            self._optimal_train(epochs=epochs)\n",
    "        self.get_dir()\n",
    "\n",
    "    def transpose_monitor(self):\n",
    "        transposed_monitor = {}\n",
    "        for key, item in self.monitor.items():\n",
    "            item_t = [[one[i] for one in item] for i in range(len(item[0]))]\n",
    "            transposed_monitor[key] = item_t\n",
    "        return transposed_monitor\n",
    "    \n",
    "    def get_dir(self):\n",
    "        inner_product = []\n",
    "        group_size = self.model.group_size\n",
    "        for i in range(len(self.monitor['v'])):\n",
    "            vec = np.array(self.monitor['v'][i][:-1])\n",
    "            n = vec.shape[0]\n",
    "            one_step = []\n",
    "            for j in range(n//self.model.group_size):\n",
    "                tmp_vec = vec[j*group_size : (j+1)*group_size]\n",
    "                tmp_support = self.sim.support[j*group_size : (j+1)*group_size]\n",
    "                res = (tmp_vec * tmp_support).sum() / (np.sqrt((tmp_vec**2).sum()) * np.sqrt((tmp_support**2).sum()))\n",
    "                one_step.append(res)\n",
    "            inner_product.append(one_step)\n",
    "        self.dir = inner_product"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "96284f61",
   "metadata": {},
   "outputs": [],
   "source": [
    "def proximal_op(w, tau):\n",
    "    group_size = 4\n",
    "    w = w.reshape(-1, group_size)\n",
    "    norms = np.sqrt((w**2).sum(axis=1))\n",
    "    norms = np.maximum(norms,tau)\n",
    "    w = (1-tau/norms)[:,None]*w\n",
    "    w = w.reshape(-1)\n",
    "    return w"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f0d7bf0e",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3949a99b",
   "metadata": {},
   "outputs": [],
   "source": [
    "snrs = []\n",
    "for run in range(30):\n",
    "#     sims = [GaussianSimulation(m, p, support = np.array([1,1,1,1,1,1,1,1,1,1,1,1])*10., std = std, seed=run)\n",
    "#            for std in range(100)]\n",
    "    snr = [GaussianSimulation(m, p, support = np.array([1,1,1,1,1,1,1,1,1,1,1,1])*10., std = std, seed=run*100).snr\n",
    "           for std in range()]\n",
    "    snrs.append(snr)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f1d0522b",
   "metadata": {},
   "outputs": [],
   "source": [
    "snrs = [[x[i] for x in snrs] for i in range(len(snrs[0]))]\n",
    "snrs = [np.mean(x) for x in snrs]\n",
    "snrs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ec545531",
   "metadata": {},
   "outputs": [],
   "source": [
    "list(zip(snrs, range(100)))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b0183010",
   "metadata": {},
   "outputs": [],
   "source": [
    "snrs = []\n",
    "for run in range(30):\n",
    "#     sims = [GaussianSimulation(m, p, support = np.array([1,1,1,1,1,1,1,1,1,1,1,1])*10., std = std, seed=run)\n",
    "#            for std in range(100)]\n",
    "    snr = [GaussianSimulation(m, p, support = np.array([1,1,1,1,1,1,1,1,1,1,1,1])*10., std = std, seed=run*100).snr\n",
    "           for std in [1,2,3,4,5,6,8,11,17,34,68]]\n",
    "    snrs.append(snr)\n",
    "snrs = [[x[i] for x in snrs] for i in range(len(snrs[0]))]\n",
    "snrs = [np.mean(x) for x in snrs]\n",
    "snrs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a8267b90",
   "metadata": {},
   "outputs": [],
   "source": [
    "penalties = 0.1 * np.exp(np.linspace(0.01, 3, 10))\n",
    "m = 150\n",
    "p = 300\n",
    "\n",
    "all_sims = []\n",
    "outer_gres = []\n",
    "outer_ires = []\n",
    "outer_rres = []\n",
    "snrs = []\n",
    "for run in range(30):\n",
    "    print('#######################')\n",
    "    print('#######################')\n",
    "    print(f'########{run}#########')\n",
    "    print('#######################')\n",
    "    print('#######################')\n",
    "    seed = run * 100\n",
    "    sims = [GaussianSimulation(m, p, support = np.array([1,1,1,1,1,1,1,1,1,1,1,1])*10., std = std, seed=seed)\n",
    "           for std in np.exp(np.linspace(0,3,10))]\n",
    "    snr = [sim.snr for sim in sims]\n",
    "    snrs.append(snr)\n",
    "    \n",
    "    gres = []\n",
    "    gres_all = []\n",
    "    for sim in sims:\n",
    "        est_err, val_error, res, val_errors = val_group_lasso(sim, penalties)\n",
    "        gres.append(est_err)\n",
    "        gres_all.append([est_err, val_error, res, val_errors])\n",
    "    outer_gres.append(gres)\n",
    "    \n",
    "    iterate_err = []\n",
    "    iterate_all = []\n",
    "    for sim in sims:\n",
    "        ratio = 10 * datadriven_ratio(sim.X, sim.y)\n",
    "        _, _, _, all_w = dual_primal(\n",
    "            sim.X.numpy(), sim.y.numpy(), step_ratio=ratio, rho=0.99, ret_all=True,\n",
    "            prox=proximal_op,\n",
    "            max_iter=100,\n",
    "            f_store=1)\n",
    "        val_err = (((sim.X_val@all_w.transpose())-sim.y_val[:,None])**2).sum(axis=0).numpy()\n",
    "        est_err = ((all_w - sim.w_star.numpy()[None,:])**2).sum(axis=1)\n",
    "        iterate_err.append(min(zip(val_err, est_err)))\n",
    "        iterate_all.append(zip(val_err, est_err))\n",
    "    final_iterate_err = [x[1] for x in iterate_err]\n",
    "    outer_ires.append(final_iterate_err)\n",
    "    \n",
    "    reparam_res = []\n",
    "    reparam_res_all = []\n",
    "    cnt = 1\n",
    "    for sim in sims:\n",
    "        model = groupModel(p=sim.p, group_size=4, depth=2)\n",
    "        init = 1e-6\n",
    "        for param in model.parameters():\n",
    "            torch.nn.init.ones_(param)\n",
    "        model.u.weight.data *= init\n",
    "        for i in range(model.num_groups):\n",
    "            model.vs[i].weight.data *= 1/np.sqrt(model.group_size)\n",
    "        trainer = groupTrainer(model, sim, lr=0.001, is_two_lr=True, is_small_train=False)\n",
    "        trainer.train(600)\n",
    "    #     est_err = ((sim.w_star - trainer.model.get_params().detach())**2).sum().item()\n",
    "        reparam_res.append(min(zip(trainer.val_err, trainer.params_est_err)))\n",
    "        reparam_res_all.append(zip(trainer.val_err, trainer.params_est_err))\n",
    "        print(reparam_res)\n",
    "\n",
    "        plt.plot(trainer.val_err)\n",
    "        plt.plot(trainer.params_est_err)\n",
    "        plt.show()\n",
    "        \n",
    "    final_reparam_res = [x[1] for x in reparam_res]\n",
    "    outer_rres.append(final_reparam_res)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7d62f146",
   "metadata": {},
   "outputs": [],
   "source": [
    "ssnrs = [[x[i] for x in snrs] for i in range(len(snrs[0]))]\n",
    "ssnrs = [np.mean(x) for x in ssnrs]\n",
    "ssnrs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f99d6093",
   "metadata": {},
   "outputs": [],
   "source": [
    "def cal(l):\n",
    "    errs = np.flip(np.log2(np.array(l)), axis=1)\n",
    "    means = np.mean(errs, axis = 0)\n",
    "    stds = np.std(errs, axis=0)#/np.sqrt(29)\n",
    "    return means, stds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bc9c98f0",
   "metadata": {},
   "outputs": [],
   "source": [
    "gmeans, gstds = cal(outer_gres)\n",
    "imeans, istds = cal(outer_ires)\n",
    "rmeans, rstds = cal(outer_rres)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2be6526f",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.rcParams.update({'text.usetex': True})\n",
    "plt.rcParams.update({'text.latex.preamble': r'\\usepackage{amsmath}'})\n",
    "plt.rcParams.update({'lines.linewidth': 3})\n",
    "plt.rcParams.update({'font.size': 15})\n",
    "plt.rcParams.update({'legend.frameon': False})"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "284b7c84",
   "metadata": {},
   "outputs": [],
   "source": [
    "from matplotlib.transforms import ScaledTranslation\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7c498840",
   "metadata": {},
   "outputs": [],
   "source": [
    "fig, ax = plt.subplots()\n",
    "\n",
    "trans1 = ax.transData + ScaledTranslation(-5/72, 0, fig.dpi_scale_trans)\n",
    "trans2 = ax.transData + ScaledTranslation(+5/72, 0, fig.dpi_scale_trans)\n",
    "\n",
    "# ax.errorbar([round(x,2) for x in ssnrs[::-1]], gmeans, yerr=gstds, \n",
    "#             capthick=1, capsize=2, fmt=':', alpha=0.75, label='group lasso')\n",
    "ax.plot([round(x,2) for x in ssnrs[::-1]], gmeans, '--o', label='PGD')\n",
    "data = {\n",
    "    'x': [round(x,2) for x in ssnrs[::-1]],\n",
    "    'y1': gmeans-gstds,\n",
    "    'y2': gmeans+gstds}\n",
    "ax.fill_between(**data, alpha=0.25)\n",
    "\n",
    "# ax.errorbar([round(x,2)+1.0 for x in ssnrs[::-1]], rmeans, yerr=rstds,\n",
    "#             capthick=1, capsize=2, fmt=':', alpha=0.75, label='reparams')\n",
    "ax.plot([round(x,2) for x in ssnrs[::-1]], rmeans, '--s', label='DGLNN')\n",
    "data = {\n",
    "    'x': [round(x,2) for x in ssnrs[::-1]],\n",
    "    'y1': rmeans-rstds,\n",
    "    'y2': rmeans+rstds}\n",
    "ax.fill_between(**data, alpha=0.25)\n",
    "\n",
    "# ax.errorbar([round(x,2)+0.5 for x in ssnrs[::-1]], imeans, yerr=istds,\n",
    "#             capthick=0.5, capsize=2, fmt=':', alpha=0.75, label='primal dual')\n",
    "ax.plot([round(x,2) for x in ssnrs[::-1]], imeans, '--*', label='Primal-Dual')\n",
    "data = {\n",
    "    'x': [round(x,2) for x in ssnrs[::-1]],\n",
    "    'y1': imeans-istds,\n",
    "    'y2': imeans+istds}\n",
    "ax.fill_between(**data, alpha=0.25)\n",
    "\n",
    "# plt.fill_between(range(6), imeans - istds, means + stds, alpha=.15)\n",
    "ax.legend()\n",
    "# plt.xticks(range(len(ssnrs)), [round(x,2) for x in ssnrs[::-1]])\n",
    "ax.set_xlabel('SNR')\n",
    "ax.set_ylabel(r'$\\log_{2} ||\\mathbf{w}_{t} - \\mathbf{w}^{\\star}||_{2}^{2}$')\n",
    "ax.set_xlim(0,35)\n",
    "fig.tight_layout()\n",
    "fig.savefig('outputs/comparisons.pdf')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cc184eec",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "326f5676",
   "metadata": {},
   "outputs": [],
   "source": [
    "penalties = 0.1 * np.exp(np.linspace(0.01, 3, 10))\n",
    "m = 150\n",
    "p = 300\n",
    "\n",
    "all_sims = []\n",
    "\n",
    "outer_gres2 = []\n",
    "outer_ires2 = []\n",
    "outer_rres2 = []\n",
    "snrs2 = []\n",
    "for run in range(30):\n",
    "    print('#######################')\n",
    "    print('#######################')\n",
    "    print(f'########{run}#########')\n",
    "    print('#######################')\n",
    "    print('#######################')\n",
    "    seed = run * 100\n",
    "    sims = [GaussianSimulation(m, p, support = np.array([1,1,1,1,1,1,1,1,1,1,1,1])*10., std = std, seed=seed)\n",
    "           for std in [1,2,3,4,5,6,8,11,17,34,68]]\n",
    "    snr = [sim.snr for sim in sims]\n",
    "    snrs.append(snr)\n",
    "    \n",
    "    gres = []\n",
    "    gres_all = []\n",
    "    for sim in sims:\n",
    "        est_err, val_error, res, val_errors = val_group_lasso(sim, penalties)\n",
    "        gres.append(est_err)\n",
    "        gres_all.append([est_err, val_error, res, val_errors])\n",
    "    outer_gres.append(gres)\n",
    "    \n",
    "    iterate_err = []\n",
    "    iterate_all = []\n",
    "    for sim in sims:\n",
    "        ratio = 10 * datadriven_ratio(sim.X, sim.y)\n",
    "        _, _, _, all_w = dual_primal(\n",
    "            sim.X.numpy(), sim.y.numpy(), step_ratio=ratio, rho=0.99, ret_all=True,\n",
    "            prox=proximal_op,\n",
    "            max_iter=100,\n",
    "            f_store=1)\n",
    "        val_err = (((sim.X_val@all_w.transpose())-sim.y_val[:,None])**2).sum(axis=0).numpy()\n",
    "        est_err = ((all_w - sim.w_star.numpy()[None,:])**2).sum(axis=1)\n",
    "        iterate_err.append(min(zip(val_err, est_err)))\n",
    "        iterate_all.append(zip(val_err, est_err))\n",
    "    final_iterate_err = [x[1] for x in iterate_err]\n",
    "    outer_ires.append(final_iterate_err)\n",
    "    \n",
    "    reparam_res = []\n",
    "    reparam_res_all = []\n",
    "    cnt = 1\n",
    "    for sim in sims:\n",
    "        model = groupModel(p=sim.p, group_size=4, depth=2)\n",
    "        init = 1e-6\n",
    "        for param in model.parameters():\n",
    "            torch.nn.init.ones_(param)\n",
    "        model.u.weight.data *= init\n",
    "        for i in range(model.num_groups):\n",
    "            model.vs[i].weight.data *= 1/np.sqrt(model.group_size)\n",
    "        trainer = groupTrainer(model, sim, lr=0.001, is_two_lr=True, is_small_train=False)\n",
    "        trainer.train(600)\n",
    "    #     est_err = ((sim.w_star - trainer.model.get_params().detach())**2).sum().item()\n",
    "        reparam_res.append(min(zip(trainer.val_err, trainer.params_est_err)))\n",
    "        reparam_res_all.append(zip(trainer.val_err, trainer.params_est_err))\n",
    "        print(reparam_res)\n",
    "\n",
    "        plt.plot(trainer.val_err)\n",
    "        plt.plot(trainer.params_est_err)\n",
    "        plt.show()\n",
    "        \n",
    "    final_reparam_res = [x[1] for x in reparam_res]\n",
    "    outer_rres.append(final_reparam_res)\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7ae2dcf5",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "im",
   "language": "python",
   "name": "im"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}