import numpy as np


class CahnHilliardSimulation(object):
    r"""
    Solve the `Cahn-Hilliard equation
    <https://en.wikipedia.org/wiki/Cahn-Hilliard_equation>`__ for
    multiple samples in arbitrary dimensions. The concentration varies
    from -1 to 1. The equation is given by

    .. math::

       \dot{\phi} = \nabla^2 \left( \phi^3 - \phi \right) - \gamma \nabla^4 \phi

    The discretiztion scheme used here is from `Chang and Rutenberg
    <http://dx.doi.org/10.1103/PhysRevE.72.055701>`__. The scheme is a
    semi-implicit discretization in time and is given by

    .. math::

       \phi_{t+\Delta t}
       + \left(1 - a_1\right) \Delta t \nabla^2 \phi_{t+\Delta t}
       + \left(1 - a_2\right) \Delta t \gamma \nabla^4 \phi_{t+\Delta t}
       = \phi_t
       - \Delta t \nabla^2 \left(a_1 \phi_t + a_2 \gamma \nabla^2 \phi_t - \phi_t^3 \right)

    where :math:`a_1=3` and :math:`a_2=0`.
    
    In 1D. 
    
    >>> np.random.seed(99)
    >>> N = 100
    >>> phi = 0.01 * (2 * np.random.random((2, N)) - 1)
    >>> ch = CahnHilliardSimulation(gamma=4)
    >>> for i in range(10000):
    ...     ch.run(phi)
    ...     phi[:] = ch.response
    >>> assert (max(phi.flat) > 2e-3) and (min(phi.flat) < -2e-3)

    In 2D.

    >>> phi = 0.01 * (2 * np.random.random((2, N, N)) - 1)
    >>> ch = CahnHilliardSimulation(gamma=4.)
    >>> for i in range(100):
    ...     ch.run(phi)
    ...     phi[:] = ch.response
    >>> assert (max(phi.flat) > 0.001) and (min(phi.flat) < -0.001)

    In 3D.

    >>> phi = 0.01 * (2 * np.random.random((2, N, N, N)) - 1)
    >>> ch = CahnHilliardSimulation(gamma=4.)
    >>> for i in range(10):
    ...     ch.run(phi)
    ...     phi[:] = ch.response

    >>> assert (max(phi.flat) > 0.0005) and (min(phi.flat) < -0.0005)

    """

    def __init__(self, dx=0.25, gamma=1., dt=0.001):
        r"""
        Args:
          dx: grid spacing
          dt: time step size
          gamma: paramater in CH equation

        """
        self.dx = dx
        self.dt = dt
        self.gamma = gamma

    def run(self, X):
        r"""
        Return the response field

        Args:
          X: Array representing the concentration field between -1 and
             1 with shape (n_samples, N, N)

        Returns:
          Array representing the microstructure at one time step ahead
          of 'X'

        """
        N = X.shape[1]
        if not np.all(np.array(X.shape[1:]) == N):
            raise RuntimeError("X must represent a square domain")

        L = self.dx * N
        k = np.arange(N)
        k[N / 2:] = (k - N / 2)[:N / 2]
        k = k * 2 * np.pi / L

        i_ = np.indices(X.shape[1:])
        ksq = np.sum(k[i_] ** 2, axis=0)[None]

        axes = np.arange(len(X.shape) - 1) + 1
        FX = np.fft.fftn(X, axes=axes)
        FX3 = np.fft.fftn(X ** 3, axes=axes)

        a1 = 3.
        a2 = 0.
        explicit = ksq * (a1 - self.gamma * a2 * ksq)
        implicit = ksq * ((1 - a1) - self.gamma * (1 - a2) * ksq)
        dt = self.dt

        Fy = (FX * (1 + dt * explicit) - ksq * dt * FX3) / (1 - dt * implicit)
        self.response = np.fft.ifftn(Fy, axes=axes).real


@@ -2249,114 +2249,8 @@
 N)%0A%0A
-        Returns:%0A          Array representing the microstructure at one time step ahead%0A          of 'X'%0A%0A
     

Fix docs