demo.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from bme_jax.main import run\n",
    "from utils.data import load_yaml_config"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "{'fasta_path': 'primates.fa',\n",
      " 'learning_rate': 1.5,\n",
      " 'max_iter': 100,\n",
      " 'n_shuffles': 100,\n",
      " 'n_steps': 5000,\n",
      " 'optimizer': 'adafactor',\n",
      " 'repo_path': './',\n",
      " 'rooted': False,\n",
      " 'substitution_model': 'F81',\n",
      " 'tol': 1e-10}\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "232 sites; 14 taxa.\n",
      "root at Mouse\n",
      "with score 1.4160958528518677\n",
      "Checking our and the discrete BME objective match:\n",
      "newick2v+reorder score: 1.5525645017623901\n",
      "fastme score: 1.5525643825531006\n",
      "diff: 1.1920928955078125e-07\n",
      "root at Mouse\n",
      "with score 1.4163050651550293\n",
      "Checking our and the discrete BME objective match:\n",
      "newick2v+reorder score: 1.5526909828186035\n",
      "nj score: 1.5526909828186035\n",
      "diff: 0.0\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 100/100 [00:16<00:00,  6.05it/s, \u001b[95m current best BME objective is =1.5525641]\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "({'D': Array([[0.        , 1.6685598 , 1.0703509 , 1.1237773 , 1.239218  ,\n",
       "          1.4387202 , 1.2580895 , 1.0462348 , 1.5057468 , 1.5333744 ,\n",
       "          1.1781396 , 1.2970101 , 1.3125371 , 1.2625202 ],\n",
       "         [1.6685598 , 0.        , 0.926419  , 0.7816145 , 1.2962402 ,\n",
       "          0.790979  , 1.3169839 , 0.7233227 , 0.6719033 , 0.6480957 ,\n",
       "          0.57571536, 0.77466655, 0.5530452 , 0.6252031 ],\n",
       "         [1.0703509 , 0.926419  , 0.        , 0.9535519 , 0.9275206 ,\n",
       "          1.075167  , 1.0046617 , 0.9647009 , 0.8805094 , 1.0375417 ,\n",
       "          1.0681767 , 1.0028745 , 1.1942139 , 1.2415973 ],\n",
       "         [1.1237773 , 0.7816145 , 0.9535519 , 0.        , 0.9751094 ,\n",
       "          0.93076867, 1.3280932 , 0.98410845, 0.8880245 , 1.1174572 ,\n",
       "          0.7647288 , 1.0347162 , 0.9507322 , 0.9710039 ],\n",
       "         [1.239218  , 1.2962402 , 0.9275206 , 0.9751094 , 0.        ,\n",
       "          1.2532077 , 0.89019126, 1.1322598 , 1.3040729 , 1.1818545 ,\n",
       "          1.0140709 , 1.0510988 , 1.2937002 , 1.7910193 ],\n",
       "         [1.4387202 , 0.790979  , 1.075167  , 0.93076867, 1.2532077 ,\n",
       "          0.        , 1.1208087 , 0.38534912, 0.792358  , 0.2802113 ,\n",
       "          0.92291445, 0.21379572, 0.7568428 , 0.9705667 ],\n",
       "         [1.2580895 , 1.3169839 , 1.0046617 , 1.3280932 , 0.89019126,\n",
       "          1.1208087 , 0.        , 1.148538  , 1.1585386 , 1.0469867 ,\n",
       "          1.1678705 , 1.0089942 , 1.5832978 , 1.431889  ],\n",
       "         [1.0462348 , 0.7233227 , 0.9647009 , 0.98410845, 1.1322598 ,\n",
       "          0.38534912, 1.148538  , 0.        , 0.6377546 , 0.40751806,\n",
       "          0.7441745 , 0.40539545, 0.87659746, 0.82588875],\n",
       "         [1.5057468 , 0.6719033 , 0.8805094 , 0.8880245 , 1.3040729 ,\n",
       "          0.792358  , 1.1585386 , 0.6377546 , 0.        , 0.67840165,\n",
       "          0.43547598, 0.6871426 , 0.45481378, 0.5048454 ],\n",
       "         [1.5333744 , 0.6480957 , 1.0375417 , 1.1174572 , 1.1818545 ,\n",
       "          0.2802113 , 1.0469867 , 0.40751806, 0.67840165, 0.        ,\n",
       "          0.67488223, 0.10338926, 0.66007954, 0.89852434],\n",
       "         [1.1781396 , 0.57571536, 1.0681767 , 0.7647288 , 1.0140709 ,\n",
       "          0.92291445, 1.1678705 , 0.7441745 , 0.43547598, 0.67488223,\n",
       "          0.        , 0.6718336 , 0.3155888 , 0.36099344],\n",
       "         [1.2970101 , 0.77466655, 1.0028745 , 1.0347162 , 1.0510988 ,\n",
       "          0.21379572, 1.0089942 , 0.40539545, 0.6871426 , 0.10338926,\n",
       "          0.6718336 , 0.        , 0.6536793 , 0.7815904 ],\n",
       "         [1.3125371 , 0.5530452 , 1.1942139 , 0.9507322 , 1.2937002 ,\n",
       "          0.7568428 , 1.5832978 , 0.87659746, 0.45481378, 0.66007954,\n",
       "          0.3155888 , 0.6536793 , 0.        , 0.2834306 ],\n",
       "         [1.2625202 , 0.6252031 , 1.2415973 , 0.9710039 , 1.7910193 ,\n",
       "          0.9705667 , 1.431889  , 0.82588875, 0.5048454 , 0.89852434,\n",
       "          0.36099344, 0.7815904 , 0.2834306 , 0.        ]], dtype=float32),\n",
       "  'int2taxa': {0: 'Mouse',\n",
       "   3: 'SquirMonk',\n",
       "   2: 'Bovine',\n",
       "   4: 'Lemur',\n",
       "   6: 'Tarsier',\n",
       "   7: 'BarbMacaq',\n",
       "   1: 'Gibbon',\n",
       "   5: 'CrabEMac',\n",
       "   8: 'Orang',\n",
       "   9: 'RhesusMac',\n",
       "   11: 'JpnMacaq',\n",
       "   10: 'Gorilla',\n",
       "   13: 'Chimp',\n",
       "   12: 'Human'},\n",
       "  'score': Array(1.5525641, dtype=float32),\n",
       "  'W': Array([[7.6257893e-15, 0.0000000e+00, 0.0000000e+00, 0.0000000e+00,\n",
       "          0.0000000e+00, 0.0000000e+00, 0.0000000e+00, 0.0000000e+00,\n",
       "          0.0000000e+00, 0.0000000e+00, 0.0000000e+00, 0.0000000e+00,\n",
       "          0.0000000e+00],\n",
       "         [7.9246192e-15, 1.0000000e+00, 0.0000000e+00, 0.0000000e+00,\n",
       "          0.0000000e+00, 0.0000000e+00, 0.0000000e+00, 0.0000000e+00,\n",
       "          0.0000000e+00, 0.0000000e+00, 0.0000000e+00, 0.0000000e+00,\n",
       "          0.0000000e+00],\n",
       "         [3.4027667e-06, 9.9998444e-01, 1.2147805e-05, 0.0000000e+00,\n",
       "          0.0000000e+00, 0.0000000e+00, 0.0000000e+00, 0.0000000e+00,\n",
       "          0.0000000e+00, 0.0000000e+00, 0.0000000e+00, 0.0000000e+00,\n",
       "          0.0000000e+00],\n",
       "         [4.3005513e-13, 2.2122440e-06, 9.9999636e-01, 1.4611186e-06,\n",
       "          0.0000000e+00, 0.0000000e+00, 0.0000000e+00, 0.0000000e+00,\n",
       "          0.0000000e+00, 0.0000000e+00, 0.0000000e+00, 0.0000000e+00,\n",
       "          0.0000000e+00],\n",
       "         [2.3078787e-06, 9.9998558e-01, 4.1257599e-06, 3.2562434e-06,\n",
       "          4.7104490e-06, 0.0000000e+00, 0.0000000e+00, 0.0000000e+00,\n",
       "          0.0000000e+00, 0.0000000e+00, 0.0000000e+00, 0.0000000e+00,\n",
       "          0.0000000e+00],\n",
       "         [2.9129851e-06, 1.4033242e-06, 2.2369452e-06, 2.2272425e-06,\n",
       "          9.9998909e-01, 2.1503372e-06, 0.0000000e+00, 0.0000000e+00,\n",
       "          0.0000000e+00, 0.0000000e+00, 0.0000000e+00, 0.0000000e+00,\n",
       "          0.0000000e+00],\n",
       "         [2.0479199e-06, 5.2891960e-06, 1.5183400e-06, 1.7045372e-06,\n",
       "          1.4224125e-06, 9.9998683e-01, 1.2330833e-06, 0.0000000e+00,\n",
       "          0.0000000e+00, 0.0000000e+00, 0.0000000e+00, 0.0000000e+00,\n",
       "          0.0000000e+00],\n",
       "         [1.3842324e-06, 9.9998361e-01, 1.7704069e-06, 1.8548440e-06,\n",
       "          1.4026143e-06, 4.5203806e-06, 7.8525204e-07, 4.7248186e-06,\n",
       "          0.0000000e+00, 0.0000000e+00, 0.0000000e+00, 0.0000000e+00,\n",
       "          0.0000000e+00],\n",
       "         [1.2462997e-06, 2.2391523e-06, 1.3638295e-06, 1.7058049e-06,\n",
       "          1.1965017e-06, 9.9998403e-01, 7.6573662e-07, 5.8913297e-06,\n",
       "          1.6069879e-06, 0.0000000e+00, 0.0000000e+00, 0.0000000e+00,\n",
       "          0.0000000e+00],\n",
       "         [1.1456499e-06, 2.0513526e-05, 1.2322088e-06, 2.0639354e-06,\n",
       "          5.2657902e-07, 9.6119095e-07, 8.4684706e-07, 8.7313157e-07,\n",
       "          9.9997097e-01, 8.9289227e-07, 0.0000000e+00, 0.0000000e+00,\n",
       "          0.0000000e+00],\n",
       "         [9.3636157e-07, 7.7212525e-07, 1.3889806e-06, 1.7150280e-06,\n",
       "          1.0631286e-06, 3.6856869e-05, 8.7927083e-07, 3.5750493e-06,\n",
       "          5.4188752e-07, 9.9995148e-01, 8.0018020e-07, 0.0000000e+00,\n",
       "          0.0000000e+00],\n",
       "         [6.0299863e-07, 8.8494871e-06, 1.1328483e-06, 4.8331657e-07,\n",
       "          8.2663490e-07, 1.0540508e-06, 1.3595572e-06, 1.1697263e-06,\n",
       "          3.1311643e-05, 1.0862691e-06, 9.9995095e-01, 1.2040060e-06,\n",
       "          0.0000000e+00],\n",
       "         [1.4301572e-06, 3.5074438e-06, 9.9235240e-07, 1.2209908e-06,\n",
       "          1.4008476e-06, 7.7147035e-07, 8.4647996e-07, 6.6314198e-07,\n",
       "          8.7431381e-06, 6.3678635e-07, 3.2549870e-05, 5.4374408e-07,\n",
       "          9.9994671e-01]], dtype=float32),\n",
       "  'scores': [Array(1.581173, dtype=float32),\n",
       "   Array(1.581173, dtype=float32),\n",
       "   Array(1.581173, dtype=float32),\n",
       "   Array(1.5630136, dtype=float32),\n",
       "   Array(1.5630136, dtype=float32),\n",
       "   Array(1.5525644, dtype=float32),\n",
       "   Array(1.5525644, dtype=float32),\n",
       "   Array(1.5525644, dtype=float32),\n",
       "   Array(1.5525644, dtype=float32),\n",
       "   Array(1.5525644, dtype=float32),\n",
       "   Array(1.5525644, dtype=float32),\n",
       "   Array(1.5525644, dtype=float32),\n",
       "   Array(1.5525644, dtype=float32),\n",
       "   Array(1.5525644, dtype=float32),\n",
       "   Array(1.5525644, dtype=float32),\n",
       "   Array(1.5525644, dtype=float32),\n",
       "   Array(1.5525644, dtype=float32),\n",
       "   Array(1.5525644, dtype=float32),\n",
       "   Array(1.5525644, dtype=float32),\n",
       "   Array(1.5525644, dtype=float32),\n",
       "   Array(1.5525644, dtype=float32),\n",
       "   Array(1.5525644, dtype=float32),\n",
       "   Array(1.5525644, dtype=float32),\n",
       "   Array(1.5525644, dtype=float32),\n",
       "   Array(1.5525644, dtype=float32),\n",
       "   Array(1.5525644, dtype=float32),\n",
       "   Array(1.5525644, dtype=float32),\n",
       "   Array(1.5525644, dtype=float32),\n",
       "   Array(1.5525644, dtype=float32),\n",
       "   Array(1.5525644, dtype=float32),\n",
       "   Array(1.5525644, dtype=float32),\n",
       "   Array(1.5525644, dtype=float32),\n",
       "   Array(1.5525644, dtype=float32),\n",
       "   Array(1.5525644, dtype=float32),\n",
       "   Array(1.5525644, dtype=float32),\n",
       "   Array(1.5525644, dtype=float32),\n",
       "   Array(1.5525644, dtype=float32),\n",
       "   Array(1.5525644, dtype=float32),\n",
       "   Array(1.5525644, dtype=float32),\n",
       "   Array(1.5525644, dtype=float32),\n",
       "   Array(1.5525644, dtype=float32),\n",
       "   Array(1.5525644, dtype=float32),\n",
       "   Array(1.5525644, dtype=float32),\n",
       "   Array(1.5525644, dtype=float32),\n",
       "   Array(1.5525644, dtype=float32),\n",
       "   Array(1.5525644, dtype=float32),\n",
       "   Array(1.5525644, dtype=float32),\n",
       "   Array(1.5525644, dtype=float32),\n",
       "   Array(1.5525644, dtype=float32),\n",
       "   Array(1.5525644, dtype=float32),\n",
       "   Array(1.5525644, dtype=float32),\n",
       "   Array(1.5525644, dtype=float32),\n",
       "   Array(1.5525644, dtype=float32),\n",
       "   Array(1.5525644, dtype=float32),\n",
       "   Array(1.5525644, dtype=float32),\n",
       "   Array(1.5525644, dtype=float32),\n",
       "   Array(1.5525644, dtype=float32),\n",
       "   Array(1.5525641, dtype=float32),\n",
       "   Array(1.5525641, dtype=float32),\n",
       "   Array(1.5525641, dtype=float32),\n",
       "   Array(1.5525641, dtype=float32),\n",
       "   Array(1.5525641, dtype=float32),\n",
       "   Array(1.5525641, dtype=float32),\n",
       "   Array(1.5525641, dtype=float32),\n",
       "   Array(1.5525641, dtype=float32),\n",
       "   Array(1.5525641, dtype=float32),\n",
       "   Array(1.5525641, dtype=float32),\n",
       "   Array(1.5525641, dtype=float32),\n",
       "   Array(1.5525641, dtype=float32),\n",
       "   Array(1.5525641, dtype=float32),\n",
       "   Array(1.5525641, dtype=float32),\n",
       "   Array(1.5525641, dtype=float32),\n",
       "   Array(1.5525641, dtype=float32),\n",
       "   Array(1.5525641, dtype=float32),\n",
       "   Array(1.5525641, dtype=float32),\n",
       "   Array(1.5525641, dtype=float32),\n",
       "   Array(1.5525641, dtype=float32),\n",
       "   Array(1.5525641, dtype=float32),\n",
       "   Array(1.5525641, dtype=float32),\n",
       "   Array(1.5525641, dtype=float32),\n",
       "   Array(1.5525641, dtype=float32),\n",
       "   Array(1.5525641, dtype=float32),\n",
       "   Array(1.5525641, dtype=float32),\n",
       "   Array(1.5525641, dtype=float32),\n",
       "   Array(1.5525641, dtype=float32),\n",
       "   Array(1.5525641, dtype=float32),\n",
       "   Array(1.5525641, dtype=float32),\n",
       "   Array(1.5525641, dtype=float32),\n",
       "   Array(1.5525641, dtype=float32),\n",
       "   Array(1.5525641, dtype=float32),\n",
       "   Array(1.5525641, dtype=float32),\n",
       "   Array(1.5525641, dtype=float32),\n",
       "   Array(1.5525641, dtype=float32),\n",
       "   Array(1.5525641, dtype=float32),\n",
       "   Array(1.5525641, dtype=float32),\n",
       "   Array(1.5525641, dtype=float32),\n",
       "   Array(1.5525641, dtype=float32),\n",
       "   Array(1.5525641, dtype=float32),\n",
       "   Array(1.5525641, dtype=float32),\n",
       "   Array(1.5525641, dtype=float32)]},\n",
       " substitution_model            F81\n",
       " optimizer               adafactor\n",
       " learning_rate                 1.5\n",
       " tree_loss_nj             1.552691\n",
       " tree_loss_fastme        1.5525644\n",
       " step_loss_nj             1.552691\n",
       " step_loss_fastme        1.5525646\n",
       " step_loss_opt_random    1.5525641\n",
       " rf_nj_opt                       2\n",
       " rf_fastme_opt                   0\n",
       " rf_nj_fastme                    2\n",
       " dtype: object)"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "cfg_ = cfg = load_yaml_config(\"cfg/bme_config_v3.yml\")\n",
    "run(cfg_)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "gradme",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.13"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}