landscape.py

'''Single species, one type of patch and matrix - basic module.'''
import numpy as np
from collections import OrderedDict
from itertools import product as iproduct

try:
    from matplotlib import pyplot as plt
except ImportError:
    pass

## example parameters
#par = OrderedDict([
#    ('r', 0.1),
#    ('K', 1.),
#    ('mu', 0.03),
#    ('Dp', 1e-4),
#    ('Dm', 1e-3),
#    # interface condition (=Dm/Dp)
#    #('g', 1.), # no longer required
#    # boundary conditions
#    ('left', [1., 0., 0.]),
#    ('right', [1., 0., 0.]),
#    ('top', [1., 0., 0.]),
#    ('bottom', [1., 0., 0.])
#    ])
# in km
#dx = 0.0286
# 35 pixels -> 1km = 1000m
# 1 pixel -> 0.0286km
# D: 10^-2 km^2/day
# r: 10-50/year ~ 0.1 day^-1
# mu: 1/30 days ~ 0.03 day^-1

## dados novos

# in km
dx = 0.02
# 1 pixel -> 0.02km
# Dm ~ 900 m^2/day ~ 10^-3 km^2/day
# Dp ~ 5 m^2/day ~ 10^-5 km^2/day
# r: 10-50/year ~ 0.1 day^-1
# mu: 1/30 days ~ 0.03 day^-1


def solve_landscape(landscape, par, dx, f_tol=None, force_positive=False, verbose=True):
    r"""Find the stationary solution for a given landscape and set of parameters.

    Uses a Newton-Krylov solver with LGMRES sparse inverse method to find a
    stationary solution (or the solution to the elliptical problem) to the
    system of equations in 2 dimensions (x is a 2-d vector):

    .. math::
        u_t(x) &= D_p \nabla^2 u(x) + ru(1-u(x)/K) = 0 \text{ in a patch} \\
        v_t(x) &= D_m \nabla^2 v(x) - \mu v(x) = 0 \text{ in the matrix}

    Notes
    -----
    This function preserves the original interface, but internally it calls the
    newer `solve_landscape_ntypes()`.


    Parameters
    ----------
    landscape : 2-d array of ints
        describe the landscape, with 1 on patches and 0 on matrix
    par : dict
        parameters (dict keys): 

        - r : reproductive rate on patches
        - K : carrying capacity on patches
        - mu : mortality rate in the matrix
        - Dp : diffusivity on patches
        - Dm : diffusivity in the matrix
        - g : habitat discontinuity parameter \gamma. See interface conditions
          below (optional, determined from Dp, Dm and alpha if absent)
        - alpha : habitat preference, only taken into account if g is not
          present. In that case, g is calculated as g = Dm * alpha /
          (Dp*(1-alpha)) (optional, ignored if g is given, set to 1/2 if both
          are absent)
        - left : (a, b, c): external boundary conditions at left border
        - right : (a, b, c): external boundary conditions at right border
        - top : (a, b, c): external boundary conditions at top border
        - bottom : (a, b, c): external boundary conditions at bottom border
    dx : float
        lenght of each edge
    f_tol : float
        tolerance for the residue, passed on to the solver routine.  Default is
        6e-6
    force_positive : bool
        make sure the solution is always non-negative - in a hacky way. Default
        False
    verbose : bool
        print residue of the solution and its maximum and minimum values

    Returns
    -------
    solution : 2-d array of the same shape of the landscape input 
        the solution


    .. rubric:: Boundary and interface conditions

    External boundaries are of the form

    .. math:: a \nabla u \cdot \hat{n} + b u + c = 0

    and may be different for left, right, top, bottom.  The derivative of u is
    taken along the normal to the boundary.

    The interfaces between patches and matrix are given by

    .. math::
        u(x) &= \gamma v(x) \\
        D_p \nabla u(x) \cdot \hat{n} &= D_m \nabla v(x) \cdot \hat{n}

    where u is in a patch and v is the solution in the matrix. Usually the
    discontinuity $\gamma$ is a result of different diffusivities and
    preference at the border (see Ovaskainen and Cornell 2003). In that case,
    given a preference $\alpha$ (between 0 and 1, exclusive) towards the patch,
    this parameter should be:

    .. math:: \gamma = \frac{D_m}{D_p} \frac{\alpha}{1-\alpha}

    This last condition is used in case $\gamma$ is not set. If $\alpha$ isn't
    set either, it's assumed $\alpha = 1/2$. These conditions are handled using
    an asymetric finite difference scheme for the 2nd derivative:

    .. math:: u_{xx}(x) = \frac{4}{3h^2} (u(x-h) - 3 u(x) + 2 u(x+h/2))

    At the interface, $u(x+h/2)$ and $v(x+h/2)$ must obey:

    .. math::
        u(x+h/2) &= \gamma v(x+h/2) \\
        D_p (u(x+h/2) - u(x))  &= D_m (v(x+h) - v(x+h/2))

    Solving this system, we arrive at the approximation at the interface:

    .. math:: u(x+h/2) = \frac{D_m v(x+h)+D_p u(x)}{D_p+D_m / \gamma}

    if u(x) is in a patch and v(x+h) is in the matrix, or

    .. math:: v(x+h/2) = \frac{D_m v(x)+D_p u(x+h)}{D_p \gamma +D_m}

    if v(x) is in the matrix and u(x+h) is in a patch.

    References
    ----------
    Ovaskainen, Otso, and Stephen J. Cornell. "Biased movement at a boundary
    and conditional occupancy times for diffusion processes." Journal of
    Applied Probability 40.3 (2003): 557-580.

    """
    p = OrderedDict([
        ('r', [-par['mu'], par['r']]),
        ('K', [np.Inf, par['K']]),
        ('D', [par['Dm'], par['Dp']]),
        ('left', par['left']),
        ('right', par['right']),
        ('top', par['top']),
        ('bottom', par['bottom'])
        ])

    return solve_landscape_ntypes(landscape, p, dx, f_tol, force_positive,
                                  verbose)


