add in flag for numeric columns to set_column_choices

Author: nickzoic

countess/core/parameters.py

@@ -14,7 +14,7 @@
 logger = logging.getLogger(__name__)
 
 
-def make_prefix_groups(strings: List[str]) -> Dict[str, List[str]]:
+def make_prefix_groups(strings: Iterable[str]) -> Dict[str, List[str]]:
     groups: Dict[str, List[str]] = {}
     for s in strings:
         if m := re.match(r"(.*?_+)([^_]+)$", s):
@@ -50,7 +50,7 @@ def get_parameters(self, key, base_dir="."):
     def get_hash_value(self):
         raise NotImplementedError(f"Implement {self.__class__.__name__}.get_hash_value()")
 
-    def set_column_choices(self, choices: List[str]):
+    def set_column_choices(self, choices: Mapping[str, bool]):
         pass
 
 
@@ -546,13 +546,18 @@ class ColumnChoiceParam(ChoiceParam):
     """A ChoiceParam which DaskTransformPlugin knows
     it should automatically update with a list of columns"""
 
-    def set_column_choices(self, choices: List[str]):
-        self.set_choices(list(choices))
+    def set_column_choices(self, choices: Mapping[str, bool]):
+        self.set_choices(list(choices.keys()))
 
     def get_column(self, df):
         return _dataframe_get_column(df, self.value)
 
 
+class NumericColumnChoiceParam(ColumnChoiceParam):
+    def set_column_choices(self, choices: Mapping[str, bool]):
+        self.set_choices([c for c, n in choices.items() if n])
+
+
 class ColumnOrNoneChoiceParam(ColumnChoiceParam):
     DEFAULT_VALUE = "— NONE —"
 
@@ -573,8 +578,8 @@ def get_column(self, df):
 
 
 class ColumnGroupChoiceParam(ChoiceParam):
-    def set_column_choices(self, choices: List[str]):
-        self.set_choices([n + "*" for n in make_prefix_groups(choices).keys()])
+    def set_column_choices(self, choices: Mapping[str, bool]):
+        self.set_choices([n + "*" for n in make_prefix_groups(choices.keys()).keys()])
 
     def get_column_prefix(self):
         return self.value.removesuffix("*")
@@ -634,7 +639,7 @@ class ColumnOrStringParam(ColumnChoiceParam):
     DEFAULT_VALUE: Any = ""
     PREFIX = "— "
 
-    def set_column_choices(self, choices: List[str]):
+    def set_column_choices(self, choices: Mapping[str, bool]):
         self.set_choices([self.PREFIX + c for c in choices])
 
     def get_column_name(self) -> Optional[str]:
@@ -678,6 +683,9 @@ def set_choices(self, choices: Iterable[str]):
 class ColumnOrIntegerParam(ColumnOrStringParam):
     DEFAULT_VALUE: int = 0
 
