def soft_sphere_neighbor_list(displacement_or_metric, box_size, species=None, sigma=1.0, epsilon=1.0, alpha=2.0, dr_threshold=0.2): """Convenience wrapper to compute soft spheres using a neighbor list.""" sigma = np.array(sigma, dtype=f32) epsilon = np.array(epsilon, dtype=f32) alpha = np.array(alpha, dtype=f32) list_cutoff = f32(np.max(sigma)) dr_threshold = f32(list_cutoff * dr_threshold) neighbor_fn = partition.neighbor_list(displacement_or_metric, box_size, list_cutoff, dr_threshold) energy_fn = smap.pair_neighbor_list( soft_sphere, space.canonicalize_displacement_or_metric(displacement_or_metric), species=species, sigma=sigma, epsilon=epsilon, alpha=alpha) return neighbor_fn, energy_fn
def periodic_displacement(side, dR): """Wraps displacement vectors into a hypercube. Args: side: Specification of hypercube size. Either, (a) float if all sides have equal length. (b) ndarray(spatial_dim) if sides have different lengths. dR: Matrix of displacements; ndarray(shape=[..., spatial_dim]). Returns: Matrix of wrapped displacements; ndarray(shape=[..., spatial_dim]). """ return np.mod(dR + side * f32(0.5), side) - f32(0.5) * side
def load_lammps_eam_parameters(f): """Reads EAM parameters from a LAMMPS file and returns relevant spline fits. This function reads single-element EAM potential fit parameters from a file in DYNAMO funcl format. In summary, the file contains: Line 1-3: comments, Line 4: Number of elements and the element type, Line 5: The number of charge values that the embedding energy is evaluated on (num_drho), interval between the charge values (drho), the number of distances the pairwise energy and the charge density is evaluated on (num_dr), the interval between these distances (dr), and the cutoff distance (cutoff). The lines that come after are the embedding function evaluated on num_drho charge values, charge function evaluated at num_dr distance values, and pairwise energy evaluated at num_dr distance values. Note that the pairwise energy is multiplied by distance (in units of eV x Angstroms). For more details of the DYNAMO file format, see: https://sites.google.com/a/ncsu.edu/cjobrien/tutorials-and-guides/eam Args: f: File handle for the EAM parameters text file. Returns: charge_fn: A function that takes an ndarray of shape [n, m] of distances between particles and returns a matrix of charge contributions. embedding_fn: Function that takes an ndarray of shape [n] of charges and returns an ndarray of shape [n] of the energy cost of embedding an atom into the charge. pairwise_fn: A function that takes an ndarray of shape [n, m] of distances and returns an ndarray of shape [n, m] of pairwise energies. cutoff: Cutoff distance for the embedding_fn and pairwise_fn. """ raw_text = f.read().split('\n') if 'setfl' not in raw_text[0]: raise ValueError( 'File format is incorrect, expected LAMMPS setfl format.') temp_params = raw_text[4].split() num_drho, num_dr = int(temp_params[0]), int(temp_params[2]) drho, dr, cutoff = float(temp_params[1]), float(temp_params[3]), float( temp_params[4]) data = np.array(map(float, raw_text[6:-1])) embedding_fn = spline(data[:num_drho], drho) charge_fn = spline(data[num_drho:num_drho + num_dr], dr) # LAMMPS EAM parameters file lists pairwise energies after multiplying by # distance, in units of eV*Angstrom. We divide the energy by distance below, distances = np.arange(num_dr) * dr # Prevent dividing by zero at zero distance, which will not # affect the calculation distances = np.where(distances == 0, f32(0.001), distances) pairwise_fn = spline( data[num_dr + num_drho:num_drho + f32(2) * num_dr] / distances, dr) return charge_fn, embedding_fn, pairwise_fn, cutoff
def bks_neighbor_list(displacement_or_metric, box_size, species, Q_sq, exp_coeff, exp_decay, attractive_coeff, repulsive_coeff, coulomb_alpha, cutoff, dr_threshold=0.8): Q_sq = np.array(Q_sq, f32) exp_coeff = np.array(exp_coeff, f32) exp_decay = np.array(exp_decay, f32) attractive_coeff = np.array(attractive_coeff, f32) repulsive_coeff = np.array(repulsive_coeff, f32) dr_threshold = f32(dr_threshold) neighbor_fn = partition.neighbor_list(displacement_or_metric, box_size, cutoff, dr_threshold) energy_fn = smap.pair_neighbor_list( bks, space.canonicalize_displacement_or_metric(displacement_or_metric), species=species, Q_sq=Q_sq, exp_coeff=exp_coeff, exp_decay=exp_decay, attractive_coeff=attractive_coeff, repulsive_coeff=repulsive_coeff, coulomb_alpha=coulomb_alpha, cutoff=cutoff) return neighbor_fn, energy_fn
def multiplicative_isotropic_cutoff(fn, r_onset, r_cutoff): """Takes an isotropic function and constructs a truncated function. Given a function f:R -> R, we construct a new function f':R -> R such that f'(r) = f(r) for r < r_onset, f'(r) = 0 for r > r_cutoff, and f(r) is C^1 everywhere. To do this, we follow the approach outlined in HOOMD Blue [1] (thanks to Carl Goodrich for the pointer). We construct a function S(r) such that S(r) = 1 for r < r_onset, S(r) = 0 for r > r_cutoff, and S(r) is C^1. Then f'(r) = S(r)f(r). Args: fn: A function that takes an ndarray of distances of shape [n, m] as well as varargs. r_onset: A float specifying the onset radius of deformation. r_cutoff: A float specifying the cutoff radius. Returns: A new function with the same signature as fn, with the properties outlined above. [1] HOOMD Blue documentation. Accessed on 05/31/2019. https://hoomd-blue.readthedocs.io/en/stable/module-md-pair.html#hoomd.md.pair.pair """ r_c = r_cutoff ** f32(2) r_o = r_onset ** f32(2) def smooth_fn(dr): r = dr ** f32(2) return np.where( dr < r_onset, f32(1), np.where( dr < r_cutoff, (r_c - r) ** f32(2) * (r_c + f32(2) * r - f32(3) * r_o) / ( r_c - r_o) ** f32(3), f32(0) ) ) @wraps(fn) def cutoff_fn(dr, *args, **kwargs): return smooth_fn(dr) * fn(dr, *args, **kwargs) return cutoff_fn
def energy(R, **kwargs): dr = metric(R, R, **kwargs) total_charge = smap._high_precision_sum(charge_fn(dr), axis=1) embedding_energy = embedding_fn(total_charge) pairwise_energy = smap._high_precision_sum( smap._diagonal_mask(pairwise_fn(dr)), axis=1) / f32(2.0) return smap._high_precision_sum(embedding_energy + pairwise_energy, axis=axis)
def morse(dr, sigma=1.0, epsilon=5.0, alpha=5.0, **unused_kwargs): """Morse interaction between particles with a minimum at r0. Args: dr: An ndarray of shape [n, m] of pairwise distances between particles. sigma: Distance between particles where the energy has a minimum. Should either be a floating point scalar or an ndarray whose shape is [n, m]. epsilon: Interaction energy scale. Should either be a floating point scalar or an ndarray whose shape is [n, m]. alpha: Range parameter. Should either be a floating point scalar or an ndarray whose shape is [n, m]. unused_kwargs: Allows extra data (e.g. time) to be passed to the energy. Returns: Matrix of energies of shape [n, m]. """ check_kwargs_time_dependence(unused_kwargs) U = epsilon * (f32(1) - np.exp(-alpha * (dr - sigma)))**f32(2) - epsilon return np.nan_to_num(np.array(U, dtype=dr.dtype))
def smooth_fn(dr): r = dr**f32(2) return np.where( dr < r_onset, f32(1), np.where(dr < r_cutoff, (r_c - r)**f32(2) * (r_c + f32(2) * r - f32(3) * r_o) / (r_c - r_o)**f32(3), f32(0)))
def lennard_jones(dr, sigma=1, epsilon=1, **unused_kwargs): """Lennard-Jones interaction between particles with a minimum at sigma. Args: dr: An ndarray of shape [n, m] of pairwise distances between particles. sigma: Distance between particles where the energy has a minimum. Should either be a floating point scalar or an ndarray whose shape is [n, m]. epsilon: Interaction energy scale. Should either be a floating point scalar or an ndarray whose shape is [n, m]. unused_kwargs: Allows extra data (e.g. time) to be passed to the energy. Returns: Matrix of energies of shape [n, m]. """ check_kwargs_time_dependence(unused_kwargs) dr = (sigma / dr)**f32(2) idr6 = dr**f32(3) idr12 = idr6**f32(2) # TODO(schsam): This seems potentially dangerous. We should do ErrorChecking # here. return np.nan_to_num(f32(4) * epsilon * (idr12 - idr6))
def soft_sphere(dr, sigma=1, epsilon=1, alpha=2, **unused_kwargs): """Finite ranged repulsive interaction between soft spheres. Args: dr: An ndarray of shape [n, m] of pairwise distances between particles. sigma: Particle diameter. Should either be a floating point scalar or an ndarray whose shape is [n, m]. epsilon: Interaction energy scale. Should either be a floating point scalar or an ndarray whose shape is [n, m]. alpha: Exponent specifying interaction stiffness. Should either be a float point scalar or an ndarray whose shape is [n, m]. unused_kwargs: Allows extra data (e.g. time) to be passed to the energy. Returns: Matrix of energies whose shape is [n, m]. """ check_kwargs_time_dependence(unused_kwargs) dr = dr / sigma U = epsilon * np.where(dr < 1.0, f32(1.0) / alpha * (f32(1.0) - dr)**alpha, f32(0.0)) return U
def main(unused_argv): key = random.PRNGKey(0) # Setup some variables describing the system. N = 500 dimension = 2 box_size = f32(25.0) # Create helper functions to define a periodic box of some size. displacement, shift = space.periodic(box_size) metric = space.metric(displacement) # Use JAX's random number generator to generate random initial positions. key, split = random.split(key) R = random.uniform(split, (N, dimension), minval=0.0, maxval=box_size, dtype=f32) # The system ought to be a 50:50 mixture of two types of particles, one # large and one small. sigma = np.array([[1.0, 1.2], [1.2, 1.4]], dtype=f32) N_2 = int(N / 2) species = np.array([0] * N_2 + [1] * N_2, dtype=i32) # Create an energy function. energy_fn = energy.soft_sphere_pair(displacement, species, sigma) force_fn = quantity.force(energy_fn) # Create a minimizer. init_fn, apply_fn = minimize.fire_descent(energy_fn, shift) opt_state = init_fn(R) # Minimize the system. minimize_steps = 50 print_every = 10 print('Minimizing.') print('Step\tEnergy\tMax Force') print('-----------------------------------') for step in range(minimize_steps): opt_state = apply_fn(opt_state) if step % print_every == 0: R = opt_state.position print('{:.2f}\t{:.2f}\t{:.2f}'.format(step, energy_fn(R), np.max(force_fn(R))))
def displacement_fn(Ra, Rb, **kwargs): _box, _inv_box = box, inv_box if 'box' in kwargs: _box = kwargs['box'] if not fractional_coordinates: _inv_box = inverse(_box) if 'new_box' in kwargs: _box = kwargs['new_box'] if not fractional_coordinates: Ra = transform(_inv_box, Ra) Rb = transform(_inv_box, Rb) dR = periodic_displacement(f32(1.0), pairwise_displacement(Ra, Rb)) return transform(_box, dR)
def displacement_fn(Ra, Rb, perturbation=None, **kwargs): _box, _inv_box = box, inv_box if 'box' in kwargs: _box = kwargs['box'] if not fractional_coordinates: _inv_box = inverse(_box) if 'new_box' in kwargs: _box = kwargs['new_box'] if not fractional_coordinates: Ra = transform(_inv_box, Ra) Rb = transform(_inv_box, Rb) dR = periodic_displacement(f32(1.0), pairwise_displacement(Ra, Rb)) dR = transform(_box, dR) if perturbation is not None: dR = raw_transform(perturbation, dR) return dR
def shift(R, dR, **unused_kwargs): check_kwargs_time_dependence(unused_kwargs) return periodic_shift(f32(1.0), R, transform(T_inv, dR))
def displacement(Ra, Rb, **unused_kwargs): check_kwargs_time_dependence(unused_kwargs) dR = periodic_displacement(f32(1.0), pairwise_displacement(Ra, Rb)) return transform(T, dR)
def shift(R, dR, t=None, **unused_kwargs): _check_time_dependence(t) check_kwargs_empty(unused_kwargs) return periodic_shift(f32(1.0), R, transform(_small_inverse(T(t)), dR))
def displacement(Ra, Rb, t=None, **unused_kwargs): _check_time_dependence(t) check_kwargs_empty(unused_kwargs) dR = periodic_displacement(f32(1.0), pairwise_displacement(Ra, Rb)) return transform(T(t), dR)
def u(R, dR): if wrapped: return periodic_shift(f32(1.0), R, dR) return R + dR
def shift(R: Array, dR: Array, **kwargs) -> Array: return periodic_shift(f32(1.0), R, transform(_small_inverse(T(**kwargs)), dR))
def displacement(Ra: Array, Rb: Array, **unused_kwargs) -> Array: dR = periodic_displacement(f32(1.0), pairwise_displacement(Ra, Rb)) return transform(T, dR)
def shift(R: Array, dR: Array, **unused_kwargs) -> Array: return periodic_shift(f32(1.0), R, transform(T_inv, dR))