Upload mandelbrot_approximation.py

I trained a neural network to learn the fractal shape from pixel coordinates using PyTorch, NumPy, and Matplotlib. This isn’t curve-fitting. This is pure function learning.

- Added positional encodings to represent spatial patterns
- Increased epochs and depth for better generalization
- And tuned resolution + loss functions

The model approximated the Mandelbrot set with stunning accuracy, no image inputs, just math.


mandelbrot_approximation.py

@@ -0,0 +1,129 @@
+# -*- coding: utf-8 -*-
+"""mandelbrot-approximation.ipynb
+
+Automatically generated by Colab.
+
+Original file is located at
+    https://colab.research.google.com/drive/1M-DGWstZXLQZdi3hkl_r5P-oSiTcSSLA
+"""
+
+import numpy as np
+
+def mandelbrot_set(xmin, xmax, ymin, ymax, width, height, max_iter=100):
+
+  x = np.linspace(xmin, xmax, width)
+  y = np.linspace(ymin, ymax, height)
+
+  X, Y = np.meshgrid(x, y)
+  C = X + 1j * Y
+
+  iterations = np.zeros_like(C, dtype=np.float32)
+  Z = np.zeros_like(C)
+  mask = np.full(C.shape, True, dtype=bool)
+
+  for i in range(max_iter):
+    Z[mask] = Z[mask]**2 + C[mask]
+    escaped = np.abs(Z) > 2
+    iterations[escaped & mask] = i / max_iter
+    mask[escaped] = False
+
+  labels = iterations
+  return X, Y, labels
+
+def add_positional_encoding(X):
+    x, y = X[:, 0], X[:, 1]
+    encoded = np.stack([
+        x, y,
+        np.sin(2 * np.pi * x), np.cos(2 * np.pi * x),
+        np.sin(2 * np.pi * y), np.cos(2 * np.pi * y)
+    ], axis=1)
+    return encoded
+
+X, Y, labels = mandelbrot_set(-1.5, -0.8, -0.2, 0.2, 1000, 1000)
+X_train = np.stack([X.ravel(), Y.ravel()], axis=1)
+X_train_encoded = add_positional_encoding(X_train)
+y_train = labels.ravel()
+
+mask = (y_train > 0.05) & (y_train < 0.95)
+X_train_boundary = X_train[mask]
+y_train_boundary = y_train[mask]
+
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from torch.optim.lr_scheduler import StepLR
+from torch.utils.data import TensorDataset, DataLoader
+import matplotlib.pyplot as plt
+
+X_train_encoded = add_positional_encoding(X_train)
+X_tensor = torch.tensor(X_train_encoded, dtype=torch.float32)
+y_tensor = torch.tensor(y_train, dtype=torch.float32)
+
+dataset = TensorDataset(X_tensor, y_tensor)
+loader = DataLoader(dataset, batch_size = 1024, shuffle = True)
+
+class MandelbrotNet(nn.Module):
+  def __init__(self):
+    super().__init__()
+    self.model = nn.Sequential(
+    nn.Linear(6, 256),
+    nn.ReLU(),
+    *[layer for _ in range(6) for layer in (nn.Linear(256, 256), nn.ReLU())],
+    nn.Linear(256, 1),
+    nn.Identity()
+)
+
+  def forward(self, x):
+    return self.model(x)
+
+model = MandelbrotNet()
+criterion = nn.MSELoss()
+optimizer = optim.Adam(model.parameters(), lr=1e-3)
+scheduler = StepLR(optimizer, step_size=30, gamma=0.5)
+
+epochs = 100
+losses = []
+
+for epoch in range(epochs):
+    total_loss = 0.0
+    for batch_X, batch_y in loader:
+        preds = model(batch_X)
+        loss = criterion(preds, batch_y.unsqueeze(1))
+        optimizer.zero_grad()
+        loss.backward()
+        optimizer.step()
+
+        total_loss += loss.item()
+
+    scheduler.step()
+    avg_loss = total_loss / len(loader)
+    losses.append(avg_loss)
+    print(f"Epoch {epoch+1}/{epochs}, Loss: {avg_loss:.4f}")
+
+plt.plot(losses)
+plt.title("Training Loss")
+plt.xlabel("Epoch")
+plt.ylabel("MSE Loss")
+plt.grid(True)
+plt.show()
+
+with torch.no_grad():
+    y_pred = model(X_tensor).squeeze().numpy()
+
+# Reshape to image grid
+pred_image = y_pred.reshape(X.shape)   # shape: (height, width)
+true_image = labels
+fig, axs = plt.subplots(1, 2, figsize=(12, 5))
+
+axs[0].imshow(true_image, cmap='inferno', extent=[-2, 1, -1.5, 1.5])
+axs[0].set_title("Ground Truth Mandelbrot Set")
+
+axs[1].imshow(pred_image, cmap='inferno', extent=[-2, 1, -1.5, 1.5])
+axs[1].set_title("Model Approximation (DL)")
+
+for ax in axs:
+    ax.set_xlabel("Re")
+    ax.set_ylabel("Im")
+
+plt.tight_layout()
+plt.show()