test_core_pipeline.py

# Copyright 2019 IBM Corporation
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import pickle
import traceback
import typing
import unittest

import numpy as np
import sklearn.datasets
import sklearn.pipeline
from sklearn.feature_selection import SelectKBest as SkSelectKBest
from sklearn.metrics import accuracy_score, r2_score
from sklearn.model_selection import train_test_split

import lale.datasets.openml
import lale.helpers
from lale.helpers import import_from_sklearn_pipeline
from lale.lib.lale import ConcatFeatures, NoOp
from lale.lib.sklearn import (
    PCA,
    AdaBoostClassifier,
    GaussianNB,
    IsolationForest,
    KNeighborsClassifier,
    LinearRegression,
    LinearSVC,
    LogisticRegression,
    Nystroem,
    OneHotEncoder,
    PassiveAggressiveClassifier,
    SelectKBest,
    SGDClassifier,
    StandardScaler,
)
from lale.lib.xgboost import XGBClassifier
from lale.operators import (
    TrainableIndividualOp,
    TrainablePipeline,
    TrainedIndividualOp,
    TrainedPipeline,
    make_choice,
    make_pipeline,
    make_union,
)


class TestCreation(unittest.TestCase):
    def setUp(self):
        data = sklearn.datasets.load_iris()
        X, y = data.data, data.target
        self.X_train, self.X_test, self.y_train, self.y_test = train_test_split(X, y)

    def test_pipeline_create(self):
        from lale.operators import Pipeline

        pipeline = Pipeline(([("pca1", PCA()), ("lr1", LogisticRegression())]))
        trained = pipeline.fit(self.X_train, self.y_train)
        predictions = trained.predict(self.X_test)
        accuracy_score(self.y_test, predictions)

    def test_pipeline_create_trainable(self):
        from lale.lib.sklearn import Pipeline as SkPipeline

        pipeline = SkPipeline(steps=[("pca1", PCA()), ("lr1", LogisticRegression())])
        self.assertIsInstance(pipeline, TrainableIndividualOp)
        trained = pipeline.fit(self.X_train, self.y_train)
        pca_trained, lr_trained = [op for _, op in trained.hyperparams()["steps"]]
        self.assertIsInstance(pca_trained, TrainedIndividualOp)
        self.assertIsInstance(lr_trained, TrainedIndividualOp)
        predictions = trained.predict(self.X_test)
        accuracy_score(self.y_test, predictions)

    def test_pipeline_create_trained(self):
        from lale.lib.sklearn import Pipeline as SkPipeline

        orig_trainable = PCA() >> LogisticRegression()
        orig_trained = orig_trainable.fit(self.X_train, self.y_train)
        self.assertIsInstance(orig_trained, TrainedPipeline)
        pca_trained, lr_trained = orig_trained.steps_list()
        pre_trained = SkPipeline(steps=[("pca1", pca_trained), ("lr1", lr_trained)])
        self.assertIsInstance(pre_trained, TrainedIndividualOp)
        predictions = pre_trained.predict(self.X_test)
        accuracy_score(self.y_test, predictions)

    def test_pipeline_clone(self):
        from sklearn.base import clone

        from lale.operators import Pipeline

        pipeline = Pipeline(([("pca1", PCA()), ("lr1", LogisticRegression())]))
        trained = pipeline.fit(self.X_train, self.y_train)
        predictions = trained.predict(self.X_test)
        orig_acc = accuracy_score(self.y_test, predictions)

        cloned_pipeline = clone(pipeline)
        trained = cloned_pipeline.fit(self.X_train, self.y_train)
        predictions = trained.predict(self.X_test)
        cloned_acc = accuracy_score(self.y_test, predictions)
        self.assertEqual(orig_acc, cloned_acc)

    def test_make_pipeline(self):
        tfm = PCA(n_components=10)
        clf = LogisticRegression(random_state=42)
        trainable = make_pipeline(tfm, clf)
        digits = sklearn.datasets.load_digits()
        trained = trainable.fit(digits.data, digits.target)
        _ = trained.predict(digits.data)

    def test_compose2(self):
        tfm = PCA(n_components=10)
        clf = LogisticRegression(random_state=42)
        trainable = tfm >> clf
        digits = sklearn.datasets.load_digits()
        trained = trainable.fit(digits.data, digits.target)
        _ = trained.predict(digits.data)

    def test_compose3(self):
        nys = Nystroem(n_components=15)
        pca = PCA(n_components=10)
        lr = LogisticRegression(random_state=42)
        trainable = nys >> pca >> lr
        digits = sklearn.datasets.load_digits()
        trained = trainable.fit(digits.data, digits.target)
        _ = trained.predict(digits.data)

    def test_pca_nys_lr(self):
        nys = Nystroem(n_components=15)
        pca = PCA(n_components=10)
        lr = LogisticRegression(random_state=42)
        trainable = make_union(nys, pca) >> lr
        digits = sklearn.datasets.load_digits()
        trained = trainable.fit(digits.data, digits.target)
        _ = trained.predict(digits.data)

    def test_compose4(self):
        digits = sklearn.datasets.load_digits()
        _ = digits
        ohe = OneHotEncoder(handle_unknown=OneHotEncoder.enum.handle_unknown.ignore)
        ohe.get_params()
        no_op = NoOp()
        pca = PCA()
        nys = Nystroem()
        lr = LogisticRegression()
        knn = KNeighborsClassifier()
        step1 = ohe | no_op
        step2 = pca | nys
        step3 = lr | knn
        model_plan = step1 >> step2 >> step3
        _ = model_plan
        # TODO: optimize on this plan and then fit and predict

    def test_compose5(self):
        ohe = OneHotEncoder(handle_unknown=OneHotEncoder.enum.handle_unknown.ignore)
        digits = sklearn.datasets.load_digits()
        lr = LogisticRegression()
        lr_trained = lr.fit(digits.data, digits.target)
        lr_trained.predict(digits.data)
        pipeline1 = ohe >> lr
        pipeline1_trained = pipeline1.fit(digits.data, digits.target)
        pipeline1_trained.predict(digits.data)

    def test_compare_with_sklearn(self):
        tfm = PCA()
        clf = LogisticRegression(
            LogisticRegression.enum.solver.saga,
            LogisticRegression.enum.multi_class.auto,
        )
        trainable = make_pipeline(tfm, clf)
        digits = sklearn.datasets.load_digits()
        trained = trainable.fit(digits.data, digits.target)
        predicted = trained.predict(digits.data)
        from sklearn.decomposition import PCA as SklearnPCA
        from sklearn.linear_model import LogisticRegression as SklearnLR

        sklearn_pipeline = sklearn.pipeline.make_pipeline(
            SklearnPCA(), SklearnLR(solver="saga", multi_class="auto")
        )
        sklearn_pipeline.fit(digits.data, digits.target)
        predicted_sklearn = sklearn_pipeline.predict(digits.data)

        lale_score = accuracy_score(digits.target, predicted)
        scikit_score = accuracy_score(digits.target, predicted_sklearn)
        self.assertEqual(lale_score, scikit_score)


