vectorized

Author: JaskRendix

pyorbital/orbital.py

@@ -92,60 +92,72 @@ class OrbitalError(Exception):
     pass
 
 
+def ecef_to_topocentric(rx, ry, rz, lat, lon, theta):
+    """Convert ECEF vector to topocentric-horizon coordinates."""
+    sin_lat = np.sin(lat)
+    cos_lat = np.cos(lat)
+    sin_theta = np.sin(theta)
+    cos_theta = np.cos(theta)
+
+    # Transform ECEF coordinates to topocentric-horizon coordinates
+    top_s = sin_lat * cos_theta * rx + sin_lat * sin_theta * ry - cos_lat * rz
+    top_e = -sin_theta * rx + cos_theta * ry
+    top_z = cos_lat * cos_theta * rx + cos_lat * sin_theta * ry + sin_lat * rz
+
+    return top_s, top_e, top_z
+
+def compute_azimuth_elevation(top_s, top_e, top_z, rg_):
+    """Compute azimuth and elevation from topocentric coordinates."""
+
+    # Azimuth is undefined when elevation is 90 degrees, 180 (pi) will be returned.
+    az_ = np.arctan2(-top_e, top_s) + np.pi
+    az_ = np.mod(az_, 2 * np.pi)  # Needed on some platforms
+
+    # Due to rounding top_z can be larger than rg_ (when el_ ~ 90).
+    top_z_divided_by_rg_ = np.clip(top_z / rg_, -1, 1)
+    el_ = np.arcsin(top_z_divided_by_rg_)
+
+    return np.rad2deg(az_), np.rad2deg(el_)
+
 def get_observer_look(sat_lon, sat_lat, sat_alt, utc_time, lon, lat, alt):
-    """Calculate observers look angle to a satellite.
+    """Calculate observer's look angle to a satellite.
 
     http://celestrak.com/columns/v02n02/
 
-    :utc_time: Observation time (datetime object)
-    :lon: Longitude of observer position on ground in degrees east
-    :lat: Latitude of observer position on ground in degrees north
-    :alt: Altitude above sea-level (geoid) of observer position on ground in km
-
-    :return: (Azimuth, Elevation)
+    :param sat_lon: Satellite longitude in degrees east
+    :param sat_lat: Satellite latitude in degrees north
+    :param sat_alt: Satellite altitude in km
+    :param utc_time: Observation time (datetime object)
+    :param lon: Observer longitude in degrees east
+    :param lat: Observer latitude in degrees north
+    :param alt: Observer altitude in km
+    :return: (Azimuth, Elevation) in degrees
     """
-    (pos_x, pos_y, pos_z), (vel_x, vel_y, vel_z) = astronomy.observer_position(
-        utc_time, sat_lon, sat_lat, sat_alt)
-
-    (opos_x, opos_y, opos_z), (ovel_x, ovel_y, ovel_z) = \
-        astronomy.observer_position(utc_time, lon, lat, alt)
+    # Get satellite and observer ECEF positions
+    (pos_x, pos_y, pos_z), _ = astronomy.observer_position(utc_time, sat_lon, sat_lat, sat_alt)
+    (opos_x, opos_y, opos_z), _ = astronomy.observer_position(utc_time, lon, lat, alt)
 
+    # Convert observer coordinates to radians
     lon = np.deg2rad(lon)
     lat = np.deg2rad(lat)
 
+    # Compute local sidereal time
     theta = (astronomy.gmst(utc_time) + lon) % (2 * np.pi)
 
+    # Vector from observer to satellite
     rx = pos_x - opos_x
     ry = pos_y - opos_y
     rz = pos_z - opos_z
-
-    sin_lat = np.sin(lat)
-    cos_lat = np.cos(lat)
-    sin_theta = np.sin(theta)
-    cos_theta = np.cos(theta)
-
-    top_s = sin_lat * cos_theta * rx + \
-        sin_lat * sin_theta * ry - cos_lat * rz
-    top_e = -sin_theta * rx + cos_theta * ry
-    top_z = cos_lat * cos_theta * rx + \
-        cos_lat * sin_theta * ry + sin_lat * rz
-
-    # Azimuth is undefined when elevation is 90 degrees, 180 (pi) will be returned.
-    az_ = np.arctan2(-top_e, top_s) + np.pi
-    az_ = np.mod(az_, 2 * np.pi)  # Needed on some platforms
-
     rg_ = np.sqrt(rx * rx + ry * ry + rz * rz)
 
-    top_z_divided_by_rg_ = top_z / rg_
+    # Convert to topocentric coordinates
+    top_s, top_e, top_z = ecef_to_topocentric(rx, ry, rz, lat, lon, theta)
 
-    # Due to rounding top_z can be larger than rg_ (when el_ ~ 90).
-    top_z_divided_by_rg_ = top_z_divided_by_rg_.clip(max=1)
-    el_ = np.arcsin(top_z_divided_by_rg_)
-
-    return np.rad2deg(az_), np.rad2deg(el_)
+    # Compute azimuth and elevation
+    return compute_azimuth_elevation(top_s, top_e, top_z, rg_)
 
 
-class Orbital(object):
+class Orbital:
     """Class for orbital computations.
 
     The *satellite* parameter is the name of the satellite to work on and is