def solve_landscape_ntypes(landscape, par, dx, f_tol=None,
        force_positive=False, skip_refine=False, return_residual=False,
        verbose=True):
    r"""Find the stationary solution for a landscape with many types of habitat.

    Uses a Newton-Krylov solver with LGMRES sparse inverse method to find a
    stationary solution (or the solution to the elliptical problem) to the
    system of equations in 2 dimensions (x is a 2-d vector):

    .. math:: 
        \frac{\partial u_i}{\partial t} = D_i \nabla^2 u_i + 
        r_i u_i\left(1-\frac{u}{K_i}\right) = 0

    Parameters
    ----------
    landscape : 2-d array of ints
        describe the landscape, with any number of habitat types
    par : dict
        parameters (dict keys):

        - r : growth rates (can be negative)
        - K : carrying capacities (cn be np.Inf)
        - mu : mortality rate in the matrix
        - D : diffusivities
        - g : dict of habitat discontinuities $\gamma_{ij}$ - see interface
          conditions below. The keys are tuples (i,j) of the habitat
          types indices (optional)
        - alpha : dict of habitat preferences, only taken into account if g is
          not present. In that case, $\gamma_{ij}$ is calculated as
          $\gamma_{ij} = D_j \alpha_{ij} / (D_i*(1-\alpha_{ij}))$ (optional)
        - left : (a, b, c): external boundary conditions at left border
        - right : (a, b, c): external boundary conditions at right border
        - top : (a, b, c): external boundary conditions at top border
        - bottom : (a, b, c): external boundary conditions at bottom border
    dx : float
        lenght of each edge
    f_tol : float
        tolerance for the residue, passed on to the solver routine.  Default is
        6e-6
    force_positive : bool
        make sure the solution is always non-negative - in a hacky way. Default
        False
    skip_refine : bool
        do not refine the grid to calculate the residual. This can greatly
        improve speed, but will generate errors (even silent wrong results) if
        the landscape has contiguous interfaces
    return_residual : bool
        returns only the residual function, without calculating the solution
    verbose : bool
        print residue of the solution and its maximum and minimum values

    Returns
    -------
    solution : 2-d array of the same shape of the landscape input 
        the solution


    .. rubric:: Boundary and interface conditions

    External boundaries are of the form

    .. math:: a \nabla u \cdot \hat{n} + b u + c = 0

    and may be different for left, right, top, bottom.  The derivative of u is
    taken along the normal to the boundary.

    The interfaces between patches and matrix are given by

    .. math::
        u_i(x) &= \gamma_{ij} u_j(x) \\
        D_i \nabla u_i(x) \cdot \hat{n} &= D_j \nabla u_j(x) \cdot \hat{n}

    Usually the discontinuity $\gamma_{ij}$ is a result of different
    diffusivities and preference at the border (see Ovaskainen and Cornell
    2003). In that case, given a preference $\alpha_{ij}$ (between 0 and 1,
    exclusive) towards $i$, this parameter should be:

    .. math:: \gamma_{ij} = \frac{D_j}{D_i} \frac{\alpha_{ij}}{1-\alpha_{ij}}

    This last condition is used in case $\gamma$ is not set. If $\alpha$ isn't
    set either, it's assumed $\alpha = 1/2$. Notice that $\alpha_{ij} +
    \alpha_{ji} = 1$, and so $\gamma_{ij} = \gamma_{ji}^{-1}$. To ensure this
    condition, the key (i,j) is always taken with $i>j$.

    These conditions are handled using an asymetric finite difference scheme
    for the 2nd derivative:

    .. math:: u_{xx}(x) = \frac{4}{3h^2} (u(x-h) - 3 u(x) + 2 u(x+h/2))

    At the interface, $u(x+h/2)$ and $v(x+h/2)$ must obey:

    .. math::
        u(x+h/2) &= \gamma v(x+h/2) \\
        D_p (u(x+h/2) - u(x))  &= D_m (v(x+h) - v(x+h/2))

    Solving this system, we arrive at the approximation at the interface:

    .. math:: u(x+h/2) = \frac{D_m v(x+h)+D_p u(x)}{D_p+D_m / \gamma}

    if u(x) is in a patch and v(x+h) is in the matrix, or

    .. math:: v(x+h/2) = \frac{D_m v(x)+D_p u(x+h)}{D_p \gamma +D_m}

    if v(x) is in the matrix and u(x+h) is in a patch.

    Example
    -------
    >>> # simple patch/matrix
    >>> from landscape import *
    >>> parn = OrderedDict([
        ('r', [-0.03, 0.1]),
        ('K', [np.Inf, 1.0]),
        ('D', [0.001, 0.0001]),
        ('left', [1.0, 0.0, 0.0]),
        ('right', [1.0, 0.0, 0.0]),
        ('top', [1.0, 0.0, 0.0]),
        ('bottom', [1.0, 0.0, 0.0])
        ])
    >>> l = np.zeros((100,100), dtype=int)
    >>> l[40:60, 40:60] = 1
    >>> sol = solve_landscape_ntypes(l, parn, dx)

    """
    from scipy.optimize import newton_krylov

    if not skip_refine:
        # refine the grid to avoid contiguous interfaces
        landscape = refine_grid(landscape)
        dx /= 2

    n = np.unique(landscape).astype(int)

    p = par.copy()
    if 'g' not in p.keys():
        p['g'] = {}
        if 'alpha' not in p.keys():
            p['alpha'] = {}
            for i, j in iproduct(n, repeat=2):
                if i > j:
                    p['alpha'][(i,j)] = 0.5
        for i, j in iproduct(n, repeat=2):
            if i > j:
                p['g'][(i,j)] = p['D'][j]/p['D'][i] * \
                                  p['alpha'][(i,j)] / (1-p['alpha'][(i,j)])

    # this ensures the consistency of the interface discontinuities
    # it ignores the values of g_ij with i < j, replacing it by 1/g_ji
    for i, j in iproduct(n, repeat=2):
        if i < j:
            p['g'][(i,j)] = 1/p['g'][(j,i)]

    (al, bl, cl) = p['left']
    (ar, br, cr) = p['right']
    (at, bt, ct) = p['top']
    (ab, bb, cb) = p['bottom']

    D = np.zeros_like(landscape, dtype=np.float_)
    r = np.zeros_like(landscape, dtype=np.float_)
    c = np.zeros_like(landscape, dtype=np.float_)
    for i in n:
        li = np.where(landscape == i)
        D[li] = p['D'][i]
        r[li] = p['r'][i]
        c[li] = - p['r'][i] / p['K'][i]

    Bx, By = find_interfaces_ntypes(landscape)
    factor = {}
    for i, j in iproduct(n, repeat=2):
        if i != j:
            factor[(i,j)] = (
                # coefficient of term u(x) in u_xx(x)
                -2. + 8./3 * p['D'][i]/(p['D'][i]+p['D'][j]/p['g'][(i,j)]),
                # coefficient of term u(x+h) in u_xx(x)
                -1. + 8./3 * p['D'][j]/(p['D'][i]+p['D'][j]/p['g'][(i,j)])
                )

    Bxleft = np.zeros_like(landscape, dtype=np.float_)
    Bxcenter = np.zeros_like(landscape, dtype=np.float_)
    Bxright = np.zeros_like(landscape, dtype=np.float_)
    Byleft = np.zeros_like(landscape, dtype=np.float_)
    Bycenter = np.zeros_like(landscape, dtype=np.float_)
    Byright = np.zeros_like(landscape, dtype=np.float_)
    for i,j in factor.keys():
        ## direction x
        # patch type i
        Bxcenter[Bx[i,j]] += factor[i,j][0]
        Bxleft[shifted_index(Bx[i,j], 0, -1)] += 1./3
        Bxright[shifted_index(Bx[i,j], 0, 1)] += factor[i,j][1]
        # patch type j
        Bxcenter[shifted_index(Bx[i,j], 0, 1)] += factor[j,i][0]
        Bxleft[Bx[i,j]] += factor[j,i][1]
        Bxright[shifted_index(Bx[i,j], 0, 2)] += 1./3
        ## direction y
        # patch type i
        Bycenter[By[i,j]] += factor[i,j][0]
        Byleft[shifted_index(By[i,j], 1, -1)] += 1./3
        Byright[shifted_index(By[i,j], 1, 1)] += factor[i,j][1]
        # patch type j
        Bycenter[shifted_index(By[i,j], 1, 1)] += factor[j,i][0]
        Byleft[By[i,j]] += factor[j,i][1]
        Byright[shifted_index(By[i,j], 1, 2)] += 1./3

    def residual(P):
        if force_positive:
            P = np.abs(P)

        d2x = np.zeros_like(P)
        d2x[1:-1,:] = P[2:,:] - 2*P[1:-1,:] + P[:-2,:]
        # external boundaries
        d2x[0,:] = P[1,:] - 2*P[0,:] + (-cl - al/dx * P[0,:])/(bl - al/dx)
        d2x[-1,:] = P[-2,:] - 2*P[-1,:] + (-cr + ar/dx * P[-1,:])/(br + ar/dx)
        # interface conditions
        d2x[1:-1,:] += (
                Bxcenter[1:-1,:] * P[1:-1,:] +
                Bxleft[:-2,:] * P[:-2,:] +
                Bxright[2:,:] * P[2:,:]
                )

        d2y = np.zeros_like(P)
        d2y[:,1:-1] = P[:,2:] - 2*P[:,1:-1] + P[:,:-2]
        # external boundaries
        d2y[:,0] = P[:,1] - 2*P[:,0] + (-cb - ab/dx * P[:,0])/(bb - ab/dx)
        d2y[:,-1] = P[:,-2] - 2*P[:,-1] + (-ct + at/dx * P[:,-1])/(bt + at/dx)
        # interface conditions
        d2y[:,1:-1] += (
                Bycenter[:,1:-1] * P[:,1:-1] +
                Byleft[:,:-2] * P[:,:-2] +
                Byright[:,2:] * P[:,2:]
                )

        return D*(d2x + d2y)/dx/dx + r*P + c*P**2

    if return_residual:
        return residual

    # solve
    guess = r.copy()
    guess[guess>0] = 1/((-c/r)[guess>0])
    guess[guess<=0] = 1e-6
    sol = newton_krylov(residual, guess, method='lgmres', f_tol=f_tol)
    if force_positive:
        sol = np.abs(sol)
    if verbose:
        print('Residual: %e' % abs(residual(sol)).max())
        print('max. pop.: %f' % sol.max())
        print('min. pop.: %f' % sol.min())
    if not skip_refine:
        sol = coarse_grid(sol)

    return sol

