Update according to feedback comments

bagherilab · Oct 7, 2023 · c786f18 · c786f18
1 parent f7f53e3
commit c786f18
Show file tree

Hide file tree

Showing 4 changed files with 35 additions and 80 deletions.
diff --git a/src/arcade_collection/output/__init__.py b/src/arcade_collection/output/__init__.py
@@ -7,8 +7,10 @@
 from .get_location_voxels import get_location_voxels
 from .merge_parsed_results import merge_parsed_results
 from .parse_cells_file import parse_cells_file
+from .parse_growth_file import parse_growth_file
 from .parse_locations_file import parse_locations_file
 
+
 convert_model_units = task(convert_model_units)
 extract_tick_json = task(extract_tick_json)
 get_location_voxels = task(get_location_voxels)

diff --git a/src/arcade_collection/output/parse_growth_file.py b/src/arcade_collection/output/parse_growth_file.py
@@ -1,7 +1,5 @@
 import json
-from os import path
 import tarfile
-from typing import List, Union
 
 import numpy as np
 import pandas as pd
@@ -22,6 +20,16 @@
     "CYCLE",
 ]
 
+CELL_STATES = [
+    "NEUTRAL",
+    "APOPTOTIC",
+    "QUIESCENT",
+    "MIGRATORY",
+    "PROLIFERATIVE",
+    "SENESCENT",
+    "NECROTIC",
+]
+
 
 @task
 def parse_growth_file(tar: tarfile.TarFile) -> pd.DataFrame:
@@ -56,26 +64,6 @@ def parse_growth_file(tar: tarfile.TarFile) -> pd.DataFrame:
     return timepoints_df
 
 
-def convert_state_to_string(state_index: int, state_list: List[str]) -> Union[str, None]:
-    """
-    Convert the numbers that represent cell state into an annotation.
-
-    Parameters
-    ----------
-    state_index :
-        The index of cell states.
-    state_list :
-        The list of cell states.
-
-    Returns
-    -------
-    :
-        The cell state annotation.
-    """
-
-    return state_list[state_index]
-
-
 def parse_growth_timepoint(timepoint: dict, seed: int) -> list:
     """
     Parse one timepoint of the simulation.
@@ -99,20 +87,16 @@ def parse_growth_timepoint(timepoint: dict, seed: int) -> list:
     time = timepoint["time"]
 
     for (location, cells) in timepoint["cells"]:
-        u = int(location[0])
-        v = int(location[1])
-        w = int(location[2])
-        z = int(location[3])
+        u, v, w, z = location
 
         for cell in cells:
-            population = cell[1]
-            state = cell[2]
-            position = cell[3]
-            volume = cell[4]
-            if len(cell[5]) == 0:
+            _, population, state, position, volume, cycles = cell
+
+            if len(cycles) == 0:
                 cycle = None
             else:
-                cycle = np.mean(cell[5])
+                cycle = np.mean(cycles)
+
             data_list = [
                 time,
                 seed,
@@ -122,7 +106,7 @@ def parse_growth_timepoint(timepoint: dict, seed: int) -> list:
                 z,
                 position,
                 population,
-                convert_state_to_string(state, ["NEU", "APO", "QUI", "MIG", "PRO", "SEN", "NEC"]),
+                CELL_STATES[state],
                 volume,
                 cycle,
             ]

diff --git a/src/arcade_collection/output/parse_params_file.py b/src/arcade_collection/output/parse_params_file.py
diff --git a/tests/arcade_collection/output/test_parse_growth_file.py b/tests/arcade_collection/output/test_parse_growth_file.py
@@ -67,22 +67,22 @@ def test_parse_growth_timepoint(self):
             "POSITION": [0, 0, 0, 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 1],
             "POPULATION": [1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 2, 0],
             "STATE": [
-                "QUI",
-                "QUI",
-                "QUI",
-                "MIG",
-                "APO",
-                "PRO",
-                "QUI",
-                "MIG",
-                "QUI",
-                "QUI",
-                "QUI",
-                "MIG",
-                "APO",
-                "PRO",
-                "QUI",
-                "NEC",
+                "QUIESCENT",
+                "QUIESCENT",
+                "QUIESCENT",
+                "MIGRATORY",
+                "APOPTOTIC",
+                "PROLIFERATIVE",
+                "QUIESCENT",
+                "MIGRATORY",
+                "QUIESCENT",
+                "QUIESCENT",
+                "QUIESCENT",
+                "MIGRATORY",
+                "APOPTOTIC",
+                "PROLIFERATIVE",
+                "QUIESCENT",
+                "NECROTIC",
             ],
             "VOLUME": [
                 2322.26,
@@ -123,9 +123,4 @@ def test_parse_growth_timepoint(self):
         }
 
         expected_df = pd.DataFrame(expected_dict)
-        print(returned_df)
-
-        print("-----")
-
-        print(expected_df)
         self.assertTrue(expected_df.equals(returned_df))