From 7148e066b442602a8f6139e19b86cea29d628453 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Al=C3=A1n=20Mu=C3=B1oz?= <amuoz@ed.ac.uk>
Date: Thu, 10 Feb 2022 17:06:35 +0000
Subject: [PATCH] add lineage functions
---
utils/lineage.py | 27 +++++++++++++++++++++++++++
1 file changed, 27 insertions(+)
create mode 100644 utils/lineage.py
diff --git a/utils/lineage.py b/utils/lineage.py
new file mode 100644
index 00000000..b72c69cd
--- /dev/null
+++ b/utils/lineage.py
@@ -0,0 +1,27 @@
+#!/usr/bin/env python3
+import numpy as np
+import pandas as pd
+
+from agora.io.bridge import groupsort
+
+
+def mb_array_to_dict(mb_array: np.ndarray):
+ """
+ Convert a lineage ndarray (trap, mother_id, daughter_id)
+ into a dictionary of lists ( mother_id ->[daughters_ids] )
+ """
+ return {
+ (trap, mo): [(trap, d[0]) for d in daughters]
+ for trap, mo_da in groupsort(mb_array).items()
+ for mo, daughters in groupsort(mo_da).items()
+ }
+
+
+def mb_array_to_indices(mb_array: np.ndarray):
+ """
+ Convert a lineage ndarray (trap, mother_id, daughter_id)
+ into a dictionary of lists ( mother_id ->[daughters_ids] )
+ """
+ return pd.MultiIndex.from_arrays(mb_array[:, :2].T).union(
+ pd.MultiIndex.from_arrays(mb_array[:, [0, 2]].T)
+ )
--
GitLab