11""" Loading Raw data """
22import logging
33from dataclasses import dataclass
4+ from datetime import datetime
45from numbers import Number
56from pathlib import Path
67from typing import List , Tuple , Union
1415from nowcasting_dataset .data_sources .metadata .metadata_model import SpaceTimeLocation
1516from nowcasting_dataset .data_sources .sun .raw_data_load_save import load_from_zarr , x_y_to_name
1617from nowcasting_dataset .data_sources .sun .sun_model import Sun
17- from nowcasting_dataset .geospatial import calculate_azimuth_and_elevation_angle
18+ from nowcasting_dataset .geospatial import calculate_azimuth_and_elevation_angle , osgb_to_lat_lon
1819
1920logger = logging .getLogger (__name__ )
2021
@@ -24,6 +25,8 @@ class SunDataSource(DataSource):
2425 """Add azimuth and elevation angles of the sun."""
2526
2627 zarr_path : Union [str , Path ]
28+ load_live : bool = False
29+ elevation_limit : int = 10
2730
2831 def __post_init__ (self ):
2932 """Post Init"""
@@ -37,7 +40,8 @@ def get_data_model_for_batch():
3740
3841 def check_input_paths_exist (self ) -> None :
3942 """Check input paths exist. If not, raise a FileNotFoundError."""
40- nd_fs_utils .check_path_exists (self .zarr_path )
43+ if not self .load_live :
44+ nd_fs_utils .check_path_exists (self .zarr_path )
4145
4246 def get_example (self , location : SpaceTimeLocation ) -> xr .Dataset :
4347 """
@@ -64,26 +68,40 @@ def get_example(self, location: SpaceTimeLocation) -> xr.Dataset:
6468 start_dt = self ._get_start_dt (t0_datetime_utc )
6569 end_dt = self ._get_end_dt (t0_datetime_utc )
6670
67- # The names of the columns get truncated when saving, therefore we need to look for the
68- # name of the columns near the location we are looking for
69- locations = np .array (
70- [[float (z .split ("," )[0 ]), float (z .split ("," )[1 ])] for z in self .azimuth .columns ]
71- )
72- location = locations [
73- np .isclose (locations [:, 0 ], x_center_osgb ) & np .isclose (locations [:, 1 ], y_center_osgb )
74- ]
75- # lets make sure there is atleast one
76- assert len (location ) > 0
77- # Take the first location, and x and y coordinates are the first and center entries in
78- # this array.
79- location = location [0 ]
80- # make name of column to pull data from. The columns name will be about
81- # something like '22222.555,3333.6666'
82- name = x_y_to_name (x = location [0 ], y = location [1 ])
83-
84- del x_center_osgb , y_center_osgb
85- azimuth = self .azimuth .loc [start_dt :end_dt ][name ]
86- elevation = self .elevation .loc [start_dt :end_dt ][name ]
71+ if not self .load_live :
72+
73+ # The names of the columns get truncated when saving, therefore we need to look for the
74+ # name of the columns near the location we are looking for
75+ locations = np .array (
76+ [[float (z .split ("," )[0 ]), float (z .split ("," )[1 ])] for z in self .azimuth .columns ]
77+ )
78+ location = locations [
79+ np .isclose (locations [:, 0 ], x_center_osgb )
80+ & np .isclose (locations [:, 1 ], y_center_osgb )
81+ ]
82+ # lets make sure there is atleast one
83+ assert len (location ) > 0
84+ # Take the first location, and x and y coordinates are the first and center entries in
85+ # this array.
86+ location = location [0 ]
87+ # make name of column to pull data from. The columns name will be about
88+ # something like '22222.555,3333.6666'
89+ name = x_y_to_name (x = location [0 ], y = location [1 ])
90+
91+ del x_center_osgb , y_center_osgb
92+ azimuth = self .azimuth .loc [start_dt :end_dt ][name ]
93+ elevation = self .elevation .loc [start_dt :end_dt ][name ]
94+
95+ else :
96+
97+ latitude , longitude = osgb_to_lat_lon (x = x_center_osgb , y = y_center_osgb )
98+
99+ datestamps = pd .date_range (start = start_dt , end = end_dt , freq = "5T" ).tolist ()
100+ azimuth_elevation = calculate_azimuth_and_elevation_angle (
101+ latitude = latitude , longitude = longitude , datestamps = datestamps
102+ )
103+ azimuth = azimuth_elevation ["azimuth" ]
104+ elevation = azimuth_elevation ["elevation" ]
87105
88106 azimuth = azimuth .to_xarray ().rename ({"index" : "time" })
89107 elevation = elevation .to_xarray ().rename ({"index" : "time" })
@@ -97,7 +115,8 @@ def _load(self):
97115
98116 logger .info (f"Loading Sun data from { self .zarr_path } " )
99117
100- self .azimuth , self .elevation = load_from_zarr (zarr_path = self .zarr_path )
118+ if not self .load_live :
119+ self .azimuth , self .elevation = load_from_zarr (zarr_path = self .zarr_path )
101120
102121 def get_locations (
103122 self , t0_datetimes_utc : pd .DatetimeIndex
@@ -112,13 +131,20 @@ def datetime_index(self) -> pd.DatetimeIndex:
112131 latitude = 51
113132 longitude = 0
114133
134+ if not self .load_live :
135+ datestamps = self .elevation .index
136+ else :
137+ datestamps = pd .date_range (
138+ datetime (2019 , 1 , 1 ), datetime (2019 , 12 , 31 , 23 , 55 ), freq = "5T"
139+ )
140+
115141 # get elevation for all datetimes
116142 azimuth_elevation = calculate_azimuth_and_elevation_angle (
117- latitude = latitude , longitude = longitude , datestamps = self . elevation . index
143+ latitude = latitude , longitude = longitude , datestamps = datestamps
118144 )
119145
120146 # only select elevations > 10
121- mask = azimuth_elevation ["elevation" ] >= 10
147+ mask = azimuth_elevation ["elevation" ] >= self . elevation_limit
122148
123149 # create warnings, so we know how many datetimes will be dropped.
124150 # Should be slightly more than half as its night time 50% of the time
@@ -128,7 +154,7 @@ def datetime_index(self) -> pd.DatetimeIndex:
128154 f"out of { len (azimuth_elevation )} as elevation is < 10"
129155 )
130156
131- datetimes = self . elevation [mask ]. index
157+ datetimes = datestamps [mask ]
132158
133159 # Sun data is only for 2019, so to expand on these by
134160 # repeating data from 2014 to 2023
0 commit comments