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()