Fix issue 96 mnlogit tests

marginaleffects/by.py

@@ -16,7 +16,7 @@ def get_by(model, estimand, newdata, by=None, wts=None):
     elif by is False:
         return estimand
 
-    if "group" in estimand.columns:
+    if "group" in estimand.columns and "group" not in by:
         by = ["group"] + by
 
     if "rowid" in estimand.columns and "rowid" in newdata.columns:

marginaleffects/model_statsmodels.py

@@ -41,7 +41,9 @@ def get_vcov(self, vcov=True):
 
         if V is not None:
             V = np.array(V)
-            if V.shape != (len(self.coef), len(self.coef)):
+            # flatten in case coef are for multi-index dataframes
+            expected_dim = len(self.coef.flatten())
+            if V.shape != (expected_dim, expected_dim):
                 raise ValueError(
                     "vcov must be a square numpy array with dimensions equal to the length of self.coef"
                 )

marginaleffects/uncertainty.py

@@ -39,6 +39,36 @@ def get_jacobian(func, coefs, eps_vcov=None):
 
 
 def get_se(J, V):
+    # J are different in python versus R
+    '''Python
+    array([[ 3.26007929e-03, -1.11398030e+05, -1.07005746e+07,
+         7.82170746e+07,  2.78495302e+07,  1.34757804e+00],
+       [ 3.00860803e+05,  1.11398009e+05,  3.15899457e+07,
+         4.00666771e+07,  1.42659178e+07, -5.02687187e+00],
+       [-3.00860806e+05,  2.03275902e-02, -2.08893711e+07,
+        -1.18283752e+08, -4.21154480e+07,  3.67929389e+00],
+       ...,
+       [ 1.94566192e-02,  6.31280874e+04,  6.06390551e+06,
+         6.95824154e+07,  2.47751212e+07,  4.27175339e-01],
+       [ 1.19234156e+05, -6.31280929e+04,  1.25194134e+07,
+        -2.27053711e+07, -8.08434379e+06, -1.51166669e+00],
+       [-1.19234176e+05,  5.47722901e-03, -1.85833189e+07,
+        -4.68770444e+07, -1.66907774e+07,  1.08449133e+00]])
+    '''
+    '''R
+                 [,1]          [,2]          [,3]          [,4]          [,5]          [,6]
+   [1,]  4.211085e-03  0.1129098874  7.622065e-01  6.449953e-03  2.390449e-01  1.167441e+00
+   [2,]  2.317994e-03 -0.0002889116  4.311471e-01  1.240511e-02  4.599770e-01  2.307350e+00
+   [3,] -1.458913e-02 -0.7474941310 -2.844880e+00  2.018562e-02  7.334876e-01  3.936196e+00
+   [4,] -1.768363e-02 -0.7516309380 -3.412940e+00  1.977302e-02  5.853414e-01  3.816193e+00
+   [5,] -6.148701e-03 -0.3694016210 -1.168253e+00  1.883557e-02  6.798748e-01  3.578758e+00
+   [6,]  3.851085e-03  0.0970890213  6.970463e-01  6.842971e-03  2.519078e-01  1.238578e+00
+   [7,] -1.807982e-02 -0.8520217380 -3.525564e+00  2.090703e-02  7.173019e-01  4.076872e+00
+   [8,] -1.324527e-02 -0.5110488046 -2.556336e+00  1.326720e-02  2.654043e-01  2.560569e+00
+   [9,]  7.328301e-03  0.1888463791  1.392377e+00  1.094683e-02  4.361391e-01  2.079898e+00
+  [10,] -4.511196e-03 -0.2639279320 -8.390823e-01  1.642779e-02  5.682850e-01  3.055570e+00
+  ...
+    '''
     se = np.sqrt(np.sum((J @ V) * J, axis=1))
     return se
 

tests/conftest.py

@@ -27,6 +27,8 @@
     null_values="NA",
 ).drop_nulls()
 
+penguins_with_nulls = pl.read_csv("tests/data/penguins.csv")
+
 quine = pl.read_csv("tests/data/quine.csv")
 
 

tests/statsmodels/test_statsmodels_mnlogit.py

