Add sysidroutine (#43)

Co-authored-by: David Vo <auscompgeek@users.noreply.github.com>

commands2/sysid/__init__.py

@@ -0,0 +1,4 @@
+from .sysidroutine import SysIdRoutine
+
+
+__all__ = ["SysIdRoutine"]

commands2/sysid/sysidroutine.py

@@ -0,0 +1,180 @@
+from dataclasses import dataclass
+from enum import Enum
+
+from wpilib.sysid import SysIdRoutineLog, State
+from ..command import Command
+from ..subsystem import Subsystem
+from wpilib import Timer
+
+from wpimath.units import seconds, volts
+
+from typing import Callable, Optional
+
+
+volts_per_second = float
+
+
+class SysIdRoutine(SysIdRoutineLog):
+    """A SysId characterization routine for a single mechanism. Mechanisms may have multiple motors.
+
+    A single subsystem may have multiple mechanisms, but mechanisms should not share test
+    routines. Each complete test of a mechanism should have its own SysIdRoutine instance, since the
+    log name of the recorded data is determined by the mechanism name.
+
+    The test state (e.g. "quasistatic-forward") is logged once per iteration during test
+    execution, and once with state "none" when a test ends. Motor frames are logged every iteration
+    during test execution.
+
+    Timestamps are not coordinated across data, so motor frames and test state tags may be
+    recorded on different log frames. Because frame alignment is not guaranteed, SysId parses the log
+    by using the test state flag to determine the timestamp range for each section of the test, and
+    then extracts the motor frames within the valid timestamp ranges. If a given test was run
+    multiple times in a single logfile, the user will need to select which of the tests to use for
+    the fit in the analysis tool.
+    """
+
+    @dataclass
+    class Config:
+        """Hardware-independent configuration for a SysId test routine.
+
+        :param rampRate:    The voltage ramp rate used for quasistatic test routines. Defaults to 1 volt
+                            per second if left null.
+        :param stepVoltage: The step voltage output used for dynamic test routines. Defaults to 7
+                            volts if left null.
+        :param timeout:     Safety timeout for the test routine commands. Defaults to 10 seconds if left
+                            null.
+        :param recordState: Optional handle for recording test state in a third-party logging
+                            solution. If provided, the test routine state will be passed to this callback instead of
+                            logged in WPILog.
+        """
+
+        rampRate: volts_per_second = 1.0
+        stepVoltage: volts = 7.0
+        timeout: seconds = 10.0
+        recordState: Optional[Callable[[State], None]] = None
+
+    @dataclass
+    class Mechanism:
+        """A mechanism to be characterized by a SysId routine.
+
+        Defines callbacks needed for the SysId test routine to control
+        and record data from the mechanism.
+
+        :param drive:     Sends the SysId-specified drive signal to the mechanism motors during test
+                          routines.
+        :param log:       Returns measured data of the mechanism motors during test routines. To return
+                          data, call `motor(string motorName)` on the supplied `SysIdRoutineLog` instance, and then
+                          call one or more of the chainable logging handles (e.g. `voltage`) on the returned
+                          `MotorLog`. Multiple motors can be logged in a single callback by calling `motor`
+                          multiple times.
+        :param subsystem: The subsystem containing the motor(s) that is (or are) being characterized.
+                          Will be declared as a requirement for the returned test commands.
+        :param name:      The name of the mechanism being tested. Will be appended to the log entry title
+                          for the routine's test state, e.g. "sysid-test-state-mechanism". Defaults to the name of
+                          the subsystem if left null.
+        """
+
+        drive: Callable[[volts], None]
+        log: Callable[[SysIdRoutineLog], None]
+        subsystem: Subsystem
+        name: Optional[str] = None
+
+        def __post_init__(self):
+            if self.name == None:
+                self.name = self.subsystem.getName()
+
+    class Direction(Enum):
+        """Motor direction for a SysId test."""
+
+        kForward = 1
+        kReverse = -1
+
+    def __init__(self, config: Config, mechanism: Mechanism):
+        """Create a new SysId characterization routine.
+
+        :param config:    Hardware-independent parameters for the SysId routine.
+        :param mechanism: Hardware interface for the SysId routine.
+        """
+        super().__init__(mechanism.subsystem.getName())
+        self.config = config
+        self.mechanism = mechanism
+        self.outputVolts = 0.0
+        self.logState = config.recordState or self.recordState
+
+    def quasistatic(self, direction: Direction) -> Command:
+        """Returns a command to run a quasistatic test in the specified direction.
+
+        The command will call the `drive` and `log` callbacks supplied at routine construction once
+        per iteration. Upon command end or interruption, the `drive` callback is called with a value of
+        0 volts.
+
+        :param direction: The direction in which to run the test.
+
+        :returns: A command to run the test.
+        """
+
+        timer = Timer()
+        if direction == self.Direction.kForward:
+            state = State.kQuasistaticForward
+        else:
+            state = State.kQuasistaticReverse
+
+        def execute():
+            self.outputVolts = direction.value * timer.get() * self.config.rampRate
+            self.mechanism.drive(self.outputVolts)
+            self.mechanism.log(self)
+            self.logState(state)
+
+        def end(interrupted: bool):
+            self.mechanism.drive(0.0)
+            self.logState(State.kNone)
+            timer.stop()
+
+        return (
+            self.mechanism.subsystem.runOnce(timer.start)
+            .andThen(self.mechanism.subsystem.run(execute))
+            .finallyDo(end)
+            .withName(
+                f"sysid-{SysIdRoutineLog.stateEnumToString(state)}-{self.mechanism.name}"
+            )
+            .withTimeout(self.config.timeout)
+        )
+
+    def dynamic(self, direction: Direction) -> Command:
+        """Returns a command to run a dynamic test in the specified direction.
+
+        The command will call the `drive` and `log` callbacks supplied at routine construction once
+        per iteration. Upon command end or interruption, the `drive` callback is called with a value of
+        0 volts.
+
+        :param direction: The direction in which to run the test.
+
+        :returns: A command to run the test.
+        """
+
+        if direction == self.Direction.kForward:
+            state = State.kDynamicForward
+        else:
+            state = State.kDynamicReverse
+
+        def command():
+            self.outputVolts = direction.value * self.config.stepVoltage
+
+        def execute():
+            self.mechanism.drive(self.outputVolts)
+            self.mechanism.log(self)
+            self.logState(state)
+
+        def end(interrupted: bool):
+            self.mechanism.drive(0.0)
+            self.logState(State.kNone)
+
+        return (
+            self.mechanism.subsystem.runOnce(command)
+            .andThen(self.mechanism.subsystem.run(execute))
+            .finallyDo(end)
+            .withName(
+                f"sysid-{SysIdRoutineLog.stateEnumToString(state)}-{self.mechanism.name}"
+            )
+            .withTimeout(self.config.timeout)
+        )

tests/test_sysidroutine.py

@@ -0,0 +1,168 @@
+import pytest
+from unittest.mock import Mock, call, ANY
+from wpilib.simulation import stepTiming, pauseTiming, resumeTiming
+from wpimath.units import volts
+from commands2 import Command, Subsystem
+from commands2.sysid import SysIdRoutine
+from wpilib.sysid import SysIdRoutineLog, State
+
+
+class Mechanism(Subsystem):
+    def recordState(self, state: State):
+        pass
+
+    def drive(self, voltage: volts):
+        pass
+
+    def log(self, log: SysIdRoutineLog):
+        pass
+
+
+@pytest.fixture
+def mechanism():
+    return Mock(spec=Mechanism)
+
+
+@pytest.fixture
+def sysid_routine(mechanism):
+    return SysIdRoutine(
+        SysIdRoutine.Config(recordState=mechanism.recordState),
+        SysIdRoutine.Mechanism(mechanism.drive, mechanism.log, Subsystem()),
+    )
+
+
+@pytest.fixture
+def quasistatic_forward(sysid_routine):
+    return sysid_routine.quasistatic(SysIdRoutine.Direction.kForward)
+
+
+@pytest.fixture
+def quasistatic_reverse(sysid_routine):
+    return sysid_routine.quasistatic(SysIdRoutine.Direction.kReverse)
+
+
+@pytest.fixture
+def dynamic_forward(sysid_routine):
+    return sysid_routine.dynamic(SysIdRoutine.Direction.kForward)
+
+
+@pytest.fixture
+def dynamic_reverse(sysid_routine):
+    return sysid_routine.dynamic(SysIdRoutine.Direction.kReverse)
+
+
+@pytest.fixture(autouse=True)
+def timing():
+    pauseTiming()
+    yield
+    resumeTiming()
+
+
+def run_command(command: Command):
+    command.initialize()
+    command.execute()
+    stepTiming(1)
+    command.execute()
+    command.end(True)
+
+
+def test_record_state_bookends_motor_logging(
+    mechanism, quasistatic_forward, dynamic_forward
+):
+    run_command(quasistatic_forward)
+
+    mechanism.assert_has_calls(
+        [
+            call.drive(ANY),
+            call.log(ANY),
+            call.recordState(State.kQuasistaticForward),
+            call.drive(ANY),
+            call.recordState(State.kNone),
+        ],
+        any_order=False,
+    )
+
+    mechanism.reset_mock()
+    run_command(dynamic_forward)
+
+    mechanism.assert_has_calls(
+        [
+            call.drive(ANY),
+            call.log(ANY),
+            call.recordState(State.kDynamicForward),
+            call.drive(ANY),
+            call.recordState(State.kNone),
+        ],
+        any_order=False,
+    )
+
+
+def test_tests_declare_correct_state(
+    mechanism,
+    quasistatic_forward,
+    quasistatic_reverse,
+    dynamic_forward,
+    dynamic_reverse,
+):
+    run_command(quasistatic_forward)
+    mechanism.recordState.assert_any_call(State.kQuasistaticForward)
+
+    run_command(quasistatic_reverse)
+    mechanism.recordState.assert_any_call(State.kQuasistaticReverse)
+
+    run_command(dynamic_forward)
+    mechanism.recordState.assert_any_call(State.kDynamicForward)
+
+    run_command(dynamic_reverse)
+    mechanism.recordState.assert_any_call(State.kDynamicReverse)
+
+
+def test_tests_output_correct_voltage(
+    mechanism,
+    quasistatic_forward,
+    quasistatic_reverse,
+    dynamic_forward,
+    dynamic_reverse,
+):
+    run_command(quasistatic_forward)
+
+    mechanism.drive.assert_has_calls(
+        [
+            call(pytest.approx(1.0)),
+            call(pytest.approx(0.0)),
+        ],
+        any_order=False,
+    )
+
+    mechanism.reset_mock()
+    run_command(quasistatic_reverse)
+
+    mechanism.drive.assert_has_calls(
+        [
+            call(pytest.approx(-1.0)),
+            call(pytest.approx(0.0)),
+        ],
+        any_order=False,
+    )
+
+    mechanism.reset_mock()
+    run_command(dynamic_forward)
+
+    mechanism.drive.assert_has_calls(
+        [
+            call(pytest.approx(7.0)),
+            call(pytest.approx(0.0)),
+        ],
+        any_order=False,
+    )
+
+    mechanism.reset_mock()
+    run_command(dynamic_reverse)
+
+    mechanism.drive.assert_has_calls(
+        [
+            call(pytest.approx(-7.0)),
+            call(pytest.approx(0.0)),
+        ],
+        any_order=False,
+    )