Computing Nearest neighbor interchange (NNI)

shuffling trees.ipynb

{
  "cells": [
    {
      "metadata": {},
      "id": "271ea6db",
      "cell_type": "markdown",
      "source": "# Computing Nearest neighbor interchange (NNI)"
    },
    {
      "metadata": {
        "trusted": false
      },
      "id": "a4a14e90",
      "cell_type": "code",
      "source": "from random import seed",
      "execution_count": 1,
      "outputs": []
    },
    {
      "metadata": {
        "trusted": false
      },
      "id": "9266f789",
      "cell_type": "code",
      "source": "from ete4 import Tree",
      "execution_count": 2,
      "outputs": []
    },
    {
      "metadata": {
        "trusted": false
      },
      "id": "9151bb90",
      "cell_type": "code",
      "source": "def iter_NNIs(tree, copies=False):\n    \"\"\"\n    Generator yielding all neighbor trees using Nearest Neighbor Interchange search.\n    \n    WARNING: only works with binary trees.\n    \n    :param False copies: if True yield copies of input Tree. Otherwise (default) it \n       yields alwaysthe same object in different configurations (if stacked in a \n       list, this will finally result in a list of all identical trees)\n    \"\"\"\n    for n in tree.iter_descendants():\n        if n.is_leaf():\n            continue\n        # yielder function out of the loop for preformance\n        if copies:\n            yielder = lambda x: x.copy()\n        else:\n            yielder = lambda x: x\n        # we only need to play with 3 nodes to get all combinations, the fourth can be forgotten\n        # find the 3 nodes\n        a, b = n.get_children()\n        if n.up.is_root():\n            if n is tree.get_children()[1]:  # root does not exist here:it's just one edge\n                continue\n            c = n.get_sisters()[0].get_children()[0]\n        else:\n            c = n.get_sisters()[0]\n        # swap nodes for first neighbor tree\n        au = a.up\n        bu = b.up\n        cu = c.up\n        a.detach()\n        c.detach()\n        au.add_child(c)\n        cu.add_child(a)\n        yield yielder(tree)\n        # rebuild tree\n        a.detach()\n        c.detach()\n        cu.add_child(c)\n        au.add_child(a)\n        # swap nodes for second neighbor tree\n        b.detach()\n        c.detach()\n        bu.add_child(c)\n        cu.add_child(b)\n        yield yielder(tree)\n        # rebuild tree\n        b.detach()\n        c.detach()\n        cu.add_child(c)\n        bu.add_child(b)",
      "execution_count": 3,
      "outputs": []
    },
    {
      "metadata": {
        "trusted": false
      },
      "id": "e9ed01be",
      "cell_type": "code",
      "source": "seed(2)\nt = Tree()\nt.populate(5, names_library=map(str, range(1, 11)))\nprint(t)",
      "execution_count": 4,
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": "\n      /-10\n   /-|\n  |   \\-9\n--|\n  |   /-8\n   \\-|\n     |   /-7\n      \\-|\n         \\-6\n"
        }
      ]
    },
    {
      "metadata": {
        "trusted": false
      },
      "id": "2611d314",
      "cell_type": "code",
      "source": "for tt in iter_NNIs(t):\n    print(tt)",
      "execution_count": 5,
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": "\n      /-9\n   /-|\n  |   \\-8\n--|\n  |      /-7\n  |   /-|\n   \\-|   \\-6\n     |\n      \\-10\n\n      /-10\n   /-|\n  |   \\-8\n--|\n  |      /-7\n  |   /-|\n   \\-|   \\-6\n     |\n      \\-9\n\n      /-10\n   /-|\n  |   \\-9\n--|\n  |      /-6\n  |   /-|\n   \\-|   \\-8\n     |\n      \\-7\n\n      /-10\n   /-|\n  |   \\-9\n--|\n  |      /-7\n  |   /-|\n   \\-|   \\-8\n     |\n      \\-6\n"
        }
      ]
    },
    {
      "metadata": {
        "trusted": false
      },
      "id": "ccd79369",
      "cell_type": "code",
      "source": "seed(2)\nt = Tree()\ntlen = 10000\nt.populate(tlen, names_library=map(str, range(1, tlen + 1)))",
      "execution_count": 6,
      "outputs": []
    },
    {
      "metadata": {},
      "id": "321b11fb",
      "cell_type": "markdown",
      "source": "## simple test"
    },
    {
      "metadata": {},
      "id": "a9d1fa59",
      "cell_type": "markdown",
      "source": "Number of neighbor trees found is:"
    },
    {
      "metadata": {
        "trusted": false
      },
      "id": "5f7a7823",
      "cell_type": "code",
      "source": "len(list(iter_NNIs(t)))",
      "execution_count": 7,
      "outputs": [
        {
          "data": {
            "text/plain": "19994"
          },
          "execution_count": 7,
          "metadata": {},
          "output_type": "execute_result"
        }
      ]
    },
    {
      "metadata": {},
      "id": "a6f13a9c",
      "cell_type": "markdown",
      "source": "should be equal to $2\\times(N-3)$"
    },
    {
      "metadata": {
        "trusted": false
      },
      "id": "7dcd368e",
      "cell_type": "code",
      "source": "2 * (tlen - 3)",
      "execution_count": 8,
      "outputs": [
        {
          "data": {
            "text/plain": "19994"
          },
          "execution_count": 8,
          "metadata": {},
          "output_type": "execute_result"
        }
      ]
    },
    {
      "metadata": {
        "trusted": false
      },
      "id": "d426e565",
      "cell_type": "code",
      "source": "",
      "execution_count": null,
      "outputs": []
    }
  ],
  "metadata": {
    "kernelspec": {
      "name": "conda-env-ete4-py",
      "display_name": "Python [conda env:ete4]",
      "language": "python"
    },
    "language_info": {
      "name": "python",
      "version": "3.9.7",
      "mimetype": "text/x-python",
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "pygments_lexer": "ipython3",
      "nbconvert_exporter": "python",
      "file_extension": ".py"
    },
    "toc": {
      "nav_menu": {},
      "number_sections": false,
      "sideBar": true,
      "skip_h1_title": false,
      "base_numbering": 1,
      "title_cell": "Table of Contents",
      "title_sidebar": "Contents",
      "toc_cell": false,
      "toc_position": {},
      "toc_section_display": true,
      "toc_window_display": false
    },
    "gist": {
      "id": "",
      "data": {
        "description": "Computing Nearest neighbor interchange (NNI)",
        "public": true
      }
    }
  },
  "nbformat": 4,
  "nbformat_minor": 5
}