class TestImportExport(unittest.TestCase):
    def setUp(self):
        data = sklearn.datasets.load_iris()
        X, y = data.data, data.target
        self.X_train, self.X_test, self.y_train, self.y_test = train_test_split(X, y)

    @classmethod
    def get_sklearn_params(cls, op):
        lale_sklearn_impl = op._impl_instance()
        wrapped_model = getattr(lale_sklearn_impl, "_wrapped_model", None)
        if wrapped_model is not None:
            lale_sklearn_impl = wrapped_model
        return lale_sklearn_impl.get_params()

    def assert_equal_predictions(self, pipeline1, pipeline2):
        trained = pipeline1.fit(self.X_train, self.y_train)
        predictions1 = trained.predict(self.X_test)

        trained = pipeline2.fit(self.X_train, self.y_train)
        predictions2 = trained.predict(self.X_test)
        for i, p1 in enumerate(predictions1):
            self.assertEqual(p1, predictions2[i])

    def test_import_from_sklearn_pipeline(self):
        from sklearn.feature_selection import f_regression
        from sklearn.pipeline import Pipeline
        from sklearn.svm import SVC as SklearnSVC

        anova_filter = SkSelectKBest(f_regression, k=3)
        clf = SklearnSVC(kernel="linear")
        sklearn_pipeline = Pipeline([("anova", anova_filter), ("svc", clf)])
        lale_pipeline = typing.cast(
            TrainablePipeline,
            import_from_sklearn_pipeline(sklearn_pipeline),
        )
        for i, pipeline_step in enumerate(sklearn_pipeline.named_steps):
            sklearn_step_params = sklearn_pipeline.named_steps[
                pipeline_step
            ].get_params()
            lale_sklearn_params = self.get_sklearn_params(lale_pipeline.steps_list()[i])
            self.assertEqual(sklearn_step_params, lale_sklearn_params)
        self.assert_equal_predictions(sklearn_pipeline, lale_pipeline)

    def test_import_from_sklearn_pipeline1(self):
        from sklearn.decomposition import PCA as SklearnPCA
        from sklearn.neighbors import KNeighborsClassifier as SklearnKNN

        sklearn_pipeline = sklearn.pipeline.make_pipeline(
            SklearnPCA(n_components=3), SklearnKNN()
        )
        lale_pipeline = typing.cast(
            TrainablePipeline,
            import_from_sklearn_pipeline(sklearn_pipeline),
        )
        for i, pipeline_step in enumerate(sklearn_pipeline.named_steps):
            sklearn_step_params = sklearn_pipeline.named_steps[
                pipeline_step
            ].get_params()
            lale_sklearn_params = self.get_sklearn_params(lale_pipeline.steps_list()[i])
            self.assertEqual(sklearn_step_params, lale_sklearn_params)
        self.assert_equal_predictions(sklearn_pipeline, lale_pipeline)

    def test_import_from_sklearn_pipeline_feature_union(self):
        from sklearn.decomposition import PCA as SklearnPCA
        from sklearn.kernel_approximation import Nystroem as SklearnNystroem
        from sklearn.neighbors import KNeighborsClassifier as SklearnKNN
        from sklearn.pipeline import FeatureUnion

        union = FeatureUnion(
            [
                ("pca", SklearnPCA(n_components=1)),
                ("nys", SklearnNystroem(n_components=2, random_state=42)),
            ]
        )
        sklearn_pipeline = sklearn.pipeline.make_pipeline(union, SklearnKNN())
        lale_pipeline = typing.cast(
            TrainablePipeline,
            import_from_sklearn_pipeline(sklearn_pipeline),
        )
        self.assertEqual(len(lale_pipeline.edges()), 3)

        self.assertIsInstance(lale_pipeline.edges()[0][0], PCA)  # type: ignore
        self.assertIsInstance(lale_pipeline.edges()[0][1], ConcatFeatures)  # type: ignore
        self.assertIsInstance(lale_pipeline.edges()[1][0], Nystroem)  # type: ignore
        self.assertIsInstance(lale_pipeline.edges()[1][1], ConcatFeatures)  # type: ignore
        self.assertIsInstance(lale_pipeline.edges()[2][0], ConcatFeatures)  # type: ignore
        self.assertIsInstance(lale_pipeline.edges()[2][1], KNeighborsClassifier)  # type: ignore
        self.assert_equal_predictions(sklearn_pipeline, lale_pipeline)

    def test_import_from_sklearn_pipeline_nested_pipeline(self):
        from sklearn.decomposition import PCA as SklearnPCA
        from sklearn.kernel_approximation import Nystroem as SklearnNystroem
        from sklearn.neighbors import KNeighborsClassifier as SklearnKNN
        from sklearn.pipeline import FeatureUnion

        union = FeatureUnion(
            [
                (
                    "selectkbest_pca",
                    sklearn.pipeline.make_pipeline(
                        SkSelectKBest(k=3), SklearnPCA(n_components=1)
                    ),
                ),
                ("nys", SklearnNystroem(n_components=2, random_state=42)),
            ]
        )
        sklearn_pipeline = sklearn.pipeline.make_pipeline(union, SklearnKNN())
        lale_pipeline = typing.cast(
            TrainablePipeline,
            import_from_sklearn_pipeline(sklearn_pipeline),
        )
        self.assertEqual(len(lale_pipeline.edges()), 4)

        # These assertions assume topological sort
        self.assertIsInstance(lale_pipeline.edges()[0][0], SelectKBest)  # type: ignore
        self.assertIsInstance(lale_pipeline.edges()[0][1], PCA)  # type: ignore
        self.assertIsInstance(lale_pipeline.edges()[1][0], PCA)  # type: ignore
        self.assertIsInstance(lale_pipeline.edges()[1][1], ConcatFeatures)  # type: ignore
        self.assertIsInstance(lale_pipeline.edges()[2][0], Nystroem)  # type: ignore
        self.assertIsInstance(lale_pipeline.edges()[2][1], ConcatFeatures)  # type: ignore
        self.assertIsInstance(lale_pipeline.edges()[3][0], ConcatFeatures)  # type: ignore
        self.assertIsInstance(lale_pipeline.edges()[3][1], KNeighborsClassifier)  # type: ignore
        self.assert_equal_predictions(sklearn_pipeline, lale_pipeline)

    def test_import_from_sklearn_pipeline_nested_pipeline1(self):
        from sklearn.decomposition import PCA as SklearnPCA
        from sklearn.kernel_approximation import Nystroem as SklearnNystroem
        from sklearn.neighbors import KNeighborsClassifier as SklearnKNN
        from sklearn.pipeline import FeatureUnion

        union = FeatureUnion(
            [
                (
                    "selectkbest_pca",
                    sklearn.pipeline.make_pipeline(
                        SkSelectKBest(k=3),
                        FeatureUnion(
                            [
                                ("pca", SklearnPCA(n_components=1)),
                                (
                                    "nested_pipeline",
                                    sklearn.pipeline.make_pipeline(
                                        SkSelectKBest(k=2), SklearnNystroem()
                                    ),
                                ),
                            ]
                        ),
                    ),
                ),
                ("nys", SklearnNystroem(n_components=2, random_state=42)),
            ]
        )
        sklearn_pipeline = sklearn.pipeline.make_pipeline(union, SklearnKNN())
        lale_pipeline = typing.cast(
            TrainablePipeline,
            import_from_sklearn_pipeline(sklearn_pipeline),
        )
        self.assertEqual(len(lale_pipeline.edges()), 8)
        # These assertions assume topological sort, which may not be unique. So the assertions are brittle.

        self.assertIsInstance(lale_pipeline.edges()[0][0], SelectKBest)  # type: ignore
        self.assertIsInstance(lale_pipeline.edges()[0][1], PCA)  # type: ignore
        self.assertIsInstance(lale_pipeline.edges()[1][0], SelectKBest)  # type: ignore
        self.assertIsInstance(lale_pipeline.edges()[1][1], SelectKBest)  # type: ignore
        self.assertIsInstance(lale_pipeline.edges()[2][0], SelectKBest)  # type: ignore
        self.assertIsInstance(lale_pipeline.edges()[2][1], Nystroem)  # type: ignore
        self.assertIsInstance(lale_pipeline.edges()[3][0], PCA)  # type: ignore
        self.assertIsInstance(lale_pipeline.edges()[3][1], ConcatFeatures)  # type: ignore
        self.assertIsInstance(lale_pipeline.edges()[4][0], Nystroem)  # type: ignore
        self.assertIsInstance(lale_pipeline.edges()[4][1], ConcatFeatures)  # type: ignore
        self.assertIsInstance(lale_pipeline.edges()[5][0], ConcatFeatures)  # type: ignore
        self.assertIsInstance(lale_pipeline.edges()[5][1], ConcatFeatures)  # type: ignore
        self.assertIsInstance(lale_pipeline.edges()[6][0], Nystroem)  # type: ignore
        self.assertIsInstance(lale_pipeline.edges()[6][1], ConcatFeatures)  # type: ignore
        self.assertIsInstance(lale_pipeline.edges()[7][0], ConcatFeatures)  # type: ignore
        self.assertIsInstance(lale_pipeline.edges()[7][1], KNeighborsClassifier)  # type: ignore
        self.assert_equal_predictions(sklearn_pipeline, lale_pipeline)

    def test_import_from_sklearn_pipeline_nested_pipeline2(self):
        from sklearn.decomposition import PCA as SklearnPCA
        from sklearn.kernel_approximation import Nystroem as SklearnNystroem
        from sklearn.neighbors import KNeighborsClassifier as SklearnKNN
        from sklearn.pipeline import FeatureUnion

        union = FeatureUnion(
            [
                (
                    "selectkbest_pca",
                    sklearn.pipeline.make_pipeline(
                        SkSelectKBest(k=3),
                        sklearn.pipeline.make_pipeline(
                            SkSelectKBest(k=2), SklearnPCA()
                        ),
                    ),
                ),
                ("nys", SklearnNystroem(n_components=2, random_state=42)),
            ]
        )
        sklearn_pipeline = sklearn.pipeline.make_pipeline(union, SklearnKNN())
        lale_pipeline = typing.cast(
            TrainablePipeline,
            import_from_sklearn_pipeline(sklearn_pipeline),
        )
        self.assertEqual(len(lale_pipeline.edges()), 5)

        self.assertIsInstance(lale_pipeline.edges()[0][0], SelectKBest)  # type: ignore
        self.assertIsInstance(lale_pipeline.edges()[0][1], SelectKBest)  # type: ignore
        self.assertIsInstance(lale_pipeline.edges()[1][0], SelectKBest)  # type: ignore
        self.assertIsInstance(lale_pipeline.edges()[1][1], PCA)  # type: ignore
        self.assertIsInstance(lale_pipeline.edges()[2][0], PCA)  # type: ignore
        self.assertIsInstance(lale_pipeline.edges()[2][1], ConcatFeatures)  # type: ignore
        self.assertIsInstance(lale_pipeline.edges()[3][0], Nystroem)  # type: ignore
        self.assertIsInstance(lale_pipeline.edges()[3][1], ConcatFeatures)  # type: ignore
        self.assertIsInstance(lale_pipeline.edges()[4][0], ConcatFeatures)  # type: ignore
        self.assertIsInstance(lale_pipeline.edges()[4][1], KNeighborsClassifier)  # type: ignore

        self.assert_equal_predictions(sklearn_pipeline, lale_pipeline)

    def test_import_from_sklearn_pipeline_noop(self):
        from sklearn.ensemble import GradientBoostingClassifier
        from sklearn.pipeline import Pipeline

        pipe = Pipeline([("noop", None), ("gbc", GradientBoostingClassifier())])
        _ = import_from_sklearn_pipeline(pipe)

    def test_import_from_sklearn_pipeline_noop1(self):
        from sklearn.ensemble import GradientBoostingClassifier
        from sklearn.pipeline import Pipeline

        pipe = Pipeline([("noop", NoOp()), ("gbc", GradientBoostingClassifier())])
        _ = import_from_sklearn_pipeline(pipe)

    def test_import_from_sklearn_pipeline_no_wrapper(self):
        from sklearn.neighbors import LocalOutlierFactor
        from sklearn.pipeline import make_pipeline as sk_make_pipeline

        sklearn_pipeline = sk_make_pipeline(PCA(), LocalOutlierFactor())
        _ = import_from_sklearn_pipeline(sklearn_pipeline, fitted=False)

    def test_import_from_sklearn_pipeline_higherorder(self):
        from sklearn.ensemble import VotingClassifier as VC
        from sklearn.feature_selection import f_regression
        from sklearn.pipeline import Pipeline
        from sklearn.svm import SVC as SklearnSVC

        anova_filter = SkSelectKBest(f_regression, k=3)
        clf = SklearnSVC(kernel="linear")
        sklearn_pipeline = Pipeline(
            [("anova", anova_filter), ("vc_svc", VC(estimators=[("clf", clf)]))]
        )
        lale_pipeline = typing.cast(
            TrainablePipeline,
            import_from_sklearn_pipeline(sklearn_pipeline),
        )
        # for i, pipeline_step in enumerate(sklearn_pipeline.named_steps):
        #     sklearn_step_params = sklearn_pipeline.named_steps[
        #         pipeline_step
        #     ].get_params()
        #     lale_sklearn_params = self.get_sklearn_params(lale_pipeline.steps_list()[i])
        #     self.assertEqual(sklearn_step_params, lale_sklearn_params)
        self.assert_equal_predictions(sklearn_pipeline, lale_pipeline)

    def test_export_to_sklearn_pipeline(self):
        lale_pipeline = PCA(n_components=3) >> KNeighborsClassifier()
        trained_lale_pipeline = lale_pipeline.fit(self.X_train, self.y_train)
        sklearn_pipeline = trained_lale_pipeline.export_to_sklearn_pipeline()
        for i, pipeline_step in enumerate(sklearn_pipeline.named_steps):
            sklearn_step_params = sklearn_pipeline.named_steps[
                pipeline_step
            ].get_params()
            lale_sklearn_params = self.get_sklearn_params(
                trained_lale_pipeline.steps_list()[i]
            )
            self.assertEqual(sklearn_step_params, lale_sklearn_params)
        self.assert_equal_predictions(sklearn_pipeline, trained_lale_pipeline)

    def test_export_to_sklearn_pipeline1(self):
        lale_pipeline = SkSelectKBest(k=3) >> KNeighborsClassifier()
        trained_lale_pipeline = lale_pipeline.fit(self.X_train, self.y_train)
        sklearn_pipeline = trained_lale_pipeline.export_to_sklearn_pipeline()
        for i, pipeline_step in enumerate(sklearn_pipeline.named_steps):
            sklearn_step_params = type(sklearn_pipeline.named_steps[pipeline_step])
            lale_sklearn_params = (
                type(trained_lale_pipeline.steps_list()[i]._impl._wrapped_model)
                if hasattr(
                    trained_lale_pipeline.steps_list()[i]._impl, "_wrapped_model"
                )
                else type(trained_lale_pipeline.steps_list()[i]._impl)
            )
            self.assertEqual(sklearn_step_params, lale_sklearn_params)
        self.assert_equal_predictions(sklearn_pipeline, trained_lale_pipeline)

    def test_export_to_sklearn_pipeline2(self):
        from sklearn.pipeline import FeatureUnion

        lale_pipeline = (
            (
                (
                    (PCA(svd_solver="randomized", random_state=42) & SkSelectKBest(k=3))
                    >> ConcatFeatures()
                )
                & Nystroem(random_state=42)
            )
            >> ConcatFeatures()
            >> KNeighborsClassifier()
        )
        trained_lale_pipeline = lale_pipeline.fit(self.X_train, self.y_train)
        sklearn_pipeline = trained_lale_pipeline.export_to_sklearn_pipeline()
        self.assertIsInstance(
            sklearn_pipeline.named_steps["featureunion"], FeatureUnion
        )
        from sklearn.neighbors import KNeighborsClassifier as SklearnKNN

        self.assertIsInstance(
            sklearn_pipeline.named_steps["kneighborsclassifier"], SklearnKNN
        )
        self.assert_equal_predictions(sklearn_pipeline, trained_lale_pipeline)

    def test_export_to_sklearn_pipeline3(self):
        from sklearn.pipeline import FeatureUnion

        lale_pipeline = (
            (
                (PCA() >> SkSelectKBest(k=2))
                & (Nystroem(random_state=42) >> SkSelectKBest(k=3))
                & (SkSelectKBest(k=3))
            )
            >> ConcatFeatures()
            >> SkSelectKBest(k=2)
            >> LogisticRegression()
        )
        trained_lale_pipeline = lale_pipeline.fit(self.X_train, self.y_train)
        sklearn_pipeline = trained_lale_pipeline.export_to_sklearn_pipeline()
        self.assertIsInstance(
            sklearn_pipeline.named_steps["featureunion"], FeatureUnion
        )
        self.assertIsInstance(
            sklearn_pipeline.named_steps["selectkbest"], SkSelectKBest
        )
        from sklearn.linear_model import LogisticRegression as SklearnLR

        self.assertIsInstance(
            sklearn_pipeline.named_steps["logisticregression"], SklearnLR
        )
        self.assert_equal_predictions(sklearn_pipeline, trained_lale_pipeline)

    def test_export_to_sklearn_pipeline4(self):
        lale_pipeline = make_pipeline(LogisticRegression())
        trained_lale_pipeline = lale_pipeline.fit(self.X_train, self.y_train)
        sklearn_pipeline = trained_lale_pipeline.export_to_sklearn_pipeline()
        from sklearn.linear_model import LogisticRegression as SklearnLR

        self.assertIsInstance(
            sklearn_pipeline.named_steps["logisticregression"], SklearnLR
        )
        self.assert_equal_predictions(sklearn_pipeline, trained_lale_pipeline)

    def test_export_to_sklearn_pipeline5(self):
        lale_pipeline = PCA() >> (XGBClassifier() | SGDClassifier())
        with self.assertRaises(ValueError):
            _ = lale_pipeline.export_to_sklearn_pipeline()

    def test_export_to_pickle(self):
        lale_pipeline = make_pipeline(LogisticRegression())
        trained_lale_pipeline = lale_pipeline.fit(self.X_train, self.y_train)
        pickle.dumps(lale_pipeline)
        pickle.dumps(trained_lale_pipeline)

    def test_import_from_sklearn_pipeline2(self):
        from sklearn.feature_selection import f_regression
        from sklearn.pipeline import Pipeline
        from sklearn.svm import SVC as SklearnSVC

        anova_filter = SkSelectKBest(f_regression, k=3)
        clf = SklearnSVC(kernel="linear")
        sklearn_pipeline = Pipeline([("anova", anova_filter), ("svc", clf)])
        sklearn_pipeline.fit(self.X_train, self.y_train)
        lale_pipeline = typing.cast(
            TrainedPipeline,
            import_from_sklearn_pipeline(sklearn_pipeline),
        )
        lale_pipeline.predict(self.X_test)

    def test_import_from_sklearn_pipeline3(self):
        from sklearn.feature_selection import f_regression
        from sklearn.pipeline import Pipeline
        from sklearn.svm import SVC as SklearnSVC

        anova_filter = SkSelectKBest(f_regression, k=3)
        clf = SklearnSVC(kernel="linear")
        sklearn_pipeline = Pipeline([("anova", anova_filter), ("svc", clf)])
        lale_pipeline = typing.cast(
            TrainablePipeline,
            import_from_sklearn_pipeline(sklearn_pipeline, fitted=False),
        )
        with self.assertRaises(
            ValueError
        ):  # fitted=False returns a Trainable, so calling predict is invalid.
            lale_pipeline.predict(self.X_test)

    def test_export_to_sklearn_pipeline_with_noop_1(self):
        lale_pipeline = NoOp() >> PCA(n_components=3) >> KNeighborsClassifier()
        trained_lale_pipeline = lale_pipeline.fit(self.X_train, self.y_train)
        sklearn_pipeline = trained_lale_pipeline.export_to_sklearn_pipeline()
        self.assert_equal_predictions(sklearn_pipeline, trained_lale_pipeline)

    def test_export_to_sklearn_pipeline_with_noop_2(self):
        lale_pipeline = PCA(n_components=3) >> NoOp() >> KNeighborsClassifier()
        trained_lale_pipeline = lale_pipeline.fit(self.X_train, self.y_train)
        sklearn_pipeline = trained_lale_pipeline.export_to_sklearn_pipeline()
        self.assert_equal_predictions(sklearn_pipeline, trained_lale_pipeline)

    def test_export_to_sklearn_pipeline_with_noop_3(self):
        # This test is probably unnecessary, but doesn't harm at this point
        lale_pipeline = PCA(n_components=3) >> KNeighborsClassifier() >> NoOp()
        trained_lale_pipeline = lale_pipeline.fit(self.X_train, self.y_train)
        _ = trained_lale_pipeline.export_to_sklearn_pipeline()

    def test_export_to_sklearn_pipeline_with_noop_4(self):
        lale_pipeline = NoOp() >> KNeighborsClassifier()
        trained_lale_pipeline = lale_pipeline.fit(self.X_train, self.y_train)
        sklearn_pipeline = trained_lale_pipeline.export_to_sklearn_pipeline()
        self.assert_equal_predictions(sklearn_pipeline, trained_lale_pipeline)


