<SYSTEM_TASK:>
Given the following text description, write Python code to implement the functionality described below step by step
<END_TASK>
<USER_TASK:>
Description:
 DistArray Julia Set
The Julia set, for a given complex number $c$, is the set of points $z$
such that $|z_{i}|$ remains bounded where $z_{i+1} = z_{i}^2 + c$.
This can be plotted by counting how many iterations are required for $|z_{i}|$ to exceed a cutoff.
Depending on the value of $c$, the Julia set may be connected and contain
a lot of points, or it could be disconnected and contain fewer points.
The points in the set will require the maximum iteration count, so
the connected sets will usually take longer to compute.
First, some imports.
Step3:  Julia set kernels
To avoid round-trips between the client and engines, our strategy will be to push a function to each engine to compute that engine's local section of the Julia set.  Here we have two options for this "kernel"
Step8:  Coordinating functions
Here we have functions that create a DistArray representing the complex plane and functions that coordinate applying the kernel on each distributed section of the DistArray.
Step9:  Using DistArray to explore the Julia set
<ASSISTANT_TASK:>
Python Code:
from __future__ import print_function
import numpy
import matplotlib.pyplot as plt
# DistArray imports
from distarray.globalapi import Context, Distribution
from distarray.globalapi.distarray import DistArray
# inline plots
%matplotlib inline
# bigger figures
plt.rcParams.update({'figure.figsize': (10, 10)})
Explanation: DistArray Julia Set
The Julia set, for a given complex number $c$, is the set of points $z$
such that $|z_{i}|$ remains bounded where $z_{i+1} = z_{i}^2 + c$.
This can be plotted by counting how many iterations are required for $|z_{i}|$ to exceed a cutoff.
Depending on the value of $c$, the Julia set may be connected and contain
a lot of points, or it could be disconnected and contain fewer points.
The points in the set will require the maximum iteration count, so
the connected sets will usually take longer to compute.
First, some imports.
End of explanation
def numpy_julia_calc(z, c, z_max, n_max):
    Calculate the Julia set using NumPy.
    Parameters
    ----------
    z : NumPy array
        array of complex values whose iterations we will count.
    c : complex
        Complex number to add at each iteration.
    z_max : float
        Magnitude of complex value that we assume goes to infinity.
    n_max : int
        Maximum number of iterations.
    
    z = numpy.asarray(z)
    counts = numpy.zeros_like(z, dtype=numpy.int32)
    hits = numpy.zeros_like(z, dtype=numpy.bool)
    mask = numpy.zeros_like(z, dtype=numpy.bool)
    n = 0
    while not numpy.all(hits) and n < n_max:
        z = z * z + c
        mask = (abs(z) > z_max) & (~hits)
        counts[mask] = n
        hits |= mask
        z[hits] = 0
        n += 1
    counts[~hits] = n_max
    return counts
def fancy_numpy_julia_calc(z, c, z_max, n_max):
    Calculate the Julia set using NumPy fancy indexing.
    Parameters
    ----------
    z : NumPy array
        array of complex values whose iterations we will count.
    c : complex
        Complex number to add at each iteration.
    z_max : float
        Magnitude of complex value that we assume goes to infinity.
    n_max : int
        Maximum number of iterations.
    
    z = numpy.asarray(z)
    counts = numpy.zeros_like(z, dtype=numpy.int32)
    hits = numpy.zeros_like(z, dtype=numpy.bool)
    mask = numpy.zeros_like(z, dtype=numpy.bool)
    n = 0
    while not numpy.all(hits) and n < n_max:
        z[~hits] = z[~hits] * z[~hits] + c
        mask = (abs(z) > z_max) & (~hits)
        counts[mask] = n
        hits |= mask
        n += 1
    counts[~hits] = n_max
    return counts
