diff --git a/benchmarks/mpas_ocean.py b/benchmarks/mpas_ocean.py index 0f6849b18..1db46e0f9 100644 --- a/benchmarks/mpas_ocean.py +++ b/benchmarks/mpas_ocean.py @@ -120,6 +120,8 @@ def time_nearest_neighbor_remapping(self): def time_inverse_distance_weighted_remapping(self): self.uxds_120["bottomDepth"].remap.inverse_distance_weighted(self.uxds_480.uxgrid) + def time_bilinear_remapping(self): + self.uxds_120["bottomDepth"].remap.bilinear(self.uxds_480.uxgrid) class RemapUpsample: @@ -136,6 +138,9 @@ def time_nearest_neighbor_remapping(self): def time_inverse_distance_weighted_remapping(self): self.uxds_480["bottomDepth"].remap.inverse_distance_weighted(self.uxds_120.uxgrid) + def time_bilinear_remapping(self): + self.uxds_480["bottomDepth"].remap.bilinear(self.uxds_120.uxgrid) + class HoleEdgeIndices(DatasetBenchmark): def time_construct_hole_edge_indices(self, resolution): diff --git a/docs/api.rst b/docs/api.rst index 63cbbac44..55f6aa768 100644 --- a/docs/api.rst +++ b/docs/api.rst @@ -369,6 +369,7 @@ UxDataArray UxDataArray.remap UxDataArray.remap.nearest_neighbor UxDataArray.remap.inverse_distance_weighted + UxDataArray.remap.bilinear UxDataset ~~~~~~~~~ @@ -380,7 +381,7 @@ UxDataset UxDataset.remap UxDataset.remap.nearest_neighbor UxDataset.remap.inverse_distance_weighted - + UxDataset.remap.bilinear Mathematical Operators ---------------------- diff --git a/docs/user-guide/remapping.ipynb b/docs/user-guide/remapping.ipynb index 861785461..d4f6b0770 100644 --- a/docs/user-guide/remapping.ipynb +++ b/docs/user-guide/remapping.ipynb @@ -1,9 +1,8 @@ { "cells": [ { - "cell_type": "markdown", - "id": "d9d3f5a8-6d3c-4a7e-9150-a2915f3e0ceb", "metadata": {}, + "cell_type": "markdown", "source": [ "# Remapping\n", "\n", @@ -11,18 +10,19 @@ "\n", "For a comprehensive overview of common remapping techniques, see the [Climate Data Guide: Regridding Overview](https://climatedataguide.ucar.edu/climate-tools/regridding-overview).\n", "\n", - "UXarray currently supports two primary remapping methods:\n", + "UXarray currently supports three remapping techniques:\n", "\n", "- **Nearest Neighbor** \n", - "- **Inverse Distance Weighted** \n" - ] + "- **Inverse Distance Weighted**\n", + "- **Bilinear**\n" + ], + "id": "7eec39631eeeb6f8" }, { - "cell_type": "code", - "execution_count": null, - "id": "7449507f-3d79-4e86-a775-3b9137153adc", "metadata": {}, + "cell_type": "code", "outputs": [], + "execution_count": null, "source": [ "import os\n", "import urllib.request\n", @@ -39,7 +39,8 @@ "hv.extension(\"matplotlib\")\n", "\n", "common_kwargs = {\"cmap\": cmocean.cm.deep, \"features\": [\"coastline\"]}" - ] + ], + "id": "bc895fe997d10515" }, { "cell_type": "markdown", @@ -53,10 +54,8 @@ }, { "cell_type": "code", - "execution_count": null, "id": "4d73a380-349d-473d-8e57-10c52102adca", "metadata": {}, - "outputs": [], "source": [ "data_var = \"bottomDepth\"\n", "\n", @@ -81,36 +80,40 @@ "}\n", "uxds_480 = ux.open_dataset(*file_path_dict[\"480km\"])\n", "uxds_120 = ux.open_dataset(*file_path_dict[\"120km\"])" - ] + ], + "outputs": [], + "execution_count": null }, { "cell_type": "code", - "execution_count": null, "id": "da0b1ff8-da1a-4c6c-9031-749b34bfad7a", "metadata": {}, - "outputs": [], "source": [ "(\n", " uxds_480[\"bottomDepth\"].plot(title=\"Bottom Depth (480km)\", **common_kwargs)\n", " + uxds_120[\"bottomDepth\"].plot(title=\"Bottom Depth (120km)\", **common_kwargs)\n", ").cols(1).opts(fig_size=200)" - ] + ], + "outputs": [], + "execution_count": null }, { "cell_type": "markdown", "id": "0d2345bc-ce03-48b5-b08c-6e9c679f3bc1", "metadata": {}, - "source": "We can view the supported remapping methods by accessing the `.remap` accessor that is part of a `UxDataArray` or `UxDataset`" + "source": [ + "We can view the supported remapping methods by accessing the `.remap` accessor that is part of a `UxDataArray` or `UxDataset`" + ] }, { "cell_type": "code", - "execution_count": null, "id": "2453895d-b41d-47fe-bc2b-42358a9acbe5", "metadata": {}, - "outputs": [], "source": [ "uxds_120.remap" - ] + ], + "outputs": [], + "execution_count": null }, { "cell_type": "markdown", @@ -148,22 +151,20 @@ }, { "cell_type": "code", - "execution_count": null, "id": "d4550735-053a-4542-b259-fb7d8c2e6fae", "metadata": {}, - "outputs": [], "source": [ "upsampling = uxds_480[\"bottomDepth\"].remap.nearest_neighbor(\n", " destination_grid=uxds_120.uxgrid, remap_to=\"faces\"\n", ")" - ] + ], + "outputs": [], + "execution_count": null }, { "cell_type": "code", - "execution_count": null, "id": "cf2cf918-62f8-4aa4-9fa1-122bc06862ce", "metadata": {}, - "outputs": [], "source": [ "(\n", " uxds_480[\"bottomDepth\"].plot(title=\"Bottom Depth (480km)\", **common_kwargs)\n", @@ -178,7 +179,9 @@ " **common_kwargs,\n", " )\n", ").cols(2).opts(fig_size=200)" - ] + ], + "outputs": [], + "execution_count": null }, { "cell_type": "markdown", @@ -200,22 +203,20 @@ }, { "cell_type": "code", - "execution_count": null, "id": "40094ba0-0dad-48d7-af70-040f088d7be5", "metadata": {}, - "outputs": [], "source": [ "downsampling = uxds_120[\"bottomDepth\"].remap.nearest_neighbor(\n", " destination_grid=uxds_480.uxgrid, remap_to=\"face centers\"\n", ")" - ] + ], + "outputs": [], + "execution_count": null }, { "cell_type": "code", - "execution_count": null, "id": "65d3f7db-9820-4c46-8d78-571a8ea48ef1", "metadata": {}, - "outputs": [], "source": [ "(\n", " uxds_120[\"bottomDepth\"].plot(title=\"Bottom Depth (120km)\", **common_kwargs)\n", @@ -230,7 +231,9 @@ " **common_kwargs,\n", " )\n", ").cols(2).opts(fig_size=200)" - ] + ], + "outputs": [], + "execution_count": null }, { "cell_type": "markdown", @@ -277,22 +280,20 @@ }, { "cell_type": "code", - "execution_count": null, "id": "400398d5-5cc0-4790-9a95-cb0f88cc1ca8", "metadata": {}, - "outputs": [], "source": [ "upsampling_idw = uxds_480[\"bottomDepth\"].remap.inverse_distance_weighted(\n", " destination_grid=uxds_120.uxgrid, remap_to=\"faces\"\n", ")" - ] + ], + "outputs": [], + "execution_count": null }, { "cell_type": "code", - "execution_count": null, "id": "8dd1d8e3-54e2-4710-b132-f69f0ba950fd", "metadata": {}, - "outputs": [], "source": [ "(\n", " uxds_480[\"bottomDepth\"].plot(title=\"Bottom Depth (480km)\", **common_kwargs)\n", @@ -309,7 +310,9 @@ " **common_kwargs,\n", " )\n", ").cols(2).opts(fig_size=200)" - ] + ], + "outputs": [], + "execution_count": null }, { "cell_type": "markdown", @@ -328,22 +331,20 @@ }, { "cell_type": "code", - "execution_count": null, "id": "7d21c42d-d368-4a34-8dcb-6643f0a0a7d1", "metadata": {}, - "outputs": [], "source": [ "downsampling_idw = uxds_120[\"bottomDepth\"].remap.inverse_distance_weighted(\n", " destination_grid=uxds_480.uxgrid, remap_to=\"faces\"\n", ")" - ] + ], + "outputs": [], + "execution_count": null }, { "cell_type": "code", - "execution_count": null, "id": "935a7a74-7d88-49ae-bc20-7b18d76795c9", "metadata": {}, - "outputs": [], "source": [ "(\n", " uxds_120[\"bottomDepth\"].plot(title=\"Bottom Depth (120km)\", **common_kwargs)\n", @@ -360,7 +361,9 @@ " **common_kwargs,\n", " )\n", ").cols(2).opts(fig_size=200)" - ] + ], + "outputs": [], + "execution_count": null }, { "cell_type": "markdown", @@ -380,10 +383,8 @@ }, { "cell_type": "code", - "execution_count": null, "id": "9e31c8ec-75a0-4898-96e7-35a4b4853ad0", "metadata": {}, - "outputs": [], "source": [ "downsampling_idw_low = uxds_120[\"bottomDepth\"].remap.inverse_distance_weighted(\n", " uxds_480.uxgrid, remap_to=\"faces\", power=1, k=2\n", @@ -391,14 +392,14 @@ "downsampling_idw_high = uxds_120[\"bottomDepth\"].remap.inverse_distance_weighted(\n", " uxds_480.uxgrid, remap_to=\"faces\", power=5, k=128\n", ")" - ] + ], + "outputs": [], + "execution_count": null }, { "cell_type": "code", - "execution_count": null, "id": "88756342-64e0-42f4-96d9-b9822e002bd7", "metadata": {}, - "outputs": [], "source": [ "(\n", " downsampling_idw_low.plot(\n", @@ -414,14 +415,14 @@ " **common_kwargs,\n", " )\n", ").cols(1).opts(fig_size=200)" - ] + ], + "outputs": [], + "execution_count": null }, { "cell_type": "code", - "execution_count": null, "id": "439ac480-c4f3-4080-a18c-8f670367c194", "metadata": {}, - "outputs": [], "source": [ "upsampling_idw_low = uxds_480[\"bottomDepth\"].remap.inverse_distance_weighted(\n", " uxds_120.uxgrid, remap_to=\"faces\", power=1, k=2\n", @@ -429,14 +430,14 @@ "upsampling_idw_high = uxds_480[\"bottomDepth\"].remap.inverse_distance_weighted(\n", " uxds_120.uxgrid, remap_to=\"faces\", power=5, k=128\n", ")" - ] + ], + "outputs": [], + "execution_count": null }, { "cell_type": "code", - "execution_count": null, "id": "0dc0497d-13bd-4b5b-9792-ad6ce26aded5", "metadata": {}, - "outputs": [], "source": [ "(\n", " upsampling_idw_low.plot(\n", @@ -452,7 +453,9 @@ " **common_kwargs,\n", " )\n", ").cols(1).opts(fig_size=200)" - ] + ], + "outputs": [], + "execution_count": null }, { "cell_type": "markdown", @@ -463,6 +466,112 @@ "\n", "In contrast, upsampling shows a much stronger response to parameter tuning. Spreading data from a coarse grid (with fewer, larger faces) onto a finer mesh (with more, smaller faces) means that increasing the neighbor count draws in values from a wider spatial context, leading to more pronounced variations in the interpolated field.\n" ] + }, + { + "cell_type": "markdown", + "id": "f1f33631-19b7-4b73-8452-7dc1e3fa48a2", + "metadata": {}, + "source": "## Bilinear" + }, + { + "cell_type": "markdown", + "id": "6bec26ce-67b6-4300-a310-63cbac2b289a", + "metadata": {}, + "source": [ + "Bilinear remapping breaks down the grid into triangles, and then finds the triangle that contains each point on the destinitation grid. It then uses the values stored at each vertex to bilinearly find a value for the point, depending on it's postion inside the triangle." + ] + }, + { + "cell_type": "markdown", + "id": "4689fe83-ea07-40a0-8fde-83b04523ef24", + "metadata": {}, + "source": [ + "Bilinear remapping finds for each point the polygon that contains it. Then, based on the position of the point within the polygon, it calculates weights. These weights are then applied to the values stored on the polygon. This gives a value to assign to the point. Bilinear remapping tends to give much more smooth results than the previous remapping methods." + ] + }, + { + "cell_type": "markdown", + "id": "467252cb-9e07-42bd-8734-15666f612387", + "metadata": {}, + "source": "### Upsampling" + }, + { + "metadata": {}, + "cell_type": "code", + "source": [ + "upsampling_bl = uxds_480[\"bottomDepth\"].remap.bilinear(\n", + " destination_grid=uxds_120.uxgrid, remap_to=\"faces\"\n", + ")" + ], + "id": "c807a61bbc5bffff", + "outputs": [], + "execution_count": null + }, + { + "metadata": {}, + "cell_type": "code", + "source": [ + "(\n", + " uxds_480[\"bottomDepth\"].plot(title=\"Bottom Depth (480km)\", **common_kwargs)\n", + " + upsampling_bl.plot(\n", + " title=\"Remapped Bottom Depth (480km to 120km)\", **common_kwargs\n", + " )\n", + " + uxds_480[\"bottomDepth\"].plot(\n", + " title=\"Zoomed (480km)\", xlim=(-10, 10), ylim=(-5, 5), **common_kwargs\n", + " )\n", + " + upsampling_bl.plot(\n", + " title=\"Zoomed Remap (480km to 120km)\",\n", + " xlim=(-10, 10),\n", + " ylim=(-5, 5),\n", + " **common_kwargs,\n", + " )\n", + ").cols(2).opts(fig_size=200)" + ], + "id": "1429cc143fcd6694", + "outputs": [], + "execution_count": null + }, + { + "cell_type": "markdown", + "id": "9ae38e84", + "metadata": {}, + "source": "### Downsampling" + }, + { + "metadata": {}, + "cell_type": "code", + "source": [ + "downsampling_bl = uxds_120[\"bottomDepth\"].remap.bilinear(\n", + " destination_grid=uxds_480.uxgrid, remap_to=\"faces\"\n", + ")" + ], + "id": "c3e12c15307c5042", + "outputs": [], + "execution_count": null + }, + { + "metadata": {}, + "cell_type": "code", + "source": [ + "(\n", + " uxds_120[\"bottomDepth\"].plot(title=\"Bottom Depth (120km)\", **common_kwargs)\n", + " + downsampling_bl.plot(\n", + " title=\"Remapped Bottom Depth (120km to 480km)\", **common_kwargs\n", + " )\n", + " + uxds_120[\"bottomDepth\"].plot(\n", + " title=\"Zoomed (120km)\", xlim=(-10, 10), ylim=(-5, 5), **common_kwargs\n", + " )\n", + " + downsampling_bl.plot(\n", + " title=\"Zoomed Remap (120km to 480km)\",\n", + " xlim=(-10, 10),\n", + " ylim=(-5, 5),\n", + " **common_kwargs,\n", + " )\n", + ").cols(2).opts(fig_size=200)" + ], + "id": "bec1b9d05d0a40ba", + "outputs": [], + "execution_count": null } ], "metadata": { @@ -481,7 +590,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.12.7" + "version": "3.12.8" } }, "nbformat": 4, diff --git a/test/test_remap.py b/test/test_remap.py index e2b1a1221..f2e6c685e 100644 --- a/test/test_remap.py +++ b/test/test_remap.py @@ -15,6 +15,9 @@ for i in (1, 2, 3) ] mpasfile_QU = ROOT / "meshfiles" / "mpas" / "QU" / "mesh.QU.1920km.151026.nc" +mpasfile_QU_2 = ROOT / "meshfiles" / "mpas" / "QU" / "oQU480.231010.nc" +outCSne30 = ROOT / "meshfiles" / "ugrid" / "outCSne30" / "outCSne30.ug" +outCSne30_var2 = ROOT / "meshfiles" / "ugrid" / "outCSne30" / "outCSne30_var2.nc" # ------------------------------------------------------------ # Helper: small 3‐point spherical grid @@ -170,13 +173,14 @@ def test_idw_power_zero_is_mean(): assert np.allclose(out.values, expected) - def test_invalid_remap_to_raises(): src_da, grid = _make_node_da([0.0, 1.0, 2.0]) with pytest.raises(ValueError): src_da.remap.nearest_neighbor(destination_grid=grid, remap_to="foobars") with pytest.raises(ValueError): src_da.remap.inverse_distance_weighted(destination_grid=grid, remap_to="123") + with pytest.raises(ValueError): + src_da.remap.bilinear(destination_grid=grid, remap_to="nothing") def test_nn_equals_idw_high_power(): @@ -194,4 +198,61 @@ def test_dataset_remap_preserves_coords(): uxds = ux.open_dataset(gridfile_geoflow, dsfiles_geoflow[0]) uxds = uxds.assign_coords(time=("time", np.arange(len(uxds.time)))) dest = ux.open_grid(mpasfile_QU) - ds_out = uxds.remap.nearest_neighbor(destination_grid=dest, remap_to="nodes") + ds_out = uxds.remap.inverse_distance_weighted(destination_grid=dest, remap_to="nodes") + assert "time" in ds_out.coords + + +# ------------------------------------------------------------ +# Bilinear tests +# ------------------------------------------------------------ +def test_b_modifies_values(): + """Bilinear remap should change the array when remap_to != source.""" + uxds = ux.open_dataset(mpasfile_QU, mpasfile_QU) + dest = ux.open_grid(gridfile_geoflow) + da_idw = uxds["latCell"].remap.inverse_distance_weighted( + destination_grid=dest, remap_to="nodes" + ) + assert not np.array_equal(uxds["latCell"].values, da_idw.values) + + +def test_b_return_types_and_counts(): + """Bilinear remap returns UxDataArray or UxDataset with correct var counts.""" + uxds = ux.open_dataset(mpasfile_QU, mpasfile_QU) + dest = ux.open_grid(gridfile_geoflow) + + da_idw = uxds["latCell"].remap.bilinear(destination_grid=dest) + + # Create a dataset with only a face centered variable + face_center_uxds = uxds["latCell"].to_dataset() + ds_idw = face_center_uxds.remap.bilinear(destination_grid=dest) + + assert isinstance(da_idw, UxDataArray) + assert isinstance(ds_idw, UxDataset) + assert set(ds_idw.data_vars) == set(face_center_uxds.data_vars) + +def test_b_edge_centers_dim_change(): + """Bilinear remap to edge centers produces an 'n_edge' dimension.""" + uxds = ux.open_dataset(mpasfile_QU, mpasfile_QU) + dest = ux.open_grid(mpasfile_QU) + da = uxds["latCell"].remap.bilinear( + destination_grid=dest, remap_to="edges" + ) + assert "n_edge" in da.dims + + +def test_b_value_errors(): + """Bilinear remapping raises a value error when the source data is not on the faces""" + uxds = ux.open_dataset(gridfile_geoflow, dsfiles_geoflow[0]) + dest = ux.open_grid(gridfile_geoflow) + + with nt.assert_raises(ValueError): + uxds['v1'].remap.bilinear(destination_grid=dest, remap_to="nodes") + +def test_b_quadrilateral(): + """Bilinear remapping on quadrilaterals, to test Newton Quadrilateral weights calculation""" + uxds = ux.open_dataset(outCSne30, outCSne30_var2) + dest = ux.open_grid(mpasfile_QU) + + out = uxds['var2'].remap.bilinear(destination_grid=dest) + + assert out.size > 0 diff --git a/uxarray/grid/geometry.py b/uxarray/grid/geometry.py index eca32c861..30af09f70 100644 --- a/uxarray/grid/geometry.py +++ b/uxarray/grid/geometry.py @@ -7,11 +7,14 @@ ERROR_TOLERANCE, INT_DTYPE, INT_FILL_VALUE, + MACHINE_EPSILON, ) from uxarray.grid.coordinates import _xyz_to_lonlat_rad from uxarray.grid.intersections import ( gca_gca_intersection, ) +from uxarray.grid.point_in_face import _face_contains_point +from uxarray.grid.utils import _get_cartesian_face_edge_nodes POLE_POINTS_XYZ = { "North": np.array([0.0, 0.0, 1.0]), @@ -1132,3 +1135,328 @@ def haversine_distance(lon_a, lat_a, lon_b, lat_b): # Return the gotten distance return distance + + +@njit(cache=True) +def barycentric_coordinates_cartesian(polygon_xyz, point_xyz): + """ + Compute the barycentric coordinates of a point inside a convex polygon, using cartesian coordinates. + + Parameters + ---------- + polygon_xyz: np.array + Cartesian coordinates of the polygon nodes + point_xyz: np.array + Cartesian coordinates of the point. + + Examples + -------- + Define a cartesian point: + >>> import numpy as np + >>> point_xyz = np.array([0.0, 0.0, 1.0], dtype=np.float64) + + Define a cartesian polygon: + >>> polygon_xyz = np.array( + ... [[0.0, 0.0, 1.0], [1.0, 0.0, 0.0], [0.0, 0.0, -1.0]], dtype=np.float64 + ... ) + + Find the weights: + >>> bar_weights = barycentric_coordinates_cartesian(polygon_xyz, point_xyz) + + Returns + ------- + weights: np.ndarray + Barycentric coordinates corresponding to each vertex in the triangle. + nodes: np.ndarray + Indices of which nodes of the polygon form the triangle containing the point + """ + + n = len(polygon_xyz) + + # If the polygon is a triangle, we can use the `triangle_line_intersection` algorithm to calculate the weights + if n == 3: + # Assign an empty array of size 3 to hold final weights + weights = np.zeros(3, dtype=np.float64) + + # Calculate the weights + triangle_weights = _triangle_line_intersection( + triangle=polygon_xyz, point=point_xyz + ) + + # Using the results, store the weights + weights[0] = 1 - triangle_weights[1] - triangle_weights[2] + weights[1] = triangle_weights[1] + weights[2] = triangle_weights[2] + + # Since all the nodes of the triangle were used, return all 3 nodes + nodes = np.array([0, 1, 2]) + + return weights, nodes + + # If the polygon is a quadrilateral, instead use the `newton_quadrilateral` algorithm to calculate the weights + elif n == 4: + # Assign an empty array of size 4 to hold final weights + weights = np.zeros(4, dtype=np.float64) + + # Calculate the weights + quadrilateral_weights = _newton_quadrilateral( + quadrilateral=polygon_xyz, point=point_xyz + ) + + # Using the results, store the weights + weights[0] = ( + 1.0 + - quadrilateral_weights[0] + - quadrilateral_weights[1] + + quadrilateral_weights[0] * quadrilateral_weights[1] + ) + weights[1] = quadrilateral_weights[0] * (1.0 - quadrilateral_weights[1]) + weights[2] = quadrilateral_weights[0] * quadrilateral_weights[1] + weights[3] = quadrilateral_weights[1] * (1.0 - quadrilateral_weights[0]) + + # Since all the nodes of the quadrilateral were used, return all 4 nodes + nodes = np.array([0, 1, 2, 3]) + + return weights, nodes + + # For polygons of any other shape, break it into triangles, + # find the triangle containing the point, + # and use`triangle_line_intersection` algorithm to calculate the weights + else: + for i in range(0, n - 2): + # Get the three points of the triangle + node_0, node_1, node_2 = ( + polygon_xyz[0], + polygon_xyz[i + 1], + polygon_xyz[i + 2], + ) + + # Get the triangle in terms of its edges for the `point_in_face` check + face_edge = _get_cartesian_face_edge_nodes( + face_idx=0, + face_node_connectivity=np.array([[0, 1, 2]]), + n_edges_per_face=np.array([3]), + node_x=np.array([node_0[0], node_1[0], node_2[0]], dtype=np.float64), + node_y=np.array([node_0[1], node_1[1], node_2[1]], dtype=np.float64), + node_z=np.array([node_0[2], node_1[2], node_2[2]], dtype=np.float64), + ) + + # Check to see if the point lies within the current triangle + contains_point = _face_contains_point( + face_edge, + point_xyz, + ) + + # If the point is in the current triangle, get the weights for that triangle + if contains_point: + # Create an empty array of size 3 to hold weights + weights = np.zeros(3, dtype=np.float64) + + # Create the triangle + triangle = np.zeros((3, 3), dtype=np.float64) + triangle[0] = node_0 + triangle[1] = node_1 + triangle[2] = node_2 + + # Calculate the weights + triangle_weights = _triangle_line_intersection( + triangle=triangle, + point=point_xyz, + ) + + # Assign the weights based off the results + weights[0] = 1 - triangle_weights[1] - triangle_weights[2] + weights[1] = triangle_weights[1] + weights[2] = triangle_weights[2] + + # Assign the current nodes as the nodes to return + nodes = np.array([0, i + 1, i + 2]) + + return weights, nodes + + # If the point doesn't reside in the polygon, raise an error + raise ValueError("Point does not reside in polygon") + + +@njit(cache=True) +def _triangle_line_intersection(triangle, point, threshold=1e12): + """ + Compute interpolation coefficients for a point relative to a triangle. + + Parameters + ---------- + triangle: np.array + Cartesian coordinates for a triangle + point: np.array + Cartesian coordinates for a point within the triangle + threshold: np.array + Condition number threshold for warning + + Examples + -------- + Calculate the weights + >>> weights_t = _triangle_line_intersection(triangle=tri_xyz, point=point_xyz) + + Using the results gotten, store the weights + >>> weights_t[0] = 1 - triangle_weights[1] - triangle_weights[2] + >>> weights_t[1] = triangle_weights[1] + >>> weights_t[2] = triangle_weights[2] + + Returns + ------- + triangle_weights: np.array + The weights of each point in the triangle + """ + + # triangle: shape (3, 3), point: shape (3,) + node_0 = triangle[0] + node_1 = triangle[1] + node_2 = triangle[2] + + # Construct matrix for barycentric interpolation + v1 = node_1 - node_0 + v2 = node_2 - node_0 + v = point - node_0 + + # Construct the matrix (columns: v1, v2, point - node_0) + matrix = np.column_stack((point, v1, v2)) + + # Estimate condition number (max column-sum norm) + conditional_number = np.sum(np.abs(matrix), axis=0) + + # Compute inverse of matrix + det = np.linalg.det(matrix) + if np.abs(det) < 1e-12: + # Singular matrix; return NaNs + return np.full(3, np.nan) + + matrix_inv = np.linalg.inv(matrix) + matrix_inv_column_sum = np.sum(np.abs(matrix_inv), axis=0) + cond_est = np.max(conditional_number) * np.max(matrix_inv_column_sum) + + # Check conditioning + if cond_est > threshold: + # Still continue, but you might choose to return NaNs + pass + + # Compute triangle weights + triangle_weights = matrix_inv @ v + + return triangle_weights + + +@njit(cache=True) +def _newton_quadrilateral(quadrilateral, point, max_iterations=150): + """ + Compute interpolation coefficients for a point relative to a quadrilateral. + + Parameters + ---------- + quadrilateral: np.array + Cartesian coordinates for a quadrilateral + point: np.array + Cartesian coordinates for a point within the triangle + max_iterations: int + Maximum number of Newton iterations + + Examples + -------- + Calculate the weights + >>> quadrilateral_weights = _newton_quadrilateral( + ... quadrilateral=quad_xyz, point=point_xyz + ... ) + >>> weights_q = np.zeros(4, dtype=np.float64) + + Using the results gotten, store the weights + >>> weights_q[0] = ( + ... 1.0 + ... - quadrilateral_weights[0] + ... - quadrilateral_weights[1] + ... + quadrilateral_weights[0] * quadrilateral_weights[1] + ... ) + >>> weights_q[1] = quadrilateral_weights[0] * (1.0 - quadrilateral_weights[1]) + >>> weights_q[2] = quadrilateral_weights[0] * quadrilateral_weights[1] + >>> weights_q[3] = quadrilateral_weights[1] * (1.0 - quadrilateral_weights[0]) + + Returns + ------- + triangle_weights: np.array + The weights of the quadrilateral + + """ + # unpack nodes + node_0 = quadrilateral[0] + node_1 = quadrilateral[1] + node_2 = quadrilateral[2] + node_3 = quadrilateral[3] + + # precompute constant vectors + A = node_0 - node_1 + node_2 - node_3 + B = node_1 - node_0 + C = node_3 - node_0 + D = point + E = node_0 - point + + # unknowns [alpha, beta, gamma] + weights = np.zeros(3, dtype=quadrilateral.dtype) + + for _ in range(max_iterations): + dA = weights[0] + dB = weights[1] + dC = weights[2] + + value_at_x = np.empty(3, dtype=quadrilateral.dtype) + for i in range(3): + value_at_x[i] = dA * dB * A[i] + dA * B[i] + dB * C[i] + dC * D[i] + E[i] + + # check convergence + norm2 = 0.0 + for i in range(3): + norm2 += value_at_x[i] * value_at_x[i] + if math.sqrt(norm2) < 1e-15: + break + + # build Jacobian J[i,j] + J = np.empty((3, 3), dtype=quadrilateral.dtype) + for i in range(3): + J[i, 0] = A[i] * dB + B[i] + J[i, 1] = A[i] * dA + C[i] + J[i, 2] = D[i] + + # compute det(J) + det = ( + J[0, 0] * (J[1, 1] * J[2, 2] - J[1, 2] * J[2, 1]) + - J[0, 1] * (J[1, 0] * J[2, 2] - J[1, 2] * J[2, 0]) + + J[0, 2] * (J[1, 0] * J[2, 1] - J[1, 1] * J[2, 0]) + ) + + # if singular (or nearly), bail out + if abs(det) < MACHINE_EPSILON: + break + + # adjugate / det + invJ = np.empty((3, 3), dtype=quadrilateral.dtype) + invJ[0, 0] = (J[1, 1] * J[2, 2] - J[1, 2] * J[2, 1]) / det + invJ[0, 1] = -(J[0, 1] * J[2, 2] - J[0, 2] * J[2, 1]) / det + invJ[0, 2] = (J[0, 1] * J[1, 2] - J[0, 2] * J[1, 1]) / det + invJ[1, 0] = -(J[1, 0] * J[2, 2] - J[1, 2] * J[2, 0]) / det + invJ[1, 1] = (J[0, 0] * J[2, 2] - J[0, 2] * J[2, 0]) / det + invJ[1, 2] = -(J[0, 0] * J[1, 2] - J[0, 2] * J[1, 0]) / det + invJ[2, 0] = (J[1, 0] * J[2, 1] - J[1, 1] * J[2, 0]) / det + invJ[2, 1] = -(J[0, 0] * J[2, 1] - J[0, 1] * J[2, 0]) / det + invJ[2, 2] = (J[0, 0] * J[1, 1] - J[0, 1] * J[1, 0]) / det + + # delta = invJ @ value_at_x + delta = np.empty(3, dtype=quadrilateral.dtype) + for i in range(3): + delta[i] = ( + invJ[i, 0] * value_at_x[0] + + invJ[i, 1] * value_at_x[1] + + invJ[i, 2] * value_at_x[2] + ) + + # Newton update + for i in range(3): + weights[i] -= delta[i] + + return weights diff --git a/uxarray/grid/grid.py b/uxarray/grid/grid.py index f8ec49aa2..3cc3a8119 100644 --- a/uxarray/grid/grid.py +++ b/uxarray/grid/grid.py @@ -2324,7 +2324,7 @@ def to_linecollection( return line_collection - def get_dual(self): + def get_dual(self, check_duplicate_nodes: bool = False): """Compute the dual for a grid, which constructs a new grid centered around the nodes, where the nodes of the primal become the face centers of the dual, and the face centers of the primal become the nodes of the @@ -2336,8 +2336,10 @@ def get_dual(self): Dual Mesh Grid constructed """ - if _check_duplicate_nodes_indices(self): - raise RuntimeError("Duplicate nodes found, cannot construct dual") + if check_duplicate_nodes: + if _check_duplicate_nodes_indices(self): + # TODO: This is very slow + raise RuntimeError("Duplicate nodes found, cannot construct dual") # Get dual mesh node face connectivity dual_node_face_conn = construct_dual(grid=self) diff --git a/uxarray/grid/neighbors.py b/uxarray/grid/neighbors.py index d836756a2..206cdbb95 100644 --- a/uxarray/grid/neighbors.py +++ b/uxarray/grid/neighbors.py @@ -995,7 +995,6 @@ def _barycentric_coordinates(nodes, point): a2 = max(_triangle_area(point, vi, vi1), ERROR_TOLERANCE) sum_wi += a0 / (a1 * a2) w.append(a0 / (a1 * a2)) - barycentric_coords = [w_i / sum_wi for w_i in w] return barycentric_coords diff --git a/uxarray/plot/utils.py b/uxarray/plot/utils.py index eaa20bf27..e97919d6c 100644 --- a/uxarray/plot/utils.py +++ b/uxarray/plot/utils.py @@ -26,7 +26,6 @@ def assign(self, backend: str): raise ValueError( f"Unsupported backend. Expected one of ['bokeh', 'matplotlib'], but received {backend}" ) - print(hv.Store.current_backend) if backend is not None and backend != hv.Store.current_backend: # only call hv.extension if it needs to be changed hv.extension(backend) diff --git a/uxarray/remap/accessor.py b/uxarray/remap/accessor.py index 9bf61ccb5..ebf74ffa4 100644 --- a/uxarray/remap/accessor.py +++ b/uxarray/remap/accessor.py @@ -7,6 +7,7 @@ from uxarray.core.dataset import UxDataset from uxarray.grid.grid import Grid +from uxarray.remap.bilinear import _bilinear from uxarray.remap.inverse_distance_weighted import _inverse_distance_weighted_remap from uxarray.remap.nearest_neighbor import _nearest_neighbor_remap @@ -88,3 +89,24 @@ def inverse_distance_weighted( return _inverse_distance_weighted_remap( self.ux_obj, destination_grid, remap_to, power, k ) + + def bilinear( + self, destination_grid: Grid, remap_to: str = "faces", **kwargs + ) -> UxDataArray | UxDataset: + """ + Perform bilinear remapping. + + Parameters + --------- + destination_grid : Grid + Destination Grid for remapping + remap_to : {'nodes', 'edges', 'faces'}, default='faces' + Which grid element receives the remapped values. + + Returns + ------- + UxDataArray or UxDataset + A new object with data mapped onto `destination_grid`. + """ + + return _bilinear(self.ux_obj, destination_grid, remap_to) diff --git a/uxarray/remap/bilinear.py b/uxarray/remap/bilinear.py new file mode 100644 index 000000000..7655a55c2 --- /dev/null +++ b/uxarray/remap/bilinear.py @@ -0,0 +1,202 @@ +from __future__ import annotations + +from typing import TYPE_CHECKING + +import numpy as np +import xarray as xr +from numba import njit + +if TYPE_CHECKING: + from uxarray.core.dataarray import UxDataArray + from uxarray.core.dataset import UxDataset + +from uxarray.constants import ERROR_TOLERANCE +from uxarray.grid import Grid +from uxarray.grid.geometry import barycentric_coordinates_cartesian + +from .utils import ( + KDTREE_DIM_MAP, + LABEL_TO_COORD, + SPATIAL_DIMS, + _assert_dimension, + _construct_remapped_ds, + _get_remap_dims, + _prepare_points, + _to_dataset, +) + + +def _bilinear( + source: UxDataArray | UxDataset, + destination_grid: Grid, + destination_dim: str = "n_face", +) -> np.ndarray: + """Bilinear Remapping between two grids, mapping data that resides on the + corner nodes, edge centers, or face centers on the source grid to the + corner nodes, edge centers, or face centers of the destination grid. + + Parameters + --------- + source_uxda : UxDataArray + Source UxDataArray + remap_to : str, default="nodes" + Location of where to map data, either "nodes", "edge centers", or "face centers" + + Returns + ------- + destination_data : np.ndarray + Data mapped to destination grid + """ + + # ensure array is a np.ndarray + _assert_dimension(destination_dim) + + # Ensure the destination grid is normalized + destination_grid.normalize_cartesian_coordinates() + + # Perform remapping on a UxDataset + ds, is_da, name = _to_dataset(source) + + # Determine which spatial dimensions we need to remap + dims_to_remap = _get_remap_dims(ds) + + for src_dim in dims_to_remap: + if src_dim != "n_face": + raise ValueError( + "Bilinear remapping is not supported for non-face centered variables" + ) + + # Construct dual for searching + dual = source.uxgrid.get_dual() + + # get destination coordinate pairs + point_xyz = _prepare_points(destination_grid, destination_dim) + + weights, indices = _barycentric_weights( + point_xyz=point_xyz, + dual=dual, + data_size=getattr(destination_grid, f"n_{KDTREE_DIM_MAP[destination_dim]}"), + source_grid=ds.uxgrid, + ) + + remapped_vars = {} + for name, da in ds.data_vars.items(): + spatial = set(da.dims) & SPATIAL_DIMS + if spatial: + source_dim = spatial.pop() + inds, w = indices, weights + + # pack indices & weights into tiny DataArrays: + indexer = xr.DataArray(inds, dims=[LABEL_TO_COORD[destination_dim], "k"]) + weight_da = xr.DataArray(w, dims=[LABEL_TO_COORD[destination_dim], "k"]) + + # gather the k neighbor values: + da_k = da.isel({source_dim: indexer}, ignore_grid=True) + + # weighted sum over the small "k" axis: + da_idw = (da_k * weight_da).sum(dim="k") + + # attach the new grid + da_idw.uxgrid = destination_grid + remapped_vars[name] = da_idw + else: + remapped_vars[name] = da + + ds_remapped = _construct_remapped_ds( + source, remapped_vars, destination_grid, destination_dim + ) + + return ds_remapped[name] if is_da else ds_remapped + + +def _barycentric_weights(point_xyz, dual, data_size, source_grid): + """Get barycentric weights and source face indices for each destination point.""" + all_weights = np.zeros((data_size, 4), dtype=np.float64) + all_indices = np.zeros((data_size, 4), dtype=int) + + # Query dual grid + face_indices, hits = dual.get_faces_containing_point(points=point_xyz) + + # Handle fallback cases (hits == 0) + fallback_mask = hits == 0 + fallback_idxs = np.where(fallback_mask)[0] + for i in fallback_idxs: + cur_inds, counts = source_grid.get_faces_containing_point(points=point_xyz[i]) + if counts == 0: + continue + all_weights[i, 0] = 1.0 + all_indices[i, 0] = int(cur_inds[0]) + + # Prepare args for the Numba function + valid_idxs = np.where(hits != 0)[0] + + # Call numba function to find weights + _calculate_weights( + valid_idxs, + point_xyz, + face_indices, + dual.node_x.values, + dual.node_y.values, + dual.node_z.values, + dual.face_node_connectivity.values, + dual.n_nodes_per_face.values, + all_weights, + all_indices, + ) + + return all_weights, all_indices + + +@njit(cache=True) +def _calculate_weights( + valid_idxs, + point_xyz, + face_indices, + x, + y, + z, + face_node_conn, + n_nodes_per_face, + all_weights, + all_indices, +): + for idx in valid_idxs: + fidx = int(face_indices[idx, 0]) + nverts = int(n_nodes_per_face[fidx]) + + polygon_xyz = np.zeros((nverts, 3), dtype=np.float64) + polygon_face_indices = np.empty(nverts, dtype=np.int32) + for j in range(nverts): + node = face_node_conn[fidx, j] + polygon_xyz[j, 0] = x[node] + polygon_xyz[j, 1] = y[node] + polygon_xyz[j, 2] = z[node] + polygon_face_indices[j] = node + + # snap check + match = _find_matching_node_index(polygon_xyz, point_xyz[idx]) + if match[0] != -1: + all_weights[idx, 0] = 1.0 + all_indices[idx, 0] = polygon_face_indices[match[0]] + continue + + weights, node_idxs = barycentric_coordinates_cartesian( + polygon_xyz, point_xyz[idx] + ) + for k in range(len(weights)): + all_weights[idx, k] = weights[k] + all_indices[idx, k] = polygon_face_indices[node_idxs[k]] + + +@njit(cache=True) +def _find_matching_node_index(nodes, point, tolerance=ERROR_TOLERANCE): + for i in range(nodes.shape[0]): + match = True + for j in range(3): # Compare each coordinate + diff = abs(nodes[i, j] - point[j]) + if diff > tolerance + tolerance * abs(point[j]): + match = False + break + if match: + return np.array([i]) # Return first matching index + return np.array([-1]) # Not found