class TestComposition(unittest.TestCase):
    def setUp(self):
        data = sklearn.datasets.load_iris()
        X, y = data.data, data.target
        self.X_train, self.X_test, self.y_train, self.y_test = train_test_split(X, y)

    def test_two_estimators_predict(self):
        pipeline = (
            StandardScaler()
            >> (PCA() & Nystroem() & LogisticRegression())
            >> ConcatFeatures()
            >> NoOp()
            >> LogisticRegression()
        )
        trained = pipeline.fit(self.X_train, self.y_train)
        trained.predict(self.X_test)

    def test_two_estimators_predict1(self):
        pipeline = (
            StandardScaler()
            >> (PCA() & Nystroem() & PassiveAggressiveClassifier())
            >> ConcatFeatures()
            >> NoOp()
            >> PassiveAggressiveClassifier()
        )
        trained = pipeline.fit(self.X_train, self.y_train)
        trained.predict(self.X_test)

    def test_two_estimators_predict_proba(self):
        pipeline = (
            StandardScaler()
            >> (PCA() & Nystroem() & LogisticRegression())
            >> ConcatFeatures()
            >> NoOp()
            >> LogisticRegression()
        )
        trained = pipeline.fit(self.X_train, self.y_train)
        trained.predict_proba(self.X_test)

    def test_two_estimators_predict_proba1(self):
        pipeline = (
            StandardScaler()
            >> (PCA() & Nystroem() & GaussianNB())
            >> ConcatFeatures()
            >> NoOp()
            >> GaussianNB()
        )
        pipeline.fit(self.X_train, self.y_train)
        pipeline.predict_proba(self.X_test)

    def test_multiple_estimators_predict_predict_proba(self):
        pipeline = (
            StandardScaler()
            >> (LogisticRegression() & PCA())
            >> ConcatFeatures()
            >> (NoOp() & LinearSVC())
            >> ConcatFeatures()
            >> KNeighborsClassifier()
        )
        pipeline.fit(self.X_train, self.y_train)
        _ = pipeline.predict_proba(self.X_test)
        _ = pipeline.predict(self.X_test)

    def test_two_transformers(self):
        tfm1 = PCA()
        tfm2 = Nystroem()
        trainable = tfm1 >> tfm2
        digits = sklearn.datasets.load_digits()
        trained = trainable.fit(digits.data, digits.target)
        _ = trained.transform(digits.data)

    def test_duplicate_instances(self):
        tfm = PCA()
        clf = LogisticRegression(
            LogisticRegression.enum.solver.lbfgs,
            LogisticRegression.enum.multi_class.auto,
        )
        with self.assertRaises(ValueError):
            _ = make_pipeline(tfm, tfm, clf)

    def test_increase_num_rows_predict(self):
        from test.mock_custom_operators import IncreaseRows

        increase_rows = IncreaseRows()
        trainable = increase_rows >> LogisticRegression()
        iris = sklearn.datasets.load_iris()
        X, y = iris.data, iris.target
        trained = trainable.fit(X, y)
        y_pred = trained.predict(X)
        self.assertEqual(len(y_pred), len(y) + increase_rows.impl.n_rows)

    def test_increase_num_rows_transform_X_y(self):
        from test.mock_custom_operators import IncreaseRows

        increase_rows_4 = IncreaseRows(n_rows=4)
        increase_rows_2 = IncreaseRows(n_rows=2)
        trainable = increase_rows_4 >> increase_rows_2
        iris = sklearn.datasets.load_iris()
        X, y = iris.data, iris.target
        trained = trainable.fit(X, y)
        output_X, output_y = trained.transform_X_y(X, y)
        self.assertEqual(output_X.shape[0], X.shape[0] + 4 + 2)
        self.assertEqual(output_X.shape[1], X.shape[1])
        self.assertEqual(output_y.shape[0], y.shape[0] + 4 + 2)

    def test_remove_last1(self):
        pipeline = (
            StandardScaler()
            >> (PCA() & Nystroem() & PassiveAggressiveClassifier())
            >> ConcatFeatures()
            >> NoOp()
            >> PassiveAggressiveClassifier()
        )
        new_pipeline = pipeline.remove_last()
        self.assertEqual(len(new_pipeline._steps), 6)
        self.assertEqual(len(pipeline._steps), 7)

    def test_remove_last2(self):
        pipeline = (
            StandardScaler()
            >> (PCA() & Nystroem() & PassiveAggressiveClassifier())
            >> ConcatFeatures()
            >> NoOp()
            >> (PassiveAggressiveClassifier() & LogisticRegression())
        )
        with self.assertRaises(ValueError):
            pipeline.remove_last()

    def test_remove_last3(self):
        pipeline = (
            StandardScaler()
            >> (PCA() & Nystroem() & PassiveAggressiveClassifier())
            >> ConcatFeatures()
            >> NoOp()
            >> PassiveAggressiveClassifier()
        )
        pipeline.remove_last().freeze_trainable()

    def test_remove_last4(self):
        pipeline = (
            StandardScaler()
            >> (PCA() & Nystroem() & PassiveAggressiveClassifier())
            >> ConcatFeatures()
            >> NoOp()
            >> PassiveAggressiveClassifier()
        )
        new_pipeline = pipeline.remove_last(inplace=True)
        self.assertEqual(len(new_pipeline._steps), 6)
        self.assertEqual(len(pipeline._steps), 6)

    def test_remove_last5(self):
        pipeline = (
            StandardScaler()
            >> (PCA() & Nystroem() & PassiveAggressiveClassifier())
            >> ConcatFeatures()
            >> NoOp()
            >> PassiveAggressiveClassifier()
        )
        pipeline.remove_last(inplace=True).freeze_trainable()


