first

fractal_generator.py

@@ -0,0 +1,82 @@
+from pathlib import Path
+
+import matplotlib.pyplot as plt
+import numpy as np
+import os
+from enum import Enum
+
+
+class FractalType(Enum):
+    Julia = 1
+    Mandelbrot = 2
+
+
+class FractalGenerator:
+    """ Creates a single fractal object and either returns it as as a numpy array, plot it or persists it as an pgn
+            image. The output of this class is used by FractalTrainingValidationSet to generate training/val sets
+            Args:
+                complex_function -- complex function to make a Julia fractal
+                n -- fractal size will ne n*n
+                xlim,ylim -- tuples with the plotting region on the complex plane
+                thr -- once a function grows larger that this number is considered to be divergent to infinity
+                max_iter -- number of compositions of the complex function with itself
+                type_ -- fractal type
+                fractal -- numpy array with the fractal
+        """
+
+    def __init__(self, n=256, xlim=(-2, 2), ylim=(-2, 2), thr=2, max_iter=10):
+        self.type_ = None
+        self.fractal = None
+        self.n = n
+        self.xlim = xlim
+        self.ylim = ylim
+        self.thr = thr
+        self.max_iter = max_iter
+
+    def create_julia(self, complex_function=lambda z: np.sin(z ** 4 + 1.41)):
+        """ Creates a fractal of the Julia family, the fractal is stored inside self.fractal """
+        fractal = np.zeros((self.n, self.n), dtype='complex')
+        x_space = np.linspace(self.xlim[0], self.xlim[1], self.n)
+        y_space = np.linspace(self.ylim[0], self.ylim[1], self.n)
+        for ix, x in enumerate(x_space):
+            for iy, y in enumerate(y_space):
+                for i in range(self.max_iter):
+                    if i == 0: z = complex(x, y)
+                    z = complex_function(z)
+                    if np.abs(z) >= self.thr: z = self.thr; break
+                fractal[ix, iy] = z
+        self.fractal = np.abs(fractal)
+        self.type_ = FractalType.Julia
+        return self
+
+    def create_mandelbrot(self):
+        """ Creates a fractal of the Mandelbrot family, the fractal is stored inside self.fractal """
+        fractal = np.zeros((self.n, self.n), dtype='complex')
+        x_space = np.linspace(self.xlim[0], self.xlim[1], self.n)
+        y_space = np.linspace(self.ylim[0], self.ylim[1], self.n)
+        for ix, x in enumerate(x_space):
+            for iy, y in enumerate(y_space):
+                for i in range(self.max_iter):
+                    if i == 0: z = 0
+                    z = z ** 2 + complex(x, y)
+                    if np.abs(z) >= self.thr: z = self.thr; break
+                fractal[ix, iy] = z
+        self.fractal = np.abs(fractal.transpose())
+        self.type_ = FractalType.Mandelbrot
+        return self
+
+    def plot(self, clim=None, **kwargs):
+        if self.fractal is None:
+            print('Nothing to plot. Generate a fractal first.')
+            return None
+        plt.matshow(self.fractal, **kwargs)
+        plt.gca().axes.get_xaxis().set_visible(False)
+        plt.gca().axes.get_yaxis().set_visible(False)
+        plt.clim(clim)
+        return plt.gcf()
+
+    def persist_plot(self, filename, container, clim=None, **kwargs):
+        if not os.path.isdir(container): os.mkdir(container)
+        self.plot(clim=clim, **kwargs)
+        plt.savefig(str(Path(container) / filename), png='png', dpi=None)
+        plt.close(plt.gcf())