Mandelbrot Set for Pi Pico in MicroPython and Assembly

mandelbrot.py

# Mandelbrot Set
import gc9a01
from machine import Pin, SPI, PWM, WDT, mem32
import framebuf
from time import sleep_ms, sleep_us, ticks_diff, ticks_us, sleep
from micropython import const
import array, gc, _thread
from usys import exit
from math import sin,cos,pi,radians,tan
from uctypes import addressof

MAXSCREEN_X = const(200)
MAXSCREEN_Y = const(200)
BLUE = const(0x1f00)
BLACK = const(0)
WHITE = const(0xffff)
GREEN = const(0xe00A)
BROWN = const(0xe091)
RED   = const(0x07e0) 
YELLOW=const(0x00fe)
MAGENTA=const(0x1FF8)

NORM_BITS = const(13) #13
NORM_FACT = const(1 << NORM_BITS)
MAXITERATE = const(96) #64
DRAW = False
QUIT = False

joyRight = Pin(17,Pin.IN)
joyDown  = Pin(18,Pin.IN)
joySel   = Pin(19,Pin.IN)
joyLeft  = Pin(20,Pin.IN)
joyUp    = Pin(21,Pin.IN)

machine.freq(270_000_000)

class Offset():
    def __init__(self):                    
        self.x = 0
        self.y = 0
        self.selold = 0
        self.seltime = 0
        self.scale = 100
        
offset = Offset()

def buttons():
    if not joyUp.value():
        offset.y-=0.02
    if not joyDown.value():
        offset.y+=0.02

    if not joyRight.value():
        offset.x+=0.02
    if not joyLeft.value():
        offset.x-=0.02
    if not joySel.value():
        if offset.seltime < 5 and offset.scale<150:
            offset.scale +=1
        elif offset.scale>3:
            offset.scale -= 1
        offset.selold = 1
    else:
        if offset.selold == 1:
            offset.seltime = 0
        else:
            offset.seltime += 1
        offset.selold = 0
    #print(offset.x)
     
        #QUIT = True
        #exit()

def fade(input_color1, input_color2):
    color1=input_color1<<8 | input_color1>>8 # byte swap to normal RBG565
    red1  =color1>>11& 0b11111               # extract red #13
    green1=color1>>6 & 0b11111               # extract green
    blue1 =color1    & 0b11111               # extract blue
    color2=input_color2<<8 | input_color2>>8 # byte swap
    red2  =color2>>11& 0b11111               # extract red
    green2=color2>>6 & 0b11111               # extract green
    blue2 =color2    & 0b11111               # extract blue
    inc_red  =(red2-  red1)/31               # find increment step
    inc_green=(green2-green1)/31
    inc_blue =(blue2- blue1)/31
    for i in range(0,32):       
        red3  =red1  +int(i*inc_red)          # build colors by steps
        green3=green1+int(i*inc_green)
        blue3 =blue1 +int(i*inc_blue)
        color3=red3<<11 | green3<<6 | blue3   # combine RGB  
        color_l2.append((color3 & 0x00FF)<<8 | (color3>>8))    # byte swap to LCD RGB565   
     

@micropython.viper
def draw_mb(realmin:int, imagmin:int, realmax:int, imagmax:int): 
    screen_addr=ptr16(screen)
    color_addr=ptr16(color_l2)
    deltareal = (realmax - realmin)>>8 #// MAXSCREEN_X
    deltaimag = (imagmax - imagmin)>>8 #// MAXSCREEN_Y
    real0 = realmin
    for py in range(MAXSCREEN_Y):
        imag0 = imagmax
        for px in range(MAXSCREEN_X):
            i = 0
            real = real0
            imag = imag0
            while i < MAXITERATE:  
                realq = (real * real) >> NORM_BITS
                imagq = (imag * imag) >> NORM_BITS
                if (realq + imagq) > (4 * NORM_FACT):
                    screen_addr[py*MAXSCREEN_X+px]=color_addr[i<<1]
                    break
                imag = ((real * imag) >> (NORM_BITS - 1)) + imag0
                real = realq - imagq + real0
                i+=4
            imag0 -= deltaimag
        real0 += deltareal
           

SCREEN_CTL    = const(0)  # 0
COLORS_CTL    = const(4)  # 1
REALMIN_CTL   = const(8)  # 2
IMAGMIN_CTL   = const(12) # 3
REALMAX_CTL   = const(16) # 4
IMAGMAX_CTL   = const(20) # 5
DELTAREAL_CTL = const(24) # 6
DELTAIMAG_CTL = const(28) # 7
REAL0_CTL     = const(32) # 8
IMAG0_CTL     = const(36) # 9
REALQ_CTL     = const(40) # 10
IMAGQ_CTL     = const(44) # 11
REAL_CTL      = const(48) # 12
IMAG_CTL      = const(52) # 13
START_CTL     = const(56) # 14
END_CTL       = const(60) # 15


