direct_marginal_likelihood.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "import statsmodels.api as sm\n",
    "from tqdm.auto import tqdm\n",
    "from scipy.special import logit,expit\n",
    "\n",
    "from src.IterativeFitting import IterativeFitting as IF\n",
    "from src.HamMethods import ham_vanilla, ham_solved\n",
    "from src.CorrFuncs import covariance_matrix, trend_est\n",
    "\n",
    "import jax.numpy as jnp\n",
    "import jax \n",
    "jax.config.update('jax_enable_x64',True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Optimization terminated successfully.\n",
      "         Current function value: 0.646961\n",
      "         Iterations 5\n"
     ]
    }
   ],
   "source": [
    "# Setting seed\n",
    "np.random.seed(23)\n",
    "\n",
    "# Initializing simulation slope and intercept parameters\n",
    "beta0 = -0.1\n",
    "beta1 = 0.02\n",
    "beta_age = 0.05\n",
    "beta_income = -0.05\n",
    "beta_sex = 1.0\n",
    "\n",
    "# Number of samples\n",
    "n = 40000\n",
    "\n",
    "# Constructing covariate vectors for additional covariates to adjust\n",
    "age_vector = np.random.uniform(low=0,high=1,size=n)\n",
    "income_vector = np.random.uniform(low=0,high=1,size=n)\n",
    "# sex_vector = np.random.binomial(n=1,p=0.6,size=n)\n",
    "\n",
    "# Assigning exposure levels and getting relevant x's\n",
    "exposure_levels = np.array([0.,1.,2.,3.])\n",
    "xs = np.random.choice(exposure_levels,size = n)\n",
    "sex_vector = np.random.binomial(n = 1,p = 0.1 + xs/4,size =n )\n",
    "\n",
    "\n",
    "# Function to generate probabilities of being a case v non-case\n",
    "other_covs = beta0 + (beta_age*age_vector) + (income_vector*beta_income) + (sex_vector*beta_sex) \n",
    "p = lambda x: (np.exp(beta1*x + other_covs))/(1 + np.exp(beta1*x + other_covs))\n",
    "\n",
    "# Actually calculating probabilities on exposures as defined\n",
    "px = p(xs)\n",
    "\n",
    "# Actually assigning to case or not\n",
    "outcomes = np.array([np.random.binomial(n=1,p=p,size=1)[0] for p in px])\n",
    "\n",
    "# Constructing and sorting dataframe of outcomes and exposure\n",
    "df = np.stack([outcomes,xs],axis=1)\n",
    "# df = df[np.argsort(df[:, 1])]\n",
    "\n",
    "# Getting xs corresponding to each exposure interval and corresponding case information\n",
    "categories = [df[np.where(df[:,1] == i)[0]] for i in exposure_levels]\n",
    "N = np.array([df.shape[0] for df in categories])\n",
    "cases = np.array([np.sum(df[:,0]) for df in categories])\n",
    "non_cases = N - cases\n",
    "a_0, A = cases[0], cases[1:]\n",
    "b_0, B = non_cases[0], non_cases[1:]\n",
    "M = np.sum(cases)\n",
    "\n",
    "# Getting crude estimates to compare\n",
    "crude = lambda x,y: (x/y) / (a_0/b_0)\n",
    "crudeors = crude(A,B)\n",
    "\n",
    "# Constructing dummy variables with respect to exposure level\n",
    "x = exposure_levels[1:]\n",
    "\n",
    "to_adjust_df = pd.DataFrame({\n",
    "    'age':age_vector,\n",
    "    'income':income_vector,\n",
    "    'sex':sex_vector,\n",
    "    'exposure':xs.astype(int)\n",
    "})\n",
    "to_adjust_df['intercept']=1\n",
    "\n",
    "to_adjust_df = pd.concat([to_adjust_df,pd.get_dummies(to_adjust_df['exposure'],prefix='exposure').astype(int).iloc[:,1:]],axis=1).drop('exposure',axis=1)\n",
    "\n",
    "results = sm.Logit(endog=df[:,0], exog=to_adjust_df).fit()\n",
    "true_C = results.cov_params().to_numpy()[4:,4:]\n",
    "true_C_inv = np.linalg.inv(true_C)\n",
    "\n",
    "# Get log-ORs and corresponding variances\n",
    "L = results.params.to_numpy()[4:]\n",
    "v = (results.bse.to_numpy()[4:])**2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Running GL\n",
    "it_fit_ex = IF(L,A,N,M)\n",
    "Ap, Bp, a0p, b0p = it_fit_ex.convexProgram()\n",
    "\n",
    "gl_C = covariance_matrix(Ap,Bp,a0p,b0p,v)\n",
    "inv_gl_C = np.linalg.inv(gl_C)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<matplotlib.collections.PathCollection at 0x7fdb07e70460>"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAYIAAAD4CAYAAADhNOGaAAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjcuMSwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/bCgiHAAAACXBIWXMAAAsTAAALEwEAmpwYAAARR0lEQVR4nO3df6zdd13H8efLtsMCurasjq6ldoalZhB15GQgI2bKWItKWhMwI8Y0BtNEIdE/bLJJZDo0mc4EY9SQBkgqUQbBOar8KGWD+IMwdscGY0htHeh6N9ikK4gUWJe3f5xv8fRy7npPz+k5vf08H8nJ/X4+38+53/e++/S8zvf7OffeVBWSpHb9wKwLkCTNlkEgSY0zCCSpcQaBJDXOIJCkxq2cdQFn45JLLqktW7bMugxJWlbuu+++/66q9Qv7l2UQbNmyhbm5uVmXIUnLSpL/HNbvrSFJapxBIEmNMwgkqXEGgSQ1ziCQpMYty08NnY0775/ntgOHePT4CS5bs5o927ay86qNsy5LkmauiSC48/55brrjQU489TQA88dPcNMdDwIYBpKa18StodsOHPpeCJxy4qmnue3AoRlVJEnnjyaC4NHjJ0bql6SWNBEEl61ZPVK/JLWkiSDYs20rq1etOK1v9aoV7Nm2dUYVSdL5o4nF4lMLwn5qSJK+XxNBAP0w8IVfkr5fE7eGJEmLMwgkqXEGgSQ1ziCQpMYZBJLUOINAkhpnEEhS4wwCSWqcQSBJjTMIJKlxBoEkNc4gkKTGGQSS1LiJBEGS7UkOJTmS5MYh+5+V5L3d/nuSbOn6n5fk40m+meQvJlGLJGk0YwdBkhXAXwKvBq4EXp/kygXD3gA8WVUvBN4G/HHX/23g94DfGbcOSdLZmcQVwdXAkap6uKq+C9wO7FgwZgewr9t+P/DKJKmq/62qf6EfCJKkGZhEEGwEHhloH+36ho6pqpPA14HnjXKQJLuTzCWZe+KJJ8YoV5I0aNksFlfV3qrqVVVv/fr1sy5Hki4YkwiCeeAFA+1NXd/QMUlWAhcDX5vAsSVJY5pEENwLXJHk8iQXATcA+xeM2Q/s6rZfC9xdVTWBY0uSxjT2H6+vqpNJ3gQcAFYA76qqh5LcAsxV1X7gncC7kxwBjtEPCwCSfBn4YeCiJDuB66vqC+PWJUlamrGDAKCqPgR8aEHfWwa2vw28bpHnbplEDZKks7NsFoslSeeGQSBJjTMIJKlxBoEkNc4gkKTGGQSS1DiDQJIaZxBIUuMMAklqnEEgSY0zCCSpcQaBJDXOIJCkxhkEktQ4g0CSGmcQSFLjDAJJapxBIEmNMwgkqXEGgSQ1ziCQpMYZBJLUOINAkhpnEEhS4wwCSWqcQSBJjTMIJKlxBoEkNc4gkKTGGQSS1LiJBEGS7UkOJTmS5MYh+5+V5L3d/nuSbBnYd1PXfyjJtknUI0laurGDIMkK4C+BVwNXAq9PcuWCYW8AnqyqFwJvA/64e+6VwA3Ai4DtwF9130+SNCWTuCK4GjhSVQ9X1XeB24EdC8bsAPZ12+8HXpkkXf/tVfWdqvoScKT7fpKkKZlEEGwEHhloH+36ho6pqpPA14HnLfG5ACTZnWQuydwTTzwxgbIlSbCMFouram9V9aqqt379+lmXI0kXjEkEwTzwgoH2pq5v6JgkK4GLga8t8bmSpHNoEkFwL3BFksuTXER/8Xf/gjH7gV3d9muBu6uquv4buk8VXQ5cAXx6AjVJ0gXjzvvnuebWu7n8xg9yza13c+f9k32/vHLcb1BVJ5O8CTgArADeVVUPJbkFmKuq/cA7gXcnOQIcox8WdOPeB3wBOAm8saqeHrcmSbpQ3Hn/PDfd8SAnnuq/NM4fP8FNdzwIwM6rhi6pjiz9N+bLS6/Xq7m5uVmXIUnn3DW33s388RPf179xzWr+9cafG+l7JbmvqnoL+5fNYrEktejRISHwTP1nwyCQpPPYZWtWj9R/NgwCSTqP7dm2ldWrTv+FC6tXrWDPtq0TO8bYi8WSpHPn1ILwbQcO8ejxE1y2ZjV7tm2d2EIxGASSdN7bedXGib7wL+StIUlqnEEgSY0zCCSpcQaBJDXOIJCkxhkEktQ4g0CSGmcQSFLjDAJJapxBIEmNMwgkqXEGgSQ1ziCQpMYZBJLUOINAkhpnEEhS4wwCSWqcQSBJjTMIJKlxBoEkNc4gkKTGGQSS1DiDQJIaZxBIUuMMAklq3FhBkGRdkoNJDndf1y4yblc35nCSXQP9f5TkkSTfHKcOSdLZG/eK4Ebgrqq6Arira58myTrgZuClwNXAzQOB8Q9dnyRpRsYNgh3Avm57H7BzyJhtwMGqOlZVTwIHge0AVfWpqnpszBokSWMYNwguHXgh/wpw6ZAxG4FHBtpHuz5J0nlg5ZkGJPkY8Pwhu9482KiqSlKTKmxIHbuB3QCbN28+V4eRpOacMQiq6rrF9iX5apINVfVYkg3A40OGzQPXDrQ3AZ8YsU6qai+wF6DX652zwJGk1ox7a2g/cOpTQLuADwwZcwC4PsnabpH4+q5PknQeGDcIbgVeleQwcF3XJkkvyTsAquoY8Fbg3u5xS9dHkj9JchR4dpKjSX5/zHokSSNK1fK7y9Lr9Wpubm7WZUjSspLkvqrqLez3J4slqXEGgSQ1ziCQpMYZBJLUOINAkhpnEEhS4wwCSWqcQSBJjTMIJKlxBoEkNc4gkKTGGQSS1DiDQJIaZxBIUuMMAklqnEEgSY0zCCSpcQaBJDXOIJCkxhkEktQ4g0CSGmcQSFLjDAJJapxBIEmNMwgkqXEGgSQ1ziCQpMYZBJLUOINAkhpnEEhS48YKgiTrkhxMcrj7unaRcbu6MYeT7Or6np3kg0m+mOShJLeOU4sk6eyMe0VwI3BXVV0B3NW1T5NkHXAz8FLgauDmgcD406r6ceAq4Jokrx6zHknSiMYNgh3Avm57H7BzyJhtwMGqOlZVTwIHge1V9a2q+jhAVX0X+Aywacx6JEkjGjcILq2qx7rtrwCXDhmzEXhkoH206/ueJGuA19C/qpAkTdHKMw1I8jHg+UN2vXmwUVWVpEYtIMlK4D3An1fVw88wbjewG2Dz5s2jHkaStIgzBkFVXbfYviRfTbKhqh5LsgF4fMiweeDagfYm4BMD7b3A4ar6szPUsbcbS6/XGzlwJEnDjXtraD+wq9veBXxgyJgDwPVJ1naLxNd3fST5Q+Bi4LfHrEOSdJbGDYJbgVclOQxc17VJ0kvyDoCqOga8Fbi3e9xSVceSbKJ/e+lK4DNJHkjy62PWI0kaUaqW312WXq9Xc3Nzsy5DkpaVJPdVVW9hvz9ZLEmNMwgkqXEGgSQ1ziCQpMYZBJLUOINAkhpnEEhS4wwCSWqcQSBJjTMIJKlxBoEkNc4gkKTGGQSS1DiDQJIaZxBIUuMMAklqnEEgSY0zCCSpcQaBJDXOIJCkxhkEktQ4g0CSGmcQSFLjDAJJapxBIEmNMwgkqXEGgSQ1ziCQpMYZBJLUOINAkho3VhAkWZfkYJLD3de1i4zb1Y05nGTXQP9Hknw2yUNJ3p5kxTj1SJJGN+4VwY3AXVV1BXBX1z5NknXAzcBLgauBmwcC45er6ieBFwPrgdeNWY8kaUTjBsEOYF+3vQ/YOWTMNuBgVR2rqieBg8B2gKr6RjdmJXARUGPWI0ka0bhBcGlVPdZtfwW4dMiYjcAjA+2jXR8ASQ4AjwP/A7x/zHokSSNaeaYBST4GPH/IrjcPNqqqkoz8jr6qtiX5QeBvgJ+jf8UwrI7dwG6AzZs3j3oYSdIizhgEVXXdYvuSfDXJhqp6LMkG+u/sF5oHrh1obwI+seAY307yAfq3moYGQVXtBfYC9Ho9byFJ0oSMe2toP3DqU0C7gA8MGXMAuD7J2m6R+HrgQJLnduFBkpXALwBfHLMeSdKIxg2CW4FXJTkMXNe1SdJL8g6AqjoGvBW4t3vc0vU9B9if5HPAA/SvJt4+Zj2SpBGlavndZen1ejU3NzfrMiRpWUlyX1X1Fvb7k8WS1DiDQJIaZxBIUuMMAklqnEEgSY0zCCSpcQaBJDXOIJCkxhkEktQ4g0CSGmcQSFLjDAJJatwZ/x6BpDO78/55bjtwiEePn+CyNavZs20rO6/aeOYnSucBg0Aa0533z3PTHQ9y4qmnAZg/foKb7ngQwDDQsuCtIWlMtx049L0QOOXEU09z24FDM6pIGo1BII3p0eMnRuqXzjcGgTSmy9asHqlfOt8YBNKY9mzbyupVK07rW71qBXu2bZ1RRdJoXCyWxnRqQdhPDWm5MgikCdh51UZf+LVseWtIkhpnEEhS4wwCSWqcQSBJjTMIJKlxqapZ1zCyJE8A/3mWT78E+O8JljMp1jUa6xqNdY3mQq3rR6tq/cLOZRkE40gyV1W9WdexkHWNxrpGY12jaa0ubw1JUuMMAklqXItBsHfWBSzCukZjXaOxrtE0VVdzawSSpNO1eEUgSRpgEEhS4y6YIEjyriSPJ/n8IvuT5M+THEnyuSQvGdi3K8nh7rFrynX9SlfPg0k+meQnB/Z9uet/IMnclOu6NsnXu2M/kOQtA/u2JznUncsbp1zXnoGaPp/k6STrun3n8ny9IMnHk3whyUNJfmvImKnPsSXWNfU5tsS6pj7HlljX1OdYkh9M8ukkn+3q+oMhY56V5L3dObknyZaBfTd1/YeSbBu5gKq6IB7AzwAvAT6/yP6fBz4MBHgZcE/Xvw54uPu6ttteO8W6Xn7qeMCrT9XVtb8MXDKj83Ut8I9D+lcA/wH8GHAR8FngymnVtWDsa4C7p3S+NgAv6bZ/CPj3hf/ds5hjS6xr6nNsiXVNfY4tpa5ZzLFuzjy3214F3AO8bMGY3wTe3m3fALy3276yO0fPAi7vzt2KUY5/wVwRVNU/AceeYcgO4K+r71PAmiQbgG3Awao6VlVPAgeB7dOqq6o+2R0X4FPApkkde5y6nsHVwJGqeriqvgvcTv/czqKu1wPvmdSxn0lVPVZVn+m2/wf4N2DhHyCY+hxbSl2zmGNLPF+LOWdz7Czqmsoc6+bMN7vmqu6x8JM8O4B93fb7gVcmSdd/e1V9p6q+BByhfw6X7IIJgiXYCDwy0D7a9S3WPwtvoP+O8pQCPprkviS7Z1DPT3eXqh9O8qKu77w4X0meTf/F9O8GuqdyvrpL8qvov2sbNNM59gx1DZr6HDtDXTObY2c6X9OeY0lWJHkAeJz+G4dF51dVnQS+DjyPCZwv/0LZeSLJz9L/R/qKge5XVNV8kh8BDib5YveOeRo+Q//3knwzyc8DdwJXTOnYS/Ea4F+ravDq4ZyfryTPpf/C8NtV9Y1Jfu9xLKWuWcyxM9Q1szm2xP+PU51jVfU08FNJ1gB/n+TFVTV0rWzSWroimAdeMNDe1PUt1j81SX4CeAewo6q+dqq/qua7r48Df8+Il3vjqKpvnLpUraoPAauSXMJ5cL46N7Dgkv1cn68kq+i/ePxNVd0xZMhM5tgS6prJHDtTXbOaY0s5X52pz7Huex8HPs733z783nlJshK4GPgakzhfk170mOUD2MLii5+/wOkLeZ/u+tcBX6K/iLe22143xbo207+n9/IF/c8Bfmhg+5PA9inW9Xz+/wcOrwb+qzt3K+kvdl7O/y/kvWhadXX7L6a/jvCcaZ2v7r/9r4E/e4YxU59jS6xr6nNsiXVNfY4tpa5ZzDFgPbCm214N/DPwiwvGvJHTF4vf122/iNMXix9mxMXiC+bWUJL30P8UwiVJjgI3019woareDnyI/qc6jgDfAn6t23csyVuBe7tvdUudfil4rut6C/37fH/VX/fhZPV/u+Cl9C8Pof8P42+r6iNTrOu1wG8kOQmcAG6o/qw7meRNwAH6n+54V1U9NMW6AH4J+GhV/e/AU8/p+QKuAX4VeLC7jwvwu/RfZGc5x5ZS1yzm2FLqmsUcW0pdMP05tgHYl2QF/Ts176uqf0xyCzBXVfuBdwLvTnKEfkjd0NX8UJL3AV8ATgJvrP5tpiXzV0xIUuNaWiOQJA1hEEhS4wwCSWqcQSBJjTMIJKlxBoEkNc4gkKTG/R+ojndilWsRFgAAAABJRU5ErkJggg==",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "plt.scatter(x,L)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 68,
   "metadata": {},
   "outputs": [],
   "source": [
    "import jax\n",
    "def marginal_like(gamma2):\n",
    "    n_groups = 3\n",
    "    I =jnp.eye(n_groups)\n",
    "    oneoneT = jnp.ones((n_groups,n_groups))\n",
    "    D = jnp.diag(v)\n",
    "    Ca = D - gamma2*I + gamma2*oneoneT\n",
    "    sign,logabsdet = jnp.linalg.slogdet(Ca)\n",
    "    invC = jnp.linalg.inv(Ca)\n",
    "    beta = x.T@invC@L/(x.T@invC@x)\n",
    "    res = L-x*beta\n",
    "    res = 0*x\n",
    "    correction =  (1/2)*jnp.log(x.T@invC@x)\n",
    "\n",
    "    return jnp.array([0.5 * (res).T@invC@(res) + 0.5 * logabsdet + correction,sign])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 70,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAXIAAAD/CAYAAADsfV27AAAAOXRFWHRTb2Z0d2FyZQBNYXRwbG90bGliIHZlcnNpb24zLjcuMSwgaHR0cHM6Ly9tYXRwbG90bGliLm9yZy/bCgiHAAAACXBIWXMAAAsTAAALEwEAmpwYAAAfLUlEQVR4nO3dd3wc1b338c9Pq2oVy03uwr1hXEUNzWBKTAtcYiAQCBBMC09IuHAB35sbLkkIIcCl5HnASQgQeg3VxJgOCcFyr9jG3XKVLNuSVXfP84dksI1klZ3d2fJ9v1778s7s7szveO2vjs6cmTHnHCIiEr9S/C5ARETCoyAXEYlzCnIRkTinIBcRiXMKchGROKcgFxGJc2EFuZl938wWm1nIzIq8KkpERFov3B75IuA84GMPahERkXZIDefDzrmlAGbmTTUiItJmYQV5e3Xt2tX169fPj12LiMSt2bNnb3fOdTtwfYtBbmYzgR5NvDTVOfdaawswsynAFIDCwkKKi4tb+1EREQHMbG1T61sMcufcRC8KcM5NA6YBFBUV6QIvIiIe0fRDEZE4F+70w3PNbANwNPCWmf3dm7JERKS1wp218irwqke1iIhIO2hoRUQkzinIRUTinIJcRCQKnHO8PHsD9cGQ59v25YQgEZFkc9+7y3no/ZWkp6Zw1uhenm5bPXIRkQh7YdZ6Hnp/JRce3pczR/X0fPsKchGRCPpkxTZuf3Uhxw3uyp3fGxmRa1MpyEVEIuTLzbu57qk5DCrI4Q8XjyMtEJnIVZCLiETAll3VXP6XL8hKD/DYjw4nLzMtYvvSwU4REY9V1tRz5ROzKK+q44Wrj6ZXflZE96ceuYiIh4Ihx0+fm8uSkl08/IOxjOzdMeL7VI9cRMRDd765hJlLt3LnOYdy0rDuUdmneuQiIh557NPVPP6PNVx1XH9+eHS/qO1XQS4i4oEZizdz51tLOP3QHtz23eFR3beCXEQkTAs2lPPT5+Yxqk8+918whpSU6N7HWEEuIhKGDTv2cMXjxXTJSedPlxaRlR6Ieg062Cki0k67quu44vFZ1NQHefaqI+mWm+FLHeqRi4i0Q10wxHVPzWHVtkoeuWQ8g7vn+laLeuQiIm3knOO//raIT1du53fnj+I7g7r6Wo965CIibfTIR6t4btZ6rp8wkMlFff0uR0EuItIWby3YxN3vLOPMUT256ZShfpcDKMhFRFptzrod/PyFeYw/pBO///7oqE8zbE5YQW5m95jZMjNbYGavmlm+R3WJiMSU9WV7uOqJYrrnZTLth+PJTIv+NMPmhNsjfxcY6ZwbBSwHbgu/JBGR2LKzqo7LH59Ffcjxl8sPp0uOP9MMmxNWkDvnZjjn6hsXPwf6hF+SiEjsqAuGuO7p2awtbZhmOLBbjt8lfYuXY+RXANM93J6IiK/2TjP8bGUpd503iqMHdvG7pCa1OI/czGYCPZp4aapz7rXG90wF6oGnD7KdKcAUgMLCwnYVKyISTY9+3DDN8CcTBnH++NgdcGgxyJ1zEw/2upn9CDgTONk55w6ynWnANICioqJm3yciEgumL9zEb6c3TDP8+SlD/C7noMI6s9PMTgduAU5wzu3xpiQREX/NX1/Ojc/PY1xhfkxNM2xOuGPkDwO5wLtmNs/MHvGgJhER32wsr+LKJ4rplpvBtEuLYmqaYXPC6pE75wZ5VYiIiN92V9dx5T5XM+waY9MMm6OLZomIAPXBEDc8O5cVWyt44vIjfL2aYVvpFH0RERpumvzhl9u485yRHDvY36sZtpWCXESS3uOfreaJf67lx8f25wdHxt/0aAW5iCS1D5Zt5X/eXMLE4d25bVJ0b5rsFQW5iCStZZt3ccOzcxnWI48HLhxDIManGTZHQS4iSWnb7hqufLyY7IwAf/5REdkZ8Tv3I34rFxFpp+q6IFP+WkxpZQ0vXn0MPTtm+V1SWBTkIpJUnHPc8tIC5q4r55FLxnFYn45+lxQ2Da2ISFJ58L2VvD6/hJtPG8rpI3v6XY4nFOQikjTeXFDC/TOXc9643lx34kC/y/GMglxEksK89eXc9MJ8Du/XibvOOwyz+Jyh0hQFuYgkvJLyKn78RDEFeRk8csl4MlJj/0JYbaGDnSKS0Cpr6vnxE8VU1wV55qojY+5+m15Qj1xEElYo5PjZ8/NYtnkXD100liFxdCGstlCQi0jCuvfdL5mxZAtTzxjBhGEFfpcTMQpyEUlIf5u7kT988BUXHdGXK77Tz+9yIkpBLiIJZ+66Hdzy8gKO7N+ZO84emVAzVJqiIBeRhFJSXsVVT86mR14mj1wynvTUxI85zVoRkYSxp7aeq55smKHy7FVH0ik73e+SokJBLiIJIRRy3PziApZs2sWfLyuKq1u1hSvxf+cQkaTw0PsreWvhJm49fRgnDevudzlRFVaQm9mdZrbAzOaZ2Qwz6+VVYSIirTV94aavr6Ey5fgBfpcTdeH2yO9xzo1yzo0B3gR+EX5JIiKtt6RkFz9/YT5jC/P5zbmJdQ2V1goryJ1zu/ZZzAZceOWIiLTe9ooarnqymI5ZaTx6yXgy0xLrGiqtFfbBTjP7NXApsBOYEHZFIiKtUFsf4rqn5rC9ooaXrjmGgrxMv0vyTYs9cjObaWaLmnicA+Ccm+qc6ws8DfzkINuZYmbFZla8bds271ogIknpl28s5os1Zfzu/FEJcZefcLTYI3fOTWzltp4G3gb+u5ntTAOmARQVFWkIRkTa7a+fr+WZf63jmhMGcs6Y3n6X47twZ60M3mfxHGBZeOWIiBzc56tKueP1xUwY2o2bTxvqdzkxIdwx8t+a2VAgBKwFrgm/JBGRpm0sr+L6p+dQ2KUDD1w0lkBK8s1QaUpYQe6c+zevChEROZiq2iBTniymtj7EHy8tIi8zze+SYoZO0ReRmOec49ZXGk6//9OlRQzsluN3STFFp+iLSMz70yereW1eCTedMoSThyfX6fetoSAXkZj26Yrt3DV9Kacf2oPrJwzyu5yYpCAXkZi1vmwPNzw7h4Hdcvj95NFJefp9ayjIRSQmVdUGufqvs6kPOaZdWkROhg7pNUd/MyISc5xz3P7qQpZubri2eP+u2X6XFNPUIxeRmPPEP9bw6tyN3HjykKS7tnh7KMhFJKbMWlPGr95aysThBdxwkg5utoaCXERixpZd1Vz39Bz6du7AfReMIUVnbraKxshFJCbUBUNc//QcKqrreerKI3XmZhsoyEUkJvzm7aUUr93BgxeNZWiP5Llxshc0tCIivnt9fgl/+WwNl3+nH2eP1q1/20pBLiK+WrFlN7e+vICiQzpx+6ThfpcTlxTkIuKbypp6rnlqNh3SA/zh4nGkBRRJ7aExchHxRcMVDReyenslT/34SLon8T03w6UffyLii6c+X8sb80u46dShHDOwq9/lxDUFuYhE3fz15fzPm0uYMLQb154w0O9y4p6CXESiaueeOq5/Zg4FuZncN1kn/XhBY+QiEjXOOW5+aT6bd1bz4jVH0yk73e+SEoJ65CISNX/5bA0zlmzh1u8OY2xhJ7/LSRgKchGJinnry7lr+lImDu/Olcf297uchOJJkJvZTWbmzEyHnkXkW3ZW1fGTxnHx339/lO7047Gwx8jNrC9wKrAu/HJEJNE457j9lYVsahwXz++gcXGvedEjvx+4BXAebEtEEsxzs9bz1sJN/PupQxmncfGICCvIzewcYKNzbr5H9YhIAlm+ZTd3vLGY4wZ35erjB/hdTsJqcWjFzGYCPZp4aSpwOw3DKi0ysynAFIDCwsI2lCgi8ai6LsgNz8wlJyOVeyeP1nzxCGoxyJ1zE5tab2aHAf2B+Y0HLvoAc8zsCOfc5ia2Mw2YBlBUVKRhGJEEd9fbS/lyy24ev/xwCnJ1HZVIavfBTufcQqBg77KZrQGKnHPbPahLROLY+8u28MQ/13LFd/pz4tCClj8gYdE8chHx1Nbd1dz84gKG9cjlltOH+l1OUvDsFH3nXD+vtiUi8ck5x80vLqCipp7nphxFZlrA75KSgnrkIuKZv36+lo+Wb+M/zxjO4O6672a0KMhFxBNfbavgN28v5YQh3bjkqEP8LiepKMhFJGx1wRA/f34emWkB7jlfp+BHmy5jKyJhe/j9lczfsJP/e/E4CnTLtqhTj1xEwrJgQzkPf7CSc8f2ZtJhPf0uJykpyEWk3Wrqg9z0wny65qTzy7MP9bucpKWhFRFpt/+duYIVWyv4y+WH0zErze9ykpZ65CLSLvPWl/PoR19xQVFfJujsTV8pyEWkzarrgtz0wjx65GUy9czhfpeT9DS0IiJt9tD7K/hqWyVPXnEEeZkaUvGbeuQi0iZLSnbx6EerOH98H44f0s3vcgQFuYi0QX0wxH+8vID8Dun85xkaUokVGloRkVZ77LPVLNy4kz/8YJzuvRlD1CMXkVZZV7qH+95dzsTh3Zl0WFM3DRO/KMhFpEXOOf7rtUUEzLjze4fqWioxRkEuIi2avmgzHy3fxs9PHUrPjll+lyMHUJCLyEFV1NRzxxuLGdEzj8uO1uVpY5EOdorIQd03Yzlbd9fwyCXjSQ2o7xeL9K2ISLOWlOzi8X+s5gdHFDK2sJPf5UgzFOQi0iTnHL98YzEds9K45bRhfpcjB6EgF5Emvb1wM1+sLuOmU4fSsYNOw49lCnIR+Zaq2iC/eXspw3vmcdERhX6XIy0IK8jN7JdmttHM5jU+JnlVmIj459GPv2JjeRW/PGsEgRTNGY91Xsxaud8593sPtiMiMaCkvIpHPvqKM0b15MgBXfwuR1pBQysisp97Zywn5OC27+oAZ7zwIsh/YmYLzOwxM2t2fpKZTTGzYjMr3rZtmwe7FRGvLd20i1fmbuDyY/rRp1MHv8uRVmoxyM1sppktauJxDvD/gIHAGGATcG9z23HOTXPOFTnnirp10zWMRWLR3e8sIy8zjetOHOR3KdIGLY6RO+cmtmZDZvZH4M2wKxIRX/zjq+18+OU2bp80TNMN40y4s1Z67rN4LrAovHJExA+hkOO305fROz+LS4/u53c50kbhzlr5nZmNARywBrg63IJEJPqmL9rMgg07uff7o8lMC/hdjrRRWEHunPuhV4WIiD9CIccD7y1nUEEO3xvb2+9ypB00/VAkyU1ftJnlWyr4PycP1sk/cUpBLpLEQiHHg++tYGC3bM44rGfLH5CYpCAXSWLvLN7Ml1t2qzce5xTkIkkqFHI8MHMFA7plc+aoXn6XI2FQkIskqRlLGnvjJ6k3Hu8U5CJJ6tGPV1HYuQNnjVZvPN4pyEWS0Oy1O5i7rpwrj+2v3ngCUJCLJKE/f7qKvMxUzh/fx+9SxAMKcpEks75sD+8s2szFRx1CdoYXtyQQvynIRZLMY5+tJsWMy3RNlYShIBdJIjur6nhh1nrOHt2LHh0z/S5HPKIgF0kiL8xaT2VtkCuO7e93KeKhuBog+9WbS3h+1vr9V1qTTzGzfZ4f/D2233vs6+f7rrdvrd9/+2bfvGfv6/s+T9nn9RTb/89vXm9Yt+/rgZT9n6c0vjeQYqSkGAEzUvd5Hgg0LAdS9v6ZQmqKkdq4PjXQsJyemkJqSgppgYbnaYGGR8NzIyM1hfRAgIy0FDJSU8hMC+z3577tl/jgnOPZWesoOqQTI3t39Lsc8VBcBfm4QzoRct8sO75ZcK6JD9Dwj/eb9++7/pttfPN83+3ss959s6+G5wesd3s/677ehgNCbu9rjlDom32FnNvvPfuva3hv0DnqQyFqgxAMOZxzBJ0jGGrYTzDUsBxq/DMY3Pu6oz7U+Hrj8/pgaL+/t3CZQWZqgKz0AFlpATqk732kkp0RIDsjleyMVHIzUsnNTCU3M43czFQ6ZqV9/cjvkE5+hzTSAvqlMFpmr93Bqm2VXHP+QL9LEY/FVZBPOqwnk3Rhn3YJ7Q31UIi6YEO41wUddcFQ48NRWx+iNhj6+s+aumDjnyGq64NU14WorgtSUxekau+jNkRVXT17aoNU1tSzsbyOPbX1VNbUs7u6npr60EHrys1IpXNOOl1zMuiSnU7X3AwKcjMoyM2ke14GPTtm0bNjJvkd0vRbQJien7We7PSALo6VgOIqyKX9UlKM9BQjPcqHRWrrQ+yurmNXdT27qurYWVVHeVUdO/fUUlZZx449tZRV1lJaWcPa0j0Ur91BWWXtt7aTlRagT6cs+nbuQGHnDhzSpQP9umYzoGs2fTp10EktLaioqeethZs4e3QvTTlMQPpGJaLSU1PokpNBl5yMVn+mtj7E9ooaNu+qZvPOakrKqygpr2bDjj1s2FHFF6vLqKip//r9GakpDOyWw5DuOQzvmcehvToyolcenbPTI9GkuPTm/BL21Ab5flFfv0uRCFCQS8xJT02hV34WvfKzmnzdOcf2ilrWlFayalsFK7dWsGJrBf9aXcbf5pV8/b7Czh0YW5jP2L75HNG/C8N65JKSpD3354vXM6ggh3GF+X6XIhGgIJe4Y2Z0y82gW24Gh/frvN9rZZW1LCnZxaKSncxbV87nq0p5rTHc8zukcWT/zhw/pBsThhY0+4Mi0azYspu568qZOmm4jjMkKAW5JJTO2ekcO7grxw7u+vW6jeVV/GtVKZ+vKuWzlaX8ffEWAIb1yOXUQ3tw5qieDOme61fJEffi7A2kphjnjtP9OBOVuebm7UVQUVGRKy4ujvp+RZxzfLWtgg+WbWPm0i18saYM52BI9xzOHduH88f3oVtu68fzY51zjhPu+ZAB3bJ5/PIj/C5HwmRms51zRQeuD7tHbmY3ANcDQeAt59wt4W5TJFLMjEEFuQwqyOWq4wewdXc17yzazOvzSrj7nWXcO+NLThnRnUuP7sdRAzrH/VDEss27WVe2h2tP1NzxRBZWkJvZBOAcYLRzrsbMCrwpSyQ6CnIzufToflx6dD9Wbq3g+VnreGn2BqYv2syYvvlcc8JATh3RPW4Pks5YvAUzOHm4/msmsnAnFV8L/NY5VwPgnNsafkki/hhUkMPUM0bwz9tO5lffG0lZZS3XPDWbSQ9+wsfLt/ldXrvMWLKZcYWdKMjVBbISWbhBPgQ4zsz+ZWYfmdnhXhQl4qfMtACXHHUI7990Ag9cOIbK2noufewLLnvsC5Zv2e13ea22YcceFpfs4rRDu/tdikRYi0MrZjYT6NHES1MbP98ZOAo4HHjBzAa4Jo6gmtkUYApAYWFhODWLREVqIIVzxvTm9JE9ePIfa3no/RWc8eAn3DhxCFcfP4DUGL9OzIzG2TmnjGjqv68kkrBmrZjZO8DdzrkPGpe/Ao5yzh3091DNWpF4VFpRwy9eW8xbCzcxuk9H7p08mkEFsTtt8cJp/6SsspYZPzvB71LEI83NWgm3S/E3YELjDoYA6cD2MLcpEpO65GTwh4vH8fAPxrKubA9nPfQZ7yza7HdZTdpRWcsXq8s47VD1xpNBuEH+GDDAzBYBzwGXNTWsIpJIzhzVi3duPJ4hPXK55qnZPPjeCmLtn/17y7YScnCqhlWSQljTD51ztcAlHtUiEje652Xy/JSjuO2Vhdz37nJWbq3gvsmjY2bc/P1lW+jZMZORvfP8LkWiQKfoi7RTZlqA+yaPZlBBDvf8/UsA7r9gjO+X1HXOUbxmB8cM7BL3JzRJ6yjIRcJgZlw/YRApZtz9zjIyUlO4+99G+XoC0cbyKrburmHcIZ18q0GiS0Eu4oFrTxxIVV2QB99bQVZ6gDvOPtS33vCcdeUAjCtUkCcLBbmIR342cTBVtfX88ZPVjOzdkck+3cRhztodZKUFGNYjdqdGirdi48iMSAIwM2797nCOGdiF/35tMSt8Ogt07rodjOrTMWYOvErk6ZsW8VAgxfjfC8bQIT3AT56ZS3VdMKr7r64Lsrhkl8bHk4yCXMRjBXmZ3Dt5NF9u2c0dbyyJ6r4XbtxJfchpfDzJKMhFIuDEoQVcfcIAnv1iHV+sLovafues3QHAWN2bM6koyEUi5MaTh1CQm8Hd7yyL2pmfc9eVU9i5A11zEucuR9IyBblIhGSlB/jpxMHMXruDmUsjf6l+5xxz1u1gnHrjSUdBLhJBk4v60r9rNvf8fRnBUGR75ToRKHkpyEUiKC2Qwr+fOpTlWyp4de7GiO5LJwIlLwW5SIRNOqwHh/XuyP3vLqemPnLTEXUiUPJSkItEmJnxs1MGs7G8ig+WRe7en0s27WJErzydCJSE9I2LRMHxg7vRJTudNxeURGwf2ytq6JGnmywnIwW5SBSkBlL47mE9eG/pVvbU1kdkH6UVtXTJSY/ItiW2KchFouSsUb2oqgtGZCpiXTDEzqo6OmcryJORglwkSg7v15nueRm8Od/74ZUdlbUAdFGQJyUFuUiUpKQYZxzWiw+/3Mau6jpPt126N8h1RmdSUpCLRNFZo3tSGwwxY/EWT7dbWtEQ5BpaSU4KcpEoGtM3nz6dsjyfvVJaWQNAVx3sTEphBbmZPW9m8xofa8xsnkd1iSQkM+PMUb34dMX2r8e1vfBNj1xDK8korCB3zl3gnBvjnBsDvAy84klVIgnsrNE9ye+QxlfbKjzbZlllLSkG+Vlpnm1T4ocn9+y0hrvMTgZO8mJ7IolsRM88Pr/tZE/PwCytrKVzdjopKf7c8Fn85dXNl48DtjjnVni0PZGEZWakBrwN3NKKGrpoWCVptRjkZjYT6NHES1Odc681Pr8IeLaF7UwBpgAUFha2sUwROZiyxh65JKcWg9w5N/Fgr5tZKnAeML6F7UwDpgEUFRVF53YpIkmitLKWEb3y/C5DfOLFIN1EYJlzboMH2xKRdiitqKGreuRJy4sgv5AWhlVEJHJq60Psqq7X1MMkFvbBTufcjzyoQ0Taaceevafnq0eerHRmp0ic23sykC6YlbwU5CJxbu/p+bpgVvJSkIvEubJKXTAr2SnIReLc9sahFV0wK3kpyEXiXFllDYEUIy9T11lJVgpykThXWlFLpw66zkoyU5CLxLnSyloNqyQ5BblInCutqNGBziSnIBeJc2WVtZp6mOQU5CJxrrSyVicDJTkFuUgcq6kPsru6XkGe5BTkInFsR2UdAJ11sDOpKchF4tj2isbT83Xlw6SmIBeJY3tPz9eVD5Obglwkju29YJamHyY3BblIHNt7CduuGlpJagpykThWWllLaoqRlxX2PWIkjinIReJYWUUtnbPTMdN1VpKZglwkjpVW1mp8XBTkIvGstLKGrjo9P+lpYE0kjuVlptG3c5bfZYjPwgpyMxsDPAJkAvXAdc65LzyoS0Ra4YkrjvC7BIkB4Q6t/A64wzk3BvhF47KIiERRuEHugLzG5x2BkjC3JyIibRTuGPmNwN/N7Pc0/FA4JuyKRESkTVoMcjObCfRo4qWpwMnAz5xzL5vZZODPwMRmtjMFmAJQWFjY7oJFRGR/5pxr/4fNdgL5zjlnDWck7HTO5bX0uaKiIldcXNzu/YqIJCMzm+2cKzpwfbhj5CXACY3PTwJWhLk9ERFpo3DHyK8CHjCzVKCaxqETERGJnrCC3Dn3KTDeo1pERKQdwhojb/dOzbYBa2mYsrhzn5f2XW7ueVdguwdlHLjvcN7b3OtNrU+2NrdlOR7b3Nbv+MDlWG6zV/+uD1xWm9vf5kOcc92+tdY559sDmNbc8kGeF0di3+G8t7nXm1qfbG1uy3I8trmt33E8tdmrf9dqc2TavO/D74tmvXGQ5eaeR2rf4by3udebWp9sbW7Lcjy2ua3f8YHLsdxmr/5dH7isNnvMl6GVcJhZsWti+k0iU5uTg9qcHCLRZr975O0xze8CfKA2Jwe1OTl43ua465GLiMj+4rFHLiIi+1CQi4jEOQW5iEicS6ggN7MUM/u1mT1kZpf5XU80mNmJZvaJmT1iZif6XU+0mFm2mRWb2Zl+1xJpZja88ft9ycyu9bueaDCz75nZH83seTM71e96osHMBpjZn83spbZ+NmaC3MweM7OtZrbogPWnm9mXZrbSzG5tYTPnAH2AOmBDpGr1ikdtdkAFDbfbS5Y2A/wH8EJkqvSOF+11zi11zl0DTAa+E8l6veBRm//mnLsKuAa4IJL1esGjNq9yzl3Zrv3HyqwVMzuehkB60jk3snFdAFgOnEJDSM0CLgICwF0HbOKKxscO59yjZvaSc+78aNXfHh61ebtzLmRm3YH7nHMXR6v+9vCozaOBLjT88NrunHszOtW3nRftdc5tNbOzgWuBvzrnnolW/e3hVZsbP3cv8LRzbk6Uym8Xj9vc5uwK9+qHnnHOfWxm/Q5YfQSw0jm3CsDMngPOcc7dBXzrV2oz2wDUNi4GI1iuJ7xo8z52ABkRKdRDHn3PJwLZwAigyszeds6FIll3e3n1HTvnXgdeN7O3gJgOco++YwN+C0yP9RAHz/8vt1nMBHkzegPr91neABx5kPe/AjxkZscBH0eysAhqU5vN7DzgNCAfeDiilUVOm9rsnJsKYGY/ovE3kohW5722fscnAufR8IP67UgWFkFt/b98Aw13G+toZoOcc49EsrgIaev33AX4NTDWzG5rDPxWifUgbxPn3B6gXWNM8co59woNP8CSjnPucb9riAbn3IfAhz6XEVXOuQeBB/2uI5qcc6U0HBNos5g52NmMjUDffZb7NK5LZGpz4rc52doLajNEsM2xHuSzgMFm1t/M0oELgdd9rinS1ObEb3OytRfU5si22evr4rb3ATwLbOKbqYNXNq6fRMOR36+AqX7XqTarzWqv2hxrbY6Z6YciItI+sT60IiIiLVCQi4jEOQW5iEicU5CLiMQ5BbmISJxTkIuIxDkFuYhInFOQi4jEOQW5iEic+/8GavcMEGQiwwAAAABJRU5ErkJggg==",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "gamma2_vals = jnp.logspace(-6,-1,100)\n",
    "j_mar = jax.jit(jax.vmap(marginal_like))\n",
    "vals = j_mar(gamma2_vals)\n",
    "plt.plot(gamma2_vals[vals[:,1]>0],vals[:,0][vals[:,1]>0])\n",
    "plt.xscale('log')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 71,
   "metadata": {},
   "outputs": [],
   "source": [
    "gamma2_opt = gamma2_vals[jnp.argmin(vals[:,0][vals[:,1]>0])]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 72,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array(0.00095455, dtype=float64)"
      ]
     },
     "execution_count": 72,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "gamma2_opt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 73,
   "metadata": {},
   "outputs": [],
   "source": [
    "n_groups = 3\n",
    "I =jnp.eye(n_groups)\n",
    "oneoneT = jnp.ones((n_groups,n_groups))\n",
    "D = jnp.diag(v)\n",
    "Cml = D - gamma2_opt*I + gamma2_opt*oneoneT"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 74,
   "metadata": {},
   "outputs": [],
   "source": [
    "#(D-true_C)/true_C"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 75,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array([[0.00087373, 0.00095455, 0.00095455],\n",
       "       [0.00095455, 0.00099766, 0.00095455],\n",
       "       [0.00095455, 0.00095455, 0.00123755]], dtype=float64)"
      ]
     },
     "execution_count": 75,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "Cml"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 76,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array([[0.    , 1.001 , 0.855 ],\n",
       "       [1.001 , 0.    , 0.5028],\n",
       "       [0.855 , 0.5028, 0.    ]], dtype=float64)"
      ]
     },
     "execution_count": 76,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.around((Cml - true_C)/true_C,4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 54,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "array([[ 0.    , -0.016 ,  0.0135],\n",
       "       [-0.016 ,  0.    , -0.1173],\n",
       "       [ 0.0135, -0.1173,  0.    ]])"
      ]
     },
     "execution_count": 54,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "np.around((gl_C - true_C)/true_C,4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 63,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array(-0.0066732, dtype=float64)"
      ]
     },
     "execution_count": 63,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x.T@jnp.linalg.inv(gl_C)@L/(x.T@jnp.linalg.inv(gl_C)@x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 64,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array(-0.00453136, dtype=float64)"
      ]
     },
     "execution_count": 64,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x.T@jnp.linalg.inv(D)@L/(x.T@jnp.linalg.inv(D)@x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 62,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Array(-0.00748569, dtype=float64)"
      ]
     },
     "execution_count": 62,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x.T@jnp.linalg.inv(Cml)@L/(x.T@jnp.linalg.inv(Cml)@x)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "base",
   "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.9.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}