Add rotation matrix args

ezzatron · Apr 26, 2024 · bd8ff8c · bd8ff8c
1 parent 3e834e5
commit bd8ff8c
Show file tree

Hide file tree

Showing 7 changed files with 158 additions and 49 deletions.
diff --git a/src/lat_lon2n_E.ts b/src/lat_lon2n_E.ts
@@ -1,10 +1,27 @@
+import type { Matrix3x3 } from "./matrix.js";
+import { ROTATION_MATRIX_e } from "./rotation.js";
 import type { Vector3 } from "./vector.js";
 
-export function lat_lon2n_E(latitude: number, longitude: number): Vector3 {
+export function lat_lon2n_E(
+  latitude: number,
+  longitude: number,
+  R_Ee: Matrix3x3 = ROTATION_MATRIX_e,
+): Vector3 {
+  const [[n00, n01, n02], [n10, n11, n12], [n20, n21, n22]] = R_Ee;
+
   const sinLat = Math.sin(latitude);
   const cosLat = Math.cos(latitude);
   const sinLon = Math.sin(longitude);
   const cosLon = Math.cos(longitude);
 
-  return [cosLat * cosLon, cosLat * sinLon, sinLat];
+  const x = sinLat;
+  const y = cosLat * sinLon;
+  const z = -cosLat * cosLon;
+
+  // flattened multiply(transpose(R_Ee), [x, y, z])
+  return [
+    n00 * x + n10 * y + n20 * z,
+    n01 * x + n11 * y + n21 * z,
+    n02 * x + n12 * y + n22 * z,
+  ];
 }
diff --git a/src/n_E2lat_lon.ts b/src/n_E2lat_lon.ts
@@ -1,14 +1,23 @@
+import type { Matrix3x3 } from "./matrix.js";
+import { ROTATION_MATRIX_e } from "./rotation.js";
 import type { Vector3 } from "./vector.js";
 
 export function n_E2lat_lon(
   n_E: Vector3,
+  R_Ee: Matrix3x3 = ROTATION_MATRIX_e,
 ): [latitude: number, longitude: number] {
   const [x, y, z] = n_E;
+  const [[n00, n01, n02], [n10, n11, n12], [n20, n21, n22]] = R_Ee;
 
-  const sinLat = z;
-  const cosLat = Math.sqrt(y ** 2 + x ** 2);
-  const cosLatSinLon = y;
-  const cosLatCosLon = x;
+  // flattened multiply(R_Ee, [x, y, z])
+  const rx = n00 * x + n01 * y + n02 * z;
+  const ry = n10 * x + n11 * y + n12 * z;
+  const rz = n20 * x + n21 * y + n22 * z;
 
-  return [Math.atan2(sinLat, cosLat), Math.atan2(cosLatSinLon, cosLatCosLon)];
+  const sinLat = rx;
+  const cosLat = Math.sqrt(ry ** 2 + rz ** 2);
+  const sinLonCosLat = ry;
+  const cosLonCosLat = -rz;
+
+  return [Math.atan2(sinLat, cosLat), Math.atan2(sinLonCosLat, cosLonCosLat)];
 }
diff --git a/src/rotation.ts b/src/rotation.ts
@@ -0,0 +1,7 @@
+import type { Matrix3x3 } from "./matrix.js";
+
+export const ROTATION_MATRIX_e: Matrix3x3 = [
+  [0, 0, 1.0],
+  [0, 1.0, 0],
+  [-1.0, 0, 0],
+];
diff --git a/test/arbitrary.ts b/test/arbitrary.ts
@@ -1,4 +1,6 @@
 import { fc } from "@fast-check/jest";
+import type { Matrix3x3 } from "../src/matrix.js";
+import type { Vector4 } from "../src/vector.js";
 
 const RADIAN = Math.PI / 180;
 
@@ -21,3 +23,49 @@ export function arbitraryLon(): fc.Arbitrary<number> {
 export function arbitraryLatLon(): fc.Arbitrary<[number, number]> {
   return fc.tuple(arbitraryLat(), arbitraryLon());
 }
+
+export function arbitraryQuaternion(): fc.Arbitrary<Vector4> {
+  // based on https://github.com/mrdoob/three.js/blob/a2e9ee8204b67f9dca79f48cf620a34a05aa8126/src/math/Quaternion.js#L592
+  // Ken Shoemake
+  // Uniform random rotations
+  // D. Kirk, editor, Graphics Gems III, pages 124-132. Academic Press, New York, 1992.
+
+  return fc
+    .tuple(
+      fc.double({ min: 0, max: Math.PI * 2, noNaN: true }),
+      fc.double({ min: 0, max: Math.PI * 2, noNaN: true }),
+      fc.double({ min: 0, max: 1, noNaN: true }),
+    )
+    .map(([theta1, theta2, x0]) => {
+      const r1 = Math.sqrt(1 - x0);
+      const r2 = Math.sqrt(x0);
+
+      const x = r1 * Math.sin(theta1);
+      const y = r1 * Math.cos(theta1);
+      const z = r2 * Math.sin(theta2);
+      const w = r2 * Math.cos(theta2);
+
+      return [x, y, z, w];
+    });
+}
+
+export function arbitrary3dRotationMatrix(): fc.Arbitrary<Matrix3x3> {
+  return arbitraryQuaternion().map(([x, y, z, w]) => {
+    // based on https://github.com/rawify/Quaternion.js/blob/c3834673b502e64e1866dbbf13568c0be93e52cc/quaternion.js#L791
+    const wx = w * x;
+    const wy = w * y;
+    const wz = w * z;
+    const xx = x * x;
+    const xy = x * y;
+    const xz = x * z;
+    const yy = y * y;
+    const yz = y * z;
+    const zz = z * z;
+
+    return [
+      [1 - 2 * (yy + zz), 2 * (xy - wz), 2 * (xz + wy)],
+      [2 * (xy + wz), 1 - 2 * (xx + zz), 2 * (yz - wx)],
+      [2 * (xz - wy), 2 * (yz + wx), 1 - 2 * (xx + yy)],
+    ];
+  });
+}
diff --git a/test/jest/lat_lon2n_E.spec.ts b/test/jest/lat_lon2n_E.spec.ts
@@ -1,6 +1,6 @@
-import { it } from "@fast-check/jest";
+import { fc, it } from "@fast-check/jest";
 import { lat_lon2n_E } from "../../src/index.js";
-import { arbitraryLatLon } from "../arbitrary.js";
+import { arbitrary3dRotationMatrix, arbitraryLatLon } from "../arbitrary.js";
 import {
   NvectorTestClient,
   createNvectorTestClient,
@@ -17,26 +17,39 @@ describe("lat_lon2n_E()", () => {
     nvectorTestClient?.close();
   });
 
-  it.prop([arbitraryLatLon()], {
-    numRuns: Infinity,
-  })("matches the Python implementation", async ([latitude, longitude]) => {
-    const expected = await nvectorTestClient.lat_lon2n_E(latitude, longitude);
+  it.prop(
+    [
+      arbitraryLatLon(),
+      fc.option(arbitrary3dRotationMatrix(), { nil: undefined }),
+    ],
+    {
+      numRuns: Infinity,
+    },
+  )(
+    "matches the Python implementation",
+    async ([latitude, longitude], R_Ee) => {
+      const expected = await nvectorTestClient.lat_lon2n_E(
+        latitude,
+        longitude,
+        R_Ee,
+      );
 
-    expect(expected).toMatchObject([
-      expect.any(Number),
-      expect.any(Number),
-      expect.any(Number),
-    ]);
+      expect(expected).toMatchObject([
+        expect.any(Number),
+        expect.any(Number),
+        expect.any(Number),
+      ]);
 
-    const actual = lat_lon2n_E(latitude, longitude);
+      const actual = lat_lon2n_E(latitude, longitude, R_Ee);
 
-    expect(actual).toMatchObject([
-      expect.any(Number),
-      expect.any(Number),
-      expect.any(Number),
-    ]);
-    expect(actual[0]).toBeCloseTo(expected[0], 10);
-    expect(actual[1]).toBeCloseTo(expected[1], 10);
-    expect(actual[2]).toBeCloseTo(expected[2], 10);
-  });
+      expect(actual).toMatchObject([
+        expect.any(Number),
+        expect.any(Number),
+        expect.any(Number),
+      ]);
+      expect(actual[0]).toBeCloseTo(expected[0], 10);
+      expect(actual[1]).toBeCloseTo(expected[1], 10);
+      expect(actual[2]).toBeCloseTo(expected[2], 10);
+    },
+  );
 });
diff --git a/test/jest/n_E2lat_lon.spec.ts b/test/jest/n_E2lat_lon.spec.ts
@@ -1,6 +1,6 @@
-import { it } from "@fast-check/jest";
+import { fc, it } from "@fast-check/jest";
 import { lat_lon2n_E, n_E2lat_lon } from "../../src/index.js";
-import { arbitraryLatLon } from "../arbitrary.js";
+import { arbitrary3dRotationMatrix, arbitraryLatLon } from "../arbitrary.js";
 import {
   NvectorTestClient,
   createNvectorTestClient,
@@ -17,19 +17,28 @@ describe("n_E2lat_lon()", () => {
     nvectorTestClient?.close();
   });
 
-  it.prop([arbitraryLatLon()], {
-    numRuns: Infinity,
-  })("matches the Python implementation", async ([latitude, longitude]) => {
-    const [x, y, z] = lat_lon2n_E(latitude, longitude);
-
-    const expected = await nvectorTestClient.n_E2lat_lon([x, y, z]);
-
-    expect(expected).toMatchObject([expect.any(Number), expect.any(Number)]);
-
-    const actual = n_E2lat_lon([x, y, z]);
-
-    expect(actual).toMatchObject([expect.any(Number), expect.any(Number)]);
-    expect(actual[0]).toBeCloseTo(expected[0], 10);
-    expect(actual[1]).toBeCloseTo(expected[1], 10);
-  });
+  it.prop(
+    [
+      arbitraryLatLon(),
+      fc.option(arbitrary3dRotationMatrix(), { nil: undefined }),
+    ],
+    {
+      numRuns: Infinity,
+    },
+  )(
+    "matches the Python implementation",
+    async ([latitude, longitude], R_Ee) => {
+      const [x, y, z] = lat_lon2n_E(latitude, longitude);
+
+      const expected = await nvectorTestClient.n_E2lat_lon([x, y, z], R_Ee);
+
+      expect(expected).toMatchObject([expect.any(Number), expect.any(Number)]);
+
+      const actual = n_E2lat_lon([x, y, z], R_Ee);
+
+      expect(actual).toMatchObject([expect.any(Number), expect.any(Number)]);
+      expect(actual[0]).toBeCloseTo(expected[0], 10);
+      expect(actual[1]).toBeCloseTo(expected[1], 10);
+    },
+  );
 });
