Merge branch 'main' into issue19

NEWS.md

@@ -1,3 +1,10 @@
+# Development
+
+New functions: 
+
+* `get_dataset()`
+* `fit_sklearn()`, `fit_statsmodels()`, `fit_linearmodels()`
+
 # 0.0.14
 
 * Thanks to Narwhals, marginaleffects can now ingest data frames in multiple formats and convert them to the Polars representation that we need internally. This no longer requires external dependencies like Pandas or DuckDB. Thanks to @artiom-matvei.

marginaleffects/__init__.py

@@ -9,6 +9,7 @@
 from .model_statsmodels import fit_statsmodels
 from .model_sklearn import fit_sklearn
 from .model_linearmodels import fit_linearmodels
+from .utils import get_dataset
 
 __all__ = [
     "avg_comparisons",
@@ -25,4 +26,5 @@
     "fit_statsmodels",
     "fit_sklearn",
     "fit_linearmodels",
+    "get_dataset",
 ]

marginaleffects/utils.py

@@ -168,3 +168,47 @@ def wrapper(*args, **kwargs):
         return validator(*args, **kwargs)
 
     return wrapper
+
+
+def get_dataset(dataset: str, docs: bool = False):
+    """
+    Download and read a dataset as a Polars DataFrame or return documentation link.
+
+    Parameters
+    ----------
+    dataset : str
+        The dataset to download. Must be one of "affairs", "airbnb", "immigration", "military", "thornton".
+    docs : bool, optional
+        If True, return the documentation URL instead of the dataset. Default is False.
+
+    Returns
+    -------
+    Union[str, pl.DataFrame]
+        A string representing the documentation URL if `docs` is True, or
+        a Polars DataFrame containing the dataset if `docs` is False.
+
+    Raises
+    ------
+    ValueError
+        If the dataset is not among the specified choices.
+    """
+    datasets = {
+        "affairs": "https://marginaleffects.com/data/affairs",
+        "airbnb": "https://marginaleffects.com/data/airbnb",
+        "immigration": "https://marginaleffects.com/data/immigration",
+        "military": "https://marginaleffects.com/data/military",
+        "thornton": "https://marginaleffects.com/data/thornton",
+    }
+
+    if dataset not in datasets:
+        raise ValueError(
+            f"Invalid dataset choice. Expected one of {list(datasets.keys())}."
+        )
+
+    base_url = datasets[dataset]
+
+    if docs:
+        return f"{base_url}.html"
+
+    df = pl.read_parquet(f"{base_url}.parquet")
+    return df