Use bytearray

dma_ws2812.py

@@ -1,17 +1,14 @@
-# Example using PIO to drive a set of WS2812 LEDs.
 
 import array, time
 from machine import Pin
 import rp2
 from time import sleep
 import dma
+import uasyncio as asyncio
 
 # Configure the number of WS2812 LEDs.
-NUM_LEDS = 256
-PIN_NUM = 0
-brightness = 0.2
 
-@rp2.asm_pio(sideset_init=rp2.PIO.OUT_LOW, out_shiftdir=rp2.PIO.SHIFT_RIGHT, autopull=True, pull_thresh=24)
+@rp2.asm_pio(sideset_init=rp2.PIO.OUT_LOW, out_shiftdir=rp2.PIO.SHIFT_LEFT, autopull=True, pull_thresh=8)
 def ws2812():
     T1 = 2
     T2 = 5
@@ -26,37 +23,34 @@ def ws2812():
     wrap()
 
 class WS2812B: 
-    def __init__(self, num_leds, pin_num, brightness, state_machine):
-        # Create the StateMachine with the ws2812 program, outputting on pin
-        self.sm = rp2.StateMachine(state_machine, ws2812, freq=8_000_000, sideset_base=Pin(pin_num))
-
-        # Start the StateMachine, it will wait for data on its FIFO.
+    def __init__(self, num_leds, pin, state_machine, brightness=0.1, invert=False):
+        self.sm = rp2.StateMachine(state_machine, ws2812, freq=8_000_000, sideset_base=Pin(pin))
         self.sm.active(1)
-
-        # Display a pattern on the LEDs via an array of LED RGB values.
-        self.ar = array.array("I", [0 for _ in range(num_leds)])
+        self.ar = bytearray(num_leds*3)
         self.num_leds = num_leds
         self.brightness = brightness
-        self.pio_dma = dma.PIO_DMA_Transfer(state_machine+4, state_machine, 32, num_leds)
+        self.invert = invert
+        self.pio_dma = dma.PIO_DMA_Transfer(state_machine+4, state_machine, 8, num_leds*3)
 
     def show(self):
-        #self.sm.put(self.ar)
         self.pio_dma.start_transfer(self.ar)
         
     def set(self, i, color):
-        self.ar[i] = int((color[1]<<16)*self.brightness) + int((color[0]<<8)*self.brightness) + int(color[2]*self.brightness)
+        self.ar[i*3] = int(color[1]*self.brightness)
+        self.ar[i*3+1] = int(color[0]*self.brightness)
+        self.ar[i*3+2] = int(color[2]*self.brightness)
     
     def fill(self, color):
-        for i in range(len(self.ar)):
-            self.set(i, color)
-
-    def wait(self):
-        return self.pio_dma.busy()
-    
-    def color_chase(self, color, wait):
         for i in range(self.num_leds):
             self.set(i, color)
-            time.sleep(wait)
+
+    def busy(self):
+        return self.pio_dma.busy()
+    
+    async def color_chase(self, color, wait):
+        for i in range(self.num_leds):
+            self.set(i, color)
+            await asyncio.sleep(wait)
             self.show()
         time.sleep(0.2)
      
@@ -74,13 +68,13 @@ class WS2812B:
         return (pos * 3, 0, 255 - pos * 3)
      
      
-    def rainbow_cycle(self, wait):
+    async def rainbow_cycle(self, wait):
         for j in range(255):
             for i in range(self.num_leds):
                 rc_index = (i * 256 // self.num_leds) + j
                 self.set(i, self.wheel(rc_index & 255))
             self.show()
-            time.sleep(wait)
+            await asyncio.sleep(wait)
 
     BLACK = (0, 0, 0)
     RED = (255, 0, 0)
@@ -92,21 +86,13 @@ class WS2812B:
     WHITE = (255, 255, 255)
     COLORS = (BLACK, RED, YELLOW, GREEN, CYAN, BLUE, PURPLE, WHITE)
 
-
-
 if __name__ == "__main__":
-    ws0 = WS2812B(256, 1, 1, 1)
-    ws1 = WS2812B(256, 2, 1, 5)
-    ws2 = WS2812B(256, 3, 1, 6)
-   
+    num_leds, pin, sm, brightness = 10, 0, 0, 1
+    ws0 = WS2812B(num_leds, pin, sm, brightness)
     while True:
         for color in ws0.COLORS:
             ws0.fill(color)
             ws0.show()
-            ws1.fill(color)
-            ws1.show()            
-            ws2.fill(color)
-            ws2.show()
             time.sleep(1)