Module nlisim.diffusion

Expand source code
import numpy as np
from scipy.sparse import csr_matrix, dok_matrix, eye
from scipy.sparse.linalg import cg

from nlisim.coordinates import Voxel
from nlisim.grid import RectangularGrid

_dtype_float64 = np.dtype('float64')


def discrete_laplacian(
    grid: RectangularGrid, mask: np.ndarray, dtype: np.dtype = _dtype_float64
) -> csr_matrix:
    """Return a discrete laplacian operator for the given restricted grid.

    This computes a standard laplacian operator as a scipy linear operator, except it is
    restricted to a grid mask.  The use case for this is to compute surface diffusion
    on a gridded variable.  The mask is generated from a category on the lung_tissue
    variable.
    """
    graph_shape = len(grid), len(grid)
    laplacian = dok_matrix(graph_shape, dtype=dtype)

    delta_z = grid.delta(0)
    delta_y = grid.delta(1)
    delta_x = grid.delta(2)

    for k, j, i in zip(*(mask).nonzero()):
        voxel = Voxel(x=i, y=j, z=k)
        voxel_index = grid.get_flattened_index(voxel)

        for neighbor in grid.get_adjacent_voxels(voxel, corners=False):
            ni = neighbor.x
            nj = neighbor.y
            nk = neighbor.z

            if not mask[nk, nj, ni]:
                continue

            neighbor_index = grid.get_flattened_index(neighbor)

            dx = delta_x[k, j, i] * (i - ni)
            dy = delta_y[k, j, i] * (j - nj)
            dz = delta_z[k, j, i] * (k - nk)
            inverse_distance2 = 1 / (dx * dx + dy * dy + dz * dz)  # units: 1/(µm^2)

            laplacian[voxel_index, voxel_index] -= inverse_distance2
            laplacian[voxel_index, neighbor_index] += inverse_distance2

    return laplacian.tocsr()


def periodic_discrete_laplacian(
    grid: RectangularGrid, mask: np.ndarray, dtype: np.dtype = _dtype_float64
) -> csr_matrix:
    """Return a laplacian operator with periodic boundary conditions.

    This computes a standard laplacian operator as a scipy linear operator, except it is
    restricted to a grid mask.  The use case for this is to compute surface diffusion
    on a gridded variable.  The mask is generated from a category on the lung_tissue
    variable.
    """
    graph_shape = len(grid), len(grid)
    z_extent, y_extent, x_extent = grid.shape
    laplacian = dok_matrix(graph_shape, dtype=dtype)

    delta_z = grid.delta(0)
    delta_y = grid.delta(1)
    delta_x = grid.delta(2)

    for k, j, i in zip(*(mask).nonzero()):
        voxel = Voxel(x=i, y=j, z=k)
        voxel_index = grid.get_flattened_index(voxel)

        for offset in [(-1, 0, 0), (1, 0, 0), (0, -1, 0), (0, 1, 0), (0, 0, -1), (0, 0, 1)]:
            # voxel coordinate displacements
            dk, dj, di = offset

            # find the neighbor for periodic boundary conditions
            neighbor: Voxel = Voxel(
                x=(i + di) % x_extent, y=(j + dj) % y_extent, z=(k + dk) % z_extent
            )

            # but maybe it isn't in the mask (i.e. air)
            if not mask[neighbor.z, neighbor.y, neighbor.x]:
                continue

            neighbor_index = grid.get_flattened_index(neighbor)

            # continuous space displacements
            dx = delta_x[k, j, i] * di
            dy = delta_y[k, j, i] * dj
            dz = delta_z[k, j, i] * dk
            inverse_distance2 = 1 / (dx * dx + dy * dy + dz * dz)  # units: 1/(µm^2)

            laplacian[voxel_index, voxel_index] -= inverse_distance2
            laplacian[voxel_index, neighbor_index] += inverse_distance2

    return laplacian.tocsr()