def solve_landscape_ntypes_nspecies(landscape, par, dx, f_tol=None,
        force_positive=False, skip_refine=False, return_residual=False,
        verbose=True):
    '''Find the stationary solution for a given landscape and set of parameters.

    Uses a Newton-Krylov solver with LGMRES sparse inverse method to find a
    stationary solution (or the solution to the elliptical problem) to the
    system of 2n equations in 2 dimensions (x is a 2-d vector):

    .. math::
        \\frac{\partial u_{ik}}{\partial t} &= D_k \\nabla^2 u_{ik} + r_{ik} u_{ik} (1-\sum_{j=1}^n \\alpha_j u_{jk}) = 0

    where i runs over the n species and k over the patch types.

    Parameters
    ----------
    landscape : 2-d array of ints
        describe the landscape, with any number of habitat types
    par : list
        the first element is the matrix (or dict with tuple keys) of
        competition coefficients, including the inverse of carrying capacities,
        and the following elements are dicts with parameters as in the
        `solve_landscape_ntypes()` function, except for the carrying capacity.
    dx : float
        length of each edge
    f_tol : float
        tolerance for the residue, passed on to the solver routine.  Default is
        6e-6
    force_positive : bool
        make sure the solution is always non-negative - in a hacky way. Default
        False
    skip_refine : bool
        do not refine the grid to calculate the residual. This can greatly
        improve speed, but will generate errors (even silent wrong results) if
        the landscape has contiguous interfaces
    return_residual : bool
        returns only the residual function, without calculating the solution
    verbose : bool
        print residue of the solution and its maximum and minimum values

    Returns
    -------
    solution : 2-d array of the same shape of the landscape input containing
        the solution

    Boundary and interface conditions
    ---------------------------------
    External boundaries are of the form

    .. math::
        a \\nabla u \cdot \hat{n} + b u + c = 0

    and may be different for left, right, top, bottom.  The derivative of u is
    taken along the normal to the boundary.

    The interfaces between patches and matrix are given by

    .. math::
        u(x) &= \gamma v(x) \\\\
        D_p \\nabla u(x) \cdot \hat{n} &= D_m \\nabla v(x) \cdot \hat{n}

    where u is a patch and v is the solution in the matrix.

    Example
    -------
    >>> from landscape import *
    >>> lA = loadtxt('landA.txt')
    >>> par = [
            [
                [[0, 0],
                 [0.1, 0.1]],
                [[0.1, 0],
                 [0.1, 0.1]]
            ],
            {'r': [-0.03, 0.1],
             'D': [0.001, 0.0001],
             'left': [1.0, 0.0, 0.0],
             'right': [1.0, 0.0, 0.0],
             'top': [1.0, 0.0, 0.0],
             'bottom': [1.0, 0.0, 0.0]},
            {'r': [0.05, 0.05],
             'D': [0.001, 0.001],
             'left': [1.0, 0.0, 0.0],
             'right': [1.0, 0.0, 0.0],
             'top': [1.0, 0.0, 0.0],
             'bottom': [1.0, 0.0, 0.0]}
            ]
    >>> sol = solve_landscape_ntypes_nspecies(lA, par, dx)

    '''
    from scipy.optimize import newton_krylov

    if not skip_refine:
        # refine the grid to avoid contiguous interfaces
        landscape = refine_grid(landscape)
        dx /= 2

    N = len(par) - 1
    n = np.unique(landscape).astype(int)

    resi = [ solve_landscape_ntypes(landscape, dict(p,
        K=np.Inf*np.ones(len(n))), dx, skip_refine=True, return_residual=True)
        for p in par[1:] ]

    sec_term = np.zeros((N, N, landscape.shape[0], landscape.shape[1]))
    for i,j in iproduct(range(N), range(N)):
        for k in n:
            lk = np.where(landscape == k)
            sec_term[i,j][lk] = par[0][k][i][j]

    def residual(P):
        if force_positive:
            P = np.abs(P)
        res = np.zeros_like(P)
        # loops are for lazy people
        for i, Pi in enumerate(P):
            res[i,:,:] = resi[i](Pi) - Pi * (sec_term[i,:,:] * P).sum(axis=0)
        return res

    if return_residual:
        return residual

    # build guess based on where growth is positive
    guess = np.zeros((N, landscape.shape[0], landscape.shape[1]))
    for i in range(N):
        for k in n:
            lk = np.where(landscape == k)
            rik = par[i+1]['r'][k]
            if rik <= 0:
                guess[i,:,:][lk] = 1e-6
            else:
                guess[i,:,:][lk] =  rik / par[0][k][i][i]

    # solve
    sol = newton_krylov(residual, guess, method='lgmres', f_tol=f_tol)
    if force_positive:
        sol = np.abs(sol)
    if verbose:
        print('Residuals:', *[ '%e' % i for i in
            np.abs(residual(sol)).max(axis=(1,2)) ])
        print('max. pops.:', *[ '%f' % i for i in sol.max(axis=(1,2)) ])
        print('min. pops.:', *[ '%f' % i for i in sol.min(axis=(1,2)) ])

    if not skip_refine:
        sol = coarse_grid(sol)

    return sol

