def gaussian_potential(center, width, amplitude):
"""A Gaussian potential
Arguments:
center center of Gaussian (scalar or vector)
width width of Gaussian (scalar or vector)
amplitude amplitude of Gaussian, (negative for repulsive)
"""
center = np.atleast_1d(center)
ndim = len(center)
width = value_to_vector(width, ndim)
amplitude = value_to_vector(amplitude, ndim)
def potential(*args):
U = np.ones_like(args[0])
for i, arg in enumerate(args):
U *= np.exp(-np.square((arg - center[i]) / width[i]))
return -amplitude * U
return potential
def harmonic_potential(center, k):
"""A harmonic potential
Arguments:
center center of harmonic potential (scalar or vector)
k spring constant of harmonic potential (scalar or vector)
"""
center = np.atleast_1d(center)
ndim = len(center)
k = value_to_vector(k, ndim)
def potential(*args):
U = np.zeros_like(args[0])
for i, arg in enumerate(args):
U += 0.5 * k[i] * (arg - center[i])**2
return U
return potential
def gaussian_pdf(center, width):
"""A Gaussian probability distribution function
Arguments:
center center of Gaussian (scalar or vector)
width width of Gaussian (scalar or vector)
"""
center = np.atleast_1d(center)
ndim = len(center)
width = value_to_vector(width, ndim)
def pdf(*args):
values = np.ones_like(args[0])
for i, arg in enumerate(args):
values *= np.exp(-np.square((arg - center[i])/width[i]))
return values/np.sum(values)
return pdf
def __init__(self,
*,
temperature,
drag,
extent,
resolution,
potential=None,
force=None,
boundary=fplanck.boundary.reflecting):
"""
Solve the Fokker-Planck equation
Arguments:
temperature temperature of the surrounding bath (scalar or vector)
drag drag coefficient (scalar or vector)
extent extent (size) of the grid (vector)
resolution spatial resolution of the grid (scalar or vector)
potential external potential function, U(ndim -> scalar)
force external force function, F(ndim -> ndim)
boundary type of boundary condition (scalar or vector, default: reflecting)
"""
self.extent = np.atleast_1d(extent)
self.ndim = self.extent.size
self.temperature = value_to_vector(temperature, self.ndim)
self.drag = value_to_vector(drag, self.ndim)
self.resolution = value_to_vector(resolution, self.ndim)
self.potential = potential
self.force = force
self.boundary = value_to_vector(boundary, self.ndim, dtype=object)
self.diffusion = constants.k * self.temperature / self.drag
self.mobility = 1 / self.drag
self.beta = 1 / (constants.k * self.temperature)
self.Ngrid = np.ceil(self.extent / resolution).astype(int)
axes = [
np.arange(self.Ngrid[i]) * self.resolution[i]
for i in range(self.ndim)
]
for axis in axes:
axis -= np.average(axis)
self.axes = axes
self.grid = np.array(np.meshgrid(*axes, indexing='ij'))
self.Rt = np.zeros_like(self.grid)
self.Lt = np.zeros_like(self.grid)
self.potential_values = np.zeros_like(self.grid[0])
self.force_values = np.zeros_like(self.grid)
if self.potential is not None:
U = self.potential(*self.grid)
self.potential_values += U
self.force_values -= np.gradient(U, *self.resolution)
for i in range(self.ndim):
dU = np.roll(U, -1, axis=i) - U
self.Rt[i] += self.diffusion[i] / self.resolution[
i]**2 * np.exp(-self.beta[i] * dU / 2)
dU = np.roll(U, 1, axis=i) - U
self.Lt[i] += self.diffusion[i] / self.resolution[
i]**2 * np.exp(-self.beta[i] * dU / 2)
if self.force is not None:
F = np.atleast_2d(self.force(*self.grid))
self.force_values += F
for i in range(self.ndim):
dU = -(np.roll(F[i], -1, axis=i) +
F[i]) / 2 * self.resolution[i]
self.Rt[i] += self.diffusion[i] / self.resolution[
i]**2 * np.exp(-self.beta[i] * dU / 2)
dU = (np.roll(F[i], 1, axis=i) + F[i]) / 2 * self.resolution[i]
self.Lt[i] += self.diffusion[i] / self.resolution[
i]**2 * np.exp(-self.beta[i] * dU / 2)
if self.force is None and self.potential is None:
for i in range(self.ndim):
self.Rt[i] = self.diffusion[i] / self.resolution[i]**2
self.Lt[i] = self.diffusion[i] / self.resolution[i]**2
for i in range(self.ndim):
if self.boundary[i] == fplanck.boundary.reflecting:
idx = slice_idx(i, self.ndim, -1)
self.Rt[i][idx] = 0
idx = slice_idx(i, self.ndim, 0)
self.Lt[i][idx] = 0
elif self.boundary[i] == fplanck.boundary.periodic:
idx = slice_idx(i, self.ndim, -1)
dU = -self.force_values[i][idx] * self.resolution[i]
self.Rt[i][idx] = self.diffusion[i] / self.resolution[
i]**2 * np.exp(-self.beta[i] * dU / 2)
idx = slice_idx(i, self.ndim, 0)
dU = self.force_values[i][idx] * self.resolution[i]
self.Lt[i][idx] = self.diffusion[i] / self.resolution[
i]**2 * np.exp(-self.beta[i] * dU / 2)
else:
raise ValueError(
f"'{self.boundary[i]}' is not a valid a boundary condition"
)
self._build_matrix()