diff --git a/test/nvector-test-api.ts b/test/nvector-test-api.ts
@@ -1,13 +1,18 @@
 import { WebSocket } from "ws";
+import type { Matrix3x3 } from "../src/matrix.js";
 import type { Vector3 } from "../src/vector.js";
 
 export type NvectorTestClient = {
   lat_lon2n_E: (
     latitude: number,
     longitude: number,
+    R_Ee?: Matrix3x3,
   ) => Promise<[x: number, y: number, z: number]>;
 
-  n_E2lat_lon: (n_E: Vector3) => Promise<[latitude: number, longitude: number]>;
+  n_E2lat_lon: (
+    n_E: Vector3,
+    R_Ee?: Matrix3x3,
+  ) => Promise<[latitude: number, longitude: number]>;
 
   close: () => void;
 };
@@ -22,21 +27,22 @@ export async function createNvectorTestClient(): Promise<NvectorTestClient> {
   });
 
   return {
-    async lat_lon2n_E(latitude, longitude) {
+    async lat_lon2n_E(latitude, longitude, R_Ee) {
       const [[x], [y], [z]] = await call<[[number], [number], [number]]>(
         "lat_lon2n_E",
-        { latitude, longitude },
+        { latitude, longitude, R_Ee },
       );
 
       return [x, y, z];
     },
 
-    async n_E2lat_lon([x, y, z]) {
+    async n_E2lat_lon([x, y, z], R_Ee) {
       const { latitude, longitude } = await call<{
         latitude: number;
         longitude: number;
       }>("n_E2lat_lon", {
         n_E: [[x], [y], [z]],
+        R_Ee,
       });
 
       return [latitude, longitude];