import numpy as np
import polars as pl
import statsmodels.formula.api as smf
from marginaleffects import *
from marginaleffects.classes import MarginaleffectsDataFrame
from .utilities import *
dat = pl.read_csv("tests/data/impartiality.csv") \
.with_columns(pl.col("impartial").cast(pl.Int8))
m = smf.logit(
"impartial ~ equal * democracy + continent",
data = dat.to_pandas()
).fit()
p = plot_predictions(m, by = ["democracy", "continent"])