UXARRAY / uxarray

Xarray-styled package for reading and directly operating on unstructured grid datasets following UGRID conventions
https://uxarray.readthedocs.io/
Apache License 2.0
142 stars 31 forks source link

Vectorize Coordinate Conversion Functions #748

Closed philipc2 closed 2 months ago

philipc2 commented 3 months ago

Closes #346

Overview

Usage Example

import uxarray as ux 
import numpy as np
from uxarray.grid.coordinates import _lonlat_rad_to_xyz, _xyz_to_lonlat_rad

# assume an initialized grid
uxgrid = ux.Grid()

lon_rad = np.deg2rad(uxgrid.node_lon)
lat_rad = np.deg2rad(uxgrid.node_lat)

# conversion is vectorized across the entire arrays of coordinates
x, y, z = _lonlat_rad_to_xyz(lon_rad, lat_rad)
philipc2 commented 2 months ago

Hi @hongyuchen1030

I'm almost done with vectorizing all the functions. The only issue I'm running into is one of the point inside polygon tests is now failing (test_pole_point_inside_polygon_from_vertice_cross)

I stepped through the function, and was able to determine that the issue somewhere in this portion of the _pole_point_inside_polygon method.

        return (
            _check_intersection(ref_edge_north, north_edges)
            + _check_intersection(ref_edge_south, south_edges)
        ) % 2 != 0

With the changes I've made here, the sum of the intersections is 2, however prior to the change the number of intersections is 1.

Do you know what might be causing this? Could it be a precision issue? Every other test is passing except for this one.

hongyuchen1030 commented 2 months ago

Hi @hongyuchen1030

I'm almost done with vectorizing all the functions. The only issue I'm running into is one of the point inside polygon tests is now failing (test_pole_point_inside_polygon_from_vertice_cross)

I stepped through the function, and was able to determine that the issue somewhere in this portion of the _pole_point_inside_polygon method.

        return (
            _check_intersection(ref_edge_north, north_edges)
            + _check_intersection(ref_edge_south, south_edges)
        ) % 2 != 0

With the changes I've made here, the sum of the intersections is 2, however prior to the change the number of intersections is 1.

Do you know what might be causing this? Could it be a precision issue? Every other test is passing except for this one.

Thank you very much for your work! I feel like the problem might be a combination of precision and the coordinate conversion style for this test case. I tried to change our error tolerance to 10-3 but it's still failing, so I don't think that's the case (at least not the only factor)

So the test case face tests a very extreme case: The face goes across the equator while still encompassing the North Pole. The face may be referred to as the "bigger face" on the sphere, while it should refer to the "smaller face". I can look into this tomorrow and see what I can find out!

But if you feel this test case is unrealistically complicate for our users, you can feel free to remove it!

philipc2 commented 2 months ago

Hi @hongyuchen1030 I'm almost done with vectorizing all the functions. The only issue I'm running into is one of the point inside polygon tests is now failing (test_pole_point_inside_polygon_from_vertice_cross) I stepped through the function, and was able to determine that the issue somewhere in this portion of the _pole_point_inside_polygon method.

        return (
            _check_intersection(ref_edge_north, north_edges)
            + _check_intersection(ref_edge_south, south_edges)
        ) % 2 != 0

With the changes I've made here, the sum of the intersections is 2, however prior to the change the number of intersections is 1. Do you know what might be causing this? Could it be a precision issue? Every other test is passing except for this one.

Thank you very much for your work! I feel like the problem might be a combination of precision and the coordinate conversion style for this test case. I tried to change our error tolerance to 10-3 but it's still failing, so I don't think that's the case (at least not the only factor)

So the test case face tests a very extreme case: The face goes across the equator while still encompassing the North Pole. The face may be referred to as the "bigger face" on the sphere, while it should refer to the "smaller face". I can look into this tomorrow and see what I can find out!

But if you feel this test case is unrealistically complicate for our users, you can feel free to remove it!

Thanks for the response! If you'd like to look into it, that would be greatly appreciated. I'll be requesting reviews soon, so if we don't get to fixing this test, we can comment it out and put together an issue for this edge case.

hongyuchen1030 commented 2 months ago

Hi @philipc2 , I think I have figured out why it will break, it's related to the coordinates conversion style

in my current node_xyz_to_lonlat_rad it has the following lines, you can tell it will make the pole point presentation uniform in a way such that northpole=[0.0, 0.5 np.pi], southpole=[0.0, -0.5 np.pi]. However, there's no such uniform function in you branch. All my functions is built on top of the uniform presentation of the north pole and south pole.

So when in your branch, one of the south pole was presented as [3.14,-0.5 * np.pi ], and it cannot be handled by my function.

    if np.absolute(dz) < (1.0 - ERROR_TOLERANCE):
        d_lon_rad = math.atan2(dy, dx)
        d_lat_rad = np.arcsin(dz)

        if d_lon_rad < 0.0:
            d_lon_rad += 2.0 * np.pi
    elif dz > 0.0:
        d_lon_rad = 0.0
        d_lat_rad = 0.5 * np.pi
    else:
        d_lon_rad = 0.0
        d_lat_rad = -0.5 * np.pi

