diff --git a/src/main/common/maths.h b/src/main/common/maths.h
index 513a36c182..881381b274 100644
--- a/src/main/common/maths.h
+++ b/src/main/common/maths.h
@@ -40,7 +40,9 @@
 #define DEGREES_TO_DECIDEGREES(angle) ((angle) * 10)
 #define DECIDEGREES_TO_DEGREES(angle) ((angle) / 10)
 #define DECIDEGREES_TO_RADIANS(angle) ((angle) / 10.0f * 0.0174532925f)
-#define DEGREES_TO_RADIANS(angle) ((angle) * 0.0174532925f)
+#define DEGREES_TO_RADIANS(angle) ((angle) * RAD)
+#define RADIANS_TO_DEGREES(angle) ((angle) / RAD)
+
 
 #define CM_S_TO_KM_H(centimetersPerSecond) ((centimetersPerSecond) * 36 / 1000)
 #define CM_S_TO_MPH(centimetersPerSecond) ((centimetersPerSecond) * 10000 / 5080 / 88)
diff --git a/src/main/common/vector.h b/src/main/common/vector.h
new file mode 100644
index 0000000000..4a0d521983
--- /dev/null
+++ b/src/main/common/vector.h
@@ -0,0 +1,149 @@
+/*
+ * This file is part of Cleanflight and Betaflight.
+ *
+ * Cleanflight and Betaflight are free software. You can redistribute
+ * this software and/or modify this software under the terms of the
+ * GNU General Public License as published by the Free Software
+ * Foundation, either version 3 of the License, or (at your option)
+ * any later version.
+ *
+ * Cleanflight and Betaflight are distributed in the hope that they
+ * will be useful, but WITHOUT ANY WARRANTY; without even the implied
+ * warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
+ * See the GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with this software.
+ *
+ * If not, see <http://www.gnu.org/licenses/>.
+ */
+
+/*
+ * some functions are taken from https://github.com/iNavFlight/inav/
+ */
+
+#pragma once
+
+#include "common/maths.h"
+
+typedef union {
+    float v[2];
+    struct {
+       float x,y;
+    };
+} fpVector2_t;
+
+
+typedef union {
+    float v[3];
+    struct {
+       float x, y, z;
+    };
+} fpVector3_t;
+
+typedef struct {
+    float m[3][3];
+} fpMat33_t;
+
+static inline fpVector3_t * vectorZero(fpVector3_t *v)
+{
+    v->x = 0.0f;
+    v->y = 0.0f;
+    v->z = 0.0f;
+    return v;
+}
+
+static inline float vectorNormSquared(const fpVector3_t * v)
+{
+    return sq(v->x) + sq(v->y) + sq(v->z);
+}
+
+static inline float vectorNorm(const fpVector3_t * v)
+{
+    return sqrtf(vectorNormSquared(v));
+}
+
+static inline fpVector3_t * vectorCrossProduct(fpVector3_t *result, const fpVector3_t *a, const fpVector3_t *b)
+{
+    fpVector3_t ab;
+
+    ab.x = a->y * b->z - a->z * b->y;
+    ab.y = a->z * b->x - a->x * b->z;
+    ab.z = a->x * b->y - a->y * b->x;
+
+    *result = ab;
+    return result;
+}
+
+static inline fpVector3_t * vectorAdd(fpVector3_t *result, const fpVector3_t *a, const fpVector3_t *b)
+{
+    fpVector3_t ab;
+
+    ab.x = a->x + b->x;
+    ab.y = a->y + b->y;
+    ab.z = a->z + b->z;
+
+    *result = ab;
+    return result;
+}
+
+static inline fpVector3_t * vectorScale(fpVector3_t *result, const fpVector3_t *a, const float b)
+{
+    fpVector3_t ab;
+
+    ab.x = a->x * b;
+    ab.y = a->y * b;
+    ab.z = a->z * b;
+
+    *result = ab;
+    return result;
+}
+
+static inline fpVector3_t * vectorNormalize(fpVector3_t *result, const fpVector3_t *v)
+{
+    float normSq = vectorNormSquared(v);
+    if (normSq > 0) {              // Hopefully sqrt(nonzero) is quite large
+        return vectorScale(result, v, 1.0f / sqrtf(normSq));
+    } else {
+        return vectorZero(result);
+    }
+}
+
+static inline fpVector3_t * matrixVectorMul(fpVector3_t * result, const fpMat33_t * mat, const fpVector3_t * a)
+{
+    fpVector3_t r;
+
+    r.x = mat->m[0][0] * a->x + mat->m[0][1] * a->y + mat->m[0][2] * a->z;
+    r.y = mat->m[1][0] * a->x + mat->m[1][1] * a->y + mat->m[1][2] * a->z;
+    r.z = mat->m[2][0] * a->x + mat->m[2][1] * a->y + mat->m[2][2] * a->z;
+
+    *result = r;
+    return result;
+}
+
+static inline fpVector3_t * matrixTrnVectorMul(fpVector3_t * result, const fpMat33_t * mat, const fpVector3_t * a)
+{
+    fpVector3_t r;
+
+    r.x = mat->m[0][0] * a->x + mat->m[1][0] * a->y + mat->m[2][0] * a->z;
+    r.y = mat->m[0][1] * a->x + mat->m[1][1] * a->y + mat->m[2][1] * a->z;
+    r.z = mat->m[0][2] * a->x + mat->m[1][2] * a->y + mat->m[2][2] * a->z;
+
+    *result = r;
+    return result;
+}
+
+static inline float vector2Cross(const fpVector2_t *a, const fpVector2_t *b)
+{
+    return a->x * b->y - a->y * b->x;
+}
+
+static inline float vector2Dot(const fpVector2_t *a, const fpVector2_t *b)
+{
+    return a->x * b->x + a->y * b->y;
+}
+
+static inline float vector2Mag(const fpVector2_t *a)
+{
+    return sqrtf(sq(a->x) + sq(a->y));
+}
diff --git a/src/main/flight/imu.c b/src/main/flight/imu.c
index a308cfd334..874bf98095 100644
--- a/src/main/flight/imu.c
+++ b/src/main/flight/imu.c
@@ -31,6 +31,7 @@
 #include "build/debug.h"
 
 #include "common/axis.h"
