Merge pull request #13 from azrael417/tkurth/tz-lock

fixing timezone to UTC

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():

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)]

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)

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)])

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)])