-Two-dimensional, size-mutable, potentially heterogeneous tabular data of type
-[pd.DataFrame](https://pandas.pydata.org/docs/reference/api/pandas.DataFrame.html)
-or its [modin](https://modin.readthedocs.io/en/stable/flow/modin/pandas/dataframe.html)
-counterpart.
+Two-dimensional, size-mutable, potentially heterogeneous tabular data.
+The type is usually [pd.DataFrame](https://pandas.pydata.org/docs/reference/api/pandas.DataFrame.html),
+but could potentially be any of the dataframe types backed by the
+selected [data engine][data-acceleration].
-Any type object from which a [dataframe][] can be created. This includes an
-[iterable](https://docs.python.org/3/glossary.html#term-iterable), a
-[dict](https://docs.python.org/3/library/functions.html#func-dict) whose
+Any type object from which a [pd.DataFrame](https://pandas.pydata.org/docs/reference/api/pandas.DataFrame.html)
+can be created. This includes an [iterable](https://docs.python.org/3/glossary.html#term-iterable),
+a [dict](https://docs.python.org/3/library/functions.html#func-dict) whose
values are 1d-arrays, a two-dimensional [list](https://docs.python.org/3/library/functions.html#func-list),
[tuple](https://docs.python.org/3/library/functions.html#func-tuple), [np.ndarray](https://numpy.org/doc/stable/reference/generated/numpy.ndarray.html) or
[sps.csr_matrix](https://docs.scipy.org/doc/scipy/reference/generated/scipy.sparse.csr_matrix.html),
-and most commonly, a [dataframe][]. This is the standard input format for
-any dataset.
+or any object that follows the [dataframe interchange protocol](https://data-apis.org/dataframe-protocol/latest/index.html).
+This is the standard input format for any dataset.
Additionally, you can provide a callable whose output is any of the
aforementioned types. This is useful when the dataset is very large and
@@ -67,13 +67,6 @@ method. Often used interchangeably with [predictor][] because of user
preference.
-
-Immutable sequence used for indexing and alignment of type [pd.Index](https://pandas.pydata.org/docs/reference/api/pandas.Index.html)
-or their [modin](https://modin.readthedocs.io/en/stable/flow/modin/pandas/dataframe.html)
-counterpart.
-
-
All values in the [`missing`][atomclassifier-missing] attribute, as
@@ -129,10 +122,10 @@ column.
series
-One-dimensional ndarray with axis labels of type
-[pd.Series](https://pandas.pydata.org/docs/reference/api/pandas.Series.html#pandas.Series)
-or its [modin](https://modin.readthedocs.io/en/stable/flow/modin/pandas/series.html)
-counterpart.
+One-dimensional ndarray with axis labels. The type is usually
+[pd.Series](https://pandas.pydata.org/docs/reference/api/pandas.Series.html#pandas.Series),
+but could potentially be any of the series types backed by the
+selected [data engine][data-acceleration].
target
diff --git a/pyproject.toml b/pyproject.toml
index aa32035f3..981280631 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -21,29 +21,24 @@ classifiers = [
dependencies = [
"beartype>=0.16.4",
"category-encoders>=2.6.3",
- "dagshub>=0.3.8",
"dill>=0.3.6",
+ "featuretools>=1.28.0",
"gplearn>=0.4.2",
"imbalanced-learn>=0.12.0",
"ipython>=8.11.0",
"ipywidgets>=8.1.1",
- "featuretools>=1.28.0",
"joblib>=1.3.1",
"matplotlib>=3.7.2",
- "mlflow>=2.7.1",
- "modin[ray]>=0.25.0",
+ "mlflow>=2.10.2",
"nltk>=3.8.1",
"numpy>=1.23.0",
"optuna>=3.4.0",
- "pandas[parquet]>=2.1.2",
+ "pandas>=2.1.2",
"plotly>=5.18.0",
- "ray[serve]>=2.9.1",
- "requests>=2.31.0",
"scikit-learn>=1.4.0",
- "scikit-learn-intelex>=2023.2.1; platform_machine == 'x86_64' or platform_machine == 'AMD64'",
"scipy>=1.10.1",
"shap>=0.43.0",
- "sktime[forecasting]>=0.24.0",
+ "sktime[forecasting]>=0.26.0",
"statsmodels>=0.14.1",
"zoofs>=0.1.26",
]
@@ -52,9 +47,18 @@ dependencies = [
full = [
"botorch>=0.8.5",
"catboost>=1.2",
+ "dagshub>=0.3.8",
+ "dask[distributed]>=2024.2.0",
"explainerdashboard>=0.4.3",
"gradio>=3.44.4",
"lightgbm>=4.1.0",
+ "modin[ray]>=0.25.0",
+ "polars>=0.20.7",
+ "pyarrow>=15.0.0",
+ "pyspark>=3.5.0",
+ "ray[serve]>=2.9.1",
+ "requests>=2.31.0",
+ "scikit-learn-intelex>=2023.2.1; platform_machine == 'x86_64' or platform_machine == 'AMD64'",
"schemdraw>=0.16",
"statsforecast>=1.6.0",
"sweetviz>=2.3.1",
@@ -70,7 +74,6 @@ dev = [
"pandas_stubs>=2.1.1.230928",
"pre-commit>=3.5.0",
"ruff>=0.1.7",
- "types-requests>=2.31.0.10",
# Testing
"nbmake>=1.4.1", # To test example notebooks
"pytest>=7.2.1",
@@ -187,4 +190,8 @@ convention = "numpy"
[tool.mypy]
ignore_missing_imports = true
-disable_error_code = ["attr-defined"]
+disable_error_code = [
+ "attr-defined",
+ "abstract", # See https://github.com/python/mypy/issues/4717
+ "override", # Transformers' methods don't always match with that of TransformerMixin
+]
diff --git a/tests/conftest.py b/tests/conftest.py
index 9721308ff..97c7858ba 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -11,7 +11,9 @@
import numpy as np
import pandas as pd
+import pyarrow as pa
import pytest
+from ray.util.joblib import register_ray
from sklearn.base import BaseEstimator
from sklearn.datasets import (
load_breast_cancer, load_diabetes, load_wine,
@@ -23,7 +25,7 @@
from sktime.split import temporal_train_test_split
from atom.data_cleaning import TransformerMixin
-from atom.utils.utils import merge, n_cols, to_df, to_pandas
+from atom.utils.utils import merge, n_cols, to_df, to_tabular
if TYPE_CHECKING:
@@ -31,7 +33,7 @@
from _pytest.monkeypatch import MonkeyPatch
- from atom.utils.types import DataFrame, Pandas, Sequence, XSelector
+ from atom.utils.types import DataFrame, Pandas, Sequence, XConstructor
class DummyTransformer(TransformerMixin, BaseEstimator):
@@ -107,6 +109,18 @@ def _mock_mlflow_log_model(mocker):
mocker.patch("mlflow.sklearn.log_model")
+@pytest.fixture(autouse=True)
+def _register_ray():
+ """Register ray as joblib backend.
+
+ Although atom does this internally, it's skipped when ray is
+ mocked. Not registering it fails the call to joblib.parallel_config
+ in basetransformer.py.
+
+ """
+ register_ray()
+
+
@pytest.fixture()
def random():
"""Return numpy's default random number generator."""
@@ -114,8 +128,8 @@ def random():
def get_train_test(
- X: XSelector | None,
- y: Sequence[Any] | DataFrame,
+ X: XConstructor | None,
+ y: Sequence[Any] | pd.DataFrame,
) -> Pandas | tuple[Pandas, Pandas]:
"""Get train and test sets from X and y.
@@ -125,7 +139,7 @@ def get_train_test(
Feature set. If None, split as time series data set.
y: sequence or DataFrame
- Target column corresponding to `X`.
+ Target column(s) corresponding to `X`.
Returns
-------
@@ -138,7 +152,8 @@ def get_train_test(
"""
if X is not None:
return train_test_split(
- merge(to_df(X), to_pandas(y, columns=[f"y{i}" for i in range(n_cols(y))])),
+ merge(to_df(X), to_tabular(y, columns=[f"y{i}" for i in range(n_cols(y))])),
+ shuffle=False,
test_size=0.3,
random_state=1,
)
@@ -154,6 +169,9 @@ def get_train_test(
X_class, y_class = load_wine(return_X_y=True, as_frame=True)
X_reg, y_reg = load_diabetes(return_X_y=True, as_frame=True)
+# Pyarrow dtypes
+X_pa = X_bin.astype(pd.ArrowDtype(pa.float64()))
+
# Multilabel classification data
X_label, y_label = make_multilabel_classification(n_samples=200, n_classes=4)
diff --git a/tests/test_atom.py b/tests/test_atom.py
index b59e857b1..f04906c3e 100644
--- a/tests/test_atom.py
+++ b/tests/test_atom.py
@@ -10,6 +10,7 @@
import numpy as np
import pandas as pd
+import pyarrow as pa
import pytest
from category_encoders.target_encoder import TargetEncoder
from pandas.testing import assert_frame_equal, assert_index_equal
@@ -32,7 +33,7 @@
from .conftest import (
X10, DummyTransformer, X10_dt, X10_nan, X10_str, X10_str2, X20_out, X_bin,
- X_class, X_ex, X_label, X_reg, X_sparse, X_text, y10, y10_label,
+ X_class, X_ex, X_label, X_pa, X_reg, X_sparse, X_text, y10, y10_label,
y10_label2, y10_sn, y10_str, y_bin, y_class, y_ex, y_fc, y_label,
y_multiclass, y_multireg, y_reg,
)
@@ -313,6 +314,13 @@ def test_inverse_transform():
assert_frame_equal(atom.inverse_transform(atom.X), X_bin)
+def test_inverse_transform_output():
+ """Assert that the output type is determined by the data engine."""
+ atom = ATOMClassifier(X_bin, y_bin, engine="pyarrow", random_state=1)
+ atom.scale()
+ assert isinstance(atom.inverse_transform(X_bin), pa.Table)
+
+
def test_load_no_atom():
"""Assert that an error is raised when the instance is not atom."""
trainer = DirectClassifier("LR", random_state=1)
@@ -442,8 +450,8 @@ def test_shrink_dense2sparse():
def test_shrink_pyarrow():
- """Assert that it works with the pyarrow data backend."""
- atom = ATOMClassifier(X_bin, y_bin, engine={"data": "pyarrow"}, random_state=1)
+ """Assert that it works with pyarrow dtypes."""
+ atom = ATOMClassifier(X_pa, y_bin, engine="pandas-pyarrow", random_state=1)
assert atom.dtypes[0].name == "double[pyarrow]"
atom.shrink()
assert atom.dtypes[0].name == "float[pyarrow]"
@@ -488,6 +496,13 @@ def test_transform_not_train_only():
assert len(atom.transform(X_bin)) == len(X_bin)
+def test_transform_output():
+ """Assert that the output type is determined by the data engine."""
+ atom = ATOMClassifier(X_bin, y_bin, engine="pyarrow", random_state=1)
+ atom.scale()
+ assert isinstance(atom.transform(X_bin), pa.Table)
+
+
# Test base transformers =========================================== >>
def test_add_after_model():
@@ -630,15 +645,15 @@ def test_add_keep_column_names():
assert atom.features.tolist() == ["x0", "x1", "x2", "x3"]
# Transformer keeps rows equal
- atom.add(DummyTransformer(strategy="equal"), get_feature_names_out=None)
+ atom.add(DummyTransformer(strategy="equal"), feature_names_out=None)
assert atom.features.tolist() == ["x0", "x1", "x2", "x3"]
# Transformer drops rows
- atom.add(DummyTransformer(strategy="drop"), get_feature_names_out=None)
+ atom.add(DummyTransformer(strategy="drop"), feature_names_out=None)
assert atom.features.tolist() == ["x0", "x2", "x3"]
# Transformer adds a new column
- atom.add(DummyTransformer(strategy="add"), columns="!x2", get_feature_names_out=None)
+ atom.add(DummyTransformer(strategy="add"), columns="!x2", feature_names_out=None)
assert atom.features.tolist() == ["x0", "x2", "x3", "x4"]
@@ -649,9 +664,9 @@ def test_raise_length_mismatch():
atom.prune(columns=[2, 4])
-def test_add_pyarrow_columns():
+def test_keep_pyarrow_dtypes():
"""Assert that columns keep the pyarrow dtype."""
- atom = ATOMClassifier(X_bin, y_bin, engine="pyarrow", random_state=1)
+ atom = ATOMClassifier(X_pa, y_bin, random_state=1)
assert isinstance(atom.dtypes[0], pd.ArrowDtype)
atom.scale()
assert isinstance(atom.dtypes[0], pd.ArrowDtype)
diff --git a/tests/test_basemodel.py b/tests/test_basemodel.py
index db2b3009f..fb1e81183 100644
--- a/tests/test_basemodel.py
+++ b/tests/test_basemodel.py
@@ -10,6 +10,7 @@
import numpy as np
import pandas as pd
+import polars as pl
import pytest
import requests
from optuna.distributions import CategoricalDistribution, IntDistribution
@@ -871,6 +872,13 @@ def test_inverse_transform():
assert_frame_equal(atom.lr.inverse_transform(atom.lr.X), X_bin)
+def test_inverse_transform_output():
+ """Assert that the output type is determined by the data engine."""
+ atom = ATOMClassifier(X_bin, y_bin, engine="polars", random_state=1)
+ atom.run("Tree")
+ assert isinstance(atom.tree.inverse_transform(X_bin), pl.DataFrame)
+
+
def test_save_estimator():
"""Assert that the save_estimator saves a pickle file."""
atom = ATOMClassifier(X_bin, y_bin, random_state=1)
@@ -918,6 +926,13 @@ def test_transform():
assert all(-3 <= v <= 3 for v in X.to_numpy().ravel()) # Data is scaled
+def test_transform_output():
+ """Assert that the output type is determined by the data engine."""
+ atom = ATOMClassifier(X_bin, y_bin, engine="polars", random_state=1)
+ atom.run("Tree")
+ assert isinstance(atom.tree.transform(X_bin), pl.DataFrame)
+
+
# Test ClassRegModel ================================================== >>
def test_classreg_get_tags():
diff --git a/tests/test_baserunner.py b/tests/test_baserunner.py
index 1a99684c0..01160c2b1 100644
--- a/tests/test_baserunner.py
+++ b/tests/test_baserunner.py
@@ -18,7 +18,7 @@
from sklearn.discriminant_analysis import LinearDiscriminantAnalysis as LDA
from atom import ATOMClassifier, ATOMForecaster, ATOMRegressor
-from atom.branch import Branch
+from atom.data import Branch
from atom.training import DirectClassifier, DirectForecaster
from atom.utils.types import SPTuple
from atom.utils.utils import NotFittedError, merge
@@ -319,7 +319,7 @@ def test_results_property_train_sizing():
assert list(atom.results.index.get_level_values(0)) == [0.2, 0.4, 0.6, 0.8, 1.0]
-# Test _set_index ================================================== >>
+# Test _get_data =================================================== >>
def test_index_is_true():
"""Assert that the indices are left as is when index=True."""
@@ -406,8 +406,6 @@ def test_duplicate_indices():
ATOMClassifier(X_bin, X_bin, index=True, random_state=1)
-# Test _get_stratify_columns======================================== >>
-
@pytest.mark.parametrize("stratify", [True, -1, "target", [-1]])
def test_stratify_options(stratify):
"""Assert that the data can be stratified among data sets."""
@@ -437,8 +435,6 @@ def test_stratify_invalid_column_str():
ATOMClassifier(X_bin, y_bin, stratify="invalid", random_state=1)
-# Test _get_data =================================================== >>
-
def test_input_is_y_without_arrays():
"""Assert that input y through parameter works."""
atom = ATOMForecaster(y=y_fc, random_state=1)
@@ -597,7 +593,7 @@ def test_input_is_train_test_with_parameter_y():
def test_input_is_train_test_for_forecast():
"""Assert that input train, test works for forecast tasks."""
- trainer = DirectForecaster("ES", random_state=1)
+ trainer = DirectForecaster("Croston", random_state=1)
trainer.run(fc_train, fc_test)
assert_series_equal(trainer.y, pd.concat([fc_train, fc_test]))
@@ -606,8 +602,8 @@ def test_input_is_3_tuples():
"""Assert that the 3 tuples input works."""
X_train = bin_train.iloc[:, :-1]
y_train = bin_train.iloc[:, -1]
- X_test = bin_test.iloc[100:-20, :-1]
- y_test = bin_test.iloc[100:-20, -1]
+ X_test = bin_test.iloc[:-20, :-1]
+ y_test = bin_test.iloc[:-20, -1]
X_holdout = bin_test.iloc[-20:, :-1]
y_holdout = bin_test.iloc[-20:, -1]
@@ -626,7 +622,7 @@ def test_input_is_train_test_holdout():
def test_4_data_provided():
- """Assert that the 4 elements input works."""
+ """Assert that the four-element input works."""
X_train = bin_train.iloc[:, :-1]
X_test = bin_test.iloc[:, :-1]
y_train = bin_train.iloc[:, -1]
@@ -638,11 +634,11 @@ def test_4_data_provided():
def test_6_data_provided():
- """Assert that the 6 elements input works."""
+ """Assert that the six-element input works."""
X_train = bin_train.iloc[:, :-1]
y_train = bin_train.iloc[:, -1]
- X_test = bin_test.iloc[100:-20, :-1]
- y_test = bin_test.iloc[100:-20, -1]
+ X_test = bin_test.iloc[:-20, :-1]
+ y_test = bin_test.iloc[:-20, -1]
X_holdout = bin_test.iloc[-20:, :-1]
y_holdout = bin_test.iloc[-20:, -1]
diff --git a/tests/test_basetrainer.py b/tests/test_basetrainer.py
index e782fdd4b..26f1af366 100644
--- a/tests/test_basetrainer.py
+++ b/tests/test_basetrainer.py
@@ -9,7 +9,6 @@
import mlflow
import pytest
-import ray
from mlflow.tracking.fluent import ActiveRun
from optuna.distributions import CategoricalDistribution, IntDistribution
from optuna.pruners import MedianPruner
@@ -376,21 +375,32 @@ def test_errors_keep():
assert trainer._models == [trainer.lda]
-@patch("atom.basetransformer.ray", MagicMock())
-@patch("atom.basetrainer.ray", MagicMock())
+# @patch("atom.basetransformer.ray", MagicMock())
+# @patch("atom.basetrainer.ray", MagicMock())
def test_parallel_with_ray():
"""Assert that parallel runs successfully with ray backend."""
trainer = DirectClassifier(
models=["LR", "LDA"],
parallel=True,
- n_jobs=1,
+ n_jobs=2,
backend="ray",
random_state=1,
)
- # Fails because Mock returns empty list
+ # Fails because MagicMock returns an empty list
with pytest.raises(RuntimeError, match=".*All models failed.*"):
trainer.run(bin_train, bin_test)
- ray.shutdown()
+
+
+def test_parallel_with_dask():
+ """Assert that parallel runs successfully with dask backend."""
+ trainer = DirectClassifier(
+ models=["LR", "LDA"],
+ parallel=True,
+ n_jobs=2,
+ backend="dask",
+ random_state=1,
+ )
+ trainer.run(bin_train, bin_test)
@patch("atom.basetrainer.Parallel", MagicMock())
diff --git a/tests/test_basetransformer.py b/tests/test_basetransformer.py
index 3cffdfcd0..9357abdb9 100644
--- a/tests/test_basetransformer.py
+++ b/tests/test_basetransformer.py
@@ -57,30 +57,11 @@ def test_device_parameter():
assert os.environ["CUDA_VISIBLE_DEVICES"] == "0"
-@patch("ray.init")
-def test_engine_parameter_modin(ray):
- """Assert that ray is initialized when modin is data backend."""
- base = BaseTransformer(device="cpu", engine="modin")
- assert base.engine.data == "modin"
- assert ray.is_called_once
-
-
-def test_engine_parameter_env_var():
- """Assert that the environment variable is set."""
- base = BaseTransformer(device="cpu", engine="pyarrow")
- assert base.engine == EngineTuple(data="pyarrow", estimator="sklearn")
- assert os.environ["ATOM_DATA_ENGINE"] == base.engine.data
-
- base = BaseTransformer(device="cpu", engine="sklearnex")
- assert base.engine == EngineTuple(data="pandas", estimator="sklearnex")
- assert os.environ["ATOM_DATA_ENGINE"] == base.engine.data
-
-
-@patch.dict("sys.modules", {"sklearnex": None})
-def test_engine_parameter_no_sklearnex():
- """Assert that an error is raised when sklearnex is not installed."""
- with pytest.raises(ModuleNotFoundError, match=".*import scikit-learn-intelex.*"):
- BaseTransformer(device="cpu", engine={"estimator": "sklearnex"})
+@pytest.mark.parametrize("engine", [None, "pandas", "sklearn", {}, EngineTuple()])
+def test_engine_parameter(engine):
+ """Assert that the engine parameter can be initialized."""
+ base = BaseTransformer(engine=engine)
+ assert base.engine == EngineTuple()
@pytest.mark.skipif(machine() not in ("x86_64", "AMD64"), reason="Only x86 support")
@@ -103,6 +84,13 @@ def test_backend_parameter_ray(ray):
assert ray.is_called_once
+@patch("dask.distributed.Client")
+def test_backend_parameter_dask(dask):
+ """Assert that dask is initialized when selected."""
+ BaseTransformer(backend="dask")
+ assert dask.is_called_once
+
+
def test_backend_parameter():
"""Assert that other backends can be specified."""
base = BaseTransformer(backend="threading")
@@ -182,53 +170,6 @@ def test_device_id_invalid():
BaseTransformer(device="gpu:2,3")
-# Test _inherit ==================================================== >>
-
-def test_inherit():
- """Assert that the inherit method passes the parameters correctly."""
- base = BaseTransformer(n_jobs=2, random_state=2)
- svc = base._inherit(RandomForestClassifier())
- assert svc.get_params()["n_jobs"] == 2
- assert svc.get_params()["random_state"] == 2
-
-
-def test_inherit_with_fixed_params():
- """Assert that fixed parameters aren't inherited."""
- base = BaseTransformer(random_state=2)
- chain = base._inherit(ClassifierChain(SVC(random_state=3)), ("base_estimator__random_state",))
- assert chain.get_params()["random_state"] == 2
- assert chain.base_estimator.get_params()["random_state"] == 3
-
-
-def test_inherit_sp():
- """Assert that the seasonal periodicity is correctly inherited."""
- atom = ATOMForecaster(y_fc, sp=[12, 24], random_state=1)
- atom.run(
- models=["bats", "tbats"],
- est_params={
- "bats": {"use_box_cox": False, "use_trend": False, "use_arma_errors": False},
- "tbats": {"use_box_cox": False, "use_trend": False, "use_arma_errors": False},
- },
- )
- assert atom.bats.estimator.get_params()["sp"] == 12 # Single seasonality
- assert atom.tbats.estimator.get_params()["sp"] == [12, 24] # Multiple seasonality
-
-
-# Test _get_est_class ============================================== >>
-
-@pytest.mark.skipif(machine() not in ("x86_64", "AMD64"), reason="Only x86 support")
-def test_get_est_class_from_engine():
- """Assert that the class can be retrieved from an engine."""
- base = BaseTransformer(device="cpu", engine={"estimator": "sklearnex"})
- assert base._get_est_class("SVC", "svm") == SVC
-
-
-def test_get_est_class_from_default():
- """Assert that the class is retrieved from sklearn when import fails."""
- base = BaseTransformer(device="cpu", engine={"estimator": "sklearnex"})
- assert base._get_est_class("GaussianNB", "naive_bayes") == GaussianNB
-
-
# Test _check_input ============================================== >>
def test_input_is_copied():
@@ -266,7 +207,7 @@ def test_column_order_is_retained():
def test_incorrect_columns():
"""Assert that an error is raised when the provided columns do not match."""
- with pytest.raises(ValueError, match=".*features are different.*"):
+ with pytest.raises(ValueError, match=".*columns are different.*"):
BaseTransformer._check_input(X_bin, columns=["1", "2"])
@@ -323,12 +264,6 @@ def test_sparse_matrices_2_tuples():
assert atom[atom.columns[0]].dtype.name == "Sparse[int64, 0]"
-def test_target_is_dict():
- """Assert that the target column is assigned correctly for a dict."""
- _, y = BaseTransformer._check_input(X10, {"a": [0] * 10})
- assert isinstance(y, pd.Series)
-
-
def test_multioutput_str():
"""Assert that multioutput can be assigned by column name."""
X, y = BaseTransformer._check_input(X_bin, ["mean radius", "worst perimeter"])
@@ -397,6 +332,53 @@ def test_X_empty_df():
assert isinstance(y, pd.Series)
+# Test _inherit ==================================================== >>
+
+def test_inherit():
+ """Assert that the inherit method passes the parameters correctly."""
+ base = BaseTransformer(n_jobs=2, random_state=2)
+ svc = base._inherit(RandomForestClassifier())
+ assert svc.get_params()["n_jobs"] == 2
+ assert svc.get_params()["random_state"] == 2
+
+
+def test_inherit_with_fixed_params():
+ """Assert that fixed parameters aren't inherited."""
+ base = BaseTransformer(random_state=2)
+ chain = base._inherit(ClassifierChain(SVC(random_state=3)), ("base_estimator__random_state",))
+ assert chain.get_params()["random_state"] == 2
+ assert chain.base_estimator.get_params()["random_state"] == 3
+
+
+def test_inherit_sp():
+ """Assert that the seasonal periodicity is correctly inherited."""
+ atom = ATOMForecaster(y_fc, sp=[12, 24], random_state=1)
+ atom.run(
+ models=["bats", "tbats"],
+ est_params={
+ "bats": {"use_box_cox": False, "use_trend": False, "use_arma_errors": False},
+ "tbats": {"use_box_cox": False, "use_trend": False, "use_arma_errors": False},
+ },
+ )
+ assert atom.bats.estimator.get_params()["sp"] == 12 # Single seasonality
+ assert atom.tbats.estimator.get_params()["sp"] == [12, 24] # Multiple seasonality
+
+
+# Test _get_est_class ============================================== >>
+
+@pytest.mark.skipif(machine() not in ("x86_64", "AMD64"), reason="Only x86 support")
+def test_get_est_class_from_engine():
+ """Assert that the class can be retrieved from an engine."""
+ base = BaseTransformer(device="cpu", engine={"estimator": "sklearnex"})
+ assert base._get_est_class("SVC", "svm") == SVC
+
+
+def test_get_est_class_from_default():
+ """Assert that the class is retrieved from sklearn when import fails."""
+ base = BaseTransformer(device="cpu", engine={"estimator": "sklearnex"})
+ assert base._get_est_class("GaussianNB", "naive_bayes") == GaussianNB
+
+
# Test log ========================================================= >>
def test_log_severity_error():
diff --git a/tests/test_branch.py b/tests/test_data.py
similarity index 86%
rename from tests/test_branch.py
rename to tests/test_data.py
index 5484168d4..9bee37d88 100644
--- a/tests/test_branch.py
+++ b/tests/test_data.py
@@ -7,13 +7,19 @@
import glob
import os
from pathlib import Path
+from unittest.mock import MagicMock, patch
+import dask.dataframe as dd
+import numpy as np
import pandas as pd
+import polars as pl
+import pyarrow as pa
import pytest
from pandas.testing import assert_frame_equal
+from sklearn.preprocessing import MinMaxScaler, StandardScaler
from atom import ATOMClassifier, ATOMRegressor
-from atom.branch import Branch, BranchManager
+from atom.data import Branch, BranchManager
from atom.training import DirectClassifier
from atom.utils.utils import merge
@@ -276,14 +282,14 @@ def test_data_properties_to_df():
"""Assert that the data attributes are converted to a df at setter."""
atom = ATOMClassifier(X_bin, y_bin, random_state=1)
atom.X = X_bin_array
- assert isinstance(atom.X, pd.DataFrame)
+ assert isinstance(atom.branch.X, pd.DataFrame)
def test_data_properties_to_series():
"""Assert that the data attributes are converted to a series at setter."""
atom = ATOMClassifier(X_bin, y_bin, random_state=1)
atom.y = y_bin_array
- assert isinstance(atom.y, pd.Series)
+ assert isinstance(atom.branch.y, pd.Series)
def test_setter_error_unequal_rows():
@@ -304,7 +310,7 @@ def test_setter_error_unequal_columns():
"""Assert that an error is raised when the setter has unequal columns."""
atom = ATOMClassifier(X_bin, y_bin, random_state=1)
new_X = atom.train
- new_X.insert(0, "new_column", 1)
+ new_X["new_column"] = 1
with pytest.raises(ValueError, match="number of columns"):
atom.train = new_X
@@ -561,6 +567,30 @@ def test_load_no_dir():
atom.branch = "main"
+def test_check_scaling_scaler_in_pipeline():
+ """Assert that check_scaling returns True when there's a scaler in the pipeline."""
+ atom = ATOMClassifier(X_bin, y=y_bin, random_state=1)
+ assert not atom.branch.check_scaling()
+ atom.add(MinMaxScaler())
+ assert atom.branch.check_scaling()
+
+
+def test_check_scaling():
+ """Assert that the check_scaling method returns whether the data is scaled."""
+ scaler = StandardScaler()
+ scaler.__class__.__name__ = "OtherName"
+
+ atom = ATOMClassifier(X_bin, y=y_bin, random_state=1)
+ atom.add(scaler)
+ assert atom.branch.check_scaling()
+
+
+def test_check_scaling_drop_binary():
+ """Assert that binary rows are dropped to check scaling."""
+ atom = ATOMClassifier(np.tile(y10, (10, 1)), y=y10, random_state=1)
+ assert atom.branch.check_scaling()
+
+
# Test BranchManager =============================================== >>
def test_branchmanager_repr():
@@ -665,3 +695,76 @@ def test_reset():
assert len(atom._branches) == 1
assert not glob.glob("joblib/atom/Branch(main).pkl")
assert atom.og is atom.branch
+
+
+# Test data engines ================================================ >>
+
+def test_numpy_engine():
+ """Assert that the numpy engine returns a numpy array."""
+ atom = ATOMClassifier(X_bin, y_bin, engine="numpy", random_state=1)
+ assert isinstance(atom.dataset, np.ndarray)
+
+
+def test_pandas_numpy_engine():
+ """Assert that the pandas engine returns numpy dtypes."""
+ atom = ATOMClassifier(X_bin, y_bin, engine="pandas", random_state=1)
+ assert all(isinstance(dtype, np.dtype) for dtype in atom.dataset.dtypes)
+ assert isinstance(atom.y.dtype, np.dtype)
+
+
+def test_pandas_pyarrow_engine():
+ """Assert that the pandas-pyarrow engine returns pyarrow dtypes."""
+ atom = ATOMClassifier(X_bin, y_bin, engine="pandas-pyarrow", random_state=1)
+ assert all(isinstance(dtype, pd.ArrowDtype) for dtype in atom.dataset.dtypes)
+ assert isinstance(atom.y.dtype, pd.ArrowDtype)
+
+
+def test_polars_engine():
+ """Assert that the polars engine returns polars types."""
+ atom = ATOMClassifier(X_bin, y_bin, engine="polars", random_state=1)
+ assert isinstance(atom.X, pl.DataFrame)
+ assert isinstance(atom.y, pl.Series)
+
+
+def test_polars_lazy_engine():
+ """Assert that the polars-lazy engine returns polars types."""
+ atom = ATOMClassifier(X_bin, y_bin, engine="polars-lazy", random_state=1)
+ assert isinstance(atom.X, pl.LazyFrame)
+ assert isinstance(atom.y, pl.Series)
+
+
+def test_pyarrow_engine():
+ """Assert that the pyarrow engine returns pyarrow types."""
+ atom = ATOMClassifier(X_bin, y_bin, engine="pyarrow", random_state=1)
+ assert isinstance(atom.X, pa.Table)
+ assert isinstance(atom.y, pa.Array)
+
+
+@patch.dict("sys.modules", {"modin": MagicMock(spec=["__spec__", "pandas"])})
+def test_modin_engine():
+ """Assert that the modin engine returns modin types."""
+ atom = ATOMClassifier(X_bin, y_bin, engine="modin", random_state=1)
+ assert "DataFrame" in str(atom.X)
+ assert "Series" in str(atom.y)
+
+
+def test_dask_engine():
+ """Assert that the dask engine returns dask types."""
+ atom = ATOMClassifier(X_bin, y_bin, engine="dask", random_state=1)
+ assert isinstance(atom.X, dd.DataFrame)
+ assert isinstance(atom.y, dd.Series)
+
+
+@patch.dict("sys.modules", {"pyspark.sql": MagicMock(spec=["__spec__", "SparkSession"])})
+def test_pyspark_engine():
+ """Assert that the pyspark engine returns pyspark types."""
+ atom = ATOMClassifier(X_bin, y_bin, engine="pyspark", random_state=1)
+ assert "createDataFrame" in str(atom.X)
+
+
+@patch.dict("sys.modules", {"pyspark": MagicMock(spec=["__spec__", "pandas"])})
+def test_pyspark_pandas_engine():
+ """Assert that the pyspark-pandas engine returns pyspark pandas types."""
+ atom = ATOMClassifier(X_bin, y_bin, engine="pyspark-pandas", random_state=1)
+ assert "DataFrame" in str(atom.X)
+ assert "Series" in str(atom.y)
diff --git a/tests/test_data_cleaning.py b/tests/test_data_cleaning.py
index c888b4525..8c5bd88ff 100644
--- a/tests/test_data_cleaning.py
+++ b/tests/test_data_cleaning.py
@@ -423,8 +423,8 @@ def test_missing_values_are_propagated():
def test_unknown_classes_are_imputed():
"""Assert that unknown classes are imputed."""
encoder = Encoder()
- encoder.fit(["a", "b", "b", "a"])
- assert encoder.transform(["c"]).iloc[0, 0] == -1.0
+ encoder.fit([["a"], ["b"], ["b"], ["a"]])
+ assert encoder.transform([["c"]]).iloc[0, 0] == -1.0
def test_ordinal_encoder():
diff --git a/tests/test_nlp.py b/tests/test_nlp.py
index b7733ae3a..e4826620a 100644
--- a/tests/test_nlp.py
+++ b/tests/test_nlp.py
@@ -5,7 +5,6 @@
"""
-from unittest.mock import MagicMock, patch
import pandas as pd
import pytest
@@ -191,21 +190,6 @@ def test_hashing():
assert "hash1" in X
-@patch.dict(
- "sys.modules",
- {
- "cuml": MagicMock(spec=["__spec__"]),
- "cuml.common.device_selection": MagicMock(spec=["set_global_device_type"]),
- "cuml.internals.memory_utils": MagicMock(spec=["set_global_output_type"]),
- "cuml.feature_extraction.text": MagicMock(),
- },
-)
-def test_gpu():
- """Assert that the gpu implementation calls the get method of matrix."""
- vectorizer = Vectorizer(device="gpu", engine="cuml")
- pytest.raises(ValueError, vectorizer.fit_transform, X_text)
-
-
def test_return_sparse():
"""Assert that the output is sparse."""
X = Vectorizer(strategy="bow", return_sparse=True).fit_transform(X_text, y10)
diff --git a/tests/test_pipeline.py b/tests/test_pipeline.py
index 586c2a737..d87681fc0 100644
--- a/tests/test_pipeline.py
+++ b/tests/test_pipeline.py
@@ -7,6 +7,7 @@
import numpy as np
import pandas as pd
+import pyarrow as pa
import pytest
from pandas.testing import assert_frame_equal
from sklearn.preprocessing import LabelEncoder, StandardScaler
@@ -216,6 +217,18 @@ def test_predict_var(pipeline_ts):
assert isinstance(pipeline_ts.predict_var(fh=range(3)), pd.DataFrame)
+def test_set_output(pipeline):
+ """Assert that the set_output method determines the data engine."""
+ pl = pipeline(model=False)
+ assert isinstance(pl.transform(X_bin), pd.DataFrame)
+
+ pl.set_output(transform="numpy")
+ assert isinstance(pl.fit_transform(X_bin, y_bin)[0], np.ndarray)
+
+ pl.set_output(transform="pyarrow")
+ assert isinstance(pl.inverse_transform(X_bin), pa.Table)
+
+
def test_score_no_parameters(pipeline_ts):
"""Assert that an error is raised when X and fh are both None."""
with pytest.raises(ValueError, match=".*cannot be both None.*"):
diff --git a/tests/test_training.py b/tests/test_training.py
index bc3d8f6c9..ce90a0acb 100644
--- a/tests/test_training.py
+++ b/tests/test_training.py
@@ -36,7 +36,7 @@ def test_models_are_restored():
)
sh.run(reg_train, reg_test)
assert "Tree" not in sh._models # The original model is deleted
- assert all(m in sh.models for m in ("Tree4", "AdaB2", "LGB1"))
+ assert all(m in sh.models for m in ("Tree4", "AdaB2", "AdaB1"))
def test_ts_int_train_sizes():
diff --git a/tests/test_utils.py b/tests/test_utils.py
index 486c94eb3..12f107676 100644
--- a/tests/test_utils.py
+++ b/tests/test_utils.py
@@ -6,7 +6,7 @@
"""
from datetime import timedelta
-from unittest.mock import Mock, patch
+from unittest.mock import patch
import numpy as np
import pytest
@@ -18,9 +18,7 @@
from atom import show_versions
from atom.pipeline import Pipeline
from atom.utils.patches import VotingClassifier, VotingRegressor
-from atom.utils.utils import (
- ClassMap, check_is_fitted, time_to_str, to_df, to_series,
-)
+from atom.utils.utils import ClassMap, check_is_fitted, time_to_str
from .conftest import X_bin, X_reg, y_bin, y_reg
@@ -158,9 +156,3 @@ def test_time_to_string():
assert time_to_str(timedelta(seconds=17).total_seconds()).startswith("17.00")
assert time_to_str(timedelta(minutes=1, seconds=2).total_seconds()) == "01m:02s"
assert time_to_str(timedelta(hours=3, minutes=8).total_seconds()) == "03h:08m:00s"
-
-
-def test_to_pandas_with_cuml():
- """Assert that cuML objects use the to_pandas method."""
- to_df(Mock(spec=["to_pandas"]), columns=[0, 1])
- to_series(Mock(spec=["to_pandas"]))