code: def main():
      """Carry out two examples of state estimation on groups.
  
      Both examples are localization problems, where only a part of the system
      is observed. The first one is a linear system, while the second one is
      non-linear.
      """
      np.random.seed(12345)
      model = LocalizationLinear()
      kalman = KalmanFilter(model)
      n_traj = 1000
      obs_freq = 50
      dt = 0.1
      init_cov = gs.array([10.0, 1.0])
      init_cov = algebra_utils.from_vector_to_diagonal_matrix(init_cov)
      prop_cov = 0.001 * gs.eye(model.dim_noise)
      obs_cov = 10 * gs.eye(model.dim_obs)
      initial_covs = (init_cov, prop_cov, obs_cov)
      kalman.initialize_covariances(*initial_covs)
  
      true_state = gs.array([0.0, 0.0])
      true_acc = gs.random.uniform(-1, 1, (n_traj, 1))
      dt_vectorized = dt * gs.ones((n_traj, 1))
      true_inputs = gs.hstack((dt_vectorized, true_acc))
  
      true_traj, inputs, observations = create_data(
          kalman, true_state, true_inputs, obs_freq
      )
  
      initial_state = np.random.multivariate_normal(
          true_state, init_cov
      )
      estimate, uncertainty = estimation(
          kalman, initial_state, inputs, observations, obs_freq
      )
  
      plt.figure()
      plt.plot(true_traj[:, 0], label="Ground Truth")
      plt.plot(estimate[:, 0], label="Kalman")
      plt.plot(
          estimate[:, 0] + uncertainty[:, 0], color="k", linestyle=":"
      )
      plt.plot(
          estimate[:, 0] - uncertainty[:, 0],
          color="k",
          linestyle=":",
          label="3_sigma envelope",
      )
      plt.plot(
          range(obs_freq, n_traj + 1, obs_freq),
          observations,
          marker="*",
          linestyle="",
          label="Observation",
      )
      plt.legend()
      plt.title("1D Localization - Position")
  
      plt.figure()
      plt.plot(true_traj[:, 1], label="Ground Truth")
      plt.plot(estimate[:, 1], label="Kalman")
      plt.plot(
          estimate[:, 1] + uncertainty[:, 1], color="k", linestyle=":"
      )
      plt.plot(
          estimate[:, 1] - uncertainty[:, 1],
          color="k",
          linestyle=":",
          label="3_sigma envelope",
      )
      plt.legend()
      plt.title("1D Localization - Speed")
  
      model = Localization()
      kalman = KalmanFilter(model)
      init_cov = gs.array([1.0, 10.0, 10.0])
      init_cov = algebra_utils.from_vector_to_diagonal_matrix(init_cov)
      prop_cov = 0.001 * gs.eye(model.dim_noise)
      obs_cov = 0.1 * gs.eye(model.dim_obs)
      initial_covs = (init_cov, prop_cov, obs_cov)
      kalman.initialize_covariances(*initial_covs)
  
      true_state = gs.zeros(model.group.dim)
      true_inputs = [
          gs.array([dt, 0.5, 0.0, 0.05]) for _ in range(n_traj)
      ]
  
      true_traj, inputs, observations = create_data(
          kalman, true_state, true_inputs, obs_freq
      )
  
      initial_state = gs.array(
          np.random.multivariate_normal(true_state, init_cov)
      )
      initial_state = gs.cast(initial_state, true_state.dtype)
      estimate, uncertainty = estimation(
          kalman, initial_state, inputs, observations, obs_freq
      )
  
      plt.figure()
      plt.plot(true_traj[:, 1], true_traj[:, 2], label="Ground Truth")
      plt.plot(estimate[:, 1], estimate[:, 2], label="Kalman")
      plt.scatter(
          observations[:, 0],
          observations[:, 1],
          s=2,
          c="k",
          label="Observation",
      )
      plt.legend()
      plt.axis("equal")
      plt.title("2D Localization")
  
      plt.show()
  
imports: import matplotlib.pyplot
  import numpy
  import geomstats.backend
  from geomstats import algebra_utils
  from geomstats.learning.kalman_filter import KalmanFilter, Localization, LocalizationLinear
label_desc: Use image node_id 3 for calling a global function with example usage: main() without return types