unit8co / darts

A python library for user-friendly forecasting and anomaly detection on time series.
https://unit8co.github.io/darts/
Apache License 2.0
8.03k stars 873 forks source link

Make ModelMode, SeasonalityMode, and TrendMode params JSON serializable #1398

Open john-ramsey opened 1 year ago

john-ramsey commented 1 year ago

Currently ModelMode, SeasonalityMode, and TrendMode are enums defined in utils which cannot be serialized in a json format without a workaround.

Currently the following block raises a TypeError (as expected)

model = Theta(2, 12, SeasonalityMode.ADDITIVE)

json.dumps(dict(model.model_params))

TypeError: Object of type SeasonalityMode is not JSON serializable

I currently use the following functon as a workaround:

def unjsonify_darts_modes(json_params):
    from darts.utils.utils import SeasonalityMode, ModelMode, TrendMode
    for param, value in json_params.items():
        if type(value) == str:
            if "ModelMode" in value:
                json_params[param] = ModelMode[value.split(".")[-1]]
            elif "TrendMode" in value:
                json_params[param] = TrendMode[value.split(".")[-1]]
            elif "SeasonalityMode" in value:
                json_params[param] = SeasonalityMode[value.split(".")[-1]]
            else:
                pass

    return json_params

Example

model = Theta(2, 12, SeasonalityMode.ADDITIVE)

model_json_params = json.dumps(dict(model.model_params), default=str)

model_dictionary = unjsonify_darts_modes(json.loads(model_json_params))

print(model.model_params == Theta(**model_dictionary).model_params)

Returns True

Proposed Solution

Modify the model objects to accept string or integer arguments in addition to Enum objects. This could be done by adding something along the lines of

if type(self.season_mode) == str:
     if self.season_mode.upper() in ("ADDITIVE", "MULTIPLICATIVE", "NONE"):
          self.season_mode = SeasonalityMode[self.season_mode.upper()]

to the individual model initialization methods. Arguably "modes" could be standardized at a lower level and this could be implemented on something like the ForecastingModel class. This would certainly make reading and historical models to logs/databases easier.

hrzn commented 1 year ago

Hi @john-ramsey and thanks for raising this. Your solution certainly looks like a good way to tackle the problem. Would you be willing to open a PR?