From abc6849f21a3c68b4438e3958f3221f184ede643 Mon Sep 17 00:00:00 2001
From: kleeman <akleeman@gmail.com>
Date: Tue, 13 Oct 2020 14:22:22 -0700
Subject: [PATCH] Add the ability to extract a child covariance function from a
 composed function

---
 .../covariance_function.hpp                   |  8 ++
 .../albatross/src/utils/covariance_utils.hpp  | 73 +++++++++++++++++++
 tests/test_covariance_functions.cc            | 70 ++++++++++++++++++
 3 files changed, 151 insertions(+)
 create mode 100644 include/albatross/src/utils/covariance_utils.hpp

diff --git a/include/albatross/src/covariance_functions/covariance_function.hpp b/include/albatross/src/covariance_functions/covariance_function.hpp
index dc181ab6..770478f7 100644
--- a/include/albatross/src/covariance_functions/covariance_function.hpp
+++ b/include/albatross/src/covariance_functions/covariance_function.hpp
@@ -305,6 +305,10 @@ class SumOfCovarianceFunctions
     return this->rhs_.state_space_representation(xs);
   }
 
+  LHS get_lhs() const { return lhs_; }
+
+  RHS get_rhs() const { return rhs_; }
+
 protected:
   LHS lhs_;
   RHS rhs_;
@@ -403,6 +407,10 @@ class ProductOfCovarianceFunctions
     return this->rhs_.state_space_representation(xs);
   }
 
+  LHS get_lhs() const { return lhs_; }
+
+  RHS get_rhs() const { return rhs_; }
+
 protected:
   LHS lhs_;
   RHS rhs_;
diff --git a/include/albatross/src/utils/covariance_utils.hpp b/include/albatross/src/utils/covariance_utils.hpp
new file mode 100644
index 00000000..6c2c69d0
--- /dev/null
+++ b/include/albatross/src/utils/covariance_utils.hpp
@@ -0,0 +1,73 @@
+/*
+ * Copyright (C) 2020 Swift Navigation Inc.
+ * Contact: Swift Navigation <dev@swiftnav.com>
+ *
+ * This source is subject to the license found in the file 'LICENSE' which must
+ * be distributed together with this source. All other rights reserved.
+ *
+ * THIS CODE AND INFORMATION IS PROVIDED "AS IS" WITHOUT WARRANTY OF ANY KIND,
+ * EITHER EXPRESSED OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE IMPLIED
+ * WARRANTIES OF MERCHANTABILITY AND/OR FITNESS FOR A PARTICULAR PURPOSE.
+ */
+
+#ifndef INCLUDE_ALBATROSS_SRC_UTILS_COVARIANCE_UTILS_HPP_
+#define INCLUDE_ALBATROSS_SRC_UTILS_COVARIANCE_UTILS_HPP_
+
+namespace albatross {
+
+namespace details {
+
+template <typename Child, typename Parent> struct cov_func_contains {
+  const static bool value = false;
+};
+
+template <typename Child> struct cov_func_contains<Child, Child> {
+  const static bool value = true;
+};
+
+template <typename Child, template <typename...> class Composite, typename LHS,
+          typename RHS>
+struct cov_func_contains<Child, Composite<LHS, RHS>> {
+  const static bool value = cov_func_contains<Child, LHS>::value ||
+                            cov_func_contains<Child, RHS>::value;
+};
+
+template <typename Parent> struct GetChildImpl {
+  template <typename Child,
+            std::enable_if_t<std::is_same<Parent, Child>::value, int> = 0>
+  static Child get_child(const Parent &parent) {
+    return parent;
+  }
+};
+
+template <typename T> struct GetChildIdentity { typedef T type; };
+
+template <template <typename...> class Composite, typename LHS, typename RHS>
+struct GetChildImpl<Composite<LHS, RHS>> {
+  template <typename Child,
+            std::enable_if_t<cov_func_contains<Child, LHS>::value, int> = 0>
+  static Child get_child(const Composite<LHS, RHS> &parent) {
+    return GetChildImpl<LHS>::template get_child<Child>(parent.get_lhs());
+  }
+
+  template <typename Child,
+            std::enable_if_t<cov_func_contains<Child, RHS>::value &&
+                                 !cov_func_contains<Child, LHS>::value,
+                             int> = 0>
+  static Child get_child(const Composite<LHS, RHS> &parent) {
+    return GetChildImpl<RHS>::template get_child<Child>(parent.get_rhs());
+  }
+};
+
+} // namespace details
+
+template <typename Child, typename Parent>
+Child get_child_covariance_function(
+    const Parent &parent,
+    details::GetChildIdentity<Child> && = details::GetChildIdentity<Child>()) {
+  return details::GetChildImpl<Parent>::template get_child<Child>(parent);
+}
+
+} // namespace albatross
+
+#endif /* INCLUDE_ALBATROSS_SRC_UTILS_COVARIANCE_UTILS_HPP_ */
diff --git a/tests/test_covariance_functions.cc b/tests/test_covariance_functions.cc
index f19336d2..257f3dea 100644
--- a/tests/test_covariance_functions.cc
+++ b/tests/test_covariance_functions.cc
@@ -11,6 +11,7 @@
  */
 #include <gtest/gtest.h>
 
