"Ray mandelbrot" (#291)

* "Ray mandelbrot"

* Added readme and test

* fix mandelbrot path in test

* include matplotlib

---------

Co-authored-by: Ritvik Rao <rsrao2@illinois.edu>

Author: rik404

.github/workflows/charm4py.yml

@@ -31,7 +31,7 @@ jobs:
           python-version: ${{ matrix.python-version }}
       - name: Install dependencies
         run: |
-          pip install setuptools cython cffi greenlet numpy torch torchvision filelock
+          pip install setuptools cython cffi greenlet numpy torch torchvision filelock matplotlib
           if [ ${{ matrix.os }} == 'macos-13' ]; then
             # pypi only distributes torch packages w/ numpy v1 for macos-x86_64
             pip install 'numpy<2'

examples/ray/mandelbrot/README.rst

@@ -0,0 +1,16 @@
+==================
+Mandelbrot example
+==================
+
+mandelbrot.py <height> <width> <max_iter> <tile_size>
+    - height - height of image
+    - width - width of image
+    - max_iter - determines the "defintion" of zoom
+    - tile_size - parallelization factor. each tile is computed in parallel
+
+    Generates mandelbrot.png - Computes a high-resolution image of the Mandelbrot set using parallel processing. It divides the image into smaller rectangular tiles and distributes these tiles across multiple worker tasks using Ray-Charm4py. Each task calculates whether each pixel in its tile belongs to the Mandelbrot set, storing the result in a memory-mapped file to avoid using excessive RAM. This approach enables efficient generation of very large fractal images by balancing performance and memory usage.
+
+.. image:: mandelbrot_ray.png
+   :alt: Mandelbrot Set Visualization
+   :align: center
+    

examples/ray/mandelbrot/mandelbrot.py

@@ -0,0 +1,71 @@
+from charm4py import ray, charm
+import numpy as np
+import matplotlib.pyplot as plt
+import os
+
+# Compute whether a point is in the Mandelbrot set
+def mandelbrot_fast(re, im, max_iter):
+    zr = zi = 0.0
+    for i in range(max_iter):
+        zr2 = zr * zr
+        zi2 = zi * zi
+        if zr2 + zi2 > 4.0:
+            return i
+        zi = 2 * zr * zi + im
+        zr = zr2 - zi2 + re
+    return max_iter
+
+# Remote task to compute a tile
+@ray.remote
+def compute_tile(x_start, x_end, y_start, y_end, width, height, max_iter):
+    tile = np.zeros((y_end - y_start, x_end - x_start), dtype=np.uint16)
+    for y in range(y_start, y_end):
+        for x in range(x_start, x_end):
+            re = 3.5 * (x / width) - 2.5
+            im = 2.0 * (y / height) - 1.0
+            tile[y - y_start, x - x_start] = mandelbrot_fast(re, im, max_iter)
+    return tile
+
+def generate_mandelbrot_image_optimized(width=12000, height=8000, max_iter=200, tile_size=1000, max_pending=1000):
+    # Pre-create the empty file with the correct size
+    total_bytes = 2 * width * height  # 2 bytes per pixel (uint16)
+    with open("output/mandelbrot_large.dat", "wb") as f:
+        f.seek(total_bytes - 1)
+        f.write(b'\0')
+    result_image = np.memmap("output/mandelbrot_large.dat", dtype=np.uint16, mode='w+', shape=(height, width))
+    pending = []
+
+    for y in range(0, height, tile_size):
+        for x in range(0, width, tile_size):
+            x_end = min(x + tile_size, width)
+            y_end = min(y + tile_size, height)
+            tile_ref = compute_tile.remote(x, x_end, y, y_end, width, height, max_iter)
+            pending.append(((x, y), tile_ref))
+
+            if len(pending) >= max_pending:
+                (x0, y0), tile = pending.pop(0)
+                tile = ray.get(tile)
+                result_image[y0:y0+tile.shape[0], x0:x0+tile.shape[1]] = tile
+
+    for (x0, y0), tile_ref in pending:
+        tile = ray.get(tile_ref)
+        result_image[y0:y0+tile.shape[0], x0:x0+tile.shape[1]] = tile
+
+    return result_image
+
+
+def main(args):
+    output_path = "output/mandelbrot_large.dat"
+    os.makedirs(os.path.dirname(output_path), exist_ok=True)
+    ray.init()
+    # Run the benchmark
+    image = generate_mandelbrot_image_optimized(width=int(args[1]), height=int(args[2]), max_iter=int(args[3]), tile_size=int(args[4]))
+    # Optional: show the result
+    plt.imshow(image, cmap='hot')
+    plt.title("Mandelbrot Set (Ray)")
+    plt.axis('off')
+    plt.savefig("mandelbrot_ray.png", dpi=300, bbox_inches='tight')
+    os.remove('output/mandelbrot_large.dat')
+    charm.exit()
+
+charm.start(main)

examples/ray/mandelbrot/mandelbrot_ray.png

@@ -312,5 +312,9 @@
         "condition": "numbaInstalled",
         "path": "examples/wave2d/wave2d.py",
         "args": "300 -1 --NO-RENDER"
+    },
+    {
+        "path": "examples/ray/mandelbrot/mandelbrot.py",
+        "args": "100 100 10 10"
     }
 ]