Merge pull request #267 from CITCOM-project/surrogate-fix

Fix bug in categorical gene types

Author: rsomers1998

causal_testing/surrogate/surrogate_search_algorithms.py

@@ -99,10 +99,14 @@ def create_gene_types(
             rel_split = str(relationship).split(" ")
 
             if rel_split[0] in var_space:
+                datatype = specification.scenario.variables.get(rel_split[0]).datatype
                 if rel_split[1] == ">=":
-                    var_space[rel_split[0]]["low"] = int(rel_split[2])
+                    var_space[rel_split[0]]["low"] = datatype(rel_split[2])
                 elif rel_split[1] == "<=":
-                    var_space[rel_split[0]]["high"] = int(rel_split[2])
+                    if datatype == int:
+                        var_space[rel_split[0]]["high"] = int(rel_split[2]) + 1
+                    else:
+                        var_space[rel_split[0]]["high"] = datatype(rel_split[2])
 
         gene_space = []
         gene_space.append(var_space[surrogate_model.treatment])

tests/surrogate_tests/test_causal_surrogate_assisted.py

@@ -3,7 +3,7 @@
 from causal_testing.specification.causal_dag import CausalDAG
 from causal_testing.specification.causal_specification import CausalSpecification
 from causal_testing.specification.scenario import Scenario
-from causal_testing.specification.variable import Input
+from causal_testing.specification.variable import Input, Output
 from causal_testing.surrogate.causal_surrogate_assisted import SimulationResult, CausalSurrogateAssistedTestCase, Simulator
 from causal_testing.surrogate.surrogate_search_algorithms import GeneticSearchAlgorithm
 from causal_testing.testing.estimators import CubicSplineRegressionEstimator
@@ -56,9 +56,9 @@ def test_surrogate_model_generation(self):
 
         causal_dag = CausalDAG(self.dag_dot_path)
         z = Input("Z", int)
-        x = Input("X", int)
+        x = Input("X", float)
         m = Input("M", int)
-        y = Input("Y", int)
+        y = Output("Y", float)
         scenario = Scenario(variables={z, x, m, y})
         specification = CausalSpecification(scenario, causal_dag)
 
@@ -75,9 +75,9 @@ def test_causal_surrogate_assisted_execution(self):
 
         causal_dag = CausalDAG(self.dag_dot_path)
         z = Input("Z", int)
-        x = Input("X", int)
+        x = Input("X", float)
         m = Input("M", int)
-        y = Input("Y", int)
+        y = Output("Y", float)
         scenario = Scenario(variables={z, x, m, y}, constraints={
             z <= 0, z >= 3,
             x <= 0, x >= 3,
@@ -107,9 +107,9 @@ def test_causal_surrogate_assisted_execution_failure(self):
 
         causal_dag = CausalDAG(self.dag_dot_path)
         z = Input("Z", int)
-        x = Input("X", int)
+        x = Input("X", float)
         m = Input("M", int)
-        y = Input("Y", int)
+        y = Output("Y", float)
         scenario = Scenario(variables={z, x, m, y}, constraints={
             z <= 0, z >= 3,
             x <= 0, x >= 3,
@@ -139,9 +139,9 @@ def test_causal_surrogate_assisted_execution_custom_aggregator(self):
 
         causal_dag = CausalDAG(self.dag_dot_path)
         z = Input("Z", int)
-        x = Input("X", int)
+        x = Input("X", float)
         m = Input("M", int)
-        y = Input("Y", int)
+        y = Output("Y", float)
         scenario = Scenario(variables={z, x, m, y}, constraints={
             z <= 0, z >= 3,
             x <= 0, x >= 3,
@@ -172,9 +172,9 @@ def test_causal_surrogate_assisted_execution_incorrect_search_config(self):
 
         causal_dag = CausalDAG(self.dag_dot_path)
         z = Input("Z", int)
-        x = Input("X", int)
+        x = Input("X", float)
         m = Input("M", int)
-        y = Input("Y", int)
+        y = Output("Y", float)
         scenario = Scenario(variables={z, x, m, y}, constraints={
             z <= 0, z >= 3,
             x <= 0, x >= 3,