|
| 1 | +import numpy as np |
| 2 | + |
| 3 | +from functools import cached_property |
| 4 | +from xarray.core.dataset import Dataset |
| 5 | + |
| 6 | +renaming_dict = {"lonCell": "xCell", |
| 7 | + "latCell": "yCell", |
| 8 | + "lonEdge": "xEdge", |
| 9 | + "latEdge": "yEdge", |
| 10 | + "lonVertex": "xVertex", |
| 11 | + "latVertex": "yVertex"} |
| 12 | + |
| 13 | +connectivity_arrays = ["cellsOnEdge", |
| 14 | + "cellsOnVertex", |
| 15 | + "verticesOnEdge", |
| 16 | + "verticesOnCell"] |
| 17 | + |
| 18 | +class Descriptor: |
| 19 | + """ |
| 20 | + Class describing unstructured MPAS meshes in order to support plotting |
| 21 | + within `matplotlib`. The class constains various methods to create |
| 22 | + `matplotlib.collections.PolyCollection` objects for variables defined at |
| 23 | + cell centers, vertices, and edges. |
| 24 | + |
| 25 | +
|
| 26 | + Attributes |
| 27 | + ---------- |
| 28 | + latlon : boolean |
| 29 | + Whethere to use the lat/lon coordinates in patch construction |
| 30 | +
|
| 31 | + NOTE: I don't think this is needed if the projection arg is |
| 32 | + properly used at initilaization |
| 33 | +
|
| 34 | + projection : cartopy.crs.CRS |
| 35 | +
|
| 36 | + transform : cartopy.crs.CRS |
| 37 | +
|
| 38 | + cell_patches : np.ndarray |
| 39 | +
|
| 40 | + edge_patches : np.ndarray |
| 41 | +
|
| 42 | + vertex_patches : np.ndarray |
| 43 | + """ |
| 44 | + def __init__(self, ds, projection=None, transform=None, use_latlon=False): |
| 45 | + """ |
| 46 | + """ |
| 47 | + self.latlon = use_latlon |
| 48 | + self.projection = projection |
| 49 | + self.transform = transform |
| 50 | + |
| 51 | + # if mesh is on a sphere, force the use of lat lon coords |
| 52 | + if ds.attrs["on_a_sphere"].strip().upper() == 'YES': |
| 53 | + self.latlon = True |
| 54 | + # also check if projection requires lat/lon coords |
| 55 | + |
| 56 | + # create a minimal dataset, stored as an attr, for patch creation |
| 57 | + self.ds = self.create_minimal_dataset(ds) |
| 58 | + |
| 59 | + # reproject the minimal dataset, even for non-spherical meshes |
| 60 | + if projection and transform: |
| 61 | + self._transform_coordinates(projection, transform) |
| 62 | + |
| 63 | + def create_minimal_dataset(self, ds): |
| 64 | + """ |
| 65 | + Create a xarray.Dataset that contains the minimal subset of |
| 66 | + coordinate / connectivity arrays needed to create pathces for plotting |
| 67 | + """ |
| 68 | + |
| 69 | + if self.latlon: |
| 70 | + coordinate_arrays = list(renaming_dict.keys()) |
| 71 | + else: |
| 72 | + coordinate_arrays = list(renaming_dict.values()) |
| 73 | + |
| 74 | + # list of coordinate / connectivity arrays needed to create patches |
| 75 | + mesh_arrays = coordinate_arrays + connectivity_arrays |
| 76 | + |
| 77 | + # get the subset of arrays from the mesh dataset |
| 78 | + minimal_ds = ds[mesh_arrays] |
| 79 | + |
| 80 | + # delete the attributes in the minimal dataset to avoid confusion |
| 81 | + minimal_ds.attrs.clear() |
| 82 | + |
| 83 | + # should zero index the connectivity arrays here. |
| 84 | + |
| 85 | + if self.latlon: |
| 86 | + |
| 87 | + # convert lat/lon coordinates from radian to degrees |
| 88 | + for loc in ["Cell", "Edge", "Vertex"]: |
| 89 | + minimal_ds[f"lon{loc}"] = np.rad2deg(minimal_ds[f"lon{loc}"]) |
| 90 | + minimal_ds[f"lat{loc}"] = np.rad2deg(minimal_ds[f"lat{loc}"]) |
| 91 | + |
| 92 | + # rename the coordinate arrays to all be named x.../y... |
| 93 | + # irrespective of whether spherical or cartesian coords are used |
| 94 | + minimal_ds = minimal_ds.rename(renaming_dict) |
| 95 | + |
| 96 | + return minimal_ds |
| 97 | + |
| 98 | + @cached_property |
| 99 | + def cell_patches(self): |
| 100 | + patches = _compute_cell_patches(self.ds) |
| 101 | + patches = self._fix_antimeridian(patches, "Cell") |
| 102 | + return patches |
| 103 | + |
| 104 | + @cached_property |
| 105 | + def edge_patches(self): |
| 106 | + patches = _compute_edge_patches(self.ds) |
| 107 | + patches = self._fix_antimeridian(patches, "Edge") |
| 108 | + return patches |
| 109 | + |
| 110 | + @cached_property |
| 111 | + def vertex_patches(self): |
| 112 | + patches = _compute_vertex_patches(self.ds) |
| 113 | + patches = self._fix_antimeridian(patches, "Vertex") |
| 114 | + return patches |
| 115 | + |
| 116 | + def _transform_coordinates(self, projection, transform): |
| 117 | + """ |
| 118 | + """ |
| 119 | + |
| 120 | + for loc in ["Cell", "Edge", "Vertex"]: |
| 121 | + |
| 122 | + transformed_coords = projection.transform_points(transform, |
| 123 | + self.ds[f"x{loc}"], self.ds[f"y{loc}"]) |
| 124 | + |
| 125 | + # transformed_coords is a np array so need to assign to the values |
| 126 | + self.ds[f"x{loc}"].values = transformed_coords[:, 0] |
| 127 | + self.ds[f"y{loc}"].values = transformed_coords[:, 1] |
| 128 | + |
| 129 | + def _fix_antimeridian(self, patches, loc, projection=None): |
| 130 | + """Correct vertices of patches that cross the antimeridian. |
| 131 | +
|
| 132 | + NOTE: Can this be a decorator? |
| 133 | + """ |
| 134 | + # coordinate arrays are transformed at initalization, so using the |
| 135 | + # transform size limit, not the projection |
| 136 | + if not projection: |
| 137 | + projection = self.projection |
| 138 | + |
| 139 | + # should be able to come up with a default size limit here, or maybe |
| 140 | + # it's already an attribute(?) Should also factor in a precomputed |
| 141 | + # axis period, as set in the attributes of the input dataset |
| 142 | + if projection: |
| 143 | + # convert to numpy array to that broadcasting below will work |
| 144 | + x_center = np.array(self.ds[f"x{loc}"]) |
| 145 | + |
| 146 | + # get distance b/w the center and vertices of the patches |
| 147 | + # NOTE: using data from masked patches array so that we compute |
| 148 | + # mask only corresponds to patches that cross the boundary, |
| 149 | + # (i.e. NOT a mask of all invalid cells). May need to be |
| 150 | + # carefull about the fillvalue depending on the transform |
| 151 | + half_distance = x_center[:, np.newaxis] - patches[...,0].data |
| 152 | + |
| 153 | + # get the size limit of the projection; |
| 154 | + size_limit = np.abs(projection.x_limits[1] - |
| 155 | + projection.x_limits[0]) / (2 * np.sqrt(2)) |
| 156 | + |
| 157 | + # left and right mask, with same number of dims as the patches |
| 158 | + l_mask = (half_distance > size_limit)[..., np.newaxis] |
| 159 | + r_mask = (half_distance < -size_limit)[..., np.newaxis] |
| 160 | + |
| 161 | + """ |
| 162 | + # Old approach masks out all patches that cross the antimeridian. |
| 163 | + # This is unnessarily restrictive. New approach corrects |
| 164 | + # the x-coordinates of vertices that lie outside the projections |
| 165 | + # bounds, which isn't perfect either |
| 166 | +
|
| 167 | + patches.mask |= l_mask |
| 168 | + patches.mask |= r_mask |
| 169 | + """ |
| 170 | + |
| 171 | + # get valid half distances for the patches that cross boundary |
| 172 | + l_offset = np.ma.MaskedArray(half_distance, |
| 173 | + ~np.any(l_mask, axis=1) | l_mask[...,0]) |
| 174 | + r_offset = np.ma.MaskedArray(half_distance, |
| 175 | + ~np.any(r_mask, axis=1) | r_mask[...,0]) |
| 176 | + |
| 177 | + # For vertices that cross the antimeridian reset the x-coordinate |
| 178 | + # of invalid vertex to be the center of the patch plus the |
| 179 | + # mean valid half distance. |
| 180 | + # |
| 181 | + # NOTE: this only fixes patches on the side of plot where they |
| 182 | + # cross the antimeridian, leaving an empty zipper like pattern |
| 183 | + # mirrored over the y-axis. |
| 184 | + patches[...,0] = np.ma.where(~l_mask[...,0], patches[...,0], |
| 185 | + x_center[:, np.newaxis] + l_offset.mean(1)[...,np.newaxis]) |
| 186 | + patches[...,0] = np.ma.where(~r_mask[...,0], patches[...,0], |
| 187 | + x_center[:, np.newaxis] + r_offset.mean(1)[...,np.newaxis]) |
| 188 | + |
| 189 | + return patches |
| 190 | + |
| 191 | + def transform_patches(self, patches, projection, transform): |
| 192 | + """ |
| 193 | + """ |
| 194 | + |
| 195 | + raise NotImplementedError("This is a place holder. Do not use.") |
| 196 | + |
| 197 | + transformed_patches = projection.transform_points(transform, |
| 198 | + patches[..., 0], patches[..., 1]) |
| 199 | + |
| 200 | + # transformation will return x,y,z values. Only need x and y |
| 201 | + patches.data[...] = transformed_patches[..., 0:2] |
| 202 | + |
| 203 | + return patches |
| 204 | + |
| 205 | +def _compute_cell_patches(ds): |
| 206 | + |
| 207 | + # get a mask of the active vertices |
| 208 | + mask = ds.verticesOnCell == 0 |
| 209 | + |
| 210 | + # get the coordinates needed to patch construction |
| 211 | + xVertex = ds.xVertex |
| 212 | + yVertex = ds.yVertex |
| 213 | + |
| 214 | + # account for zero indexing |
| 215 | + verticesOnCell = ds.verticesOnCell - 1 |
| 216 | + |
| 217 | + # reshape/expand the vertices coordinate arrays |
| 218 | + x_vert = np.ma.MaskedArray(xVertex[verticesOnCell], mask=mask) |
| 219 | + y_vert = np.ma.MaskedArray(yVertex[verticesOnCell], mask=mask) |
| 220 | + |
| 221 | + verts = np.ma.stack((x_vert, y_vert), axis=-1) |
| 222 | + |
| 223 | + return verts |
| 224 | + |
| 225 | +def _compute_edge_patches(ds, latlon=False): |
| 226 | + |
| 227 | + # account for zeros indexing |
| 228 | + cellsOnEdge = ds.cellsOnEdge - 1 |
| 229 | + verticesOnEdge = ds.verticesOnEdge - 1 |
| 230 | + |
| 231 | + # is this masking sufficent ? |
| 232 | + cellMask = cellsOnEdge <= 0 |
| 233 | + vertexMask = verticesOnEdge <= 0 |
| 234 | + |
| 235 | + # get the coordinates needed to patch construction |
| 236 | + xCell = ds.xCell |
| 237 | + yCell = ds.yCell |
| 238 | + xVertex = ds.xVertex |
| 239 | + yVertex = ds.yVertex |
| 240 | + |
| 241 | + # get subset of cell coordinate arrays corresponding to edge patches |
| 242 | + xCell = np.ma.MaskedArray(xCell[cellsOnEdge], mask=cellMask) |
| 243 | + yCell = np.ma.MaskedArray(yCell[cellsOnEdge], mask=cellMask) |
| 244 | + # get subset of vertex coordinate arrays corresponding to edge patches |
| 245 | + xVertex = np.ma.MaskedArray(xVertex[verticesOnEdge], mask=vertexMask) |
| 246 | + yVertex = np.ma.MaskedArray(yVertex[verticesOnEdge], mask=vertexMask) |
| 247 | + |
| 248 | + x_vert = np.ma.stack((xCell[:,0], xVertex[:,0], |
| 249 | + xCell[:,1], xVertex[:,1]), axis=-1) |
| 250 | + |
| 251 | + y_vert = np.ma.stack((yCell[:,0], yVertex[:,0], |
| 252 | + yCell[:,1], yVertex[:,1]), axis=-1) |
| 253 | + |
| 254 | + |
| 255 | + verts = np.ma.stack((x_vert, y_vert), axis=-1) |
| 256 | + |
| 257 | + return verts |
| 258 | + |
| 259 | +def _compute_vertex_patches(ds, latlon=False): |
| 260 | + |
| 261 | + # get a mask of the active vertices |
| 262 | + mask = ds.cellsOnVertex == 0 |
| 263 | + |
| 264 | + # get the coordinates needed to patch construction |
| 265 | + xCell = ds.xCell |
| 266 | + yCell = ds.yCell |
| 267 | + |
| 268 | + # account for zero indexing |
| 269 | + cellsOnVertex = ds.cellsOnVertex - 1 |
| 270 | + |
| 271 | + # reshape/expand the vertices coordinate arrays |
| 272 | + x_vert = np.ma.MaskedArray(xCell[cellsOnVertex], mask=mask) |
| 273 | + y_vert = np.ma.MaskedArray(yCell[cellsOnVertex], mask=mask) |
| 274 | + |
| 275 | + verts = np.ma.stack((x_vert, y_vert), axis=-1) |
| 276 | + |
| 277 | + return verts |
0 commit comments