def solve_landscape_nspecies(landscape, par, dx, f_tol=None,
        force_positive=False, skip_refine=False, verbose=True):
    '''Find the stationary solution for a given landscape and set of parameters.

    Uses a Newton-Krylov solver with LGMRES sparse inverse method to find a
    stationary solution (or the solution to the elliptical problem) to the
    system of 2n equations in 2 dimensions (x is a 2-d vector):

    .. math::
        \\frac{\partial u_i}{\partial t} &= D_p \\nabla^2 u_i + r_i u_i (1-\sum_{j=1}^n \\alpha_j u_j) = 0 \\text{ in a patch} \\\\
        \\frac{\partial v_i}{\partial t} &= D_m \\nabla^2 v_i - \mu_i v_i = 0 \\text{ in the matrix}

    Notes
    -----
    This function preserves the original interface, but internally it calls the
    newer `solve_landscape_ntypes_nspecies()`.

    Parameters
    ----------
    landscape : a 2-d array (of ints) describing the landscape, with 1 on
        patches and 0 on matrix
    par : a ordered dict containing parameters in the following order:
        r: list of reproductive rates on patches
        alpha: matrix of interaction parameters on patches (diagonals are minus the inverse of carrying capacity)
        mu: list of mortality rates in the matrix
        Dp: list of diffusivities on patches
        Dm: list of diffusivities in the matrix
        g: habitat preference parameter \gamma, usually less than one. See interface conditions below
        left: (a, b, c): external boundary conditions at left border
        right: (a, b, c): external boundary conditions at right border
        top: (a, b, c): external boundary conditions at top border
        bottom: (a, b, c): external boundary conditions at bottom border

    dx : float
        length of each edge
    f_tol : float
        tolerance for the residue, passed on to the solver routine.  Default is
        6e-6
    force_positive : bool
        make sure the solution is always non-negative - in a hacky way. Default
        False
    skip_refine : bool
        do not refine the grid to calculate the residual. This can greatly
        improve speed, but will generate errors (even silent wrong results) if
        the landscape has contiguous interfaces
    verbose : bool
        print residue of the solution and its maximum and minimum values

    Returns
    -------
    solution : 2-d array of the same shape of the landscape input containing
        the solution

    Boundary and interface conditions
    ---------------------------------
    External boundaries are of the form

    .. math::
        a \\nabla u \cdot \hat{n} + b u + c = 0

    and may be different for left, right, top, bottom.  The derivative of u is
    taken along the normal to the boundary.

    The interfaces between patches and matrix are given by

    .. math::
        u(x) &= \gamma v(x) \\\\
        D_p \\nabla u(x) \cdot \hat{n} &= D_m \\nabla v(x) \cdot \hat{n}

    where u is a patch and v is the solution in the matrix. These conditions
    are handled using an assymetric finite difference scheme for the 2nd
    derivative:

    .. math::
        u_xx(x) = (4/3/h**2) (u(x-h) - 3 u(x) + 2 u(x+h/2))

    with the approximations at the interface:

    .. math::
        u(x+h/2) = (Dm*v(x+h)+Dp*u(x))/(Dp+Dm*g)

    if u(x) is in a patch and v(x+h) is in the matrix, or

    .. math::
        v(x+h/2) = g*(Dm*v(x)+Dp*u(x+h))/(Dp+Dm*g)

    if v(x) is in the matrix and u(x+h) is in a patch.

    Example
    -------
    >>> from landscape import *
    >>> parn = OrderedDict([
        ('rp', [0.1, 0.05]),
        ('rm', [-0.03, 0.05]),
        ('alphap', [[0.1, 0], [0.1, 0.1]]),
        ('alpham', [[0., 0.], [0.1, 0.1]]),
        ('Dp', [1e-4, 1e-3]),
        ('Dm', [1e-3, 1e-3]),
        # interface condition can be omitted!
        #('g', [.1, 1.]),
        # boundary conditions
        ('left', [1., 0., 0.]),
        ('right', [1., 0., 0.]),
        ('top', [1., 0., 0.]),
        ('bottom', [1., 0., 0.])
        ])
    >>> lA = loadtxt('landA.txt')
    >>> sol = solve_landscape_nspecies(lA, parn, dx)

    '''
    p0 = [ par['alpham'], par['alphap'] ]
    p1 = dict(
            r = [par['rm'][0], par['rp'][0]],
            D = [par['Dm'][0], par['Dp'][0]],
            **{ k: par[k] for k in ['left', 'right', 'top', 'bottom'] }
            )
    p2 = dict(
            r = [par['rm'][1], par['rp'][1]],
            D = [par['Dm'][1], par['Dp'][1]],
            **{ k: par[k] for k in ['left', 'right', 'top', 'bottom'] }
            )
    newpar = [p0, p1, p2]

    return solve_landscape_ntypes_nspecies(landscape, newpar, dx, f_tol=f_tol,
            skip_refine=skip_refine, verbose=verbose)

