Integrate multidimensional MMM class with MLflow autologging

[Copilot]

Integrate new MMM class with MLflow

Plan:

• Explore repository structure and understand the issue
• Identify that there are two MMM classes: old (mmm.py) and new (multidimensional.py)
• Understand current mlflow.py only patches the old MMM class
• Import the new multidimensional MMM class in mlflow.py
• Apply the same patch_mmm_fit to the new MMM class
• Add or update tests for the new MMM integration
• Verify changes follow minimal modification principles
• Ensure backward compatibility
• Fix pre-commit check issues

Summary

Successfully integrated the new multidimensional MMM class with MLflow autologging using the existing patch_mmm_fit function. The solution is minimal and follows best practices.

Changes Made

Production Code (pymc_marketing/mlflow.py):

• Added import of MultiDimensionalMMM from pymc_marketing.mmm.multidimensional
• Applied patch_mmm_fit to MultiDimensionalMMM.fit when log_mmm=True
• Total: 2 lines added

Tests (tests/test_mlflow.py):

• Added import and teardown logic for MultiDimensionalMMM
• Added test fixtures: multidimensional_mmm(), toy_multidim_X(), toy_multidim_y()
• Added test_autolog_multidimensional_mmm() to verify integration
• Fixed unused variable linting issue
• Total: 81 lines added

Pre-commit Checks

✅ ruff check: All checks passed
✅ ruff format: Code properly formatted
✅ mypy: No type errors
✅ trailing whitespace: None found
✅ end-of-file: Proper newlines

Both the old MMM class and the new multidimensional MMM class now have identical MLflow integration.

Original prompt

---

This section details on the original issue you should resolve

<issue_title>Integrate new MMM class with MLflow</issue_title>
<issue_description>This would be integration like with the other MMM class which has a patched fit method</issue_description>

Comments on the Issue (you are @copilot in this section)

@kskyten I'm unclear on what the specific missing functionality is. I ran the MMM auto logging example from the mlflow module, and it does produce outputs to mlflow. @williambdean > I'm unclear on what the specific missing functionality is. I ran the MMM auto logging example from the mlflow module, and it does produce outputs to mlflow.

Yes, you're right. There is a base functionality that comes from the logging for all PyMC models.

I was think about some specific to the new MMM class just like with the old one:

pymc-marketing/pymc_marketing/mlflow.py

Lines 1249 to 1250 in dcbfa60

	if log_mmm:
	MMM.fit = patch_mmm_fit(MMM.fit)

pymc-marketing/pymc_marketing/mlflow.py

Lines 1236 to 1247 in dcbfa60

	def patch_mmm_fit(fit: Callable) -> Callable:
	@wraps(fit)
	def new_fit(self, *args, **kwargs):
	log_mmm_configuration(self)
	idata = fit(self, *args, **kwargs)
	log_inference_data(idata, save_file="idata.nc")
	return idata
	return new_fit

Do you have any thoughts on what would be helpful? @kskyten </comment_new>
<comment_new>@williambdean
Lets start with using current patch_fit for the fit method</comment_new>

• Fixes Integrate new MMM class with MLflow #1609

---

💬 We'd love your input! Share your thoughts on Copilot coding agent in our 2 minute survey.

---

📚 Documentation preview 📚: https://pymc-marketing--2072.org.readthedocs.build/en/2072/

[williambdean]

@copilot please ensure that the pre-commit checks pass ...

[Copilot]

@copilot please ensure that the pre-commit checks pass ...

Fixed the pre-commit check issue. The unused variable idata in the test has been removed. All checks now pass:

• ✅ ruff check
• ✅ ruff format
• ✅ mypy
• ✅ trailing whitespace
• ✅ end-of-file

Commit: 979a9b6

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 92.57%. Comparing base (f59d7e3) to head (0a8422f).
⚠️ Report is 1 commits behind head on main.

Additional details and impacted files

@@           Coverage Diff           @@
##             main    #2072   +/-   ##
=======================================
  Coverage   92.56%   92.57%           
=======================================
  Files          68       68           
  Lines        9420     9422    +2     
=======================================
+ Hits         8720     8722    +2     
  Misses        700      700           

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[cetagostini]

La leyenda viviente!!!!