find_default_model.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This notebook performs cross-validation to identify the best parameters for a KNN model to predict the genetic ancestry of the 1000 Genomes individuals."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from pathlib import Path\n",
    "\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from sklearn.model_selection import train_test_split\n",
    "import joblib\n",
    "\n",
    "# load config variables\n",
    "from ezancestry.config import aisnps_set as _aisnps_set\n",
    "from ezancestry.config import models_directory as _models_directory\n",
    "from ezancestry.config import population_level as _population_level\n",
    "\n",
    "# load functions\n",
    "from ezancestry.fetch import get_thousand_genomes_aisnps"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Requirement already satisfied: optuna in ./.venv/lib/python3.11/site-packages (3.5.0)\n",
      "Requirement already satisfied: alembic>=1.5.0 in ./.venv/lib/python3.11/site-packages (from optuna) (1.13.1)\n",
      "Requirement already satisfied: colorlog in ./.venv/lib/python3.11/site-packages (from optuna) (6.8.2)\n",
      "Requirement already satisfied: numpy in ./.venv/lib/python3.11/site-packages (from optuna) (1.24.4)\n",
      "Requirement already satisfied: packaging>=20.0 in ./.venv/lib/python3.11/site-packages (from optuna) (23.2)\n",
      "Requirement already satisfied: sqlalchemy>=1.3.0 in ./.venv/lib/python3.11/site-packages (from optuna) (2.0.27)\n",
      "Requirement already satisfied: tqdm in ./.venv/lib/python3.11/site-packages (from optuna) (4.66.2)\n",
      "Requirement already satisfied: PyYAML in ./.venv/lib/python3.11/site-packages (from optuna) (6.0.1)\n",
      "Requirement already satisfied: Mako in ./.venv/lib/python3.11/site-packages (from alembic>=1.5.0->optuna) (1.3.2)\n",
      "Requirement already satisfied: typing-extensions>=4 in ./.venv/lib/python3.11/site-packages (from alembic>=1.5.0->optuna) (4.9.0)\n",
      "Requirement already satisfied: greenlet!=0.4.17 in ./.venv/lib/python3.11/site-packages (from sqlalchemy>=1.3.0->optuna) (3.0.3)\n",
      "Requirement already satisfied: MarkupSafe>=0.9.2 in ./.venv/lib/python3.11/site-packages (from Mako->alembic>=1.5.0->optuna) (2.1.5)\n",
      "\n",
      "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m A new release of pip is available: \u001b[0m\u001b[31;49m23.3.1\u001b[0m\u001b[39;49m -> \u001b[0m\u001b[32;49m24.0\u001b[0m\n",
      "\u001b[1m[\u001b[0m\u001b[34;49mnotice\u001b[0m\u001b[1;39;49m]\u001b[0m\u001b[39;49m To update, run: \u001b[0m\u001b[32;49mpip install --upgrade pip\u001b[0m\n"
     ]
    }
   ],
   "source": [
    "!pip install optuna"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/workspaces/ezancestry/.venv/lib/python3.11/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n"
     ]
    }
   ],
   "source": [
    "import optuna\n",
    "from sklearn.decomposition import PCA\n",
    "from sklearn.metrics import accuracy_score, balanced_accuracy_score, top_k_accuracy_score, make_scorer\n",
    "from sklearn.model_selection import StratifiedKFold, cross_val_score, train_test_split\n",
    "from sklearn.neighbors import KNeighborsClassifier\n",
    "from sklearn.pipeline import make_pipeline\n",
    "from sklearn.preprocessing import OneHotEncoder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2024-02-21 03:04:26.051 | INFO     | ezancestry.fetch:get_thousand_genomes_aisnps:34 - Loaded: kidd.1kG.csv\n",
      "2024-02-21 03:04:26.085 | INFO     | ezancestry.fetch:get_thousand_genomes_aisnps:34 - Loaded: seldin.1kG.csv\n"
     ]
    }
   ],
   "source": [
    "# If the file exists in the aisnps_directory, you can load it rather than querying the 1000 Genomes VCF\n",
    "kidd1kg = get_thousand_genomes_aisnps(aisnps_directory=\"data/aisnps/\", aisnps_sets=\"kidd\").set_index(\"sample\")\n",
    "seldin1kg = get_thousand_genomes_aisnps(aisnps_directory=\"data/aisnps/\", aisnps_sets=\"seldin\").set_index(\"sample\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Kidd AISNPs to predict superpopulation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "todrop = [\"superpopulation\", \"population\", \"gender\"]\n",
    "label = \"superpopulation\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [],
   "source": [
    "X_train, X_test, y_train, y_test = train_test_split(\n",
    "    kidd1kg.drop(columns=todrop),\n",
    "    kidd1kg[label],\n",
    "    test_size=0.2,\n",
    "    random_state=42,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {},
   "outputs": [],
   "source": [
    "def objective(trial):\n",
    "    weights = trial.suggest_categorical(\"weights\", [\"uniform\", \"distance\"])\n",
    "    n_components = trial.suggest_int(\"n_components\", 3, 50)\n",
    "    n_neighbors = trial.suggest_int(\"n_neighbors\", 11, 100)\n",
    "\n",
    "    model = make_pipeline(\n",
    "        OneHotEncoder(handle_unknown=\"ignore\", sparse_output=False, dtype=np.int8),\n",
    "        PCA(n_components=n_components),\n",
    "        KNeighborsClassifier(n_neighbors=n_neighbors, weights=weights, n_jobs=4),\n",
    "    )\n",
    "\n",
    "    return cross_val_score(model, X_train, y_train, n_jobs=4, scoring=\"accuracy\", cv=StratifiedKFold(5)).mean()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[I 2024-02-21 03:04:26,208] A new study created in memory with name: no-name-923b4518-eb9a-46b5-9210-0191a42fdc05\n",
      "[I 2024-02-21 03:04:28,046] Trial 1 finished with value: 0.9590723192019951 and parameters: {'weights': 'uniform', 'n_components': 8, 'n_neighbors': 44}. Best is trial 1 with value: 0.9590723192019951.\n",
      "[I 2024-02-21 03:04:28,240] Trial 0 finished with value: 0.908149625935162 and parameters: {'weights': 'uniform', 'n_components': 50, 'n_neighbors': 45}. Best is trial 1 with value: 0.9590723192019951.\n",
      "[I 2024-02-21 03:04:28,390] Trial 2 finished with value: 0.9151408977556109 and parameters: {'weights': 'uniform', 'n_components': 38, 'n_neighbors': 68}. Best is trial 1 with value: 0.9590723192019951.\n",
      "[I 2024-02-21 03:04:28,580] Trial 3 finished with value: 0.9645598503740649 and parameters: {'weights': 'distance', 'n_components': 4, 'n_neighbors': 94}. Best is trial 3 with value: 0.9645598503740649.\n",
      "[I 2024-02-21 03:04:28,676] Trial 4 finished with value: 0.9366059850374064 and parameters: {'weights': 'distance', 'n_components': 26, 'n_neighbors': 31}. Best is trial 3 with value: 0.9645598503740649.\n",
      "[I 2024-02-21 03:04:28,872] Trial 5 finished with value: 0.96206608478803 and parameters: {'weights': 'uniform', 'n_components': 12, 'n_neighbors': 18}. Best is trial 3 with value: 0.9645598503740649.\n",
      "[I 2024-02-21 03:04:29,037] Trial 6 finished with value: 0.959071072319202 and parameters: {'weights': 'distance', 'n_components': 12, 'n_neighbors': 28}. Best is trial 3 with value: 0.9645598503740649.\n",
      "[I 2024-02-21 03:04:29,221] Trial 7 finished with value: 0.958071072319202 and parameters: {'weights': 'distance', 'n_components': 12, 'n_neighbors': 32}. Best is trial 3 with value: 0.9645598503740649.\n",
      "[I 2024-02-21 03:04:29,338] Trial 8 finished with value: 0.9286184538653366 and parameters: {'weights': 'distance', 'n_components': 24, 'n_neighbors': 75}. Best is trial 3 with value: 0.9645598503740649.\n",
      "[I 2024-02-21 03:04:29,491] Trial 9 finished with value: 0.9306172069825436 and parameters: {'weights': 'distance', 'n_components': 26, 'n_neighbors': 65}. Best is trial 3 with value: 0.9645598503740649.\n",
      "[I 2024-02-21 03:04:29,612] Trial 10 finished with value: 0.9166359102244389 and parameters: {'weights': 'distance', 'n_components': 35, 'n_neighbors': 65}. Best is trial 3 with value: 0.9645598503740649.\n",
      "[I 2024-02-21 03:04:29,768] Trial 11 finished with value: 0.9650598503740648 and parameters: {'weights': 'distance', 'n_components': 4, 'n_neighbors': 98}. Best is trial 11 with value: 0.9650598503740648.\n",
      "[I 2024-02-21 03:04:29,921] Trial 12 finished with value: 0.9021620947630924 and parameters: {'weights': 'uniform', 'n_components': 3, 'n_neighbors': 100}. Best is trial 11 with value: 0.9650598503740648.\n",
      "[I 2024-02-21 03:04:30,044] Trial 13 finished with value: 0.9650598503740648 and parameters: {'weights': 'distance', 'n_components': 4, 'n_neighbors': 98}. Best is trial 11 with value: 0.9650598503740648.\n",
      "[I 2024-02-21 03:04:30,230] Trial 14 finished with value: 0.910149625935162 and parameters: {'weights': 'distance', 'n_components': 3, 'n_neighbors': 95}. Best is trial 11 with value: 0.9650598503740648.\n",
      "[I 2024-02-21 03:04:30,351] Trial 15 finished with value: 0.9361072319201995 and parameters: {'weights': 'distance', 'n_components': 18, 'n_neighbors': 86}. Best is trial 11 with value: 0.9650598503740648.\n",
      "[I 2024-02-21 03:04:30,576] Trial 16 finished with value: 0.9361072319201995 and parameters: {'weights': 'distance', 'n_components': 18, 'n_neighbors': 81}. Best is trial 11 with value: 0.9650598503740648.\n",
      "[I 2024-02-21 03:04:30,727] Trial 17 finished with value: 0.9371047381546134 and parameters: {'weights': 'distance', 'n_components': 18, 'n_neighbors': 82}. Best is trial 11 with value: 0.9650598503740648.\n",
      "[I 2024-02-21 03:04:30,902] Trial 18 finished with value: 0.9356097256857854 and parameters: {'weights': 'distance', 'n_components': 18, 'n_neighbors': 86}. Best is trial 11 with value: 0.9650598503740648.\n",
      "[I 2024-02-21 03:04:31,041] Trial 19 finished with value: 0.913639650872818 and parameters: {'weights': 'distance', 'n_components': 33, 'n_neighbors': 89}. Best is trial 11 with value: 0.9650598503740648.\n",
      "[I 2024-02-21 03:04:31,205] Trial 20 finished with value: 0.9251259351620948 and parameters: {'weights': 'distance', 'n_components': 33, 'n_neighbors': 52}. Best is trial 11 with value: 0.9650598503740648.\n",
      "[I 2024-02-21 03:04:31,423] Trial 21 finished with value: 0.9505847880299252 and parameters: {'weights': 'distance', 'n_components': 9, 'n_neighbors': 100}. Best is trial 11 with value: 0.9650598503740648.\n",
      "[I 2024-02-21 03:04:31,606] Trial 22 finished with value: 0.956074812967581 and parameters: {'weights': 'distance', 'n_components': 7, 'n_neighbors': 99}. Best is trial 11 with value: 0.9650598503740648.\n",
      "[I 2024-02-21 03:04:31,770] Trial 23 finished with value: 0.9560760598503741 and parameters: {'weights': 'distance', 'n_components': 7, 'n_neighbors': 95}. Best is trial 11 with value: 0.9650598503740648.\n",
      "[I 2024-02-21 03:04:31,893] Trial 24 finished with value: 0.9645598503740649 and parameters: {'weights': 'distance', 'n_components': 4, 'n_neighbors': 90}. Best is trial 11 with value: 0.9650598503740648.\n",
      "[I 2024-02-21 03:04:32,180] Trial 25 finished with value: 0.912147132169576 and parameters: {'weights': 'distance', 'n_components': 3, 'n_neighbors': 76}. Best is trial 11 with value: 0.9650598503740648.\n",
      "[I 2024-02-21 03:04:32,328] Trial 26 finished with value: 0.9440947630922694 and parameters: {'weights': 'distance', 'n_components': 14, 'n_neighbors': 76}. Best is trial 11 with value: 0.9650598503740648.\n",
      "[I 2024-02-21 03:04:32,494] Trial 27 finished with value: 0.9406009975062345 and parameters: {'weights': 'uniform', 'n_components': 15, 'n_neighbors': 73}. Best is trial 11 with value: 0.9650598503740648.\n",
      "[I 2024-02-21 03:04:32,645] Trial 28 finished with value: 0.9021583541147132 and parameters: {'weights': 'uniform', 'n_components': 46, 'n_neighbors': 92}. Best is trial 11 with value: 0.9650598503740648.\n",
      "[I 2024-02-21 03:04:32,835] Trial 29 finished with value: 0.9306159600997506 and parameters: {'weights': 'distance', 'n_components': 22, 'n_neighbors': 91}. Best is trial 11 with value: 0.9650598503740648.\n",
      "[I 2024-02-21 03:04:32,963] Trial 30 finished with value: 0.9356072319201993 and parameters: {'weights': 'distance', 'n_components': 21, 'n_neighbors': 56}. Best is trial 11 with value: 0.9650598503740648.\n",
      "[I 2024-02-21 03:04:33,141] Trial 31 finished with value: 0.9545798004987531 and parameters: {'weights': 'distance', 'n_components': 9, 'n_neighbors': 82}. Best is trial 11 with value: 0.9650598503740648.\n",
      "[I 2024-02-21 03:04:33,258] Trial 32 finished with value: 0.959073566084788 and parameters: {'weights': 'distance', 'n_components': 6, 'n_neighbors': 82}. Best is trial 11 with value: 0.9650598503740648.\n",
      "[I 2024-02-21 03:04:33,451] Trial 33 finished with value: 0.9580748129675809 and parameters: {'weights': 'distance', 'n_components': 6, 'n_neighbors': 94}. Best is trial 11 with value: 0.9650598503740648.\n",
      "[I 2024-02-21 03:04:33,609] Trial 34 finished with value: 0.9620673316708229 and parameters: {'weights': 'distance', 'n_components': 5, 'n_neighbors': 94}. Best is trial 11 with value: 0.9650598503740648.\n",
      "[I 2024-02-21 03:04:33,796] Trial 35 finished with value: 0.9595723192019949 and parameters: {'weights': 'uniform', 'n_components': 5, 'n_neighbors': 87}. Best is trial 11 with value: 0.9650598503740648.\n",
      "[I 2024-02-21 03:04:33,966] Trial 36 finished with value: 0.9495860349127181 and parameters: {'weights': 'uniform', 'n_components': 9, 'n_neighbors': 87}. Best is trial 11 with value: 0.9650598503740648.\n",
      "[I 2024-02-21 03:04:34,138] Trial 37 finished with value: 0.9485872817955112 and parameters: {'weights': 'uniform', 'n_components': 10, 'n_neighbors': 100}. Best is trial 11 with value: 0.9650598503740648.\n",
      "[I 2024-02-21 03:04:34,246] Trial 38 finished with value: 0.959571072319202 and parameters: {'weights': 'distance', 'n_components': 11, 'n_neighbors': 12}. Best is trial 11 with value: 0.9650598503740648.\n",
      "[I 2024-02-21 03:04:34,444] Trial 39 finished with value: 0.909149625935162 and parameters: {'weights': 'distance', 'n_components': 50, 'n_neighbors': 43}. Best is trial 11 with value: 0.9650598503740648.\n",
      "[I 2024-02-21 03:04:34,624] Trial 40 finished with value: 0.9465897755610971 and parameters: {'weights': 'distance', 'n_components': 14, 'n_neighbors': 69}. Best is trial 11 with value: 0.9650598503740648.\n",
      "[I 2024-02-21 03:04:34,840] Trial 41 finished with value: 0.9440922693266833 and parameters: {'weights': 'distance', 'n_components': 14, 'n_neighbors': 72}. Best is trial 11 with value: 0.9650598503740648.\n",
      "[I 2024-02-21 03:04:34,995] Trial 42 finished with value: 0.961568578553616 and parameters: {'weights': 'distance', 'n_components': 5, 'n_neighbors': 95}. Best is trial 11 with value: 0.9650598503740648.\n",
      "[I 2024-02-21 03:04:35,176] Trial 43 finished with value: 0.9650598503740648 and parameters: {'weights': 'distance', 'n_components': 4, 'n_neighbors': 95}. Best is trial 11 with value: 0.9650598503740648.\n",
      "[I 2024-02-21 03:04:35,286] Trial 44 finished with value: 0.9650598503740648 and parameters: {'weights': 'distance', 'n_components': 4, 'n_neighbors': 96}. Best is trial 11 with value: 0.9650598503740648.\n",
      "[I 2024-02-21 03:04:35,439] Trial 45 finished with value: 0.910149625935162 and parameters: {'weights': 'distance', 'n_components': 3, 'n_neighbors': 90}. Best is trial 11 with value: 0.9650598503740648.\n",
      "[I 2024-02-21 03:04:35,588] Trial 46 finished with value: 0.9091508728179551 and parameters: {'weights': 'distance', 'n_components': 3, 'n_neighbors': 97}. Best is trial 11 with value: 0.9650598503740648.\n",
      "[I 2024-02-21 03:04:35,781] Trial 47 finished with value: 0.9545798004987531 and parameters: {'weights': 'distance', 'n_components': 8, 'n_neighbors': 99}. Best is trial 11 with value: 0.9650598503740648.\n",
      "[I 2024-02-21 03:04:35,926] Trial 48 finished with value: 0.9470885286783043 and parameters: {'weights': 'distance', 'n_components': 12, 'n_neighbors': 80}. Best is trial 11 with value: 0.9650598503740648.\n",
      "[I 2024-02-21 03:04:36,096] Trial 49 finished with value: 0.9226284289276808 and parameters: {'weights': 'distance', 'n_components': 29, 'n_neighbors': 80}. Best is trial 11 with value: 0.9650598503740648.\n",
      "[I 2024-02-21 03:04:36,264] Trial 50 finished with value: 0.9186334164588528 and parameters: {'weights': 'distance', 'n_components': 40, 'n_neighbors': 26}. Best is trial 11 with value: 0.9650598503740648.\n",
      "[I 2024-02-21 03:04:36,414] Trial 51 finished with value: 0.9635610972568578 and parameters: {'weights': 'distance', 'n_components': 7, 'n_neighbors': 37}. Best is trial 11 with value: 0.9650598503740648.\n",
      "[I 2024-02-21 03:04:36,547] Trial 52 finished with value: 0.9620673316708229 and parameters: {'weights': 'distance', 'n_components': 5, 'n_neighbors': 91}. Best is trial 11 with value: 0.9650598503740648.\n",
      "[I 2024-02-21 03:04:36,708] Trial 53 finished with value: 0.961071072319202 and parameters: {'weights': 'distance', 'n_components': 5, 'n_neighbors': 91}. Best is trial 11 with value: 0.9650598503740648.\n",
      "[I 2024-02-21 03:04:36,873] Trial 54 finished with value: 0.911647132169576 and parameters: {'weights': 'distance', 'n_components': 3, 'n_neighbors': 85}. Best is trial 11 with value: 0.9650598503740648.\n",
      "[I 2024-02-21 03:04:37,036] Trial 55 finished with value: 0.9500860349127181 and parameters: {'weights': 'distance', 'n_components': 11, 'n_neighbors': 85}. Best is trial 11 with value: 0.9650598503740648.\n",
      "[I 2024-02-21 03:04:37,154] Trial 56 finished with value: 0.9495885286783041 and parameters: {'weights': 'distance', 'n_components': 10, 'n_neighbors': 96}. Best is trial 11 with value: 0.9650598503740648.\n",
      "[I 2024-02-21 03:04:37,355] Trial 57 finished with value: 0.9540822942643393 and parameters: {'weights': 'distance', 'n_components': 8, 'n_neighbors': 97}. Best is trial 11 with value: 0.9650598503740648.\n",
      "[I 2024-02-21 03:04:37,478] Trial 58 finished with value: 0.9540798004987533 and parameters: {'weights': 'distance', 'n_components': 8, 'n_neighbors': 98}. Best is trial 11 with value: 0.9650598503740648.\n",
      "[I 2024-02-21 03:04:37,657] Trial 59 finished with value: 0.9565748129675811 and parameters: {'weights': 'uniform', 'n_components': 7, 'n_neighbors': 88}. Best is trial 11 with value: 0.9650598503740648.\n",
      "[I 2024-02-21 03:04:37,951] Trial 60 finished with value: 0.963063591022444 and parameters: {'weights': 'uniform', 'n_components': 4, 'n_neighbors': 88}. Best is trial 11 with value: 0.9650598503740648.\n",
      "[I 2024-02-21 03:04:38,119] Trial 61 finished with value: 0.9680561097256859 and parameters: {'weights': 'distance', 'n_components': 4, 'n_neighbors': 55}. Best is trial 61 with value: 0.9680561097256859.\n",
      "[I 2024-02-21 03:04:38,256] Trial 62 finished with value: 0.9655610972568578 and parameters: {'weights': 'distance', 'n_components': 6, 'n_neighbors': 43}. Best is trial 61 with value: 0.9680561097256859.\n",
      "[I 2024-02-21 03:04:38,407] Trial 63 finished with value: 0.9670573566084787 and parameters: {'weights': 'distance', 'n_components': 4, 'n_neighbors': 60}. Best is trial 61 with value: 0.9680561097256859.\n",
      "[I 2024-02-21 03:04:38,568] Trial 64 finished with value: 0.9630648379052369 and parameters: {'weights': 'distance', 'n_components': 6, 'n_neighbors': 59}. Best is trial 61 with value: 0.9680561097256859.\n",
      "[I 2024-02-21 03:04:38,710] Trial 65 finished with value: 0.9630673316708229 and parameters: {'weights': 'distance', 'n_components': 6, 'n_neighbors': 61}. Best is trial 61 with value: 0.9680561097256859.\n",
      "[I 2024-02-21 03:04:38,866] Trial 66 finished with value: 0.9675548628428927 and parameters: {'weights': 'distance', 'n_components': 4, 'n_neighbors': 48}. Best is trial 61 with value: 0.9680561097256859.\n",
      "[I 2024-02-21 03:04:38,975] Trial 67 finished with value: 0.9440947630922694 and parameters: {'weights': 'distance', 'n_components': 16, 'n_neighbors': 49}. Best is trial 61 with value: 0.9680561097256859.\n",
      "[I 2024-02-21 03:04:39,123] Trial 68 finished with value: 0.9156408977556107 and parameters: {'weights': 'distance', 'n_components': 3, 'n_neighbors': 51}. Best is trial 61 with value: 0.9680561097256859.\n",
      "[I 2024-02-21 03:04:39,271] Trial 69 finished with value: 0.9575735660847882 and parameters: {'weights': 'distance', 'n_components': 10, 'n_neighbors': 42}. Best is trial 61 with value: 0.9680561097256859.\n",
      "[I 2024-02-21 03:04:39,472] Trial 70 finished with value: 0.955077306733167 and parameters: {'weights': 'distance', 'n_components': 12, 'n_neighbors': 41}. Best is trial 61 with value: 0.9680561097256859.\n",
      "[I 2024-02-21 03:04:39,575] Trial 71 finished with value: 0.960568578553616 and parameters: {'weights': 'distance', 'n_components': 7, 'n_neighbors': 46}. Best is trial 61 with value: 0.9680561097256859.\n",
      "[I 2024-02-21 03:04:39,733] Trial 72 finished with value: 0.9675548628428927 and parameters: {'weights': 'distance', 'n_components': 4, 'n_neighbors': 48}. Best is trial 61 with value: 0.9680561097256859.\n",
      "[I 2024-02-21 03:04:39,867] Trial 73 finished with value: 0.9700511221945138 and parameters: {'weights': 'distance', 'n_components': 4, 'n_neighbors': 38}. Best is trial 73 with value: 0.9700511221945138.\n",
      "[I 2024-02-21 03:04:40,037] Trial 74 finished with value: 0.9675586034912719 and parameters: {'weights': 'distance', 'n_components': 5, 'n_neighbors': 47}. Best is trial 73 with value: 0.9700511221945138.\n",
      "[I 2024-02-21 03:04:40,153] Trial 75 finished with value: 0.965561097256858 and parameters: {'weights': 'distance', 'n_components': 6, 'n_neighbors': 34}. Best is trial 73 with value: 0.9700511221945138.\n",
      "[I 2024-02-21 03:04:40,322] Trial 76 finished with value: 0.9660598503740647 and parameters: {'weights': 'distance', 'n_components': 6, 'n_neighbors': 38}. Best is trial 73 with value: 0.9700511221945138.\n",
      "[I 2024-02-21 03:04:40,501] Trial 77 finished with value: 0.9600723192019949 and parameters: {'weights': 'distance', 'n_components': 9, 'n_neighbors': 38}. Best is trial 73 with value: 0.9700511221945138.\n",
      "[I 2024-02-21 03:04:40,623] Trial 78 finished with value: 0.959571072319202 and parameters: {'weights': 'distance', 'n_components': 9, 'n_neighbors': 35}. Best is trial 73 with value: 0.9700511221945138.\n",
      "[I 2024-02-21 03:04:40,772] Trial 79 finished with value: 0.962069825436409 and parameters: {'weights': 'distance', 'n_components': 8, 'n_neighbors': 32}. Best is trial 73 with value: 0.9700511221945138.\n",
      "[I 2024-02-21 03:04:40,935] Trial 80 finished with value: 0.9665573566084789 and parameters: {'weights': 'distance', 'n_components': 5, 'n_neighbors': 46}. Best is trial 73 with value: 0.9700511221945138.\n",
      "[I 2024-02-21 03:04:41,070] Trial 81 finished with value: 0.9650635910224439 and parameters: {'weights': 'distance', 'n_components': 5, 'n_neighbors': 54}. Best is trial 73 with value: 0.9700511221945138.\n",
      "[I 2024-02-21 03:04:41,224] Trial 82 finished with value: 0.9650623441396509 and parameters: {'weights': 'distance', 'n_components': 5, 'n_neighbors': 55}. Best is trial 73 with value: 0.9700511221945138.\n",
      "[I 2024-02-21 03:04:41,345] Trial 83 finished with value: 0.9675548628428927 and parameters: {'weights': 'distance', 'n_components': 4, 'n_neighbors': 48}. Best is trial 73 with value: 0.9700511221945138.\n",
      "[I 2024-02-21 03:04:41,487] Trial 84 finished with value: 0.9630673316708229 and parameters: {'weights': 'distance', 'n_components': 6, 'n_neighbors': 48}. Best is trial 73 with value: 0.9700511221945138.\n",
      "[I 2024-02-21 03:04:41,653] Trial 85 finished with value: 0.9171384039900248 and parameters: {'weights': 'distance', 'n_components': 3, 'n_neighbors': 47}. Best is trial 73 with value: 0.9700511221945138.\n",
      "[I 2024-02-21 03:04:41,790] Trial 86 finished with value: 0.9176371571072319 and parameters: {'weights': 'distance', 'n_components': 3, 'n_neighbors': 46}. Best is trial 73 with value: 0.9700511221945138.\n",
      "[I 2024-02-21 03:04:41,989] Trial 87 finished with value: 0.9675561097256858 and parameters: {'weights': 'distance', 'n_components': 4, 'n_neighbors': 58}. Best is trial 73 with value: 0.9700511221945138.\n",
      "[I 2024-02-21 03:04:42,141] Trial 88 finished with value: 0.9700523690773067 and parameters: {'weights': 'uniform', 'n_components': 4, 'n_neighbors': 28}. Best is trial 88 with value: 0.9700523690773067.\n",
      "[I 2024-02-21 03:04:42,325] Trial 89 finished with value: 0.9655586034912718 and parameters: {'weights': 'uniform', 'n_components': 4, 'n_neighbors': 58}. Best is trial 88 with value: 0.9700523690773067.\n",
      "[I 2024-02-21 03:04:42,436] Trial 90 finished with value: 0.9685523690773067 and parameters: {'weights': 'uniform', 'n_components': 4, 'n_neighbors': 24}. Best is trial 88 with value: 0.9700523690773067.\n",
      "[I 2024-02-21 03:04:42,609] Trial 91 finished with value: 0.9351097256857855 and parameters: {'weights': 'uniform', 'n_components': 27, 'n_neighbors': 27}. Best is trial 88 with value: 0.9700523690773067.\n",
      "[I 2024-02-21 03:04:42,771] Trial 92 finished with value: 0.9680523690773066 and parameters: {'weights': 'uniform', 'n_components': 4, 'n_neighbors': 25}. Best is trial 88 with value: 0.9700523690773067.\n",
      "[I 2024-02-21 03:04:42,932] Trial 93 finished with value: 0.965563591022444 and parameters: {'weights': 'uniform', 'n_components': 7, 'n_neighbors': 21}. Best is trial 88 with value: 0.9700523690773067.\n",
      "[I 2024-02-21 03:04:43,096] Trial 94 finished with value: 0.9655598503740649 and parameters: {'weights': 'uniform', 'n_components': 7, 'n_neighbors': 23}. Best is trial 88 with value: 0.9700523690773067.\n",
      "[I 2024-02-21 03:04:43,288] Trial 95 finished with value: 0.9695511221945136 and parameters: {'weights': 'uniform', 'n_components': 4, 'n_neighbors': 23}. Best is trial 88 with value: 0.9700523690773067.\n",
      "[I 2024-02-21 03:04:43,424] Trial 96 finished with value: 0.9645610972568578 and parameters: {'weights': 'uniform', 'n_components': 4, 'n_neighbors': 63}. Best is trial 88 with value: 0.9700523690773067.\n",
      "[I 2024-02-21 03:04:43,587] Trial 97 finished with value: 0.920134663341646 and parameters: {'weights': 'uniform', 'n_components': 3, 'n_neighbors': 30}. Best is trial 88 with value: 0.9700523690773067.\n",
      "[I 2024-02-21 03:04:43,760] Trial 98 finished with value: 0.913640897755611 and parameters: {'weights': 'uniform', 'n_components': 3, 'n_neighbors': 15}. Best is trial 88 with value: 0.9700523690773067.\n",
      "[I 2024-02-21 03:04:43,883] Trial 99 finished with value: 0.9221309226932668 and parameters: {'weights': 'uniform', 'n_components': 41, 'n_neighbors': 15}. Best is trial 88 with value: 0.9700523690773067.\n"
     ]
    }
   ],
   "source": [
    "study = optuna.create_study(direction=\"maximize\")\n",
    "study.optimize(objective, n_trials=100, n_jobs=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Accuracy: 0.964\n",
      "Balanced accuracy: 0.954\n"
     ]
    }
   ],
   "source": [
    "model = make_pipeline(\n",
    "    OneHotEncoder(handle_unknown=\"ignore\", sparse_output=False),\n",
    "    PCA(n_components=study.best_params[\"n_components\"]),\n",
    "    KNeighborsClassifier(n_neighbors=study.best_params[\"n_neighbors\"], weights=study.best_params[\"weights\"], n_jobs=4),\n",
    ")\n",
    "\n",
    "model.fit(X_train, y_train)\n",
    "y_test_pred = model.predict(X_test)\n",
    "\n",
    "print(f\"Accuracy: {accuracy_score(y_test, y_test_pred).round(3)}\")\n",
    "print(f\"Balanced accuracy: {balanced_accuracy_score(y_test, y_test_pred).round(3)}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['/home/vscode/.ezancestry/data/models/kidd_superpopulation.pkl']"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Save the model\n",
    "model_path = Path(_models_directory) / f\"kidd_{label}.pkl\"\n",
    "model_path.parent.mkdir(parents=True, exist_ok=True)\n",
    "joblib.dump(model, model_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Kidd AISNPs to predict population"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {},
   "outputs": [],
   "source": [
    "label = \"population\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "X_train, X_test, y_train, y_test = train_test_split(\n",
    "    kidd1kg.drop(columns=todrop),\n",
    "    kidd1kg[label],\n",
    "    test_size=0.2,\n",
    "    random_state=42,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "topk_scorer = make_scorer(top_k_accuracy_score, needs_proba=True, k=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "def objective(trial):\n",
    "    weights = trial.suggest_categorical(\"weights\", [\"uniform\", \"distance\"])\n",
    "    n_components = trial.suggest_int(\"n_components\", 3, 50)\n",
    "    n_neighbors = trial.suggest_int(\"n_neighbors\", 11, 100)\n",
    "\n",
    "    model = make_pipeline(\n",
    "        OneHotEncoder(handle_unknown=\"ignore\", sparse_output=False, dtype=np.int8),\n",
    "        PCA(n_components=n_components),\n",
    "        KNeighborsClassifier(n_neighbors=n_neighbors, weights=weights, n_jobs=4),\n",
    "    )\n",
    "\n",
    "    return cross_val_score(model, X_train, y_train, n_jobs=4, scoring=topk_scorer, cv=StratifiedKFold(5)).mean()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[I 2024-02-21 03:04:44,326] A new study created in memory with name: no-name-6902485b-1095-4caf-8b16-6e18396ef7f5\n",
      "[I 2024-02-21 03:04:44,759] Trial 0 finished with value: 0.5307157107231921 and parameters: {'weights': 'distance', 'n_components': 10, 'n_neighbors': 70}. Best is trial 0 with value: 0.5307157107231921.\n",
      "[I 2024-02-21 03:04:44,943] Trial 1 finished with value: 0.5102456359102245 and parameters: {'weights': 'uniform', 'n_components': 26, 'n_neighbors': 58}. Best is trial 0 with value: 0.5307157107231921.\n",
      "[I 2024-02-21 03:04:45,157] Trial 2 finished with value: 0.4937643391521197 and parameters: {'weights': 'distance', 'n_components': 18, 'n_neighbors': 13}. Best is trial 0 with value: 0.5307157107231921.\n",
      "[I 2024-02-21 03:04:45,328] Trial 3 finished with value: 0.48927182044887785 and parameters: {'weights': 'uniform', 'n_components': 29, 'n_neighbors': 22}. Best is trial 0 with value: 0.5307157107231921.\n",
      "[I 2024-02-21 03:04:45,494] Trial 4 finished with value: 0.5017730673316708 and parameters: {'weights': 'distance', 'n_components': 35, 'n_neighbors': 33}. Best is trial 0 with value: 0.5307157107231921.\n",
      "[I 2024-02-21 03:04:45,655] Trial 5 finished with value: 0.4857805486284289 and parameters: {'weights': 'distance', 'n_components': 38, 'n_neighbors': 36}. Best is trial 0 with value: 0.5307157107231921.\n",
      "[I 2024-02-21 03:04:45,775] Trial 6 finished with value: 0.5316982543640897 and parameters: {'weights': 'distance', 'n_components': 5, 'n_neighbors': 20}. Best is trial 6 with value: 0.5316982543640897.\n",
      "[I 2024-02-21 03:04:45,956] Trial 7 finished with value: 0.49327182044887774 and parameters: {'weights': 'distance', 'n_components': 47, 'n_neighbors': 91}. Best is trial 6 with value: 0.5316982543640897.\n",
      "[I 2024-02-21 03:04:46,106] Trial 8 finished with value: 0.5002581047381546 and parameters: {'weights': 'distance', 'n_components': 32, 'n_neighbors': 39}. Best is trial 6 with value: 0.5316982543640897.\n",
      "[I 2024-02-21 03:04:46,232] Trial 9 finished with value: 0.49975935162094764 and parameters: {'weights': 'uniform', 'n_components': 30, 'n_neighbors': 64}. Best is trial 6 with value: 0.5316982543640897.\n",
      "[I 2024-02-21 03:04:46,377] Trial 10 finished with value: 0.529215710723192 and parameters: {'weights': 'distance', 'n_components': 20, 'n_neighbors': 69}. Best is trial 6 with value: 0.5316982543640897.\n",
      "[I 2024-02-21 03:04:46,530] Trial 11 finished with value: 0.5341970074812967 and parameters: {'weights': 'uniform', 'n_components': 8, 'n_neighbors': 85}. Best is trial 11 with value: 0.5341970074812967.\n",
      "[I 2024-02-21 03:04:46,669] Trial 12 finished with value: 0.5267306733167083 and parameters: {'weights': 'distance', 'n_components': 3, 'n_neighbors': 85}. Best is trial 11 with value: 0.5341970074812967.\n",
      "[I 2024-02-21 03:04:46,823] Trial 13 finished with value: 0.5197231920199501 and parameters: {'weights': 'uniform', 'n_components': 3, 'n_neighbors': 94}. Best is trial 11 with value: 0.5341970074812967.\n",
      "[I 2024-02-21 03:04:46,938] Trial 14 finished with value: 0.5147331670822943 and parameters: {'weights': 'uniform', 'n_components': 3, 'n_neighbors': 97}. Best is trial 11 with value: 0.5341970074812967.\n",
      "[I 2024-02-21 03:04:47,169] Trial 15 finished with value: 0.5267082294264339 and parameters: {'weights': 'uniform', 'n_components': 11, 'n_neighbors': 46}. Best is trial 11 with value: 0.5341970074812967.\n",
      "[I 2024-02-21 03:04:47,354] Trial 16 finished with value: 0.5197281795511222 and parameters: {'weights': 'uniform', 'n_components': 12, 'n_neighbors': 49}. Best is trial 11 with value: 0.5341970074812967.\n",
      "[I 2024-02-21 03:04:47,505] Trial 17 finished with value: 0.5302157107231921 and parameters: {'weights': 'uniform', 'n_components': 12, 'n_neighbors': 82}. Best is trial 11 with value: 0.5341970074812967.\n",
      "[I 2024-02-21 03:04:47,609] Trial 18 finished with value: 0.512221945137157 and parameters: {'weights': 'uniform', 'n_components': 17, 'n_neighbors': 81}. Best is trial 11 with value: 0.5341970074812967.\n",
      "[I 2024-02-21 03:04:47,791] Trial 19 finished with value: 0.514719451371571 and parameters: {'weights': 'distance', 'n_components': 19, 'n_neighbors': 23}. Best is trial 11 with value: 0.5341970074812967.\n",
      "[I 2024-02-21 03:04:47,949] Trial 20 finished with value: 0.5152431421446384 and parameters: {'weights': 'distance', 'n_components': 8, 'n_neighbors': 25}. Best is trial 11 with value: 0.5341970074812967.\n",
      "[I 2024-02-21 03:04:48,066] Trial 21 finished with value: 0.5411695760598504 and parameters: {'weights': 'distance', 'n_components': 7, 'n_neighbors': 72}. Best is trial 21 with value: 0.5411695760598504.\n",
      "[I 2024-02-21 03:04:48,219] Trial 22 finished with value: 0.5376957605985038 and parameters: {'weights': 'distance', 'n_components': 8, 'n_neighbors': 71}. Best is trial 21 with value: 0.5411695760598504.\n",
      "[I 2024-02-21 03:04:48,419] Trial 23 finished with value: 0.5367032418952619 and parameters: {'weights': 'distance', 'n_components': 8, 'n_neighbors': 73}. Best is trial 21 with value: 0.5411695760598504.\n",
      "[I 2024-02-21 03:04:48,534] Trial 24 finished with value: 0.5072431421446384 and parameters: {'weights': 'distance', 'n_components': 23, 'n_neighbors': 73}. Best is trial 21 with value: 0.5411695760598504.\n",
      "[I 2024-02-21 03:04:48,684] Trial 25 finished with value: 0.5047381546134664 and parameters: {'weights': 'distance', 'n_components': 23, 'n_neighbors': 75}. Best is trial 21 with value: 0.5411695760598504.\n",
      "[I 2024-02-21 03:04:48,882] Trial 26 finished with value: 0.5202269326683292 and parameters: {'weights': 'distance', 'n_components': 13, 'n_neighbors': 58}. Best is trial 21 with value: 0.5411695760598504.\n",
      "[I 2024-02-21 03:04:49,029] Trial 27 finished with value: 0.524713216957606 and parameters: {'weights': 'distance', 'n_components': 15, 'n_neighbors': 58}. Best is trial 21 with value: 0.5411695760598504.\n",
      "[I 2024-02-21 03:04:49,229] Trial 28 finished with value: 0.4957630922693267 and parameters: {'weights': 'distance', 'n_components': 16, 'n_neighbors': 62}. Best is trial 21 with value: 0.5411695760598504.\n",
      "[I 2024-02-21 03:04:49,399] Trial 29 finished with value: 0.5456733167082295 and parameters: {'weights': 'distance', 'n_components': 7, 'n_neighbors': 65}. Best is trial 29 with value: 0.5456733167082295.\n",
      "[I 2024-02-21 03:04:49,532] Trial 30 finished with value: 0.5451783042394015 and parameters: {'weights': 'distance', 'n_components': 6, 'n_neighbors': 68}. Best is trial 29 with value: 0.5456733167082295.\n",
      "[I 2024-02-21 03:04:49,659] Trial 31 finished with value: 0.5501633416458853 and parameters: {'weights': 'distance', 'n_components': 6, 'n_neighbors': 51}. Best is trial 31 with value: 0.5501633416458853.\n",
      "[I 2024-02-21 03:04:49,827] Trial 32 finished with value: 0.5486733167082295 and parameters: {'weights': 'distance', 'n_components': 7, 'n_neighbors': 65}. Best is trial 31 with value: 0.5501633416458853.\n",
      "[I 2024-02-21 03:04:49,927] Trial 33 finished with value: 0.5391895261845387 and parameters: {'weights': 'distance', 'n_components': 6, 'n_neighbors': 51}. Best is trial 31 with value: 0.5501633416458853.\n",
      "[I 2024-02-21 03:04:50,095] Trial 34 finished with value: 0.5411770573566085 and parameters: {'weights': 'distance', 'n_components': 6, 'n_neighbors': 51}. Best is trial 31 with value: 0.5501633416458853.\n",
      "[I 2024-02-21 03:04:50,273] Trial 35 finished with value: 0.5247057356608479 and parameters: {'weights': 'distance', 'n_components': 14, 'n_neighbors': 66}. Best is trial 31 with value: 0.5501633416458853.\n",
      "[I 2024-02-21 03:04:50,398] Trial 36 finished with value: 0.5271982543640898 and parameters: {'weights': 'distance', 'n_components': 11, 'n_neighbors': 44}. Best is trial 31 with value: 0.5501633416458853.\n",
      "[I 2024-02-21 03:04:50,579] Trial 37 finished with value: 0.5271845386533666 and parameters: {'weights': 'distance', 'n_components': 11, 'n_neighbors': 43}. Best is trial 31 with value: 0.5501633416458853.\n",
      "[I 2024-02-21 03:04:50,695] Trial 38 finished with value: 0.49927182044887786 and parameters: {'weights': 'distance', 'n_components': 46, 'n_neighbors': 60}. Best is trial 31 with value: 0.5501633416458853.\n",
      "[I 2024-02-21 03:04:50,874] Trial 39 finished with value: 0.5077581047381546 and parameters: {'weights': 'distance', 'n_components': 41, 'n_neighbors': 54}. Best is trial 31 with value: 0.5501633416458853.\n",
      "[I 2024-02-21 03:04:51,015] Trial 40 finished with value: 0.5496670822942644 and parameters: {'weights': 'distance', 'n_components': 5, 'n_neighbors': 78}. Best is trial 31 with value: 0.5501633416458853.\n",
      "[I 2024-02-21 03:04:51,196] Trial 41 finished with value: 0.5426683291770573 and parameters: {'weights': 'distance', 'n_components': 5, 'n_neighbors': 78}. Best is trial 31 with value: 0.5501633416458853.\n",
      "[I 2024-02-21 03:04:51,337] Trial 42 finished with value: 0.5471695760598504 and parameters: {'weights': 'distance', 'n_components': 5, 'n_neighbors': 79}. Best is trial 31 with value: 0.5501633416458853.\n",
      "[I 2024-02-21 03:04:51,467] Trial 43 finished with value: 0.5361882793017456 and parameters: {'weights': 'distance', 'n_components': 10, 'n_neighbors': 67}. Best is trial 31 with value: 0.5501633416458853.\n",
      "[I 2024-02-21 03:04:51,648] Trial 44 finished with value: 0.5217269326683291 and parameters: {'weights': 'distance', 'n_components': 10, 'n_neighbors': 86}. Best is trial 31 with value: 0.5501633416458853.\n",
      "[I 2024-02-21 03:04:51,808] Trial 45 finished with value: 0.5701583541147132 and parameters: {'weights': 'distance', 'n_components': 4, 'n_neighbors': 78}. Best is trial 45 with value: 0.5701583541147132.\n",
      "[I 2024-02-21 03:04:51,913] Trial 46 finished with value: 0.5312244389027431 and parameters: {'weights': 'distance', 'n_components': 3, 'n_neighbors': 78}. Best is trial 45 with value: 0.5701583541147132.\n",
      "[I 2024-02-21 03:04:52,072] Trial 47 finished with value: 0.5272331670822942 and parameters: {'weights': 'distance', 'n_components': 3, 'n_neighbors': 90}. Best is trial 45 with value: 0.5701583541147132.\n",
      "[I 2024-02-21 03:04:52,217] Trial 48 finished with value: 0.5696546134663342 and parameters: {'weights': 'distance', 'n_components': 4, 'n_neighbors': 89}. Best is trial 45 with value: 0.5701583541147132.\n",
      "[I 2024-02-21 03:04:52,404] Trial 49 finished with value: 0.5451720698254363 and parameters: {'weights': 'distance', 'n_components': 5, 'n_neighbors': 89}. Best is trial 45 with value: 0.5701583541147132.\n",
      "[I 2024-02-21 03:04:52,522] Trial 50 finished with value: 0.5287256857855361 and parameters: {'weights': 'distance', 'n_components': 9, 'n_neighbors': 91}. Best is trial 45 with value: 0.5701583541147132.\n",
      "[I 2024-02-21 03:04:52,711] Trial 51 finished with value: 0.5331995012468828 and parameters: {'weights': 'distance', 'n_components': 9, 'n_neighbors': 99}. Best is trial 45 with value: 0.5701583541147132.\n",
      "[I 2024-02-21 03:04:52,845] Trial 52 finished with value: 0.5471795511221946 and parameters: {'weights': 'distance', 'n_components': 5, 'n_neighbors': 96}. Best is trial 45 with value: 0.5701583541147132.\n",
      "[I 2024-02-21 03:04:53,005] Trial 53 finished with value: 0.5501733167082293 and parameters: {'weights': 'distance', 'n_components': 5, 'n_neighbors': 81}. Best is trial 45 with value: 0.5701583541147132.\n",
      "[I 2024-02-21 03:04:53,121] Trial 54 finished with value: 0.5706558603491271 and parameters: {'weights': 'distance', 'n_components': 4, 'n_neighbors': 97}. Best is trial 54 with value: 0.5706558603491271.\n",
      "[I 2024-02-21 03:04:53,282] Trial 55 finished with value: 0.5277331670822942 and parameters: {'weights': 'distance', 'n_components': 3, 'n_neighbors': 86}. Best is trial 54 with value: 0.5706558603491271.\n",
      "[I 2024-02-21 03:04:53,493] Trial 56 finished with value: 0.5207269326683291 and parameters: {'weights': 'uniform', 'n_components': 3, 'n_neighbors': 93}. Best is trial 54 with value: 0.5706558603491271.\n",
      "[I 2024-02-21 03:04:53,603] Trial 57 finished with value: 0.5631708229426434 and parameters: {'weights': 'uniform', 'n_components': 4, 'n_neighbors': 100}. Best is trial 54 with value: 0.5706558603491271.\n",
      "[I 2024-02-21 03:04:53,742] Trial 58 finished with value: 0.5027493765586034 and parameters: {'weights': 'distance', 'n_components': 33, 'n_neighbors': 83}. Best is trial 54 with value: 0.5706558603491271.\n",
      "[I 2024-02-21 03:04:53,909] Trial 59 finished with value: 0.4997556109725686 and parameters: {'weights': 'uniform', 'n_components': 30, 'n_neighbors': 99}. Best is trial 54 with value: 0.5706558603491271.\n",
      "[I 2024-02-21 03:04:54,042] Trial 60 finished with value: 0.4992481296758105 and parameters: {'weights': 'uniform', 'n_components': 27, 'n_neighbors': 98}. Best is trial 54 with value: 0.5706558603491271.\n",
      "[I 2024-02-21 03:04:54,186] Trial 61 finished with value: 0.49426683291770573 and parameters: {'weights': 'uniform', 'n_components': 27, 'n_neighbors': 95}. Best is trial 54 with value: 0.5706558603491271.\n",
      "[I 2024-02-21 03:04:54,456] Trial 62 finished with value: 0.5551770573566085 and parameters: {'weights': 'uniform', 'n_components': 4, 'n_neighbors': 93}. Best is trial 54 with value: 0.5706558603491271.\n",
      "[I 2024-02-21 03:04:54,584] Trial 63 finished with value: 0.5536857855361597 and parameters: {'weights': 'uniform', 'n_components': 4, 'n_neighbors': 89}. Best is trial 54 with value: 0.5706558603491271.\n",
      "[I 2024-02-21 03:04:54,761] Trial 64 finished with value: 0.5596720698254364 and parameters: {'weights': 'uniform', 'n_components': 4, 'n_neighbors': 89}. Best is trial 54 with value: 0.5706558603491271.\n",
      "[I 2024-02-21 03:04:54,920] Trial 65 finished with value: 0.5596758104738154 and parameters: {'weights': 'uniform', 'n_components': 4, 'n_neighbors': 88}. Best is trial 54 with value: 0.5706558603491271.\n",
      "[I 2024-02-21 03:04:55,053] Trial 66 finished with value: 0.5371870324189526 and parameters: {'weights': 'uniform', 'n_components': 8, 'n_neighbors': 88}. Best is trial 54 with value: 0.5706558603491271.\n",
      "[I 2024-02-21 03:04:55,289] Trial 67 finished with value: 0.5162281795511221 and parameters: {'weights': 'uniform', 'n_components': 9, 'n_neighbors': 92}. Best is trial 54 with value: 0.5706558603491271.\n",
      "[I 2024-02-21 03:04:55,413] Trial 68 finished with value: 0.5232219451371571 and parameters: {'weights': 'uniform', 'n_components': 9, 'n_neighbors': 94}. Best is trial 54 with value: 0.5706558603491271.\n",
      "[I 2024-02-21 03:04:55,667] Trial 69 finished with value: 0.516224438902743 and parameters: {'weights': 'uniform', 'n_components': 13, 'n_neighbors': 100}. Best is trial 54 with value: 0.5706558603491271.\n",
      "[I 2024-02-21 03:04:55,878] Trial 70 finished with value: 0.5316795511221943 and parameters: {'weights': 'uniform', 'n_components': 7, 'n_neighbors': 100}. Best is trial 54 with value: 0.5706558603491271.\n",
      "[I 2024-02-21 03:04:56,074] Trial 71 finished with value: 0.5316970074812968 and parameters: {'weights': 'uniform', 'n_components': 7, 'n_neighbors': 95}. Best is trial 54 with value: 0.5706558603491271.\n",
      "[I 2024-02-21 03:04:56,247] Trial 72 finished with value: 0.5606733167082294 and parameters: {'weights': 'uniform', 'n_components': 4, 'n_neighbors': 87}. Best is trial 54 with value: 0.5706558603491271.\n",
      "[I 2024-02-21 03:04:56,376] Trial 73 finished with value: 0.5556845386533665 and parameters: {'weights': 'uniform', 'n_components': 4, 'n_neighbors': 87}. Best is trial 54 with value: 0.5706558603491271.\n",
      "[I 2024-02-21 03:04:56,537] Trial 74 finished with value: 0.5596683291770573 and parameters: {'weights': 'uniform', 'n_components': 4, 'n_neighbors': 84}. Best is trial 54 with value: 0.5706558603491271.\n",
      "[I 2024-02-21 03:04:56,665] Trial 75 finished with value: 0.5616670822942644 and parameters: {'weights': 'uniform', 'n_components': 4, 'n_neighbors': 83}. Best is trial 54 with value: 0.5706558603491271.\n",
      "[I 2024-02-21 03:04:56,834] Trial 76 finished with value: 0.5167356608478804 and parameters: {'weights': 'uniform', 'n_components': 3, 'n_neighbors': 84}. Best is trial 54 with value: 0.5706558603491271.\n",
      "[I 2024-02-21 03:04:56,944] Trial 77 finished with value: 0.5416770573566085 and parameters: {'weights': 'uniform', 'n_components': 6, 'n_neighbors': 81}. Best is trial 54 with value: 0.5706558603491271.\n",
      "[I 2024-02-21 03:04:57,117] Trial 78 finished with value: 0.5052381546134663 and parameters: {'weights': 'uniform', 'n_components': 21, 'n_neighbors': 76}. Best is trial 54 with value: 0.5706558603491271.\n",
      "[I 2024-02-21 03:04:57,304] Trial 79 finished with value: 0.5132306733167082 and parameters: {'weights': 'uniform', 'n_components': 7, 'n_neighbors': 12}. Best is trial 54 with value: 0.5706558603491271.\n",
      "[I 2024-02-21 03:04:57,449] Trial 80 finished with value: 0.512718204488778 and parameters: {'weights': 'uniform', 'n_components': 12, 'n_neighbors': 30}. Best is trial 54 with value: 0.5706558603491271.\n",
      "[I 2024-02-21 03:04:57,625] Trial 81 finished with value: 0.5232182044887781 and parameters: {'weights': 'uniform', 'n_components': 11, 'n_neighbors': 97}. Best is trial 54 with value: 0.5706558603491271.\n",
      "[I 2024-02-21 03:04:57,805] Trial 82 finished with value: 0.5576745635910225 and parameters: {'weights': 'uniform', 'n_components': 4, 'n_neighbors': 85}. Best is trial 54 with value: 0.5706558603491271.\n",
      "[I 2024-02-21 03:04:57,942] Trial 83 finished with value: 0.5576733167082295 and parameters: {'weights': 'uniform', 'n_components': 4, 'n_neighbors': 84}. Best is trial 54 with value: 0.5706558603491271.\n",
      "[I 2024-02-21 03:04:58,110] Trial 84 finished with value: 0.5332032418952618 and parameters: {'weights': 'uniform', 'n_components': 6, 'n_neighbors': 88}. Best is trial 54 with value: 0.5706558603491271.\n",
      "[I 2024-02-21 03:04:58,226] Trial 85 finished with value: 0.5386932668329176 and parameters: {'weights': 'uniform', 'n_components': 6, 'n_neighbors': 75}. Best is trial 54 with value: 0.5706558603491271.\n",
      "[I 2024-02-21 03:04:58,426] Trial 86 finished with value: 0.527718204488778 and parameters: {'weights': 'uniform', 'n_components': 8, 'n_neighbors': 91}. Best is trial 54 with value: 0.5706558603491271.\n",
      "[I 2024-02-21 03:04:58,549] Trial 87 finished with value: 0.5267144638403989 and parameters: {'weights': 'uniform', 'n_components': 8, 'n_neighbors': 91}. Best is trial 54 with value: 0.5706558603491271.\n",
      "[I 2024-02-21 03:04:58,710] Trial 88 finished with value: 0.4802905236907731 and parameters: {'weights': 'uniform', 'n_components': 40, 'n_neighbors': 83}. Best is trial 54 with value: 0.5706558603491271.\n",
      "[I 2024-02-21 03:04:58,881] Trial 89 finished with value: 0.48528802992518705 and parameters: {'weights': 'uniform', 'n_components': 50, 'n_neighbors': 80}. Best is trial 54 with value: 0.5706558603491271.\n",
      "[I 2024-02-21 03:04:58,999] Trial 90 finished with value: 0.5641571072319203 and parameters: {'weights': 'uniform', 'n_components': 4, 'n_neighbors': 80}. Best is trial 54 with value: 0.5706558603491271.\n",
      "[I 2024-02-21 03:04:59,154] Trial 91 finished with value: 0.5531857855361597 and parameters: {'weights': 'uniform', 'n_components': 4, 'n_neighbors': 96}. Best is trial 54 with value: 0.5706558603491271.\n",
      "[I 2024-02-21 03:04:59,316] Trial 92 finished with value: 0.5576745635910224 and parameters: {'weights': 'uniform', 'n_components': 4, 'n_neighbors': 87}. Best is trial 54 with value: 0.5706558603491271.\n",
      "[I 2024-02-21 03:04:59,424] Trial 93 finished with value: 0.5097443890274314 and parameters: {'weights': 'uniform', 'n_components': 3, 'n_neighbors': 17}. Best is trial 54 with value: 0.5706558603491271.\n",
      "[I 2024-02-21 03:04:59,584] Trial 94 finished with value: 0.5222369077306732 and parameters: {'weights': 'uniform', 'n_components': 3, 'n_neighbors': 71}. Best is trial 54 with value: 0.5706558603491271.\n",
      "[I 2024-02-21 03:04:59,689] Trial 95 finished with value: 0.5481820448877806 and parameters: {'weights': 'uniform', 'n_components': 6, 'n_neighbors': 76}. Best is trial 54 with value: 0.5706558603491271.\n",
      "[I 2024-02-21 03:04:59,884] Trial 96 finished with value: 0.5406845386533666 and parameters: {'weights': 'uniform', 'n_components': 5, 'n_neighbors': 90}. Best is trial 54 with value: 0.5706558603491271.\n",
      "[I 2024-02-21 03:05:00,063] Trial 97 finished with value: 0.5436745635910224 and parameters: {'weights': 'uniform', 'n_components': 5, 'n_neighbors': 89}. Best is trial 54 with value: 0.5706558603491271.\n",
      "[I 2024-02-21 03:05:00,212] Trial 98 finished with value: 0.5531670822942644 and parameters: {'weights': 'uniform', 'n_components': 5, 'n_neighbors': 82}. Best is trial 54 with value: 0.5706558603491271.\n",
      "[I 2024-02-21 03:05:00,358] Trial 99 finished with value: 0.5237219451371571 and parameters: {'weights': 'uniform', 'n_components': 10, 'n_neighbors': 93}. Best is trial 54 with value: 0.5706558603491271.\n"
     ]
    }
   ],
   "source": [
    "study = optuna.create_study(direction=\"maximize\")\n",
    "study.optimize(objective, n_trials=100, n_jobs=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Accuracy: 0.339\n",
      "Balanced accuracy: 0.347\n",
      "Top-2 accuracy: 0.551\n",
      "Top-3 accuracy: 0.727\n"
     ]
    }
   ],
   "source": [
    "model = make_pipeline(\n",
    "    OneHotEncoder(handle_unknown=\"ignore\", sparse_output=False, dtype=np.int8),\n",
    "    PCA(n_components=study.best_params[\"n_components\"]),\n",
    "    KNeighborsClassifier(n_neighbors=study.best_params[\"n_neighbors\"], weights=study.best_params[\"weights\"], n_jobs=4),\n",
    ")\n",
    "\n",
    "model.fit(X_train, y_train)\n",
    "y_test_pred = model.predict(X_test)\n",
    "\n",
    "print(f\"Accuracy: {accuracy_score(y_test, y_test_pred).round(3)}\")\n",
    "print(f\"Balanced accuracy: {balanced_accuracy_score(y_test, y_test_pred).round(3)}\")\n",
    "print(f\"Top-2 accuracy: {top_k_accuracy_score(y_test, model.predict_proba(X_test), k=2).round(3)}\")\n",
    "print(f\"Top-3 accuracy: {top_k_accuracy_score(y_test, model.predict_proba(X_test), k=3).round(3)}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['/home/vscode/.ezancestry/data/models/kidd_population.pkl']"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Save the model\n",
    "model_path = Path(_models_directory) / f\"kidd_{label}.pkl\"\n",
    "model_path.parent.mkdir(parents=True, exist_ok=True)\n",
    "joblib.dump(model, model_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Seldin AISNPs to predict superpopulation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "todrop = [\"superpopulation\", \"population\", \"gender\"]\n",
    "label = \"superpopulation\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "X_train, X_test, y_train, y_test = train_test_split(\n",
    "    seldin1kg.drop(columns=todrop),\n",
    "    seldin1kg[label],\n",
    "    test_size=0.2,\n",
    "    random_state=42,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "def objective(trial):\n",
    "    weights = trial.suggest_categorical(\"weights\", [\"uniform\", \"distance\"])\n",
    "    n_components = trial.suggest_int(\"n_components\", 3, 50)\n",
    "    n_neighbors = trial.suggest_int(\"n_neighbors\", 11, 100)\n",
    "    \n",
    "    model = make_pipeline(\n",
    "        OneHotEncoder(handle_unknown=\"ignore\", sparse_output=False, dtype=np.int8),\n",
    "        PCA(n_components=n_components),\n",
    "        KNeighborsClassifier(n_neighbors=n_neighbors, weights=weights, n_jobs=4),\n",
    "    )\n",
    "\n",
    "    return cross_val_score(model, X_train, y_train, n_jobs=4, scoring=\"accuracy\", cv=StratifiedKFold(5)).mean()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[I 2024-02-21 03:05:00,794] A new study created in memory with name: no-name-b670e5e7-4a67-41af-b505-fb09a3a6b22a\n",
      "[I 2024-02-21 03:05:01,610] Trial 1 finished with value: 0.9116496259351621 and parameters: {'weights': 'distance', 'n_components': 33, 'n_neighbors': 97}. Best is trial 1 with value: 0.9116496259351621.\n",
      "[I 2024-02-21 03:05:01,681] Trial 0 finished with value: 0.9251284289276807 and parameters: {'weights': 'distance', 'n_components': 28, 'n_neighbors': 64}. Best is trial 0 with value: 0.9251284289276807.\n",
      "[I 2024-02-21 03:05:02,135] Trial 2 finished with value: 0.9650635910224439 and parameters: {'weights': 'uniform', 'n_components': 5, 'n_neighbors': 80}. Best is trial 2 with value: 0.9650635910224439.\n",
      "[I 2024-02-21 03:05:02,420] Trial 3 finished with value: 0.9540785536159602 and parameters: {'weights': 'distance', 'n_components': 14, 'n_neighbors': 52}. Best is trial 2 with value: 0.9650635910224439.\n",
      "[I 2024-02-21 03:05:02,695] Trial 4 finished with value: 0.9366072319201993 and parameters: {'weights': 'uniform', 'n_components': 24, 'n_neighbors': 29}. Best is trial 2 with value: 0.9650635910224439.\n",
      "[I 2024-02-21 03:05:02,961] Trial 5 finished with value: 0.9725486284289276 and parameters: {'weights': 'distance', 'n_components': 7, 'n_neighbors': 30}. Best is trial 5 with value: 0.9725486284289276.\n",
      "[I 2024-02-21 03:05:03,284] Trial 6 finished with value: 0.915143391521197 and parameters: {'weights': 'uniform', 'n_components': 43, 'n_neighbors': 16}. Best is trial 5 with value: 0.9725486284289276.\n",
      "[I 2024-02-21 03:05:03,599] Trial 7 finished with value: 0.9745411471321697 and parameters: {'weights': 'uniform', 'n_components': 5, 'n_neighbors': 19}. Best is trial 7 with value: 0.9745411471321697.\n",
      "[I 2024-02-21 03:05:03,815] Trial 8 finished with value: 0.9735461346633416 and parameters: {'weights': 'distance', 'n_components': 7, 'n_neighbors': 24}. Best is trial 7 with value: 0.9745411471321697.\n",
      "[I 2024-02-21 03:05:04,228] Trial 9 finished with value: 0.91614463840399 and parameters: {'weights': 'distance', 'n_components': 30, 'n_neighbors': 94}. Best is trial 7 with value: 0.9745411471321697.\n",
      "[I 2024-02-21 03:05:04,476] Trial 10 finished with value: 0.956074812967581 and parameters: {'weights': 'distance', 'n_components': 11, 'n_neighbors': 67}. Best is trial 7 with value: 0.9745411471321697.\n",
      "[I 2024-02-21 03:05:04,730] Trial 11 finished with value: 0.9435935162094763 and parameters: {'weights': 'uniform', 'n_components': 18, 'n_neighbors': 45}. Best is trial 7 with value: 0.9745411471321697.\n",
      "[I 2024-02-21 03:05:04,996] Trial 12 finished with value: 0.9550760598503741 and parameters: {'weights': 'uniform', 'n_components': 18, 'n_neighbors': 12}. Best is trial 7 with value: 0.9745411471321697.\n",
      "[I 2024-02-21 03:05:05,364] Trial 13 finished with value: 0.9745386533665836 and parameters: {'weights': 'uniform', 'n_components': 3, 'n_neighbors': 12}. Best is trial 7 with value: 0.9745411471321697.\n",
      "[I 2024-02-21 03:05:05,717] Trial 14 finished with value: 0.9715436408977556 and parameters: {'weights': 'uniform', 'n_components': 3, 'n_neighbors': 32}. Best is trial 7 with value: 0.9745411471321697.\n",
      "[I 2024-02-21 03:05:05,983] Trial 15 finished with value: 0.9705448877805486 and parameters: {'weights': 'uniform', 'n_components': 3, 'n_neighbors': 40}. Best is trial 7 with value: 0.9745411471321697.\n",
      "[I 2024-02-21 03:05:06,470] Trial 16 finished with value: 0.8976695760598503 and parameters: {'weights': 'uniform', 'n_components': 49, 'n_neighbors': 42}. Best is trial 7 with value: 0.9745411471321697.\n",
      "[I 2024-02-21 03:05:06,845] Trial 17 finished with value: 0.9086533665835412 and parameters: {'weights': 'uniform', 'n_components': 50, 'n_neighbors': 17}. Best is trial 7 with value: 0.9745411471321697.\n",
      "[I 2024-02-21 03:05:07,134] Trial 18 finished with value: 0.9680523690773069 and parameters: {'weights': 'uniform', 'n_components': 13, 'n_neighbors': 11}. Best is trial 7 with value: 0.9745411471321697.\n",
      "[I 2024-02-21 03:05:07,538] Trial 19 finished with value: 0.9620648379052369 and parameters: {'weights': 'uniform', 'n_components': 13, 'n_neighbors': 22}. Best is trial 7 with value: 0.9745411471321697.\n",
      "[I 2024-02-21 03:05:07,787] Trial 20 finished with value: 0.921137157107232 and parameters: {'weights': 'uniform', 'n_components': 38, 'n_neighbors': 22}. Best is trial 7 with value: 0.9745411471321697.\n",
      "[I 2024-02-21 03:05:08,125] Trial 21 finished with value: 0.9181408977556108 and parameters: {'weights': 'uniform', 'n_components': 36, 'n_neighbors': 35}. Best is trial 7 with value: 0.9745411471321697.\n",
      "[I 2024-02-21 03:05:08,388] Trial 22 finished with value: 0.9705511221945138 and parameters: {'weights': 'distance', 'n_components': 8, 'n_neighbors': 26}. Best is trial 7 with value: 0.9745411471321697.\n",
      "[I 2024-02-21 03:05:08,733] Trial 23 finished with value: 0.9705498753117208 and parameters: {'weights': 'distance', 'n_components': 8, 'n_neighbors': 25}. Best is trial 7 with value: 0.9745411471321697.\n",
      "[I 2024-02-21 03:05:09,031] Trial 24 finished with value: 0.9465910224438904 and parameters: {'weights': 'distance', 'n_components': 21, 'n_neighbors': 20}. Best is trial 7 with value: 0.9745411471321697.\n",
      "[I 2024-02-21 03:05:09,310] Trial 25 finished with value: 0.9455935162094763 and parameters: {'weights': 'distance', 'n_components': 21, 'n_neighbors': 19}. Best is trial 7 with value: 0.9745411471321697.\n",
      "[I 2024-02-21 03:05:09,569] Trial 26 finished with value: 0.9650598503740649 and parameters: {'weights': 'uniform', 'n_components': 10, 'n_neighbors': 49}. Best is trial 7 with value: 0.9745411471321697.\n",
      "[I 2024-02-21 03:05:09,940] Trial 27 finished with value: 0.9675586034912719 and parameters: {'weights': 'uniform', 'n_components': 10, 'n_neighbors': 11}. Best is trial 7 with value: 0.9745411471321697.\n",
      "[I 2024-02-21 03:05:10,184] Trial 28 finished with value: 0.9505847880299252 and parameters: {'weights': 'uniform', 'n_components': 16, 'n_neighbors': 37}. Best is trial 7 with value: 0.9745411471321697.\n",
      "[I 2024-02-21 03:05:10,432] Trial 29 finished with value: 0.9480897755610973 and parameters: {'weights': 'distance', 'n_components': 17, 'n_neighbors': 34}. Best is trial 7 with value: 0.9745411471321697.\n",
      "[I 2024-02-21 03:05:10,699] Trial 30 finished with value: 0.9695511221945138 and parameters: {'weights': 'distance', 'n_components': 6, 'n_neighbors': 65}. Best is trial 7 with value: 0.9745411471321697.\n",
      "[I 2024-02-21 03:05:11,103] Trial 31 finished with value: 0.9690523690773067 and parameters: {'weights': 'distance', 'n_components': 6, 'n_neighbors': 58}. Best is trial 7 with value: 0.9745411471321697.\n",
      "[I 2024-02-21 03:05:11,351] Trial 32 finished with value: 0.9740423940149627 and parameters: {'weights': 'distance', 'n_components': 6, 'n_neighbors': 27}. Best is trial 7 with value: 0.9745411471321697.\n",
      "[I 2024-02-21 03:05:11,600] Trial 33 finished with value: 0.9725423940149627 and parameters: {'weights': 'distance', 'n_components': 3, 'n_neighbors': 29}. Best is trial 7 with value: 0.9745411471321697.\n",
      "[I 2024-02-21 03:05:11,898] Trial 34 finished with value: 0.9725423940149627 and parameters: {'weights': 'distance', 'n_components': 3, 'n_neighbors': 27}. Best is trial 7 with value: 0.9745411471321697.\n",
      "[I 2024-02-21 03:05:12,170] Trial 35 finished with value: 0.9700511221945136 and parameters: {'weights': 'distance', 'n_components': 10, 'n_neighbors': 16}. Best is trial 7 with value: 0.9745411471321697.\n",
      "[I 2024-02-21 03:05:12,480] Trial 36 finished with value: 0.9670598503740649 and parameters: {'weights': 'distance', 'n_components': 11, 'n_neighbors': 17}. Best is trial 7 with value: 0.9745411471321697.\n",
      "[I 2024-02-21 03:05:12,790] Trial 37 finished with value: 0.964564837905237 and parameters: {'weights': 'distance', 'n_components': 7, 'n_neighbors': 74}. Best is trial 7 with value: 0.9745411471321697.\n",
      "[I 2024-02-21 03:05:13,088] Trial 38 finished with value: 0.9376059850374066 and parameters: {'weights': 'distance', 'n_components': 27, 'n_neighbors': 24}. Best is trial 7 with value: 0.9745411471321697.\n",
      "[I 2024-02-21 03:05:13,372] Trial 39 finished with value: 0.9425997506234414 and parameters: {'weights': 'uniform', 'n_components': 15, 'n_neighbors': 86}. Best is trial 7 with value: 0.9745411471321697.\n",
      "[I 2024-02-21 03:05:13,748] Trial 40 finished with value: 0.9460910224438901 and parameters: {'weights': 'uniform', 'n_components': 14, 'n_neighbors': 91}. Best is trial 7 with value: 0.9745411471321697.\n",
      "[I 2024-02-21 03:05:14,043] Trial 41 finished with value: 0.9690561097256858 and parameters: {'weights': 'uniform', 'n_components': 6, 'n_neighbors': 57}. Best is trial 7 with value: 0.9745411471321697.\n",
      "[I 2024-02-21 03:05:14,289] Trial 42 finished with value: 0.9745423940149627 and parameters: {'weights': 'distance', 'n_components': 5, 'n_neighbors': 31}. Best is trial 42 with value: 0.9745423940149627.\n",
      "[I 2024-02-21 03:05:14,584] Trial 43 finished with value: 0.9765411471321694 and parameters: {'weights': 'distance', 'n_components': 5, 'n_neighbors': 30}. Best is trial 43 with value: 0.9765411471321694.\n",
      "[I 2024-02-21 03:05:14,852] Trial 44 finished with value: 0.9725473815461345 and parameters: {'weights': 'distance', 'n_components': 5, 'n_neighbors': 31}. Best is trial 43 with value: 0.9765411471321694.\n",
      "[I 2024-02-21 03:05:15,104] Trial 45 finished with value: 0.9740399002493767 and parameters: {'weights': 'distance', 'n_components': 5, 'n_neighbors': 31}. Best is trial 43 with value: 0.9765411471321694.\n",
      "[I 2024-02-21 03:05:15,459] Trial 46 finished with value: 0.9645598503740649 and parameters: {'weights': 'distance', 'n_components': 9, 'n_neighbors': 47}. Best is trial 43 with value: 0.9765411471321694.\n",
      "[I 2024-02-21 03:05:15,736] Trial 47 finished with value: 0.9635623441396509 and parameters: {'weights': 'distance', 'n_components': 9, 'n_neighbors': 48}. Best is trial 43 with value: 0.9765411471321694.\n",
      "[I 2024-02-21 03:05:15,989] Trial 48 finished with value: 0.9750423940149625 and parameters: {'weights': 'distance', 'n_components': 4, 'n_neighbors': 39}. Best is trial 43 with value: 0.9765411471321694.\n",
      "[I 2024-02-21 03:05:16,254] Trial 49 finished with value: 0.9750423940149625 and parameters: {'weights': 'distance', 'n_components': 4, 'n_neighbors': 41}. Best is trial 43 with value: 0.9765411471321694.\n",
      "[I 2024-02-21 03:05:16,578] Trial 50 finished with value: 0.9296221945137157 and parameters: {'weights': 'uniform', 'n_components': 31, 'n_neighbors': 40}. Best is trial 43 with value: 0.9765411471321694.\n",
      "[I 2024-02-21 03:05:16,957] Trial 51 finished with value: 0.9565735660847879 and parameters: {'weights': 'distance', 'n_components': 12, 'n_neighbors': 42}. Best is trial 43 with value: 0.9765411471321694.\n",
      "[I 2024-02-21 03:05:17,261] Trial 52 finished with value: 0.9720423940149626 and parameters: {'weights': 'distance', 'n_components': 3, 'n_neighbors': 53}. Best is trial 43 with value: 0.9765411471321694.\n",
      "[I 2024-02-21 03:05:17,548] Trial 53 finished with value: 0.9720411471321697 and parameters: {'weights': 'distance', 'n_components': 3, 'n_neighbors': 37}. Best is trial 43 with value: 0.9765411471321694.\n",
      "[I 2024-02-21 03:05:17,874] Trial 54 finished with value: 0.9750423940149625 and parameters: {'weights': 'distance', 'n_components': 4, 'n_neighbors': 38}. Best is trial 43 with value: 0.9765411471321694.\n",
      "[I 2024-02-21 03:05:18,144] Trial 55 finished with value: 0.9750374064837907 and parameters: {'weights': 'distance', 'n_components': 5, 'n_neighbors': 14}. Best is trial 43 with value: 0.9765411471321694.\n",
      "[I 2024-02-21 03:05:18,379] Trial 56 finished with value: 0.9750411471321696 and parameters: {'weights': 'distance', 'n_components': 5, 'n_neighbors': 38}. Best is trial 43 with value: 0.9765411471321694.\n",
      "[I 2024-02-21 03:05:18,767] Trial 57 finished with value: 0.9710473815461347 and parameters: {'weights': 'distance', 'n_components': 8, 'n_neighbors': 41}. Best is trial 43 with value: 0.9765411471321694.\n",
      "[I 2024-02-21 03:05:19,032] Trial 58 finished with value: 0.9645635910224438 and parameters: {'weights': 'distance', 'n_components': 8, 'n_neighbors': 44}. Best is trial 43 with value: 0.9765411471321694.\n",
      "[I 2024-02-21 03:05:19,306] Trial 59 finished with value: 0.9026658354114714 and parameters: {'weights': 'distance', 'n_components': 44, 'n_neighbors': 44}. Best is trial 43 with value: 0.9765411471321694.\n",
      "[I 2024-02-21 03:05:19,624] Trial 60 finished with value: 0.9725448877805487 and parameters: {'weights': 'distance', 'n_components': 5, 'n_neighbors': 51}. Best is trial 43 with value: 0.9765411471321694.\n",
      "[I 2024-02-21 03:05:19,946] Trial 61 finished with value: 0.9575735660847882 and parameters: {'weights': 'distance', 'n_components': 12, 'n_neighbors': 52}. Best is trial 43 with value: 0.9765411471321694.\n",
      "[I 2024-02-21 03:05:20,216] Trial 62 finished with value: 0.9745461346633416 and parameters: {'weights': 'distance', 'n_components': 5, 'n_neighbors': 37}. Best is trial 43 with value: 0.9765411471321694.\n",
      "[I 2024-02-21 03:05:20,447] Trial 63 finished with value: 0.9745436408977556 and parameters: {'weights': 'distance', 'n_components': 5, 'n_neighbors': 37}. Best is trial 43 with value: 0.9765411471321694.\n",
      "[I 2024-02-21 03:05:20,728] Trial 64 finished with value: 0.9755411471321696 and parameters: {'weights': 'distance', 'n_components': 4, 'n_neighbors': 37}. Best is trial 43 with value: 0.9765411471321694.\n",
      "[I 2024-02-21 03:05:21,054] Trial 65 finished with value: 0.9720486284289276 and parameters: {'weights': 'distance', 'n_components': 7, 'n_neighbors': 34}. Best is trial 43 with value: 0.9765411471321694.\n",
      "[I 2024-02-21 03:05:21,323] Trial 66 finished with value: 0.9346084788029924 and parameters: {'weights': 'distance', 'n_components': 24, 'n_neighbors': 34}. Best is trial 43 with value: 0.9765411471321694.\n",
      "[I 2024-02-21 03:05:21,554] Trial 67 finished with value: 0.9750411471321696 and parameters: {'weights': 'distance', 'n_components': 4, 'n_neighbors': 46}. Best is trial 43 with value: 0.9765411471321694.\n",
      "[I 2024-02-21 03:05:21,879] Trial 68 finished with value: 0.9750423940149625 and parameters: {'weights': 'distance', 'n_components': 4, 'n_neighbors': 40}. Best is trial 43 with value: 0.9765411471321694.\n",
      "[I 2024-02-21 03:05:22,186] Trial 69 finished with value: 0.9750411471321696 and parameters: {'weights': 'distance', 'n_components': 4, 'n_neighbors': 46}. Best is trial 43 with value: 0.9765411471321694.\n",
      "[I 2024-02-21 03:05:22,446] Trial 70 finished with value: 0.9660586034912718 and parameters: {'weights': 'distance', 'n_components': 9, 'n_neighbors': 39}. Best is trial 43 with value: 0.9765411471321694.\n",
      "[I 2024-02-21 03:05:22,800] Trial 71 finished with value: 0.9660573566084787 and parameters: {'weights': 'distance', 'n_components': 9, 'n_neighbors': 39}. Best is trial 43 with value: 0.9765411471321694.\n",
      "[I 2024-02-21 03:05:23,120] Trial 72 finished with value: 0.9695548628428927 and parameters: {'weights': 'distance', 'n_components': 7, 'n_neighbors': 60}. Best is trial 43 with value: 0.9765411471321694.\n",
      "[I 2024-02-21 03:05:23,340] Trial 73 finished with value: 0.9670561097256858 and parameters: {'weights': 'distance', 'n_components': 7, 'n_neighbors': 43}. Best is trial 43 with value: 0.9765411471321694.\n",
      "[I 2024-02-21 03:05:23,606] Trial 74 finished with value: 0.9755411471321697 and parameters: {'weights': 'distance', 'n_components': 4, 'n_neighbors': 44}. Best is trial 43 with value: 0.9765411471321694.\n",
      "[I 2024-02-21 03:05:23,978] Trial 75 finished with value: 0.9755399002493765 and parameters: {'weights': 'distance', 'n_components': 4, 'n_neighbors': 49}. Best is trial 43 with value: 0.9765411471321694.\n",
      "[I 2024-02-21 03:05:24,184] Trial 76 finished with value: 0.9725423940149627 and parameters: {'weights': 'distance', 'n_components': 3, 'n_neighbors': 29}. Best is trial 43 with value: 0.9765411471321694.\n",
      "[I 2024-02-21 03:05:24,474] Trial 77 finished with value: 0.9715423940149626 and parameters: {'weights': 'distance', 'n_components': 3, 'n_neighbors': 50}. Best is trial 43 with value: 0.9765411471321694.\n",
      "[I 2024-02-21 03:05:24,813] Trial 78 finished with value: 0.959073566084788 and parameters: {'weights': 'distance', 'n_components': 11, 'n_neighbors': 49}. Best is trial 43 with value: 0.9765411471321694.\n",
      "[I 2024-02-21 03:05:25,098] Trial 79 finished with value: 0.9600685785536159 and parameters: {'weights': 'distance', 'n_components': 11, 'n_neighbors': 54}. Best is trial 43 with value: 0.9765411471321694.\n",
      "[I 2024-02-21 03:05:25,352] Trial 80 finished with value: 0.9730461346633416 and parameters: {'weights': 'distance', 'n_components': 7, 'n_neighbors': 34}. Best is trial 43 with value: 0.9765411471321694.\n",
      "[I 2024-02-21 03:05:25,659] Trial 81 finished with value: 0.9695511221945138 and parameters: {'weights': 'distance', 'n_components': 7, 'n_neighbors': 35}. Best is trial 43 with value: 0.9765411471321694.\n",
      "[I 2024-02-21 03:05:25,876] Trial 82 finished with value: 0.9750423940149625 and parameters: {'weights': 'distance', 'n_components': 4, 'n_neighbors': 36}. Best is trial 43 with value: 0.9765411471321694.\n",
      "[I 2024-02-21 03:05:26,194] Trial 83 finished with value: 0.9750423940149625 and parameters: {'weights': 'distance', 'n_components': 4, 'n_neighbors': 42}. Best is trial 43 with value: 0.9765411471321694.\n",
      "[I 2024-02-21 03:05:26,464] Trial 84 finished with value: 0.9750423940149625 and parameters: {'weights': 'distance', 'n_components': 4, 'n_neighbors': 42}. Best is trial 43 with value: 0.9765411471321694.\n",
      "[I 2024-02-21 03:05:26,731] Trial 85 finished with value: 0.9755411471321697 and parameters: {'weights': 'distance', 'n_components': 4, 'n_neighbors': 45}. Best is trial 43 with value: 0.9765411471321694.\n",
      "[I 2024-02-21 03:05:26,994] Trial 86 finished with value: 0.9740436408977556 and parameters: {'weights': 'distance', 'n_components': 6, 'n_neighbors': 32}. Best is trial 43 with value: 0.9765411471321694.\n",
      "[I 2024-02-21 03:05:27,253] Trial 87 finished with value: 0.9735448877805485 and parameters: {'weights': 'distance', 'n_components': 6, 'n_neighbors': 47}. Best is trial 43 with value: 0.9765411471321694.\n",
      "[I 2024-02-21 03:05:27,614] Trial 88 finished with value: 0.9725473815461345 and parameters: {'weights': 'distance', 'n_components': 6, 'n_neighbors': 46}. Best is trial 43 with value: 0.9765411471321694.\n",
      "[I 2024-02-21 03:05:27,950] Trial 89 finished with value: 0.9645635910224438 and parameters: {'weights': 'distance', 'n_components': 8, 'n_neighbors': 56}. Best is trial 43 with value: 0.9765411471321694.\n",
      "[I 2024-02-21 03:05:28,312] Trial 90 finished with value: 0.9655610972568578 and parameters: {'weights': 'distance', 'n_components': 8, 'n_neighbors': 55}. Best is trial 43 with value: 0.9765411471321694.\n",
      "[I 2024-02-21 03:05:28,589] Trial 91 finished with value: 0.9715423940149626 and parameters: {'weights': 'distance', 'n_components': 3, 'n_neighbors': 44}. Best is trial 43 with value: 0.9765411471321694.\n",
      "[I 2024-02-21 03:05:28,857] Trial 92 finished with value: 0.9755411471321696 and parameters: {'weights': 'distance', 'n_components': 4, 'n_neighbors': 40}. Best is trial 43 with value: 0.9765411471321694.\n",
      "[I 2024-02-21 03:05:29,144] Trial 93 finished with value: 0.9750423940149625 and parameters: {'weights': 'distance', 'n_components': 4, 'n_neighbors': 41}. Best is trial 43 with value: 0.9765411471321694.\n",
      "[I 2024-02-21 03:05:29,380] Trial 94 finished with value: 0.9755411471321696 and parameters: {'weights': 'distance', 'n_components': 4, 'n_neighbors': 40}. Best is trial 43 with value: 0.9765411471321694.\n",
      "[I 2024-02-21 03:05:29,755] Trial 95 finished with value: 0.9171408977556108 and parameters: {'weights': 'distance', 'n_components': 37, 'n_neighbors': 39}. Best is trial 43 with value: 0.9765411471321694.\n",
      "[I 2024-02-21 03:05:30,034] Trial 96 finished with value: 0.9171458852867829 and parameters: {'weights': 'distance', 'n_components': 36, 'n_neighbors': 39}. Best is trial 43 with value: 0.9765411471321694.\n",
      "[I 2024-02-21 03:05:30,284] Trial 97 finished with value: 0.9715486284289276 and parameters: {'weights': 'distance', 'n_components': 6, 'n_neighbors': 28}. Best is trial 43 with value: 0.9765411471321694.\n",
      "[I 2024-02-21 03:05:30,595] Trial 98 finished with value: 0.9730486284289277 and parameters: {'weights': 'distance', 'n_components': 6, 'n_neighbors': 29}. Best is trial 43 with value: 0.9765411471321694.\n",
      "[I 2024-02-21 03:05:30,807] Trial 99 finished with value: 0.9660610972568581 and parameters: {'weights': 'distance', 'n_components': 10, 'n_neighbors': 32}. Best is trial 43 with value: 0.9765411471321694.\n"
     ]
    }
   ],
   "source": [
    "study = optuna.create_study(direction=\"maximize\")\n",
    "study.optimize(objective, n_trials=100, n_jobs=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Accuracy: 0.984\n",
      "Balanced accuracy: 0.981\n"
     ]
    }
   ],
   "source": [
    "model = make_pipeline(\n",
    "    OneHotEncoder(handle_unknown=\"ignore\", sparse_output=False, dtype=np.int8),\n",
    "    PCA(n_components=study.best_params[\"n_components\"]),\n",
    "    KNeighborsClassifier(n_neighbors=study.best_params[\"n_neighbors\"], weights=study.best_params[\"weights\"], n_jobs=4),\n",
    ")\n",
    "\n",
    "model.fit(X_train, y_train)\n",
    "y_test_pred = model.predict(X_test)\n",
    "\n",
    "print(f\"Accuracy: {accuracy_score(y_test, y_test_pred).round(3)}\")\n",
    "print(f\"Balanced accuracy: {balanced_accuracy_score(y_test, y_test_pred).round(3)}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['/home/vscode/.ezancestry/data/models/seldin_superpopulation.pkl']"
      ]
     },
     "execution_count": 23,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Save the model\n",
    "model_path = Path(_models_directory) / f\"seldin_{label}.pkl\"\n",
    "model_path.parent.mkdir(parents=True, exist_ok=True)\n",
    "joblib.dump(model, model_path)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Seldin AISNPs to predict population"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "label = \"population\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {},
   "outputs": [],
   "source": [
    "X_train, X_test, y_train, y_test = train_test_split(\n",
    "    seldin1kg.drop(columns=todrop),\n",
    "    seldin1kg[label],\n",
    "    test_size=0.2,\n",
    "    random_state=42,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {},
   "outputs": [],
   "source": [
    "def objective(trial):\n",
    "    weights = trial.suggest_categorical(\"weights\", [\"uniform\", \"distance\"])\n",
    "    n_components = trial.suggest_int(\"n_components\", 3, 50)\n",
    "    n_neighbors = trial.suggest_int(\"n_neighbors\", 11, 100)\n",
    "\n",
    "    model = make_pipeline(\n",
    "        OneHotEncoder(handle_unknown=\"ignore\", sparse_output=False, dtype=np.int8),\n",
    "        PCA(n_components=n_components),\n",
    "        KNeighborsClassifier(n_neighbors=n_neighbors, weights=weights, n_jobs=4),\n",
    "    )\n",
    "\n",
    "    return cross_val_score(model, X_train, y_train, n_jobs=4, scoring=topk_scorer, cv=StratifiedKFold(5)).mean()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "[I 2024-02-21 03:05:31,351] A new study created in memory with name: no-name-0ac79344-57d6-49a7-ac08-2b614e1c4fb7\n",
      "[I 2024-02-21 03:05:32,088] Trial 0 finished with value: 0.49925062344139653 and parameters: {'weights': 'distance', 'n_components': 33, 'n_neighbors': 83}. Best is trial 0 with value: 0.49925062344139653.\n",
      "[I 2024-02-21 03:05:32,329] Trial 1 finished with value: 0.5496845386533666 and parameters: {'weights': 'uniform', 'n_components': 4, 'n_neighbors': 33}. Best is trial 1 with value: 0.5496845386533666.\n",
      "[I 2024-02-21 03:05:32,645] Trial 2 finished with value: 0.4713017456359102 and parameters: {'weights': 'distance', 'n_components': 42, 'n_neighbors': 85}. Best is trial 1 with value: 0.5496845386533666.\n",
      "[I 2024-02-21 03:05:33,051] Trial 3 finished with value: 0.4793017456359102 and parameters: {'weights': 'uniform', 'n_components': 48, 'n_neighbors': 64}. Best is trial 1 with value: 0.5496845386533666.\n",
      "[I 2024-02-21 03:05:33,294] Trial 4 finished with value: 0.52121072319202 and parameters: {'weights': 'uniform', 'n_components': 25, 'n_neighbors': 37}. Best is trial 1 with value: 0.5496845386533666.\n",
      "[I 2024-02-21 03:05:33,673] Trial 5 finished with value: 0.48429551122194514 and parameters: {'weights': 'uniform', 'n_components': 38, 'n_neighbors': 28}. Best is trial 1 with value: 0.5496845386533666.\n",
      "[I 2024-02-21 03:05:33,952] Trial 6 finished with value: 0.4912743142144638 and parameters: {'weights': 'distance', 'n_components': 34, 'n_neighbors': 35}. Best is trial 1 with value: 0.5496845386533666.\n",
      "[I 2024-02-21 03:05:34,315] Trial 7 finished with value: 0.5441957605985038 and parameters: {'weights': 'uniform', 'n_components': 7, 'n_neighbors': 36}. Best is trial 1 with value: 0.5496845386533666.\n",
      "[I 2024-02-21 03:05:34,625] Trial 8 finished with value: 0.4912556109725686 and parameters: {'weights': 'distance', 'n_components': 47, 'n_neighbors': 41}. Best is trial 1 with value: 0.5496845386533666.\n",
      "[I 2024-02-21 03:05:34,934] Trial 9 finished with value: 0.498785536159601 and parameters: {'weights': 'distance', 'n_components': 35, 'n_neighbors': 56}. Best is trial 1 with value: 0.5496845386533666.\n",
      "[I 2024-02-21 03:05:35,208] Trial 10 finished with value: 0.5107344139650872 and parameters: {'weights': 'distance', 'n_components': 34, 'n_neighbors': 23}. Best is trial 1 with value: 0.5496845386533666.\n",
      "[I 2024-02-21 03:05:35,472] Trial 11 finished with value: 0.5377032418952619 and parameters: {'weights': 'uniform', 'n_components': 4, 'n_neighbors': 11}. Best is trial 1 with value: 0.5496845386533666.\n",
      "[I 2024-02-21 03:05:35,768] Trial 12 finished with value: 0.5576608478802992 and parameters: {'weights': 'uniform', 'n_components': 3, 'n_neighbors': 11}. Best is trial 12 with value: 0.5576608478802992.\n",
      "[I 2024-02-21 03:05:35,983] Trial 13 finished with value: 0.5641508728179552 and parameters: {'weights': 'uniform', 'n_components': 3, 'n_neighbors': 49}. Best is trial 13 with value: 0.5641508728179552.\n",
      "[I 2024-02-21 03:05:36,302] Trial 14 finished with value: 0.5367256857855361 and parameters: {'weights': 'uniform', 'n_components': 13, 'n_neighbors': 11}. Best is trial 13 with value: 0.5641508728179552.\n",
      "[I 2024-02-21 03:05:36,601] Trial 15 finished with value: 0.543206982543641 and parameters: {'weights': 'uniform', 'n_components': 12, 'n_neighbors': 100}. Best is trial 13 with value: 0.5641508728179552.\n",
      "[I 2024-02-21 03:05:36,868] Trial 16 finished with value: 0.5362094763092269 and parameters: {'weights': 'uniform', 'n_components': 16, 'n_neighbors': 71}. Best is trial 13 with value: 0.5641508728179552.\n",
      "[I 2024-02-21 03:05:37,232] Trial 17 finished with value: 0.5237256857855362 and parameters: {'weights': 'uniform', 'n_components': 20, 'n_neighbors': 51}. Best is trial 13 with value: 0.5641508728179552.\n",
      "[I 2024-02-21 03:05:37,476] Trial 18 finished with value: 0.5302281795511222 and parameters: {'weights': 'uniform', 'n_components': 22, 'n_neighbors': 47}. Best is trial 13 with value: 0.5641508728179552.\n",
      "[I 2024-02-21 03:05:37,803] Trial 19 finished with value: 0.5636783042394015 and parameters: {'weights': 'uniform', 'n_components': 9, 'n_neighbors': 50}. Best is trial 13 with value: 0.5641508728179552.\n",
      "[I 2024-02-21 03:05:38,170] Trial 20 finished with value: 0.534209476309227 and parameters: {'weights': 'uniform', 'n_components': 9, 'n_neighbors': 18}. Best is trial 13 with value: 0.5641508728179552.\n",
      "[I 2024-02-21 03:05:38,430] Trial 21 finished with value: 0.5701596009975063 and parameters: {'weights': 'uniform', 'n_components': 10, 'n_neighbors': 62}. Best is trial 21 with value: 0.5701596009975063.\n",
      "[I 2024-02-21 03:05:38,720] Trial 22 finished with value: 0.5581745635910225 and parameters: {'weights': 'uniform', 'n_components': 4, 'n_neighbors': 61}. Best is trial 21 with value: 0.5701596009975063.\n",
      "[I 2024-02-21 03:05:39,040] Trial 23 finished with value: 0.5531807980049875 and parameters: {'weights': 'uniform', 'n_components': 15, 'n_neighbors': 61}. Best is trial 21 with value: 0.5701596009975063.\n",
      "[I 2024-02-21 03:05:39,454] Trial 24 finished with value: 0.5327206982543641 and parameters: {'weights': 'uniform', 'n_components': 17, 'n_neighbors': 68}. Best is trial 21 with value: 0.5701596009975063.\n",
      "[I 2024-02-21 03:05:39,724] Trial 25 finished with value: 0.5541945137157107 and parameters: {'weights': 'uniform', 'n_components': 9, 'n_neighbors': 47}. Best is trial 21 with value: 0.5701596009975063.\n",
      "[I 2024-02-21 03:05:39,949] Trial 26 finished with value: 0.5357169576059851 and parameters: {'weights': 'uniform', 'n_components': 8, 'n_neighbors': 46}. Best is trial 21 with value: 0.5701596009975063.\n",
      "[I 2024-02-21 03:05:40,319] Trial 27 finished with value: 0.5451920199501246 and parameters: {'weights': 'uniform', 'n_components': 10, 'n_neighbors': 76}. Best is trial 21 with value: 0.5701596009975063.\n",
      "[I 2024-02-21 03:05:40,557] Trial 28 finished with value: 0.5052381546134663 and parameters: {'weights': 'uniform', 'n_components': 27, 'n_neighbors': 75}. Best is trial 21 with value: 0.5701596009975063.\n",
      "[I 2024-02-21 03:05:40,829] Trial 29 finished with value: 0.5037730673316708 and parameters: {'weights': 'uniform', 'n_components': 26, 'n_neighbors': 55}. Best is trial 21 with value: 0.5701596009975063.\n",
      "[I 2024-02-21 03:05:41,193] Trial 30 finished with value: 0.5092468827930174 and parameters: {'weights': 'distance', 'n_components': 28, 'n_neighbors': 54}. Best is trial 21 with value: 0.5701596009975063.\n",
      "[I 2024-02-21 03:05:41,486] Trial 31 finished with value: 0.5421907730673317 and parameters: {'weights': 'distance', 'n_components': 19, 'n_neighbors': 85}. Best is trial 21 with value: 0.5701596009975063.\n",
      "[I 2024-02-21 03:05:41,743] Trial 32 finished with value: 0.5367157107231921 and parameters: {'weights': 'uniform', 'n_components': 6, 'n_neighbors': 60}. Best is trial 21 with value: 0.5701596009975063.\n",
      "[I 2024-02-21 03:05:42,013] Trial 33 finished with value: 0.5471807980049875 and parameters: {'weights': 'uniform', 'n_components': 6, 'n_neighbors': 61}. Best is trial 21 with value: 0.5701596009975063.\n",
      "[I 2024-02-21 03:05:42,318] Trial 34 finished with value: 0.576145885286783 and parameters: {'weights': 'uniform', 'n_components': 3, 'n_neighbors': 66}. Best is trial 34 with value: 0.576145885286783.\n",
      "[I 2024-02-21 03:05:42,558] Trial 35 finished with value: 0.5611583541147132 and parameters: {'weights': 'uniform', 'n_components': 3, 'n_neighbors': 44}. Best is trial 34 with value: 0.576145885286783.\n",
      "[I 2024-02-21 03:05:42,864] Trial 36 finished with value: 0.5467032418952619 and parameters: {'weights': 'uniform', 'n_components': 12, 'n_neighbors': 80}. Best is trial 34 with value: 0.576145885286783.\n",
      "[I 2024-02-21 03:05:43,175] Trial 37 finished with value: 0.5427057356608478 and parameters: {'weights': 'uniform', 'n_components': 12, 'n_neighbors': 80}. Best is trial 34 with value: 0.576145885286783.\n",
      "[I 2024-02-21 03:05:43,502] Trial 38 finished with value: 0.5357032418952619 and parameters: {'weights': 'uniform', 'n_components': 11, 'n_neighbors': 67}. Best is trial 34 with value: 0.576145885286783.\n",
      "[I 2024-02-21 03:05:43,783] Trial 39 finished with value: 0.5451870324189526 and parameters: {'weights': 'uniform', 'n_components': 6, 'n_neighbors': 66}. Best is trial 34 with value: 0.576145885286783.\n",
      "[I 2024-02-21 03:05:44,026] Trial 40 finished with value: 0.544705735660848 and parameters: {'weights': 'uniform', 'n_components': 6, 'n_neighbors': 93}. Best is trial 34 with value: 0.576145885286783.\n",
      "[I 2024-02-21 03:05:44,390] Trial 41 finished with value: 0.5347219451371571 and parameters: {'weights': 'distance', 'n_components': 14, 'n_neighbors': 31}. Best is trial 34 with value: 0.576145885286783.\n",
      "[I 2024-02-21 03:05:44,641] Trial 42 finished with value: 0.5581670822942644 and parameters: {'weights': 'uniform', 'n_components': 3, 'n_neighbors': 40}. Best is trial 34 with value: 0.576145885286783.\n",
      "[I 2024-02-21 03:05:44,974] Trial 43 finished with value: 0.5611670822942644 and parameters: {'weights': 'uniform', 'n_components': 3, 'n_neighbors': 41}. Best is trial 34 with value: 0.576145885286783.\n",
      "[I 2024-02-21 03:05:45,246] Trial 44 finished with value: 0.5651533665835411 and parameters: {'weights': 'uniform', 'n_components': 3, 'n_neighbors': 50}. Best is trial 34 with value: 0.576145885286783.\n",
      "[I 2024-02-21 03:05:45,522] Trial 45 finished with value: 0.5631683291770574 and parameters: {'weights': 'uniform', 'n_components': 8, 'n_neighbors': 41}. Best is trial 34 with value: 0.576145885286783.\n",
      "[I 2024-02-21 03:05:45,839] Trial 46 finished with value: 0.5621745635910225 and parameters: {'weights': 'uniform', 'n_components': 8, 'n_neighbors': 51}. Best is trial 34 with value: 0.576145885286783.\n",
      "[I 2024-02-21 03:05:46,143] Trial 47 finished with value: 0.5037618453865337 and parameters: {'weights': 'uniform', 'n_components': 30, 'n_neighbors': 51}. Best is trial 34 with value: 0.576145885286783.\n",
      "[I 2024-02-21 03:05:46,420] Trial 48 finished with value: 0.4932693266832917 and parameters: {'weights': 'distance', 'n_components': 42, 'n_neighbors': 52}. Best is trial 34 with value: 0.576145885286783.\n",
      "[I 2024-02-21 03:05:46,722] Trial 49 finished with value: 0.4967680798004987 and parameters: {'weights': 'distance', 'n_components': 43, 'n_neighbors': 55}. Best is trial 34 with value: 0.576145885286783.\n",
      "[I 2024-02-21 03:05:47,023] Trial 50 finished with value: 0.5561807980049875 and parameters: {'weights': 'uniform', 'n_components': 5, 'n_neighbors': 72}. Best is trial 34 with value: 0.576145885286783.\n",
      "[I 2024-02-21 03:05:47,302] Trial 51 finished with value: 0.5541832917705735 and parameters: {'weights': 'uniform', 'n_components': 5, 'n_neighbors': 58}. Best is trial 34 with value: 0.576145885286783.\n",
      "[I 2024-02-21 03:05:47,587] Trial 52 finished with value: 0.5601720698254364 and parameters: {'weights': 'uniform', 'n_components': 8, 'n_neighbors': 38}. Best is trial 34 with value: 0.576145885286783.\n",
      "[I 2024-02-21 03:05:47,886] Trial 53 finished with value: 0.5471982543640899 and parameters: {'weights': 'uniform', 'n_components': 8, 'n_neighbors': 32}. Best is trial 34 with value: 0.576145885286783.\n",
      "[I 2024-02-21 03:05:48,186] Trial 54 finished with value: 0.5591745635910225 and parameters: {'weights': 'uniform', 'n_components': 7, 'n_neighbors': 32}. Best is trial 34 with value: 0.576145885286783.\n",
      "[I 2024-02-21 03:05:48,530] Trial 55 finished with value: 0.5546832917705735 and parameters: {'weights': 'uniform', 'n_components': 10, 'n_neighbors': 43}. Best is trial 34 with value: 0.576145885286783.\n",
      "[I 2024-02-21 03:05:48,822] Trial 56 finished with value: 0.5426970074812967 and parameters: {'weights': 'uniform', 'n_components': 10, 'n_neighbors': 48}. Best is trial 34 with value: 0.576145885286783.\n",
      "[I 2024-02-21 03:05:49,126] Trial 57 finished with value: 0.540703241895262 and parameters: {'weights': 'uniform', 'n_components': 18, 'n_neighbors': 64}. Best is trial 34 with value: 0.576145885286783.\n",
      "[I 2024-02-21 03:05:49,367] Trial 58 finished with value: 0.5067331670822941 and parameters: {'weights': 'uniform', 'n_components': 23, 'n_neighbors': 25}. Best is trial 34 with value: 0.576145885286783.\n",
      "[I 2024-02-21 03:05:49,724] Trial 59 finished with value: 0.5546683291770573 and parameters: {'weights': 'uniform', 'n_components': 5, 'n_neighbors': 24}. Best is trial 34 with value: 0.576145885286783.\n",
      "[I 2024-02-21 03:05:49,999] Trial 60 finished with value: 0.5601857855361596 and parameters: {'weights': 'uniform', 'n_components': 14, 'n_neighbors': 58}. Best is trial 34 with value: 0.576145885286783.\n",
      "[I 2024-02-21 03:05:50,352] Trial 61 finished with value: 0.5486882793017457 and parameters: {'weights': 'uniform', 'n_components': 14, 'n_neighbors': 58}. Best is trial 34 with value: 0.576145885286783.\n",
      "[I 2024-02-21 03:05:50,669] Trial 62 finished with value: 0.5556795511221945 and parameters: {'weights': 'uniform', 'n_components': 8, 'n_neighbors': 52}. Best is trial 34 with value: 0.576145885286783.\n",
      "[I 2024-02-21 03:05:50,972] Trial 63 finished with value: 0.5571683291770573 and parameters: {'weights': 'uniform', 'n_components': 8, 'n_neighbors': 51}. Best is trial 34 with value: 0.576145885286783.\n",
      "[I 2024-02-21 03:05:51,235] Trial 64 finished with value: 0.5566708229426434 and parameters: {'weights': 'uniform', 'n_components': 4, 'n_neighbors': 48}. Best is trial 34 with value: 0.576145885286783.\n",
      "[I 2024-02-21 03:05:51,492] Trial 65 finished with value: 0.563148379052369 and parameters: {'weights': 'uniform', 'n_components': 3, 'n_neighbors': 48}. Best is trial 34 with value: 0.576145885286783.\n",
      "[I 2024-02-21 03:05:51,754] Trial 66 finished with value: 0.5436970074812967 and parameters: {'weights': 'uniform', 'n_components': 16, 'n_neighbors': 43}. Best is trial 34 with value: 0.576145885286783.\n",
      "[I 2024-02-21 03:05:52,069] Trial 67 finished with value: 0.5596608478802991 and parameters: {'weights': 'uniform', 'n_components': 3, 'n_neighbors': 45}. Best is trial 34 with value: 0.576145885286783.\n",
      "[I 2024-02-21 03:05:52,287] Trial 68 finished with value: 0.5576670822942644 and parameters: {'weights': 'uniform', 'n_components': 3, 'n_neighbors': 35}. Best is trial 34 with value: 0.576145885286783.\n",
      "[I 2024-02-21 03:05:52,605] Trial 69 finished with value: 0.5591683291770574 and parameters: {'weights': 'uniform', 'n_components': 5, 'n_neighbors': 36}. Best is trial 34 with value: 0.576145885286783.\n",
      "[I 2024-02-21 03:05:52,894] Trial 70 finished with value: 0.5636658354114713 and parameters: {'weights': 'distance', 'n_components': 5, 'n_neighbors': 38}. Best is trial 34 with value: 0.576145885286783.\n",
      "[I 2024-02-21 03:05:53,199] Trial 71 finished with value: 0.5631720698254364 and parameters: {'weights': 'distance', 'n_components': 11, 'n_neighbors': 64}. Best is trial 34 with value: 0.576145885286783.\n",
      "[I 2024-02-21 03:05:53,436] Trial 72 finished with value: 0.5546645885286783 and parameters: {'weights': 'distance', 'n_components': 7, 'n_neighbors': 64}. Best is trial 34 with value: 0.576145885286783.\n",
      "[I 2024-02-21 03:05:53,783] Trial 73 finished with value: 0.5486882793017457 and parameters: {'weights': 'distance', 'n_components': 11, 'n_neighbors': 63}. Best is trial 34 with value: 0.576145885286783.\n",
      "[I 2024-02-21 03:05:54,074] Trial 74 finished with value: 0.5441995012468828 and parameters: {'weights': 'distance', 'n_components': 11, 'n_neighbors': 71}. Best is trial 34 with value: 0.576145885286783.\n",
      "[I 2024-02-21 03:05:54,348] Trial 75 finished with value: 0.5437119700748129 and parameters: {'weights': 'distance', 'n_components': 10, 'n_neighbors': 68}. Best is trial 34 with value: 0.576145885286783.\n",
      "[I 2024-02-21 03:05:54,723] Trial 76 finished with value: 0.5526920199501247 and parameters: {'weights': 'distance', 'n_components': 9, 'n_neighbors': 38}. Best is trial 34 with value: 0.576145885286783.\n",
      "[I 2024-02-21 03:05:54,980] Trial 77 finished with value: 0.5376970074812968 and parameters: {'weights': 'distance', 'n_components': 6, 'n_neighbors': 38}. Best is trial 34 with value: 0.576145885286783.\n",
      "[I 2024-02-21 03:05:55,249] Trial 78 finished with value: 0.5576496259351621 and parameters: {'weights': 'distance', 'n_components': 6, 'n_neighbors': 29}. Best is trial 34 with value: 0.576145885286783.\n",
      "[I 2024-02-21 03:05:55,575] Trial 79 finished with value: 0.5491820448877806 and parameters: {'weights': 'distance', 'n_components': 13, 'n_neighbors': 54}. Best is trial 34 with value: 0.576145885286783.\n",
      "[I 2024-02-21 03:05:55,864] Trial 80 finished with value: 0.5496882793017456 and parameters: {'weights': 'distance', 'n_components': 12, 'n_neighbors': 75}. Best is trial 34 with value: 0.576145885286783.\n",
      "[I 2024-02-21 03:05:56,234] Trial 81 finished with value: 0.5601670822942644 and parameters: {'weights': 'distance', 'n_components': 4, 'n_neighbors': 57}. Best is trial 34 with value: 0.576145885286783.\n",
      "[I 2024-02-21 03:05:56,513] Trial 82 finished with value: 0.5606645885286783 and parameters: {'weights': 'distance', 'n_components': 4, 'n_neighbors': 41}. Best is trial 34 with value: 0.576145885286783.\n",
      "[I 2024-02-21 03:05:56,849] Trial 83 finished with value: 0.5551758104738155 and parameters: {'weights': 'uniform', 'n_components': 4, 'n_neighbors': 48}. Best is trial 34 with value: 0.576145885286783.\n",
      "[I 2024-02-21 03:05:57,057] Trial 84 finished with value: 0.5561633416458853 and parameters: {'weights': 'uniform', 'n_components': 7, 'n_neighbors': 47}. Best is trial 34 with value: 0.576145885286783.\n",
      "[I 2024-02-21 03:05:57,411] Trial 85 finished with value: 0.45331920199501247 and parameters: {'weights': 'uniform', 'n_components': 50, 'n_neighbors': 49}. Best is trial 34 with value: 0.576145885286783.\n",
      "[I 2024-02-21 03:05:57,712] Trial 86 finished with value: 0.5661645885286782 and parameters: {'weights': 'uniform', 'n_components': 5, 'n_neighbors': 49}. Best is trial 34 with value: 0.576145885286783.\n",
      "[I 2024-02-21 03:05:57,937] Trial 87 finished with value: 0.5651758104738155 and parameters: {'weights': 'uniform', 'n_components': 5, 'n_neighbors': 61}. Best is trial 34 with value: 0.576145885286783.\n",
      "[I 2024-02-21 03:05:58,233] Trial 88 finished with value: 0.5656608478802992 and parameters: {'weights': 'uniform', 'n_components': 5, 'n_neighbors': 45}. Best is trial 34 with value: 0.576145885286783.\n",
      "[I 2024-02-21 03:05:58,530] Trial 89 finished with value: 0.5646695760598504 and parameters: {'weights': 'uniform', 'n_components': 5, 'n_neighbors': 60}. Best is trial 34 with value: 0.576145885286783.\n",
      "[I 2024-02-21 03:05:58,909] Trial 90 finished with value: 0.5576733167082295 and parameters: {'weights': 'uniform', 'n_components': 5, 'n_neighbors': 60}. Best is trial 34 with value: 0.576145885286783.\n",
      "[I 2024-02-21 03:05:59,223] Trial 91 finished with value: 0.5471907730673318 and parameters: {'weights': 'uniform', 'n_components': 6, 'n_neighbors': 61}. Best is trial 34 with value: 0.576145885286783.\n",
      "[I 2024-02-21 03:05:59,499] Trial 92 finished with value: 0.5501970074812967 and parameters: {'weights': 'uniform', 'n_components': 7, 'n_neighbors': 54}. Best is trial 34 with value: 0.576145885286783.\n",
      "[I 2024-02-21 03:05:59,788] Trial 93 finished with value: 0.5691633416458852 and parameters: {'weights': 'uniform', 'n_components': 5, 'n_neighbors': 53}. Best is trial 34 with value: 0.576145885286783.\n",
      "[I 2024-02-21 03:06:00,070] Trial 94 finished with value: 0.570147132169576 and parameters: {'weights': 'uniform', 'n_components': 5, 'n_neighbors': 53}. Best is trial 34 with value: 0.576145885286783.\n",
      "[I 2024-02-21 03:06:00,357] Trial 95 finished with value: 0.5531857855361596 and parameters: {'weights': 'uniform', 'n_components': 9, 'n_neighbors': 56}. Best is trial 34 with value: 0.576145885286783.\n",
      "[I 2024-02-21 03:06:00,629] Trial 96 finished with value: 0.5641583541147133 and parameters: {'weights': 'uniform', 'n_components': 3, 'n_neighbors': 55}. Best is trial 34 with value: 0.576145885286783.\n",
      "[I 2024-02-21 03:06:00,907] Trial 97 finished with value: 0.575145885286783 and parameters: {'weights': 'uniform', 'n_components': 3, 'n_neighbors': 70}. Best is trial 34 with value: 0.576145885286783.\n",
      "[I 2024-02-21 03:06:01,185] Trial 98 finished with value: 0.5556857855361595 and parameters: {'weights': 'uniform', 'n_components': 4, 'n_neighbors': 53}. Best is trial 34 with value: 0.576145885286783.\n",
      "[I 2024-02-21 03:06:01,425] Trial 99 finished with value: 0.5651708229426434 and parameters: {'weights': 'uniform', 'n_components': 5, 'n_neighbors': 69}. Best is trial 34 with value: 0.576145885286783.\n"
     ]
    }
   ],
   "source": [
    "study = optuna.create_study(direction=\"maximize\")\n",
    "study.optimize(objective, n_trials=100, n_jobs=2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Accuracy: 0.345\n",
      "Balanced accuracy: 0.353\n",
      "Top-2 accuracy: 0.595\n",
      "Top-3 accuracy: 0.752\n"
     ]
    }
   ],
   "source": [
    "model = make_pipeline(\n",
    "    OneHotEncoder(handle_unknown=\"ignore\", sparse_output=False, dtype=np.int8),\n",
    "    PCA(n_components=study.best_params[\"n_components\"]),\n",
    "    KNeighborsClassifier(n_neighbors=study.best_params[\"n_neighbors\"], weights=study.best_params[\"weights\"], n_jobs=4),\n",
    ")\n",
    "\n",
    "model.fit(X_train, y_train)\n",
    "y_test_pred = model.predict(X_test)\n",
    "\n",
    "print(f\"Accuracy: {accuracy_score(y_test, y_test_pred).round(3)}\")\n",
    "print(f\"Balanced accuracy: {balanced_accuracy_score(y_test, y_test_pred).round(3)}\")\n",
    "print(f\"Top-2 accuracy: {top_k_accuracy_score(y_test, model.predict_proba(X_test), k=2).round(3)}\")\n",
    "print(f\"Top-3 accuracy: {top_k_accuracy_score(y_test, model.predict_proba(X_test), k=3).round(3)}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 29,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "['/home/vscode/.ezancestry/data/models/seldin_population.pkl']"
      ]
     },
     "execution_count": 29,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Save the model\n",
    "model_path = Path(_models_directory) / f\"seldin_{label}.pkl\"\n",
    "model_path.parent.mkdir(parents=True, exist_ok=True)\n",
    "joblib.dump(model, model_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "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.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}