@micropython.asm_thumb
def draw_mb_asm(r0):
    #b(EXIT)
    ldr(r1,[r0,REALMAX_CTL])
    ldr(r2,[r0,REALMIN_CTL])
    str(r2,[r0,REAL0_CTL])       # real0 = realmin
    sub(r1,r1,r2)
    asr(r1,r1,8)
    str(r1,[r0,DELTAREAL_CTL])   # deltareal = (realmax - realmin)>>8
    ldr(r1,[r0,IMAGMAX_CTL])
    ldr(r2,[r0,IMAGMIN_CTL])
    sub(r1,r1,r2)
    asr(r1,r1,8)
    str(r1,[r0,DELTAIMAG_CTL])   # deltaimag = (imagmax - imagmin)>>8
    #mov(r1,0)
    ldr(r1,[r0,START_CTL])
    label(LOOP_PY)               # r1 = py
    ldr(r7,[r0,IMAGMAX_CTL])
    str(r7,[r0,IMAG0_CTL])       # imag0 = imagmax
    mov(r2,0)
    label(LOOP_PX)               # r2 = px
    mov(r3,0)                    # r3 = i
    ldr(r7,[r0,REAL0_CTL])
    str(r7,[r0,REAL_CTL])        # real = real0
    ldr(r7,[r0,IMAG0_CTL])    
    str(r7,[r0,IMAG_CTL])        # imag = imag0
    label(LOOP_I)
    ldr(r4,[r0,REAL_CTL])
    mul(r4,r4)
    asr(r4,r4,13)
    str(r4,[r0,REALQ_CTL])       # realq = (real * real) >> NORM_BITS
    ldr(r5,[r0,IMAG_CTL])
    mul(r5,r5)
    asr(r5,r5,13)    
    str(r5,[r0,IMAGQ_CTL])       # imagq = (imag * imag) >> NORM_BITS
    add(r6,r4,r5)                # realq + imagq    
    mov(r7,1)
    lsl(r7,r7,15)                #                    4 * NORM_FACT
    cmp(r6,r7)                   # (realq + imagq) > (4 * NORM_FACT)
    ble(SKIP_DRAW)
    ldr(r4,[r0,SCREEN_CTL])
    mov(r5,200)
    mul(r5,r1)
    add(r5,r5,r2)                # py*MAXSCREEN_X+px
    add(r5,r5,r5)                # double for strh
    add(r4,r4,r5)
    ldr(r5,[r0,COLORS_CTL])
    lsl(r6,r3,1)
    add(r6,r6,r6)                # double for ldrh
    add(r5,r5,r6)
    ldrh(r5,[r5,0])    
    strh(r5,[r4,0])              # screen_addr[py*MAXSCREEN_X+px]=color_addr[i<<1]
    b(BREAK_I)    
    label(SKIP_DRAW)
    ldr(r4,[r0,IMAG_CTL])
    ldr(r5,[r0,REAL_CTL])
    mul(r4,r5)                   # real * imag
    asr(r4,r4,12)                #              >> (NORM_BITS - 1)
    ldr(r6,[r0,IMAG0_CTL])
    add(r4,r4,r6)                #
    str(r4,[r0,IMAG_CTL])        # imag = ((real * imag) >> (NORM_BITS - 1)) + imag0
    ldr(r4,[r0,REALQ_CTL])
    ldr(r5,[r0,IMAGQ_CTL])
    ldr(r6,[r0,REAL0_CTL])
    sub(r4,r4,r5)
    add(r4,r4,r6)
    str(r4,[r0,REAL_CTL])        # real = realq - imagq + real0
    add(r3,r3,4)                 # i+=4
    cmp(r3,96)
    blt(LOOP_I)
    label(BREAK_I)    
    ldr(r6,[r0,IMAG0_CTL])
    ldr(r7,[r0,DELTAIMAG_CTL])
    sub(r6,r6,r7)
    str(r6,[r0,IMAG0_CTL])       # imag0 -= deltaimag
    add(r2,r2,1)
    cmp(r2,200)
    blt(LOOP_PX)
    ldr(r6,[r0,REAL0_CTL])
    ldr(r7,[r0,DELTAREAL_CTL])
    add(r6,r6,r7)
    str(r6,[r0,REAL0_CTL])       # real0 += deltareal    
    add(r1,r1,1)
    ldr(r7,[r0,END_CTL])     # 200
    cmp(r1,r7)
    blt(LOOP_PY)
    label(EXIT)
    
    