def find_interfaces(landscape):
    '''Helper function that marks where are the internal boundaries.'''
    B = 2*landscape - 1
    Bx = B[:-1,:]*(- B[1:,:] * B[:-1,:] + 1)//2
    Bxpm = (Bx + 1)//2
    Bxmp = (-Bx + 1)//2

    By = B[:,:-1]*(- B[:,1:] * B[:,:-1] + 1)//2
    Bypm = (By + 1)//2
    Bymp = (-By + 1)//2
    return Bxpm, Bxmp, Bypm, Bymp

def find_interfaces_ntypes(landscape):
    '''Determines internal boundaries for landscapes with many habitat types.

    Parameters
    ----------
    landscape : 2-d array

    Returns
    -------
    Bx, By : tuple of two dicts
        each key is a tuple (i,j) corresponding to the numbers of the habitat
        types, and the value corresponds to the indices where a boundary
        between them appears, along either the x or y direction

    '''
    n = np.unique(landscape).astype(int)
    A = 2**(landscape.astype(int))
    Ax = A[1:,:] - A[:-1,:]
    Ay = A[:,1:] - A[:,:-1]
    Bx = {}
    By = {}
    for i, j in iproduct(n, repeat=2):
        if i != j:
            Bx[(i,j)] = np.where(Ax == 2**j - 2**i)
            By[(i,j)] = np.where(Ay == 2**j - 2**i)

    return Bx, By

