dggstools.rhpx.rhpxutils

  1"""
  2Helper functions.
  3"""
  4
  5from collections import namedtuple
  6from typing import List, Any
  7
  8import fiona.transform
  9import rasterio.crs
 10import shapely.geometry
 11import shapely.ops
 12from rasterio.transform import Affine
 13
 14from rhealpixdggs.dggs import RHEALPixDGGS, Cell
 15from rhealpixdggs.ellipsoids import WGS84_A, WGS84_F, WGS84_ELLIPSOID
 16from dggstools.rhpx.utils.rasterutils import *
 17
 18# ROBERT GIBB, ALEXANDER RAICHEV, AND MICHAEL SPETH. THE RHEALPIX DISCRETE GLOBAL GRID SYSTEM (2013)
 19RHEALPIX_MEAN_AREAL_DISTORTION = 1.178  # Min, max and median (which actually is 1.177) are equal to the mean, as
 20# rhealpix is an equiareal projection
 21
 22RHEALPixDGGSNamedTuple = namedtuple("RHEALPixDGGSNamedTuple", ["ellipsoid", "n_side", "north_square", "south_square"])
 23
 24
 25def rdggs_to_namedtuple(rdggs: RHEALPixDGGS) -> RHEALPixDGGSNamedTuple:
 26    assert rdggs.ellipsoid.a == WGS84_A and rdggs.ellipsoid.f == WGS84_F, \
 27        "Only the WGS84 Ellipsoid can be used for now"
 28    return RHEALPixDGGSNamedTuple("WGS84", rdggs.N_side, rdggs.north_square,
 29                                  rdggs.south_square)
 30
 31def namedtuple_to_rdggs(rdggs_namedtuple: RHEALPixDGGSNamedTuple) -> RHEALPixDGGS:
 32    assert rdggs_namedtuple.ellipsoid == "WGS84", "Only the WGS84 Ellipsoid can be used for now"
 33    return RHEALPixDGGS(ellipsoid=WGS84_ELLIPSOID, N_side=rdggs_namedtuple.n_side,
 34                        north_square=rdggs_namedtuple.north_square, south_square=rdggs_namedtuple.south_square)
 35
 36
 37def pyproj_crs_to_rdggs(crs: pyproj.CRS, N_side: int) -> RHEALPixDGGS:
 38    # n_side is not part of the rHEALPix projection, but it is part of the rHEALPIX DGGS. Here it must be a parameter
 39    crs_dict = crs.to_dict()
 40    assert crs_dict["proj"] == "rhealpix", "crs must be rHEALPix"
 41    assert crs.ellipsoid.name == "WGS 84" or crs.ellipsoid.name == "WGS84" or crs.ellipsoid.name == "WGS_84", \
 42        "Only the WGS84 Ellipsoid is supported"
 43    assert N_side == 2 or N_side == 3, f"N_side must be 2 or 3 but it is {N_side}"
 44
 45    return RHEALPixDGGS(ellipsoid=WGS84_ELLIPSOID, N_side=N_side, north_square=crs_dict["north_square"],
 46                        south_square=crs_dict["south_square"])
 47
 48
 49def cellidstr_to_suid(cellid: str) -> List:
 50    return list(cellid[0]) + [int(digit) for digit in cellid[1:]]
 51
 52def cellid_resolution_idx(cellid: str) -> int:
 53    return len(cellid) - 1
 54
 55def get_parent_cellid(cellid: str) -> str:
 56    if len(cellid) <= 1:
 57        return ""
 58    else:
 59        return cellid[:-1]
 60
 61def get_ascendant_cellid_at_resolution_idx(cellid: str, res: int) -> str | None:
 62    cellid_res = cellid_resolution_idx(cellid)
 63    if cellid_res  <= res:
 64        return None
 65    else:
 66        if res < cellid_res:
 67            return cellid[:res+1]
 68        else:
 69            return None
 70
 71def get_ascendant_cellids_up_to_resolution_idx(cellid: str, res: int) -> list[str]:
 72    cellid_res = cellid_resolution_idx(cellid)
 73    if cellid_res <= res:
 74        return []
 75    else:
 76        result = []
 77        for i in range(res, cellid_res):
 78            result.extend([cellid[:i+1]])
 79        return result
 80
 81
 82def get_children_cellids(cellid: str, rhpx: RHEALPixDGGS) -> List[str]:
 83    return [cellid + str(i) for i in range(rhpx.N_side**2)]
 84
 85def get_descendant_cellids_at_resolution_idx(cellid: str, rhpx: RHEALPixDGGS, res: int) -> list[str]:
 86    """
 87    Returns the descendant cellids of cellid at exactly resolution res (not up to res, at res).
 88    """
 89    if res <= cellid_resolution_idx(cellid):
 90        return []
 91    elif res == cellid_resolution_idx(cellid) + 1:
 92        return get_children_cellids(cellid, rhpx)
 93    else:
 94        result = []
 95        for cid in get_descendant_cellids_at_resolution_idx(cellid, rhpx, res-1):
 96            result.extend(get_children_cellids(cid, rhpx))
 97        return result
 98
 99def get_descendant_cellids_up_to_resolution_idx(cellid: str, rhpx: RHEALPixDGGS, res: int) -> list[str]:
