Adding annotate.py

pycytominer/annotate.py

@@ -0,0 +1,61 @@
+"""
+Annotates profiles with metadata information
+"""
+
+import numpy as np
+import pandas as pd
+
+
+def annotate(
+    profiles,
+    platemap,
+    join_on=["Metadata_well_position", "Metadata_Well"],
+    output_file="none",
+    add_metadata_id_to_platemap=True,
+):
+    """
+    Exclude features that have correlations above a certain threshold
+
+    Arguments:
+    profiles - either pandas DataFrame or a file that stores profile data
+    platemap - either pandas DataFrame or a file that stores platemap metadata
+    join_on - list of length two indicating which variables to merge profiles and plate
+              [default: ["Metadata_well_position", "Metadata_Well"]]. The first element
+              indicates variable(s) in platemap and the second element indicates
+              variable(s) in profiles to merge using.
+              Note the setting of `add_metadata_id_to_platemap`
+    output_file - [default: "none"] if provided, will write annotated profiles to file
+                  if not specified, will return the annotated profiles. We recommend
+                  that this output file be suffixed with "_augmented.csv".
+    add_metadata_id_to_platemap - boolean if the platemap variables should be recoded
+
+    Return:
+    Pandas DataFrame of annotated profiles or written to file
+    """
+    # Load Data
+    if not isinstance(profiles, pd.DataFrame):
+        try:
+            profiles = pd.read_csv(profiles)
+        except FileNotFoundError:
+            raise FileNotFoundError("{} profile file not found".format(profiles))
+
+    if not isinstance(platemap, pd.DataFrame):
+        try:
+            platemap = pd.read_csv(platemap, sep="\t")
+        except FileNotFoundError:
+            raise FileNotFoundError("{} platemap file not found".format(platemap))
+
+    if add_metadata_id_to_platemap:
+        platemap.columns = [
+            "Metadata_{}".format(x) if not x.startswith("Metadata_") else x
+            for x in platemap.columns
+        ]
+
+    annotated = platemap.merge(
+        profiles, left_on=join_on[0], right_on=join_on[1], how="inner"
+    ).drop(join_on[0], axis='columns')
+
+    if output_file != "none":
+        annotated.to_csv(output_file, index=False)
+    else:
+        return annotated

pycytominer/tests/test_annotate.py

@@ -0,0 +1,74 @@
+import tempfile
+import random
+import pandas as pd
+from pycytominer.annotate import annotate
+
+random.seed(123)
+
+# Get temporary directory
+tmpdir = tempfile.gettempdir()
+
+# Lauch a sqlite connection
+output_file = "{}/test.csv".format(tmpdir)
+
+# Build data to use in tests
+data_df = pd.concat(
+    [
+        pd.DataFrame({"g": "a", "x": [1, 3, 8], "y": [5, 3, 1]}),
+        pd.DataFrame({"g": "b", "x": [1, 3, 5], "y": [8, 3, 1]}),
+    ]
+).reset_index(drop=True)
+
+# Build data to use in tests
+data_df = pd.concat(
+    [
+        pd.DataFrame(
+            {"Metadata_Well": ["A01", "A02", "A03"], "x": [1, 3, 8], "y": [5, 3, 1]}
+        ),
+        pd.DataFrame(
+            {"Metadata_Well": ["B01", "B02", "B03"], "x": [1, 3, 5], "y": [8, 3, 1]}
+        ),
+    ]
+).reset_index(drop=True)
+
+platemap_df = pd.DataFrame(
+    {
+        "well_position": ["A01", "A02", "A03", "B01", "B02", "B03"],
+        "gene": ["x", "y", "z"] * 2,
+    }
+).reset_index(drop=True)
+
+
+def test_annotate():
+    result = annotate(
+        profiles=data_df,
+        platemap=platemap_df,
+        join_on=["Metadata_well_position", "Metadata_Well"],
+    )
+
+    expected_result = platemap_df.merge(
+        data_df, left_on="Metadata_well_position", right_on="Metadata_Well"
+    ).drop(["Metadata_well_position"], axis="columns")
+
+    pd.testing.assert_frame_equal(result, expected_result)
+
+
+def test_annotate_write():
+    _ = annotate(
+        profiles=data_df,
+        platemap=platemap_df,
+        join_on=["Metadata_well_position", "Metadata_Well"],
+        add_metadata_id_to_platemap=False,
+        output_file=output_file
+    )
+
+    result = annotate(
+        profiles=data_df,
+        platemap=platemap_df,
+        join_on=["Metadata_well_position", "Metadata_Well"],
+        add_metadata_id_to_platemap=False,
+        output_file="none"
+    )
+    expected_result = pd.read_csv(output_file)
+
+    pd.testing.assert_frame_equal(result, expected_result)