{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "a33a92ba",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import sys\n",
    "os.environ[\"R_HOME\"] = r\"C:\\Program Files\\R\\R-4.4.2\"  # Use raw string for Windows paths\n",
    "sys.path.append(r\"c:\\Document\\Serieux\\Travail\\python_work\\cEBNM_torch\\py\")\n",
    "\n",
    "\n",
    "from utils import * \n",
    "# Core libraries\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    " \n",
    "\n",
    "# scikit-learn\n",
    "from sklearn.linear_model import LassoCV, RidgeCV\n",
    "from sklearn.metrics import mean_squared_error\n",
    "\n",
    "# rpy2 for R integration\n",
    "import rpy2.robjects as ro\n",
    "from rpy2.robjects import numpy2ri, r\n",
    "from rpy2.robjects.packages import importr\n",
    "\n",
    "# Activate NumPy-to-R conversion\n",
    "numpy2ri.activate()\n",
    "ashr = importr(\"ashr\")\n",
    "\n",
    "# torch (if used later in the project)\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.optim as optim\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "\n",
    " \n",
    "# Local utilities (if you have a utils.py file)\n",
    "from utils_mix import *\n",
    "from empirical_mdn import emdn_posterior_means\n",
    "from utils_mix import *\n",
    "from numerical_routine import *\n",
    "from distribution_operation import *\n",
    "from posterior_computation import posterior_mean_norm\n",
    "from cash_solver import Cash_posterior_means\n",
    "from ash import call_r_ash_fit_all_with_postmean\n",
    "import pandas as pd\n",
    "from pathlib import Path\n",
    "from sklearn.metrics import mean_squared_error, mean_absolute_error\n",
    "from sklearn.linear_model import LassoCV\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "0dda98f6",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "def col_scale(X, with_mean=True, with_std=True):\n",
    "    scaler = StandardScaler(with_mean=with_mean, with_std=with_std)\n",
    "    X_scaled = scaler.fit_transform(X)\n",
    "    return X_scaled, scaler.scale_\n",
    "\n",
    "\n",
    "def call_r_ash_fit_all_with_postmean(beta, sigma):\n",
    "    \"\"\"\n",
    "    Calls R's ash function and returns:\n",
    "    - log mixture weights\n",
    "    - mixture standard deviations\n",
    "    - posterior mean of beta\n",
    "    \"\"\"\n",
    "    from rpy2.rinterface_lib.sexp import NULLType\n",
    "\n",
    "    sebetahat = np.full_like(beta, sigma)\n",
    "    ash_obj = ashr.ash(betahat=beta, sebetahat=sebetahat, mixcompdist=\"normal\")\n",
    "\n",
    "    fitted_g = ash_obj.rx2(\"fitted_g\")\n",
    "    pi_r = np.array(fitted_g.rx2(\"pi\"), dtype=np.float32)\n",
    "    scale_r = np.array(fitted_g.rx2(\"sd\"), dtype=np.float32)\n",
    "\n",
    "    posterior_mean_r = ash_obj.rx2(\"result\").rx2(\"PosteriorMean\")\n",
    "    if isinstance(posterior_mean_r, NULLType):\n",
    "        raise RuntimeError(\"R ash() returned NULL for result$PosteriorMean\")\n",
    "    posterior_mean = np.array(posterior_mean_r, dtype=np.float32)\n",
    "\n",
    "    log_pi = np.log(np.clip(pi_r, 1e-12, 1.0))\n",
    "    return log_pi, scale_r, posterior_mean\n",
    "\n",
    "\n",
    "def get_data_loglik_normal(betahat, sebetahat, location, scale):\n",
    "    var = sebetahat[:, None] ** 2 + scale[None, :] ** 2\n",
    "    return -0.5 * (np.log(2 * np.pi * var) + (betahat[:, None] - location) ** 2 / var)\n",
    "\n",
    "\n",
    "def apply_log_sum_exp(data_loglik, log_pi):\n",
    "    combined_loglik = data_loglik + log_pi\n",
    "    return combined_loglik - logsumexp(combined_loglik, axis=1)[:, None]\n",
    "\n",
    "\n",
    "class PosteriorMeanNorm:\n",
    "    def __init__(self, post_mean, post_mean2, post_sd):\n",
    "        self.post_mean = post_mean\n",
    "        self.post_mean2 = post_mean2\n",
    "        self.post_sd = post_sd\n",
    "\n",
    "\n",
    "def posterior_mean_norm(betahat, sebetahat, log_pi, scale, location=None):\n",
    "    if location is None:\n",
    "        location = np.zeros_like(scale)\n",
    "\n",
    "    data_loglik = get_data_loglik_normal(betahat, sebetahat, location, scale)\n",
    "    log_post_assignment = apply_log_sum_exp(data_loglik, log_pi)\n",
    "\n",
    "    n, K = betahat.shape[0], scale.shape[0]\n",
    "    var = np.zeros((n, K))\n",
    "\n",
    "    if scale[0] == 0:\n",
    "        for i in range(n):\n",
    "            var[i, :] = np.concatenate(([0], 1 / ((1 / sebetahat[i]**2) + (1 / scale[1:]**2))))\n",
    "    else:\n",
    "        for i in range(n):\n",
    "            var[i, :] = 1 / ((1 / sebetahat[i]**2) + (1 / scale**2))\n",
    "\n",
    "    temp = np.zeros((n, K))\n",
    "    for i in range(n):\n",
    "        temp[i, :] = (var[i, :] / sebetahat[i]**2) * betahat[i] + location * (1 - var[i, :] / sebetahat[i]**2)\n",
    "\n",
    "    post_mean = np.sum(np.exp(log_post_assignment) * temp, axis=1)\n",
    "    post_mean2 = np.sum(np.exp(log_post_assignment) * (var + temp**2), axis=1)\n",
    "    post_sd = np.sqrt(post_mean2 - post_mean**2)\n",
    "    return PosteriorMeanNorm(post_mean, post_mean2, post_sd)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "53aaaa7e",
   "metadata": {},
   "outputs": [],
   "source": [
    "def nash_mdn(X, y, sideinfo=None, maxit=4, damping=0.99, eb_kwargs=None):\n",
    "    k=0 \n",
    "    X = X.astype(np.float32)\n",
    "    y = y.astype(np.float32).flatten()\n",
    "    X, csd = col_scale(X)\n",
    "    y, ysd = col_scale(y.reshape(-1, 1))\n",
    "    y = y.flatten()\n",
    "\n",
    "    n, p = X.shape\n",
    "    sqn = np.sqrt(n).astype(np.float32)\n",
    "    beta = LassoCV(cv=5, max_iter=5000).fit(X, y).coef_\n",
    "\n",
    "    log_pi, scale, location = None, None, None\n",
    "    td_beta = []\n",
    "    model_param = None\n",
    "    pi0=0\n",
    "    \n",
    "    for o in range(maxit):\n",
    "        betahat_list, sebetahat_list = [], []\n",
    "        for k in range(p):\n",
    "            r = y - X @ beta\n",
    "            beta_new = beta.copy()\n",
    "            xk = X[:, k]\n",
    "            r_k = r + xk * beta_new[k]\n",
    "            betahat_k = np.dot(xk, r_k) / n\n",
    "            sebetahat_k = max(np.std(r_k) / sqn, 1e-4)\n",
    "            if 0<1: \n",
    "\n",
    "                beta_k_new = betahat_k\n",
    "                beta_k_new = damping * beta_k_new + (1 - damping) * beta_new[k]\n",
    "             \n",
    "            else :\n",
    "                print(o)\n",
    "                pm = posterior_mean_norm(\n",
    "                    betahat=np.array([betahat_k]),\n",
    "                    sebetahat=np.array([sebetahat_k]),\n",
    "                    log_pi=np.log( np.clip( result.pi_np[k, :], 1e-12, 1.0)),\n",
    "                    location=result.location[k, :],\n",
    "                    scale=result.scale[k, :]\n",
    "                    )\n",
    "             \n",
    "                beta_k_new = pm.post_mean[0]\n",
    "\n",
    "            r += xk * (beta_new[k]  - beta_k_new)\n",
    "            beta_k_new = damping * beta_k_new + (1 - damping) * beta_new[k]\n",
    "            beta_new[k]=beta_k_new\n",
    "            betahat_list.append(betahat_k)\n",
    "            sebetahat_list.append(sebetahat_k) \n",
    "\n",
    "\n",
    "        betahat_arr = np.array(betahat_list)\n",
    "        avg_se = np.mean(sebetahat_list)\n",
    "        beta = beta_new.copy()\n",
    "        res_sq_final = np.sum((y - X @ beta ) ** 2)\n",
    "        sigma_0_term = np.dot(beta, beta - betahat_arr )   \n",
    " \n",
    "        beta = beta_new.copy()\n",
    "        res_sq_final = np.sum((y - X @ beta ) ** 2)\n",
    "        sigma_0_term = np.dot(beta, beta - betahat_arr )\n",
    "        denom = n + p * (1 - pi0)\n",
    "        sigma_0 = np.sqrt(max((res_sq_final + sigma_0_term) / denom, 1e-8)) / sqn\n",
    "\n",
    "        res_sq = np.sum((y - X @ betahat_arr) ** 2)\n",
    "        drift_term = np.dot(betahat_arr,  betahat_arr) / np.sqrt(n + p)\n",
    "        s = np.sqrt((res_sq + drift_term)) / sqn\n",
    "\n",
    "        # Update drift_comp\n",
    "        drift_comp = (1 / (n / s**2 + 1 / sigma_0**2)) * (n / s**2)\n",
    "        if o<1: \n",
    "            # Update beta with ash\n",
    "            ash_input = betahat_arr   \n",
    "            result =emdn_posterior_means(X= sideinfo,\n",
    "                                            betahat= ash_input,\n",
    "                                            sebetahat=np.array(sebetahat_list)   )\n",
    "            emdn_posterior_means(X= sideinfo,\n",
    "                                          betahat= ash_input,\n",
    "                                          sebetahat= np.full_like(ash_input,sigma_0), **(eb_kwargs or {}))\n",
    "        else :\n",
    "            ash_input = drift_comp * betahat_arr + (1 - drift_comp) * result.post_mean\n",
    "            result =emdn_posterior_means(X= sideinfo,\n",
    "                                          betahat= ash_input,\n",
    "                                          sebetahat=np.full_like(ash_input,sigma_0)   )\n",
    "            log_pi= result.pi_np\n",
    "            beta =   result.post_mean\n",
    "\n",
    "\n",
    "  \n",
    "    out=   ysd * result.post_mean/ csd\n",
    "    return out\n",
    " \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "fdf2943f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 10/50, Loss: -0.6314\n",
      "Epoch 20/50, Loss: -1.7690\n",
      "Epoch 30/50, Loss: -3.1476\n",
      "Epoch 40/50, Loss: -4.0074\n",
      "Epoch 50/50, Loss: -0.6389\n",
      "Epoch 10/50, Loss: 1.2506\n",
      "Epoch 20/50, Loss: -3.6690\n",
      "Epoch 30/50, Loss: 1.0733\n",
      "Epoch 40/50, Loss: -0.4153\n",
      "Epoch 50/50, Loss: -1.5665\n",
      "Epoch 10/50, Loss: -0.4815\n",
      "Epoch 20/50, Loss: 2.0001\n",
      "Epoch 30/50, Loss: 2.3950\n",
      "Epoch 40/50, Loss: -0.1239\n",
      "Epoch 50/50, Loss: -3.0335\n",
      "  Fold 1: RMSE = 6.2169, MAD = 5.4324\n",
      "['TCGA']: mean RMSE = 6.2169, mean MAD = 5.4324\n",
      "\n",
      "Epoch 10/50, Loss: -0.6447\n",
      "Epoch 20/50, Loss: -2.5323\n",
      "Epoch 30/50, Loss: -1.8570\n",
      "Epoch 40/50, Loss: -1.7571\n",
      "Epoch 50/50, Loss: -3.7888\n",
      "Epoch 10/50, Loss: -1.1766\n",
      "Epoch 20/50, Loss: -2.5223\n",
      "Epoch 30/50, Loss: -0.9274\n",
      "Epoch 40/50, Loss: -0.8696\n",
      "Epoch 50/50, Loss: -2.9100\n",
      "Epoch 10/50, Loss: 0.9553\n",
      "Epoch 20/50, Loss: 1.0464\n",
      "Epoch 30/50, Loss: -1.1275\n",
      "Epoch 40/50, Loss: -3.0328\n",
      "Epoch 50/50, Loss: -0.0421\n",
      "  Fold 2: RMSE = 835.9014, MAD = 835.8764\n",
      "['TCGA']: mean RMSE = 421.0592, mean MAD = 420.6544\n",
      "\n",
      "Epoch 10/50, Loss: -2.2140\n",
      "Epoch 20/50, Loss: -3.0080\n",
      "Epoch 30/50, Loss: -3.4499\n",
      "Epoch 40/50, Loss: -3.9511\n",
      "Epoch 50/50, Loss: -3.8555\n",
      "Epoch 10/50, Loss: 3.2423\n",
      "Epoch 20/50, Loss: -0.4772\n",
      "Epoch 30/50, Loss: -1.6535\n",
      "Epoch 40/50, Loss: -3.6148\n",
      "Epoch 50/50, Loss: 1.7079\n",
      "Epoch 10/50, Loss: -2.3708\n",
      "Epoch 20/50, Loss: -0.2761\n",
      "Epoch 30/50, Loss: -1.3918\n",
      "Epoch 40/50, Loss: -0.9312\n",
      "Epoch 50/50, Loss: -0.3683\n",
      "  Fold 3: RMSE = 771.0442, MAD = 771.0238\n",
      "['TCGA']: mean RMSE = 537.7208, mean MAD = 537.4442\n",
      "\n",
      "Epoch 10/50, Loss: -2.1685\n",
      "Epoch 20/50, Loss: -1.1218\n",
      "Epoch 30/50, Loss: -2.6704\n",
      "Epoch 40/50, Loss: -3.0851\n",
      "Epoch 50/50, Loss: -0.6703\n",
      "Epoch 10/50, Loss: -0.2297\n",
      "Epoch 20/50, Loss: -1.1813\n",
      "Epoch 30/50, Loss: -2.8212\n",
      "Epoch 40/50, Loss: -1.6575\n",
      "Epoch 50/50, Loss: -2.3410\n",
      "Epoch 10/50, Loss: -0.8046\n",
      "Epoch 20/50, Loss: -2.0277\n",
      "Epoch 30/50, Loss: -2.2414\n",
      "Epoch 40/50, Loss: -1.1667\n",
      "Epoch 50/50, Loss: -1.7928\n",
      "  Fold 4: RMSE = 7.9050, MAD = 7.5916\n",
      "['TCGA']: mean RMSE = 405.2669, mean MAD = 404.9810\n",
      "\n",
      "Epoch 10/50, Loss: 0.1078\n",
      "Epoch 20/50, Loss: -0.3374\n",
      "Epoch 30/50, Loss: -1.8506\n",
      "Epoch 40/50, Loss: -3.2432\n",
      "Epoch 50/50, Loss: -3.7864\n",
      "Epoch 10/50, Loss: -0.1034\n",
      "Epoch 20/50, Loss: -2.2688\n",
      "Epoch 30/50, Loss: -2.2383\n",
      "Epoch 40/50, Loss: -0.4979\n",
      "Epoch 50/50, Loss: -3.7657\n",
      "Epoch 10/50, Loss: -0.0816\n",
      "Epoch 20/50, Loss: -1.7053\n",
      "Epoch 30/50, Loss: -0.9755\n",
      "Epoch 40/50, Loss: -0.7624\n",
      "Epoch 50/50, Loss: -2.1981\n",
      "  Fold 5: RMSE = 361.3753, MAD = 361.3643\n",
      "['TCGA']: mean RMSE = 396.4886, mean MAD = 396.2577\n",
      "\n",
      "Epoch 10/50, Loss: -0.3458\n",
      "Epoch 20/50, Loss: -1.2697\n",
      "Epoch 30/50, Loss: -2.5027\n",
      "Epoch 40/50, Loss: 2.2453\n",
      "Epoch 50/50, Loss: -2.9852\n",
      "Epoch 10/50, Loss: -0.5555\n",
      "Epoch 20/50, Loss: -0.4976\n",
      "Epoch 30/50, Loss: -0.4540\n",
      "Epoch 40/50, Loss: -1.1971\n",
      "Epoch 50/50, Loss: -0.4387\n",
      "Epoch 10/50, Loss: -1.7552\n",
      "Epoch 20/50, Loss: -2.0704\n",
      "Epoch 30/50, Loss: -1.3084\n",
      "Epoch 40/50, Loss: -0.8122\n",
      "Epoch 50/50, Loss: -3.3772\n",
      "  Fold 6: RMSE = 183.3416, MAD = 183.3150\n",
      "['TCGA']: mean RMSE = 360.9641, mean MAD = 360.7672\n",
      "\n",
      "Epoch 10/50, Loss: -2.1702\n",
      "Epoch 20/50, Loss: -2.0631\n",
      "Epoch 30/50, Loss: 0.7039\n",
      "Epoch 40/50, Loss: -1.3071\n",
      "Epoch 50/50, Loss: -3.7788\n",
      "Epoch 10/50, Loss: -1.2956\n",
      "Epoch 20/50, Loss: -2.0150\n",
      "Epoch 30/50, Loss: 0.7106\n",
      "Epoch 40/50, Loss: -1.2822\n",
      "Epoch 50/50, Loss: 0.3958\n",
      "Epoch 10/50, Loss: -1.3323\n",
      "Epoch 20/50, Loss: 0.1502\n",
      "Epoch 30/50, Loss: -0.5149\n",
      "Epoch 40/50, Loss: -3.6626\n",
      "Epoch 50/50, Loss: -1.4593\n",
      "  Fold 7: RMSE = 395.0608, MAD = 395.0515\n",
      "['TCGA']: mean RMSE = 365.8350, mean MAD = 365.6650\n",
      "\n",
      "  Fold 8: missing file — skipping.\n",
      "  Fold 9: missing file — skipping.\n",
      "  Fold 10: missing file — skipping.\n"
     ]
    }
   ],
   "source": [
    "dataset_name=[\"TCGA\"]\n",
    "rmses, mads = [], []\n",
    "for k in range(1, 11):\n",
    "        try:\n",
    "            base_path = Path(\"C:/Document/Serieux/Travail/Data_analysis_and_papers/nash_experiement/data_split\")\n",
    "            result_path = Path(\"C:/Document/Serieux/Travail/Data_analysis_and_papers/nash_experiement/results_realdata\")\n",
    "            datasets = [\"TCGA\"] \n",
    "            dataset_path = Path(\"C:/Document/Serieux/Travail/Data_analysis_and_papers/nash_experiement/data_split/TCGA\")\n",
    "            X_train = pd.read_csv(dataset_path / f\"X_train{k}.csv\").values.astype(np.float32)\n",
    "            y_train = pd.read_csv(dataset_path / f\"y_train{k}.csv\").values.astype(np.float32).flatten()\n",
    "            X_test = pd.read_csv(dataset_path / f\"X_test{k}.csv\").values.astype(np.float32)\n",
    "            y_test = pd.read_csv(dataset_path / f\"y_test{k}.csv\").values.astype(np.float32).flatten()\n",
    " \n",
    "            df = pd.read_csv(dataset_path / \"infocov.csv\")\n",
    "         \n",
    "\n",
    "# One-hot encode the only column\n",
    "            one_hot = pd.get_dummies(df[df.columns[0]])\n",
    "\n",
    "# Convert to NumPy\n",
    "            onehot_array = one_hot.to_numpy(dtype=np.float32)\n",
    "            sideinfo_array = onehot_array\n",
    "            mu= combined = np.concatenate([ y_test,y_train]).mean()\n",
    "            y_test=  y_test-mu\n",
    "            y_train=y_train-mu\n",
    "# Now pass sideinfo to your model:\n",
    "            beta_est = nash_mdn(X=X_train, y=y_train, sideinfo=sideinfo_array, maxit=2)\n",
    "\n",
    "            y_pred_test = X_test @ beta_est\n",
    "\n",
    "        # Metrics\n",
    "            rmse = np.sqrt(mean_squared_error(y_test, y_pred_test))\n",
    "            mad = mean_absolute_error(y_test, y_pred_test)\n",
    "            rmses.append(rmse)\n",
    "            mads.append(mad)\n",
    "            print(f\"  Fold {k}: RMSE = {rmse:.4f}, MAD = {mad:.4f}\")\n",
    "\n",
    "    # Save results\n",
    "            if rmses:\n",
    "                df = pd.DataFrame({'RMSE': rmses, 'MAD': mads})\n",
    "                df.to_csv(result_path / f\"{dataset_name}_nash_mdn.csv\", index=False)\n",
    "                print(f\"{dataset_name}: mean RMSE = {np.mean(rmses):.4f}, mean MAD = {np.mean(mads):.4f}\\n\")\n",
    "            else:\n",
    "                print(f\"{dataset_name}: no valid folds to evaluate.\\n\")\n",
    "\n",
    "# View result \n",
    "        except FileNotFoundError:\n",
    "            print(f\"  Fold {k}: missing file — skipping.\")\n",
    "            continue"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "17d5b60e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "     x\n",
      "0  309\n",
      "1  309\n",
      "2  266\n",
      "3  220\n",
      "4   32\n",
      "Columns: Index(['x'], dtype='object')\n"
     ]
    }
   ],
   "source": [
    "dataset_name=[\"TCGA\"]\n",
    "k=1\n",
    "\n",
    "\n",
    "base_path = Path(\"C:/Document/Serieux/Travail/Data_analysis_and_papers/nash_experiement/data_split\")\n",
    "result_path = Path(\"C:/Document/Serieux/Travail/Data_analysis_and_papers/nash_experiement/results_realdata\")\n",
    "datasets = [\"TCGA\"]\n",
    "dataset_path = Path(\"C:/Document/Serieux/Travail/Data_analysis_and_papers/nash_experiement/data_split/TCGA\")\n",
    "X_train = pd.read_csv(dataset_path / f\"X_train{k}.csv\").values.astype(np.float32)\n",
    "y_train = pd.read_csv(dataset_path / f\"y_train{k}.csv\").values.astype(np.float32).flatten()\n",
    "X_test = pd.read_csv(dataset_path / f\"X_test{k}.csv\").values.astype(np.float32)\n",
    "y_test = pd.read_csv(dataset_path / f\"y_test{k}.csv\").values.astype(np.float32).flatten()\n",
    "\n",
    "# ✅ FIXED: keep as DataFrame\n",
    "df = pd.read_csv(dataset_path / \"infocov.csv\")\n",
    "print(df.head())\n",
    "print(\"Columns:\", df.columns)\n",
    "\n",
    "# One-hot encode the only column\n",
    "one_hot = pd.get_dummies(df[df.columns[0]])\n",
    "\n",
    "# Convert to NumPy\n",
    "onehot_array = one_hot.to_numpy(dtype=np.float32)\n",
    "group_labels = one_hot.columns.to_list()\n",
    "\n",
    "onehot_array = one_hot.to_numpy(dtype=np.float32)\n",
    "sideinfo_array = onehot_array\n",
    " \n",
    "# Now pass sideinfo to your model:\n",
    "from sklearn.preprocessing import StandardScaler\n",
    "\n",
    "# Initialize the scaler\n",
    "scaler = StandardScaler()\n",
    "X_train_scaled = scaler.fit_transform(X_train)\n",
    "\n",
    "# Use the same scaler to transform X_test\n",
    "X_test_scaled = scaler.transform(X_test)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "4a12eecc",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[0., 0., 0., ..., 0., 0., 1.],\n",
       "       [0., 0., 0., ..., 0., 0., 1.],\n",
       "       [0., 0., 0., ..., 0., 0., 0.],\n",
       "       ...,\n",
       "       [0., 0., 0., ..., 0., 0., 1.],\n",
       "       [0., 0., 0., ..., 0., 0., 1.],\n",
       "       [0., 0., 0., ..., 0., 0., 1.]], dtype=float32)"
      ]
     },
     "execution_count": 8,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "onehot_array"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "id": "d30842ac",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 10/20,   Variance Loss: -524.0538\n",
      "Epoch 20/20,   Variance Loss: -524.0687\n",
      "Epoch 10/20,   Variance Loss: -1124.8631\n",
      "Epoch 20/20,   Variance Loss: -1125.0511\n"
     ]
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "Cell \u001b[1;32mIn[28], line 81\u001b[0m\n\u001b[0;32m     79\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m :\n\u001b[0;32m     80\u001b[0m     ash_input \u001b[38;5;241m=\u001b[39m drift_comp \u001b[38;5;241m*\u001b[39m betahat_arr \u001b[38;5;241m+\u001b[39m (\u001b[38;5;241m1\u001b[39m \u001b[38;5;241m-\u001b[39m drift_comp) \u001b[38;5;241m*\u001b[39m result\u001b[38;5;241m.\u001b[39mpost_mean\n\u001b[1;32m---> 81\u001b[0m     result \u001b[38;5;241m=\u001b[39m\u001b[43mCash_posterior_means\u001b[49m\u001b[43m(\u001b[49m\u001b[43mX\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43m \u001b[49m\u001b[43msideinfo\u001b[49m\u001b[43m \u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m     82\u001b[0m \u001b[43m                                  \u001b[49m\u001b[43mbetahat\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43m  \u001b[49m\u001b[43mash_input\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;66;43;03m#ash_input ,\u001b[39;49;00m\n\u001b[0;32m     83\u001b[0m \u001b[43m                                  \u001b[49m\u001b[43msebetahat\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43m \u001b[49m\u001b[43mnp\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mfull_like\u001b[49m\u001b[43m(\u001b[49m\u001b[43mash_input\u001b[49m\u001b[43m,\u001b[49m\u001b[43msigma_0\u001b[49m\u001b[43m)\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m/\u001b[39;49m\u001b[43msqn\u001b[49m\u001b[43m   \u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m     84\u001b[0m \u001b[43m                                  \u001b[49m\u001b[43mn_epochs\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;241;43m20\u001b[39;49m\u001b[43m   \u001b[49m\u001b[43m)\u001b[49m\n\u001b[0;32m     85\u001b[0m     log_pi\u001b[38;5;241m=\u001b[39m result\u001b[38;5;241m.\u001b[39mpi_np\n",
      "File \u001b[1;32mc:\\Document\\Serieux\\Travail\\python_work\\cEBNM_torch\\py\\cash_solver.py:124\u001b[0m, in \u001b[0;36mCash_posterior_means\u001b[1;34m(X, betahat, sebetahat, n_epochs, n_layers, num_classes, hidden_dim, batch_size, lr, model_param, penalty)\u001b[0m\n\u001b[0;32m    121\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m inputs, targets, noise_std \u001b[38;5;129;01min\u001b[39;00m dataloader:\n\u001b[0;32m    123\u001b[0m     residuals \u001b[38;5;241m=\u001b[39mtargets\u001b[38;5;241m.\u001b[39mdetach()\u001b[38;5;241m.\u001b[39mnumpy() \n\u001b[1;32m--> 124\u001b[0m     batch_loglik \u001b[38;5;241m=\u001b[39m \u001b[43mget_data_loglik_normal\u001b[49m\u001b[43m(\u001b[49m\u001b[43mbetahat\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mresiduals\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m    125\u001b[0m \u001b[43m                                      \u001b[49m\u001b[43msebetahat\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mnoise_std\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m    126\u001b[0m \u001b[43m                                      \u001b[49m\u001b[43mlocation\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[38;5;241;43m0\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mscale\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m    127\u001b[0m \u001b[43m                                      \u001b[49m\u001b[43mscale\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mscale\u001b[49m\u001b[43m)\u001b[49m\n\u001b[0;32m    128\u001b[0m     optimizer_cash\u001b[38;5;241m.\u001b[39mzero_grad()\n\u001b[0;32m    129\u001b[0m     outputs \u001b[38;5;241m=\u001b[39m model_cash(inputs)\n",
      "File \u001b[1;32mc:\\Document\\Serieux\\Travail\\python_work\\cEBNM_torch\\py\\distribution_operation.py:27\u001b[0m, in \u001b[0;36mget_data_loglik_normal\u001b[1;34m(betahat, sebetahat, location, scale)\u001b[0m\n\u001b[0;32m     25\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mlen\u001b[39m(location\u001b[38;5;241m.\u001b[39mshape)\u001b[38;5;241m==\u001b[39m\u001b[38;5;241m1\u001b[39m:\n\u001b[0;32m     26\u001b[0m   \u001b[38;5;28;01mfor\u001b[39;00m i \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mrange\u001b[39m(betahat\u001b[38;5;241m.\u001b[39mshape[\u001b[38;5;241m0\u001b[39m]):\n\u001b[1;32m---> 27\u001b[0m     out[i,] \u001b[38;5;241m=\u001b[39m\u001b[43mconvolved_logpdf_normal\u001b[49m\u001b[43m(\u001b[49m\u001b[43mbetahat\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mbetahat\u001b[49m\u001b[43m[\u001b[49m\u001b[43mi\u001b[49m\u001b[43m]\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m     28\u001b[0m \u001b[43m                                     \u001b[49m\u001b[43msebetahat\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43msebetahat\u001b[49m\u001b[43m[\u001b[49m\u001b[43mi\u001b[49m\u001b[43m]\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\n\u001b[0;32m     29\u001b[0m \u001b[43m                                     \u001b[49m\u001b[43mlocation\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mlocation\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m     30\u001b[0m \u001b[43m                                     \u001b[49m\u001b[43mscale\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mscale\u001b[49m\u001b[43m)\u001b[49m\n\u001b[0;32m     31\u001b[0m \u001b[38;5;28;01melif\u001b[39;00m \u001b[38;5;28mlen\u001b[39m(location\u001b[38;5;241m.\u001b[39mshape)\u001b[38;5;241m==\u001b[39m\u001b[38;5;241m2\u001b[39m:\n\u001b[0;32m     32\u001b[0m   \u001b[38;5;28;01mfor\u001b[39;00m i \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mrange\u001b[39m(betahat\u001b[38;5;241m.\u001b[39mshape[\u001b[38;5;241m0\u001b[39m]):\n",
      "File \u001b[1;32mc:\\Document\\Serieux\\Travail\\python_work\\cEBNM_torch\\py\\distribution_operation.py:11\u001b[0m, in \u001b[0;36mconvolved_logpdf_normal\u001b[1;34m(betahat, sebetahat, location, scale)\u001b[0m\n\u001b[0;32m      8\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21mconvolved_logpdf_normal\u001b[39m(  betahat,sebetahat, location, scale):\n\u001b[0;32m      9\u001b[0m   \u001b[38;5;66;03m# Calculate the standard deviation\u001b[39;00m\n\u001b[0;32m     10\u001b[0m   \u001b[38;5;66;03m#location ois wether a 1 d np with same length as scale \u001b[39;00m\n\u001b[1;32m---> 11\u001b[0m     sd \u001b[38;5;241m=\u001b[39m np\u001b[38;5;241m.\u001b[39msqrt(\u001b[43msebetahat\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m2\u001b[39;49m \u001b[38;5;241m+\u001b[39m  scale\u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39m\u001b[38;5;241m2\u001b[39m)\n\u001b[0;32m     13\u001b[0m     \u001b[38;5;66;03m# Calculate the log probability density\u001b[39;00m\n\u001b[0;32m     14\u001b[0m     logp \u001b[38;5;241m=\u001b[39m norm\u001b[38;5;241m.\u001b[39mlogpdf(betahat, loc\u001b[38;5;241m=\u001b[39m location, scale\u001b[38;5;241m=\u001b[39msd)\n",
      "File \u001b[1;32mc:\\Users\\willi\\anaconda3\\envs\\ml_env\\lib\\site-packages\\torch\\_tensor.py:39\u001b[0m, in \u001b[0;36m_handle_torch_function_and_wrap_type_error_to_not_implemented.<locals>.wrapped\u001b[1;34m(*args, **kwargs)\u001b[0m\n\u001b[0;32m     37\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m has_torch_function(args):\n\u001b[0;32m     38\u001b[0m         \u001b[38;5;28;01mreturn\u001b[39;00m handle_torch_function(wrapped, args, \u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)\n\u001b[1;32m---> 39\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m f(\u001b[38;5;241m*\u001b[39margs, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)\n\u001b[0;32m     40\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mTypeError\u001b[39;00m:\n\u001b[0;32m     41\u001b[0m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mNotImplemented\u001b[39m\n",
      "\u001b[1;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "source": [
    "y=y_train\n",
    "sideinfo= sideinfo_array\n",
    "maxit=10\n",
    "damping=0.99\n",
    "    # Standardize\n",
    "X=X_train_scaled\n",
    "y, ysd = col_scale(y.reshape(-1, 1))\n",
    "y = y.flatten()\n",
    "\n",
    "n, p = X.shape\n",
    "sqn = np.sqrt(n).astype(np.float32)\n",
    "\n",
    "    # Initialize via Lasso\n",
    "lasso = LassoCV(cv=5, max_iter=5000)\n",
    "lasso.fit(X, y)\n",
    "beta = lasso.coef_\n",
    "\n",
    "    # Initialize prior as None\n",
    "log_pi, scale, location = None, None, None\n",
    "td_beta = []\n",
    "pi0=0\n",
    "for o in range(maxit):\n",
    "        betahat_list, sebetahat_list = [], []\n",
    "        for k in range(p):\n",
    "            r = y - X @ beta\n",
    "            beta_new = beta.copy()\n",
    "            xk = X[:, k]\n",
    "            r_k = r + xk * beta_new[k]\n",
    "            betahat_k = np.dot(xk, r_k) / n\n",
    "            sebetahat_k = max(np.std(r_k) / sqn, 1e-4)\n",
    "            if 0<1: \n",
    "\n",
    "                beta_k_new = betahat_k\n",
    "                beta_k_new = damping * beta_k_new + (1 - damping) * beta_new[k]\n",
    "             \n",
    "            else :\n",
    "                print(o)\n",
    "                pm = posterior_mean_norm(\n",
    "                    betahat=np.array([betahat_k]),\n",
    "                    sebetahat=np.array([sebetahat_k]),\n",
    "                    log_pi=np.log( np.clip( result.pi_np[k, :], 1e-12, 1.0)),\n",
    "                    location=result.location[k, :],\n",
    "                    scale=result.scale[k, :]\n",
    "                    )\n",
    "             \n",
    "                beta_k_new = pm.post_mean[0]\n",
    "\n",
    "            r += xk * (beta_new[k]  - beta_k_new)\n",
    "            beta_k_new = damping * beta_k_new + (1 - damping) * beta_new[k]\n",
    "            beta_new[k]=beta_k_new\n",
    "            betahat_list.append(betahat_k)\n",
    "            sebetahat_list.append(sebetahat_k) \n",
    "\n",
    "\n",
    "        betahat_arr = np.array(betahat_list)\n",
    "        avg_se = np.mean(sebetahat_list)\n",
    "        beta = beta_new.copy()\n",
    "        res_sq_final = np.sum((y - X @ betahat_arr ) ** 2)\n",
    "        sigma_0_term = np.dot(beta, beta - betahat_arr )   \n",
    " \n",
    "        beta = beta_new.copy()\n",
    "        res_sq_final = np.sum((y - X @ beta ) ** 2)\n",
    "        sigma_0_term = np.dot(beta, beta - betahat_arr )\n",
    "        denom = n + p * (1 - pi0)\n",
    "        sigma_0 = np.sqrt(max((res_sq_final + sigma_0_term) / denom, 1e-8)) / sqn\n",
    "\n",
    "        res_sq = np.sum((y - X @ betahat_arr) ** 2)\n",
    "        drift_term = np.dot(betahat_arr,  betahat_arr) / np.sqrt(n + p)\n",
    "        s = np.sqrt((res_sq + drift_term)) / sqn\n",
    "\n",
    "        # Update drift_comp\n",
    "        drift_comp =    (1 / (n / s**2 + 1 / sigma_0**2)) * (n / s**2)\n",
    "        if o<1: \n",
    "            # Update beta with ash\n",
    "            ash_input = betahat_arr   \n",
    "            result =Cash_posterior_means(X= sideinfo ,\n",
    "                                            betahat= betahat_arr ,\n",
    "                                            sebetahat=np.array(sebetahat_list) ,n_epochs=20  ) \n",
    "        else :\n",
    "            ash_input = drift_comp * betahat_arr + (1 - drift_comp) * result.post_mean\n",
    "            result =Cash_posterior_means(X= sideinfo ,\n",
    "                                          betahat=  ash_input, #ash_input ,\n",
    "                                          sebetahat= np.full_like(ash_input,sigma_0) /sqn   ,\n",
    "                                          n_epochs=20   )\n",
    "            log_pi= result.pi_np\n",
    "            #beta =   result.post_mean"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "590aaae1",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 10/50,   Variance Loss: -524.0518\n",
      "Epoch 20/50,   Variance Loss: -524.0056\n",
      "Epoch 30/50,   Variance Loss: -524.1042\n",
      "Epoch 40/50,   Variance Loss: -524.1706\n",
      "Epoch 50/50,   Variance Loss: -524.1362\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "c:\\Document\\Serieux\\Travail\\python_work\\cEBNM_torch\\py\\cash_solver.py:159: RuntimeWarning: divide by zero encountered in log\n",
      "  log_pi=np.log(all_pi_values_np[i, :]),\n"
     ]
    }
   ],
   "source": [
    "result =Cash_posterior_means(X= sideinfo ,\n",
    "                                            betahat= betahat_arr ,\n",
    "                                            sebetahat=np.array(sebetahat_list) ,n_epochs=50  ) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "id": "e763f2a6",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[<matplotlib.lines.Line2D at 0x1c7e26cb700>]"
      ]
     },
     "execution_count": 29,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 48980 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.plot(betahat_arr )\n",
    "plt.plot(result.post_mean)\n",
    "plt.plot(lasso.coef_)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 47,
   "id": "cfcfd121",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0.2247384994149241\n",
      "0.16365185\n"
     ]
    },
    {
     "data": {
      "image/png": "[encoded data, 33060 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.scatter(X@result.post_mean, y)\n",
    "plt.scatter(X@lasso.coef_, y)\n",
    "print(np.sqrt(((X@result.post_mean - y )**2).mean()))\n",
    "print(np.sqrt(((X@lasso.coef_- y )**2).mean()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "id": "84d9a9e4",
   "metadata": {},
   "outputs": [],
   "source": [
    "tt ,oo= col_scale(y_test.reshape(-1, 1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 44,
   "id": "0868a941",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([1.28168141])"
      ]
     },
     "execution_count": 44,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "oo"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 46,
   "id": "034e5af9",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1.0409785181536713"
      ]
     },
     "execution_count": 46,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 23340 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.scatter(  X_test@result.post_mean , tt)\n",
    "np.sqrt(((X_test_scaled@(result.post_mean )/oo  - tt /oo)**2).mean())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "29d144b4",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from sklearn.linear_model import LassoCV\n",
    "from sklearn.preprocessing import StandardScaler\n",
    "\n",
    "def col_scale(X, with_mean=True, with_std=True):\n",
    "    \"\"\"Column-wise standardization.\"\"\"\n",
    "    scaler = StandardScaler(with_mean=with_mean, with_std=with_std)\n",
    "    X_scaled = scaler.fit_transform(X)\n",
    "    return X_scaled, scaler.scale_, scaler.mean_\n",
    "\n",
    "def nash_mdn(X, y, sideinfo, emdn_posterior_means, maxit=10, damping=0.99, eb_kwargs=None):\n",
    "    \"\"\"\n",
    "    Neural Adaptive Shrinkage with Mixture Density Network posterior.\n",
    "    \n",
    "    Args:\n",
    "        X: (n, p) design matrix\n",
    "        y: (n,) response vector\n",
    "        sideinfo: (p, d) side information for MDN prior\n",
    "        emdn_posterior_means: callable returning posterior means from MDN\n",
    "        maxit: maximum number of iterations\n",
    "        damping: damping factor for stability\n",
    "        eb_kwargs: additional args for emdn_posterior_means\n",
    "        \n",
    "    Returns:\n",
    "        Estimated beta (de-standardized)\n",
    "    \"\"\"\n",
    "    # Standardize X and y\n",
    "    X, csd, cmean = col_scale(X)\n",
    "    y, ysd, ymean = col_scale(y.reshape(-1, 1))\n",
    "    y = y.flatten()\n",
    "    n, p = X.shape\n",
    "    sqn = np.sqrt(n)\n",
    "\n",
    "    # Lasso initialization\n",
    "    lasso = LassoCV(cv=5, max_iter=5000).fit(X, y)\n",
    "    beta = lasso.coef_\n",
    "\n",
    "    # First iteration estimate\n",
    "    log_pi, scale, location = None, None, None\n",
    "    drift_comp = 0.5\n",
    "    sigma_0 = 1.0  # Will be updated\n",
    "    pi0 = 0.0\n",
    "\n",
    "    for o in range(maxit):\n",
    "        betahat_list, sebetahat_list = [], []\n",
    "        r = y - X @ beta\n",
    "        beta_new = beta.copy()\n",
    "\n",
    "        for k in range(p):\n",
    "            xk = X[:, k]\n",
    "            r_k = r + xk * beta_new[k]\n",
    "            betahat_k = np.dot(xk, r_k) / n\n",
    "            sebetahat_k = max(np.std(r_k, ddof=1) / sqn, 1e-4)\n",
    "            betahat_list.append(betahat_k)\n",
    "            sebetahat_list.append(sebetahat_k)\n",
    "            r -= xk * (betahat_k - beta_new[k])\n",
    "            beta_new[k] = betahat_k\n",
    "\n",
    "        betahat_arr = np.array(betahat_list)\n",
    "        sebetahat_arr = np.array(sebetahat_list)\n",
    "\n",
    "        # Residual scale and drift update\n",
    "        res_sq = np.sum((y - X @ beta_new) ** 2)\n",
    "        drift_term = np.dot(betahat_arr, betahat_arr - beta_new) / np.sqrt(n + p)\n",
    "        s = np.sqrt((res_sq + drift_term)) / sqn\n",
    "        sigma_0 = np.sqrt((res_sq + np.sum((beta_new - betahat_arr) ** 2)) / (n + p * (1 - pi0))) / sqn\n",
    "        drift_comp = (1 / (n / s**2 + 1 / sigma_0**2)) * (n / s**2)\n",
    "\n",
    "        # Drifted posterior input\n",
    "        if o == 0:\n",
    "            post_input = betahat_arr\n",
    "        else:\n",
    "            post_input = drift_comp * betahat_arr + (1 - drift_comp) * result.post_mean\n",
    "\n",
    "        # MDN posterior update\n",
    "        result = emdn_posterior_means(\n",
    "            X=sideinfo,\n",
    "            betahat=post_input,\n",
    "            sebetahat=np.full_like(post_input, sigma_0),\n",
    "            **(eb_kwargs or {})\n",
    "        )\n",
    "\n",
    "        # Posterior mean update with damping\n",
    "        beta = damping * result.post_mean + (1 - damping) * beta\n",
    "\n",
    "        # Optional: extract pi0 from MDN output\n",
    "        if hasattr(result, \"pi_np\"):\n",
    "            pi0 = result.pi_np[:, 0].mean() if result.pi_np.ndim == 2 else result.pi_np[0]\n",
    "\n",
    "    # Return de-standardized beta\n",
    "    return ysd * beta / csd\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 49,
   "id": "9b666aca",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from sklearn.linear_model import LassoCV\n",
    "from sklearn.preprocessing import StandardScaler\n",
    "\n",
    "def col_scale(X, with_mean=True, with_std=True):\n",
    "    \"\"\"Column-wise standardization.\"\"\"\n",
    "    scaler = StandardScaler(with_mean=with_mean, with_std=with_std)\n",
    "    X_scaled = scaler.fit_transform(X)\n",
    "    return X_scaled, scaler.scale_, scaler.mean_"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dfe1e6d3",
   "metadata": {},
   "outputs": [],
   "source": [
    "def nash_noinfo_vectorized(X, y, maxit=100):\n",
    "    X, csd, cmean = col_scale(X)\n",
    "    y, ysd, ymean = col_scale(y.reshape(-1, 1))\n",
    "    y = y.flatten()\n",
    "\n",
    "    n, p = X.shape\n",
    "    sqn = np.sqrt(n)\n",
    "\n",
    "    # Initialize with Lasso\n",
    "    lasso = LassoCV(cv=10, max_iter=5000).fit(X, y)\n",
    "    beta = lasso.coef_\n",
    "\n",
    "    td_beta = []\n",
    "    drift_comp = 0.5\n",
    "\n",
    "    for o in range(maxit):\n",
    "        full_r = y - X @ beta\n",
    "        beta_matrix = np.tile(beta, (n, 1))\n",
    "        rj = full_r[:, None] + X * beta_matrix\n",
    "\n",
    "        beta_temp = np.sum(X * rj, axis=0) / n\n",
    "        centered_rj = rj - beta_temp\n",
    "        sd_temp = np.std(centered_rj, axis=0, ddof=1) / sqn\n",
    "\n",
    "        if o == 0:\n",
    "            beta = beta_temp.copy()\n",
    "            sigma_0 = np.mean(sd_temp)\n",
    "\n",
    "        b_post = drift_comp * beta_temp + (1 - drift_comp) * beta\n",
    "        resid = y - X @ b_post\n",
    "        s = np.sqrt((np.sum(resid**2) + np.sum((beta_temp - b_post) * b_post)) / np.sqrt(n + p)) / sqn\n",
    "\n",
    "        drift_comp = (1 / (n / s**2 + 1 / sigma_0**2)) * (n / s**2)\n",
    "        ash_input = drift_comp * beta_temp + (1 - drift_comp) * beta\n",
    "\n",
    "        beta, beta2, pi0 = call_r_ash(ash_input, np.full(p, sd_temp ))\n",
    "\n",
    "        resid2 = y - X @ beta\n",
    "        sigma_0 = np.sqrt((np.sum(resid2**2) + np.sum((b_post - beta) * beta)) / (n + p * (1 - pi0))) / sqn\n",
    "\n",
    "        td_beta.append(ysd * beta / csd)\n",
    "\n",
    "    print(\"Final residual variance:\", np.var(y - X @ beta))\n",
    "    print({\"s\": s, \"mean_sd\": np.mean(sd_temp), \"sigma_0\": sigma_0, \"drift_comp\": drift_comp})\n",
    "    return ysd * beta / csd\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 51,
   "id": "265c2427",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 10/50, Loss: 1.1956\n",
      "Epoch 20/50, Loss: -2.4792\n",
      "Epoch 30/50, Loss: -0.7214\n"
     ]
    },
    {
     "ename": "KeyboardInterrupt",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
      "Cell \u001b[1;32mIn[51], line 51\u001b[0m\n\u001b[0;32m     48\u001b[0m     post_input \u001b[38;5;241m=\u001b[39m drift_comp \u001b[38;5;241m*\u001b[39m betahat_arr \u001b[38;5;241m+\u001b[39m (\u001b[38;5;241m1\u001b[39m \u001b[38;5;241m-\u001b[39m drift_comp) \u001b[38;5;241m*\u001b[39m result\u001b[38;5;241m.\u001b[39mpost_mean\n\u001b[0;32m     50\u001b[0m \u001b[38;5;66;03m# MDN posterior update\u001b[39;00m\n\u001b[1;32m---> 51\u001b[0m result \u001b[38;5;241m=\u001b[39m \u001b[43memdn_posterior_means\u001b[49m\u001b[43m(\u001b[49m\n\u001b[0;32m     52\u001b[0m \u001b[43m    \u001b[49m\u001b[43mX\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43msideinfo\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m     53\u001b[0m \u001b[43m    \u001b[49m\u001b[43mbetahat\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mpost_input\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m     54\u001b[0m \u001b[43m    \u001b[49m\u001b[43msebetahat\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mnp\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mfull_like\u001b[49m\u001b[43m(\u001b[49m\u001b[43mpost_input\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43msd_temp\u001b[49m\u001b[43m \u001b[49m\u001b[43m)\u001b[49m\u001b[43m  \u001b[49m\n\u001b[0;32m     55\u001b[0m \u001b[43m\u001b[49m\u001b[43m)\u001b[49m\n\u001b[0;32m     57\u001b[0m \u001b[38;5;66;03m# Posterior mean update with damping\u001b[39;00m\n\u001b[0;32m     58\u001b[0m beta \u001b[38;5;241m=\u001b[39m damping \u001b[38;5;241m*\u001b[39m result\u001b[38;5;241m.\u001b[39mpost_mean \u001b[38;5;241m+\u001b[39m (\u001b[38;5;241m1\u001b[39m \u001b[38;5;241m-\u001b[39m damping) \u001b[38;5;241m*\u001b[39m beta\n",
      "File \u001b[1;32mc:\\Document\\Serieux\\Travail\\python_work\\cEBNM_torch\\py\\empirical_mdn.py:103\u001b[0m, in \u001b[0;36memdn_posterior_means\u001b[1;34m(X, betahat, sebetahat, n_epochs, n_layers, n_gaussians, hidden_dim, batch_size, lr, model_param)\u001b[0m\n\u001b[0;32m    101\u001b[0m model\u001b[38;5;241m.\u001b[39mtrain()\n\u001b[0;32m    102\u001b[0m running_loss \u001b[38;5;241m=\u001b[39m \u001b[38;5;241m0.0\u001b[39m\n\u001b[1;32m--> 103\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m inputs, targets, noise_std \u001b[38;5;129;01min\u001b[39;00m dataloader:\n\u001b[0;32m    104\u001b[0m     optimizer\u001b[38;5;241m.\u001b[39mzero_grad()\n\u001b[0;32m    105\u001b[0m     pi, mu, log_sigma \u001b[38;5;241m=\u001b[39m model(inputs)\n",
      "File \u001b[1;32mc:\\Users\\willi\\anaconda3\\envs\\ml_env\\lib\\site-packages\\torch\\utils\\data\\dataloader.py:701\u001b[0m, in \u001b[0;36m_BaseDataLoaderIter.__next__\u001b[1;34m(self)\u001b[0m\n\u001b[0;32m    698\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_sampler_iter \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[0;32m    699\u001b[0m     \u001b[38;5;66;03m# TODO(https://github.com/pytorch/pytorch/issues/76750)\u001b[39;00m\n\u001b[0;32m    700\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_reset()  \u001b[38;5;66;03m# type: ignore[call-arg]\u001b[39;00m\n\u001b[1;32m--> 701\u001b[0m data \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_next_data\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[0;32m    702\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_num_yielded \u001b[38;5;241m+\u001b[39m\u001b[38;5;241m=\u001b[39m \u001b[38;5;241m1\u001b[39m\n\u001b[0;32m    703\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m (\n\u001b[0;32m    704\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_dataset_kind \u001b[38;5;241m==\u001b[39m _DatasetKind\u001b[38;5;241m.\u001b[39mIterable\n\u001b[0;32m    705\u001b[0m     \u001b[38;5;129;01mand\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_IterableDataset_len_called \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m\n\u001b[0;32m    706\u001b[0m     \u001b[38;5;129;01mand\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_num_yielded \u001b[38;5;241m>\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_IterableDataset_len_called\n\u001b[0;32m    707\u001b[0m ):\n",
      "File \u001b[1;32mc:\\Users\\willi\\anaconda3\\envs\\ml_env\\lib\\site-packages\\torch\\utils\\data\\dataloader.py:757\u001b[0m, in \u001b[0;36m_SingleProcessDataLoaderIter._next_data\u001b[1;34m(self)\u001b[0m\n\u001b[0;32m    755\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m \u001b[38;5;21m_next_data\u001b[39m(\u001b[38;5;28mself\u001b[39m):\n\u001b[0;32m    756\u001b[0m     index \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_next_index()  \u001b[38;5;66;03m# may raise StopIteration\u001b[39;00m\n\u001b[1;32m--> 757\u001b[0m     data \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_dataset_fetcher\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mfetch\u001b[49m\u001b[43m(\u001b[49m\u001b[43mindex\u001b[49m\u001b[43m)\u001b[49m  \u001b[38;5;66;03m# may raise StopIteration\u001b[39;00m\n\u001b[0;32m    758\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_pin_memory:\n\u001b[0;32m    759\u001b[0m         data \u001b[38;5;241m=\u001b[39m _utils\u001b[38;5;241m.\u001b[39mpin_memory\u001b[38;5;241m.\u001b[39mpin_memory(data, \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_pin_memory_device)\n",
      "File \u001b[1;32mc:\\Users\\willi\\anaconda3\\envs\\ml_env\\lib\\site-packages\\torch\\utils\\data\\_utils\\fetch.py:52\u001b[0m, in \u001b[0;36m_MapDatasetFetcher.fetch\u001b[1;34m(self, possibly_batched_index)\u001b[0m\n\u001b[0;32m     50\u001b[0m         data \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mdataset\u001b[38;5;241m.\u001b[39m__getitems__(possibly_batched_index)\n\u001b[0;32m     51\u001b[0m     \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m---> 52\u001b[0m         data \u001b[38;5;241m=\u001b[39m [\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mdataset[idx] \u001b[38;5;28;01mfor\u001b[39;00m idx \u001b[38;5;129;01min\u001b[39;00m possibly_batched_index]\n\u001b[0;32m     53\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m     54\u001b[0m     data \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mdataset[possibly_batched_index]\n",
      "File \u001b[1;32mc:\\Users\\willi\\anaconda3\\envs\\ml_env\\lib\\site-packages\\torch\\utils\\data\\_utils\\fetch.py:52\u001b[0m, in \u001b[0;36m<listcomp>\u001b[1;34m(.0)\u001b[0m\n\u001b[0;32m     50\u001b[0m         data \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mdataset\u001b[38;5;241m.\u001b[39m__getitems__(possibly_batched_index)\n\u001b[0;32m     51\u001b[0m     \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m---> 52\u001b[0m         data \u001b[38;5;241m=\u001b[39m [\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mdataset[idx] \u001b[38;5;28;01mfor\u001b[39;00m idx \u001b[38;5;129;01min\u001b[39;00m possibly_batched_index]\n\u001b[0;32m     53\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m     54\u001b[0m     data \u001b[38;5;241m=\u001b[39m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mdataset[possibly_batched_index]\n",
      "\u001b[1;31mKeyboardInterrupt\u001b[0m: "
     ]
    }
   ],
   "source": [
    "maxit=2\n",
    "    # Standardize X and y\n",
    "X, csd, cmean = col_scale(X)\n",
    "y, ysd, ymean = col_scale(y.reshape(-1, 1))\n",
    "y = y.flatten()\n",
    "n, p = X.shape\n",
    "sqn = np.sqrt(n)\n",
    "\n",
    "    # Lasso initialization\n",
    "lasso = LassoCV(cv=5, max_iter=5000).fit(X, y)\n",
    "beta = lasso.coef_\n",
    "\n",
    "    # First iteration estimate\n",
    "log_pi, scale, location = None, None, None\n",
    "drift_comp = 0.5\n",
    "sigma_0 = 1.0  # Will be updated\n",
    "pi0 = 0.0\n",
    "\n",
    "for o in range(maxit):\n",
    "        betahat_list, sebetahat_list = [], []\n",
    "        r = y - X @ beta\n",
    "        beta_new = beta.copy()\n",
    "\n",
    "        for k in range(p):\n",
    "            xk = X[:, k]\n",
    "            r_k = r + xk * beta_new[k]\n",
    "            betahat_k = np.dot(xk, r_k) / n\n",
    "            sebetahat_k = max(np.std(r_k, ddof=1) / sqn, 1e-4)\n",
    "            betahat_list.append(betahat_k)\n",
    "            sebetahat_list.append(sebetahat_k)\n",
    "            r -= xk * (betahat_k - beta_new[k])\n",
    "            beta_new[k] = betahat_k\n",
    "\n",
    "        betahat_arr = np.array(betahat_list)\n",
    "        sebetahat_arr = np.array(sebetahat_list)\n",
    "\n",
    "        # Residual scale and drift update\n",
    "        res_sq = np.sum((y - X @ beta_new) ** 2)\n",
    "        drift_term = np.dot(betahat_arr, betahat_arr - beta_new) / np.sqrt(n + p)\n",
    "        s = np.sqrt((res_sq + drift_term)) / sqn\n",
    "        sigma_0 = np.sqrt((res_sq + np.sum((beta_new - betahat_arr) ** 2)) / (n + p * (1 - pi0))) / sqn\n",
    "        drift_comp = (1 / (n / s**2 + 1 / sigma_0**2)) * (n / s**2)\n",
    "        sd_temp = sebetahat_arr.mean() / sqn\n",
    "        # Drifted posterior input\n",
    "        if o == 0:\n",
    "            post_input = betahat_arr\n",
    "        else:\n",
    "            post_input = drift_comp * betahat_arr + (1 - drift_comp) * result.post_mean\n",
    "\n",
    "        # MDN posterior update\n",
    "        result = emdn_posterior_means(\n",
    "            X=sideinfo,\n",
    "            betahat=post_input,\n",
    "            sebetahat=np.full_like(post_input, sd_temp )  \n",
    "        )\n",
    "\n",
    "        # Posterior mean update with damping\n",
    "        beta = damping * result.post_mean + (1 - damping) * beta\n",
    "\n",
    "        # Optional: extract pi0 from MDN output\n",
    "        if hasattr(result, \"pi_np\"):\n",
    "            pi0 = result.pi_np[:, 0].mean() if result.pi_np.ndim == 2 else result.pi_np[0]\n",
    "\n",
    "    # Return de-standardized beta\n",
    " \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 52,
   "id": "fdb1e83d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 10/50, Loss: -0.2138\n",
      "Epoch 20/50, Loss: 0.2851\n",
      "Epoch 30/50, Loss: -1.8157\n",
      "Epoch 40/50, Loss: -3.1090\n",
      "Epoch 50/50, Loss: -0.5162\n"
     ]
    }
   ],
   "source": [
    "result = emdn_posterior_means(\n",
    "            X=sideinfo,\n",
    "            betahat=post_input,\n",
    "            sebetahat=np.full_like(post_input, sd_temp )  \n",
    "        )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 70,
   "id": "0e462bf7",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[<matplotlib.lines.Line2D at 0x1c78bd40730>]"
      ]
     },
     "execution_count": 70,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 31896 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.plot( result.location[:,0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 109,
   "id": "f2edc2a7",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "[<matplotlib.lines.Line2D at 0x1c78b815d80>]"
      ]
     },
     "execution_count": 109,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 32056 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.plot(result.post_mean)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 66,
   "id": "18ed259b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "6.918259143829347"
      ]
     },
     "execution_count": 66,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 20400 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.scatter(  (X_test@result.post_mean -(X_test@result.post_mean).mean()) , (y_test-y_test.mean()))\n",
    "np.sqrt(((  (X_test@result.post_mean -(X_test@result.post_mean).mean()) -y_test.mean())  )**2).mean()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 73,
   "id": "e2c3f81a",
   "metadata": {},
   "outputs": [],
   "source": [
    "def nash_mdn(X, y, sideinfo=None, maxit=4, damping=0.99, eb_kwargs=None):\n",
    "    k=0 \n",
    "    X = X.astype(np.float32)\n",
    "    y = y.astype(np.float32).flatten()\n",
    "    X, csd = col_scale(X)\n",
    "    y, ysd = col_scale(y.reshape(-1, 1))\n",
    "    y = y.flatten()\n",
    "\n",
    "    n, p = X.shape\n",
    "    sqn = np.sqrt(n).astype(np.float32)\n",
    "    beta = LassoCV(cv=5, max_iter=5000).fit(X, y).coef_\n",
    "\n",
    "    log_pi, scale, location = None, None, None\n",
    "    td_beta = []\n",
    "    model_param = None\n",
    "    pi0=0\n",
    "    \n",
    "    for o in range(maxit):\n",
    "        betahat_list, sebetahat_list = [], []\n",
    "        for k in range(p):\n",
    "            r = y - X @ beta\n",
    "            beta_new = beta.copy()\n",
    "            xk = X[:, k]\n",
    "            r_k = r + xk * beta_new[k]\n",
    "            betahat_k = np.dot(xk, r_k) / n\n",
    "            sebetahat_k = max(np.std(r_k) / sqn, 1e-4)\n",
    "            if 0<1: \n",
    "\n",
    "                beta_k_new = betahat_k\n",
    "                beta_k_new = damping * beta_k_new + (1 - damping) * beta_new[k]\n",
    "             \n",
    "            else :\n",
    "                print(o)\n",
    "                pm = posterior_mean_norm(\n",
    "                    betahat=np.array([betahat_k]),\n",
    "                    sebetahat=np.array([sebetahat_k]),\n",
    "                    log_pi=np.log( np.clip( result.pi_np[k, :], 1e-12, 1.0)),\n",
    "                    location=result.location[k, :],\n",
    "                    scale=result.scale[k, :]\n",
    "                    )\n",
    "             \n",
    "                beta_k_new = pm.post_mean[0]\n",
    "\n",
    "            r += xk * (beta_new[k]  - beta_k_new)\n",
    "            beta_k_new = damping * beta_k_new + (1 - damping) * beta_new[k]\n",
    "            beta_new[k]=beta_k_new\n",
    "            betahat_list.append(betahat_k)\n",
    "            sebetahat_list.append(sebetahat_k) \n",
    "\n",
    "\n",
    "        betahat_arr = np.array(betahat_list)\n",
    "        avg_se = np.mean(sebetahat_list)\n",
    "        beta = beta_new.copy()\n",
    "        res_sq_final = np.sum((y - X @ beta ) ** 2)\n",
    "        sigma_0_term = np.dot(beta, beta - betahat_arr )   \n",
    " \n",
    "        beta = beta_new.copy()\n",
    "        res_sq_final = np.sum((y - X @ beta ) ** 2)\n",
    "        sigma_0_term = np.dot(beta, beta - betahat_arr )\n",
    "        denom = n + p * (1 - pi0)\n",
    "        sigma_0 = np.sqrt(max((res_sq_final + sigma_0_term) / denom, 1e-8)) / sqn\n",
    "\n",
    "        res_sq = np.sum((y - X @ betahat_arr) ** 2)\n",
    "        drift_term = np.dot(betahat_arr,  betahat_arr) / np.sqrt(n + p)\n",
    "        s = np.sqrt((res_sq + drift_term)) / sqn\n",
    "\n",
    "        # Update drift_comp\n",
    "        drift_comp = (1 / (n / s**2 + 1 / sigma_0**2)) * (n / s**2)\n",
    "        if o<1: \n",
    "            # Update beta with ash\n",
    "            ash_input = betahat_arr   \n",
    "            result =emdn_posterior_means(X= sideinfo,\n",
    "                                            betahat= ash_input,\n",
    "                                            sebetahat=np.array(sebetahat_list)   )\n",
    "            emdn_posterior_means(X= sideinfo,\n",
    "                                          betahat= ash_input,\n",
    "                                          sebetahat= np.full_like(ash_input,sigma_0), **(eb_kwargs or {}))\n",
    "        else :\n",
    "            ash_input = drift_comp * betahat_arr + (1 - drift_comp) * result.post_mean\n",
    "            result =emdn_posterior_means(X= sideinfo,\n",
    "                                          betahat= ash_input,\n",
    "                                          sebetahat=np.full_like(ash_input,sigma_0)   )\n",
    "            log_pi= result.pi_np\n",
    "            beta =   result.post_mean\n",
    "\n",
    "\n",
    "  \n",
    "    out=   ysd * result.post_mean/ csd\n",
    "    return out\n",
    " \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 77,
   "id": "dabab76f",
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "def col_scale(X, with_mean=True, with_std=True):\n",
    "    scaler = StandardScaler(with_mean=with_mean, with_std=with_std)\n",
    "    X_scaled = scaler.fit_transform(X)\n",
    "    return X_scaled, scaler.scale_\n",
    "\n",
    "\n",
    "def call_r_ash_fit_all_with_postmean(beta, sigma):\n",
    "    \"\"\"\n",
    "    Calls R's ash function and returns:\n",
    "    - log mixture weights\n",
    "    - mixture standard deviations\n",
    "    - posterior mean of beta\n",
    "    \"\"\"\n",
    "    from rpy2.rinterface_lib.sexp import NULLType\n",
    "\n",
    "    sebetahat = np.full_like(beta, sigma)\n",
    "    ash_obj = ashr.ash(betahat=beta, sebetahat=sebetahat, mixcompdist=\"normal\")\n",
    "\n",
    "    fitted_g = ash_obj.rx2(\"fitted_g\")\n",
    "    pi_r = np.array(fitted_g.rx2(\"pi\"), dtype=np.float32)\n",
    "    scale_r = np.array(fitted_g.rx2(\"sd\"), dtype=np.float32)\n",
    "\n",
    "    posterior_mean_r = ash_obj.rx2(\"result\").rx2(\"PosteriorMean\")\n",
    "    if isinstance(posterior_mean_r, NULLType):\n",
    "        raise RuntimeError(\"R ash() returned NULL for result$PosteriorMean\")\n",
    "    posterior_mean = np.array(posterior_mean_r, dtype=np.float32)\n",
    "\n",
    "    log_pi = np.log(np.clip(pi_r, 1e-12, 1.0))\n",
    "    return log_pi, scale_r, posterior_mean\n",
    "\n",
    "\n",
    "def get_data_loglik_normal(betahat, sebetahat, location, scale):\n",
    "    var = sebetahat[:, None] ** 2 + scale[None, :] ** 2\n",
    "    return -0.5 * (np.log(2 * np.pi * var) + (betahat[:, None] - location) ** 2 / var)\n",
    "\n",
    "\n",
    "def apply_log_sum_exp(data_loglik, log_pi):\n",
    "    combined_loglik = data_loglik + log_pi\n",
    "    return combined_loglik - logsumexp(combined_loglik, axis=1)[:, None]\n",
    "\n",
    "\n",
    "class PosteriorMeanNorm:\n",
    "    def __init__(self, post_mean, post_mean2, post_sd):\n",
    "        self.post_mean = post_mean\n",
    "        self.post_mean2 = post_mean2\n",
    "        self.post_sd = post_sd\n",
    "\n",
    "\n",
    "def posterior_mean_norm(betahat, sebetahat, log_pi, scale, location=None):\n",
    "    if location is None:\n",
    "        location = np.zeros_like(scale)\n",
    "\n",
    "    data_loglik = get_data_loglik_normal(betahat, sebetahat, location, scale)\n",
    "    log_post_assignment = apply_log_sum_exp(data_loglik, log_pi)\n",
    "\n",
    "    n, K = betahat.shape[0], scale.shape[0]\n",
    "    var = np.zeros((n, K))\n",
    "\n",
    "    if scale[0] == 0:\n",
    "        for i in range(n):\n",
    "            var[i, :] = np.concatenate(([0], 1 / ((1 / sebetahat[i]**2) + (1 / scale[1:]**2))))\n",
    "    else:\n",
    "        for i in range(n):\n",
    "            var[i, :] = 1 / ((1 / sebetahat[i]**2) + (1 / scale**2))\n",
    "\n",
    "    temp = np.zeros((n, K))\n",
    "    for i in range(n):\n",
    "        temp[i, :] = (var[i, :] / sebetahat[i]**2) * betahat[i] + location * (1 - var[i, :] / sebetahat[i]**2)\n",
    "\n",
    "    post_mean = np.sum(np.exp(log_post_assignment) * temp, axis=1)\n",
    "    post_mean2 = np.sum(np.exp(log_post_assignment) * (var + temp**2), axis=1)\n",
    "    post_sd = np.sqrt(post_mean2 - post_mean**2)\n",
    "    return PosteriorMeanNorm(post_mean, post_mean2, post_sd)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 78,
   "id": "1fd6fd5e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Epoch 10/50, Loss: -0.5917\n",
      "Epoch 20/50, Loss: -1.0131\n",
      "Epoch 30/50, Loss: -2.5789\n",
      "Epoch 40/50, Loss: -1.0748\n",
      "Epoch 50/50, Loss: -1.0967\n",
      "Epoch 10/50, Loss: 0.7220\n",
      "Epoch 20/50, Loss: -1.4069\n",
      "Epoch 30/50, Loss: -1.5993\n",
      "Epoch 40/50, Loss: -1.1915\n",
      "Epoch 50/50, Loss: -3.6560\n",
      "Epoch 10/50, Loss: -2.8376\n",
      "Epoch 20/50, Loss: -1.0421\n",
      "Epoch 30/50, Loss: -0.7526\n",
      "Epoch 40/50, Loss: -2.5841\n",
      "Epoch 50/50, Loss: -1.8400\n"
     ]
    }
   ],
   "source": [
    "beta_mdn =  nash_mdn(X=X_train,y=y_train, sideinfo=sideinfo_array,maxit=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 93,
   "id": "f2487c46",
   "metadata": {},
   "outputs": [],
   "source": [
    "X_scaled = scaler.fit_transform(X_train)\n",
    " \n",
    "y, ysd = col_scale(y_train.reshape(-1, 1))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 104,
   "id": "5494e0b2",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([1.2893138])"
      ]
     },
     "execution_count": 104,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ysd"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 103,
   "id": "08cccccb",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1.2512114872369255"
      ]
     },
     "execution_count": 103,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.sqrt(np.mean( (X_scaled@beta_mdn/9- y)**2 ))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 106,
   "id": "25173a1b",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[1.        , 0.83974611],\n",
       "       [0.83974611, 1.        ]])"
      ]
     },
     "execution_count": 106,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.corrcoef(X_test@beta_mdn, y_test)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 108,
   "id": "83f77ffc",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[1.        , 0.84936762],\n",
       "       [0.84936762, 1.        ]])"
      ]
     },
     "execution_count": 108,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.corrcoef(X_test@lasso.coef_ , y_test)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 97,
   "id": "3810fd5e",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.collections.PathCollection at 0x1c78b8148e0>"
      ]
     },
     "execution_count": 97,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 31056 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.scatter(X_scaled@beta_mdn/5, y)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 105,
   "id": "5545e938",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.collections.PathCollection at 0x1c78b7fea40>"
      ]
     },
     "execution_count": 105,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "[encoded data, 17260 characters]",
      "text/plain": [
       "<Figure size 640x480 with 1 Axes>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.scatter( lasso.coef_,beta_mdn)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "ml_env",
   "language": "python",
   "name": "python3"
  },
  "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.10.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
