Repository: cytomining/pycytominer

Issue:
Load only some columns during compartment loading in SingleCells
Right now, it is possible to subsample the data afterward to certain rows, and after doing something like merge_single_cells you can subset to certain columns, but there is no way to subset before loading in everything. This is a hassle when data is large - @rsenft1 is running into it in some of her projects. If we were to add this functionality, would a PR be welcomed?

Comments:

Patch:
diff --git a/pycytominer/annotate.py b/pycytominer/annotate.py
index 816e2e5..e64e833 100644
--- a/pycytominer/annotate.py
+++ b/pycytominer/annotate.py
@@ -79,7 +79,9 @@ def annotate(
 
     annotated = platemap.merge(
         profiles, left_on=join_on[0], right_on=join_on[1], how="inner"
-    ).drop(join_on[0], axis="columns")
+    )
+    if join_on[0] != join_on[1]:
+        annotated = annotated.drop(join_on[0], axis="columns")
 
     # Add specific Connectivity Map (CMAP) formatting
     if format_broad_cmap:
@@ -95,6 +97,9 @@ def annotate(
             ), "external metadata at {} does not exist".format(external_metadata)
 
             external_metadata = pd.read_csv(external_metadata)
+    else:
+        # Make a copy of the external metadata to avoid modifying the original column names
+        external_metadata = external_metadata.copy()
 
     if isinstance(external_metadata, pd.DataFrame):
         external_metadata.columns = [
diff --git a/pycytominer/cyto_utils/cells.py b/pycytominer/cyto_utils/cells.py
index 08029fe..e26d792 100644
--- a/pycytominer/cyto_utils/cells.py
+++ b/pycytominer/cyto_utils/cells.py
@@ -52,7 +52,7 @@ class SingleCells(object):
     image_feature_categories : list of str, optional
         List of categories of features from the image table to add to the profiles.
     features: str or list of str, default "infer"
-        List of features that should be aggregated.
+        List of features that should be loaded or aggregated.
     load_image_data : bool, default True
         Whether or not the image data should be loaded into memory.
     image_table_name : str, default "image"
@@ -415,10 +415,14 @@ class SingleCells(object):
         return num_rows
 
     def get_sql_table_col_names(self, table):
-        """Get feature and metadata columns lists."""
+        """Get column names from the database."""
         ptr = self.conn.execute(f"SELECT * FROM {table} LIMIT 1").cursor
         col_names = [obj[0] for obj in ptr.description]
 
+        return col_names
+
+    def split_column_categories(self, col_names):
+        """Split a list of column names into feature and metadata columns lists."""
         feat_cols = []
         meta_cols = []
         for col in col_names:
@@ -445,7 +449,10 @@ class SingleCells(object):
 
         # Get data useful to pre-alloc memory
         num_cells = self.count_sql_table_rows(compartment)
-        meta_cols, feat_cols = self.get_sql_table_col_names(compartment)
+        col_names = self.get_sql_table_col_names(compartment)
+        if self.features != "infer":  # allow to get only some features
+            col_names = [x for x in col_names if x in self.features]
+        meta_cols, feat_cols = self.split_column_categories(col_names)
         num_meta, num_feats = len(meta_cols), len(feat_cols)
 
         # Use pre-allocated np.array for data
@@ -613,7 +620,11 @@ class SingleCells(object):
             sql=f"select {cols} from {compartment} limit 1",
             con=self.conn,
         )