100    """
101    Returns the descendant cellids of cellid up to resolution res (not at res, up to res). It does
102    not include itself.
103    """
104    result = []
105    for i in range(cellid_resolution_idx(cellid), res+1):
106        result.extend(get_descendant_cellids_at_resolution_idx(cellid, rhpx, i))
107    return result
108
109
110def get_gdf_attrs_from_rhealpix_file(input_file_path: str) -> dict[str, Any]:
111    result = {}
112    profile = get_raster_profile(input_file_path)
113
114    with rasterio.open(input_file_path) as raster:
115        result["left"], result["top"], result["right"], result["bottom"], resx, resy = (
116            get_bbox_from_raster_profile(profile))
117        tags = raster.tags()
118        if "n_side" in tags:
119            n_side = int(tags["n_side"])
120            rdggs = pyproj_crs_to_rdggs(pyproj.CRS(profile["crs"]), n_side)
121            rdggs_helper = RHEALPixDGGSHelper(rdggs)
122            result["rhealpixdggs"] = {"n_side": rdggs.N_side,
123                                      "north_square": rdggs.north_square,
124                                      "south_square": rdggs.south_square,
125                                      "max_areal_resolution": rdggs.max_areal_resolution,
126                                      "max_resolution": rdggs.max_resolution,
127                                      "ellipsoid": rdggs.ellipsoid}
128            resolution_idx_x, _ = rdggs_helper.get_closest_resolution(abs(resx))
129            resolution_idx_y, _ = rdggs_helper.get_closest_resolution(abs(resy))  # resy is often a negative number
130            assert resolution_idx_x == resolution_idx_y, \
131                f"{input_file_path} is not a proper rhealpix file. Its cells are not squares."
132            result["res_idx"] = resolution_idx_x
133        else:
134            result["rhealpixdggs"] = {}
135            result["res_idx"] = -1  # Using -1 to indicate that the resolution is not known
136
137        result["res"] = resx
138        result["height"] = raster.height
139        result["width"] = raster.width
140        result["nbands"] = raster.count
141        result["nodata"] = raster.nodata
142        result["nodatavals"] = raster.nodatavals
143        result["dtypes"] = raster.dtypes
144
145    return result
146
147
148def gdf_attrs_to_rdggs(gdf_attrs: dict) -> RHEALPixDGGS:
149    return RHEALPixDGGS(ellipsoid=WGS84_ELLIPSOID,
150                        N_side=gdf_attrs["rhealpixdggs"]["n_side"],
151                        north_square=gdf_attrs["rhealpixdggs"]["north_square"],
152                        south_square=gdf_attrs["rhealpixdggs"]["south_square"])
153
154
155class RHEALPixDGGSHelper:
156
157    def __init__(self, rdggs: RHEALPixDGGS):
158        self.rdggs = rdggs
159
160    def rhealpixdef_to_proj_string(self) -> str:
161        assert self.rdggs.ellipsoid.a == WGS84_A and self.rdggs.ellipsoid.f == WGS84_F, \
162            "Only the WGS84 Ellipsoid can be used for now"
163        # The current value of ellipsoid.f in the library corresponds to that of the GRS80 ellipsoid
164        # and not to the WGS84 one (at least according to the values in Proj.4). It does not seem to initialize
165        # trouble right now (I have tried changing it and I have seen no changes) but maybe it could explain
166        # some future problems... Proj.4 defaults to GRS80 for rhealpix.
167
168        # The planar origin lon_0 and lat_0 is by default 0,0, so the planar coordinates 0,0 will fall on
169        # the Q3 cell. Notice that this is different from some depictions of the planar rHEALPix in some
170        # papers. We can change Lon_0 because proj supports that, but not lat_0
171        assert self.rdggs.ellipsoid.lat_0 == 0, f"Proj does not support a lat_0 parameter " \
172                                                f"(and {self.rdggs.ellipsoid.lat_0} was requested). This will not work."
173
174        # by default is +ellps=GRS80, but that is OK as long as the WGS84 in the rhealpixdggs library
175        # has the parameters of the GRS80. Besides that, I have not been able to detect differences
176        # using one or the other
177        if self.rdggs.ellipsoid.lon_0 == 0 and self.rdggs.ellipsoid.lat_0 == 0:
178            return f"+proj=rhealpix +south_square={self.rdggs.south_square} +north_square={self.rdggs.north_square}"
179        else:
180            return f"+proj=rhealpix +south_square={self.rdggs.south_square} +north_square={self.rdggs.north_square} " \
181                   f"+lon_0={self.rdggs.ellipsoid.lon_0}"
182
183    def rhealpixdef_to_wkt_string(self) -> str:
184        return pyproj.CRS(self.rhealpixdef_to_proj_string()).to_wkt()
185
186    def rhealpixdef_to_proj_ellps(self) -> str:
187        assert self.rdggs.ellipsoid.a == WGS84_A and self.rdggs.ellipsoid.f == WGS84_F, \
188            "Only the WGS84 Ellipsoid can be used for now"
189        return "WGS84"
190
191    def rhealpixdef_to_pyproj_crs(self) -> pyproj.CRS:
192        return pyproj.CRS(self.rhealpixdef_to_proj_string())
193
194
195    def cell_widths_for_all_resolutions(self) -> List[float]:
196        return [self.rdggs.cell_width(i) for i in range(self.rdggs.max_resolution)]
197
198    def get_closest_higher_resolution(self, base_resolution: float) -> Tuple[int, float]:
199        for i in range(self.rdggs.max_resolution):
200            if self.rdggs.cell_width(i) < base_resolution:
201                return i, self.rdggs.cell_width(i)
202
203    def get_closest_lower_resolution(self, base_resolution: float) -> Tuple[int, float]:
204        for i in range(self.rdggs.max_resolution):
205            if self.rdggs.cell_width(i) < base_resolution:
206                return i - 1, self.rdggs.cell_width(i - 1)
207
208    def get_closest_resolution(self, base_resolution: float) -> Tuple[int, float]:
209        for i in range(self.rdggs.max_resolution):
210            if self.rdggs.cell_width(i) < base_resolution:
211                higher = i, self.rdggs.cell_width(i)
212                lower = i - 1, self.rdggs.cell_width(i - 1)
213                if (lower[1] - base_resolution) < (base_resolution - higher[1]):
214                    return lower
215                else:
216                    return higher
217
218    def planar_boundary(self):
219        # Generating a non planar boundary would not work (or at least is far more tricky).
220        # The N and S cap corners are all aligned in geodetic coordinates, so that "squares" become lines...
221        cell_geoms = []
222        for cell in self.rdggs.grid(0):  # Every cell of resolution 0 (i.e., the faces of the cube)
223            (xmin, xmax), (ymin, ymax) = cell.xy_range()
224            cell_geoms.append(shapely.geometry.box(xmin, ymin, xmax, ymax))
225
226        return shapely.ops.unary_union(cell_geoms)
227
228    def project_and_clip_to_rhealpix(self, geom: dict, crs: str = None):
229        """
230        Project geom to rdggs, clip against the auids (the faces of the cube) and return the resulting geometry.
231
232        geom is a fiona/shapely geometry
233        crs is a proj4 string.
234        """
235        dst_crs = self.rhealpixdef_to_proj_string()
236        unfolded_cube = self.planar_boundary()
237        geom_rhealpix = fiona.transform.transform_geom(crs, dst_crs, geom)
238        return shapely.geometry.shape(geom_rhealpix).intersection(unfolded_cube)
239
240    def align_transform(self, transform: Affine, dst_resolution_idx: int) -> Affine:
241        """
242        Takes a transform not aligned to the rhealpix grid, and returns the closest one which is
243        aligned.
244
245        In theory, this could be substituted by something like this:
246        transform, width, height = rasterio.warp.aligned_target(transform, width, height, dst_resolution)
247        where width and height are taken from the output of calculate_default_transform.
248        Although this solution seems to align the result properly in the horizontal axis, it does not seem to
249        align properly in the vertical axis: the center of the output pixels do not correspond to the
250        centroids of the cells (it seems to align them with the edges of those cells). And this in turn
251        creates other problems interpreting each of those pixels as a cell (edges are problematic for this).
252        So, until proven otherwise, we will assume that this apparently too complex solution is the right one.
253        """
254        # This code takes the top-left coordinates of transform, sees to which rhealpix cell they belong to, and uses
255        # the closest vertex of that cell as the new top-left for the transform
256        current_left = transform[2]
257        current_top = transform[5]
258        current_topleft_cell = self.rdggs.cell_from_point(dst_resolution_idx, (current_left, current_top))
259        if current_topleft_cell is not None:
260            new_left, new_top = self._get_closest_vertex_in_rhealpix_cell(current_left, current_top, current_topleft_cell)
261        else:
262            # if current_topleft_cell is None, it means that current_left, current_top falls outside the N square
263            # (or maybe the S for some southern hemisphere datasets). We can realign taking as reference a cell that
264            # falls on the Equator for the left:
265            alternative_topleft_cell = self.rdggs.cell_from_point(dst_resolution_idx, (current_left, 0))
266            new_left, _ = self._get_closest_vertex_in_rhealpix_cell(current_left, current_top, alternative_topleft_cell)
267
268            # And with one that falls in the N square for the top
269            c = Cell(self.rdggs, ['N', 0])
270            (xmin, xmax), (ymin, ymax) = c.xy_range()
271            # I need to add/subtract a small tolerance for this to work
272            alternative_topleft_cell = self.rdggs.cell_from_point(dst_resolution_idx, (xmin+0.000001, current_top-0.000001))
273
274            if alternative_topleft_cell is None:
275                # Our final option is trying with the S square
276                c = Cell(self.rdggs, ['S', 0])
277                (xmin, xmax), (ymin, ymax) = c.xy_range()
278                # I need to add/subtract a small tolerance for this to work
279
280                alternative_topleft_cell = self.rdggs.cell_from_point(dst_resolution_idx,
281                                                                      (xmin+0.000001, current_top-0.000001))
282            _, new_top = self._get_closest_vertex_in_rhealpix_cell(current_left, current_top, alternative_topleft_cell)
283
284        return Affine.translation(new_left - current_left, new_top - current_top) * transform
285
286    def align_bounds(self, bounds:Tuple[float, float, float, float], dst_resolution_idx: int) -> Affine:
287        """
288        Takes a bounds tuple(left,top,right,bottom) not aligned to the rhealpix grid, and returns the
289        closest one which is aligned using the align_transform method.
290        """
291        return self.align_transform(Affine(bounds[2]-bounds[0], 0, bounds[0], 0, bounds[3]-bounds[1], bounds[1]),
292                                    dst_resolution_idx)
293
294
295    def get_bbox_in_rhealpix(self, input_crs: rasterio.crs.CRS, left: float, top: float, right: float,
296                             bottom: float) -> Tuple[float, float, float, float]:
297        """
298        Takes the left, top, right, bottom coordinates of a bounding box, coordinates in input_crs, and
299        returns them reprojected to rhealpix as defined in rdggs (left, top, right, bottom).
300        """
301        dst_crs = self.rhealpixdef_to_proj_string()
302        transformer = pyproj.Transformer.from_crs(input_crs, dst_crs, always_xy=True)
303        left_rheal, top_rheal = transformer.transform(left, top)
304        right_rheal, bottom_rheal = transformer.transform(right, bottom)
305
306        # It is possible, I think, that the left,top and right,bottom corners in the original CRS are not
307        # the left,top and right,bottom corners after transforming to rhealpix (because the N and S squares are
308        # rotated). So we generate the other two corners and take the leftmost, rightmost, topmost and bottom-most values:
309        left_right_candidate, top_bottom_candidate = transformer.transform(left, bottom)
310        left_rheal = min(left_rheal, left_right_candidate)
311        right_rheal = max(right_rheal, left_right_candidate)
312        top_rheal = max(top_rheal, top_bottom_candidate)
313        bottom_rheal = min(bottom_rheal, top_bottom_candidate)
314
315        left_right_candidate, top_bottom_candidate = transformer.transform(right, top)
316        left_rheal = min(left_rheal, left_right_candidate)
317        right_rheal = max(right_rheal, left_right_candidate)
318        top_rheal = max(top_rheal, top_bottom_candidate)
319        bottom_rheal = min(bottom_rheal, top_bottom_candidate)
320
321        return left_rheal, top_rheal, right_rheal, bottom_rheal
322
323    def _get_closest_vertex_in_rhealpix_cell(self, left: float, top: float, cell: Cell) -> Tuple[float, float]:
324        """
325         Returns the vertex in cell which is closer (cartesian distance) to the point left, top.
326         We are assuming plane cells.
327         """
328
329        def cartesian_dist(x1, y1, x2, y2):
330            return ((x2 - x1) ** 2 + (y2 - y1) ** 2) ** 0.5
331
332        distances_and_vertices = [(cartesian_dist(left, top, x, y), x, y) for x, y in cell.vertices()]
333        _, closest_x, closest_y = min(distances_and_vertices)
334        return closest_x, closest_y