Update diagnostics/MotorDiagnostics.py

diagnostics/MotorDiagnostics.py

@@ -0,0 +1,142 @@
+from pybricks.hubs import PrimeHub
+from pybricks.pupdevices import Motor, ColorSensor, UltrasonicSensor, ForceSensor
+from pybricks.parameters import Button, Color, Direction, Port, Side, Stop
+from pybricks.robotics import DriveBase
+from pybricks.tools import wait, StopWatch
+import umath
+
+hub = PrimeHub()
+class MotorDiagnostics:
+  def __init__(self):
+      self.testmotor = None
+      self.port_map = {
+        "A": Port.A,
+        "B": Port.B,
+        "C": Port.C,
+        "D": Port.D,
+        "E": Port.E,
+        "F": Port.F,
+      }
+
+  def stdev(self, vals):
+    data = vals
+    if len(data) < 2:
+        return 0
+    # Calculate the mean
+    mean = sum(data) / len(data)
+
+    # Calculate the variance (sum of squared differences from the mean, divided by n-1 for sample standard deviation)
+    variance = sum([(x - mean) ** 2 for x in data]) / float(len(data) - 1)
+
+    # Calculate the standard deviation (square root of the variance)
+    
+    std_dev_manual = umath.sqrt(variance)
+    
+
+    return (std_dev_manual)
+  def health_score(self, desired, avg_speed, stdev_speed, avg_load):
+    # Speed accuracy: penalize % error
+    accuracy = max(0, 100 - abs(avg_speed - desired) / desired * 100)
+
+    # Stability: penalize deviation relative to desired
+    stability = max(0, 100 - (stdev_speed / desired) * 100)
+
+    # Normalize load: map 10–20 as baseline (≈0%), 200 as max (≈100%)
+    baseline = 15      # midpoint of idle range
+    max_observed = 200 # heavy load/stall
+    normalized_load = max(0, avg_load - baseline)
+    load_pct = min(100, (normalized_load / (max_observed - baseline)) * 100)
+
+    load_score = max(0, 100 - load_pct)
+
+    # Final score: average of the three
+    return (accuracy + stability + load_score) / 3
+
+
+    # Final score: average of the three
+    return (accuracy + stability + load_score) / 3
+  def initializeMotor(self):
+    valid_ports = {"A", "B", "C", "D", "E", "F"}
+    while True:
+      motorinput = input(
+        "This test will run your motor at 3 speeds: 180, 540, and 1000 degrees per second.\n"
+        "Please make sure your motor is not under any load (for example, your hand) during the test.\n"
+        "If you want to test the wheel's load, note that this will affect the load measurements.\n"
+        "Enter the port for the motor you would like to test (A, B, C, D, E, or F): "
+      ).strip().upper()
+
+      try:
+        # Only create a new Motor if we don't already have one
+        if self.testmotor is None:
+            self.testmotor = Motor(self.port_map[motorinput])
+            print(f"Motor initialized on port {motorinput}.")
+        else:
+            print(f"Reusing existing motor on port {motorinput}.")
+        break
+
+      except OSError as e:
+        if e.errno == 16:  # EBUSY
+          print(f"Port {motorinput} is already in use. Reusing existing motor.")
+          # Do not overwrite self.testmotor here — keep the existing reference
+          break
+        else:
+          print(f"Error initializing motor on port {motorinput}: {e}")
+          print("Make sure a motor is actually connected to this port.")
+          self.testmotor = None
+
+  def testSpeed(self, speed):
+    self.testmotor.reset_angle(0)
+    
+    motorspeeds = []
+    motorloads = []
+    target_angle = speed * 3
+    print("\n", speed, "DEGREES PER SECOND TEST")
+    self.testmotor.run_angle(speed, target_angle, Stop.HOLD, False)
+    stopwatchmotor = StopWatch()
+    while stopwatchmotor.time() < 3000:
+        wait(20)
+        motorspeeds.append(self.testmotor.speed())
+        motorloads.append(self.testmotor.load())
+    
+    max_speed, max_accel, max_torque = self.testmotor.control.limits()
+
+
+    print("Desired motor speed: ", str(speed))
+    if motorspeeds:
+      avg = sum(motorspeeds) / len(motorspeeds)
+      print("Average motor speed:", avg)
+      print("Motor speed deviation:", str(self.stdev(motorspeeds)))
+    else:
+      print("No speed samples collected.")
+      avg = 0
+    if motorloads:
+      avgload = sum(motorloads) / len(motorloads)
+
+      print("Average motor load:", avgload)
+      print("Motor load deviation:", str(self.stdev(motorloads)))
+    else:
+      print("No load samples collected.")
+      avgload = 0
+    score = self.health_score(speed, avg, self.stdev(motorspeeds), avgload)
+    print("Health score for this test:", str(score) + "%")
+    return score
+  def fullTest(self):
+    self.initializeMotor()
+    print("Load measurements are in mNm. Speed measurements are in degrees per second.")
+    max_speed, max_accel, max_torque = self.testmotor.control.limits()
+    print("Maximum motor speed:", max_speed)
+    test180 = self.testSpeed(180)
+    test540 = self.testSpeed(540)
+    test1000 = self.testSpeed(1000)
+    print("\n FINAL MOTOR STATISTICS")
+    final = (test180 + test540 + test1000) / 3
+    print("Final motor health score:", str(final) + "%")
+    if final < 80:
+      print("Your motor is in need of attention. Make sure to clean it regularly and charge the Prime Hub.")
+    elif final < 90:
+      print("Your motor is in OK condition. Make sure to clean it regularly and charge the Prime Hub.")
+    elif final < 97:
+      print("Your motor is in great condition!")
+    else:
+      print("Your motor is in AMAZING condition!!!")
+    self.testmotor.stop()