After I add an extra lon = 0.0 here below in your _xyz_to_lonlat_rad, it passes. I

    if scalar:
        if np.abs(z) > 1.0 - ERROR_TOLERANCE:
            lat = np.sign(z) * np.pi / 2
            lon = 0.0

I feel like this will be a good test case for the new coordinates conversion function.

philipc2 commented 2 months ago

Hi @philipc2 , I think I have figured out why it will break, it's related to the coordinates conversion style

in my current node_xyz_to_lonlat_rad it has the following lines, you can tell it will make the pole point presentation uniform in a way such that northpole=[0.0, 0.5 np.pi], southpole=[0.0, -0.5 np.pi]. However, there's no such uniform function in you branch. All my functions is built on top of the uniform presentation of the north pole and south pole.

So when in your branch, one of the south pole was presented as [3.14,-0.5 * np.pi ], and it cannot be handled by my function.

    if np.absolute(dz) < (1.0 - ERROR_TOLERANCE):
        d_lon_rad = math.atan2(dy, dx)
        d_lat_rad = np.arcsin(dz)

        if d_lon_rad < 0.0:
            d_lon_rad += 2.0 * np.pi
    elif dz > 0.0:
        d_lon_rad = 0.0
        d_lat_rad = 0.5 * np.pi
    else:
        d_lon_rad = 0.0
        d_lat_rad = -0.5 * np.pi

After I add an extra lon = 0.0 here below in your _xyz_to_lonlat_rad, it passes. I

    if scalar:
        if np.abs(z) > 1.0 - ERROR_TOLERANCE:
            lat = np.sign(z) * np.pi / 2
            lon = 0.0

I feel like this will be a good test case for the new coordinates conversion function.

Thank you for the explanation! It's working for me locally now too.

philipc2 commented 2 months ago

For combability, the coordinate conversion functions here are all written in pure Numpy, no longer depending on Numba in certain cases.

The main reason for this is that many Numpy functions (such as linalg.norm) when used with Numba don't support optional parameters.

Once Rachel joins us this summer, we can do some further Numba exploration as was described in #202

philipc2 commented 2 months ago

Some preliminary performance analysis of running the code vectorized compared to iteratively.

image

image

About a 150x improvement for Cartesian to Spherical and 50x for Spherical to Cartesian.

erogluorhan commented 2 months ago

Performance improvements sound great; I am excited to review it as soon as I can (aiming at today)

rajeeja commented 2 months ago

For combability, the coordinate conversion functions here are all written in pure Numpy, no longer depending on Numba in certain cases.

The main reason for this is that many Numpy functions (such as linalg.norm) when used with Numba don't support optional parameters.

Once Rachel joins us this summer, we can do some further Numba exploration as was described in #202

Good work.

After some profiling and understanding the number of calls to certain functions - I believe in introducing some numba decorator again, that were removed from this PR - specifically the one's in area calculation; it might speedup things.

Getting things compatible and decorating all function->function->function with Numba can be a pain. It is known that Numba struggles with functions that use dynamic data structures.

rajeeja commented 2 months ago

Some preliminary performance analysis of running the code vectorized compared to iteratively.

image

image

About a 150x improvement for Cartesian to Spherical and 50x for Spherical to Cartesian.

Is my understanding correct: Removing numba and using numpy vectorization caused this?

rajeeja commented 2 months ago

There is some numba in grid/ geometry, neighbors and connectivity - all mostly non-nested functions. A note for #202

philipc2 commented 2 months ago

There is some numba in grid/ geometry, neighbors and connectivity - all mostly non-nested functions. A note for #202

Thanks for that! Can you reference that in the issue directly too?

philipc2 commented 2 months ago

Some preliminary performance analysis of running the code vectorized compared to iteratively. image image About a 150x improvement for Cartesian to Spherical and 50x for Spherical to Cartesian.

Is my understanding correct: Removing numba and using numpy vectorization caused this?

Correct, the vectorization timings can be achieved when running this function on the entire arrays, as opposed to on each value itteratievely.

x, y, z = _lonlat_rad_to_xyz(lon_rad, lat_rad)
philipc2 commented 2 months ago

For combability, the coordinate conversion functions here are all written in pure Numpy, no longer depending on Numba in certain cases. The main reason for this is that many Numpy functions (such as linalg.norm) when used with Numba don't support optional parameters. Once Rachel joins us this summer, we can do some further Numba exploration as was described in #202

Good work.

After some profiling and understanding the number of calls to certain functions - I believe in introducing some numba decorator again, that were removed from this PR - specifically the one's in area calculation; it might speedup things.

Getting things compatible and decorating all function->function->function with Numba can be a pain. It is known that Numba struggles with functions that use dynamic data structures.

Yeah, I'm hoping we can explore this more this summer, since there is still room for improvement with the performance.