-        typeof_str = ", ".join([f"typeof({x})" for x in compartment_row1.columns])
+        all_columns = compartment_row1.columns
+        if self.features != "infer": # allow to get only some features
+            all_columns = [x for x in all_columns if x in self.features]
+
+        typeof_str = ", ".join([f"typeof({x})" for x in all_columns])
         compartment_dtypes = pd.read_sql(
             sql=f"select {typeof_str} from {compartment} limit 1",
             con=self.conn,
diff --git a/pycytominer/cyto_utils/load.py b/pycytominer/cyto_utils/load.py
index 5bb17ad..7224144 100644
--- a/pycytominer/cyto_utils/load.py
+++ b/pycytominer/cyto_utils/load.py
@@ -74,6 +74,9 @@ def load_platemap(platemap, add_metadata_id=True):
             platemap = pd.read_csv(platemap, sep=delim)
         except FileNotFoundError:
             raise FileNotFoundError(f"{platemap} platemap file not found")
+    else:
+        # Setting platemap to a copy to prevent column name changes from back-propagating
+        platemap = platemap.copy()
 
     if add_metadata_id:
         platemap.columns = [


Test patch:
diff --git a/pycytominer/tests/test_annotate.py b/pycytominer/tests/test_annotate.py
index 1866fbd..50b2166 100644
--- a/pycytominer/tests/test_annotate.py
+++ b/pycytominer/tests/test_annotate.py
@@ -52,11 +52,32 @@ def test_annotate():
     pd.testing.assert_frame_equal(result, expected_result)
 
 
+def test_annotate_platemap_naming():
+
+    # Test annotate with the same column name in platemap and data.
+    platemap_modified_df = PLATEMAP_DF.copy().rename(
+        columns={"well_position": "Metadata_Well"}
+    )
+
+    expected_result = platemap_modified_df.merge(
+        DATA_DF, left_on="Metadata_Well", right_on="Metadata_Well"
+    ).rename(columns={"gene": "Metadata_gene"})
+
+    result = annotate(
+        profiles=DATA_DF,
+        platemap=platemap_modified_df,
+        join_on=["Metadata_Well", "Metadata_Well"],
+    )
+
+    pd.testing.assert_frame_equal(result, expected_result)
+
+
 def test_annotate_output():
+
     annotate(
         profiles=DATA_DF,
         platemap=PLATEMAP_DF,
-        join_on=["Metadata_well_position", "Metadata_Well"],
+        join_on=["well_position", "Metadata_Well"],
         add_metadata_id_to_platemap=False,
         output_file=OUTPUT_FILE,
     )
@@ -64,7 +85,7 @@ def test_annotate_output():
     result = annotate(
         profiles=DATA_DF,
         platemap=PLATEMAP_DF,
-        join_on=["Metadata_well_position", "Metadata_Well"],
+        join_on=["well_position", "Metadata_Well"],
         add_metadata_id_to_platemap=False,
         output_file="none",
     )
@@ -74,11 +95,12 @@ def test_annotate_output():
 
 
 def test_annotate_output_compress():
+
     compress_file = pathlib.Path(f"{TMPDIR}/test_annotate_compress.csv.gz")
     annotate(
         profiles=DATA_DF,
         platemap=PLATEMAP_DF,
-        join_on=["Metadata_well_position", "Metadata_Well"],
+        join_on=["well_position", "Metadata_Well"],
         add_metadata_id_to_platemap=False,
         output_file=compress_file,
         compression_options={"method": "gzip"},
@@ -87,7 +109,7 @@ def test_annotate_output_compress():
     result = annotate(
         profiles=DATA_DF,
         platemap=PLATEMAP_DF,
-        join_on=["Metadata_well_position", "Metadata_Well"],
+        join_on=["well_position", "Metadata_Well"],
         add_metadata_id_to_platemap=False,
         output_file="none",
     )
diff --git a/pycytominer/tests/test_cyto_utils/test_cells.py b/pycytominer/tests/test_cyto_utils/test_cells.py
index 295db95..c9ec8ed 100644
--- a/pycytominer/tests/test_cyto_utils/test_cells.py
+++ b/pycytominer/tests/test_cyto_utils/test_cells.py
@@ -193,6 +193,19 @@ AP_IMAGE_DIFF_NAME = SingleCells(
     sql_file=IMAGE_DIFF_FILE, load_image_data=False, image_feature_categories=["Count"]
 )
 
+SUBSET_FEATURES = [
+    "TableNumber",
+    "ImageNumber",
+    "ObjectNumber",
+    "Cells_Parent_Nuclei",
+    "Cytoplasm_Parent_Cells",
+    "Cytoplasm_Parent_Nuclei",
+    "Cells_a",
+    "Cytoplasm_a",
+    "Nuclei_a",
+]
+AP_SUBSET = SingleCells(sql_file=TMP_SQLITE_FILE, features=SUBSET_FEATURES)
+
 
 def test_SingleCells_init():
     """
@@ -291,10 +304,17 @@ def test_sc_count_sql_table():
 def test_get_sql_table_col_names():
     # Iterate over initialized compartments
     for compartment in AP.compartments:
-        meta_cols, feat_cols = AP.get_sql_table_col_names(table=compartment)
-        assert meta_cols == ["ObjectNumber", "ImageNumber", "TableNumber"]
-        for i in ["a", "b", "c", "d"]:
-            assert f"{compartment.capitalize()}_{i}" in feat_cols
+        expected_meta_cols = ["ObjectNumber", "ImageNumber", "TableNumber"]
+        expected_feat_cols = [f"{compartment.capitalize()}_{i}" for i in ["a", "b", "c", "d"]]
+        if compartment == 'cytoplasm':
+            expected_feat_cols += ["Cytoplasm_Parent_Cells","Cytoplasm_Parent_Nuclei"]
+        col_name_result = AP.get_sql_table_col_names(table=compartment)
+        assert sorted(col_name_result) == sorted(expected_feat_cols+expected_meta_cols)
+        meta_cols, feat_cols = AP.split_column_categories(
+            col_name_result
+        )
+        assert meta_cols == expected_meta_cols
+        assert feat_cols == expected_feat_cols
 
 
 def test_merge_single_cells():
@@ -417,6 +437,13 @@ def test_merge_single_cells():
     )
 
 
+def test_merge_single_cells_subset():
+    sc_merged_df = AP_SUBSET.merge_single_cells()
+    assert (sc_merged_df.shape[1]) == 13
+    non_meta_cols = [x for x in sc_merged_df.columns if "Metadata" not in x]
+    assert len(non_meta_cols) == len([x for x in non_meta_cols if x in SUBSET_FEATURES])
+
+
 def test_merge_single_cells_subsample():
 
     for subsample_frac in [0.1, 0.5, 0.9]: