spatial/r3: implement Rotate

Updates gonum#1513.
BarrensZeppelin · Dec 1, 2020 · 9c06200 · 9c06200
1 parent 8d1ffe1
commit 9c06200
Show file tree

Hide file tree

Showing 2 changed files with 92 additions and 1 deletion.
diff --git a/spatial/r3/vector.go b/spatial/r3/vector.go
@@ -4,7 +4,11 @@
 
 package r3
 
-import "math"
+import (
+	"math"
+
+	"gonum.org/v1/gonum/num/quat"
+)
 
 // Vec is a 3D vector.
 type Vec struct {
@@ -49,6 +53,11 @@ func (p Vec) Cross(q Vec) Vec {
 	}
 }
 
+// Rotate returns a new vector, rotated by alpha around the provided axis.
+func (p Vec) Rotate(alpha float64, axis Vec) Vec {
+	return NewRotation(alpha, axis).Rotate(p)
+}
+
 // Norm returns the Euclidean norm of p
 //  |p| = sqrt(p_x^2 + p_y^2 + p_z^2).
 func Norm(p Vec) float64 {
@@ -79,3 +88,50 @@ func Cos(p, q Vec) float64 {
 type Box struct {
 	Min, Max Vec
 }
+
+// TODO: possibly useful additions to the current rotation API:
+//  - create rotations from Euler angles (NewRotationFromEuler?)
+//  - create rotations from rotation matrices (NewRotationFromMatrix?)
+//  - return the equivalent Euler angles from a Rotation
+//  - return the equivalent rotation matrix from a Rotation
+//
+// Euler angles have issues (see [1] for a discussion).
+// We should think carefully before adding them in.
+// [1]: http://www.euclideanspace.com/maths/geometry/rotations/conversions/quaternionToEuler/
+
+// Rotation describes a rotation in space.
+type Rotation quat.Number
+
+// NewRotation creates a rotation by alpha, around axis.
+func NewRotation(alpha float64, axis Vec) Rotation {
+	if alpha == 0 {
+		return Rotation{Real: 1}
+	}
+	q := raise(axis)
+	sin, cos := math.Sincos(0.5 * alpha)
+	q = quat.Scale(sin/quat.Abs(q), q)
+	q.Real += cos
+	if len := quat.Abs(q); len != 1 {
+		q = quat.Scale(1/len, q)
+	}
+
+	return Rotation(q)
+}
+
+// Rotate returns the rotated vector according to the definition of rot.
+func (r Rotation) Rotate(p Vec) Vec {
+	if r.isIdentity() {
+		return p
+	}
+	qq := quat.Number(r)
+	pp := quat.Mul(quat.Mul(qq, raise(p)), quat.Conj(qq))
+	return Vec{X: pp.Imag, Y: pp.Jmag, Z: pp.Kmag}
+}
+
+func (r Rotation) isIdentity() bool {
+	return r == Rotation{Real: 1}
+}
+
+func raise(p Vec) quat.Number {
+	return quat.Number{Imag: p.X, Jmag: p.Y, Kmag: p.Z}
+}
diff --git a/spatial/r3/vector_test.go b/spatial/r3/vector_test.go
@@ -236,6 +236,32 @@ func TestCos(t *testing.T) {
 	}
 }
 
+func TestRotate(t *testing.T) {
+	const tol = 1e-14
+	for _, test := range []struct {
+		v, axis Vec
+		alpha   float64
+		want    Vec
+	}{
+		{Vec{1, 0, 0}, Vec{1, 0, 0}, math.Pi / 2, Vec{1, 0, 0}},
+		{Vec{1, 0, 0}, Vec{1, 0, 0}, 0, Vec{1, 0, 0}},
+		{Vec{1, 0, 0}, Vec{1, 0, 0}, 2 * math.Pi, Vec{1, 0, 0}},
+		{Vec{1, 0, 0}, Vec{0, 0, 0}, math.Pi / 2, Vec{math.NaN(), math.NaN(), math.NaN()}},
+		{Vec{1, 0, 0}, Vec{0, 1, 0}, math.Pi / 2, Vec{0, 0, -1}},
+		{Vec{1, 0, 0}, Vec{0, 1, 0}, math.Pi, Vec{-1, 0, 0}},
+		{Vec{2, 0, 0}, Vec{0, 1, 0}, math.Pi, Vec{-2, 0, 0}},
+		{Vec{1, 2, 3}, Vec{1, 1, 1}, 2. / 3. * math.Pi, Vec{3, 1, 2}},
+	} {
+		got := test.v.Rotate(test.alpha, test.axis)
+		if !vecApproxEqual(got, test.want, tol) {
+			t.Errorf(
+				"rotate(%v, %v, %v)= %v, want=%v",
+				test.v, test.alpha, test.axis, got, test.want,
+			)
+		}
+	}
+}
+
 func vecIsNaN(v Vec) bool {
 	return math.IsNaN(v.X) && math.IsNaN(v.Y) && math.IsNaN(v.Z)
 }
@@ -250,3 +276,12 @@ func vecEqual(a, b Vec) bool {
 	}
 	return a == b
 }
+
+func vecApproxEqual(a, b Vec, tol float64) bool {
+	if vecIsNaNAny(a) || vecIsNaNAny(b) {
+		return vecIsNaN(a) && vecIsNaN(b)
+	}
+	return scalar.EqualWithinAbs(a.X, b.X, tol) &&
+		scalar.EqualWithinAbs(a.Y, b.Y, tol) &&
+		scalar.EqualWithinAbs(a.Z, b.Z, tol)
+}