def shifted_index(x, index, shift):
    """Shifts the indices of elements from an array.

    Helper function to deal with indices returned by `np.where()`.

    Parameters
    ----------
    x : tuple of arrays
    index : integer
        Position of the shifted array in th tuple
    shift : integer
        Value by which to shift the array

    Returns
    -------
    out : tuple of shifted arrays (same shape as input)

    Notes
    -----
    This function doesn't affect the original arrays, it makes a copy, modifies
    and then returns it. For that reason, it's not very efficient.

    """
    import copy
    newx = copy.deepcopy(list(x))
    newx[index] += shift
    return tuple(newx)

def solve_multiple_parameters(variables, values, landscape, par, dx,
        f_tol=None, force_positive=False, verbose=True, multiprocess=True):
    '''Solve given landscape for several combinations of parameters.

    Solves a given landscape with a set of common parameters and for a range of
    values for some variables, optionally using multiprocessing to speed it up.

    Parameters
    ----------
    variables : list of strings
        names of the varied parameters
    values : list of lists (or tuples)
        each item contains a list of values corresponding to the parameters
        given in the `variables` list
    landscape : 2-d array
        zeroes and ones describing the landscape
    par : ordered dict
        common values for all the problem parameters. See documentation for
        `solve_landscape()`
    dx: float
        length of each edge
    f_tol : float
        tolerance for the residue, passed on to the solver routine. Default is
        6e-6
    verbose : bool
        print residue of the solution and its maximum and minimum values.
        Notice that the order of appearance of each output is not the same as
        the input if multiprocess is True.
    multiprocess : bool
        determines whether to use multiprocessing to use multiple cores. True
        by default, in which case the total number of CPUs minus one are used

    Returns
    -------
    solutions: list of 2-d arrays
        solutions of each set of parameters, in the same ordering of the input
        values

    Example
    -------
    >>> from landscape import *
    >>> lA = image_to_landscape('landA.tif')
    >>> D = np.arange(0.01, 0.05, 0.01)
    >>> r = np.arange(0.1, 0.5, 0.1)
    >>> values = [ (x, y, y) for x, y in iproduct(r, D) ]
    >>> sols = solve_multiple_parameters(['r', 'Dp', 'Dm'], values, lA, par, dx)

    '''
    from functools import partial
    # this is not compatible with python 2.6 when using multiprocessing, due to
    # bug http://bugs.python.org/issue5228
    solve_landscape_wrapper = partial(solve_landscape, landscape, dx=dx,
            f_tol=f_tol, verbose=verbose)
    works = [ OrderedDict(par, **dict(zip(variables, p))) for p in values ]

    if multiprocess:
        from multiprocessing import Pool, cpu_count
        cpus = cpu_count()
        pool = Pool(cpus if cpus == 1 else cpus - 1)
        solutions = pool.map(solve_landscape_wrapper, works)
        pool.close()
        pool.join()
    else:
        solutions = map(solve_landscape_wrapper, works)

    return solutions