+#include "common/vector.h"
 
 #include "pg/pg.h"
 #include "pg/pg_ids.h"
@@ -198,7 +199,7 @@ static float invSqrt(float x)
     return 1.0f / sqrtf(x);
 }
 
-static void imuMahonyAHRSupdate(float dt, float gx, float gy, float gz,
+STATIC_UNIT_TESTED void imuMahonyAHRSupdate(float dt, float gx, float gy, float gz,
                                 bool useAcc, float ax, float ay, float az,
                                 bool useMag,
                                 float cogYawGain, float courseOverGround, const float dcmKpGain)
@@ -212,43 +213,59 @@ static void imuMahonyAHRSupdate(float dt, float gx, float gy, float gz,
     float ex = 0, ey = 0, ez = 0;
     if (cogYawGain != 0.0f) {
         // Used in a GPS Rescue to boost IMU yaw gain when course over ground and velocity to home differ significantly
-        while (courseOverGround >  M_PIf) {
-            courseOverGround -= (2.0f * M_PIf);
-        }
-        while (courseOverGround < -M_PIf) {
-            courseOverGround += (2.0f * M_PIf);
-        }
-        const float ez_ef = cogYawGain * (- sin_approx(courseOverGround) * rMat[0][0] - cos_approx(courseOverGround) * rMat[1][0]);
-        ex = rMat[2][0] * ez_ef;
-        ey = rMat[2][1] * ez_ef;
-        ez = rMat[2][2] * ez_ef;
+
+        // Compute heading vector in EF from scalar CoG. CoG is clockwise from North
+        // Note that Earth frame X is pointing north and sin/cos argument is anticlockwise
+        const fpVector2_t cog_ef = {.x = cos_approx(-courseOverGround), .y = sin_approx(-courseOverGround)};
+#define THRUST_COG 1
+#if THRUST_COG
+        const fpVector2_t heading_ef = {.x = rMat[X][Z], .y = rMat[Y][Z]};  // body Z axis (up) - direction of thrust vector
+#else
+        const fpVector2_t heading_ef = {.x = rMat[0][0], .y = rMat[1][0]};  // body X axis. Projected vector magnitude is reduced as pitch increases
+#endif
+        // cross product = 1 * |heading| * sin(angle) (magnitude of Z vector in 3D)
+        // order operands so that rotation is in direction of zero error
+        const float cross = vector2Cross(&heading_ef, &cog_ef);
+        // dot product, 1 * |heading| * cos(angle)
+        const float dot = vector2Dot(&heading_ef, &cog_ef);
+        // use cross product / sin(angle) when error < 90deg (cos > 0),
+        //   |heading| if error is larger (cos < 0)
+        const float heading_mag = vector2Mag(&heading_ef);
+        float ez_ef = (dot > 0) ? cross : (cross < 0 ? -1.0f : 1.0f) * heading_mag;
+#if THRUST_COG
+        // increase gain for small tilt (just heuristic; sqrt is cheap on F4+)
+        ez_ef /= sqrtf(heading_mag);
+#endif
+        ez_ef *= cogYawGain;          // apply gain parameter
+        // covert to body frame
+        ex += rMat[2][0] * ez_ef;
+        ey += rMat[2][1] * ez_ef;
+        ez += rMat[2][2] * ez_ef;
     }
 
 #ifdef USE_MAG
     // Use measured magnetic field vector
-    float mx = mag.magADC[X];
-    float my = mag.magADC[Y];
-    float mz = mag.magADC[Z];
-    float recipMagNorm = sq(mx) + sq(my) + sq(mz);
+    fpVector3_t mag_bf = {{mag.magADC[X], mag.magADC[Y], mag.magADC[Z]}};
+    float recipMagNorm = vectorNormSquared(&mag_bf);
     if (useMag && recipMagNorm > 0.01f) {
         // Normalise magnetometer measurement
-        recipMagNorm = invSqrt(recipMagNorm);
-        mx *= recipMagNorm;
-        my *= recipMagNorm;
-        mz *= recipMagNorm;
+        vectorNormalize(&mag_bf, &mag_bf);
 
         // For magnetometer correction we make an assumption that magnetic field is perpendicular to gravity (ignore Z-component in EF).
         // This way magnetic field will only affect heading and wont mess roll/pitch angles
 
-        // (hx; hy; 0) - measured mag field vector in EF (assuming Z-component is zero)
+        // (hx; hy; 0) - measured mag field vector in EF (forcing Z-component to zero)
         // (bx; 0; 0) - reference mag field vector heading due North in EF (assuming Z-component is zero)
-        const float hx = rMat[0][0] * mx + rMat[0][1] * my + rMat[0][2] * mz;
-        const float hy = rMat[1][0] * mx + rMat[1][1] * my + rMat[1][2] * mz;
-        const float bx = sqrtf(hx * hx + hy * hy);
+        fpVector3_t mag_ef;
+        matrixVectorMul(&mag_ef, (const fpMat33_t*)&rMat, &mag_bf);  // BF->EF
+        mag_ef.z = 0.0f;                // project to XY plane (optimized away)
 
+        fpVector2_t north_ef = {{ 1.0f, 0.0f }};
         // magnetometer error is cross product between estimated magnetic north and measured magnetic north (calculated in EF)
-        const float ez_ef = -(hy * bx);
-
+        // increase gain on large misalignment
+        const float dot = vector2Dot((fpVector2_t*)&mag_ef, &north_ef);
+        const float cross = vector2Cross((fpVector2_t*)&mag_ef, &north_ef);
+        const float ez_ef = (dot > 0) ? cross : (cross < 0 ? -1.0f : 1.0f) * vector2Mag((fpVector2_t*)&mag_ef);
         // Rotate mag error vector back to BF and accumulate
         ex += rMat[2][0] * ez_ef;
         ey += rMat[2][1] * ez_ef;
diff --git a/src/test/unit/flight_imu_unittest.cc b/src/test/unit/flight_imu_unittest.cc
index 2ec9304f47..90ed5cb6d2 100644
--- a/src/test/unit/flight_imu_unittest.cc
+++ b/src/test/unit/flight_imu_unittest.cc
@@ -26,6 +26,7 @@ extern "C" {
 
     #include "common/axis.h"
     #include "common/maths.h"
+    #include "common/vector.h"
 
     #include "config/feature.h"
     #include "pg/pg.h"
@@ -57,7 +58,10 @@ extern "C" {
 
     void imuComputeRotationMatrix(void);
     void imuUpdateEulerAngles(void);
-
+    void imuMahonyAHRSupdate(float dt, float gx, float gy, float gz,
+                             bool useAcc, float ax, float ay, float az,
+                             bool useMag,
+                             float cogYawGain, float courseOverGround, const float dcmKpGain);
     extern quaternion q;
     extern float rMat[3][3];
     extern bool attitudeIsEstablished;
@@ -76,6 +80,15 @@ extern "C" {
 
 const float sqrt2over2 = sqrtf(2) / 2.0f;
 
+void quaternion_from_axis_angle(quaternion* q, float angle, float x, float y, float z) {
+    fpVector3_t a = {{x, y, z}};
+    vectorNormalize(&a, &a);
+    q->w = cos(angle / 2);
+    q->x = a.x * sin(angle / 2);
+    q->y = a.y * sin(angle / 2);
+    q->z = a.z * sin(angle / 2);
+}
+
 TEST(FlightImuTest, TestCalculateRotationMatrix)
 {
     #define TOL 1e-6
@@ -201,6 +214,161 @@ TEST(FlightImuTest, TestSmallAngle)
     EXPECT_FALSE(isUpright());
 }
 
+testing::AssertionResult DoubleNearWrapPredFormat(const char* expr1, const char* expr2,
+                                                  const char* abs_error_expr, const char* wrap_expr, double val1,
+                                                  double val2, double abs_error, double wrap) {
+    const double diff = remainder(val1 - val2, wrap);
+    if (fabs(diff) <= abs_error) return testing::AssertionSuccess();
+
+    return testing::AssertionFailure()
+        << "The difference between " << expr1 << " and " << expr2 << " is "
+        << diff << " (wrapped to 0 .. " << wrap_expr << ")"
+        << ", which exceeds " << abs_error_expr << ", where\n"
+        << expr1 << " evaluates to " << val1 << ",\n"
+        << expr2 << " evaluates to " << val2 << ", and\n"
+        << abs_error_expr << " evaluates to " << abs_error << ".";
+}
+
+#define EXPECT_NEAR_DEG(val1, val2, abs_error)                \
+    EXPECT_PRED_FORMAT4(DoubleNearWrapPredFormat, val1, val2, \
+                        abs_error, 360.0)
+
+#define EXPECT_NEAR_RAD(val1, val2, abs_error)                \
+    EXPECT_PRED_FORMAT4(DoubleNearWrapPredFormat, val1, val2, \
+                        abs_error, 2 * M_PI)
+
+
+
+class MahonyFixture : public ::testing::Test {
+protected:
+    fpVector3_t gyro;
+    bool useAcc;
+    fpVector3_t acc;
+    bool useMag;
+    fpVector3_t magEF;
+    float cogGain;
+    float cogDeg;
+    float dcmKp;
+    float dt;
+    void SetUp() override {
+        vectorZero(&gyro);
+        useAcc = false;
+        vectorZero(&acc);
+        cogGain = 0.0;   // no cog
+        cogDeg  = 0.0;
+        dcmKp = .25;     // default dcm_kp
+        dt = 1e-2;       // 100Hz update
+
+        imuConfigure(0, 0);
+        // level, poiting north
+        setOrientationAA(0, {{1,0,0}});        // identity
+    }
+    virtual void setOrientationAA(float angleDeg, fpVector3_t axis) {
+        quaternion_from_axis_angle(&q, DEGREES_TO_RADIANS(angleDeg), axis.x, axis.y, axis.z);
+        imuComputeRotationMatrix();
+    }
+
+    float wrap(float angle) {
+        angle = fmod(angle, 360);
+        if (angle < 0) angle += 360;
+        return angle;
+    }
+    float angleDiffNorm(fpVector3_t *a, fpVector3_t* b, fpVector3_t weight = {{1,1,1}}) {
+        fpVector3_t tmp;
+        vectorScale(&tmp, b, -1);
+        vectorAdd(&tmp, &tmp, a);
+        for (int i = 0; i < 3; i++)
+            tmp.v[i] *= weight.v[i];
+        for (int i = 0; i < 3; i++)
+            tmp.v[i] = std::remainder(tmp.v[i], 360.0);
+        return vectorNorm(&tmp);
+    }
+    // run Mahony for some time
+    // return time it took to get within 1deg from target
+    float imuIntegrate(float runTime, fpVector3_t * target) {
+        float alignTime = -1;
+        for (float t = 0; t < runTime; t += dt) {
+            //     if (fmod(t, 1) < dt) printf("MagBF=%.2f %.2f %.2f\n", magBF.x, magBF.y, magBF.z);
+            imuMahonyAHRSupdate(dt,
+                                gyro.x, gyro.y, gyro.z,
+                                useAcc, acc.x, acc.y, acc.z,
+                                useMag,                             // no mag now
+                                cogGain, DEGREES_TO_RADIANS(cogDeg),   // use Cog, param direction
+                                dcmKp);
+            imuUpdateEulerAngles();
+            // if (fmod(t, 1) < dt) printf("%3.1fs - %3.1f %3.1f %3.1f\n", t, attitude.values.roll / 10.0f, attitude.values.pitch / 10.0f, attitude.values.yaw / 10.0f);
+            // remember how long it took
+            if (alignTime < 0) {
+                fpVector3_t rpy = {{attitude.values.roll / 10.0f, attitude.values.pitch / 10.0f, attitude.values.yaw / 10.0f}};
+                float error = angleDiffNorm(&rpy, target);
+                if (error < 1)
+                    alignTime = t;
+            }
+        }
+        return alignTime;
+    }
+};
+
+class YawTest: public MahonyFixture, public testing::WithParamInterface<float> {
+};
+
+TEST_P(YawTest, TestCogAlign)
+{
+    cogGain = 1.0;
+    cogDeg = GetParam();
+    const float rollDeg = 30;    // 30deg pitch forward
+    setOrientationAA(rollDeg, {{0, 1, 0}});
+    fpVector3_t expect = {{0, rollDeg, wrap(cogDeg)}};
+    // integrate IMU. about 25s is enough in worst case
+    float alignTime = imuIntegrate(80, &expect);
+
+    imuUpdateEulerAngles();
+    // quad stays level
+    EXPECT_NEAR_DEG(attitude.values.roll / 10.0, expect.x, .1);
+    EXPECT_NEAR_DEG(attitude.values.pitch / 10.0, expect.y, .1);
+    // yaw is close to CoG direction
+    EXPECT_NEAR_DEG(attitude.values.yaw / 10.0, expect.z, 1);  // error < 1 deg
+    if (alignTime >= 0) {
+        printf("[          ] Aligned to %.f deg in %.2fs\n", cogDeg, alignTime);
+    }
+}
+
+TEST_P(YawTest, TestMagAlign)
+{
+    float initialAngle = GetParam();
+
+    // level, rotate to param heading
+    quaternion_from_axis_angle(&q, -DEGREES_TO_RADIANS(initialAngle), 0, 0, 1);
+    imuComputeRotationMatrix();
+
+    fpVector3_t expect = {{0, 0, 0}};    // expect zero yaw
+
+    fpVector3_t magBF = {{1, 0, .5}};    // use arbitrary Z component, point north
+
+    for (int i = 0; i < 3; i++)
+        mag.magADC[i] = magBF.v[i];
+
+    useMag = true;
+    // integrate IMU. about 25s is enough in worst case
+    float alignTime = imuIntegrate(30, &expect);
+
+    imuUpdateEulerAngles();
+    // quad stays level
+    EXPECT_NEAR_DEG(attitude.values.roll / 10.0, expect.x, .1);
+    EXPECT_NEAR_DEG(attitude.values.pitch / 10.0, expect.y, .1);
+    // yaw is close to north (0 deg)
+    EXPECT_NEAR_DEG(attitude.values.yaw / 10.0, expect.z, 1.0);  // error < 1 deg
+    if (alignTime >= 0) {
+        printf("[          ] Aligned from %.f deg in %.2fs\n", initialAngle, alignTime);
+    }
+}
+
+INSTANTIATE_TEST_SUITE_P(
+  TestAngles, YawTest,
+  ::testing::Values(
+      0, 45, -45, 90, 180, 270, 720+45
+      ));
+
 // STUBS
 
 extern "C" {