Source code for arbdmodel.interactions

from abc import abstractmethod, ABCMeta
import os, sys
import scipy
import numpy as np
from shutil import copyfile
from .grid import writeDx
""" Module providing classes used to describe potentials in ARBD """


## Abstract classes that others inherit from
[docs] class AbstractPotential(metaclass=ABCMeta): """ Abstract base class for implementing interaction potentials. This class defines the basic structure for creating and managing interaction potential functions that can be written to files and used in simulations. Parameters ---------- range_ : tuple of (float, float or None) The range of distances for the potential (min, max). If max is None, potential extends to infinity. resolution : float, optional The spacing between points in the discretized potential. Default is 0.1. max_force : float or None, optional Maximum allowed force magnitude. If provided, forces exceeding this value will be capped. max_potential : float or None, optional Maximum allowed potential energy. If provided, the potential will be modified to ensure it doesn't exceed this value. zero : str, optional Method to set the zero of the potential. Options are: - 'min': Zero is at the minimum value - 'first': Zero is at the first point - 'last': Zero is at the last point periodic : bool Whether the potential is periodic (False by default). Note ----- Subclasses must implement at minimum the `potential` and `filename` methods. """ """ Abstract class for writing potentials """ def __init__(self, range_=(0,None), resolution=0.1, max_force = None, max_potential = None, zero='last'): self.range_ = range_ self.resolution = resolution self.max_force = max_force self.max_potential = max_potential self.zero = zero @property def range_(self): return self.__range_ @range_.setter def range_(self,value): assert(len(value) == 2) self.__range_ = tuple(value) @property def periodic(self): return False
[docs] @abstractmethod def potential(self, r, types=None): raise NotImplementedError
def __remove_nans(self, u): z = lambda x: np.argwhere(x).flatten() nans=np.isnan(u) u[nans] = scipy.interpolate.interp1d(z(~nans), u[~nans], fill_value='extrapolate')(z(nans)) def __remove_nans_bak(self, u): left = np.isnan(u[:-1]) right = np.where(left)[0]+1 while np.any(np.isnan(u[right])): right[np.isnan(u[right])] += 1 u[:-1][left] = u[right]
[docs] def filename(self, types=None): raise NotImplementedError('Inherited potential objects should overload this function')
def _cap_potential(self, r, u): self.__remove_nans(u) if self.zero == 'min': u = u - np.min(u) elif self.zero == 'first': u = u - u[0] elif self.zero == 'last': u = u - u[-1] else: raise ValueError('Unrecognized option for "zero" argument') max_force = self.max_force if max_force is not None: assert(max_force > 0) f = np.diff(u)/np.diff(r) f[f>max_force] = max_force f[f<-max_force] = -max_force u[0] = 0 u[1:] = np.cumsum(f*np.diff(r)) if self.max_potential is not None: f = np.diff(u)/np.diff(r) ids = np.where( 0.5*(u[1:]+u[:-1]) > self.max_potential )[0] # w = np.sqrt(2*self.max_potential/self.k) # drAvg = 0.5*(np.abs(dr[ids]) + np.abs(dr[ids+1])) # f[ids] = f[ids] * np.exp(-(drAvg-w)/(w)) f[ids] = 0 u[0] = 0 u[1:] = np.cumsum(f*np.diff(r)) if self.zero == 'min': u = u - np.min(u) elif self.zero == 'first': u = u - u[0] elif self.zero == 'last': u = u - u[-1] else: raise ValueError('Unrecognized option for "zero" argument') return u
[docs] def write_file(self, filename=None, types=None): """ Write potential energy values as a function of distance to a file. This method evaluates the potential function over a range of distances and writes the distance-potential pairs to a text file. Parameters ---------- filename : str, optional The path to the output file. If None, a default filename will be generated using the `filename` method with the given types. types : tuple or list, optional The particle types for which to calculate the potential. Used for both filename generation (if filename is None) and potential evaluation. If None, defaults to the types defined in the interaction model. Note ----- The potential values are capped using the `_cap_potential` method to handle potential singularities or extreme values. The output file contains two columns: distance (r) and potential energy (u). Any 'divide by zero' or 'invalid value' numpy warnings are ignored during the potential calculation. """ if filename is None: filename = self.filename(types) rmin,rmax = self.range_ r = np.arange(rmin, rmax+self.resolution, self.resolution) with np.errstate(divide='ignore',invalid='ignore'): u = self.potential(r, types) u = self._cap_potential(r,u) np.savetxt(filename, np.array([r,u]).T)
def __hash__(self): return hash((self.range_, self.resolution, self.max_force, self.max_potential, self.zero)) def __eq__(self, other): for a in ("range_", "resolution", "max_force", "max_potential", "zero"): if getattr(self,a) != getattr(other,a): return False return type(self).__name__ == type(other).__name__ ## Want the same objects after copy operations def __copy__(self): return self def __deepcopy__(self, memo): return self
## Concrete nonbonded pontentials
[docs] class LennardJones(AbstractPotential):
[docs] def potential(self, r, types): typeA, typeB = types epsilon = np.sqrt( typeA.epsilon**2 + typeB.epsilon**2 ) r0 = 0.5 * (typeA.radius + typeB.radius) r6 = (r0/r)**6 r12 = r6**2 u = 4 * epsilon * (r12-r6) # u[0] = u[1] # Remove NaN return u
[docs] class HalfHarmonic(AbstractPotential):
[docs] def potential(self, r, types): typeA, typeB = types k = 10 # kcal/mol AA**2 r0 = (typeA.radius + typeB.radius) u = 0.5 * k * (r-r0)**2 u[r > r0] = np.zeros( np.shape(u[r > r0]) ) return u
[docs] class TabulatedNonbonded(AbstractPotential): def __init__(self, tableFile, *args, **kwargs): self.tableFile = tableFile AbstractPotential.__init__(self,*args,**kwargs) ## TODO: check that tableFile exists and is regular file
[docs] def potential(self, r, types): raise NotImplementedError('This should probably not be implemented')
[docs] def write_file(self, filename, types): if filename != self.tableFile: copyfile(self.tableFile, filename)
[docs] class BoundaryPotential(AbstractPotential): """Boundary potential for confining simulations.""" def __init__(self, cell_vectors, cell_origin=None, well_depth=1.0, resolution=2.0, blur=5.0,filename="boundary.dx", **kwargs): """Initialize boundary potential. Args: cell_vectors: List of 3 cell basis vectors [[x1,y1,z1], [x2,y2,z2], [x3,y3,z3]] cell_origin: Cell origin coordinates [x,y,z] well_depth: Depth of potential well resolution: Grid resolution blur: Smoothing parameter **kwargs: Additional parameters for AbstractPotential """ default_range = (0, np.max([np.linalg.norm(v) for v in cell_vectors])) if 'range_' not in kwargs: kwargs['range_'] = default_range if 'resolution' not in kwargs: kwargs['resolution'] = resolution super().__init__(**kwargs) self.cell_vectors = cell_vectors self.cell_origin = cell_origin or [0, 0, 0] self.well_depth = well_depth self.blur = blur self.grid_cache = None self.filename=filename def _find_boundary_resolution(self): """Find appropriate boundary resolution for the cell vectors.""" v1, v2, v3 = self.cell_vectors min_PBC_length = min([max(v1), max(v2), max(v3)]) if min_PBC_length < self.resolution: dd = round(min_PBC_length / 2) else: dd = self.resolution n1 = round(np.linalg.norm(v1)/dd) n2 = round(np.linalg.norm(v2)/dd) n3 = round(np.linalg.norm(v3)/dd) return dd, n1, n2, n3 def _find_boundary_end_points(self): """Find boundary endpoints from cell vectors and origin.""" v1, v2, v3 = self.cell_vectors xyz_travel = [sum([v1[ind], v2[ind], v3[ind]]) for ind in range(3)] xyz_min = [self.cell_origin[ind] - xyz_travel[ind] * 0.5 for ind in range(3)] xyz_max = [self.cell_origin[ind] + xyz_travel[ind] * 0.5 for ind in range(3)] return (xyz_min[0], xyz_min[1], xyz_min[2], xyz_max[0], xyz_max[1], xyz_max[2]) def _create_boundary_region(self, mx, my, mz, n1, n2, n3): """Create boundary region points for the cell.""" v1, v2, v3 = self.cell_vectors v1_scaled = np.array(v1) / n1 v2_scaled = np.array(v2) / n2 v3_scaled = np.array(v3) / n3 region_pts = [] start_pt = np.array([mx, my, mz]) for i in range(n1 + 1): for j in range(n2 + 1): for k in range(n3 + 1): region_pts.append(start_pt + i * v1_scaled + j * v2_scaled + k * v3_scaled) return region_pts def _convert_to_mesh_points(self, region_pts, min_wallX, min_wallY, min_wallZ, dx, dy, dz, grid_shape): """Convert region points to grid indices.""" mesh_pts = [] for pt in region_pts: indx = round((pt[0] - min_wallX) / dx) indy = round((pt[1] - min_wallY) / dy) indz = round((pt[2] - min_wallZ) / dz) if (0 <= indx < grid_shape[0] and 0 <= indy < grid_shape[1] and 0 <= indz < grid_shape[2]): mesh_pts.append((indx, indy, indz)) return mesh_pts def _create_rectangular_mesh(self, wall_range_x, wall_range_y, wall_range_z, dd): """Create a rectangular mesh grid.""" x = np.linspace(wall_range_x[0], wall_range_x[1], int((wall_range_x[1]-wall_range_x[0])/dd) + 1) y = np.linspace(wall_range_y[0], wall_range_y[1], int((wall_range_y[1]-wall_range_y[0])/dd) + 1) z = np.linspace(wall_range_z[0], wall_range_z[1], int((wall_range_z[1]-wall_range_z[0])/dd) + 1) dx = np.mean(np.diff(x)) dy = np.mean(np.diff(y)) dz = np.mean(np.diff(z)) X, Y, Z = np.meshgrid(x, y, z, indexing='ij') return X, Y, Z, dx, dy, dz def _blur_3d_grid(self, grid, blur_size, dd): """Apply 3D Gaussian blur to a grid.""" # Create Gaussian kernel sideLen = 2*int(blur_size*3)+1 gauss = scipy.signal.gaussian(sideLen, blur_size/dd) i = np.arange(sideLen) i, j, k = np.meshgrid(i, i, i, indexing='ij') kernel = gauss[i]*gauss[j]*gauss[k] kernel = kernel/kernel.sum() # Apply convolution return scipy.signal.fftconvolve(grid, kernel, mode='same')
[docs] def potential(self, r, types=None): """Calculate potential at positions. This isn't used directly for grid potentials in ARBD, but implemented for consistency with the AbstractPotential interface. """ # Calculate distance from boundary max_dist = np.max([np.linalg.norm(v) for v in self.cell_vectors]) / 2 potential = np.zeros_like(r) for i, ri in enumerate(r): # Simplified: calculate distance from origin in normalized coordinates dist_vec = np.array([ ri/np.linalg.norm(self.cell_vectors[0]), ri/np.linalg.norm(self.cell_vectors[1]), ri/np.linalg.norm(self.cell_vectors[2]) ]) # Distance from boundary (simplified) dist_from_boundary = max_dist - np.max(np.abs(dist_vec)) if dist_from_boundary < 0: potential[i] = self.well_depth * (1 - np.exp(dist_from_boundary/self.blur)) return potential
[docs] def write_file(self, filename=None, types=None): """Write boundary potential to grid file.""" # Find optimal resolution dd, n1, n2, n3 = self._find_boundary_resolution() # Find boundary endpoints mx, my, mz, Mx, My, Mz = self._find_boundary_end_points() # Set up wall ranges with buffer scale = 1.5 # Scale factor for wall range wall_range_x = (scale*(mx-2*self.blur), scale*(Mx+2*self.blur)) wall_range_y = (scale*(my-2*self.blur), scale*(My+2*self.blur)) wall_range_z = (scale*(mz-2*self.blur), scale*(Mz+2*self.blur)) # Create mesh X, Y, Z, dx, dy, dz = self._create_rectangular_mesh( wall_range_x, wall_range_y, wall_range_z, dd) # Create boundary points region_pts = self._create_boundary_region( mx, my, mz, n1, n2, n3) # Convert to mesh points mesh_pts = self._convert_to_mesh_points( region_pts, wall_range_x[0], wall_range_y[0], wall_range_z[0], dx, dy, dz, X.shape) # Create potential pot = np.zeros(np.shape(X)) for pt in mesh_pts: pot[pt[0], pt[1], pt[2]] = -self.well_depth # Negative for attraction # Apply Gaussian blur pot_blur = self._blur_3d_grid(pot, self.blur, dd) # Write to DX file origin_out = [wall_range_x[0], wall_range_y[0], wall_range_z[0]] delta = [dx, dy, dz] writeDx(self.filename, pot_blur, origin_out, delta) return self.filename
[docs] class NullPotential(AbstractPotential): def __init__(self, range_=(0,1), resolution=0.5, filename_prefix='./potentials/', *args, **kwargs): self.filename_prefix = filename_prefix AbstractPotential.__init__(self, range_=range_, resolution=resolution, *args,**kwargs)
[docs] def potential(self, r, types): return np.zeros(r.shape)
[docs] def filename(self, types=None): return f"{self.filename_prefix}nullpot.dat"
## Bonded potentials
[docs] class HarmonicBondedPotential(AbstractPotential): def __init__(self, k, r0, filename_prefix='./potentials/', *args, **kwargs): self.k = k self.r0 = r0 self.filename_prefix = filename_prefix if 'zero' not in kwargs: kwargs['zero'] = 'min' AbstractPotential.__init__(self, *args, **kwargs) @property def kscale(self): return 1.0 @property @abstractmethod def type_(self): return "none"
[docs] def filename(self, types=None): assert(types is None) return f"{self.filename_prefix}{self.type_}-{self.k*self.kscale:.3f}-{self.r0:.3f}.dat"
[docs] def potential(self, r, types=None): dr = r-self.r0 if self.periodic == True: r_span = self.range_[1]-self.range_[0] assert(r_span > 0) dr = np.mod( dr+0.5*r_span, r_span) - 0.5*r_span return 0.5*self.k*dr**2
def __hash__(self): assert(self.type_ != "None") return hash((self.type_, self.k, self.r0, self.filename_prefix, self.periodic, AbstractPotential.__hash__(self))) # def __eq__(self, other): for a in ("type_", "k", "r0", "filename_prefix", "periodic"): if getattr(self,a) != getattr(other,a): return False return AbstractPotential.__eq__(self,other)
[docs] class HarmonicBond(HarmonicBondedPotential): def __init__(self, k, r0, correct_geometry=False, temperature=295, *args, **kwargs): if 'range_' not in kwargs: kwargs['range_'] = (0,50) HarmonicBondedPotential.__init__(self, k, r0, *args, **kwargs) self.correct_geometry = correct_geometry self.temperature = temperature @property def kscale(self): return 1.0 @property def type_(self): return 'gbond' if self.correct_geometry else 'bond'
[docs] def potential(self, r, types=None): u = HarmonicBondedPotential.potential(self, r, types) if self.correct_geometry: with np.errstate(divide='ignore',invalid='ignore'): du = 2*0.58622592*np.log(r) * self.temperature/295 du[np.logical_not(np.isfinite(du))] = 0 u = u+du return u
[docs] class HarmonicAngle(HarmonicBondedPotential): def __init__(self, *args, **kwargs): if 'range_' not in kwargs: kwargs['range_'] = (0,181) HarmonicBondedPotential.__init__(self, *args, **kwargs) @property def kscale(self): return (180.0/np.pi)**2 @property def type_(self): return 'angle'
[docs] class HarmonicDihedral(HarmonicBondedPotential): def __init__(self, *args, **kwargs): if 'range_' not in kwargs: kwargs['range_'] = (-180,180) HarmonicBondedPotential.__init__(self, *args, **kwargs) @property def kscale(self): return (180.0/np.pi)**2 @property def type_(self): return 'dihedral' @property def periodic(self): return True
[docs] class HarmonicVectorAngle(HarmonicBondedPotential): def __init__(self, *args, **kwargs): if 'range_' not in kwargs: kwargs['range_'] = (0,180) HarmonicBondedPotential.__init__(self, *args, **kwargs) @property def kscale(self): return (180.0/np.pi)**2 @property def type_(self): return 'vecangle'
[docs] class WLCSKPotential(HarmonicBondedPotential, metaclass=ABCMeta): """ ## https://aip.scitation.org/doi/full/10.1063/1.4968020 """ def __init__(self, d, lp, kT, **kwargs): ## Note, we're leveraging the HarmonicBondedPotential base class and set k to lp here, but it isn't proper HarmonicBondedPotential.__init__(self, d, lp, **kwargs) self.d = d # separation self.lp = lp # persistence length self.kT = kT
[docs] def filename(self,types=None): assert(types is None) return "%s%s-%.3f-%.3f.dat" % (self.filename_prefix, self.type_, self.d, self.lp)
def __hash__(self): return hash((self.d, self.lp, self.kT, HarmonicBondedPotential.__hash__(self))) def __eq__(self, other): for a in ("d", "lp", "kT"): if self.__dict__[a] != other.__dict__[a]: return False return HarmonicBondedPotential.__eq__(self,other)
[docs] class WLCSKBond(WLCSKPotential): """ ## https://aip.scitation.org/doi/full/10.1063/1.4968020 """ def __init__(self, d, lp, kT, **kwargs): if 'range_' not in kwargs: kwargs['range_'] = (0,50) if 'resolution' not in kwargs: kwargs['resolution'] = 0.02 if 'max_force' not in kwargs: kwargs['max_force'] = 100 WLCSKPotential.__init__(self, d, lp, kT, **kwargs) @property def type_(self): return "wlcbond"
[docs] def potential(self, r, types=None): dr = r nk = self.d / (2*self.lp) q2 = (dr / self.d)**2 a1,a2 = 1, -7.0/(2*nk) a3 = 3.0/32 - 3.0/(8*nk) - 6.0/(4*nk**2) p0,p1,p2,p3,p4 = 13.0/32, 3.4719,2.5064,-1.2906,0.6482 a4 = (p0 + p1/(2*nk) + p2*(2*nk)**-2) / (1+p3/(2*nk)+p4*(2*nk)**-2) with np.errstate(divide='ignore',invalid='ignore'): u = self.kT * nk * ( a1/(1-q2) - a2*np.log(1-q2) + a3*q2 - 0.5*a4*q2*(q2-2) ) return u
[docs] class WLCSKAngle(WLCSKPotential): """ ## https://aip.scitation.org/doi/full/10.1063/1.4968020 """ def __init__(self, d, lp, kT, **kwargs): if 'range_' not in kwargs: kwargs['range_'] = (0,181) if 'resolution' not in kwargs: kwargs['resolution'] = 0.5 WLCSKPotential.__init__(self, d, lp, kT, **kwargs) @property def type_(self): return "wlcangle"
[docs] def potential(self, r, types=None): dr = r - 180 nk = self.d / (2*self.lp) p1,p2,p3,p4 = -1.237, 0.8105, -1.0243, 0.4595 C = (1 + p1*(2*nk) + p2*(2*nk)**2) / (2*nk+p3*(2*nk)**2+p4*(2*nk)**3) u = self.kT * C * (1-np.cos(dr * np.pi / 180)) return u