def test_add_intercept_raises_chunks():
X = da.random.uniform(size=(10, 4), chunks=(4, 2))
with pytest.raises(ValueError) as m:
add_intercept(X)
assert m.match("Chunking is only allowed")
def test_add_intercept_raises_ndim():
X = da.random.uniform(size=10, chunks=5)
with pytest.raises(ValueError) as m:
add_intercept(X)
assert m.match("'X' should have 2 dimensions")
def test_add_intercept_dask_dataframe():
X = dd.from_pandas(pd.DataFrame({"A": [1, 2, 3]}), npartitions=2)
result = add_intercept(X)
expected = dd.from_pandas(
pd.DataFrame(
{"intercept": [1, 1, 1], "A": [1, 2, 3]}, columns=["intercept", "A"]
),
npartitions=2,
)
assert_eq(result, expected)
df = dd.from_pandas(pd.DataFrame({"intercept": [1, 2, 3]}), npartitions=2)
with pytest.raises(ValueError):
add_intercept(df)
def test_add_intercept_unknown_ndim():
X = dd.from_pandas(pd.DataFrame(np.ones((10, 5))), 2).values
result = add_intercept(X)
expected = np.ones((10, 6))
da.utils.assert_eq(result, expected)
