Issue with Saving Model in PyMC Marketing (TimedeltaIndex Error)

[stdp95]

Hi everyone,

I’m encountering an issue when trying to save a PyMC Marketing Model in Databricks. Here’s the context:
I successfully ran the model, and when I attempt to save it using the .save() method, I get the following error:
AttributeError: 'TimedeltaIndex' object has no attribute 'unit'File <command-2594243823868214>, line 1 ----> 1 mmm.save("model.nc")
The stack trace indicates an issue in xarray while attempting to use to_netcdf. Specifically, it seems related to the handling of a TimedeltaIndex in time_deltas.unit.

Additional Details:
Environment: Running the code on Databricks.
PyMC-Marketing Version: pymc_marketing-0.10.0.

Has anyone encountered this error or knows how to resolve it?

Thank you for any insights!

[wd60622]

Hi. Could you share the full traceback? The save method is a light wrapper around arviz save function, is it possible to save any other way?

[stdp95]

Here is the full traceback

AttributeError: 'TimedeltaIndex' object has no attribute 'unit' File <command-2594243823868231>, line 1 ----> 1 mmm.save("mymodel.nc") File /local_disk0/.ephemeral_nfs/envs/pythonEnv-97ed79d0-d8a1-4947-a0db-e89c3646d3bb/lib/python3.10/site-packages/pymc_marketing/model_builder.py:445, in ModelBuilder.save(self, fname) 443 if self.idata is not None and "posterior" in self.idata: 444 file = Path(str(fname)) --> 445 self.idata.to_netcdf(str(file)) 446 else: 447 raise RuntimeError("The model hasn't been fit yet, call .fit() first") File /local_disk0/.ephemeral_nfs/envs/pythonEnv-97ed79d0-d8a1-4947-a0db-e89c3646d3bb/lib/python3.10/site-packages/arviz/data/inference_data.py:514, in InferenceData.to_netcdf(self, filename, compress, groups, engine, base_group, overwrite_existing) 508 if compress: 509 kwargs["encoding"] = { 510 var_name: {"zlib": True} 511 for var_name, values in data.variables.items() 512 if _compressible_dtype(values.dtype) 513 } --> 514 data.to_netcdf(filename, mode=mode, group=f"{base_group}/{group}", **kwargs) 515 data.close() 516 mode = "a" File /local_disk0/.ephemeral_nfs/envs/pythonEnv-97ed79d0-d8a1-4947-a0db-e89c3646d3bb/lib/python3.10/site-packages/xarray/core/dataset.py:2373, in Dataset.to_netcdf(self, path, mode, format, group, engine, encoding, unlimited_dims, compute, invalid_netcdf, auto_complex) 2370 encoding = {} 2371 from xarray.backends.api import to_netcdf -> 2373 return to_netcdf( # type: ignore[return-value] # mypy cannot resolve the overloads:( 2374 self, 2375 path, 2376 mode=mode, 2377 format=format, 2378 group=group, 2379 engine=engine, 2380 encoding=encoding, 2381 unlimited_dims=unlimited_dims, 2382 compute=compute, 2383 multifile=False, 2384 invalid_netcdf=invalid_netcdf, 2385 auto_complex=auto_complex, 2386 ) File /local_disk0/.ephemeral_nfs/envs/pythonEnv-97ed79d0-d8a1-4947-a0db-e89c3646d3bb/lib/python3.10/site-packages/xarray/backends/api.py:1905, in to_netcdf(dataset, path_or_file, mode, format, group, engine, encoding, unlimited_dims, compute, multifile, invalid_netcdf, auto_complex) 1900 # TODO: figure out how to refactor this logic (here and in save_mfdataset) 1901 # to avoid this mess of conditionals 1902 try: 1903 # TODO: allow this work (setting up the file for writing array data) 1904 # to be parallelized with dask -> 1905 dump_to_store( 1906 dataset, store, writer, encoding=encoding, unlimited_dims=unlimited_dims 1907 ) 1908 if autoclose: 1909 store.close() File /local_disk0/.ephemeral_nfs/envs/pythonEnv-97ed79d0-d8a1-4947-a0db-e89c3646d3bb/lib/python3.10/site-packages/xarray/backends/api.py:1952, in dump_to_store(dataset, store, writer, encoder, encoding, unlimited_dims) 1949 if encoder: 1950 variables, attrs = encoder(variables, attrs) -> 1952 store.store(variables, attrs, check_encoding, writer, unlimited_dims=unlimited_dims) File /local_disk0/.ephemeral_nfs/envs/pythonEnv-97ed79d0-d8a1-4947-a0db-e89c3646d3bb/lib/python3.10/site-packages/xarray/backends/common.py:454, in AbstractWritableDataStore.store(self, variables, attributes, check_encoding_set, writer, unlimited_dims) 451 if writer is None: 452 writer = ArrayWriter() --> 454 variables, attributes = self.encode(variables, attributes) 456 self.set_attributes(attributes) 457 self.set_dimensions(variables, unlimited_dims=unlimited_dims) File /local_disk0/.ephemeral_nfs/envs/pythonEnv-97ed79d0-d8a1-4947-a0db-e89c3646d3bb/lib/python3.10/site-packages/xarray/backends/common.py:638, in WritableCFDataStore.encode(self, variables, attributes) 635 def encode(self, variables, attributes): 636 # All NetCDF files get CF encoded by default, without this attempting 637 # to write times, for example, would fail. --> 638 variables, attributes = cf_encoder(variables, attributes) 639 variables = { 640 k: ensure_dtype_not_object(v, name=k) for k, v in variables.items() 641 } 642 variables = {k: self.encode_variable(v) for k, v in variables.items()} File /local_disk0/.ephemeral_nfs/envs/pythonEnv-97ed79d0-d8a1-4947-a0db-e89c3646d3bb/lib/python3.10/site-packages/xarray/conventions.py:759, in cf_encoder(variables, attributes) 756 # add encoding for time bounds variables if present. 757 _update_bounds_encoding(variables) --> 759 new_vars = {k: encode_cf_variable(v, name=k) for k, v in variables.items()} 761 # Remove attrs from bounds variables (issue #2921) 762 for var in new_vars.values(): File /local_disk0/.ephemeral_nfs/envs/pythonEnv-97ed79d0-d8a1-4947-a0db-e89c3646d3bb/lib/python3.10/site-packages/xarray/conventions.py:759, in <dictcomp>(.0) 756 # add encoding for time bounds variables if present. 757 _update_bounds_encoding(variables) --> 759 new_vars = {k: encode_cf_variable(v, name=k) for k, v in variables.items()} 761 # Remove attrs from bounds variables (issue #2921) 762 for var in new_vars.values(): File /local_disk0/.ephemeral_nfs/envs/pythonEnv-97ed79d0-d8a1-4947-a0db-e89c3646d3bb/lib/python3.10/site-packages/xarray/conventions.py:101, in encode_cf_variable(var, needs_copy, name) 89 ensure_not_multiindex(var, name=name) 91 for coder in [ 92 CFDatetimeCoder(), 93 times.CFTimedeltaCoder(), (...) 99 variables.BooleanCoder(), 100 ]: --> 101 var = coder.encode(var, name=name) 103 for attr_name in CF_RELATED_DATA: 104 pop_to(var.encoding, var.attrs, attr_name) File /local_disk0/.ephemeral_nfs/envs/pythonEnv-97ed79d0-d8a1-4947-a0db-e89c3646d3bb/lib/python3.10/site-packages/xarray/coding/times.py:1117, in CFDatetimeCoder.encode(self, variable, name) 1115 calendar = encoding.pop("calendar", None) 1116 dtype = encoding.get("dtype", None) -> 1117 (data, units, calendar) = encode_cf_datetime(data, units, calendar, dtype) 1119 safe_setitem(attrs, "units", units, name=name) 1120 safe_setitem(attrs, "calendar", calendar, name=name) File /local_disk0/.ephemeral_nfs/envs/pythonEnv-97ed79d0-d8a1-4947-a0db-e89c3646d3bb/lib/python3.10/site-packages/xarray/coding/times.py:858, in encode_cf_datetime(dates, units, calendar, dtype) 856 return _lazily_encode_cf_datetime(dates, units, calendar, dtype) 857 else: --> 858 return _eagerly_encode_cf_datetime(dates, units, calendar, dtype) File /local_disk0/.ephemeral_nfs/envs/pythonEnv-97ed79d0-d8a1-4947-a0db-e89c3646d3bb/lib/python3.10/site-packages/xarray/coding/times.py:933, in _eagerly_encode_cf_datetime(dates, units, calendar, dtype, allow_units_modification) 928 time_delta = needed_time_delta 930 # get resolution of TimedeltaIndex and align time_delta 931 # todo: check, if this works in any case 932 num = _division( --> 933 time_deltas, time_delta.astype(f"=m8[{time_deltas.unit}]"), floor_division 934 ) 935 num = reshape(num.values, dates.shape) 937 except (OutOfBoundsDatetime, OverflowError, ValueError):

[stdp95]

I’ve tried using pickle but encountered issues as well—it didn’t seem to fit the scenario. If there are alternative methods or formats you’d recommend for saving the model or its data, I’d be happy to try them.