Mandelbrot script for python

import numpy as np
import matplotlib.pyplot as plt

class Mandelbrot:
    def __init__(self, xmin, xmax, ymin, ymax, width, height, max_iter):
        self.xmin = xmin
        self.xmax = xmax
        self.ymin = ymin
        self.ymax = ymax
        self.width = width
        self.height = height
        self.max_iter = max_iter

    def mandelbrot(self, c):
        z = 0
        n = 0
        while abs(z) <= 2 and n < self.max_iter:
            z = z*z + c
            n += 1
        return n

    def mandelbrot_set(self):
        x = np.linspace(self.xmin, self.xmax, self.width)
        y = np.linspace(self.ymin, self.ymax, self.height)
        mandelbrot_matrix = np.zeros((self.height, self.width), dtype=int)

        for i in range(self.height):
            for j in range(self.width):
                mandelbrot_matrix[i, j] = self.mandelbrot(complex(x[j], y[i]))

        return mandelbrot_matrix

    def plot(self):
        mandelbrot_matrix = self.mandelbrot_set()
        plt.imshow(mandelbrot_matrix.T, extent=(self.xmin, self.xmax, self.ymin, self.ymax), cmap='magma', aspect='auto')
        plt.colorbar()
        plt.title("Mandelbrot Fractal")
        plt.xlabel("Real")
        plt.ylabel("Imaginary")
        plt.show()

def main_menu():
    print("Welcome to Mandelbrot Fractal Generator!")
    print("1. Generate Mandelbrot Fractal")
    print("2. Exit")
    choice = input("Enter your choice: ")
    return choice

if __name__ == "__main__":
    while True:
        choice = main_menu()
        if choice == "1":
            xmin = float(input("Enter xmin: "))
            xmax = float(input("Enter xmax: "))
            ymin = float(input("Enter ymin: "))
            ymax = float(input("Enter ymax: "))
            width = int(input("Enter width: "))
            height = int(input("Enter height: "))
            max_iter = int(input("Enter max iterations: "))

            mandelbrot = Mandelbrot(xmin, xmax, ymin, ymax, width, height, max_iter)
            mandelbrot.plot()
        elif choice == "2":
            print("Exiting...")
            break
        else:
            print("Invalid choice. Please try again.")