Skip to content
Snippets Groups Projects
Commit 15c953e2 authored by James Graham's avatar James Graham
Browse files

Perf improve for measuring bonds - more numba

parent dd07c2ad
No related branches found
No related tags found
No related merge requests found
...@@ -15,7 +15,8 @@ except ImportError: ...@@ -15,7 +15,8 @@ except ImportError:
pass pass
from .util import stat_moments, sliding, dist_with_pbc, transpose_and_sample from .util import stat_moments, sliding, dist_with_pbc, transpose_and_sample
from .util import extend_graph_chain, cross, backup_file from .util import extend_graph_chain, backup_file
from .util import vector_len, vector_cross, vector_angle, vector_angle_signed
from .parsers.cfg import CFG from .parsers.cfg import CFG
np.seterr(all="raise") np.seterr(all="raise")
...@@ -59,8 +60,10 @@ class Bond: ...@@ -59,8 +60,10 @@ class Bond:
rad2 = np.pi * np.pi / (180. * 180.) rad2 = np.pi * np.pi / (180. * 180.)
conv = {2: lambda: rt / var, conv = {2: lambda: rt / var,
3: lambda: 25. if angle_default_fc else rt * np.sin(np.radians(mean))**2 / (var * rad2), 3: lambda: 25. if angle_default_fc else rt * np.sin(np.radians(mean))**2 / (var * rad2),
# 3: lambda: 25. if angle_default_fc else rt / (var * rad2),
4: lambda: rt / (var * rad2)} 4: lambda: rt / (var * rad2)}
# conv = {2: lambda: rt / (2 * var),
# 3: lambda: 25. if angle_default_fc else rt / (2 * np.sin(np.radians(mean))**2 * var * rad2),
# 4: lambda: rt / (var * rad2)}
self.eqm = mean self.eqm = mean
try: try:
...@@ -80,27 +83,6 @@ class Bond: ...@@ -80,27 +83,6 @@ class Bond:
return "<Bond containing atoms {0}>".format(", ".join(self.atoms)) return "<Bond containing atoms {0}>".format(", ".join(self.atoms))
def angle(a, b):
"""
Calculate the angle between two vectors.
:param a: First vector
:param b: Second vector
:return: Angle in radians
"""
dot = a[0]*b[0] + a[1]*b[1] + a[2]*b[2]
mag = math.sqrt((a[0]*a[0] + a[1]*a[1] + a[2]*a[2]) * (b[0]*b[0] + b[1]*b[1] + b[2]*b[2]))
try:
return math.acos(dot / mag)
except ValueError as e:
val = dot / mag
if abs(val) - 1 < 0.01:
# If within 1% of acceptable value correct it, else reraise
return math.acos(max(-1, min(1, dot / mag)))
e.args = (e.args[0] + " in acos(" + str(dot / mag) + ")",)
raise
class BondSet: class BondSet:
""" """
Class used to perform bond measurements in a Frame. Class used to perform bond measurements in a Frame.
...@@ -309,26 +291,22 @@ class BondSet: ...@@ -309,26 +291,22 @@ class BondSet:
""" """
def calc_length(atoms): def calc_length(atoms):
vec = dist_with_pbc(atoms[0].coords, atoms[1].coords, frame.box) vec = dist_with_pbc(atoms[0].coords, atoms[1].coords, frame.box)
return math.sqrt(vec[0]*vec[0] + vec[1]*vec[1] + vec[2]*vec[2]) return vector_len(vec)
def calc_angle(atoms): def calc_angle(atoms):
veca = dist_with_pbc(atoms[0].coords, atoms[1].coords, frame.box) veca = dist_with_pbc(atoms[0].coords, atoms[1].coords, frame.box)
vecb = dist_with_pbc(atoms[1].coords, atoms[2].coords, frame.box) vecb = dist_with_pbc(atoms[1].coords, atoms[2].coords, frame.box)
ret = np.degrees(np.pi - angle(veca, vecb)) return math.degrees(math.pi - vector_angle(veca, vecb))
return ret
def calc_dihedral(atoms): def calc_dihedral(atoms):
veca = dist_with_pbc(atoms[0].coords, atoms[1].coords, frame.box) veca = dist_with_pbc(atoms[0].coords, atoms[1].coords, frame.box)
vecb = dist_with_pbc(atoms[1].coords, atoms[2].coords, frame.box) vecb = dist_with_pbc(atoms[1].coords, atoms[2].coords, frame.box)
vecc = dist_with_pbc(atoms[2].coords, atoms[3].coords, frame.box) vecc = dist_with_pbc(atoms[2].coords, atoms[3].coords, frame.box)
c1 = cross(veca, vecb) c1 = vector_cross(veca, vecb)
c2 = cross(vecb, vecc) c2 = vector_cross(vecb, vecc)
c3 = cross(c1, c2)
ang = np.degrees(angle(c1, c2)) return math.degrees(vector_angle_signed(c1, c2, vecb))
direction = np.dot(vecb, c3)
return ang if direction > 0 else -ang
calc = {2: calc_length, calc = {2: calc_length,
3: calc_angle, 3: calc_angle,
...@@ -338,7 +316,7 @@ class BondSet: ...@@ -338,7 +316,7 @@ class BondSet:
try: try:
mol_meas = self._molecules[res.name] mol_meas = self._molecules[res.name]
except KeyError: except KeyError:
# Bonds have not been specified for molecule - probably water # Bonds have not been specified for molecule - probably water - ignore this residue
continue continue
adj_res = {"-": prev_res, adj_res = {"-": prev_res,
...@@ -350,8 +328,7 @@ class BondSet: ...@@ -350,8 +328,7 @@ class BondSet:
atoms = [adj_res.get(name[0], res)[name.lstrip("-+")] for name in bond.atoms] atoms = [adj_res.get(name[0], res)[name.lstrip("-+")] for name in bond.atoms]
val = calc[len(atoms)](atoms) val = calc[len(atoms)](atoms)
bond.values.append(val) bond.values.append(val)
except (NotImplementedError, TypeError, FloatingPointError): except (NotImplementedError, TypeError):
# NotImplementedError is raised if form is not implemented
# TypeError is raised when residues on end of chain calc bond to next # TypeError is raised when residues on end of chain calc bond to next
pass pass
...@@ -366,7 +343,7 @@ class BondSet: ...@@ -366,7 +343,7 @@ class BondSet:
if progress: if progress:
try: try:
total = sum(map(len, self._molecules.values())) total = sum(map(len, self._molecules.values()))
bond_iter_wrap = tqdm(bond_iter, total=total) bond_iter_wrap = tqdm(bond_iter, total=total, ncols=80)
except NameError: except NameError:
pass pass
......
...@@ -280,7 +280,7 @@ class Progress: ...@@ -280,7 +280,7 @@ class Progress:
else: else:
try: try:
self._quiet = True self._quiet = True
collections.deque(tqdm(self, total=len(self)-1), maxlen=0) collections.deque(tqdm(self, total=len(self)-1, ncols=80), maxlen=0)
except NameError: except NameError:
self._quiet = False self._quiet = False
collections.deque(self, maxlen=0) collections.deque(self, maxlen=0)
......
...@@ -5,6 +5,7 @@ This module contains some general purpose utility functions used in PyCGTOOL. ...@@ -5,6 +5,7 @@ This module contains some general purpose utility functions used in PyCGTOOL.
import os import os
import itertools import itertools
import random import random
import math
import numpy as np import numpy as np
np.seterr(all="raise") np.seterr(all="raise")
...@@ -22,13 +23,14 @@ def jit_dummy(*args, **kwargs): ...@@ -22,13 +23,14 @@ def jit_dummy(*args, **kwargs):
return wrap return wrap
try: try:
import numba
from numba import jit from numba import jit
except ImportError: except ImportError:
jit = jit_dummy jit = jit_dummy
@jit @jit(numba.float32[3](numba.float32[3], numba.float32[3]))
def cross(u, v): def vector_cross(u, v):
""" """
Return vector cross product of two 3d vectors as numpy array. Return vector cross product of two 3d vectors as numpy array.
...@@ -43,6 +45,54 @@ def cross(u, v): ...@@ -43,6 +45,54 @@ def cross(u, v):
return res return res
@jit(numba.float32(numba.float32[3], numba.float32[3]))
def vector_dot(u, v):
"""
Return vector dot product of two 3d vectors.
:param u: First 3d vector
:param v: Second 3d vector
:return: Dot product of two vectors
"""
return u[0]*v[0] + u[1]*v[1] + u[2]*v[2]
@jit(numba.float32(numba.float32[3]))
def vector_len(v):
return math.sqrt(v[0]*v[0] + v[1]*v[1] + v[2]*v[2])
@jit(numba.float32(numba.float32[3], numba.float32[3]))
def vector_angle(a, b):
"""
Calculate the angle between two vectors.
:param a: First vector
:param b: Second vector
:return: Angle in radians
"""
mag = vector_len(a) * vector_len(b)
dot = vector_dot(a, b) / mag
# Previously checked to within 1%, but this prevented numba optimisation
ang = math.acos(max(-1, min(1, dot)))
return ang
@jit
def vector_angle_signed(a, b, ref=np.array([0., 0., 1.])):
"""
Calculate the signed angle between two vectors.
:param a: First vector
:param b: Second vector
:param ref: Optional reference vector, will use global z-axis if not provided
:return: Signed angle in radians
"""
ang = vector_angle(a, b)
signum = math.copysign(1, vector_dot(vector_cross(a, b), ref))
return ang * signum
def tuple_equivalent(tuple1, tuple2): def tuple_equivalent(tuple1, tuple2):
""" """
Check if two node tuples are equivalent. Assumes undirected edges. Check if two node tuples are equivalent. Assumes undirected edges.
......
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Please register or to comment