Part way to fixing categoricals

Author: jmafoster1

causal_testing/generation/abstract_causal_test_case.py

@@ -93,7 +93,6 @@ def _generate_concrete_tests(
             var = self.scenario.variables[name]
             samples[var.name] = lhsmdu.inverseTransformSample(var.distribution, samples[var.name])
 
-        print(samples["ego_vehicle"])
         for index, row in samples.iterrows():
             optimizer = z3.Optimize()
             for c in self.scenario.constraints:

causal_testing/specification/variable.py

@@ -22,7 +22,7 @@ def z3_types(datatype):
     if datatype in types:
         return types[datatype]
     if issubclass(datatype, Enum):
-        dtype, _ = EnumSort(datatype.__name__, [x.name for x in datatype])
+        dtype, _ = EnumSort(datatype.__name__, [x.value for x in datatype])
         return lambda x: Const(x, dtype)
     if hasattr(datatype, "to_z3"):
         return datatype.to_z3()
@@ -162,15 +162,14 @@ def cast(self, val: Any) -> T:
         if (isinstance(val, float) or isinstance(val, int)) and (self.datatype == int or self.datatype == float):
             return self.datatype(val)
         if issubclass(self.datatype, Enum) and isinstance(val, DatatypeRef):
-            return self.datatype[str(val)]
+            return self.datatype(str(val))
         return self.datatype(str(val))
 
     def z3_val(self, z3_var, val: Any) -> T:
         native_val = self.cast(val)
-        print(val, type(val), native_val, type(native_val))
         if isinstance(native_val, Enum):
             values = [z3_var.sort().constructor(c)() for c in range(z3_var.sort().num_constructors())]
-            values = [v for v in values if type(val)[str(v)] == val]
+            values = [v for v in values if type(val)(str(v)) == val]
             assert len(values) == 1, f"Expected {values} to be length 1"
             return values[0]
         return native_val

causal_testing/testing/estimators.py

@@ -453,6 +453,9 @@ def _run_linear_regression(self) -> RegressionResultsWrapper:
         cols += [x for x in self.adjustment_set if x not in cols]
         treatment_and_adjustments_cols = reduced_df[cols + ["Intercept"]]
         outcome_col = reduced_df[list(self.outcome)]
+        for col in treatment_and_adjustments_cols:
+            if str(treatment_and_adjustments_cols.dtypes[col]) == "object":
+                treatment_and_adjustments_cols = pd.get_dummies(treatment_and_adjustments_cols, columns=[col], drop_first=True)
         regression = sm.OLS(outcome_col, treatment_and_adjustments_cols)
         model = regression.fit()
         return model