From 030eea18986ce5a8c9d7668e1211d1284917c68c Mon Sep 17 00:00:00 2001
From: bmccandless <bmccandless@chanzuckerberg.com>
Date: Tue, 26 May 2020 16:06:33 -0700
Subject: [PATCH] Update run_diffexp.py (#1494)

You can now select labels within groups instead of random rows.
---
 server/test/run_diffexp.py | 55 +++++++++++++++++++++++++++++---------
 1 file changed, 43 insertions(+), 12 deletions(-)

diff --git a/server/test/run_diffexp.py b/server/test/run_diffexp.py
index 9faba5ac..d3fad73a 100644
--- a/server/test/run_diffexp.py
+++ b/server/test/run_diffexp.py
@@ -15,8 +15,10 @@ from server.data_cxg.cxg_adaptor import CxgAdaptor
 def main():
     parser = argparse.ArgumentParser("A command to test diffexp")
     parser.add_argument("dataset", help="name of a dataset to load")
-    parser.add_argument("-na", "--numA", type=int, required=True, help="number of rows in group A")
-    parser.add_argument("-nb", "--numB", type=int, required=True, help="number of rows in group B")
+    parser.add_argument("-na", "--numA", type=int, help="number of rows in group A")
+    parser.add_argument("-nb", "--numB", type=int, help="number of rows in group B")
+    parser.add_argument("-va", "--varA", help="obs variable:value to use for group A")
+    parser.add_argument("-vb", "--varB", help="obs variable:value to use for group B")
     parser.add_argument("-t", "--trials", default=1, type=int, help="number of trials")
     parser.add_argument(
         "-a", "--alg", choices=("default", "generic", "cxg"), default="default", help="algorithm to use"
@@ -41,22 +43,34 @@ def main():
         if isinstance(adaptor, CxgAdaptor):
             adaptor.open_array("X").schema.dump()
 
-    numA = args.numA
-    numB = args.numB
+    random.seed(args.seed)
+    np.random.seed(args.seed)
     rows = adaptor.get_shape()[0]
 
-    random.seed(args.seed)
+    if args.numA:
+        filterA = random.sample(range(rows), args.numA)
+    elif args.varA:
+        vname, vval = args.varA.split(":")
+        filterA = get_filter_from_obs(adaptor, vname, vval)
+    else:
+        print("must supply numA or varA")
+        sys.exit(1)
 
-    if not args.new_selection:
-        samples = random.sample(range(rows), numA + numB)
-        filterA = samples[:numA]
-        filterB = samples[numA:]
+    if args.numB:
+        filterB = random.sample(range(rows), args.numB)
+    elif args.varB:
+        vname, vval = args.varB.split(":")
+        filterB = get_filter_from_obs(adaptor, vname, vval)
+    else:
+        print("must supply numB or varB")
+        sys.exit(1)
 
     for i in range(args.trials):
         if args.new_selection:
-            samples = random.sample(range(rows), numA + numB)
-            filterA = samples[:numA]
-            filterB = samples[numA:]
+            if args.numA:
+                filterA = random.sample(range(rows), args.numA)
+            if args.numB:
+                filterB = random.sample(range(rows), args.numB)
 
         maskA = np.zeros(rows, dtype=bool)
         maskA[filterA] = True
@@ -82,5 +96,22 @@ def main():
             print(res)
 
 
+def get_filter_from_obs(adaptor, obsname, obsval):
+    attrs = adaptor.get_obs_columns()
+    if obsname not in attrs:
+        print(f"Unknown obs attr {obsname}: expected on of {attrs}")
+        sys.exit(1)
+    obsvals = adaptor.query_obs_array(obsname)[:]
+    obsval = type(obsvals[0])(obsval)
+
+    vfilter = np.where(obsvals == obsval)[0]
+    if len(vfilter) == 0:
+        u = np.unique(obsvals)
+        print(f"Unknown value in variable {obsname}:{obsval}: expected one of {list(u)}")
+        sys.exit(1)
+
+    return vfilter
+
+
 if __name__ == "__main__":
     main()