Source code for espaloma.mm.torsion

# =============================================================================
# IMPORTS
# =============================================================================
import espaloma as esp


# =============================================================================
# MODULE FUNCTIONS
# =============================================================================
[docs]def periodic_torsion( x, k, periodicity=list(range(1, 7)), phases=[0.0 for _ in range(6)] ): """Periodic torsion potential Parameters ---------- x : `torch.Tensor`, `shape = (batch_size, 1)` Dihedral value. k : `torch.Tensor`, `shape = (batch_size, n_phases)` Force constants. periodicity : `torch.Tensor`, `shape = (batch_size, n_phases)` Periodicities phases : `torch.Tensor`, `shape = (batch_size, n_phases)` Phase offsets Returns ------- u : `torch.Tensor`, `shape = (batch_size, 1)` Energy. """ # NOTE: # 0.5 because all torsions are calculated twice out = 0.5 * esp.mm.functional.periodic( x=x, k=k, periodicity=periodicity, phases=phases, ) # assert(out.shape == (len(x), 1)) return out
[docs]def angle_angle( u_angle_left, u_angle_right, k_angle_angle, ): u_angle_left = u_angle_left - u_angle_left.min(dim=-1, keepdims=True)[0] u_angle_right = ( u_angle_right - u_angle_right.min(dim=-1, keepdims=True)[0] ) return k_angle_angle * (u_angle_left ** 0.5) * (u_angle_right ** 0.5)
[docs]def angle_torsion( u_angle_left, u_angle_right, u_torsion, k_angle_torsion, ): u_angle_left = u_angle_left - u_angle_left.min(dim=-1, keepdims=True)[0] u_angle_right = ( u_angle_right - u_angle_right.min(dim=-1, keepdims=True)[0] ) return ( k_angle_torsion * (u_angle_left ** 0.5) * u_torsion + k_angle_torsion * (u_angle_right ** 0.5) * u_torsion )
[docs]def angle_angle_torsion( u_angle_left, u_angle_right, u_torsion, k_angle_angle_torsion, ): u_angle_left = u_angle_left - u_angle_left.min(dim=-1, keepdims=True)[0] u_angle_right = ( u_angle_right - u_angle_right.min(dim=-1, keepdims=True)[0] ) return ( k_angle_angle_torsion * (u_angle_left ** 0.5) * (u_angle_right ** 0.5) * u_torsion )
[docs]def bond_torsion( u_bond_left, u_bond_right, u_bond_center, u_torsion, k_side_torsion, k_center_torsion, ): u_bond_left = u_bond_left - u_bond_left.min(dim=-1, keepdims=True)[0] u_bond_right = u_bond_right - u_bond_right.min(dim=-1, keepdims=True)[0] u_bond_center = ( u_bond_center - u_bond_center.min(dim=-1, keepdims=True)[0] ) return ( k_side_torsion * u_torsion * (u_bond_left ** 0.5) + k_side_torsion * u_torsion * (u_bond_right ** 0.5) + k_center_torsion * u_torsion * (u_bond_center ** 0.5) )