class TestAutoPipeline(unittest.TestCase):
    def _fit_predict(self, prediction_type, all_X, all_y, verbose=True):
        if verbose:
            _file_name, _line, fn_name, _text = traceback.extract_stack()[-2]
            print(f"--- TestAutoPipeline.{fn_name}() ---")
        from lale.lib.lale import AutoPipeline

        train_X, test_X, train_y, test_y = train_test_split(all_X, all_y)
        trainable = AutoPipeline(
            prediction_type=prediction_type, max_evals=10, verbose=verbose
        )
        trained = trainable.fit(train_X, train_y)
        predicted = trained.predict(test_X)
        if prediction_type == "regression":
            score = f"r2 score {r2_score(test_y, predicted):.2f}"
        else:
            score = f"accuracy {accuracy_score(test_y, predicted):.1%}"
        if verbose:
            print(score)
            pipe = trained.get_pipeline()
            assert pipe is not None
            print(pipe.pretty_print(show_imports=False))

    def test_sklearn_iris(self):
        # classification, only numbers, no missing values
        all_X, all_y = sklearn.datasets.load_iris(return_X_y=True)
        self._fit_predict("classification", all_X, all_y)

    def test_sklearn_digits(self):
        # classification, numbers but some appear categorical, no missing values
        all_X, all_y = sklearn.datasets.load_digits(return_X_y=True)
        self._fit_predict("classification", all_X, all_y)

    def test_sklearn_boston(self):
        # regression, categoricals+numbers, no missing values
        from lale.datasets.util import load_boston

        all_X, all_y = load_boston(return_X_y=True)
        self._fit_predict("regression", all_X, all_y)

    def test_sklearn_diabetes(self):
        # regression, categoricals+numbers, no missing values
        all_X, all_y = sklearn.datasets.load_diabetes(return_X_y=True)
        self._fit_predict("regression", all_X, all_y)

    def test_openml_creditg(self):
        # classification, categoricals+numbers incl. string, no missing values
        (orig_train_X, orig_train_y), _ = lale.datasets.openml.fetch(
            "credit-g", "classification", preprocess=False
        )
        subsample_X, _, subsample_y, _ = train_test_split(
            orig_train_X, orig_train_y, train_size=0.05
        )
        self._fit_predict("classification", subsample_X, subsample_y)

    def test_missing_iris(self):
        # classification, only numbers, synthetically added missing values
        all_X, all_y = sklearn.datasets.load_iris(return_X_y=True)
        with_missing_X = lale.helpers.add_missing_values(all_X)
        with self.assertRaisesRegex(ValueError, "Input.*contains NaN"):
            lr_trainable = LogisticRegression()
            _ = lr_trainable.fit(with_missing_X, all_y)
        self._fit_predict("classification", with_missing_X, all_y)

    def test_missing_boston(self):
        # regression, categoricals+numbers, synthetically added missing values
        from lale.datasets.util import load_boston

        all_X, all_y = load_boston(return_X_y=True)
        with_missing_X = lale.helpers.add_missing_values(all_X)
        with self.assertRaisesRegex(ValueError, "Input.*contains NaN"):
            lr_trainable = LinearRegression()
            _ = lr_trainable.fit(with_missing_X, all_y)
        self._fit_predict("regression", with_missing_X, all_y)

    def test_missing_creditg(self):
        # classification, categoricals+numbers incl. string, synth. missing
        (orig_train_X, orig_train_y), _ = lale.datasets.openml.fetch(
            "credit-g", "classification", preprocess=False
        )
        subsample_X, _, subsample_y, _ = train_test_split(
            orig_train_X, orig_train_y, train_size=0.05
        )
        with_missing_X = lale.helpers.add_missing_values(subsample_X)
        self._fit_predict("classification", with_missing_X, subsample_y)


