def train_test_split_time(data, first_test_day):
  """Split data into training_data (before first_test_day) and test_data (after).

  Args:
    data: an xr.Dataset containing the ground-truth infections.
    first_test_day: a time.coord at which to split the data into train and test
      sets.

  Returns:
    train_data: an xr.Dataset containing the training data for a model.
    test_data: an xr.Dataset containing the testing data for a model.
  """
  # everything before first_test_day
  train_data = data.sel(time=slice(None, first_test_day - 1)).copy()
  # everything after first_test_day
  test_data = data.sel(time=slice(first_test_day, None)).copy()
  # drop all variables that won't exist on real test data
  things_included = ['location', 'time']
  if hasattr(test_data, 'dynamic_covariates'):
    things_included.append('dynamic_covariates')
  if 'original_time' in test_data.data_vars.keys():
    things_included.append('original_time')
  test_data = test_data[things_included]

  train_data.attrs['split'] = 'train'
  train_data.attrs['split_spec'] = 'train_test_split_time'
  test_data.attrs['split'] = 'test'
  test_data.attrs['split_spec'] = 'train_test_split_time'

  # datetime64 is not serializable as an attribute, causing to_netcdf to crash
  # if we include it in attrs. However, making it a variable works :-/.
  train_data = train_data.assign(first_test_day=first_test_day)
  test_data = test_data.assign(first_test_day=first_test_day)

  return train_data, test_data