From bd703c33e3aee61e997549c51638222b419cb85f Mon Sep 17 00:00:00 2001
From: jimmy <me@jimmy.nz>
Date: Sat, 30 Dec 2023 13:39:55 +0000
Subject: [PATCH] Use DMA

---
 dma_ws2812.py | 149 +++++++++++++++++++++++++++++---------------------
 1 file changed, 88 insertions(+), 61 deletions(-)

diff --git a/dma_ws2812.py b/dma_ws2812.py
index ac32888..c182095 100644
--- a/dma_ws2812.py
+++ b/dma_ws2812.py
@@ -1,43 +1,18 @@
-# Example based on pio_ws2812.py, but modified to use
-# a DMA channel to push out the data to the PIO
+# 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
 
 # Configure the number of WS2812 LEDs.
-NUM_LEDS = 9
+NUM_LEDS = 256
+PIN_NUM = 0
+brightness = 0.2
 
-SM = 0
-
-PIO0_BASE = 0x50200000
-PIO1_BASE = 0x50300000
-PIO0_BASE_TXF0 = PIO0_BASE+0x10
-PIO0_BASE_TXF1 = PIO0_BASE+0x14
-PIO0_BASE_TXF2 = PIO0_BASE+0x18
-PIO0_BASE_TXF3 = PIO0_BASE+0x1c
-
-PIO1_BASE_TXF0 = PIO1_BASE+0x10
-PIO1_BASE_TXF1 = PIO1_BASE+0x14
-PIO1_BASE_TXF2 = PIO1_BASE+0x18
-PIO1_BASE_TXF3 = PIO1_BASE+0x1c
-
-
-if SM < 4:
-    TFX = PIO0_BASE+0x10 + SM * 4
-else:
-    TFX = PIO1_BASE+0x10 + (SM-4) * 4
-    
-
-
-@rp2.asm_pio(
-    sideset_init=rp2.PIO.OUT_LOW,
-    out_shiftdir=rp2.PIO.SHIFT_LEFT,
-    autopull=True,
-    pull_thresh=24,
-)
+@rp2.asm_pio(sideset_init=rp2.PIO.OUT_LOW, out_shiftdir=rp2.PIO.SHIFT_RIGHT, autopull=True, pull_thresh=24)
 def ws2812():
-    # fmt: off
     T1 = 2
     T2 = 5
     T3 = 3
@@ -49,39 +24,91 @@ def ws2812():
     label("do_zero")
     nop()                   .side(0)    [T2 - 1]
     wrap()
-    # fmt: on
 
+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))
 
-# Create the StateMachine with the ws2812 program, outputting on Pin(2).
-sm = rp2.StateMachine(SM, ws2812, freq=8_000_000, sideset_base=Pin(0))
+        # Start the StateMachine, it will wait for data on its FIFO.
+        self.sm.active(1)
 
-# Start the StateMachine, it will wait for data on its FIFO.
-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.num_leds = num_leds
+        self.brightness = brightness
+        self.pio_dma = dma.PIO_DMA_Transfer(state_machine+4, state_machine, 32, num_leds)
 
-# Dummy data
-ar = array.array("I", [0 for _ in range(NUM_LEDS)])
-r = 0xff
-g = 0xff
-b = 0xff
-ar[0] = 0xff000000
-ar[1] = 0x00ff0000
-ar[2] = 0x0000ff00
-ar[3] = r<<24 + g<<16 + b<<8
+    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)
 
+    def fill(self, color):
+        for i in range(len(self.ar)):
+            self.set(i, color)
 
-# starts the transfer, using DMA channel 0
-def dma_out():
-    import dma, uctypes
-    dma.init_channels()
-    d0=dma.CHANNELS[0]
-    d0.CTRL_TRIG.EN = 0
-    d0.TRANS_COUNT = NUM_LEDS
-    d0.READ_ADDR = uctypes.addressof(ar)
-    d0.WRITE_ADDR = TFX
-    d0.CTRL_TRIG.INCR_WRITE = 0
-    d0.CTRL_TRIG.INCR_READ = 1
-    d0.CTRL_TRIG.DATA_SIZE = 2
-    d0.CTRL_TRIG.EN = 1
+    def wait(self):
+        return self.pio_dma.busy()
     
-dma_out()
-time.sleep(1)
+    def color_chase(self, color, wait):
+        for i in range(self.num_leds):
+            self.set(i, color)
+            time.sleep(wait)
+            self.show()
+        time.sleep(0.2)
+     
+    def wheel(self, pos):
+        # Input a value 0 to 255 to get a color value.
+        # The colours are a transition r - g - b - back to r.
+        if pos < 0 or pos > 255:
+            return (0, 0, 0)
+        if pos < 85:
+            return (255 - pos * 3, pos * 3, 0)
+        if pos < 170:
+            pos -= 85
+            return (0, 255 - pos * 3, pos * 3)
+        pos -= 170
+        return (pos * 3, 0, 255 - pos * 3)
+     
+     
+    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)
+
+    BLACK = (0, 0, 0)
+    RED = (255, 0, 0)
+    YELLOW = (255, 150, 0)
+    GREEN = (0, 255, 0)
+    CYAN = (0, 255, 255)
+    BLUE = (0, 0, 255)
+    PURPLE = (180, 0, 255)
+    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)
+   
+    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)
+            
+
+
+