diff --git a/qcore/coordinates.py b/qcore/coordinates.py index f28626b1..f3f586b2 100644 --- a/qcore/coordinates.py +++ b/qcore/coordinates.py @@ -17,6 +17,8 @@ [0]: https://www.linz.govt.nz/guidance/geodetic-system/coordinate-systems-used-new-zealand/projections/new-zealand-transverse-mercator-2000-nztm2000 """ +from typing import Union + import numpy as np import pyproj @@ -84,3 +86,33 @@ def nztm_to_wgs_depth(nztm_coordinates: np.ndarray) -> np.ndarray: array([[-36.8509, 174.7645, 100], [-41.2924, 174.7787, 100]]) """ return np.array(_NZTM2WGS.transform(*nztm_coordinates.T)).T + + +def distance_between_wgs_depth_coordinates( + point_a: np.ndarray, point_b: np.ndarray +) -> Union[float, np.ndarray]: + """Return the distance between two points in lat, lon, depth format. + + Valid only for points that can be converted into NZTM format. + + Parameters + ---------- + point_a : np.ndarray + The first point (lat, lon, optionally depth). May have shape + (2,), (3,), (n, 2), (n, 3). + + point_b : np.ndarray + The second point (lat, lon, optionally depth). May have shape + (2, ), (3,), (n, 2), (n, 3) + + Returns + ------- + float or np.ndarray + The distance (in metres) between point_a and point_b. Will + return an array of floats if input contains multiple points + """ + if len(point_a.shape) > 1: + return np.linalg.norm( + wgs_depth_to_nztm(point_a) - wgs_depth_to_nztm(point_b), axis=1 + ) + return np.linalg.norm(wgs_depth_to_nztm(point_a) - wgs_depth_to_nztm(point_b))