calculate_wsd.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "12d523b2",
   "metadata": {},
   "outputs": [],
   "source": [
    "from scipy.stats import wasserstein_distance as wsd\n",
    "from scipy.stats import beta as beta_distribution\n",
    "\n",
    "from ddsm import noise_factory\n",
    "import torch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "333f0bcc",
   "metadata": {},
   "outputs": [],
   "source": [
    "num_cat = 9\n",
    "num_samples = 500000"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b33722fc",
   "metadata": {},
   "outputs": [],
   "source": [
    "def noise_factory(N, n_time_steps, a, b, total_time=4, order=100,\n",
    "                  time_steps=1000, speed_balanced=True, logspace=False,\n",
    "                  mode=\"independent\", \n",
    "                  #mode=\"path\", \n",
    "                  device=\"cuda\",\n",
    "                  noise_only=False,\n",
    "                  boundary_mode='clamp')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "c89acd76",
   "metadata": {},
   "outputs": [],
   "source": [
    "total_time = 1.0\n",
    "order = 100\n",
    "steps_per_tick = 200\n",
    "speed_balance = True #if True s=2/(a+b) is used; otherwise s=1\n",
    "logspace = False\n",
    "mode ='independent'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ad571de5",
   "metadata": {},
   "outputs": [],
   "source": [
    "device=\"cpu\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "28006c49",
   "metadata": {},
   "outputs": [],
   "source": [
    "alpha = torch.ones(1)\n",
    "beta =  (num_cat-1)*torch.ones(1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "c4964555",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([1])\n",
      "torch.Size([1])\n"
     ]
    }
   ],
   "source": [
    "print(alpha.shape)\n",
    "print(beta.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "a933c150",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "clamp with steps=100: wsd=0.007017955881073237\n",
      "clamp with steps=100: wsd=0.007150544142611445\n",
      "clamp with steps=100: wsd=0.0068189085119101725\n",
      "clamp with steps=100: wsd=0.007123989489040336\n",
      "clamp with steps=100: wsd=0.006950329654692467\n",
      "\n",
      "clamp with steps=100: mean=0.00701234582811594 and speed_balance=True\n",
      "clamp with steps=100: std=0.0001349616068182513 and speed_balance=True\n",
      "\n",
      "reflect_boundaries with steps=100: wsd=0.006987178854938313\n",
      "reflect_boundaries with steps=100: wsd=0.006732128456849396\n",
      "reflect_boundaries with steps=100: wsd=0.006797766843601802\n",
      "reflect_boundaries with steps=100: wsd=0.0070380988440061605\n",
      "reflect_boundaries with steps=100: wsd=0.006947154300943437\n",
      "\n",
      "reflect_boundaries with steps=100: mean=0.006900465581566095 and speed_balance=True\n",
      "reflect_boundaries with steps=100: std=0.00012992908887099475 and speed_balance=True\n",
      "\n",
      "reflect with steps=100: wsd=0.007094939547602649\n",
      "reflect with steps=100: wsd=0.0068946464727944815\n",
      "reflect with steps=100: wsd=0.006971135644897272\n",
      "reflect with steps=100: wsd=0.007211670349788356\n",
      "reflect with steps=100: wsd=0.007185494115574644\n",
      "\n",
      "reflect with steps=100: mean=0.00707157701253891 and speed_balance=True\n",
      "reflect with steps=100: std=0.0001365117495879531 and speed_balance=True\n",
      "\n",
      "clamp with steps=400: wsd=0.006973087794559582\n",
      "clamp with steps=400: wsd=0.006801493289006876\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n",
      "KeyboardInterrupt\n",
      "\n"
     ]
    }
   ],
   "source": [
    "iteration = 5\n",
    "for num_time_steps in [100,400,800,1200,1600,2000]:\n",
    "    for boundary_mode in ['clamp', 'reflect_boundaries', 'reflect']:\n",
    "        distances = torch.zeros(iteration)\n",
    "        for i in range(iteration):\n",
    "            v_one, v_zero = noise_factory(num_samples, \n",
    "                                            num_time_steps,\n",
    "                                            alpha,\n",
    "                                            beta,\n",
    "                                            total_time=total_time,\n",
    "                                            order=order,\n",
    "                                            time_steps=steps_per_tick,\n",
    "                                            logspace=logspace,\n",
    "                                            speed_balanced=speed_balance,\n",
    "                                            mode=mode,\n",
    "                                            device=device,\n",
    "                                            boundary_mode=boundary_mode,\n",
    "                                            noise_only=True)\n",
    "\n",
    "            vT_approx = torch.cat([v_one[:,-1,0],v_zero[:,-1,0]])\n",
    "            vT_true = beta_distribution.rvs(alpha, beta, size=vT_approx.shape[0])\n",
    "            d = wsd(vT_approx,vT_true)\n",
    "            print(f'{boundary_mode} with steps={num_time_steps}: wsd={d}')\n",
    "            distances[i] = d\n",
    "\n",
    "        #print('v shape one', v_one.shape)\n",
    "        #print('v shape zero', v_zero.shape)\n",
    "\n",
    "        print(f'\\n{boundary_mode} with steps={num_time_steps}: mean={torch.mean(distances)} and speed_balance={speed_balance}')\n",
    "        print(f'{boundary_mode} with steps={num_time_steps}: std={torch.std(distances)} and speed_balance={speed_balance}\\n')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6e7b250d",
   "metadata": {},
   "source": [
    "## Ground truth"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 74,
   "id": "ea70f234",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "test1 (1000000,)\n",
      "test2 (1000000,)\n",
      "distance 0.00015255260316868874\n",
      "test1 (1000000,)\n",
      "test2 (1000000,)\n",
      "distance 0.00021186835176490015\n",
      "test1 (1000000,)\n",
      "test2 (1000000,)\n",
      "distance 0.00011741920541712615\n",
      "test1 (1000000,)\n",
      "test2 (1000000,)\n",
      "distance 0.00019672195663265597\n",
      "test1 (1000000,)\n",
      "test2 (1000000,)\n",
      "distance 0.00014337657853888966\n",
      "test1 (1000000,)\n",
      "test2 (1000000,)\n",
      "distance 0.00022759701250514168\n",
      "test1 (1000000,)\n",
      "test2 (1000000,)\n",
      "distance 0.00029576271299861255\n",
      "test1 (1000000,)\n",
      "test2 (1000000,)\n",
      "distance 0.00010474523625257931\n",
      "test1 (1000000,)\n",
      "test2 (1000000,)\n",
      "distance 0.0001795923039257112\n",
      "test1 (1000000,)\n",
      "test2 (1000000,)\n",
      "distance 0.00011724995625918759\n"
     ]
    }
   ],
   "source": [
    "iteration = 10\n",
    "distances = torch.zeros(iteration)\n",
    "for i in range(iteration):\n",
    "    test1 = beta_distribution.rvs(alpha, beta, size=1000000)\n",
    "    test2 = beta_distribution.rvs(alpha, beta, size=1000000)\n",
    "    d = wsd(test1,test2)\n",
    "    print('distance',d)\n",
    "    distances[i] = d"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 75,
   "id": "2c7e6e99",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "mean tensor(0.0002)\n",
      "std tensor(5.9900e-05)\n"
     ]
    }
   ],
   "source": [
    "print('mean',torch.mean(distances))\n",
    "print('std',torch.std(distances))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7bb43f95",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "gfdm",
   "language": "python",
   "name": "gfdm"
  },
  "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.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}