diff --git a/.gitignore b/.gitignore index 15201ac..cf5fd3f 100644 --- a/.gitignore +++ b/.gitignore @@ -82,6 +82,9 @@ target/ profile_default/ ipython_config.py +# Agents +.pi-subagents/ + # pyenv # For a library or package, you might want to ignore these files since the code is # intended to run in multiple environments; otherwise, check them in: diff --git a/src/orbital/translation/steps/div.py b/src/orbital/translation/steps/div.py index c054363..e430c7c 100644 --- a/src/orbital/translation/steps/div.py +++ b/src/orbital/translation/steps/div.py @@ -1,11 +1,7 @@ """Defines the translation step for the Div operation.""" -import typing - -import ibis - from ..translator import Translator -from ..variables import NumericVariablesGroup, ValueVariablesGroup, VariablesGroup +from ..variables import ValueVariablesGroup class DivTranslator(Translator): @@ -29,8 +25,10 @@ def process(self) -> None: """Performs the translation and set the output variable.""" # https://onnx.ai/onnx/operators/onnx__Div.html - left_keys, left_values = self._operand_values(self.inputs[0]) - right_keys, right_values = self._operand_values(self.inputs[1]) + left_keys, left_values = self._variables.consume_operand_values(self.inputs[0]) + right_keys, right_values = self._variables.consume_operand_values( + self.inputs[1] + ) # The first operand that is a group dictates the width of the result # and, at the very end, the names of the resulting columns. @@ -68,30 +66,3 @@ def process(self) -> None: } ) ) - - def _operand_values( - self, name: str - ) -> tuple[typing.Optional[list[str]], list[ibis.expr.types.NumericValue]]: - """Consume an operand and return its column names and its values. - - The column names are ``None`` for anything that is not a group of columns, - which is what tells apart a group of one column from a plain single value. - - :param name: Name of the variable or constant to consume. - """ - # Classifiers normalize the probabilities of each class with - # Div(scores, ReduceSum(Abs(scores))), so an operand can also be a - # column computed by a previous node instead of a constant. - operand = self._variables.consume(name) - if isinstance(operand, VariablesGroup): - group = NumericVariablesGroup(operand) - return list(group.keys()), list(group.values()) - - values = [] - for value in operand if isinstance(operand, (list, tuple)) else [operand]: - if isinstance(value, (int, float)): - value = ibis.literal(value) - elif not isinstance(value, ibis.expr.types.NumericValue): - raise ValueError("Div: the operands must be numeric values.") - values.append(value) - return None, typing.cast(list[ibis.expr.types.NumericValue], values) diff --git a/src/orbital/translation/steps/sub.py b/src/orbital/translation/steps/sub.py index 85ae7b3..5e2e1e9 100644 --- a/src/orbital/translation/steps/sub.py +++ b/src/orbital/translation/steps/sub.py @@ -1,70 +1,68 @@ """Implementation of the Sub operator.""" -import typing - -import ibis - -from orbital.translation.variables import ( - NumericVariablesGroup, - ValueVariablesGroup, - VariablesGroup, -) - from ..translator import Translator +from ..variables import ValueVariablesGroup class SubTranslator(Translator): """Processes a Sub node and updates the variables with the output expression. - Given the node to translate, the variables and constants available for - the translation context, generates a query expression that processes - the input variables and produces a new output variable that computes - based on the Sub operation. + Both operands are treated symmetrically: each one can be a group of columns, + a single column, a constant scalar or a constant list of values. + + When any of the two operands is a group of columns, the result is a group of + columns too and it borrows its column names from the first operand that is a + group, as those names end up being the names of the resulting SQL columns. + That group also dictates the width of the result: any other operand must + either provide exactly as many values, or a single value that is subtracted + from (or subtracts) every column of the group. + + When neither operand is a group, both must be single values and the result + is a single column. """ def process(self) -> None: """Performs the translation and set the output variable.""" # https://onnx.ai/onnx/operators/onnx__Sub.html - assert len(self._inputs) == 2, "The Sub node must have exactly 2 inputs." - - first_operand = self._variables.consume(self._inputs[0]) - second_operand = self._variables.get_initializer_value(self._inputs[1]) - if second_operand is None or not isinstance(second_operand, (list, tuple)): - raise NotImplementedError( - "Sub: Second input (divisor) must be a constant list." - ) - type_check_var = first_operand - if isinstance(type_check_var, dict): - type_check_var = next(iter(type_check_var.values()), None) - if not isinstance(type_check_var, ibis.expr.types.NumericValue): - raise ValueError("Sub: The first operand must be a numeric value.") + left_keys, left_values = self._variables.consume_operand_values(self.inputs[0]) + right_keys, right_values = self._variables.consume_operand_values( + self.inputs[1] + ) - sub_values = list(second_operand) - if isinstance(first_operand, VariablesGroup): - first_operand = NumericVariablesGroup(first_operand) - struct_fields = list(first_operand.keys()) - assert len(sub_values) == len(struct_fields), ( - f"The number of values in the initializer ({len(sub_values)}) must match the number of fields ({len(struct_fields)}" - ) - self.set_output( - ValueVariablesGroup( - { - field: ( - self._optimizer.fold_operation( - first_operand[field] - sub_values[i] - ) - ) - for i, field in enumerate(struct_fields) - } + # The first operand that is a group dictates the width of the result + # and, at the very end, the names of the resulting columns. + keys = left_keys if left_keys is not None else right_keys + if keys is None: + # Simple case, no columns group involved, so we just subtract the two values. + if len(left_values) != 1 or len(right_values) != 1: + raise ValueError( + "Sub: when no operand is a group of columns, each operand must contain only one value." ) + self.set_output( + self._optimizer.fold_operation(left_values[0] - right_values[0]) ) - else: - if len(sub_values) != 1: + return + + for values in (left_values, right_values): + if len(values) not in (1, len(keys)): raise ValueError( - "When the first operand is a single column, the second operand must contain exactly 1 value" + "Sub: the number of values of each operand must match the number of columns of the resulting group." ) - first_operand = typing.cast(ibis.expr.types.NumericValue, first_operand) - self.set_output( - self._optimizer.fold_operation(first_operand - sub_values[0]) + + # A single value is shared by all columns of the resulting group. + if len(left_values) == 1: + left_values = left_values * len(keys) + if len(right_values) == 1: + right_values = right_values * len(keys) + + self.set_output( + ValueVariablesGroup( + { + key: self._optimizer.fold_operation(left_value - right_value) + for key, left_value, right_value in zip( + keys, left_values, right_values + ) + } ) + ) diff --git a/src/orbital/translation/variables.py b/src/orbital/translation/variables.py index d074dea..c6fb2ed 100644 --- a/src/orbital/translation/variables.py +++ b/src/orbital/translation/variables.py @@ -144,6 +144,33 @@ def consume( self._consumed.add(name) return self._variables[name] + def consume_operand_values( + self, name: str + ) -> tuple[typing.Optional[list[str]], list[ibis.expr.types.NumericValue]]: + """Consume an operand of a mathematical operation as column names and values. + + The column names are ``None`` for anything that is not a group of columns, + which is what tells apart a group of one column from a plain single value. + + :param name: Name of the variable or constant to consume. + """ + # Classifiers normalize the probabilities of each class with + # Div(scores, ReduceSum(Abs(scores))) for example, so an operand can + # also be a column computed by a previous node instead of a constant. + operand = self.consume(name) + if isinstance(operand, VariablesGroup): + group = NumericVariablesGroup(operand) + return list(group.keys()), list(group.values()) + + values = [] + for value in operand if isinstance(operand, (list, tuple)) else [operand]: + if isinstance(value, (int, float)): + value = ibis.literal(value) + elif not isinstance(value, ibis.expr.types.NumericValue): + raise ValueError("The operands must be numeric values.") + values.append(value) + return None, typing.cast(list[ibis.expr.types.NumericValue], values) + def peek_variable( self, name: str, default: typing.Optional[ibis.Expr] = None ) -> typing.Union[ibis.Expr, VariablesGroup, None]: diff --git a/tests/test_pipeline_e2e.py b/tests/test_pipeline_e2e.py index 57f1b27..92b3683 100644 --- a/tests/test_pipeline_e2e.py +++ b/tests/test_pipeline_e2e.py @@ -3,7 +3,7 @@ from sklearn.compose import ColumnTransformer from sklearn.feature_selection import SelectKBest, f_regression from sklearn.linear_model import ElasticNet, LinearRegression, LogisticRegression -from sklearn.neural_network import MLPRegressor +from sklearn.neural_network import MLPClassifier, MLPRegressor from sklearn.pipeline import Pipeline from sklearn.preprocessing import MinMaxScaler, OneHotEncoder, StandardScaler @@ -272,3 +272,50 @@ def test_tanh_mlp_regression(self, iris_data, db_connection): rtol=1e-4, atol=1e-4, ) + + def test_binary_mlp_classifier(self, iris_data, db_connection): + """Test a binary MLP classifier, whose negative class is Sub(1.0, probability).""" + df, feature_names = iris_data + conn, dialect = db_connection + + binary_df = df[df["target"].isin([0, 1])].copy() + + sklearn_pipeline = Pipeline( + [ + ("scaler", StandardScaler()), + ( + "mlp", + MLPClassifier( + hidden_layer_sizes=(8,), + # lbfgs converges on this small dataset, + # avoiding ConvergenceWarning noise from adam. + solver="lbfgs", + max_iter=500, + random_state=0, + ), + ), + ] + ) + + X = binary_df[feature_names] + y = binary_df["target"] + sklearn_pipeline.fit(X, y) + sklearn_proba = pd.DataFrame( + sklearn_pipeline.predict_proba(X), + columns=sklearn_pipeline.classes_, + index=binary_df.index, + ) + + features = {fname: types.FloatColumnType() for fname in feature_names} + parsed_pipeline = orbital.parse_pipeline(sklearn_pipeline, features=features) + + sql = orbital.export_sql("data", parsed_pipeline, dialect=dialect) + sql_results = execute_sql(sql, conn, dialect, binary_df) + + for class_label in sklearn_pipeline.classes_: + np.testing.assert_allclose( + sql_results[f"output_probability.{class_label}"].values.flatten(), + sklearn_proba[class_label].values.flatten(), + rtol=1e-4, + atol=1e-4, + ) diff --git a/tests/test_pipeline_steps.py b/tests/test_pipeline_steps.py index 176c10b..e866129 100644 --- a/tests/test_pipeline_steps.py +++ b/tests/test_pipeline_steps.py @@ -761,60 +761,76 @@ def test_sub_group_columns(self): assert list(backend.execute(result["col_a"])) == [7.0, 17.0, 27.0] assert list(backend.execute(result["col_b"])) == [50.0, 150.0, 250.0] - def test_sub_invalid_non_numeric(self): - """Test SubTranslator raises error for non-numeric operand.""" - table = ibis.memtable({"input": [1.0, 2.0, 3.0]}) + def test_sub_group_columns_broadcast_single_value(self): + """Test SubTranslator with a group of columns and a single value.""" + table = ibis.memtable( + { + "col_a": [10.0, 20.0, 30.0], + "col_b": [100.0, 200.0, 300.0], + } + ) model = onnx.parser.parse_graph(""" agraph (float[N] input) => (float[N] output) - + { output = Sub(input, sub_value) } """) - variables = GraphVariables(table, model) - variables["input"] = "not_a_numeric_value" # type: ignore[assignment] + variables = GraphVariables(ibis.memtable({"input": [1.0]}), model) + variables["input"] = NumericVariablesGroup( + { + "col_a": table["col_a"], + "col_b": table["col_b"], + } + ) translator = SubTranslator( table, model.node[0], variables, self.optimizer, TranslationOptions() ) + translator.process() - with pytest.raises(ValueError, match="first operand must be a numeric value"): - translator.process() + result = variables.peek_variable("output") + assert isinstance(result, ValueVariablesGroup) - def test_sub_single_column_requires_single_value(self): - """Test SubTranslator raises error when single column given multiple values.""" - table = ibis.memtable({"input": [1.0, 2.0, 3.0]}) + backend = ibis.duckdb.connect() + assert list(backend.execute(result["col_a"])) == [0.0, 10.0, 20.0] + assert list(backend.execute(result["col_b"])) == [90.0, 190.0, 290.0] + + def test_sub_constant_minus_column(self): + """Test SubTranslator computes the complement of a column. + + Binary classifiers export the negative class probability + as Sub(1.0, positive_class_probability). + """ + table = ibis.memtable({"input": [0.25, 0.5, 0.75]}) model = onnx.parser.parse_graph(""" agraph (float[N] input) => (float[N] output) - + { - output = Sub(input, sub_values) + output = Sub(unity, input) } """) variables = GraphVariables(table, model) - translator = SubTranslator( table, model.node[0], variables, self.optimizer, TranslationOptions() ) + translator.process() - with pytest.raises(ValueError, match="must contain exactly 1 value"): - translator.process() + result = variables.peek_variable("output") - def test_sub_mismatched_column_count(self): - """Test SubTranslator raises error when column count doesn't match.""" - table = ibis.memtable( - { - "col_a": [1.0, 2.0, 3.0], - "col_b": [10.0, 20.0, 30.0], - } - ) + backend = ibis.duckdb.connect() + assert list(backend.execute(result)) == [0.75, 0.5, 0.25] + + def test_sub_constant_minus_group_columns(self): + """Test SubTranslator computes the complement of a group of columns.""" + table = ibis.memtable({"col_a": [0.25, 0.5], "col_b": [0.1, 0.2]}) model = onnx.parser.parse_graph(""" agraph (float[N] input) => (float[N] output) - + { - output = Sub(input, sub_values) + output = Sub(unity, input) } """) @@ -829,16 +845,57 @@ def test_sub_mismatched_column_count(self): translator = SubTranslator( table, model.node[0], variables, self.optimizer, TranslationOptions() ) + translator.process() - with pytest.raises(AssertionError, match="must match the number of fields"): - translator.process() + result = variables.peek_variable("output") + assert isinstance(result, ValueVariablesGroup) + assert list(result.keys()) == ["col_a", "col_b"] - def test_sub_second_operand_not_constant(self): - """Test SubTranslator raises error when second operand is not a constant.""" + backend = ibis.duckdb.connect() + assert list(backend.execute(result["col_a"])) == [0.75, 0.5] + assert list(backend.execute(result["col_b"])) == [0.9, 0.8] + + def test_sub_column_minus_computed_column(self): + """Test SubTranslator subtracts a column computed by a previous node.""" table = ibis.memtable( { - "input": [1.0, 2.0, 3.0], - "other": [5.0, 5.0, 5.0], + "input": [10.0, 20.0, 30.0], + "other": [1.0, 2.0, 3.0], + } + ) + model = onnx.parser.parse_graph(""" + agraph (float[N] input, float[N] other) => (float[N] output) { + output = Sub(input, other) + } + """) + + variables = GraphVariables(table, model) + translator = SubTranslator( + table, model.node[0], variables, self.optimizer, TranslationOptions() + ) + translator.process() + + result = variables.peek_variable("output") + + # The subtrahend must be consumed, or it would leak into the results. + assert "other" not in variables + + backend = ibis.duckdb.connect() + assert list(backend.execute(result)) == [9.0, 18.0, 27.0] + + def test_sub_group_columns_minus_group_columns(self): + """Test SubTranslator subtracts two groups of columns. + + The resulting columns are named after the first group. + """ + table = ibis.memtable( + { + "input": [1.0, 1.0], + "other": [1.0, 1.0], + "minuend_a": [10.0, 20.0], + "minuend_b": [100.0, 200.0], + "subtrahend_a": [1.0, 2.0], + "subtrahend_b": [10.0, 20.0], } ) model = onnx.parser.parse_graph(""" @@ -848,12 +905,81 @@ def test_sub_second_operand_not_constant(self): """) variables = GraphVariables(table, model) + variables["input"] = NumericVariablesGroup( + { + "col_a": table["minuend_a"], + "col_b": table["minuend_b"], + } + ) + variables["other"] = NumericVariablesGroup( + { + "sub_a": table["subtrahend_a"], + "sub_b": table["subtrahend_b"], + } + ) translator = SubTranslator( table, model.node[0], variables, self.optimizer, TranslationOptions() ) + translator.process() - with pytest.raises(NotImplementedError, match="must be a constant list"): + result = variables.peek_variable("output") + assert isinstance(result, ValueVariablesGroup) + assert list(result.keys()) == ["col_a", "col_b"] + + backend = ibis.duckdb.connect() + assert list(backend.execute(result["col_a"])) == [9.0, 18.0] + assert list(backend.execute(result["col_b"])) == [90.0, 180.0] + + def test_sub_single_column_requires_single_value(self): + """Test SubTranslator raises error when single column given multiple values.""" + table = ibis.memtable({"input": [1.0, 2.0, 3.0]}) + model = onnx.parser.parse_graph(""" + agraph (float[N] input) => (float[N] output) + + { + output = Sub(input, sub_values) + } + """) + + variables = GraphVariables(table, model) + + translator = SubTranslator( + table, model.node[0], variables, self.optimizer, TranslationOptions() + ) + + with pytest.raises(ValueError, match="must contain only one value"): + translator.process() + + def test_sub_mismatched_column_count(self): + """Test SubTranslator raises error when column count doesn't match values.""" + table = ibis.memtable( + { + "col_a": [1.0, 2.0, 3.0], + "col_b": [10.0, 20.0, 30.0], + } + ) + model = onnx.parser.parse_graph(""" + agraph (float[N] input) => (float[N] output) + + { + output = Sub(input, sub_values) + } + """) + + variables = GraphVariables(ibis.memtable({"input": [1.0]}), model) + variables["input"] = NumericVariablesGroup( + { + "col_a": table["col_a"], + "col_b": table["col_b"], + } + ) + + translator = SubTranslator( + table, model.node[0], variables, self.optimizer, TranslationOptions() + ) + + with pytest.raises(ValueError, match="must match the number of columns"): translator.process()