mandelbrot.py
3774 bytes
1#! /usr/bin/env python
2# -*- coding: utf-8 -*-
3# vim:fenc=utf-8
4#
5# Copyright © 2019 pavle <pavle.portic@tilda.center>
6#
7# Distributed under terms of the BSD-3-Clause license.
8
9
10import numpy as np # type: ignore
11import curses
12
13import colors
14
15MAX_ITERS = 100
16ITERATIONS = 0
17COLOR_COUNT = 101
18MONOCHROME = False
19ZOOM_LEVEL = 0
20X_SCALE = 1.75
21Y_SCALE = 1
22X_MID = -0.75
23Y_MID = 0
24
25
26def get_points(cols, rows):
27 x_min = X_MID - X_SCALE * (1 / (2 ** ZOOM_LEVEL))
28 x_max = X_MID + X_SCALE * (1 / (2 ** ZOOM_LEVEL))
29 y_min = Y_MID - Y_SCALE * (1 / (2 ** ZOOM_LEVEL))
30 y_max = Y_MID + Y_SCALE * (1 / (2 ** ZOOM_LEVEL))
31 x = np.linspace(x_min, x_max, cols)
32 y = np.linspace(y_min, y_max, rows)
33 return x, y
34
35
36def calculate_point(x, y):
37 global ITERATIONS
38 z = 0
39 p = (x - 0.25) ** 2 + y ** 2
40 if p * (p + (x - 0.25)) < 0.25 * (y ** 2):
41 return -1
42
43 if (x + 1) ** 2 + y ** 2 <= 0.0625:
44 return -1
45
46 for i in range(MAX_ITERS):
47 ITERATIONS += 1
48 z = z ** 2 + complex(x, y)
49 if abs(z) >= 2:
50 return i
51
52 return -1
53
54
55def calculate_set(cols, rows):
56 global ITERATIONS
57 ITERATIONS = 0
58 x, y = get_points(cols, rows)
59 matrix = np.zeros([rows, cols])
60 for i in range(rows):
61 for j in range(cols):
62 matrix[i, j] = calculate_point(x[j], y[i])
63
64 return matrix
65
66
67def print_set(screen, cols, rows, matrix, char):
68 screen.clear()
69 for i in range(rows):
70 for j in range(cols):
71 iteration = matrix[i, j]
72 if iteration == -1:
73 screen.addstr(i, j, ' ', curses.color_pair(1))
74 else:
75 color_index = (int(iteration) + 2) % COLOR_COUNT
76 color_index = 2 if color_index < 2 else color_index
77 screen.addstr(i, j, ' ', curses.color_pair(color_index))
78
79 screen.addstr(
80 rows,
81 0,
82 f'Max Iterations: {MAX_ITERS}, Zoom: {ZOOM_LEVEL}',
83 curses.color_pair(COLOR_COUNT),
84 )
85 screen.refresh()
86
87
88def handle_keyboard(char):
89 global X_MID
90 global Y_MID
91 global ZOOM_LEVEL
92 global MAX_ITERS
93 global MONOCHROME
94 global COLOR_COUNT
95
96 if char == 113: # Quit
97 return True
98 elif char == 258: # Down
99 Y_MID += Y_SCALE * (1 / (2 ** (ZOOM_LEVEL + 1)))
100 return False
101 elif char == 259: # Up
102 Y_MID -= Y_SCALE * (1 / (2 ** (ZOOM_LEVEL + 1)))
103 return False
104 elif char == 260: # Left
105 X_MID -= X_SCALE * (1 / (2 ** (ZOOM_LEVEL + 1)))
106 return False
107 elif char == 261: # Right
108 X_MID += X_SCALE * (1 / (2 ** (ZOOM_LEVEL + 1)))
109 return False
110 elif char == 40: # Zoom out
111 if ZOOM_LEVEL > 0:
112 ZOOM_LEVEL -= 1
113 return False
114 elif char == 41: # Zoom in
115 ZOOM_LEVEL += 1
116 return False
117 elif char == 61: # Zoom reset
118 ZOOM_LEVEL = 0
119 return False
120 elif char == 43: # Iterations up
121 MAX_ITERS += 1
122 return False
123 elif char == 45: # Iterations down
124 MAX_ITERS = 1 if MAX_ITERS == 1 else MAX_ITERS - 1
125 return False
126 elif char == 99: # Toggle colors
127 if MONOCHROME:
128 COLOR_COUNT = colors.init_101_colors()
129 MONOCHROME = False
130 else:
131 COLOR_COUNT = colors.init_101_monochrome()
132 MONOCHROME = True
133
134 return False
135
136
137def main(screen):
138 global COLOR_COUNT
139 rows, cols = screen.getmaxyx()
140 rows -= 1
141 COLOR_COUNT = colors.init_101_colors()
142 char = 0
143 while True:
144 matrix = calculate_set(cols, rows)
145 print_set(screen, cols, rows, matrix, char)
146 char = screen.getch()
147 if handle_keyboard(char):
148 break
149
150
151if __name__ == '__main__':
152 curses.wrapper(main)