def solve_multiple_landscapes(landscapes, par, dx, f_tol=None, verbose=True,
        force_positive=True, multiprocess=True):
    '''Solve several landscapes.

    Solves a set of landscape with a given set of parameters, optionally using
    multiprocessing to speed things up.

    Parameters
    ----------
    landscape : list of 2-d arrays
        zeroes and ones describing the landscape
    par : ordered dict
        values for all the problem parameters. See documentation for
        `solve_landscape()`
    dx : float
        lenght of each edge
    f_tol : float
        tolerance for the residue, passed on to the solver routine. Default is
        6e-6
    verbose : bool
        print residue of the solution and its maximum and minimum values.
        Notice that the order of appearance of each output is not the same as
        the input if multiprocess is True.
    multiprocess : bool
        whether to use multiprocessing to use multiple cores. True by default,
        in which case the total number of CPUs minus one are used

    Returns
    -------
    solutions : list os 2-d arrays
        solutions to each landscape, in the same ordering of the input values

    Example
    -------
    >>> from landscape import *
    >>> l = [ random_landscape(i*100, 0.8, 100) for i in range(30,70,10) ]
    >>> sols = solve_multiple_landscapes(l, par, dx)

    '''
    from functools import partial
    # this is not compatible with python 2.6 when using multiprocessing, due to
    # bug http://bugs.python.org/issue5228
    solve_landscape_wrapper = partial(solve_landscape, par=par, dx=dx,
            f_tol=f_tol, force_positive=force_positive, verbose=verbose)

    if multiprocess:
        from multiprocessing import Pool, cpu_count
        cpus = cpu_count()
        pool = Pool(cpus if cpus == 1 else cpus - 1)
        solutions = pool.map(solve_landscape_wrapper, landscapes)
        pool.close()
        pool.join()
    else:
        solutions = map(solve_landscape_wrapper, landscapes)

    return solutions

def refine_grid(landscape, n=2):
    '''Increase the resolution of a landscape grid by a factor of n.'''
    return np.repeat(np.repeat(landscape, n, axis=1), n, axis=0)

def coarse_grid(landscape):
    '''Decrease resolution of a grid by a factor of 2.

    Each square of 2 by 2 pixels over the last two axes is replaced by a single
    pixel with its average value.

    '''
    a = (landscape[...,::2,:] + landscape[...,1::2,:]) / 2
    return (a[...,:,::2] + a[...,:,1::2]) / 2

