playground.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 52,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from scipy.special import logit, expit\n",
    "from scipy.special import digamma, polygamma\n",
    "from scipy.optimize import fsolve\n",
    "\n",
    "from scipy import stats\n",
    "\n",
    "def find_logbeta_params(mean_logprob, var_logprob):\n",
    "    def equations(vars):\n",
    "        alpha, beta = vars\n",
    "        # Mean equation\n",
    "        eq1 = digamma(alpha) - digamma(alpha + beta) - mean_logprob\n",
    "        # Variance equation\n",
    "        eq2 = polygamma(1, alpha) - polygamma(1, alpha + beta) - var_logprob\n",
    "        return [eq1, eq2]\n",
    "    \n",
    "    # Initial guess for α,β\n",
    "    alpha0, beta0 = 1.0, 1.0\n",
    "    alpha, beta = fsolve(equations, [alpha0, beta0])\n",
    "    \n",
    "    return alpha, beta\n",
    "\n",
    "def logbeta_logprobs(mean_logprob, var_logprob, n_samples=10000):\n",
    "    alpha, beta = find_logbeta_params(mean_logprob, var_logprob)\n",
    "    probs = stats.beta.rvs(alpha, beta, size=n_samples)\n",
    "    logprobs = np.log(probs)\n",
    "    \n",
    "    print(f\"Fitted alpha, beta: {alpha:.3f}, {beta:.3f}\")\n",
    "    print(f\"Target mean, var: {mean_logprob:.3f}, {var_logprob:.3f}\")\n",
    "    print(f\"Achieved mean, var: {np.mean(logprobs):.3f}, {np.var(logprobs):.3f}\")\n",
    "    \n",
    "    return logprobs\n",
    "def compute_p2(p1, var_logprob, n, num_samples=10000):\n",
    "    # Validate p1 is between 0 and 1\n",
    "    if not 0 <= p1 <= 1:\n",
    "        raise ValueError(\"p1 must be between 0 and 1\")\n",
    "    \n",
    "\n",
    "    p11 = logbeta_logprobs(p1, var_logprob, num_samples)\n",
    "\n",
    "    \n",
    "    # Compute p11 and p12\n",
    "    p11 = expit(logit(p1) + epsilon1)\n",
    "    p12 = expit(logit(p1) + epsilon2)\n",
    "    \n",
    "    # Compute p2\n",
    "    p2 = p11 * p12 + (1 - p11) / n\n",
    "    \n",
    "    return p2, p11, p12"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def old_loss_two_hop(p1, n=100):\n",
    "    return p1**2 + ((1-p1)**2) / (n-1)\n",
    "\n",
    "\n",
    "def better_loss_two_hop(p1, n=100):\n",
    "    return p1**2 + (1-p1) / n\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 55,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(0.003994149414941494,\n",
       " 0.004,\n",
       " 0.005664439517671612,\n",
       " 0.0748176954313564,\n",
       " 0.07419483546284834)"
      ]
     },
     "execution_count": 55,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "p_1 = 0.0625\n",
    "\n",
    "old_loss_two_hop(p_1, 10000), better_loss_two_hop(p_1, 10000), compute_p2(p_1, np.abs(logit(p_1))/4, 10000)[0].mean(), compute_p2(p_1, np.abs(logit(p_1))/4, 10000)[1].mean(), compute_p2(p_1, np.abs(logit(p_1))/4, 10000)[2].mean()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 38,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Estimated p1: 0.3640, sqrt p2: 0.3873\n"
     ]
    }
   ],
   "source": [
    "from scipy import optimize\n",
    "\n",
    "def find_p1(target_p2, sigma, n, num_samples=10000):\n",
    "    def objective(p1):\n",
    "        np.random.seed(42)  # Fix seed for consistent optimization\n",
    "        p2_samples = compute_p2(p1[0], sigma, n, num_samples)\n",
    "        return np.abs(np.mean(p2_samples) - target_p2)\n",
    "    \n",
    "    # Try multiple starting points\n",
    "    best_result = None\n",
    "    best_score = float('inf')\n",
    "    \n",
    "    for start in [0.1, 0.3, 0.5, 0.7, 0.9]:\n",
    "        result = optimize.minimize(\n",
    "            objective, \n",
    "            x0=[start],\n",
    "            bounds=[(0, 1)],\n",
    "            method='L-BFGS-B',\n",
    "            options={'ftol': 1e-8}\n",
    "        )\n",
    "        \n",
    "        if result.fun < best_score:\n",
    "            best_score = result.fun\n",
    "            best_result = result\n",
    "    \n",
    "    return best_result.x[0]\n",
    "\n",
    "# Example usage\n",
    "target_p2 = 0.15\n",
    "sigma = 1.0\n",
    "n = 10000\n",
    "estimated_p1 = find_p1(target_p2, sigma, n)\n",
    "print(f\"Estimated p1: {estimated_p1:.4f}, sqrt p2: {np.sqrt(target_p2):.4f}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "-1.0986122886681098"
      ]
     },
     "execution_count": 41,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "logit(0.25)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "For mean_logprob = -0.010\n",
      "Min variance: 0.000\n",
      "Max variance: 1.739\n",
      "For mean_logprob = -0.500\n",
      "Min variance: 0.000\n",
      "Max variance: 86.706\n",
      "For mean_logprob = -1.000\n",
      "Min variance: 0.000\n",
      "Max variance: 172.913\n",
      "For mean_logprob = -2.000\n",
      "Min variance: 0.000\n",
      "Max variance: 343.826\n",
      "For mean_logprob = -3.000\n",
      "Min variance: 0.000\n",
      "Max variance: 334.176\n",
      "For mean_logprob = -4.000\n",
      "Min variance: 0.000\n",
      "Max variance: 441.569\n"
     ]
    }
   ],
   "source": [
    "def explore_logbeta_bounds(mean_logprob):\n",
    "    # Try a range of alpha values\n",
    "    alphas = np.logspace(-2, 4, 100)\n",
    "    vars = []\n",
    "    for alpha in alphas:\n",
    "        # For each alpha, find beta that gives our target mean\n",
    "        def mean_eq(beta):\n",
    "            return digamma(alpha) - digamma(alpha + beta) - mean_logprob\n",
    "        \n",
    "        try:\n",
    "            beta = fsolve(mean_eq, [1.0])[0]\n",
    "            if beta > 0:  # Valid beta only\n",
    "                var = polygamma(1, alpha) - polygamma(1, alpha + beta)\n",
    "                vars.append(var)\n",
    "        except:\n",
    "            continue\n",
    "    \n",
    "    if vars:\n",
    "        print(f\"For mean_logprob = {mean_logprob:.3f}\")\n",
    "        print(f\"Min variance: {min(vars):.3f}\")\n",
    "        print(f\"Max variance: {max(vars):.3f}\")\n",
    "        return min(vars), max(vars)\n",
    "    return None\n",
    "\n",
    "# Test some means\n",
    "for mean in [-0.01,-0.5, -1.0, -2.0, -3.0, -4.0]:\n",
    "    explore_logbeta_bounds(mean)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.11.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}