Add plot configs for looking at 2kHz IMU data

Needed to add the ability to specify specific signals for the x-axis to
be able to actually understand the IMU signals with their timestamps.

Change-Id: I0f3c4b41a10ad6a98bbd71751ffcca70a90e394d
diff --git a/frc971/analysis/plot.py b/frc971/analysis/plot.py
index d325bd7..1eba716 100644
--- a/frc971/analysis/plot.py
+++ b/frc971/analysis/plot.py
@@ -8,7 +8,7 @@
 import sys
 
 from frc971.analysis.py_log_reader import LogReader
-from frc971.analysis.plot_config_pb2 import PlotConfig, Signal
+from frc971.analysis.plot_config_pb2 import PlotConfig, Signal, Line
 from google.protobuf import text_format
 
 import numpy as np
@@ -58,6 +58,10 @@
                     total_acceleration = np.sqrt(
                         msg[accel_x]**2 + msg[accel_y]**2 + msg[accel_z]**2)
                     new_msg['total_acceleration'] = total_acceleration
+                timestamp = 'monotonic_timestamp_ns'
+                if timestamp in msg:
+                    timestamp_sec = msg[timestamp] * 1e-9
+                    new_msg['monotonic_timestamp_sec'] = timestamp_sec
                 entries.append((entry[0], entry[1], new_msg))
             if 'CalcIMU' in self.data:
                 raise RuntimeError('CalcIMU is already a member of data.')