+#include "../include/albatross/src/utils/covariance_utils.hpp"
 #include <albatross/Core>
 #include <albatross/CovarianceFunctions>
 
@@ -250,4 +251,73 @@ TEST(test_covariance_functions, test_state_space_representation) {
       std::is_same<decltype(xs_and_ys), std::vector<variant<X, Y>>>::value));
 }
 
+template <typename Child, typename Parent>
+bool can_get_child(const Parent &parent) {
+  auto child = get_child_covariance_function<Child>(parent);
+  return std::is_same<Child, decltype(child)>::value;
+}
+
+TEST(test_covariance_function, test_get_child_covariance_function) {
+
+  DummyCovariance dummy;
+  const auto dummy_copy = get_child_covariance_function<DummyCovariance>(dummy);
+  EXPECT_TRUE(can_get_child<DummyCovariance>(dummy_copy));
+
+  IndependentNoise<double> noise;
+  noise.sigma_independent_noise = M_2_PI;
+  EXPECT_TRUE(can_get_child<IndependentNoise<double>>(noise));
+  EXPECT_EQ(get_child_covariance_function<IndependentNoise<double>>(noise)
+                .sigma_independent_noise,
+            M_2_PI);
+
+  Polynomial<2> poly;
+  EXPECT_TRUE(can_get_child<Polynomial<2>>(poly));
+
+  const auto sum = noise + dummy;
+  EXPECT_TRUE(can_get_child<IndependentNoise<double>>(sum));
+  EXPECT_EQ(get_child_covariance_function<IndependentNoise<double>>(sum)
+                .sigma_independent_noise,
+            M_2_PI);
+  EXPECT_TRUE(can_get_child<DummyCovariance>(sum));
+
+  const auto double_sum = sum + poly;
+  EXPECT_TRUE(can_get_child<IndependentNoise<double>>(double_sum));
+  EXPECT_EQ(get_child_covariance_function<IndependentNoise<double>>(double_sum)
+                .sigma_independent_noise,
+            M_2_PI);
+  EXPECT_TRUE(can_get_child<DummyCovariance>(double_sum));
+  EXPECT_TRUE(can_get_child<Polynomial<2>>(double_sum));
+
+  const auto prod = noise * dummy;
+  EXPECT_TRUE(can_get_child<IndependentNoise<double>>(prod));
+  EXPECT_EQ(get_child_covariance_function<IndependentNoise<double>>(prod)
+                .sigma_independent_noise,
+            M_2_PI);
+  EXPECT_TRUE(can_get_child<DummyCovariance>(prod));
+
+  const auto double_prod = prod * poly;
+  EXPECT_TRUE(can_get_child<IndependentNoise<double>>(double_prod));
+  EXPECT_EQ(get_child_covariance_function<IndependentNoise<double>>(double_prod)
+                .sigma_independent_noise,
+            M_2_PI);
+  EXPECT_TRUE(can_get_child<DummyCovariance>(double_prod));
+  EXPECT_TRUE(can_get_child<Polynomial<2>>(double_prod));
+
+  const auto prod_sum = prod * sum;
+  EXPECT_TRUE(can_get_child<IndependentNoise<double>>(prod_sum));
+  EXPECT_EQ(get_child_covariance_function<IndependentNoise<double>>(prod_sum)
+                .sigma_independent_noise,
+            M_2_PI);
+  EXPECT_TRUE(can_get_child<DummyCovariance>(prod_sum));
+
+  const auto double_prod_sum = double_prod * sum;
+  EXPECT_TRUE(can_get_child<IndependentNoise<double>>(double_prod_sum));
+  EXPECT_EQ(
+      get_child_covariance_function<IndependentNoise<double>>(double_prod_sum)
+          .sigma_independent_noise,
+      M_2_PI);
+  EXPECT_TRUE(can_get_child<DummyCovariance>(double_prod_sum));
+  EXPECT_TRUE(can_get_child<Polynomial<2>>(double_prod_sum));
+}
+
 } // namespace albatross