improved lineage graph logic

caveclient/chunkedgraph.py

@@ -668,12 +668,51 @@ def get_lineage_graph(
 
         url = self._endpoints["handle_lineage_graph"].format_map(endpoint_mapping)
         data = json.dumps({"root_ids": root_id}, cls=BaseEncoder)
-        r = handle_response(self.session.post(url, data=data, params=params))
+        lineage_graph = handle_response(
+            self.session.post(url, data=data, params=params)
+        )
+
+        rm_node_ids = []
+        rm_node_idx = []
+        rm_link_idx = []
+
+        for node_idx, node in enumerate(lineage_graph["nodes"]):
+            if (
+                timestamp_past is not None
+                and node["timestamp"] < timestamp_past.timestamp()
+            ):
+                rm_node_ids.append(node["id"])
+                rm_node_idx.append(node_idx)
+            if (
+                timestamp_future is not None
+                and node["timestamp"] >= timestamp_future.timestamp()
+            ):
+                rm_node_ids.append(node["id"])
+                rm_node_idx.append(node_idx)
+
+        rm_node_ids = set(rm_node_ids)
+        for link_idx, link in enumerate(lineage_graph["links"]):
+            if link["source"] in rm_node_ids or link["target"] in rm_node_ids:
+                lineage_graph["links"].remove(link)
+                rm_link_idx.append(link_idx)
+
+        rm_link_idx = set(rm_link_idx)
+        rm_node_idx = set(rm_node_idx)
+        lineage_graph["links"] = [
+            link
+            for idx, link in enumerate(lineage_graph["links"])
+            if idx not in rm_link_idx
+        ]
+        lineage_graph["nodes"] = [
+            node
+            for idx, node in enumerate(lineage_graph["nodes"])
+            if idx not in rm_node_idx
+        ]
 
         if as_nx_graph:
-            return nx.node_link_graph(r)
+            return nx.node_link_graph(lineage_graph)
         else:
-            return r
+            return lineage_graph
 
     def get_latest_roots(self, root_id, timestamp_future=None):
         """Returns root ids that are the latest successors of a given root id.

debug_tabular_changelog.ipynb

@@ -0,0 +1,199 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "id": "f79c5211",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import caveclient\n",
+    "# from pychunkedgraph.graph import chunkedgraph, attributes, segmenthistory"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "id": "1b503d6e",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "client = caveclient.CAVEclient(\"fanc_production_mar2021\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "35f6019e",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import requests \n",
+    "import numpy as np\n",
+    "\n",
+    "def get_userinfo_dict(user_ids, auth_token):\n",
+    "    AUTH_URL = \"global.daf-apis.com/auth\"\n",
+    "\n",
+    "#     if AUTH_URL is None:\n",
+    "#         raise cg_exceptions.ChunkedGraphError(\"No AUTH_URL defined\")\n",
+    "\n",
+    "    users_request = requests.get(\n",
+    "        f\"https://{AUTH_URL}/api/v1/user?id={','.join(map(str, np.unique(user_ids)))}\",\n",
+    "        headers={\"authorization\": \"Bearer \" + auth_token},\n",
+    "        timeout=5,\n",
+    "    )\n",
+    "    return {x[\"id\"]: x[\"name\"] for x in users_request.json()}, {\n",
+    "        x[\"id\"]: x[\"pi\"] for x in users_request.json()\n",
+    "    }"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "9edb2271",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "# get_userinfo_dict([1, 2, 3], \"5bf9d876a194e9217d53961bec09d42f\")\n",
+    "get_userinfo_dict([1, 2, 3], \"04f6375fdae394939fa01bad9d4ea8c2\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "id": "7b63d5f5",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "array([ True])"
+      ]
+     },
+     "execution_count": 5,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "client.chunkedgraph.is_latest_roots([648518346506629278])"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 25,
+   "id": "a7443d48",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from pychunkedgraph.graph.client.bigtable.client import Client as BigTableClient\n",
+    "from pychunkedgraph.graph.client.bigtable import get_client_info\n",
+    "from collections import namedtuple\n",
+    "_backend_clientinfo_fields = (\"TYPE\", \"CONFIG\")\n",
+    "_backend_clientinfo_defaults = (None, None)\n",
+    "BackendClientInfo = namedtuple(\n",
+    "    \"BackendClientInfo\",\n",
+    "    _backend_clientinfo_fields,\n",
+    "    defaults=_backend_clientinfo_defaults,\n",
+    ")\n",
+    "\n",
+    "client_info = BackendClientInfo(\n",
+    "    CONFIG=get_client_info(project=\"fanc-fly\",\n",
+    "                           instance=\"chunkedgraphs\", \n",
+    "                           admin=False, \n",
+    "                           read_only=True)\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 26,
+   "id": "4c248860",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "cg = chunkedgraph.ChunkedGraph(graph_id=\"mar2021_prod\", client_info=client_info)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 36,
+   "id": "1e06aaf6",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "user_id = 2\n",
+    "root_ids = [648518346506629278]\n",
+    "filtered = False\n",
+    "\n",
+    "history = segmenthistory.SegmentHistory(\n",
+    "    cg,\n",
+    "    root_ids,\n",
+    ")\n",
+    "if filtered:\n",
+    "    tab = history.tabular_changelogs_filtered\n",
+    "else:\n",
+    "    tab = history.tabular_changelogs\n",
+    "all_user_ids = []\n",
+    "for tab_k in tab.keys():\n",
+    "    all_user_ids.extend(np.array(tab[tab_k][\"user_id\"]).reshape(-1))\n",
+    "\n",
+    "all_user_ids = []\n",
+    "for tab_k in tab.keys():\n",
+    "    all_user_ids.extend(np.array(tab[tab_k][\"user_id\"]).reshape(-1))\n",
+    "\n",
+    "all_user_ids = np.unique(all_user_ids)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 37,
+   "id": "e7a95676",
+   "metadata": {},
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "array([], dtype=float64)"
+      ]
+     },
+     "execution_count": 37,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
+   "source": [
+    "all_user_ids"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "b1617b4d",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "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.7.12"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}