Add closed loop option to drivetrain config.

Change-Id: Icff6a814f6a75843bbb0d441d12cc33b5387b843
diff --git a/frc971/control_loops/drivetrain/drivetrain_config.h b/frc971/control_loops/drivetrain/drivetrain_config.h
index 04cfa6c..21dd874 100644
--- a/frc971/control_loops/drivetrain/drivetrain_config.h
+++ b/frc971/control_loops/drivetrain/drivetrain_config.h
@@ -15,10 +15,18 @@
   SIMPLE_SHIFTER = 1,  // Switch gears without speedmatch logic.
 };
 
+enum class LoopType : int32_t {
+  OPEN_LOOP = 0,  // Only use open loop logic.
+  CLOSED_LOOP = 1,  // Add in closed loop calculation.
+};
+
 struct DrivetrainConfig {
   // Shifting method we are using.
   ShifterType shifter_type;
 
+  // Type of loop to use.
+  LoopType loop_type;
+
   // Polydrivetrain functions returning various controller loops with plants.
   ::std::function<StateFeedbackLoop<4, 2, 2>()> make_drivetrain_loop;
   ::std::function<StateFeedbackLoop<2, 2, 2>()> make_v_drivetrain_loop;
diff --git a/frc971/control_loops/drivetrain/drivetrain_lib_test.cc b/frc971/control_loops/drivetrain/drivetrain_lib_test.cc
index 554a93e..afa8cf1 100644
--- a/frc971/control_loops/drivetrain/drivetrain_lib_test.cc
+++ b/frc971/control_loops/drivetrain/drivetrain_lib_test.cc
@@ -34,26 +34,19 @@
 const DrivetrainConfig &GetDrivetrainConfig() {
   static DrivetrainConfig kDrivetrainConfig{
       ::frc971::control_loops::drivetrain::ShifterType::HALL_EFFECT_SHIFTER,
+      ::frc971::control_loops::drivetrain::LoopType::CLOSED_LOOP,
 
       ::y2016::control_loops::drivetrain::MakeDrivetrainLoop,
       ::y2016::control_loops::drivetrain::MakeVelocityDrivetrainLoop,
       ::y2016::control_loops::drivetrain::MakeKFDrivetrainLoop,
 
       ::y2016::control_loops::drivetrain::kDt,
-      ::y2016::control_loops::drivetrain::kStallTorque,
-      ::y2016::control_loops::drivetrain::kStallCurrent,
-      ::y2016::control_loops::drivetrain::kFreeSpeedRPM,
-      ::y2016::control_loops::drivetrain::kFreeCurrent,
-      ::y2016::control_loops::drivetrain::kJ,
-      ::y2016::control_loops::drivetrain::kMass,
       ::y2016::control_loops::drivetrain::kRobotRadius,
       ::y2016::control_loops::drivetrain::kWheelRadius,
-      ::y2016::control_loops::drivetrain::kR,
       ::y2016::control_loops::drivetrain::kV,
-      ::y2016::control_loops::drivetrain::kT,
+
       ::y2016::control_loops::drivetrain::kHighGearRatio,
       ::y2016::control_loops::drivetrain::kLowGearRatio,
-
       kThreeStateDriveShifter,
       kThreeStateDriveShifter};
 
diff --git a/frc971/control_loops/drivetrain/polydrivetrain.cc b/frc971/control_loops/drivetrain/polydrivetrain.cc
index 9dad2a8..6e08b84 100644
--- a/frc971/control_loops/drivetrain/polydrivetrain.cc
+++ b/frc971/control_loops/drivetrain/polydrivetrain.cc
@@ -260,8 +260,10 @@
 }
 
 void PolyDrivetrain::Update() {
-  loop_->mutable_X_hat()(0, 0) = kf_.X_hat()(1, 0);
-  loop_->mutable_X_hat()(1, 0) = kf_.X_hat()(3, 0);
+  if (dt_config_.loop_type == LoopType::CLOSED_LOOP) {
+    loop_->mutable_X_hat()(0, 0) = kf_.X_hat()(1, 0);
+    loop_->mutable_X_hat()(1, 0) = kf_.X_hat()(3, 0);
+  }
 
   // TODO(austin): Observer for the current velocity instead of difference
   // calculations.
@@ -318,6 +320,11 @@
     for (int i = 0; i < 2; i++) {
       loop_->mutable_U()[i] = ::aos::Clip(U_ideal[i], -12, 12);
     }
+
+    if (dt_config_.loop_type == LoopType::OPEN_LOOP) {
+      loop_->mutable_X_hat() =
+          loop_->A() * loop_->X_hat() + loop_->B() * loop_->U();
+    }
   } else {
     const double current_left_velocity =
         (position_.left_encoder - last_position_.left_encoder) /
diff --git a/y2014/control_loops/drivetrain/drivetrain_base.cc b/y2014/control_loops/drivetrain/drivetrain_base.cc
index e78fa22..1176718 100644
--- a/y2014/control_loops/drivetrain/drivetrain_base.cc
+++ b/y2014/control_loops/drivetrain/drivetrain_base.cc
@@ -15,6 +15,7 @@
 const DrivetrainConfig &GetDrivetrainConfig() {
   static DrivetrainConfig kDrivetrainConfig{
       ::frc971::control_loops::drivetrain::ShifterType::HALL_EFFECT_SHIFTER,
+      ::frc971::control_loops::drivetrain::LoopType::CLOSED_LOOP,
 
       ::y2014::control_loops::drivetrain::MakeDrivetrainLoop,
       ::y2014::control_loops::drivetrain::MakeVelocityDrivetrainLoop,
diff --git a/y2016/control_loops/drivetrain/drivetrain_base.cc b/y2016/control_loops/drivetrain/drivetrain_base.cc
index 0065199..95b8907 100644
--- a/y2016/control_loops/drivetrain/drivetrain_base.cc
+++ b/y2016/control_loops/drivetrain/drivetrain_base.cc
@@ -20,6 +20,7 @@
 const DrivetrainConfig &GetDrivetrainConfig() {
   static DrivetrainConfig kDrivetrainConfig{
       ::frc971::control_loops::drivetrain::ShifterType::HALL_EFFECT_SHIFTER,
+      ::frc971::control_loops::drivetrain::LoopType::CLOSED_LOOP,
 
       ::y2016::control_loops::drivetrain::MakeDrivetrainLoop,
       ::y2016::control_loops::drivetrain::MakeVelocityDrivetrainLoop,