diff --git a/datasets/era5/test_time.py b/datasets/era5/test_time.py index 8f36925..27f3a80 100644 --- a/datasets/era5/test_time.py +++ b/datasets/era5/test_time.py @@ -14,12 +14,13 @@ # limitations under the License. import datetime +import pytz from datasets.era5 import time def test_datetime_range(): - times = time.datetime_range(2018, datetime.timedelta(hours=6), 2) - assert times == [datetime.datetime(2018, 1, 1, 0), datetime.datetime(2018, 1, 1, 6)] + times = time.datetime_range(2018, datetime.timedelta(hours=6), 2, tzinfo=pytz.utc) + assert times == [datetime.datetime(2018, 1, 1, 0, tzinfo=pytz.utc), datetime.datetime(2018, 1, 1, 6, tzinfo=pytz.utc)] def test_filename_to_year(): diff --git a/datasets/era5/time.py b/datasets/era5/time.py index 360191e..8d8d120 100644 --- a/datasets/era5/time.py +++ b/datasets/era5/time.py @@ -16,7 +16,7 @@ from typing import List import os import datetime - +import pytz def filename_to_year(path: str) -> int: filename = os.path.basename(path) @@ -26,5 +26,5 @@ def filename_to_year(path: str) -> int: def datetime_range( year: int, time_step: datetime.timedelta, n: int ) -> List[datetime.datetime]: - initial_time = datetime.datetime(year=year, month=1, day=1) + initial_time = datetime.datetime(year=year, month=1, day=1, tzinfo=pytz.utc) return [initial_time + time_step * i for i in range(n)] diff --git a/makani/third_party/climt/zenith_angle.py b/makani/third_party/climt/zenith_angle.py index 6fc67fe..be9e9f4 100644 --- a/makani/third_party/climt/zenith_angle.py +++ b/makani/third_party/climt/zenith_angle.py @@ -30,6 +30,7 @@ # modified 2024: vectorization over coordinates and JIT compilation added import datetime +import pytz import numpy as np from typing import Union, Tuple, TypeVar @@ -45,7 +46,7 @@ def _days_from_2000(model_time: np.ndarray) -> np.ndarray: """Get the days since year 2000.""" # compute total days - time_diff = model_time - datetime.datetime(2000, 1, 1, 12, 0) + time_diff = model_time - datetime.datetime(2000, 1, 1, 12, 0, tzinfo=pytz.utc) result = np.asarray(time_diff).astype("timedelta64[us]") / np.timedelta64(1, "D") result = result.astype(dtype) @@ -233,7 +234,9 @@ def cos_zenith_angle( lon_grid, lat_grid = np.meshgrid(lon, lat) # model time - model_time = np.asarray([datetime.datetime(2002, 1, 1, 12, 0, 0), datetime.datetime(2002, 6, 1, 12, 0, 0), datetime.datetime(2003, 1, 1, 12, 0, 0)]) + model_time = np.asarray([datetime.datetime(2002, 1, 1, 12, 0, 0, tzinfo=pytz.utc), + datetime.datetime(2002, 6, 1, 12, 0, 0, tzinfo=pytz.utc), + datetime.datetime(2003, 1, 1, 12, 0, 0, tzinfo=pytz.utc)]) # test _days_from_2000 days = _days_from_2000(model_time) diff --git a/makani/utils/dataloaders/dali_es_helper_2d.py b/makani/utils/dataloaders/dali_es_helper_2d.py index 092dc7a..e42de03 100644 --- a/makani/utils/dataloaders/dali_es_helper_2d.py +++ b/makani/utils/dataloaders/dali_es_helper_2d.py @@ -32,6 +32,7 @@ # we need this for the zenith angle feature import datetime +import pytz # import splitting logic from modulus.distributed.utils import compute_split_shapes @@ -359,7 +360,7 @@ def _compute_zenith_angle(self, zen_inp, zen_tar, local_idx, year_idx): # compute hours into the year year = self.years[year_idx] - jan_01_epoch = datetime.datetime(year, 1, 1, 0, 0, 0) + jan_01_epoch = datetime.datetime(year, 1, 1, 0, 0, 0, tzinfo=pytz.utc) # zenith angle for input inp_times = np.asarray([jan_01_epoch + datetime.timedelta(hours=idx * self.dhours) for idx in range(local_idx - self.dt * self.n_history, local_idx + 1, self.dt)]) diff --git a/makani/utils/dataloaders/data_loader_multifiles.py b/makani/utils/dataloaders/data_loader_multifiles.py index fa89298..b6ff09e 100644 --- a/makani/utils/dataloaders/data_loader_multifiles.py +++ b/makani/utils/dataloaders/data_loader_multifiles.py @@ -27,6 +27,7 @@ # for the zenith angle import datetime +import pytz # for grid conversion from makani.utils.grids import GridConverter @@ -204,7 +205,7 @@ def _compute_zenith_angle(self, local_idx, year_idx): # compute hours into the year year = self.years[year_idx] - jan_01_epoch = datetime.datetime(year, 1, 1, 0, 0, 0) + jan_01_epoch = datetime.datetime(year, 1, 1, 0, 0, 0, tzinfo=pytz.utc) # zenith angle for input inp_times = np.asarray([jan_01_epoch + datetime.timedelta(hours=idx * self.dhours) for idx in range(local_idx - self.dt * self.n_history, local_idx + 1, self.dt)])