def image_to_landscape(image, n=2):
    '''Converts an image to a 2d array of integers from 0 to n-1.

    Most image formats (including TIFF) requires PIL.

    Parameters
    ----------
    image : string
        image filename

    Returns
    -------
    landscape : 2-d array
        all values are integers from 0 to n-1, where zeroes correspond to
        lighter shades in the original image, and darker shades to higher
        values, so that each value corresponds to a distinct type of habitat

    '''
    from matplotlib.image import imread
    from scipy.ndimage.measurements import histogram

    M = imread(image).sum(axis=2)[::-1,:]
    L = np.zeros_like(M)
    nbins = max(20, n*4)
    freq = histogram(M, M.min(), M.max(), nbins)
    bins = np.linspace(M.min(), M.max(), nbins)
    peaks = sorted(freq.argsort()[-n:])
    delims = [bins[0]] + [ bins[(peaks[i+1] + peaks[i])//2] for i in range(n-1) ] + [bins[-1]+1]
    for i in range(n):
        L[ (M > delims[i]) & (M < delims[i+1]) ] = i
    #M = np.around(M/M.max())
    #return 1 - M.astype(np.int16)
    return L

def swap_values(landscape, swaps):
    '''Substitute values of a landscape


    Parameters
    ----------
    landscape : 2-d array
    swaps: list of tuples
        each item if of the form (original_value, new_value)

    Returns
    -------
    landscape : 2-d array

    Example
    -------
    >>> # rotation of 3 values
    >>> newl = swap_values(l, [(0,1), (1,2), (2,0)])

    '''
    newl = landscape.copy()
    for i,j in swaps:
        newl[ landscape == i ] = j
    return newl

def random_landscape(cover, frag, size, radius=1, norm='taxicab'):
    '''Generates random square landscapes with a given cover and fragmentation.

    Parameters
    ----------
    cover : integer smaller than size**2, total number of patch elements
    frag : float between 0 and 1 (both exclusive), gives the level of
        fragmentation, with 0 being a single clump and 1 a totally random
        (scattered) landscape
    size : integer, lenght of the landscape (total area is size squared)
    radius : integer, maximum distance at which an element is considered to be
        a neighbour of another
    norm : one of 'maximum' or 'taxicab', defines the norm used for measuring
        distances between elements of the landscape

    Returns
    -------
    landscape : a square 2-d array of zeroes and ones

    References
    ----------
    Lenore Fahrig, Relative Effects of Habitat Loss and Fragmentation on
    Population Extinction, The Journal of Wildlife Management, Vol. 61, No. 3
    (Jul., 1997), pp. 603-610

    '''
    from numpy.random import rand, randint

    if cover > size**2:
        raise ValueError('Cover (=%i) must be smaller than total area (=%i).' %
                (cover, size**2)) 
    M = np.zeros((size+2*radius, size+2*radius), dtype=int)
    n = 0

    if norm.lower() == 'maximum':
        while n < cover:
            i, j = randint(radius, size+radius, 2)
            if M[i, j]:
                continue
            if radius == 0 or any(M[i-radius:i+radius+1, j-radius:j+radius+1]):
                M[i, j] = 1
                n += 1
            elif rand() < frag:
                M[i, j] = 1
                n += 1

    elif norm.lower() == 'taxicab':
        indexes = []
        for x in range(-radius, radius+1):
            indexes += [(x, y) for y in range(-radius+abs(x), radius+1-abs(x))]
        indexes.remove((0, 0))
        indexes = np.array(indexes)
        while n < cover:
            i, j = randint(radius, size+radius, 2)
            if M[i, j]:
                continue
            if radius == 0 or any(M[i+indexes[:,0], j+indexes[:,1]]):
                M[i, j] = 1
                n += 1
            elif rand() < frag:
                M[i, j] = 1
                n += 1
    else:
        raise ValueError('Expected either "taxicab" or "maximum" for keyword '
                'arg "norm", got %s.' % norm)

    if radius == 0:
        return M
    return M[radius:-radius,radius:-radius]

def random_landscape_inv(cover, frag, size, radius=1, norm='taxicab'):
    land0 = 1 - random_landscape(size**2-cover, frag, size, radius, norm)
    count, labeled = statistics_landscape(land0, labels=True)

    npatches = len(count) - 1
    while npatches > 1./frag:
        p = count.argmin()

def popcount_patches(landscape, solution):
    '''Labels patches and sums total area and population in each one.

    Parameters
    ----------
    landscape : 2-d array 
        ones and zeroes representing patches and matrix
    solution : 2-d array
        solution to the system

    Returns
    -------
    array : 2-d array 
        the first line is the area and the second is the total population. The
        first column is the total area and population in the matrix.

    '''
    count, labeled = statistics_landscape(landscape, labels=True)

    pops = []
    for i in range(0, len(count)):
        marked = where(labeled == i)
        pops.append(solution[marked].sum())
    return np.array([count, pops])

def statistics_landscape(landscape, labels=False):
    '''Labels patches and sums total area in each one.

    Parameters
    ----------
    landscape : 2-d array
        ones and zeroes representing patches and matrix
    labels : bool
        if True, return the labeled 2-d array landscape

    Returns
    -------
    array : 2-d array 
        the first line is the area and the second is the total population. The
        first column is the total area and population in the matrix.
    labeled : tuple
        if `labels` is True, returns a tuple (array, labeled), where the second
        element is the labeled 2-d array landscape

    '''
    from scipy.ndimage import label
    labeled, npatches = label(landscape)
    count = []
    for i in range(0, npatches+1):
        marked = np.where(labeled == i)
        area = len(marked[0])
        count.append(area)

    if labels:
        return np.array(count), labeled
    return np.array(count)

def fit_exponential(count):
    from scipy.stats import kstest

    m = 1./np.mean(count)
    exp_cdf = lambda x: m * exp(-m*x)
    return kstest(count, exp_cdf)

def plot_landscape(landscape, solution, dx=dx, K=1.):
    #from matplotlib.pyplot import imshow, colorbar, ylabel, xlabel, yticks, xticks, subplot, cm
    extent = (0, solution.shape[1]*dx, 0, solution.shape[0]*dx)

    plt.subplot(211)
    plt.imshow(landscape, cmap=plt.cm.binary, origin='lower',
               interpolation='none', extent=extent)
    plt.ylabel('y (km)')
    locs, labels = plt.xticks()
    plt.xticks(locs, ['']*len(locs))

    plt.subplot(212)
    plt.imshow(solution, cmap=plt.cm.jet, origin='lower', interpolation='none',
               extent=extent, vmin=0, vmax = K)
    plt.xlabel('x (km)')
    plt.ylabel('y (km)')

    plt.colorbar()

def plot_density_area(pops, labels, xlog=True, ylog=False):
    for y, label in zip(pops, labels):
        plt.plot(y[0][1:], y[1][1:]/y[0][1:], '.', label=label)

    plt.xlabel('area')
    plt.ylabel('avg. density')
    if xlog and ylog:
        plt.loglog()
    elif xlog:
        plt.xscale('log')
    elif ylog:
        plt.yscale('log')
    plt.legend(loc='best', frameon=False)

def plot_nspecies(landscape, solutions, dx=dx):
    from matplotlib.gridspec import GridSpec
    n = len(solutions)
    K = np.ceil(max([ s.max() for s in solutions ]))
    extent = (0, landscape.shape[1]*dx, 0, landscape.shape[0]*dx)

    fig = plt.figure(figsize=(1.4, n+1))
    gs = GridSpec(n+2, 1, height_ratios=(n+1)*[10]+[1])
    plt.subplot(gs[0])
    plt.imshow(landscape, cmap=plt.cm.binary, origin='lower',
            interpolation='none', extent=extent)
    plt.tick_params(labelbottom=False, labelleft=False)

    for i, s in enumerate(solutions):
        plt.subplot(gs[i+1])
        plt.imshow(s, cmap=plt.cm.jet, origin='lower', interpolation='none',
                extent=extent, vmin=0, vmax = K)
        plt.tick_params(labelbottom=False, labelleft=False)
    ax=plt.subplot(gs[n+1])
    plt.colorbar(orientation='horizontal', cax=ax)
    plt.tight_layout(0.1)
    return fig