def mandelbrot(start,end,step):
    global DRAW
    FPS = 60
    for scale in range(start,end,step):
        gticks = ticks_us()
        buttons()
        realmin = -4.5 * offset.scale / 100 - offset.y # -2.0, -3.7
        imagmin = -2 * offset.scale / 100 +   offset.x # -1.2
        realmax =  0.7 * offset.scale / 100 - offset.y #  0.7
        imagmax =  2 * offset.scale / 100 +   offset.x #  1.2
        realmin1 = -2.5 * offset.scale / 100 - offset.y # -2.0
        imagmin1 = -2 * offset.scale / 100 +   offset.x # -1.2
        realmax1 =  1.3 * offset.scale / 100 - offset.y #  0.7
        imagmax1 =  2* offset.scale / 100 +    offset.x #  1.2
        #draw_mb(int(realmin*NORM_FACT),int(imagmin*NORM_FACT), int(realmax*NORM_FACT), int(imagmax*NORM_FACT))        
        mb_control[REALMIN_CTL>>2] = int(realmin*NORM_FACT)
        mb_control[IMAGMIN_CTL>>2] = int(imagmin*NORM_FACT)
        mb_control[REALMAX_CTL>>2] = int(realmax*NORM_FACT)
        mb_control[IMAGMAX_CTL>>2] = int(imagmax*NORM_FACT)
        mb_control1[REALMIN_CTL>>2] = int(realmin1*NORM_FACT)
        mb_control1[IMAGMIN_CTL>>2] = int(imagmin1*NORM_FACT)
        mb_control1[REALMAX_CTL>>2] = int(realmax1*NORM_FACT)
        mb_control1[IMAGMAX_CTL>>2] = int(imagmax1*NORM_FACT)
        DRAW = True
        draw_mb_asm(mb_control)
        while DRAW:
            pass
        screen.text(str(FPS),80,10,0xffff)
        tft.blit_buffer(screen, 20, 20, MAXSCREEN_X, MAXSCREEN_Y)
        screen.fill(0)
        FPS = 1_000_000//ticks_diff(ticks_us(),gticks)
        #print(FPS)
        #QUIT = True
        #exit()

def core1():
    global DRAW, QUIT
    while not QUIT:
        if DRAW:
            draw_mb_asm(mb_control1)
            DRAW = False

   

if __name__=='__main__':
    spi = SPI(1, baudrate=63_000_000, sck=Pin(10), mosi=Pin(11))
    tft = gc9a01.GC9A01(
        spi,
        240,
        240,
        reset=Pin(12, Pin.OUT),
        cs=Pin(9, Pin.OUT),
        dc=Pin(8, Pin.OUT),
        backlight=Pin(13, Pin.OUT),
        rotation=0)
    tft.init()
    tft.rotation(0)
    tft.fill(gc9a01.BLACK)
    sleep(0.5)
    machine.mem32[0x40008048] = 1<<11 # enable peri_ctrl clock
    display_buffer=bytearray(MAXSCREEN_X * MAXSCREEN_Y * 2)
    screen=framebuf.FrameBuffer(display_buffer, MAXSCREEN_X , MAXSCREEN_Y, framebuf.RGB565)
    color_l2 = array.array('H', [])
    fade(BLACK,RED)
    fade(RED,YELLOW)
    fade(YELLOW,GREEN)
    fade(GREEN,YELLOW)
    fade(YELLOW,RED)
    fade(RED,BLACK)
    mb_control  = array.array('i',(addressof(screen),addressof(color_l2),REALMIN_CTL,IMAGMIN_CTL,REALMAX_CTL,IMAGMAX_CTL,DELTAREAL_CTL,
                              DELTAIMAG_CTL,REAL0_CTL,IMAG0_CTL,REALQ_CTL,IMAGQ_CTL,REAL_CTL,IMAG_CTL,0,100,0,0))
    mb_control1 = array.array('i',(addressof(screen),addressof(color_l2),REALMIN_CTL,IMAGMIN_CTL,REALMAX_CTL,IMAGMAX_CTL,DELTAREAL_CTL,
                              DELTAIMAG_CTL,REAL0_CTL,IMAG0_CTL,REALQ_CTL,IMAGQ_CTL,REAL_CTL,IMAG_CTL,100,200,0,0))
 
    _thread.start_new_thread(core1, ())

while True: 
    mandelbrot(150,1,-1)
    #QUIT = True
    #exit()
    mandelbrot(1,150,1)