1# Copyright (C) 2004, Thomas Hamelryck (thamelry@binf.ku.dk)
2#
3# This file is part of the Biopython distribution and governed by your
4# choice of the "Biopython License Agreement" or the "BSD 3-Clause License".
5# Please see the LICENSE file that should have been included as part of this
6# package.
7
8"""Vector class, including rotation-related functions."""
9
10
11import numpy  # type: ignore
12from typing import Tuple, Optional
13
14
15def m2rotaxis(m):
16    """Return angles, axis pair that corresponds to rotation matrix m.
17
18    The case where ``m`` is the identity matrix corresponds to a singularity
19    where any rotation axis is valid. In that case, ``Vector([1, 0, 0])``,
20    is returned.
21    """
22    eps = 1e-5
23
24    # Check for singularities a la http://www.euclideanspace.com/maths/geometry/rotations/conversions/matrixToAngle/
25    if (
26        abs(m[0, 1] - m[1, 0]) < eps
27        and abs(m[0, 2] - m[2, 0]) < eps
28        and abs(m[1, 2] - m[2, 1]) < eps
29    ):
30        # Singularity encountered. Check if its 0 or 180 deg
31        if (
32            abs(m[0, 1] + m[1, 0]) < eps
33            and abs(m[0, 2] + m[2, 0]) < eps
34            and abs(m[1, 2] + m[2, 1]) < eps
35            and abs(m[0, 0] + m[1, 1] + m[2, 2] - 3) < eps
36        ):
37            angle = 0
38        else:
39            angle = numpy.pi
40    else:
41        # Angle always between 0 and pi
42        # Sense of rotation is defined by axis orientation
43        t = 0.5 * (numpy.trace(m) - 1)
44        t = max(-1, t)
45        t = min(1, t)
46        angle = numpy.arccos(t)
47
48    if angle < 1e-15:
49        # Angle is 0
50        return 0.0, Vector(1, 0, 0)
51    elif angle < numpy.pi:
52        # Angle is smaller than pi
53        x = m[2, 1] - m[1, 2]
54        y = m[0, 2] - m[2, 0]
55        z = m[1, 0] - m[0, 1]
56        axis = Vector(x, y, z)
57        axis.normalize()
58        return angle, axis
59    else:
60        # Angle is pi - special case!
61        m00 = m[0, 0]
62        m11 = m[1, 1]
63        m22 = m[2, 2]
64        if m00 > m11 and m00 > m22:
65            x = numpy.sqrt(m00 - m11 - m22 + 0.5)
66            y = m[0, 1] / (2 * x)
67            z = m[0, 2] / (2 * x)
68        elif m11 > m00 and m11 > m22:
69            y = numpy.sqrt(m11 - m00 - m22 + 0.5)
70            x = m[0, 1] / (2 * y)
71            z = m[1, 2] / (2 * y)
72        else:
73            z = numpy.sqrt(m22 - m00 - m11 + 0.5)
74            x = m[0, 2] / (2 * z)
75            y = m[1, 2] / (2 * z)
76        axis = Vector(x, y, z)
77        axis.normalize()
78        return numpy.pi, axis
79
80
81def vector_to_axis(line, point):
82    """Vector to axis method.
83
84    Return the vector between a point and
85    the closest point on a line (ie. the perpendicular
86    projection of the point on the line).
87
88    :type line: L{Vector}
89    :param line: vector defining a line
90
91    :type point: L{Vector}
92    :param point: vector defining the point
93    """
94    line = line.normalized()
95    np = point.norm()
96    angle = line.angle(point)
97    return point - line ** (np * numpy.cos(angle))
98
99
100def rotaxis2m(theta, vector):
101    """Calculate left multiplying rotation matrix.
102
103    Calculate a left multiplying rotation matrix that rotates
104    theta rad around vector.
105
106    :type theta: float
107    :param theta: the rotation angle
108
109    :type vector: L{Vector}
110    :param vector: the rotation axis
111
112    :return: The rotation matrix, a 3x3 Numeric array.
113
114    Examples
115    --------
116    >>> from numpy import pi
117    >>> from Bio.PDB.vectors import rotaxis2m
118    >>> from Bio.PDB.vectors import Vector
119    >>> m = rotaxis2m(pi, Vector(1, 0, 0))
120    >>> Vector(1, 2, 3).left_multiply(m)
121    <Vector 1.00, -2.00, -3.00>
122
123    """
124    vector = vector.normalized()
125    c = numpy.cos(theta)
126    s = numpy.sin(theta)
127    t = 1 - c
128    x, y, z = vector.get_array()
129    rot = numpy.zeros((3, 3))
130    # 1st row
131    rot[0, 0] = t * x * x + c
132    rot[0, 1] = t * x * y - s * z
133    rot[0, 2] = t * x * z + s * y
134    # 2nd row
135    rot[1, 0] = t * x * y + s * z
136    rot[1, 1] = t * y * y + c
137    rot[1, 2] = t * y * z - s * x
138    # 3rd row
139    rot[2, 0] = t * x * z - s * y
140    rot[2, 1] = t * y * z + s * x
141    rot[2, 2] = t * z * z + c
142    return rot
143
144
145rotaxis = rotaxis2m
146
147
148def refmat(p, q):
149    """Return a (left multiplying) matrix that mirrors p onto q.
150
151    :type p,q: L{Vector}
152    :return: The mirror operation, a 3x3 Numeric array.
153
154    Examples
155    --------
156    >>> from Bio.PDB.vectors import refmat
157    >>> p, q = Vector(1, 2, 3), Vector(2, 3, 5)
158    >>> mirror = refmat(p, q)
159    >>> qq = p.left_multiply(mirror)
160    >>> print(q)
161    <Vector 2.00, 3.00, 5.00>
162    >>> print(qq)
163    <Vector 1.21, 1.82, 3.03>
164
165    """
166    p = p.normalized()
167    q = q.normalized()
168    if (p - q).norm() < 1e-5:
169        return numpy.identity(3)
170    pq = p - q
171    pq.normalize()
172    b = pq.get_array()
173    b.shape = (3, 1)
174    i = numpy.identity(3)
175    ref = i - 2 * numpy.dot(b, numpy.transpose(b))
176    return ref
177
178
179def rotmat(p, q):
180    """Return a (left multiplying) matrix that rotates p onto q.
181
182    :param p: moving vector
183    :type p: L{Vector}
184
185    :param q: fixed vector
186    :type q: L{Vector}
187
188    :return: rotation matrix that rotates p onto q
189    :rtype: 3x3 Numeric array
190
191    Examples
192    --------
193    >>> from Bio.PDB.vectors import rotmat
194    >>> p, q = Vector(1, 2, 3), Vector(2, 3, 5)
195    >>> r = rotmat(p, q)
196    >>> print(q)
197    <Vector 2.00, 3.00, 5.00>
198    >>> print(p)
199    <Vector 1.00, 2.00, 3.00>
200    >>> p.left_multiply(r)
201    <Vector 1.21, 1.82, 3.03>
202
203    """
204    rot = numpy.dot(refmat(q, -p), refmat(p, -p))
205    return rot
206
207
208def calc_angle(v1, v2, v3):
209    """Calculate angle method.
210
211    Calculate the angle between 3 vectors
212    representing 3 connected points.
213
214    :param v1, v2, v3: the tree points that define the angle
215    :type v1, v2, v3: L{Vector}
216
217    :return: angle
218    :rtype: float
219    """
220    v1 = v1 - v2
221    v3 = v3 - v2
222    return v1.angle(v3)
223
224
225def calc_dihedral(v1, v2, v3, v4):
226    """Calculate dihedral angle method.
227
228    Calculate the dihedral angle between 4 vectors
229    representing 4 connected points. The angle is in
230    ]-pi, pi].
231
232    :param v1, v2, v3, v4: the four points that define the dihedral angle
233    :type v1, v2, v3, v4: L{Vector}
234    """
235    ab = v1 - v2
236    cb = v3 - v2
237    db = v4 - v3
238    u = ab ** cb
239    v = db ** cb
240    w = u ** v
241    angle = u.angle(v)
242    # Determine sign of angle
243    try:
244        if cb.angle(w) > 0.001:
245            angle = -angle
246    except ZeroDivisionError:
247        # dihedral=pi
248        pass
249    return angle
250
251
252class Vector:
253    """3D vector."""
254
255    def __init__(self, x, y=None, z=None):
256        """Initialize the class."""
257        if y is None and z is None:
258            # Array, list, tuple...
259            if len(x) != 3:
260                raise ValueError("Vector: x is not a list/tuple/array of 3 numbers")
261            self._ar = numpy.array(x, "d")
262        else:
263            # Three numbers
264            self._ar = numpy.array((x, y, z), "d")
265
266    def __repr__(self):
267        """Return vector 3D coordinates."""
268        x, y, z = self._ar
269        return "<Vector %.2f, %.2f, %.2f>" % (x, y, z)
270
271    def __neg__(self):
272        """Return Vector(-x, -y, -z)."""
273        a = -self._ar
274        return Vector(a)
275
276    def __add__(self, other):
277        """Return Vector+other Vector or scalar."""
278        if isinstance(other, Vector):
279            a = self._ar + other._ar
280        else:
281            a = self._ar + numpy.array(other)
282        return Vector(a)
283
284    def __sub__(self, other):
285        """Return Vector-other Vector or scalar."""
286        if isinstance(other, Vector):
287            a = self._ar - other._ar
288        else:
289            a = self._ar - numpy.array(other)
290        return Vector(a)
291
292    def __mul__(self, other):
293        """Return Vector.Vector (dot product)."""
294        return sum(self._ar * other._ar)
295
296    def __truediv__(self, x):
297        """Return Vector(coords/a)."""
298        a = self._ar / numpy.array(x)
299        return Vector(a)
300
301    def __pow__(self, other):
302        """Return VectorxVector (cross product) or Vectorxscalar."""
303        if isinstance(other, Vector):
304            a, b, c = self._ar
305            d, e, f = other._ar
306            c1 = numpy.linalg.det(numpy.array(((b, c), (e, f))))
307            c2 = -numpy.linalg.det(numpy.array(((a, c), (d, f))))
308            c3 = numpy.linalg.det(numpy.array(((a, b), (d, e))))
309            return Vector(c1, c2, c3)
310        else:
311            a = self._ar * numpy.array(other)
312            return Vector(a)
313
314    def __getitem__(self, i):
315        """Return value of array index i."""
316        return self._ar[i]
317
318    def __setitem__(self, i, value):
319        """Assign values to array index i."""
320        self._ar[i] = value
321
322    def __contains__(self, i):
323        """Validate if i is in array."""
324        return i in self._ar
325
326    def norm(self):
327        """Return vector norm."""
328        return numpy.sqrt(sum(self._ar * self._ar))
329
330    def normsq(self):
331        """Return square of vector norm."""
332        return abs(sum(self._ar * self._ar))
333
334    def normalize(self):
335        """Normalize the Vector object.
336
337        Changes the state of ``self`` and doesn't return a value.
338        If you need to chain function calls or create a new object
339        use the ``normalized`` method.
340        """
341        if self.norm():
342            self._ar = self._ar / self.norm()
343
344    def normalized(self):
345        """Return a normalized copy of the Vector.
346
347        To avoid allocating new objects use the ``normalize`` method.
348        """
349        v = self.copy()
350        v.normalize()
351        return v
352
353    def angle(self, other):
354        """Return angle between two vectors."""
355        n1 = self.norm()
356        n2 = other.norm()
357        c = (self * other) / (n1 * n2)
358        # Take care of roundoff errors
359        c = min(c, 1)
360        c = max(-1, c)
361        return numpy.arccos(c)
362
363    def get_array(self):
364        """Return (a copy of) the array of coordinates."""
365        return numpy.array(self._ar)
366
367    def left_multiply(self, matrix):
368        """Return Vector=Matrix x Vector."""
369        a = numpy.dot(matrix, self._ar)
370        return Vector(a)
371
372    def right_multiply(self, matrix):
373        """Return Vector=Vector x Matrix."""
374        a = numpy.dot(self._ar, matrix)
375        return Vector(a)
376
377    def copy(self):
378        """Return a deep copy of the Vector."""
379        return Vector(self._ar)
380
381
382"""Homogeneous matrix geometry routines.
383
384Rotation, translation, scale, and coordinate transformations.
385
386Robert T. Miller 2019
387"""
388
389
390def homog_rot_mtx(angle_rads: float, axis: str) -> numpy.array:
391    """Generate a 4x4 single-axis numpy rotation matrix.
392
393    :param float angle_rads: the desired rotation angle in radians
394    :param char axis: character specifying the rotation axis
395    """
396    cosang = numpy.cos(angle_rads)
397    sinang = numpy.sin(angle_rads)
398
399    if "z" == axis:
400        return numpy.array(
401            (
402                (cosang, -sinang, 0, 0),
403                (sinang, cosang, 0, 0),
404                (0, 0, 1, 0),
405                (0, 0, 0, 1),
406            ),
407            dtype=numpy.float64,
408        )
409    elif "y" == axis:
410        return numpy.array(
411            (
412                (cosang, 0, sinang, 0),
413                (0, 1, 0, 0),
414                (-sinang, 0, cosang, 0),
415                (0, 0, 0, 1),
416            ),
417            dtype=numpy.float64,
418        )
419    else:
420        return numpy.array(
421            (
422                (1, 0, 0, 0),
423                (0, cosang, -sinang, 0),
424                (0, sinang, cosang, 0),
425                (0, 0, 0, 1),
426            ),
427            dtype=numpy.float64,
428        )
429
430
431def set_Z_homog_rot_mtx(angle_rads: float, mtx: numpy.ndarray):
432    """Update existing Z rotation matrix to new angle."""
433    cosang = numpy.cos(angle_rads)
434    sinang = numpy.sin(angle_rads)
435
436    mtx[0][0] = mtx[1][1] = cosang
437    mtx[1][0] = sinang
438    mtx[0][1] = -sinang
439
440
441def set_Y_homog_rot_mtx(angle_rads: float, mtx: numpy.ndarray):
442    """Update existing Y rotation matrix to new angle."""
443    cosang = numpy.cos(angle_rads)
444    sinang = numpy.sin(angle_rads)
445
446    mtx[0][0] = mtx[2][2] = cosang
447    mtx[0][2] = sinang
448    mtx[2][0] = -sinang
449
450
451def set_X_homog_rot_mtx(angle_rads: float, mtx: numpy.ndarray):
452    """Update existing X rotation matrix to new angle."""
453    cosang = numpy.cos(angle_rads)
454    sinang = numpy.sin(angle_rads)
455
456    mtx[1][1] = mtx[2][2] = cosang
457    mtx[2][1] = sinang
458    mtx[1][2] = -sinang
459
460
461def homog_trans_mtx(x: float, y: float, z: float) -> numpy.array:
462    """Generate a 4x4 numpy translation matrix.
463
464    :param x, y, z: translation in each axis
465    """
466    return numpy.array(
467        ((1, 0, 0, x), (0, 1, 0, y), (0, 0, 1, z), (0, 0, 0, 1)), dtype=numpy.float64
468    )
469
470
471def set_homog_trans_mtx(x: float, y: float, z: float, mtx: numpy.ndarray):
472    """Update existing translation matrix to new values."""
473    mtx[0][3] = x
474    mtx[1][3] = y
475    mtx[2][3] = z
476
477
478def homog_scale_mtx(scale: float) -> numpy.array:
479    """Generate a 4x4 numpy scaling matrix.
480
481    :param float scale: scale multiplier
482    """
483    return numpy.array(
484        [[scale, 0, 0, 0], [0, scale, 0, 0], [0, 0, scale, 0], [0, 0, 0, 1]],
485        dtype=numpy.float64,
486    )
487
488
489def _get_azimuth(x: float, y: float) -> float:
490    sign_y = -1.0 if y < 0.0 else 1.0
491    sign_x = -1.0 if x < 0.0 else 1.0
492    return (
493        numpy.arctan2(y, x)
494        if (0 != x and 0 != y)
495        else (numpy.pi / 2.0 * sign_y)  # +/-90 if X=0, Y!=0
496        if 0 != y
497        else numpy.pi
498        if sign_x < 0.0  # 180 if Y=0, X < 0
499        else 0.0  # 0 if Y=0, X >= 0
500    )
501
502
503def get_spherical_coordinates(xyz: numpy.array) -> Tuple[float, float, float]:
504    """Compute spherical coordinates (r, azimuth, polar_angle) for X,Y,Z point.
505
506    :param array xyz: column vector (3 row x 1 column numpy array)
507    :return: tuple of r, azimuth, polar_angle for input coordinate
508    """
509    r = numpy.linalg.norm(xyz)
510    if 0 == r:
511        return (0, 0, 0)
512    azimuth = _get_azimuth(xyz[0], xyz[1])
513    polar_angle = numpy.arccos(xyz[2] / r)
514
515    return (r, azimuth, polar_angle)
516
517
518gtm = numpy.identity(4, dtype=numpy.float64)
519gmrz = numpy.identity(4, dtype=numpy.float64)
520gmry = numpy.identity(4, dtype=numpy.float64)
521gmrz2 = numpy.identity(4, dtype=numpy.float64)
522
523
524def coord_space(
525    a0: numpy.ndarray, a1: numpy.ndarray, a2: numpy.ndarray, rev: bool = False
526) -> Tuple[numpy.ndarray, Optional[numpy.ndarray]]:
527    """Generate transformation matrix to coordinate space defined by 3 points.
528
529    New coordinate space will have:
530        acs[0] on XZ plane
531        acs[1] origin
532        acs[2] on +Z axis
533
534    :param numpy column array x3 acs: X,Y,Z column input coordinates x3
535    :param bool rev: if True, also return reverse transformation matrix
536        (to return from coord_space)
537    :returns: 4x4 numpy array, x2 if rev=True
538    """
539    # dbg = False
540    # if dbg:
541    #    print(a0.transpose())
542    #    print(a1.transpose())
543    #    print(a2.transpose())
544
545    # a0 = acs[0]
546    # a1 = acs[1]
547    # a2 = acs[2]
548
549    global gtm
550    global gmry
551    global gmrz, gmrz2
552
553    tm = gtm
554    mry = gmry
555    mrz = gmrz
556    mrz2 = gmrz2
557
558    # tx acs[1] to origin
559    # tm = homog_trans_mtx(-a1[0][0], -a1[1][0], -a1[2][0])
560    set_homog_trans_mtx(-a1[0], -a1[1], -a1[2], tm)
561
562    # directly translate a2 using a1
563    p = a2 - a1
564    sc = get_spherical_coordinates(p)
565
566    # if dbg:
567    #    print("p", p.transpose())
568    #    print("sc", sc)
569
570    # mrz = homog_rot_mtx(-sc[1], "z")  # rotate translated a2 -azimuth about Z
571    set_Z_homog_rot_mtx(-sc[1], mrz)
572    # mry = homog_rot_mtx(-sc[2], "y")  # rotate translated a2 -polar_angle about Y
573    set_Y_homog_rot_mtx(-sc[2], mry)
574
575    # mt completes a1-a2 on Z-axis, still need to align a0 with XZ plane
576    # mt = mry @ mrz @ tm  # python 3.5 and later
577    mt = gmry.dot(gmrz.dot(gtm))
578
579    # if dbg:
580    #    print("tm:\n", tm)
581    #    print("mrz:\n", mrz)
582    #    print("mry:\n", mry)
583    #    # print("mt ", mt)
584
585    p = mt.dot(a0)
586
587    # if dbg:
588    #    print("mt:\n", mt, "\na0:\n", a0, "\np:\n", p)
589
590    # need azimuth of translated a0
591    # sc2 = get_spherical_coordinates(p)
592    # print(sc2)
593    azimuth2 = _get_azimuth(p[0], p[1])
594
595    # rotate a0 -azimuth2 about Z to align with X
596    # mrz2 = homog_rot_mtx(-azimuth2, "z")
597    set_Z_homog_rot_mtx(-azimuth2, mrz2)
598
599    # mt = mrz2 @ mt
600    mt = gmrz2.dot(mt)
601
602    # if dbg:
603    #    print("mt:", mt, "\na0:", a0, "\np:", p)
604    #    # print(p, "\n", azimuth2, "\n", mrz2, "\n", mt)
605
606    # if dbg:
607    #    print("mt:\n", mt)
608    #    print("<<<<<<==============================")
609
610    if not rev:
611        return mt, None
612
613    # rev=True, so generate the reverse transformation
614
615    # rotate a0 theta about Z, reversing alignment with X
616    # mrz2 = homog_rot_mtx(azimuth2, "z")
617    set_Z_homog_rot_mtx(azimuth2, mrz2)
618    # rotate a2 phi about Y
619    # mry = homog_rot_mtx(sc[2], "y")
620    set_Y_homog_rot_mtx(sc[2], mry)
621    # rotate a2 theta about Z
622    # mrz = homog_rot_mtx(sc[1], "z")
623    set_Z_homog_rot_mtx(sc[1], mrz)
624    # translation matrix origin to a1
625    # tm = homog_trans_mtx(a1[0][0], a1[1][0], a1[2][0])
626    set_homog_trans_mtx(a1[0], a1[1], a1[2], tm)
627
628    # mr = tm @ mrz @ mry @ mrz2
629    mr = gtm.dot(gmrz.dot(gmry.dot(gmrz2)))
630    # mr = numpy.dot(tm, numpy.dot(mrz, numpy.dot(mry, mrz2)))
631
632    return mt, mr
633
634
635def multi_rot_Z(angle_rads: numpy.ndarray) -> numpy.ndarray:
636    """Create [entries] numpy Z rotation matrices for [entries] angles.
637
638    :param entries: int number of matrices generated.
639    :param angle_rads: numpy array of angles
640    :returns: entries x 4 x 4 homogeneous rotation matrices
641    """
642    rz = numpy.empty((angle_rads.shape[0], 4, 4))
643    rz[...] = numpy.identity(4)
644    rz[:, 0, 0] = rz[:, 1, 1] = numpy.cos(angle_rads)
645    rz[:, 1, 0] = numpy.sin(angle_rads)
646    rz[:, 0, 1] = -rz[:, 1, 0]
647    return rz
648
649
650def multi_rot_Y(angle_rads: numpy.ndarray) -> numpy.ndarray:
651    """Create [entries] numpy Y rotation matrices for [entries] angles.
652
653    :param entries: int number of matrices generated.
654    :param angle_rads: numpy array of angles
655    :returns: entries x 4 x 4 homogeneous rotation matrices
656    """
657    ry = numpy.empty((angle_rads.shape[0], 4, 4))
658    ry[...] = numpy.identity(4)
659    ry[:, 0, 0] = ry[:, 2, 2] = numpy.cos(angle_rads)
660    ry[:, 0, 2] = numpy.sin(angle_rads)
661    ry[:, 2, 0] = -ry[:, 0, 2]
662
663    return ry
664