Explanation: Julia set kernels
To avoid round-trips between the client and engines, our strategy will be to push a function to each engine to compute that engine's local section of the Julia set.  Here we have two options for this "kernel": a version that avoids NumPy fancy indexing, and one that uses it.
End of explanation
def create_complex_plane(context, resolution, dist, re_ax, im_ax):
    Create a DistArray containing points on the complex plane.
    Parameters
    ----------
    context : DistArray Context
    resolution : 2-tuple
        The number of points along Re and Im axes.
    dist : 2-element sequence or dict
        dist_type for of the DistArray Distribution.
    re_ax : 2-tuple
        The (lower, upper) range of the Real axis.
    im_ax : 2-tuple
        The (lower, upper) range of the Imaginary axis.
    
    import numpy as np
    def fill_complex_plane(arr, re_ax, im_ax, resolution):
        Fill in points on the complex coordinate plane.
        re_step = float(re_ax[1] - re_ax[0]) / resolution[0]
        im_step = float(im_ax[1] - im_ax[0]) / resolution[1]
        for i in arr.distribution[0].global_iter:
            for j in arr.distribution[1].global_iter:
                arr.global_index[i, j] = complex(re_ax[0] + re_step * i,
                                                 im_ax[0] + im_step * j)
    # Create an empty distributed array.
    distribution = Distribution(context, (resolution[0], resolution[1]),
                                dist=dist)
    complex_plane = context.empty(distribution, dtype=np.complex64)
    context.apply(fill_complex_plane,
                  (complex_plane.key, re_ax, im_ax, resolution))
    return complex_plane
def local_julia_calc(la, c, z_max, n_max, kernel):
    Calculate the number of iterations for the point to escape.
    Parameters
    ----------
    la : LocalArray
        LocalArray of complex values whose iterations we will count.
    c : complex
        Complex number to add at each iteration.
    z_max : float
        Magnitude of complex value that we assume goes to infinity.
    n_max : int
        Maximum number of iterations.
    kernel : function
        Kernel to use for computation of the Julia set.  Options are 'fancy',
        'numpy', or 'cython'.
    
    from distarray.localapi import LocalArray
    counts = kernel(la, c, z_max, n_max)
    res = LocalArray(la.distribution, buf=counts)
    return proxyize(res)  # noqa
def distributed_julia_calc(distarray, c, z_max, n_max,
                           kernel=fancy_numpy_julia_calc):
    Calculate the Julia set for an array of points in the complex plane.
    Parameters
    ----------
    distarray : DistArray
        DistArray of complex values whose iterations we will count.
    c : complex
        Complex number to add at each iteration.
    z_max : float
        Magnitude of complex value that we assume goes to infinity.
    n_max : int
        Maximum number of iterations.
    kernel: function
        Kernel to use for computation of the Julia set.  Options are 'fancy',
        'numpy', or 'cython'.
    
    context = distarray.context
    iters_key = context.apply(local_julia_calc,
                              (distarray.key, c, z_max, n_max),
                              {'kernel': kernel})
    iters_da = DistArray.from_localarrays(iters_key[0], context=context,
                                          dtype=numpy.int32)
    return iters_da
Explanation: Coordinating functions
Here we have functions that create a DistArray representing the complex plane and functions that coordinate applying the kernel on each distributed section of the DistArray.
End of explanation
context = Context()  # handles communicating with the engines
# create the complex plane
resolution = (512, 512)
dist = 'cc'  # different distributions of the data among the engines
             # one character per dimension
             # 'c' stands for 'cyclic', 'b' for block
re_ax = (-1.5, 1.5)
im_ax = (-1.5, 1.5)
complex_plane = create_complex_plane(context, resolution, dist, re_ax, im_ax)
# calculate the Julia set
z_max = 2.0
n_max = 100
c = complex(0.285, 0.01)
iters_da = distributed_julia_calc(complex_plane, c, z_max, n_max, kernel=numpy_julia_calc)
# plot it!
plt.matshow(iters_da.toarray(), cmap='hot')
context.close()
Explanation: Using DistArray to explore the Julia set
End of explanation
<END_TASK>