@@ -77,28 +81,57 @@
         to understanding how some internal filters work."""
         self.calculate_imu_signals()
 
-    def plot_signal(self, axes: matplotlib.axes.Axes, signal: Signal):
-        if not signal.channel in self.data:
-            raise ValueError("No channel alias " + signal.channel)
-        field_path = signal.field.split('.')
-        monotonic_time = []
-        signal_data = []
-        for entry in self.data[signal.channel]:
-            monotonic_time.append(entry[0] * 1e-9)
-            value = entry[2]
-            for name in field_path:
-                # If the value wasn't populated in a given message, fill in
-                # NaN rather than crashing.
-                if name in value:
-                    value = value[name]
-                else:
-                    value = float("nan")
-                    break
-            # Catch NaNs and convert them to floats.
-            value = float(value)
-            signal_data.append(value)
-        label_name = signal.channel + "." + signal.field
-        axes.plot(monotonic_time, signal_data, marker='o', label=label_name)
+    def extract_field(self, message: dict, field: str):
+        """Extracts a field with the given name from the message.
+
+        message will be a dictionary with field names as the keys and then
+        values, lists, or more dictionaries as the values. field is the full
+        path to the field to extract, with periods separating sub-messages."""
+        field_path = field.split('.')
+        value = message
+        for name in field_path:
+            # If the value wasn't populated in a given message, fill in
+            # NaN rather than crashing.
+            if name in value:
+                value = value[name]
+            else:
+                return None
+        # Catch NaNs and convert them to floats.
+        return float(value)
+
+    def plot_line(self, axes: matplotlib.axes.Axes, line: Line):
+        if not line.HasField('y_signal'):
+            raise ValueError("No y_channel specified for line.")
+        y_signal = line.y_signal
+        if not y_signal.channel in self.data:
+            raise ValueError("No channel alias " + y_signal.channel)
+        x_signal = line.x_signal if line.HasField('x_signal') else None
+        if x_signal is not None and not x_signal.channel in self.data:
+            raise ValueError("No channel alias " + x_signal.channel)
+        y_messages = self.data[y_signal.channel]
+        x_messages = self.data[
+            x_signal.channel] if x_signal is not None else None
+        if x_messages is not None and len(x_messages) != len(y_messages):
+            raise ValueError(
+                "X and Y signal lengths don't match. X channel is " +
+                x_signal.channel + " Y channel is " + y_signal.channel)
+        x_data = []
+        y_data = []
+        for ii in range(len(y_messages)):
+            y_entry = y_messages[ii]
+            if x_signal is None:
+                x_data.append(y_entry[0] * 1e-9)
+            else:
+                x_entry = x_messages[ii]
+                x_data.append(self.extract_field(x_entry[2], x_signal.field))
+            y_data.append(self.extract_field(y_entry[2], y_signal.field))
+            if x_data[-1] is None and y_data[-1] is not None:
+                raise ValueError(
+                    "Only one of the x and y signals is present. X " +
+                    x_signal.channel + "." + x_signal.field + " Y " +
+                    y_signal.channel + "." + y_signal.field)
+        label_name = y_signal.channel + "." + y_signal.field
+        axes.plot(x_data, y_data, marker='o', label=label_name)
 
     def plot(self):
         shared_axis = None
@@ -110,8 +143,8 @@
                     num_subplots, 1, ii + 1, sharex=shared_axis)
                 shared_axis = shared_axis or axes
                 axes_config = figure_config.axes[ii]
-                for signal in axes_config.signal:
-                    self.plot_signal(axes, signal)
+                for line in axes_config.line:
+                    self.plot_line(axes, line)
                 # Make the legend transparent so we can see behind it.
                 legend = axes.legend(framealpha=0.5)
                 axes.set_xlabel("Monotonic Time (sec)")
diff --git a/frc971/analysis/plot_config.proto b/frc971/analysis/plot_config.proto
index 9aaea89..c4c61c2 100644
--- a/frc971/analysis/plot_config.proto
+++ b/frc971/analysis/plot_config.proto
@@ -24,10 +24,26 @@
   optional string field = 2;
 }
 
+// A single line to plot.
+message Line {
+  // The signal to plot on the y-axis.
+  optional Signal y_signal = 1;
+  // If set, we will use this signal for the x-axis of the plot. By default, we
+  // will use the monotonic sent time of the message. This is helpful for both
+  // plotting against non-time based signals (e.g., plotting x/y robot position)
+  // as well as plotting against times other than the message sent time (e.g.,
+  // for the IMU where the sample capture time is separate from the actual
+  // sent time.
+  // Note that the x and y signals should have exactly the same number of
+  // entries--otherwise, we need to write logic to handle resampling one signal
+  // to a different rate.
+  optional Signal x_signal = 2;
+}
+
 // Message representing a single pyplot Axes, with specifications for exactly
 // which signals to show in the supplied subplot.
 message Axes {
-  repeated Signal signal = 1;
+  repeated Line line = 1;
   optional string ylabel = 2;
 }
 
diff --git a/frc971/analysis/plot_configs/drivetrain.pb b/frc971/analysis/plot_configs/drivetrain.pb
index 2b15220..0ac4fff 100644
--- a/frc971/analysis/plot_configs/drivetrain.pb
+++ b/frc971/analysis/plot_configs/drivetrain.pb
@@ -31,39 +31,53 @@
 
 figure {
   axes {
-    signal {
-      channel: "JoystickState"
-      field: "test_mode"
+    line {
+      y_signal {
+        channel: "JoystickState"
+        field: "test_mode"
+      }
     }
-    signal {
-      channel: "JoystickState"
-      field: "autonomous"
+    line {
+      y_signal {
+        channel: "JoystickState"
+        field: "autonomous"
+      }
     }
-    signal {
-      channel: "RobotState"
-      field: "browned_out"
+    line {
+      y_signal {
+        channel: "RobotState"
+        field: "browned_out"
+      }
     }
-    signal {
-      channel: "JoystickState"
-      field: "enabled"
+    line {
+      y_signal {
+        channel: "JoystickState"
+        field: "enabled"
+      }
     }
     ylabel: "[bool]"
   }
   axes {
-    signal {
-      channel: "RobotState"
-      field: "voltage_battery"
+    line {
+      y_signal {
+        channel: "RobotState"
+        field: "voltage_battery"
+      }
     }
     ylabel: "[V]"
   }
   axes {
-    signal {
-      channel: "Status"
-      field: "line_follow_logging.frozen"
+    line {
+      y_signal {
+        channel: "Status"
+        field: "line_follow_logging.frozen"
+      }
     }
-    signal {
-      channel: "Goal"
-      field: "quickturn"
+    line {
+      y_signal {
+        channel: "Goal"
+        field: "quickturn"
+      }
     }
     ylabel: "[bool]"
   }
@@ -71,59 +85,83 @@
 
 figure {
   axes {
-    signal {
-      channel: "Status"
-      field: "poly_drive_logging.ff_left_voltage"
+    line {
+      y_signal {
+        channel: "Status"
+        field: "poly_drive_logging.ff_left_voltage"
+      }
     }
-    signal {
-      channel: "Status"
-      field: "poly_drive_logging.ff_right_voltage"
+    line {
+      y_signal {
+        channel: "Status"
+        field: "poly_drive_logging.ff_right_voltage"
+      }
     }
-    signal {
-      channel: "Output"
-      field: "left_voltage"
+    line {
+      y_signal {
+        channel: "Output"
+        field: "left_voltage"
+      }
     }
-    signal {
-      channel: "Output"
-      field: "right_voltage"
+    line {
+      y_signal {
+        channel: "Output"
+        field: "right_voltage"
+      }
     }
     ylabel: "[V]"
   }
   axes {
-    signal {
-      channel: "Status"
-      field: "theta"
+    line {
+      y_signal {
+        channel: "Status"
+        field: "theta"
+      }
     }
     ylabel: "[rad]"
   }
   axes {
-    signal {
-      channel: "Status"
-      field: "robot_speed"
+    line {
+      y_signal {
+        channel: "Status"
+        field: "robot_speed"
+      }
     }
-    signal {
-      channel: "Status"
-      field: "trajectory_logging.left_velocity"
+    line {
+      y_signal {
+        channel: "Status"
+        field: "trajectory_logging.left_velocity"
+      }
     }
-    signal {
-      channel: "Status"
-      field: "poly_drive_logging.goal_left_velocity"
+    line {
+      y_signal {
+        channel: "Status"
+        field: "poly_drive_logging.goal_left_velocity"
+      }
     }
-    signal {
-      channel: "Status"
-      field: "estimated_left_velocity"
+    line {
+      y_signal {
+        channel: "Status"
+        field: "estimated_left_velocity"
+      }
     }
-    signal {
-      channel: "Status"
-      field: "trajectory_logging.right_velocity"
+    line {
+      y_signal {
+        channel: "Status"
+        field: "trajectory_logging.right_velocity"
+      }
     }
-    signal {
-      channel: "Status"
-      field: "poly_drive_logging.goal_right_velocity"
+    line {
+      y_signal {
+        channel: "Status"
+        field: "poly_drive_logging.goal_right_velocity"
+      }
     }
-    signal {
-      channel: "Status"
-      field: "estimated_right_velocity"
+    line {
+      y_signal {
+        channel: "Status"
+        field: "estimated_right_velocity"
+      }
     }
     ylabel: "[m/s]"
   }
@@ -131,33 +169,129 @@
 
 figure {
   axes {
-    signal {
-      channel: "IMU"
-      field: "gyro_x"
+    line {
+      y_signal {
+        channel: "IMU"
+        field: "gyro_x"
+      }
+      x_signal {
+        channel: "CalcIMU"
+        field: "monotonic_timestamp_sec"
+      }
     }
-    signal {
-      channel: "IMU"
-      field: "gyro_y"
+    line {
+      y_signal {
+        channel: "IMU"
+        field: "gyro_y"
+      }
+      x_signal {
+        channel: "CalcIMU"
+        field: "monotonic_timestamp_sec"
+      }
     }
-    signal {
-      channel: "IMU"
-      field: "gyro_z"
+    line {
+      y_signal {
+        channel: "IMU"
+        field: "gyro_z"
+      }
+      x_signal {
+        channel: "CalcIMU"
+        field: "monotonic_timestamp_sec"
+      }
     }
     ylabel: "rad / sec"
   }
   axes {
-    signal {
-      channel: "IMU"
-      field: "accelerometer_x"
+    line {
+      y_signal {
+        channel: "CalcIMU"
+        field: "total_acceleration"
+      }
+      x_signal {
+        channel: "CalcIMU"
+        field: "monotonic_timestamp_sec"
+      }
     }
-    signal {
-      channel: "IMU"
-      field: "accelerometer_y"
+    line {
+      y_signal {
+        channel: "IMU"
+        field: "accelerometer_x"
+      }
+      x_signal {
+        channel: "CalcIMU"
+        field: "monotonic_timestamp_sec"
+      }
     }
-    signal {
-      channel: "IMU"
-      field: "accelerometer_z"
+    line {
+      y_signal {
+        channel: "IMU"
+        field: "accelerometer_y"
+      }
+      x_signal {
+        channel: "CalcIMU"
+        field: "monotonic_timestamp_sec"
+      }
+    }
+    line {
+      y_signal {
+        channel: "IMU"
+        field: "accelerometer_z"
+      }
+      x_signal {
+        channel: "CalcIMU"
+        field: "monotonic_timestamp_sec"
+      }
     }
     ylabel: "g"
   }
 }
+
+figure {
+  axes {
+    line {
+      y_signal {
+        channel: "Status"
+        field: "down_estimator.yaw"
+      }
+    }
+    line {
+      y_signal {
+        channel: "Status"
+        field: "down_estimator.lateral_pitch"
+      }
+    }
+    line {
+      y_signal {
+        channel: "Status"
+        field: "down_estimator.longitudinal_pitch"
+      }
+    }
+    ylabel: "rad"
+  }
+  axes {
+    line {
+      y_signal {
+        channel: "Status"
+        field: "down_estimator.quaternion_x"
+      }
+    }
+    line {
+      y_signal {
+        channel: "Status"
+        field: "down_estimator.quaternion_y"
+      }
+    }
+    line {
+      y_signal {
+        channel: "Status"
+        field: "down_estimator.quaternion_z"
+      }
+    }
+    line {
+      y_signal {
+        channel: "Status"
+        field: "down_estimator.quaternion_w"
+      }
+    }
+  }
+}
diff --git a/frc971/analysis/plot_configs/gyro.pb b/frc971/analysis/plot_configs/gyro.pb
index f246835..760752c 100644
--- a/frc971/analysis/plot_configs/gyro.pb
+++ b/frc971/analysis/plot_configs/gyro.pb
@@ -6,36 +6,78 @@
 
 figure {
   axes {
-    signal {
-      channel: "IMU"
-      field: "gyro_x"
+    line {
+      y_signal {
+        channel: "IMU"
+        field: "gyro_x"
+      }
+      x_signal {
+        channel: "CalcIMU"
+        field: "monotonic_timestamp_sec"
+      }
     }
-    signal {
-      channel: "IMU"
-      field: "gyro_y"
+    line {
+      y_signal {
+        channel: "IMU"
+        field: "gyro_y"
+      }
+      x_signal {
+        channel: "CalcIMU"
+        field: "monotonic_timestamp_sec"
+      }
     }
-    signal {
-      channel: "IMU"
-      field: "gyro_z"
+    line {
+      y_signal {
+        channel: "IMU"
+        field: "gyro_z"
+      }
+      x_signal {
+        channel: "CalcIMU"
+        field: "monotonic_timestamp_sec"
+      }
     }
     ylabel: "rad / sec"
   }
   axes {
-    signal {
-      channel: "CalcIMU"
-      field: "total_acceleration"
+    line {
+      y_signal {
+        channel: "CalcIMU"
+        field: "total_acceleration"
+      }
+      x_signal {
+        channel: "CalcIMU"
+        field: "monotonic_timestamp_sec"
+      }
     }
-    signal {
-      channel: "IMU"
-      field: "accelerometer_x"
+    line {
+      y_signal {
+        channel: "IMU"
+        field: "accelerometer_x"
+      }
+      x_signal {
+        channel: "CalcIMU"
+        field: "monotonic_timestamp_sec"
+      }
     }
-    signal {
-      channel: "IMU"
-      field: "accelerometer_y"
+    line {
+      y_signal {
+        channel: "IMU"
+        field: "accelerometer_y"
+      }
+      x_signal {
+        channel: "CalcIMU"
+        field: "monotonic_timestamp_sec"
+      }
     }
-    signal {
-      channel: "IMU"
-      field: "accelerometer_z"
+    line {
+      y_signal {
+        channel: "IMU"
+        field: "accelerometer_z"
+      }
+      x_signal {
+        channel: "CalcIMU"
+        field: "monotonic_timestamp_sec"
+      }
     }
     ylabel: "g"
   }