laikad: calc_pos_fix numpy implementation (commaai#24865)

* Replace posfix with gauss newton method

* Cleanup

* Check if glonass is in the list

* Fix

* also return residual

* Add residuals

* Update selfdrive/locationd/laikad.py

Co-authored-by: Willem Melching <willem.melching@gmail.com>

* Cleanup

Co-authored-by: Gijs Koning <gijs-koning@live.nl>

selfdrive/locationd/laikad.py

@@ -7,16 +7,17 @@
 import numpy as np
 from collections import defaultdict
 
+import sympy
 from numpy.linalg import linalg
 
 from cereal import log, messaging
 from common.params import Params, put_nonblocking
 from laika import AstroDog
-from laika.constants import SECS_IN_HR, SECS_IN_MIN
+from laika.constants import EARTH_ROTATION_RATE, SECS_IN_HR, SECS_IN_MIN, SPEED_OF_LIGHT
 from laika.ephemeris import Ephemeris, EphemerisType, convert_ublox_ephem
 from laika.gps_time import GPSTime
 from laika.helpers import ConstellationId
-from laika.raw_gnss import GNSSMeasurement, calc_pos_fix, correct_measurements, process_measurements, read_raw_ublox
+from laika.raw_gnss import GNSSMeasurement, correct_measurements, process_measurements, read_raw_ublox
 from selfdrive.locationd.models.constants import GENERATED_DIR, ObservationKind
 from selfdrive.locationd.models.gnss_kf import GNSSKalman
 from selfdrive.locationd.models.gnss_kf import States as GStates
@@ -38,6 +39,7 @@ def __init__(self, valid_const=("GPS", "GLONASS"), auto_update=False, valid_ephe
     self.last_cached_t = None
     self.save_ephemeris = save_ephemeris
     self.load_cache()
+    self.posfix_functions = {constellation: get_posfix_sympy_fun(constellation) for constellation in (ConstellationId.GPS, ConstellationId.GLONASS)}
 
   def load_cache(self):
     cache = Params().get(EPHEMERIS_CACHE)
@@ -66,7 +68,9 @@ def process_ublox_msg(self, ublox_msg, ublox_mono_time: int, block=False):
         self.fetch_orbits(latest_msg_t + SECS_IN_MIN, block)
       new_meas = read_raw_ublox(report)
       processed_measurements = process_measurements(new_meas, self.astro_dog)
-      pos_fix = calc_pos_fix(processed_measurements, min_measurements=4)
+
+      min_measurements = 5 if any(p.constellation_id == ConstellationId.GLONASS for p in processed_measurements) else 4
+      pos_fix = calc_pos_fix_gauss_newton(processed_measurements, self.posfix_functions, min_measurements=min_measurements)
 
       t = ublox_mono_time * 1e-9
       kf_pos_std = None
@@ -84,7 +88,7 @@ def process_ublox_msg(self, ublox_msg, ublox_mono_time: int, block=False):
       if est_pos is not None:
         corrected_measurements = correct_measurements(processed_measurements, est_pos, self.astro_dog)
 
-      self.update_localizer(pos_fix, t, corrected_measurements)
+      self.update_localizer(est_pos, t, corrected_measurements)
       kf_valid = all(self.kf_valid(t))
 
       ecef_pos = self.gnss_kf.x[GStates.ECEF_POS].tolist()
@@ -110,7 +114,7 @@ def process_ublox_msg(self, ublox_msg, ublox_mono_time: int, block=False):
     # elif ublox_msg.which == 'ionoData':
     # todo add this. Needed to better correct messages offline. First fix ublox_msg.cc to sent them.
 
-  def update_localizer(self, pos_fix, t: float, measurements: List[GNSSMeasurement]):
+  def update_localizer(self, est_pos, t: float, measurements: List[GNSSMeasurement]):
     # Check time and outputs are valid
     valid = self.kf_valid(t)
     if not all(valid):
@@ -123,11 +127,10 @@ def update_localizer(self, pos_fix, t: float, measurements: List[GNSSMeasurement
       else:
         cloudlog.error("Gnss kalman std too far")
 
-      if len(pos_fix) == 0:
+      if est_pos is None:
         cloudlog.info("Position fix not available when resetting kalman filter")
         return
-      post_est = pos_fix[0][:3].tolist()
-      self.init_gnss_localizer(post_est)
+      self.init_gnss_localizer(est_pos.tolist())
     if len(measurements) > 0:
       kf_add_observations(self.gnss_kf, t, measurements)
     else:
@@ -227,6 +230,90 @@ def deserialize_hook(dct):
   return dct
 
 
+def calc_pos_fix_gauss_newton(measurements, posfix_functions, x0=None, signal='C1C', min_measurements=6):
+  '''
+  Calculates gps fix using gauss newton method
+  To solve the problem a minimal of 4 measurements are required.
+    If Glonass is included 5 are required to solve for the additional free variable.
+  returns:
+  0 -> list with positions
+  '''
+  if x0 is None:
+    x0 = [0, 0, 0, 0, 0]
+  n = len(measurements)
+  if n < min_measurements:
+    return []
+
+  Fx_pos = pr_residual(measurements, posfix_functions, signal=signal)
+  x = gauss_newton(Fx_pos, x0)
+  residual, _ = Fx_pos(x, weight=1.0)
+  return x, residual
+
+
+def pr_residual(measurements, posfix_functions, signal='C1C'):
+  def Fx_pos(inp, weight=None):
+    vals, gradients = [], []
+
+    for meas in measurements:
+      pr = meas.observables[signal]
+      pr += meas.sat_clock_err * SPEED_OF_LIGHT
+
+      w = (1 / meas.observables_std[signal]) if weight is None else weight
+
+      val, *gradient = posfix_functions[meas.constellation_id](*inp, pr, *meas.sat_pos, w)
+      vals.append(val)
+      gradients.append(gradient)
+    return np.asarray(vals), np.asarray(gradients)
+
+  return Fx_pos
+
+
+def gauss_newton(fun, b, xtol=1e-8, max_n=25):
+  for _ in range(max_n):
+    # Compute function and jacobian on current estimate
+    r, J = fun(b)
+
+    # Update estimate
+    delta = np.linalg.pinv(J) @ r
+    b -= delta
+
+    # Check step size for stopping condition
+    if np.linalg.norm(delta) < xtol:
+      break
+  return b
+
+
+def get_posfix_sympy_fun(constellation):
+  # Unknowns
+  x, y, z = sympy.Symbol('x'), sympy.Symbol('y'), sympy.Symbol('z')
+  bc = sympy.Symbol('bc')
+  bg = sympy.Symbol('bg')
+  var = [x, y, z, bc, bg]
+
+  # Knowns
+  pr = sympy.Symbol('pr')
+  sat_x, sat_y, sat_z = sympy.Symbol('sat_x'), sympy.Symbol('sat_y'), sympy.Symbol('sat_z')
+  weight = sympy.Symbol('weight')
+
+  theta = EARTH_ROTATION_RATE * (pr - bc) / SPEED_OF_LIGHT
+  val = sympy.sqrt(
+    (sat_x * sympy.cos(theta) + sat_y * sympy.sin(theta) - x) ** 2 +
+    (sat_y * sympy.cos(theta) - sat_x * sympy.sin(theta) - y) ** 2 +
+    (sat_z - z) ** 2
+  )
+
+  if constellation == ConstellationId.GLONASS:
+    res = weight * (val - (pr - bc - bg))
+  elif constellation == ConstellationId.GPS:
+    res = weight * (val - (pr - bc))
+  else:
+    raise NotImplementedError(f"Constellation {constellation} not supported")
+
+  res = [res] + [sympy.diff(res, v) for v in var]
+
+  return sympy.lambdify([x, y, z, bc, bg, pr, sat_x, sat_y, sat_z, weight], res)
+
+
 def main():
   sm = messaging.SubMaster(['ubloxGnss'])
   pm = messaging.PubMaster(['gnssMeasurements'])