@@ -0,0 +1,216 @@
+import polars as pl
+import pytest
+import statsmodels.formula.api as smf
+from polars.testing import assert_series_equal
+from tests.conftest import penguins, penguins_with_nulls
+from marginaleffects import *
+
+
+island_categories = ["Biscoe", "Dream", "Torgersen"]
+island_mapping = {island: code for code, island in enumerate(island_categories)}
+dat = penguins.drop_nulls(
+    ["species", "island", "bill_length_mm", "flipper_length_mm"]
+).with_columns(
+    pl.col("island").replace_strict(island_mapping),
+    pl.col("flipper_length_mm").cast(pl.Float64),
+)
+mod = smf.mnlogit("island ~ bill_length_mm + flipper_length_mm", dat.to_pandas()).fit()
+
+
+# @pytest.mark.skip(reason="statsmodels vcov is weird")
+def test_predictions_01():
+    """
+    R code to generate the csv file data
+    # Load necessary packages
+    library(nnet)
+    library(dplyr)
+    library(marginaleffects)
+
+    # Load and prepare the data
+    penguins <- read.csv("https://vincentarelbundock.github.io/Rdatasets/csv/palmerpenguins/penguins.csv")
+    penguins_clean <- penguins %>%
+    select(island, bill_length_mm, flipper_length_mm) %>%
+    na.omit()
+    penguins_clean$island <- relevel(factor(penguins_clean$island), ref = "Biscoe")
+
+    # Fit the model using nnet!
+    model_r <- multinom(island ~ bill_length_mm  + flipper_length_mm, data = penguins_clean)
+
+    # Extract and display coefficients in case the models do not match
+    # coef_r <- coef(model_r)
+    # print("Coefficients from R model:")
+    # print(coef_r)
+
+    predictions(model_r)
+
+    """
+    penguins_clean = penguins_with_nulls.select(
+        ["island", "bill_length_mm", "flipper_length_mm"]
+    ).drop_nulls()
+
+    # Define island categories and create a mapping
+    island_categories = ["Biscoe", "Dream", "Torgersen"]
+    island_mapping = {island: code for code, island in enumerate(island_categories)}
+
+    # Map 'island' to integer codes
+    penguins_clean = penguins_clean.with_columns(
+        pl.col("island").replace_strict(island_mapping)
+    )
+
+    model_py = smf.mnlogit(
+        "island ~ bill_length_mm + flipper_length_mm", data=penguins_clean
+    ).fit()
+    unknown = predictions(model_py).sort(by=["rowid", "group"])
+
+    known = pl.read_csv("tests/r/test_statsmodels_mnlogit_predictions_01.csv")
+    known = known.with_columns(
+        pl.col("group").replace(island_mapping, return_dtype=pl.Int8).alias("group")
+    )
+    known = known.sort(by=["rowid", "group"])
+
+    assert_series_equal(known["estimate"], unknown["estimate"], rtol=1e-2)
+
+
+# @pytest.mark.skip(reason="statsmodels vcov is weird")
+def test_predictions_02():
+    unknown = predictions(mod, by="species").sort(["group", "species"])
+    known = (
+        pl.read_csv("tests/r/test_statsmodels_mnlogit_predictions_02.csv")
+        .with_columns(pl.col("group").replace_strict(island_mapping))
+        .sort(["group", "species"])
+    )
+    assert_series_equal(known["estimate"], unknown["estimate"], rtol=1e-1)
+
+
+def test_comparisons_01():
+    penguins_clean = penguins_with_nulls.select(
+        ["island", "bill_length_mm", "flipper_length_mm"]
+    ).drop_nulls()
+
+    # Define island categories and create a mapping
+    island_categories = ["Biscoe", "Dream", "Torgersen"]
+    island_mapping = {island: code for code, island in enumerate(island_categories)}
+
+    # Map 'island' to integer codes
+    penguins_clean = penguins_clean.with_columns(
+        pl.col("island").replace_strict(island_mapping)
+    )
+
+    mod = smf.mnlogit(
+        "island ~ bill_length_mm + flipper_length_mm", data=penguins_clean
+    ).fit()
+
+    unknown = (
+        comparisons(mod)
+        .with_columns(pl.col("group").replace(island_mapping))
+        .sort(["rowid", "term", "group"])
+        .filter(pl.col("term") != "flipper_length_mm")
+    )
+    known = (
+        pl.read_csv("tests/r/test_statsmodels_mnlogit_comparisons_01.csv")
+        .with_columns(pl.col("group").replace(island_mapping))
+        .sort(["rowid", "term", "group"])
+        .filter(pl.col("term") != "flipper_length_mm")
+    )
+    # new_column_names = {col: col.replace('.', '_') for col in known.columns}
+    # known = known.rename(new_column_names)
+
+    assert_series_equal(known["estimate"].head(), unknown["estimate"].head(), rtol=3e-1)
+
+    assert_series_equal(known["estimate"], unknown["estimate"], atol=1e-1)
+
+
+def test_comparisons_03():
+    penguins_clean = penguins_with_nulls.select(
+        ["island", "bill_length_mm", "flipper_length_mm"]
+    ).drop_nulls()
+
+    # Define island categories and create a mapping
+    island_categories = ["Biscoe", "Dream", "Torgersen"]
+    island_mapping = {island: code for code, island in enumerate(island_categories)}
+
+    # Map 'island' to integer codes
+    penguins_clean = penguins_clean.with_columns(
+        pl.col("island").replace_strict(island_mapping)
+    )
+
+    mod = smf.mnlogit(
+        "island ~ flipper_length_mm", data=penguins_clean
+    ).fit()
+
+    unknown = (
+        comparisons(mod)
+        .with_columns(pl.col("group").replace(island_mapping))
+        .sort(["rowid", "term", "group"])
+    )
+    known = (
+        pl.read_csv("tests/r/test_statsmodels_mnlogit_comparisons_03.csv")
+        .with_columns(pl.col("group").replace(island_mapping))
+        .sort(["rowid", "term", "group"])
+    )
+
+    assert_series_equal(known["estimate"].head(), unknown["estimate"].head(), atol=1e-3)
+
+    assert_series_equal(known["estimate"], unknown["estimate"], atol=1e-3)
+
+
+def test_comparisons_04():
+    penguins_clean = penguins_with_nulls.select(
+        ["island", "bill_length_mm", "flipper_length_mm"]
+    ).drop_nulls()
+
+    # Define island categories and create a mapping
+    island_categories = ["Biscoe", "Dream", "Torgersen"]
+    island_mapping = {island: code for code, island in enumerate(island_categories)}
+
+    # Map 'island' to integer codes
+    penguins_clean = penguins_clean.with_columns(
+        pl.col("island").replace_strict(island_mapping)
+    )
+
+    mod = smf.mnlogit(
+        "island ~ flipper_length_mm + bill_length_mm", data=penguins_clean
+    ).fit()
+
+    unknown = (
+        comparisons(mod)
+        .with_columns(pl.col("group").replace(island_mapping))
+        .sort(["rowid", "term", "group"])
+    )
+    known = (
+        pl.read_csv("tests/r/test_statsmodels_mnlogit_comparisons_04.csv")
+        .with_columns(pl.col("group").replace(island_mapping))
+        .sort(["rowid", "term", "group"])
+    )
+
+    assert_series_equal(known["estimate"].head(), unknown["estimate"].head(), atol=1e-3)
+
+    assert_series_equal(known["estimate"], unknown["estimate"], atol=1e-3)
+
+# Function to print visual comparison
+def compare_polars_tables(known, unknown, index=0):
+    headers = ["Column", "Table known Value", "Table unknown Value", "Difference"]
+    row_format = "{:<25} {:<25} {:<25} {:<15}"
+
+    print(row_format.format(*headers))
+    print("-" * 60)
+
+    for col in known.columns:
+        val1 = known[col][index]
+        val2 = unknown[col][index]
+        difference = "Yes" if val1 != val2 else "No"
+
+        print(row_format.format(col, val1, val2, difference))
+    print(row_format.format("index", index, index, "No"))
+
+# # @pytest.mark.skip(reason="statsmodels vcov is weird")
+# def test_comparisons_02():
+#     unknown = (
+#         comparisons(mod, by=["group", "species"])
+#         .with_columns(pl.col("group").map_dict(r))
+#         .sort(["term", "group", "species"])
+#     )
+#     known = pl.read_csv("tests/r/test_statsmodels_mnlogit_comparisons_02.csv").sort(
+#         ["term", "group", "species"]
+#     )
+#     assert_series_equal(known["estimate"], unknown["estimate"], rtol=1e-2)