From 15c953e2c8a4cb7a86a124d073991ccd6b8604fc Mon Sep 17 00:00:00 2001 From: James Graham <J.A.Graham@soton.ac.uk> Date: Fri, 5 Aug 2016 15:02:22 +0100 Subject: [PATCH] Perf improve for measuring bonds - more numba --- pycgtool/bondset.py | 49 +++++++++++---------------------------- pycgtool/interface.py | 2 +- pycgtool/util.py | 54 +++++++++++++++++++++++++++++++++++++++++-- 3 files changed, 66 insertions(+), 39 deletions(-) diff --git a/pycgtool/bondset.py b/pycgtool/bondset.py index b61b11d..73bbfa8 100644 --- a/pycgtool/bondset.py +++ b/pycgtool/bondset.py @@ -15,7 +15,8 @@ except ImportError: pass 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 np.seterr(all="raise") @@ -59,8 +60,10 @@ class Bond: rad2 = np.pi * np.pi / (180. * 180.) 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 / (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 try: @@ -80,27 +83,6 @@ class Bond: 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 used to perform bond measurements in a Frame. @@ -309,26 +291,22 @@ class BondSet: """ def calc_length(atoms): 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): veca = dist_with_pbc(atoms[0].coords, atoms[1].coords, frame.box) vecb = dist_with_pbc(atoms[1].coords, atoms[2].coords, frame.box) - ret = np.degrees(np.pi - angle(veca, vecb)) - return ret + return math.degrees(math.pi - vector_angle(veca, vecb)) def calc_dihedral(atoms): veca = dist_with_pbc(atoms[0].coords, atoms[1].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) - c1 = cross(veca, vecb) - c2 = cross(vecb, vecc) - c3 = cross(c1, c2) + c1 = vector_cross(veca, vecb) + c2 = vector_cross(vecb, vecc) - ang = np.degrees(angle(c1, c2)) - direction = np.dot(vecb, c3) - return ang if direction > 0 else -ang + return math.degrees(vector_angle_signed(c1, c2, vecb)) calc = {2: calc_length, 3: calc_angle, @@ -338,7 +316,7 @@ class BondSet: try: mol_meas = self._molecules[res.name] except KeyError: - # Bonds have not been specified for molecule - probably water + # Bonds have not been specified for molecule - probably water - ignore this residue continue adj_res = {"-": prev_res, @@ -350,8 +328,7 @@ class BondSet: atoms = [adj_res.get(name[0], res)[name.lstrip("-+")] for name in bond.atoms] val = calc[len(atoms)](atoms) bond.values.append(val) - except (NotImplementedError, TypeError, FloatingPointError): - # NotImplementedError is raised if form is not implemented + except (NotImplementedError, TypeError): # TypeError is raised when residues on end of chain calc bond to next pass @@ -366,7 +343,7 @@ class BondSet: if progress: try: 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: pass diff --git a/pycgtool/interface.py b/pycgtool/interface.py index 210a486..34b44bd 100644 --- a/pycgtool/interface.py +++ b/pycgtool/interface.py @@ -280,7 +280,7 @@ class Progress: else: try: 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: self._quiet = False collections.deque(self, maxlen=0) diff --git a/pycgtool/util.py b/pycgtool/util.py index 6292059..0f964c6 100644 --- a/pycgtool/util.py +++ b/pycgtool/util.py @@ -5,6 +5,7 @@ This module contains some general purpose utility functions used in PyCGTOOL. import os import itertools import random +import math import numpy as np np.seterr(all="raise") @@ -22,13 +23,14 @@ def jit_dummy(*args, **kwargs): return wrap try: + import numba from numba import jit except ImportError: jit = jit_dummy -@jit -def cross(u, v): +@jit(numba.float32[3](numba.float32[3], numba.float32[3])) +def vector_cross(u, v): """ Return vector cross product of two 3d vectors as numpy array. @@ -43,6 +45,54 @@ def cross(u, v): 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): """ Check if two node tuples are equivalent. Assumes undirected edges. -- GitLab