Python Matplotlib custom zoom function isn't working

Question

I'm making code for an interactive Mandelbrot set display, and I wanted the zoom function to update the image every time it is used. So I'm using Matplotlib event handling to create a custom zoom function.

The zoom function that is already there will pixelate the image as you zoom in, because it zooms in on the same image. My zoom should update the image every time one drags the mouse.

But it doesn't work. Every time I zoom once, the plot doesn't change shape, so the image itself stretches. Zoom again, and it just zooms there on the original picture. The ticks don't even change the entire time. Try it -- here's my code with a picture provided by Trenton McKinney:

# Get necessary packages
import numpy as np
from matplotlib import pyplot as plt

# Calculation
COLORS = [(1, 0, 0)]
for transition in [[0, 1, 0], [-1, 0, 0], [0, 0, 1], [0, -1, 0], [1, 0, 0]]:
    for _ in range(10):
        r, g, b = COLORS[-1]
        r = round(r + transition[0] / 10, 1)
        g = round(g + transition[1] / 10, 1)
        b = round(b + transition[2] / 10, 1)
        COLORS.append((r, g, b))
        
def color(n):
    z = complex(0, 0)
    c = n
    for n in range(len(COLORS)):
        z = z ** 2 + c
        if abs(z) > 2:
            return COLORS[n]
    return (0, 0, 0)


def get_data(c1, c2):
    d = []
    for i in np.linspace(c1.imag, c2.imag, num = 500):
        rvals = [complex(r, i) for r in np.linspace(c1.real, c2.real, num = 500)]
        d.append(list(map(color, rvals)))
    return d


# Interface
corner1 = (-2, -2)
corner2 = (2, 2)

def display_data(c1, c2):
    data = get_data(c1, c2)
    image.set_data(data)
    figure.show()
    
    
def startzoom(event):
    global corner1, corner2
    if event.button:
        x = event.xdata
        y = event.ydata
        corner1 = corner2 = (x, y)
        
        
def stopzoom(event):
    global corner1, corner2
    if event.button:
        x = event.xdata
        y = event.ydata
        corner2 = (x, y)
        if corner1 == corner2:
            corner1, corner2 = (-2, -2), (2, 2)
        display_data(complex(*corner1), complex(*corner2))
        

# Put it all together
data = get_data(complex(*corner1), complex(*corner2))
figure, axes = plt.subplots()
axes.set_title("Mandelbrot Set")
axes.set_xlabel("Real")
axes.set_ylabel("Imaginary")
image = axes.imshow(list(reversed(data)), extent = [-2, 2, -2, 2])
figure.canvas.mpl_connect("button_press_event", startzoom)
figure.canvas.mpl_connect("button_release_event", stopzoom)
figure.show()

What should I do?

Answer 1

I edited your display_data function as below:

def display_data(c1, c2):
    data = get_data(c1, c2)
    axes.imshow(list(reversed(data)), extent = [corner1[0], corner2[0], corner1[1], corner2[1]])
    plt.show()

You should update extent parameter values each time a user click on the plot, in order to update image borders.
Moreover, I added axes.invert_xaxis() and axes.invert_yaxis() so the zoom works if every direction: from bottom-left to top-right, from top-left to bottom-right, from top-right to bottom-left and from bottom-right to top-left.

Complete Code

import numpy as np
from matplotlib import pyplot as plt

# Calculation
COLORS = [(1, 0, 0)]
for transition in [[0, 1, 0], [-1, 0, 0], [0, 0, 1], [0, -1, 0], [1, 0, 0]]:
    for _ in range(10):
        r, g, b = COLORS[-1]
        r = round(r + transition[0] / 10, 1)
        g = round(g + transition[1] / 10, 1)
        b = round(b + transition[2] / 10, 1)
        COLORS.append((r, g, b))

def color(n):
    z = complex(0, 0)
    c = n
    for n in range(len(COLORS)):
        z = z ** 2 + c
        if abs(z) > 2:
            return COLORS[n]
    return (0, 0, 0)


def get_data(c1, c2):
    d = []
    for i in np.linspace(c1.imag, c2.imag, num = 500):
        rvals = [complex(r, i) for r in np.linspace(c1.real, c2.real, num = 500)]
        d.append(list(map(color, rvals)))
    return d


# Interface
corner1 = (-2, -2)
corner2 = (2, 2)


def display_data(c1, c2):
    data = get_data(c1, c2)
    axes.imshow(list(reversed(data)), extent = [corner1[0], corner2[0], corner1[1], corner2[1]])
    plt.show()


def startzoom(event):
    global corner1, corner2
    if event.button:
        x = event.xdata
        y = event.ydata
        corner1 = corner2 = (x, y)


def stopzoom(event):
    global corner1, corner2
    if event.button:
        x = event.xdata
        y = event.ydata
        corner2 = (x, y)
        if corner1 == corner2:
            corner1, corner2 = (-2, -2), (2, 2)
        display_data(complex(*corner1), complex(*corner2))


# Put it all together
data = get_data(complex(*corner1), complex(*corner2))
figure, axes = plt.subplots()
axes.set_title("Mandelbrot Set")
axes.set_xlabel("Real")
axes.set_ylabel("Imaginary")
image = axes.imshow(list(reversed(data)), extent = [-2, 2, -2, 2])
figure.canvas.mpl_connect("button_press_event", startzoom)
figure.canvas.mpl_connect("button_release_event", stopzoom)
plt.show()