Fixing bradcast bug

mlair/helpers/geofunctions.py

@@ -80,3 +80,74 @@ def haversine_dist(lat1: xr_int_float, lon1: xr_int_float,
 
     return dist
 
+
+def bearing_angle(lat1: xr_int_float, lon1: xr_int_float,
+                   lat2: xr_int_float, lon2: xr_int_float,
+                   to_radians: bool = True, return_deg: bool = True) -> xr.DataArray:
+    """
+    Calculate initial bearing angle (forward azimuth) between multiple points.
+    Implemented by following
+    http://www.movable-type.co.uk/scripts/latlong.html
+
+    θ = atan2( sin Δλ ⋅ cos φ2 , cos φ1 ⋅ sin φ2 − sin φ1 ⋅ cos φ2 ⋅ cos Δλ )
+    where 	φ1,λ1 is the start point, φ2,λ2 the end point (Δλ is the difference in longitude)
+
+    :param lat1: Latitude(-s) of first location
+    :param lon1: Longitude(-s) of first location
+    :param lat2: Latitude(-s) of second location
+    :param lon2: Longitude(-s) of second location
+    :param to_radians: Flag if conversion from degree to radiant is required for input
+    :param return_deg: Flag if conversion from radiant to degree is required for return value
+    :return: initial bearing angle (forward azimuth) in degree
+
+    """
+    lat1_orig = lat1
+    lon1_orig = lon1
+    lat2_orig = lat2
+    lon2_orig = lon2
+
+    if to_radians:
+        (lat1, lon1), (lat2, lon2) = deg2rad_all_points(lat1, lon1, lat2, lon2)
+
+    lat1 = convert2xrda(lat1, use_1d_default=True)
+    lon1 = convert2xrda(lon1, use_1d_default=True)
+    lat2 = convert2xrda(lat2, use_1d_default=True)
+    lon2 = convert2xrda(lon2, use_1d_default=True)
+
+    assert lat1.shape == lon1.shape
+    assert lat2.shape == lon2.shape
+    assert isinstance(lat1, xr.DataArray)
+    assert isinstance(lon1, xr.DataArray)
+    assert isinstance(lat2, xr.DataArray)
+    assert isinstance(lon2, xr.DataArray)
+
+    # broadcast lats and lons to calculate distances in a vectorized manner.
+    lat1, lat2 = xr.broadcast(lat1, lat2)
+    lon1, lon2 = xr.broadcast(lon1, lon2)
+
+    # from https://www.igismap.com/formula-to-find-bearing-or-heading-angle-between-two-points-latitude-longitude/
+    dlon = lon1 - lon2
+    x = da.cos(lat2) * da.sin(dlon)
+    y = da.cos(lat1) * da.sin(lat2) - da.sin(lat1) * da.cos(lat2) * da.cos(dlon)
+    theta_rad = da.arctan2(x, y)
+
+    if return_deg is True:
+        theta_deg = (theta_rad * 180 / np.pi + 360) % 360
+        # theta_deg = da.rad2deg(theta_rad)
+        return theta_deg
+    else:
+        return theta_rad
+
+
+if __name__ == '__main__':
+    kansas_StLouis = bearing_angle(lat1=39.099912, lon1=-94.581213,
+                                   lat2=38.627089, lon2=-90.200203,
+                                   to_radians=True, return_deg=True)
+    London_Rio = bearing_angle(lat1=51.5, lon1=0,
+                               lat2=-22.97, lon2=-43.18,
+                               to_radians=True, return_deg=True)
+    t1_1 = bearing_angle(lat1=0., lon1=0., lat2=50., lon2=0., to_radians=True, return_deg=True)
+    t1_2 = bearing_angle(lat1=50., lon1=0., lat2=0., lon2=0., to_radians=True, return_deg=True)
+    t1_exp = 0.
+    t2 = bearing_angle(lat1=0., lon1=0., lat2=0., lon2=90., to_radians=True, return_deg=True)
+    t2_exp = 90.