def apply_diffusion(
    variable: np.ndarray,
    laplacian: csr_matrix,
    diffusivity: float,
    dt: float,
    tolerance: float = 1e-10,
) -> np.ndarray:
    """Apply diffusion to a variable.

    Solves Laplace's equation in 3D using Crank-Nicholson.  The variable is
    advanced in time by `dt` time units using the conjugate gradient method.

    Note that, due to numerical error, we cannot guarantee that the quantity
    of the molecule will remain constant.

    The intended use case for this method is to perform "surface diffusion" generated
    by a mask from the `lung_tissue` variable, e.g.

        surface_mask = lung_tissue == TissueTypes.EPITHELIUM
        laplacian = discrete_laplacian(grid, mask)
        iron_concentration[:] = apply_diffusion(iron_concentration, laplacian, diffusivity, dt)
    """
    a = eye(*laplacian.shape) - (diffusivity * dt / 2.0) * laplacian
    b = eye(*laplacian.shape) + (diffusivity * dt / 2.0) * laplacian
    var_next, info = cg(a, b @ variable.ravel(), tol=tolerance)
    if info > 0:
        raise Exception(f'CG failed (after {info} iterations)')
    elif info < 0:
        raise Exception(f'CG failed ({info})')

    return np.maximum(0.0, var_next.reshape(variable.shape))

Functions

def apply_diffusion(variable: numpy.ndarray, laplacian: scipy.sparse.csr.csr_matrix, diffusivity: float, dt: float, tolerance: float = 1e-10) ‑> numpy.ndarray

Apply diffusion to a variable.

Solves Laplace's equation in 3D using Crank-Nicholson. The variable is advanced in time by dt time units using the conjugate gradient method.

Note that, due to numerical error, we cannot guarantee that the quantity of the molecule will remain constant.

The intended use case for this method is to perform "surface diffusion" generated by a mask from the lung_tissue variable, e.g.

surface_mask = lung_tissue == TissueTypes.EPITHELIUM
laplacian = discrete_laplacian(grid, mask)
iron_concentration[:] = apply_diffusion(iron_concentration, laplacian, diffusivity, dt)
Expand source code
def apply_diffusion(
    variable: np.ndarray,
    laplacian: csr_matrix,
    diffusivity: float,
    dt: float,
    tolerance: float = 1e-10,
) -> np.ndarray:
    """Apply diffusion to a variable.

    Solves Laplace's equation in 3D using Crank-Nicholson.  The variable is
    advanced in time by `dt` time units using the conjugate gradient method.

    Note that, due to numerical error, we cannot guarantee that the quantity
    of the molecule will remain constant.

    The intended use case for this method is to perform "surface diffusion" generated
    by a mask from the `lung_tissue` variable, e.g.

        surface_mask = lung_tissue == TissueTypes.EPITHELIUM
        laplacian = discrete_laplacian(grid, mask)
        iron_concentration[:] = apply_diffusion(iron_concentration, laplacian, diffusivity, dt)
    """
    a = eye(*laplacian.shape) - (diffusivity * dt / 2.0) * laplacian
    b = eye(*laplacian.shape) + (diffusivity * dt / 2.0) * laplacian
    var_next, info = cg(a, b @ variable.ravel(), tol=tolerance)
    if info > 0:
        raise Exception(f'CG failed (after {info} iterations)')
    elif info < 0:
        raise Exception(f'CG failed ({info})')

    return np.maximum(0.0, var_next.reshape(variable.shape))
def discrete_laplacian(grid: RectangularGrid, mask: numpy.ndarray, dtype: numpy.dtype = dtype('float64')) ‑> scipy.sparse.csr.csr_matrix

Return a discrete laplacian operator for the given restricted grid.

This computes a standard laplacian operator as a scipy linear operator, except it is restricted to a grid mask. The use case for this is to compute surface diffusion on a gridded variable. The mask is generated from a category on the lung_tissue variable.