+    def set_column_choices(self, choices: Mapping[str, bool]):
+        self.set_choices([self.PREFIX + c for c, n in choices.items() if n])
+
     def __init__(
         self,
         label: str,
@@ -755,7 +763,7 @@ def set_parameter(self, key: str, value: Union[bool, int, float, str], base_dir:
             elif isinstance(param, ScalarParam):
                 param.set_value(value)
 
-    def set_column_choices(self, choices: List[str]):
+    def set_column_choices(self, choices: Mapping[str, bool]):
         logger.debug("HasSubParametersMixin.set_column_choices %s", choices)
         for p in self.params.values():
             p.set_column_choices(choices)
@@ -856,7 +864,7 @@ def get_hash_value(self):
             digest.update(p.get_hash_value().encode("utf-8"))
         return digest.hexdigest()
 
-    def set_column_choices(self, choices: List[str]):
+    def set_column_choices(self, choices: Mapping[str, bool]):
         self.param.set_column_choices(choices)
         for p in self.params:
             p.set_column_choices(choices)
@@ -897,10 +905,10 @@ def get_parameters(self, key, base_dir="."):
             yield f"{key}.{n}._label", p.label
             yield from p.get_parameters(f"{key}.{n}", base_dir)
 
-    def set_column_choices(self, choices: List[str]):
+    def set_column_choices(self, choices: Mapping[str, bool]):
         params_by_label = {p.label: p for p in self.params}
         self.params = []
-        for label in choices:
+        for label in choices.keys():
             if label in params_by_label:
                 self.params.append(params_by_label[label])
             else:
@@ -913,6 +921,11 @@ def get_column_params(self):
             yield p.label, p
 
 
+class PerNumericColumnArrayParam(PerColumnArrayParam):
+    def set_column_choices(self, choices: Mapping[str, bool]):
+        super().set_column_choices({k: True for k, v in choices.items() if v})
+
+
 class FileArrayParam(ArrayParam):
     """FileArrayParam is an ArrayParam arranged per-file.  Using this class
     really just marks it as expecting to be populated from an open file

countess/core/plugins.py

@@ -36,7 +36,7 @@
     HasSubParametersMixin,
     MultiParam,
 )
-from countess.utils.duckdb import duckdb_combine, duckdb_escape_literal
+from countess.utils.duckdb import duckdb_combine, duckdb_dtype_is_numeric, duckdb_escape_literal
 from countess.utils.files import clean_filename
 from countess.utils.pyarrow import python_type_to_arrow_dtype
 
@@ -137,7 +137,10 @@ def prepare_multi(self, ddbc: DuckDBPyConnection, sources: Mapping[str, DuckDBPy
         self.prepare(ddbc, duckdb_combine(ddbc, list(sources.values())))
 
     def prepare(self, ddbc: DuckDBPyConnection, source: Optional[DuckDBPyRelation]) -> None:
-        self.set_column_choices([] if source is None else source.columns)
+        if source is None:
+            self.set_column_choices({})
+        else:
+            self.set_column_choices({c: duckdb_dtype_is_numeric(d) for c, d in zip(source.columns, source.dtypes)})
 
     def execute_multi(
         self, ddbc: DuckDBPyConnection, sources: Mapping[str, DuckDBPyRelation], row_limit: Optional[int] = None

countess/plugins/correlation.py

@@ -4,7 +4,7 @@
 from duckdb import DuckDBPyConnection, DuckDBPyRelation
 
 from countess import VERSION
-from countess.core.parameters import ColumnChoiceParam, ColumnOrNoneChoiceParam
+from countess.core.parameters import NumericColumnChoiceParam, ColumnOrNoneChoiceParam
 from countess.core.plugins import DuckdbSimplePlugin
 from countess.utils.duckdb import duckdb_escape_identifier
 
@@ -18,8 +18,8 @@ class CorrelationPlugin(DuckdbSimplePlugin):
     link = "https://countess-project.github.io/CountESS/included-plugins/#correlation-tool"
 
     group = ColumnOrNoneChoiceParam("Group")
-    column1 = ColumnChoiceParam("Column X")
-    column2 = ColumnChoiceParam("Column Y")
+    column1 = NumericColumnChoiceParam("Column X")
+    column2 = NumericColumnChoiceParam("Column Y")
 
     columns: list[str] = []
     dataframes: list[pd.DataFrame] = []

countess/plugins/score_scale.py

@@ -8,6 +8,7 @@
     ArrayParam,
     ChoiceParam,
     ColumnChoiceParam,
+    NumericColumnChoiceParam,
     ColumnOrNoneChoiceParam,
     FloatParam,
     StringParam,
@@ -45,7 +46,7 @@ class ScoreScalingPlugin(DuckdbSimplePlugin):
     description = "Scaled Scores using variant classification"
     version = VERSION
 
-    score_col = ColumnChoiceParam("Score Column")
+    score_col = NumericColumnChoiceParam("Score Column")
     classifiers = ArrayParam("Variant Classifiers", ScaleClassParam("Class"), min_size=2, max_size=2, read_only=True)
     group_col = ColumnOrNoneChoiceParam("Group By")
 

countess/plugins/vampseq.py

@@ -4,43 +4,27 @@
 from duckdb import DuckDBPyConnection, DuckDBPyRelation
 
 from countess import VERSION
-from countess.core.parameters import (
-    ArrayParam,
-    ColumnChoiceParam,
-    FloatParam,
-    TabularMultiParam,
-    PerColumnArrayParam,
-)
+from countess.core.parameters import FloatParam, PerNumericColumnArrayParam, TabularMultiParam
 from countess.core.plugins import DuckdbSimplePlugin
-from countess.utils.duckdb import duckdb_escape_identifier, duckdb_escape_literal, duckdb_dtype_is_numeric
+from countess.utils.duckdb import duckdb_escape_identifier, duckdb_escape_literal
 
 logger = logging.getLogger(__name__)
 
 
 class CountColumnParam(TabularMultiParam):
     weight = FloatParam("Weight")
 
+
 class VampSeqScorePlugin(DuckdbSimplePlugin):
     name = "VAMP-seq Scoring"
     description = "Calculate scores from weighed bin counts"
     version = VERSION
 
-    columns = PerColumnArrayParam("Columns", CountColumnParam("Column"))
-
-    def prepare(self, ddbc: DuckDBPyConnection, source: Optional[DuckDBPyRelation]) -> None:
-        # Override prepare to only select numeric columns.
-        if source is None:
-            self.set_column_choices([])
-        else:
-            self.set_column_choices([
-                c for c, d in zip(source.columns, source.dtypes)
-                if duckdb_dtype_is_numeric(d)
-            ])
+    columns = PerNumericColumnArrayParam("Columns", CountColumnParam("Column"))
 
     def execute(
         self, ddbc: DuckDBPyConnection, source: DuckDBPyRelation, row_limit: Optional[int] = None
     ) -> Optional[DuckDBPyRelation]:
-
         weighted_columns = {
             duckdb_escape_identifier(name): duckdb_escape_literal(param.weight.value)
             for name, param in self.columns.get_column_params()
@@ -50,8 +34,8 @@ def execute(
         if not weighted_columns:
             return source
 
-        weighted_counts = ' + '.join(f"{k} * {v}" for k, v in weighted_columns.items())
-        total_counts = ' + '.join(k for k in weighted_columns.keys())
+        weighted_counts = " + ".join(f"{k} * {v}" for k, v in weighted_columns.items())
+        total_counts = " + ".join(k for k in weighted_columns.keys())
 
         proj = f"*, ({weighted_counts}) / ({total_counts}) as score"
 

tests/plugins/test_variant.py

@@ -11,7 +11,7 @@ def test_variant_ref_value():
     plugin.set_parameter("reference", "AGAAGTAGAGG")
     plugin.set_parameter("variant.seq_type", "g")
     plugin.set_parameter("variant.output", "out")
-    plugin.set_column_choices(["seq"])
+    plugin.set_column_choices({"seq": False})
 
     assert plugin.transform({"seq": "TGAAGTAGAGG"})["out"] == "g.1A>T"
     assert plugin.transform({"seq": "AGAAGTTGTGG"})["out"] == "g.[7A>T;9A>T]"
@@ -24,7 +24,7 @@ def test_variant_ref_column():
     plugin.set_parameter("reference", "— ref")
     plugin.set_parameter("variant.seq_type", "g")
     plugin.set_parameter("variant.output", "out")
-    plugin.set_column_choices(["seq", "ref"])
+    plugin.set_column_choices({"seq": False, "ref": False})
 
     assert plugin.transform({"ref": "TACACACAG", "seq": "TACAGACAG"})["out"] == "g.5C>G"
     assert plugin.transform({"ref": "ATGGTTGGTTC", "seq": "ATGGTTGGTGGTTCG"})["out"] == "g.[7_9dup;11_12insG]"
@@ -37,7 +37,7 @@ def test_variant_ref_offset():
     plugin.set_parameter("variant.offset", "— offs")
     plugin.set_parameter("variant.seq_type", "g")
     plugin.set_parameter("variant.output", "out")
-    plugin.set_column_choices(["seq", "offs"])
+    plugin.set_column_choices({"seq": False, "offs": True})
 
     assert plugin.transform({"seq": "TGAAGTAGAGG", "offs": "0"})["out"] == "g.1A>T"
     assert plugin.transform({"seq": "AGAAGTTGTGG", "offs": "10"})["out"] == "g.[17A>T;19A>T]"
@@ -65,7 +65,7 @@ def test_variant_ref_offset_minus():
     plugin.set_parameter("variant.seq_type", "g")
     plugin.set_parameter("variant.minus_strand", True)
     plugin.set_parameter("variant.output", "out")
-    plugin.set_column_choices(["seq", "offs"])
+    plugin.set_column_choices({"seq": False, "offs": True})
 
     assert plugin.transform({"seq": "TGAAGTAGAGG"})["out"] == "g.1011T>A"
     assert plugin.transform({"seq": "AGAAGTTGTGG"})["out"] == "g.[1003T>A;1005T>A]"
@@ -79,7 +79,7 @@ def test_variant_too_many():
     plugin.set_parameter("variant.seq_type", "g")
     plugin.set_parameter("variant.output", "out")
     plugin.set_parameter("variant.maxlen", 2)
-    plugin.set_column_choices(["seq", "offs"])
+    plugin.set_column_choices({"seq": False, "offs": True})
 
     assert plugin.transform({"seq": "TGAAGTAGAGG"})["out"] == "g.1A>T"
     assert plugin.transform({"seq": "AGAAGTTGTGG"})["out"] == "g.[7A>T;9A>T]"

tests/test_parameters.py

@@ -226,7 +226,7 @@ def test_coindex():
 def test_columnorintegerparam():
     df = pd.DataFrame([[1, 2], [3, 4]], columns=["a", "b"])
     cp = ColumnOrIntegerParam("x")
-    cp.set_column_choices(["a", "b"])
+    cp.set_column_choices({"a": True, "b": True})
 
     assert cp.get_column_name() is None
 
@@ -240,7 +240,7 @@ def test_columnorintegerparam():
     assert cp.get_column_name() == "a"
     assert isinstance(cp.get_column_or_value(df, False), pd.Series)
 
-    cp.set_column_choices(["c", "d"])
+    cp.set_column_choices({"c": True, "d": True})
     assert cp.choice is None
 
     cp.value = "hello"
@@ -250,7 +250,7 @@ def test_columnorintegerparam():
 def test_columngroup():
     df = pd.DataFrame([], columns=["one_two", "one_three", "two_one", "two_two", "two_three", "three_four_five"])
     cp = ColumnGroupOrNoneChoiceParam("x")
-    cp.set_column_choices(df.columns)
+    cp.set_column_choices({c: True for c in df.columns})
     assert cp.is_none()
     assert "one_*" in cp.choices
     assert "two_*" in cp.choices
@@ -330,11 +330,11 @@ def test_pcap():
     pp = IntegerParam("x")
     ap = PerColumnArrayParam("y", param=pp)
 
-    ap.set_column_choices(["a", "b", "c"])
+    ap.set_column_choices({"a": True, "b": True, "c": True })
     assert len(ap) == 3
     apa, apb, apc = list(ap)
 
-    ap.set_column_choices(["c", "d", "b", "a"])
+    ap.set_column_choices({"c": True, "d": True, "b": True, "a": True})
     assert len(ap) == 4
     assert ap[0] is apc
     assert ap[2] is apb