@@ -336,7 +348,7 @@ def get_orbit_number(self, utc_time, tbus_style=False, as_float=False):
 
         return orbit
 
-    def get_next_passes(self, utc_time, length, lon, lat, alt, tol=0.001, horizon=0):
+    def get_next_passes(self, utc_time: dt.datetime, length: float, lon: float, lat: float, alt: float, tol: float = 0.001, horizon: float = 0) -> list[tuple[dt.datetime, dt.datetime, dt.datetime]]:
         """Calculate passes for the next hours for a given start time and a given observer.
 
         Original by Martin.
@@ -522,23 +534,61 @@ def _elevation_inv(self, utc_time, lon, lat, alt, horizon, minutes):
         """Compute the inverse of elevation."""
         return -self._elevation(utc_time, lon, lat, alt, horizon, minutes)
 
+    def get_lonlatalt_vectorized(self, utc_times):
+        """Vectorized version of get_lonlatalt for multiple times."""
+        if isinstance(utc_times, str) or not hasattr(utc_times, '__iter__'):
+            raise TypeError("Input must be a datetime or iterable of datetimes.")
+        if utc_times is None or len(utc_times) == 0:
+            return np.array([]), np.array([]), np.array([])
+        if not isinstance(utc_times, (list, np.ndarray)):
+            utc_times = [utc_times]
+        results = [self.get_lonlatalt(t) for t in utc_times]
+        lon, lat, alt = zip(*results)
+        return np.array(lon), np.array(lat), np.array(alt)
+
+    def get_observer_look_vectorized(self, utc_times, lon, lat, alt):
+        """Vectorized observer look angles."""
+        if isinstance(utc_times, str) or not hasattr(utc_times, '__iter__'):
+            raise TypeError("Input must be a datetime or iterable of datetimes.")
+        if utc_times is None or len(utc_times) == 0:
+            return np.array([]), np.array([])
+        if not isinstance(utc_times, (list, np.ndarray)):
+            utc_times = [utc_times]
+        results = [self.get_observer_look(t, lon, lat, alt) for t in utc_times]
+        az, el = zip(*results)
+        return np.array(az), np.array(el)
+
+    def get_orbit_number_vectorized(self, utc_times, tbus_style=False):
+        """Vectorized orbit number calculation."""
+        if isinstance(utc_times, str) or not hasattr(utc_times, '__iter__'):
+            raise TypeError("Input must be a datetime or iterable of datetimes.")
+        if utc_times is None or len(utc_times) == 0:
+            return []
+        if not isinstance(utc_times, (list, np.ndarray)):
+            utc_times = [utc_times]
+        return [self.get_orbit_number(t, tbus_style=tbus_style) for t in utc_times]
+
 
 def _get_root(fun, start, end, tol=0.01):
-    """Root finding scheme."""
-    x_0 = end
-    x_1 = start
-    fx_0 = fun(end)
-    fx_1 = fun(start)
+    """Find a root of `fun` in [start, end] using Brent's method with tolerance `tol`."""
+    x_0 = float(end)
+    x_1 = float(start)
+    fx_0 = fun(x_0)
+    fx_1 = fun(x_1)
+
+    if fx_0 * fx_1 > 0:
+        raise ValueError("Function values at interval endpoints must have opposite signs.")
+
+    # Swap for better convergence
     if abs(fx_0) < abs(fx_1):
         fx_0, fx_1 = fx_1, fx_0
         x_0, x_1 = x_1, x_0
 