class TestOperatorChoice(unittest.TestCase):
    def test_make_choice_with_instance(self):
        from sklearn.datasets import load_iris

        iris = load_iris()
        X, y = iris.data, iris.target
        tfm = PCA() | Nystroem() | NoOp()
        with self.assertRaises(AttributeError):
            # we are trying to trigger a runtime error here, so we ignore the static warning
            _ = tfm.fit(X, y)  # type: ignore
        _ = (OneHotEncoder | NoOp) >> tfm >> (LogisticRegression | KNeighborsClassifier)
        _ = (
            (OneHotEncoder | NoOp)
            >> (PCA | Nystroem)
            >> (LogisticRegression | KNeighborsClassifier)
        )
        _ = (
            make_choice(OneHotEncoder, NoOp)
            >> make_choice(PCA, Nystroem)
            >> make_choice(LogisticRegression, KNeighborsClassifier)
        )


class TestScore(unittest.TestCase):
    def setUp(self):
        data = sklearn.datasets.load_iris()
        X, y = data.data, data.target
        self.X_train, self.X_test, self.y_train, self.y_test = train_test_split(X, y)

    def test_trained_pipeline(self):
        trainable_pipeline = StandardScaler() >> LogisticRegression()
        trained_pipeline = trainable_pipeline.fit(self.X_train, self.y_train)
        score = trained_pipeline.score(self.X_test, self.y_test)
        predictions = trained_pipeline.predict(self.X_test)

        accuracy = accuracy_score(self.y_test, predictions)
        self.assertEqual(accuracy, score)

    def test_trainable_pipeline(self):
        trainable_pipeline = StandardScaler() >> LogisticRegression()
        trainable_pipeline.fit(self.X_train, self.y_train)
        score = trainable_pipeline.score(self.X_test, self.y_test)
        predictions = trainable_pipeline.predict(self.X_test)

        accuracy = accuracy_score(self.y_test, predictions)
        self.assertEqual(accuracy, score)

    def test_planned_pipeline(self):
        planned_pipeline = StandardScaler >> LogisticRegression
        with self.assertRaises(AttributeError):
            planned_pipeline.score(self.X_test, self.y_test)  # type: ignore


