77# granted to it by virtue of its status as an intergovernmental organisation
88# nor does it submit to any jurisdiction.
99
10- from typing import Any
10+ import logging
11+ from dataclasses import dataclass
12+ from typing import Any , ClassVar
1113
14+ import xarray as xr
1215from omegaconf .listconfig import ListConfig
1316
17+ _logger = logging .getLogger (__name__ )
18+ _logger .setLevel (logging .INFO )
19+
1420
1521def to_list (obj : Any ) -> list :
1622 """
@@ -30,3 +36,100 @@ def to_list(obj: Any) -> list:
3036 elif not isinstance (obj , list ):
3137 obj = [obj ]
3238 return obj
39+
40+
41+ class RegionLibrary :
42+ """
43+ Predefined bounding boxes for known regions.
44+ """
45+
46+ REGIONS : ClassVar [dict [str , tuple [float , float , float , float ]]] = {
47+ "global" : (- 90.0 , 90.0 , - 180.0 , 180.0 ),
48+ "nhem" : (0.0 , 90.0 , - 180.0 , 180.0 ),
49+ "shem" : (- 90.0 , 0.0 , - 180.0 , 180.0 ),
50+ "tropics" : (- 30.0 , 30.0 , - 180.0 , 180.0 ),
51+ }
52+
53+
54+ @dataclass (frozen = True )
55+ class RegionBoundingBox :
56+ lat_min : float
57+ lat_max : float
58+ lon_min : float
59+ lon_max : float
60+
61+ def __post_init__ (self ):
62+ """Validate the bounding box coordinates."""
63+ self .validate ()
64+
65+ def validate (self ):
66+ """Validate the bounding box coordinates."""
67+ if not (- 90 <= self .lat_min <= 90 and - 90 <= self .lat_max <= 90 ):
68+ raise ValueError (
69+ f"Latitude bounds must be between -90 and 90. Got: { self .lat_min } , { self .lat_max } "
70+ )
71+ if not (- 180 <= self .lon_min <= 180 and - 180 <= self .lon_max <= 180 ):
72+ raise ValueError (
73+ f"Longitude bounds must be between -180 and 180. Got: { self .lon_min } , { self .lon_max } "
74+ )
75+ if self .lat_min >= self .lat_max :
76+ raise ValueError (
77+ f"Latitude minimum must be less than maximum. Got: { self .lat_min } , { self .lat_max } "
78+ )
79+ if self .lon_min >= self .lon_max :
80+ raise ValueError (
81+ f"Longitude minimum must be less than maximum. Got: { self .lon_min } , { self .lon_max } "
82+ )
83+
84+ def contains (self , lat : float , lon : float ) -> bool :
85+ """Check if a lat/lon point is within the bounding box."""
86+ return (self .lat_min <= lat <= self .lat_max ) and (
87+ self .lon_min <= lon <= self .lon_max
88+ )
89+
90+ def apply_mask (
91+ self ,
92+ data : xr .Dataset | xr .DataArray ,
93+ lat_name : str = "lat" ,
94+ lon_name : str = "lon" ,
95+ data_dim : str = "ipoint" ,
96+ ) -> xr .Dataset | xr .DataArray :
97+ """Filter Dataset or DataArray by spatial bounding box on 'ipoint' dimension.
98+ Parameters
99+ ----------
100+ data :
101+ The data to filter.
102+ lat_name:
103+ Name of the latitude coordinate in the data.
104+ lon_name:
105+ Name of the longitude coordinate in the data.
106+ data_dim:
107+ Name of the dimension that contains the lat/lon coordinates.
108+
109+ Returns
110+ -------
111+ Filtered data with only points within the bounding box.
112+ """
113+ # lat/lon coordinates should be 1D and aligned with ipoint
114+ lat = data [lat_name ]
115+ lon = data [lon_name ]
116+
117+ mask = (
118+ (lat >= self .lat_min )
119+ & (lat <= self .lat_max )
120+ & (lon >= self .lon_min )
121+ & (lon <= self .lon_max )
122+ )
123+
124+ return data .sel ({data_dim : mask })
125+
126+ @classmethod
127+ def from_region_name (cls , region : str ) -> "RegionBoundingBox" :
128+ region = region .lower ()
129+ try :
130+ return cls (* RegionLibrary .REGIONS [region ])
131+ except KeyError as err :
132+ raise ValueError (
133+ f"Region '{ region } ' is not supported. "
134+ f"Available regions: { ', ' .join (RegionLibrary .REGIONS .keys ())} "
135+ ) from err
0 commit comments