Adding arguments to model builder (#1867)

* Adding arguments to model builder

* changes

Author: cetagostini

pymc_marketing/model_builder.py

@@ -456,13 +456,20 @@ def set_idata_attrs(
         idata.attrs = attrs
         return idata
 
-    def save(self, fname: str) -> None:
+    def save(self, fname: str, **kwargs) -> None:
         """Save the model's inference data to a file.
 
         Parameters
         ----------
         fname : str
             The name and path of the file to save the inference data with model parameters.
+        **kwargs
+            Additional keyword arguments to pass to arviz.InferenceData.to_netcdf().
+            Common options include:
+            - engine : str, optional (default "netcdf4")
+                Library to use for writing files.
+            - groups : list of str, optional
+                Groups to save to netcdf. If None, all groups are saved.
 
         Returns
         -------
@@ -483,14 +490,20 @@ def save(self, fname: str) -> None:
         >>>         super().__init__()
         >>> model = MyModel()
         >>> model.fit(X, y)
+        >>> # Basic save
+        >>> model.save("model_results.nc")
+        >>>
+        >>> # Save with specific options
         >>> model.save(
-        ...     "model_results.nc"
-        ... )  # This will call the overridden method in MyModel
+        ...     "model_results.nc",
+        ...     engine="netcdf4",
+        ...     groups=["posterior", "log_likelihood"],
+        ... )
 
         """
         if self.idata is not None and "posterior" in self.idata:
             file = Path(str(fname))
-            self.idata.to_netcdf(str(file))
+            self.idata.to_netcdf(str(file), **kwargs)
         else:
             raise RuntimeError("The model hasn't been fit yet, call .fit() first")
 

tests/mmm/test_mmm.py

@@ -765,6 +765,63 @@ def test_save_load(self, mmm_fitted: MMM):
         assert model.sampler_config == model2.sampler_config
         os.remove("test_save_load")
 
+    def test_save_load_with_kwargs(self, mmm_fitted: MMM):
+        """Test save/load functionality with kwargs (engine and groups)."""
+        model = mmm_fitted
+
+        # Use kwargs to test functionality - ArviZ supports engine and groups
+        # Note: ArviZ's to_netcdf has limited compression support compared to xarray
+        compression_kwargs = {
+            "engine": "h5netcdf",  # Alternative engine that may have better compression
+        }
+
+        model.save("test_save_load_kwargs", **compression_kwargs)
+
+        # Load and verify
+        model2 = MMM.load("test_save_load_kwargs")
+        assert model.date_column == model2.date_column
+        assert model.control_columns == model2.control_columns
+        assert model.channel_columns == model2.channel_columns
+        assert model.adstock.l_max == model2.adstock.l_max
+        assert model.validate_data == model2.validate_data
+        assert model.yearly_seasonality == model2.yearly_seasonality
+        assert model.model_config == model2.model_config
+        assert model.sampler_config == model2.sampler_config
+
+        os.remove("test_save_load_kwargs")
+
+    def test_save_load_engine_comparison(self, mmm_fitted: MMM):
+        """Test save/load with different engines and kwargs options."""
+        model = mmm_fitted
+
+        # Save with default engine
+        model.save("test_save_load_default")
+
+        # Save with h5netcdf engine (demonstrates kwargs functionality)
+        engine_kwargs = {
+            "engine": "h5netcdf",
+        }
+        model.save("test_save_load_h5netcdf", **engine_kwargs)
+
+        # Verify both files exist
+        assert os.path.exists("test_save_load_default")
+        assert os.path.exists("test_save_load_h5netcdf")
+
+        # Verify both can be loaded successfully and have the same data
+        model_default = MMM.load("test_save_load_default")
+        model_h5netcdf = MMM.load("test_save_load_h5netcdf")
+
+        # Both should have the same model configuration
+        assert (
+            model.model_config
+            == model_default.model_config
+            == model_h5netcdf.model_config
+        )
+
+        # Clean up
+        os.remove("test_save_load_default")
+        os.remove("test_save_load_h5netcdf")
+
     def test_fail_id_after_load(self, monkeypatch, toy_X, toy_y):
         # This is the new behavior for the property
         def mock_property(self):

tests/test_model_builder.py

@@ -13,6 +13,7 @@
 #   limitations under the License.
 import hashlib
 import json
+import os
 import re
 import sys
 import tempfile
@@ -221,6 +222,73 @@ def test_save_without_fit_raises_runtime_error():
         model_builder.save("saved_model")
 
 
+def test_save_with_kwargs(fitted_model_instance):
+    """Test that kwargs are properly passed to to_netcdf"""
+    import unittest.mock as mock
+
+    with mock.patch.object(fitted_model_instance.idata, "to_netcdf") as mock_to_netcdf:
+        temp = tempfile.NamedTemporaryFile(mode="w", encoding="utf-8", delete=False)
+
+        # Test with kwargs supported by InferenceData.to_netcdf()
+        kwargs = {"engine": "netcdf4", "groups": ["posterior", "log_likelihood"]}
+
+        fitted_model_instance.save(temp.name, **kwargs)
+
+        # Verify to_netcdf was called with the correct arguments
+        mock_to_netcdf.assert_called_once_with(temp.name, **kwargs)
+        temp.close()
+
+
+def test_save_with_kwargs_integration(fitted_model_instance):
+    """Test save function with actual kwargs (integration test)"""
+
+    temp = tempfile.NamedTemporaryFile(mode="w", encoding="utf-8", delete=False)
+    temp_path = temp.name
+    temp.close()
+
+    try:
+        # Test with specific groups - this tests that kwargs are passed through
+        fitted_model_instance.save(temp_path, groups=["posterior"])
+
+        # Verify file was created successfully
+        assert os.path.exists(temp_path)
+
+        # Verify we can read the file and it contains the expected groups
+        from pymc_marketing.utils import from_netcdf
+
+        loaded_idata = from_netcdf(temp_path)
+        assert "posterior" in loaded_idata.groups()
+        # Should only have posterior since we specified groups=["posterior"]
+        assert "fit_data" not in loaded_idata.groups()
+
+    finally:
+        # Clean up
+        if os.path.exists(temp_path):
+            os.unlink(temp_path)
+
+
+def test_save_kwargs_backward_compatibility(fitted_model_instance):
+    """Test that save function still works without kwargs (backward compatibility)"""
+    temp = tempfile.NamedTemporaryFile(mode="w", encoding="utf-8", delete=False)
+    temp_path = temp.name
+    temp.close()
+
+    try:
+        # Test without any kwargs (original behavior)
+        fitted_model_instance.save(temp_path)
+
+        # Verify file was created and can be loaded
+        assert os.path.exists(temp_path)
+        loaded_model = ModelBuilderTest.load(temp_path)
+        assert loaded_model.idata is not None
+        assert "posterior" in loaded_model.idata.groups()
+
+    finally:
+        # Clean up
+        if os.path.exists(temp_path):
+            os.unlink(temp_path)
+
+
 def test_empty_sampler_config_fit(toy_X, toy_y, mock_pymc_sample):
     sampler_config = {}
     model_builder = ModelBuilderTest(sampler_config=sampler_config)