class TestScoreSamples(unittest.TestCase):
    def setUp(self):
        data = sklearn.datasets.load_iris()
        X, y = data.data, data.target
        self.X_train, self.X_test, self.y_train, self.y_test = train_test_split(X, y)
        import warnings

        warnings.filterwarnings("ignore")

    def test_trained_pipeline(self):
        trainable_pipeline = StandardScaler() >> IsolationForest()
        trained_pipeline = trainable_pipeline.fit(self.X_train, self.y_train)
        _ = trained_pipeline.score_samples(self.X_test)

    def test_trainable_pipeline(self):
        trainable_pipeline = StandardScaler() >> IsolationForest()
        trainable_pipeline.fit(self.X_train, self.y_train)
        with self.assertWarns(DeprecationWarning):
            _ = trainable_pipeline.score_samples(self.X_test)

    def test_planned_pipeline(self):
        planned_pipeline = StandardScaler >> IsolationForest
        with self.assertRaises(AttributeError):
            planned_pipeline.score_samples(self.X_test)  # type: ignore

    def test_with_incompatible_estimator(self):
        trainable_pipeline = StandardScaler() >> LogisticRegression()
        trained_pipeline = trainable_pipeline.fit(self.X_train, self.y_train)
        with self.assertRaises(AttributeError):
            _ = trained_pipeline.score_samples(self.X_test)


