More General Utilities With Representation Of Time In idata Objects

[AFg6K7h4fhy2]

Generalize utilities to add dates and times to az.InferenceData objects. A demonstration of these utilities can be found in #40 .

[AFg6K7h4fhy2]

#28 can be covered in this PR given the PR's scope.

[AFg6K7h4fhy2]

@damonbayer Can you take a quick look at the changes?

EDIT: Adding rendered file now under notebooks. → See the pdf file.

[AFg6K7h4fhy2]

@dylanhmorris Current situation is that I would like you first to look at this file, then decide which option is closest to what you believe is best, then suggest changes to that option (code-changes or broader changes). I've selected Option 3 (pass date, group, and dim_name as a tuple). Finally, this file begins to set up the tests for Option 3 (but this could be applied to the other options as well). I would appreciate some suggestions on this file. A rendered version of the test file as a pdf with additional output can be found here.

[AFg6K7h4fhy2]

In communications between DHM and the author, a decision has been made to cease using pl.date_range and pl.datetime_range and instead use np.arange with utilities for datetime64[ns] and datetime64[D] objects. As such, some code and tests need to be updated.

[AFg6K7h4fhy2]

The author's remark:

Even if one uses the approach involving np.arange w/ timedelta64 "D", assign_coords converts to [ns] precision, keeping the "times" even if one is operating in "dates". Bite the bullet here an just always accept "times" for coords? I want to canvas assign_coords once more.

from datetime import date, datetime, timedelta

import arviz as az
import numpy as np
import polars as pl
import xarray as xr


coords = {"obs_dim_0": [0, 1, 2, 3, 5]}
data = xr.Dataset(
    {"obs": ("obs_dim_0", np.array([10, 20, 30, 40, 50]))}, coords=coords
)
idata = az.InferenceData(observed_data=data)

# test parameters
group = "observed_data"
variable = "obs"
dimension = "obs_dim_0"
start_date = date(2022, 8, 1)

expected_datetimes = np.array([
    datetime(2022, 8, 1),
    datetime(2022, 8, 2),
    datetime(2022, 8, 3),
    datetime(2022, 8, 4),
    datetime(2022, 8, 5),
]).astype("datetime64[D]")

idata_group = getattr(idata, group)
variable_data = idata_group[variable]
interval_size = variable_data.sizes[dimension]

out = np.arange(
    start_date, interval_size, np.timedelta64(1, "D")
)

print(out)

idata_group = idata_group.assign_coords({dimension: out})

setattr(idata, group, idata_group)

print(idata.observed_data.coords[dimension].values)

Out:

['2022-08-01' '2022-08-02' '2022-08-03' '2022-08-04' '2022-08-05']
['2022-08-01T00:00:00.000000000' '2022-08-02T00:00:00.000000000'
 '2022-08-03T00:00:00.000000000' '2022-08-04T00:00:00.000000000'
 '2022-08-05T00:00:00.000000000']
<ipython-input-2-6ff7ebed4ad4>:35: UserWarning: Converting non-nanosecond precision datetime values to nanosecond precision. This behavior can eventually be relaxed in xarray, as it is an artifact from pandas which is now beginning to support non-nanosecond precision values. This warning is caused by passing non-nanosecond np.datetime64 or np.timedelta64 values to the DataArray or Variable constructor; it can be silenced by converting the values to nanosecond precision ahead of time.
  idata_group = idata_group.assign_coords({dimension: out})

Warning:

UserWarning: Converting non-nanosecond precision datetime values to nanosecond precision. This behavior can eventually be relaxed in xarray, as it is an artifact from pandas which is now beginning to support non-nanosecond precision values. This warning is caused by passing non-nanosecond np.datetime64 or np.timedelta64 values to the DataArray or Variable constructor; it can be silenced by converting the values to nanosecond precision ahead of time

[dylanhmorris]

Yes, if xarray currently only supports [ns] precision, we should do the same in this function (and relax it if and when xarray changes)

[AFg6K7h4fhy2]

Re: #35 (comment)

With this, the author will still operate with using two paths, one for datetime64[D] and another for datetime64[ns]. However, it will be noted in a docstring / type hint that the output (i.e. the assigned coords) is datetime64[ns]. Rather than remove what will operate as intended as some point in the future (i.e. when xarray behavior is relaxed w/ non [ns] precision), the author will retain the separation, which does not seem to harm usability.

[dylanhmorris]

Sounds good, but to avoid the warning, you should coerce the [D] resolution arrays to [ns] resolution before calling assign_coords()

[AFg6K7h4fhy2]

Re: #35 (comment)

Sounds good; I've been proceeding with properly placed .astype("datetime64[ns]")s.

[dylanhmorris]

Looking very good, thanks @AFg6K7h4fhy2. A few final things.

[dylanhmorris]

Thanks for your hard work, @AFg6K7h4fhy2!