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