class TestPredictLogProba(unittest.TestCase):
    def setUp(self):
        data = sklearn.datasets.load_iris()
        X, y = data.data, data.target
        self.X_train, self.X_test, self.y_train, self.y_test = train_test_split(X, y)
        import warnings

        warnings.filterwarnings("ignore")

    def test_trained_pipeline(self):
        trainable_pipeline = StandardScaler() >> AdaBoostClassifier()
        trained_pipeline = trainable_pipeline.fit(self.X_train, self.y_train)
        _ = trained_pipeline.predict_log_proba(self.X_test)

    def test_trainable_pipeline(self):
        trainable_pipeline = StandardScaler() >> AdaBoostClassifier()
        trainable_pipeline.fit(self.X_train, self.y_train)
        with self.assertWarns(DeprecationWarning):
            _ = trainable_pipeline.predict_log_proba(self.X_test)

    def test_planned_pipeline(self):
        planned_pipeline = StandardScaler >> AdaBoostClassifier
        with self.assertRaises(AttributeError):
            planned_pipeline.predict_log_proba(self.X_test)  # type: ignore

    def test_with_incompatible_estimator(self):
        trainable_pipeline = StandardScaler() >> IsolationForest()
        trained_pipeline = trainable_pipeline.fit(self.X_train, self.y_train)
        with self.assertRaises(AttributeError):
            _ = trained_pipeline.predict_log_proba(self.X_test)

    def test_with_incompatible_estimator_1(self):
        trainable_pipeline = IsolationForest()
        trained_pipeline = trainable_pipeline.fit(self.X_train, self.y_train)
        with self.assertRaises(AttributeError):
            _ = trained_pipeline.predict_log_proba(self.X_test)


