fix: serialization for arbitrary HSGP instance to time_varying_intercept/media in MMM

[jamespooley]

Description

This is a draft PR being opened to let you know the bug is being worked on.

The code changes here get rid of post-fitting serialization errors like TypeError: Object of type SoftPlusHSGP is not JSON serializable when passing HSGP instances to the time_varying_intercept or time_varying_media parameters of the multidimensional.MMM class.

There’s not an open issue for this yet, but there’s also a related deserialization issue when passing arbitrary HSGP instances. The tl;dr is that after fitting and .save()-ing a model with them, there’s an error when .load()-ing the model.

One (admittedly not great) hotfix to get rid of that error is to replace

time_varying_intercept: Annotated[
            StrictBool | InstanceOf[HSGPBase],

in this block of code (and similarly with time_verying_media) with

time_varying_intercept: Annotated[
            StrictBool | InstanceOf[HSGPBase] | dict,

But even though this gets rid of the .load()-ing error, the time-varying structure for the saved model isn’t recovered after the .load(), and the time-varying components have the default configuration.

So I’m wondering if the logic there needs to be updated to use HSGPKwargs logic like that here or maybe something new like hsgp_from_dict() like how saturation_from_dict() is used?

If so, that’s likely another PR, but I wanted to raise it here since they’re both related to the “pass arbitrary HSGP instances” functionality.

Related Issue

• Closes Serialization errors when passing HSGP instances to time_varying_media #1935
• Related to Support HSGP instance for time_varying_media in MMM #1881 and Support HSGP instance for time_varying_intercept in MMM #1894

Checklist

• Checked that the pre-commit linting/style checks pass. Feel free to comment pre-commit.ci autofix to auto-fix.
• Included tests that prove the fix is effective or that the new feature works
• Added necessary documentation (docstrings and/or example notebooks) using numpydoc format.
• If you are a pro: each commit corresponds to a relevant logical change

---

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

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 92.51%. Comparing base (e9b1b45) to head (fe4ea45).
⚠️ Report is 3 commits behind head on main.

Additional details and impacted files

@@             Coverage Diff             @@
##             main    #2016       +/-   ##
===========================================
+ Coverage   36.79%   92.51%   +55.72%     
===========================================
  Files          68       68               
  Lines        9373     9396       +23     
===========================================
+ Hits         3449     8693     +5244     
+ Misses       5924      703     -5221     

☔ 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.

[williambdean]

Thanks for the PR @jamespooley

Can you take a look to see if the model will load (adding a test would be appreciated)

The loading logic is here:

pymc-marketing/pymc_marketing/mmm/multidimensional.py

Lines 589 to 591 in 9c429ed

	"time_varying_intercept": json.loads(
	attrs.get("time_varying_intercept", "false")
	),

[jamespooley]

@williambdean Drafted some tests to show that the loading logic itself...

pymc-marketing/pymc_marketing/mmm/multidimensional.py

Lines 589 to 591 in 9c429ed

	"time_varying_intercept": json.loads(
	attrs.get("time_varying_intercept", "false")
	),

pymc-marketing/pymc_marketing/mmm/multidimensional.py

Line 593 in 9c429ed

"time_varying_media": json.loads(attrs.get("time_varying_media", "false")),

...works fine (i.e., returns the same information in the user-specific HSGP instance) with the updated code.

But I also added a temporary test to clarify the issue that when calling MMM.load() on a successfully saved MMM with a HSGP instance passed to either of the time_varying_* parameters raises a ValidationError since that loading logic will return a dict, which currently isn't then converted into an instance of HSGPBase.

As mentioned, that error will go away by adding a | dict to the relevant portions of the following:

pymc-marketing/pymc_marketing/mmm/multidimensional.py

Lines 278 to 288 in 9c429ed

	time_varying_intercept: Annotated[
	StrictBool | InstanceOf[HSGPBase],
	Field(
	description=(
	"Whether to use a time-varying intercept, or pass an HSGP instance "
	"(e.g., SoftPlusHSGP) specifying dims and priors."
	),
	),
	] = False,
	time_varying_media: Annotated[
	StrictBool | InstanceOf[HSGPBase],

But it seems like the proper solution would be something along these lines, just for the time_varying_* parameters that can accept HSGPs:

pymc-marketing/pymc_marketing/mmm/multidimensional.py

Lines 585 to 586 in 9c429ed

	"adstock": adstock_from_dict(json.loads(attrs["adstock"])),
	"saturation": saturation_from_dict(json.loads(attrs["saturation"])),

Any thoughts here?

[williambdean]

But it seems like the proper solution would be something along these lines, just for the time_varying_* parameters that can accept HSGPs:

pymc-marketing/pymc_marketing/mmm/multidimensional.py

Lines 585 to 586 in 9c429ed

	"adstock": adstock_from_dict(json.loads(attrs["adstock"])),
	"saturation": saturation_from_dict(json.loads(attrs["saturation"])),

Any thoughts here?

Good point. Let's just start with the most obvious cases at the moment. We can make a function like that hsgp_from_dict which will just wrap these cases at the moment. Let's continue with further implementation in additional PRs

[jamespooley]

@williambdean Sounds good to me.

The following wouldn’t be pretty and I wouldn’t be super proud of it, but because there isn’t currently an equivalent of {"lookup_name": "logistic", ...} for HSGPs after .to_dict()-ing like there is for SaturationTransformation, modifying the HSGP serialization to be something like this…

attrs["time_varying_media"] = json.dumps(
    self.time_varying_media
    if not isinstance(self.time_varying_media, HSGPBase)
    else {
        **self.time_varying_media.to_dict(),
        **{"hsgp_class": self.time_varying_media.__class__.__name__},
    }
)

…and then a dead simple implementation of hsgp_from_dict like the following would likely hotfix the the MMM.load() issue.

from pymc_marketing.mmm.hsgp import HSGP, HSGPPeriodic, SoftPlusHSGP

def hsgp_from_dict(data: dict | bool):
    if isinstance(data, bool):
        return data

    HSGP_CLASSES= {
        "HSGP": HSGP,
        "SoftPlusHSGP": SoftPlusHSGP,
        "HSGPPeriodic": HSGPPeriodic,
    }

    data = data.copy()
    cls = HSGP_CLASSES[data.pop("hsgp_class")]

    return cls.from_dict(data)

It's certainly hacky, but don’t know if you think an approach along these lines is too hacky even for the first pass at this. But it could work until a proper implementation is tackled in future PRs.

If you're aligned on something like ☝️ as a "good enough for now" approach, I can clean up and test the implementation, update the code/tests, and then finalize this PR.

[williambdean]

Looks good. Thanks for these alterations!

[williambdean]

I see this is still a draft PR, @jamespooley. I will merge and we can continue working in a separate PR if needed.

[jamespooley]

Sounds good, @williambdean. Glad I could make a tiny contribution here, and I'd happily pitch in on future PRs related to this.

[williambdean]

This is no tiny PR! Thanks @jamespooley 🚀