add an f2py version

zingale · zingale · commit ed423d047cf3 · 2024-04-23T09:02:13.000-04:00
diff --git a/examples/extensions/f2py/README b/examples/extensions/f2py/README
@@ -0,0 +1,7 @@
+Note: python >= 3.12 needs ninja and meson to build.
+
+Build this as:
+
+f2py -c mandel.f90 -m mandel_f2py
+
+then you can `import mandel_f2py` in python.
diff --git a/examples/extensions/f2py/mandel.f90 b/examples/extensions/f2py/mandel.f90
@@ -0,0 +1,62 @@
+subroutine mandelbrot(N, xmin, xmax, ymin, ymax, max_iter, m)
+
+  implicit none
+
+  integer, intent(in) :: N
+  double precision, intent(in) :: xmin, xmax, ymin, ymax
+  integer, intent(in) :: max_iter
+
+  integer, intent(out) :: m(N, N)
+
+  double complex, parameter :: i_unit = (0, 1)
+
+!f2py depend(N) :: m
+!f2py intent(out) :: m
+
+  integer :: i, j, niter
+  double precision :: x(N), y(N)
+  double precision :: dx, dy
+  double complex, allocatable :: c(:, :)
+  double complex, allocatable :: z(:, :)
+
+  ! compute coordinates
+  dx = (xmax - xmin) / (N - 1)
+  dy = (ymax - ymin) / (N - 1)
+
+  do i = 1, N
+     x(i) = xmin + i * dx
+     y(i) = ymin + i * dy
+  enddo
+
+  allocate(c(N, N))
+
+  do j = 1, N
+     do i = 1, N
+        c(i, j) = x(i) + i_unit * y(j)
+     enddo
+  enddo
+
+  m(:, :) = 0
+
+  allocate(z(N, N))
+  z(:, :) = 0.0
+
+  do niter = 1, max_iter
+
+     do j = 1, N
+        do i = 1, N
+
+           if (m(i, j) == 0) then
+              z(i, j) = z(i, j) * z(i, j) + c(i, j)
+
+              if (abs(z(i,j)) > 2) then
+                 m(i, j) = niter
+              endif
+           endif
+
+        enddo
+     enddo
+
+  enddo
+
+end subroutine mandelbrot
diff --git a/examples/extensions/f2py/test_mandel.py b/examples/extensions/f2py/test_mandel.py
@@ -0,0 +1,27 @@
+import matplotlib.pyplot as plt
+import numpy as np
+
+import mandel_f2py
+
+import time
+
+start = time.time()
+
+xmin = -2.0
+xmax = 2.0
+ymin = -2.0
+ymax = 2.0
+
+max_iter = 50
+
+m = mandel_f2py.mandelbrot(1024, xmin, xmax, ymin, ymax, max_iter)
+
+print(f"execution time = {time.time() - start}\n")
+
+
+fig, ax = plt.subplots()
+ax.imshow(np.transpose(m), origin="lower",
+          extent=[xmin, xmax, ymin, ymax])
+
+fig.tight_layout()
+fig.savefig("test.png")
diff --git a/examples/extensions/python/mandel.py b/examples/extensions/python/mandel.py
@@ -17,7 +17,7 @@ def mandelbrot(N,
     m = np.zeros((N, N), dtype=np.int32)
 
     for i in range(max_iter):
-        z = z**2 + c
+        z[m == 0] = z[m == 0]**2 + c[m == 0]
 
         m[np.logical_and(np.abs(z) > 2, m == 0)] = i
 