Expand source code
def discrete_laplacian(
    grid: RectangularGrid, mask: np.ndarray, dtype: np.dtype = _dtype_float64
) -> csr_matrix:
    """Return a discrete laplacian operator for the given restricted grid.

    This computes a standard laplacian operator as a scipy linear operator, except it is
    restricted to a grid mask.  The use case for this is to compute surface diffusion
    on a gridded variable.  The mask is generated from a category on the lung_tissue
    variable.
    """
    graph_shape = len(grid), len(grid)
    laplacian = dok_matrix(graph_shape, dtype=dtype)

    delta_z = grid.delta(0)
    delta_y = grid.delta(1)
    delta_x = grid.delta(2)

    for k, j, i in zip(*(mask).nonzero()):
        voxel = Voxel(x=i, y=j, z=k)
        voxel_index = grid.get_flattened_index(voxel)

        for neighbor in grid.get_adjacent_voxels(voxel, corners=False):
            ni = neighbor.x
            nj = neighbor.y
            nk = neighbor.z

            if not mask[nk, nj, ni]:
                continue

            neighbor_index = grid.get_flattened_index(neighbor)

            dx = delta_x[k, j, i] * (i - ni)
            dy = delta_y[k, j, i] * (j - nj)
            dz = delta_z[k, j, i] * (k - nk)
            inverse_distance2 = 1 / (dx * dx + dy * dy + dz * dz)  # units: 1/(µm^2)

            laplacian[voxel_index, voxel_index] -= inverse_distance2
            laplacian[voxel_index, neighbor_index] += inverse_distance2

    return laplacian.tocsr()
def periodic_discrete_laplacian(grid: RectangularGrid, mask: numpy.ndarray, dtype: numpy.dtype = dtype('float64')) ‑> scipy.sparse.csr.csr_matrix

Return a laplacian operator with periodic boundary conditions.

This computes a standard laplacian operator as a scipy linear operator, except it is restricted to a grid mask. The use case for this is to compute surface diffusion on a gridded variable. The mask is generated from a category on the lung_tissue variable.

Expand source code
def periodic_discrete_laplacian(
    grid: RectangularGrid, mask: np.ndarray, dtype: np.dtype = _dtype_float64
) -> csr_matrix:
    """Return a laplacian operator with periodic boundary conditions.

    This computes a standard laplacian operator as a scipy linear operator, except it is
    restricted to a grid mask.  The use case for this is to compute surface diffusion
    on a gridded variable.  The mask is generated from a category on the lung_tissue
    variable.
    """
    graph_shape = len(grid), len(grid)
    z_extent, y_extent, x_extent = grid.shape
    laplacian = dok_matrix(graph_shape, dtype=dtype)

    delta_z = grid.delta(0)
    delta_y = grid.delta(1)
    delta_x = grid.delta(2)

    for k, j, i in zip(*(mask).nonzero()):
        voxel = Voxel(x=i, y=j, z=k)
        voxel_index = grid.get_flattened_index(voxel)

        for offset in [(-1, 0, 0), (1, 0, 0), (0, -1, 0), (0, 1, 0), (0, 0, -1), (0, 0, 1)]:
            # voxel coordinate displacements
            dk, dj, di = offset

            # find the neighbor for periodic boundary conditions
            neighbor: Voxel = Voxel(
                x=(i + di) % x_extent, y=(j + dj) % y_extent, z=(k + dk) % z_extent
            )

            # but maybe it isn't in the mask (i.e. air)
            if not mask[neighbor.z, neighbor.y, neighbor.x]:
                continue

            neighbor_index = grid.get_flattened_index(neighbor)

            # continuous space displacements
            dx = delta_x[k, j, i] * di
            dy = delta_y[k, j, i] * dj
            dz = delta_z[k, j, i] * dk
            inverse_distance2 = 1 / (dx * dx + dy * dy + dz * dz)  # units: 1/(µm^2)

            laplacian[voxel_index, voxel_index] -= inverse_distance2
            laplacian[voxel_index, neighbor_index] += inverse_distance2

    return laplacian.tocsr()