class TestPartialFit(unittest.TestCase):
    def setUp(self):
        data = sklearn.datasets.load_iris()
        X, y = data.data, data.target
        self.X_train, self.X_test, self.y_train, self.y_test = train_test_split(X, y)
        import warnings

        warnings.filterwarnings("ignore")

    def test_first_call(self):
        trainable_pipeline = StandardScaler()
        trained_pipeline = trainable_pipeline.fit(self.X_train, self.y_train)
        new_pipeline = trained_pipeline.freeze_trained() >> SGDClassifier()
        new_trained_pipeline = new_pipeline.partial_fit(
            self.X_train, self.y_train, classes=[0, 1, 2]
        )
        _ = new_trained_pipeline.predict(self.X_test)

    def test_multiple_calls_with_classes(self):
        trainable_pipeline = StandardScaler()
        trained_pipeline = trainable_pipeline.fit(self.X_train, self.y_train)
        new_pipeline = trained_pipeline.freeze_trained() >> SGDClassifier()
        new_trained_pipeline = new_pipeline.partial_fit(
            self.X_train, self.y_train, classes=[0, 1, 2]
        )
        new_trained_pipeline = new_trained_pipeline.partial_fit(
            self.X_test, self.y_test, classes=[0, 1, 2]
        )
        _ = new_trained_pipeline.predict(self.X_test)

    def _last_impl_has(self, op, attr):
        last = op.get_last()
        assert last is not None
        return hasattr(last._impl, attr)

    def test_second_call_without_classes(self):
        trainable_pipeline = StandardScaler()
        trained_pipeline = trainable_pipeline.fit(self.X_train, self.y_train)
        new_pipeline = trained_pipeline.freeze_trained() >> SGDClassifier()
        new_trained_pipeline = new_pipeline.partial_fit(
            self.X_train, self.y_train, classes=[0, 1, 2]
        )
        # Once SGDClassifier is trained, it has a classes_ attribute.
        self.assertTrue(self._last_impl_has(new_trained_pipeline, "classes_"))
        new_trained_pipeline = new_trained_pipeline.partial_fit(
            self.X_test, self.y_test
        )
        _ = new_trained_pipeline.predict(self.X_test)

    def test_second_call_with_different_classes(self):
        trainable_pipeline = StandardScaler()
        trained_pipeline = trainable_pipeline.fit(self.X_train, self.y_train)
        new_pipeline = trained_pipeline.freeze_trained() >> SGDClassifier()
        new_trained_pipeline = new_pipeline.partial_fit(
            self.X_train, self.y_train, classes=[0, 1, 2]
        )
        # Once SGDClassifier is trained, it has a classes_ attribute.
        self.assertTrue(self._last_impl_has(new_trained_pipeline, "classes_"))
        subset_labels = self.y_test[np.where(self.y_test != 0)]
        subset_X = self.X_test[0 : len(subset_labels)]
        new_trained_pipeline = new_trained_pipeline.partial_fit(subset_X, subset_labels)
        _ = new_trained_pipeline.predict(self.X_test)

    def test_second_call_with_different_classes_trainable(self):
        trainable_pipeline = StandardScaler()
        trained_pipeline = trainable_pipeline.fit(self.X_train, self.y_train)
        new_pipeline = trained_pipeline.freeze_trained() >> SGDClassifier()
        new_pipeline.partial_fit(self.X_train, self.y_train, classes=[0, 1, 2])
        # Once SGDClassifier is trained, it has a classes_ attribute.
        self.assertTrue(self._last_impl_has(new_pipeline._trained, "classes_"))
        subset_labels = self.y_test[np.where(self.y_test != 0)]
        subset_X = self.X_test[0 : len(subset_labels)]
        new_trained_pipeline = new_pipeline.partial_fit(subset_X, subset_labels)
        _ = new_trained_pipeline.predict(self.X_test)

    def test_call_on_trainable(self):
        trainable_pipeline = StandardScaler()
        trained_pipeline = trainable_pipeline.fit(self.X_train, self.y_train)
        new_pipeline = trained_pipeline.freeze_trained() >> SGDClassifier()
        new_pipeline.partial_fit(self.X_train, self.y_train, classes=[0, 1, 2])
        new_pipeline.pretty_print()
        new_trained_pipeline = new_pipeline.partial_fit(
            self.X_test, self.y_test, classes=[0, 1, 2]
        )
        self.assertEqual(new_trained_pipeline, new_pipeline._trained)
        _ = new_trained_pipeline.predict(self.X_test)
        new_pipeline.partial_fit(self.X_train, self.y_train, classes=[0, 1, 2])

    def test_call_on_trainable_with_freeze_trained_prefix(self):
        trainable_pipeline = StandardScaler()
        trained_pipeline = trainable_pipeline.fit(self.X_train, self.y_train)
        new_pipeline = trained_pipeline >> SGDClassifier()
        new_pipeline.partial_fit(self.X_train, self.y_train, classes=[0, 1, 2])
        new_pipeline.pretty_print()
        new_trained_pipeline = new_pipeline.partial_fit(
            self.X_test, self.y_test, classes=[0, 1, 2]
        )
        self.assertEqual(new_trained_pipeline, new_pipeline._trained)
        _ = new_trained_pipeline.predict(self.X_test)
        new_pipeline.partial_fit(self.X_train, self.y_train, classes=[0, 1, 2])

    def test_call_on_trainable_with_freeze_trained_prefix_false(self):
        trainable_pipeline = StandardScaler()
        trained_pipeline = trainable_pipeline.fit(self.X_train, self.y_train)
        new_pipeline = trained_pipeline >> SGDClassifier()
        with self.assertRaises(ValueError):
            new_pipeline.partial_fit(
                self.X_train,
                self.y_train,
                freeze_trained_prefix=False,
                classes=[0, 1, 2],
            )

    def test_call_on_trained_with_freeze_trained_prefix(self):
        trainable_pipeline = StandardScaler() >> SGDClassifier()
        trained_pipeline = trainable_pipeline.fit(self.X_train, self.y_train)
        new_pipeline = trained_pipeline
        new_pipeline.partial_fit(self.X_train, self.y_train, classes=[0, 1, 2])
        new_pipeline.pretty_print()
        new_trained_pipeline = new_pipeline.partial_fit(
            self.X_test, self.y_test, classes=[0, 1, 2]
        )
        _ = new_trained_pipeline.predict(self.X_test)
        new_pipeline.partial_fit(self.X_train, self.y_train, classes=[0, 1, 2])

    def test_call_on_trained_with_freeze_trained_prefix_false(self):
        trainable_pipeline = StandardScaler() >> SGDClassifier()
        trained_pipeline = trainable_pipeline.fit(self.X_train, self.y_train)
        new_pipeline = trained_pipeline
        with self.assertRaises(ValueError):
            new_pipeline.partial_fit(
                self.X_train,
                self.y_train,
                freeze_trained_prefix=False,
                classes=[0, 1, 2],
            )