-    x_n = optimize.brentq(fun, x_0, x_1)
-    return x_n
+    return optimize.brentq(fun, x_0, x_1, xtol=tol, rtol=tol)
 
 
-def _get_max_parab(fun, start, end, tol=0.01):
-    """Successive parabolic interpolation."""
+def _get_max_parab(fun, start, end, tol=0.01, max_iter=50):
+    """Find peak of `fun` in [start, end] using parabolic interpolation with fallback."""
     a = float(start)
     c = float(end)
     b = (a + c) / 2.0
@@ -547,25 +597,35 @@ def _get_max_parab(fun, start, end, tol=0.01):
     f_b = fun(b)
     f_c = fun(c)
 
-    x = b
-    with np.errstate(invalid="raise"):
-        while True:
-            try:
-                x = x - 0.5 * (((b - a) ** 2 * (f_b - f_c)
-                                - (b - c) ** 2 * (f_b - f_a)) /
-                               ((b - a) * (f_b - f_c) - (b - c) * (f_b - f_a)))
-            except FloatingPointError:
-                return b
-            if abs(b - x) <= tol:
-                return x
-            f_x = fun(x)
-            # sometimes the estimation diverges... return best guess
-            if f_x > f_b:
-                logger.info("Parabolic interpolation did not converge, returning best guess so far.")
-                return b
-
-            a, b, c = (a + x) / 2.0, x, (x + c) / 2.0
-            f_a, f_b, f_c = fun(a), f_x, fun(c)
+    # Handle flat or symmetric cases
+    if f_a == f_b == f_c:
+        return b
+
+    for _ in range(max_iter):
+        denom = ((b - a) * (f_b - f_c) - (b - c) * (f_b - f_a))
+        if denom == 0:
+            return b  # fallback
+
+        try:
+            x = b - 0.5 * (((b - a)**2 * (f_b - f_c) - (b - c)**2 * (f_b - f_a)) / denom)
+        except ZeroDivisionError:
+            return b
+
+        if abs(b - x) <= tol:
+            return x
+
+        f_x = fun(x)
+
+        # Divergence or invalid result
+        if f_x > f_b or not (a <= x <= c):
+            return b
+
+        # Update bracket
+        a, b, c = (a + x) / 2.0, x, (x + c) / 2.0
+        f_a, f_b, f_c = fun(a), f_x, fun(c)
+
+    # Max iterations reached
+    return b
 
 
 class OrbitElements(object):
@@ -1271,17 +1331,20 @@ def kep2xyz(kep):
 
 
 if __name__ == "__main__":
-    obs_lon, obs_lat = np.deg2rad((12.4143, 55.9065))
-    obs_alt = 0.02
+    # Observer's location in degrees
+    obs_lon, obs_lat = 12.4143, 55.9065
+    obs_alt = 0.02  # Altitude in km
     o = Orbital(satellite="METOP-B")
 
     t_start = dt.datetime.now()
-    t_stop = t_start + dt.timedelta(minutes=20)
-    t = t_start
-    while t < t_stop:
-        t += dt.timedelta(seconds=15)
-        lon, lat, alt = o.get_lonlatalt(t)
-        lon, lat = np.rad2deg((lon, lat))
-        az, el = o.get_observer_look(t, obs_lon, obs_lat, obs_alt)
-        ob = o.get_orbit_number(t, tbus_style=True)
-        print(az, el, ob)
+    t_end = t_start + dt.timedelta(minutes=20)
+    time_step = 15  # seconds
+    times = [t_start + dt.timedelta(seconds=i * time_step)
+             for i in range(int((t_end - t_start).total_seconds() // time_step))]
+
+    lon, lat, alt = o.get_lonlatalt_vectorized(times)
+    az, el = o.get_observer_look_vectorized(times, obs_lon, obs_lat, obs_alt)
+    ob = o.get_orbit_number_vectorized(times, tbus_style=True)
+
+    for i, t in enumerate(times):
+        print(f"Time: {t.strftime('%H:%M:%S')} | Az: {az[i]:.2f}° | El: {el[i]:.2f}° | Orbit: {ob[i]}")