Fix singularities in atan2 at 0, 0 in swerve physics

The derivative of atan2 is
   -y / (x^2 + y^2), x / (x^2 + y^2)

This explodes at 0, 0.  Be much more gentle at 0, 0 in our calculation
to reduce the rate of change.  This makes the solver converge
*significantly* faster at 0, 0.  We still can't go sideways yet from 0,
0.

Change-Id: If5496714ec1fd36fac76913d208e803257dcd02f
Signed-off-by: Austin Schuh <austin.linux@gmail.com>
diff --git a/frc971/control_loops/swerve/BUILD b/frc971/control_loops/swerve/BUILD
index fcd60a1..670bcc4 100644
--- a/frc971/control_loops/swerve/BUILD
+++ b/frc971/control_loops/swerve/BUILD
@@ -208,6 +208,7 @@
         ":physics_test_utils",
         "@pip//casadi",
         "@pip//numpy",
+        "@pip//scipy",
     ],
 )
 
diff --git a/frc971/control_loops/swerve/generate_physics.cc b/frc971/control_loops/swerve/generate_physics.cc
index 2e3fe73..0e8b013 100644
--- a/frc971/control_loops/swerve/generate_physics.cc
+++ b/frc971/control_loops/swerve/generate_physics.cc
@@ -412,7 +412,7 @@
     result_py->emplace_back("    sin = casadi.sin");
     result_py->emplace_back("    sign = casadi.sign");
     result_py->emplace_back("    cos = casadi.cos");
-    result_py->emplace_back("    atan2 = casadi.atan2");
+    result_py->emplace_back("    atan2 = soft_atan2");
     result_py->emplace_back("    fmax = casadi.fmax");
     result_py->emplace_back("    fabs = casadi.fabs");
 
@@ -456,7 +456,7 @@
     result_py->emplace_back("    sin = casadi.sin");
     result_py->emplace_back("    sign = casadi.sign");
     result_py->emplace_back("    cos = casadi.cos");
-    result_py->emplace_back("    atan2 = casadi.atan2");
+    result_py->emplace_back("    atan2 = soft_atan2");
     result_py->emplace_back("    fmax = casadi.fmax");
     result_py->emplace_back("    fabs = casadi.fabs");
 
@@ -574,6 +574,14 @@
     // result_py.emplace_back("        [X[STATE_MOMENT, 0]],");
     result_py.emplace_back("    ])");
     result_py.emplace_back("");
+    constexpr double kLogGain = 1.0 / 0.05;
+    result_py.emplace_back("def soft_atan2(y, x):");
+    result_py.emplace_back("    return casadi.arctan2(");
+    result_py.emplace_back("        y,");
+    result_py.emplace_back("        casadi.logsumexp(casadi.SX(numpy.array(");
+    result_py.emplace_back(absl::Substitute(
+        "            [1.0, casadi.fabs(x) * $0.0]))) / $0.0)", kLogGain));
+    result_py.emplace_back("");
 
     result_py.emplace_back("# Returns the derivative of our state vector");
     result_py.emplace_back("# [thetas0, thetad0, omegas0, omegad0,");
@@ -839,15 +847,17 @@
     // Velocity of the contact patch in field coordinates
     DenseMatrix temp_matrix = DenseMatrix(2, 1);
     DenseMatrix temp_matrix2 = DenseMatrix(2, 1);
+    DenseMatrix temp_matrix3 = DenseMatrix(2, 1);
     result.contact_patch_velocity = DenseMatrix(2, 1);
 
     mul_dense_dense(R(theta_), result.mounting_location, temp_matrix);
     add_dense_dense(angle_cross(temp_matrix, omega_), robot_velocity,
                     temp_matrix2);
     mul_dense_dense(R(add(theta_, result.thetas)),
-                    DenseMatrix(2, 1, {neg(caster_), integer(0)}), temp_matrix);
+                    DenseMatrix(2, 1, {neg(caster_), integer(0)}),
+                    temp_matrix3);
     add_dense_dense(temp_matrix2,
-                    angle_cross(temp_matrix, add(omega_, result.omegas)),
+                    angle_cross(temp_matrix3, add(omega_, result.omegas)),
                     result.contact_patch_velocity);
 
     VLOG(1);
diff --git a/frc971/control_loops/swerve/physics_test.py b/frc971/control_loops/swerve/physics_test.py
index 2ce82ef..9aa6457 100644
--- a/frc971/control_loops/swerve/physics_test.py
+++ b/frc971/control_loops/swerve/physics_test.py
@@ -3,6 +3,7 @@
 import numpy
 import sys, os
 import casadi
+import scipy
 from numpy.testing import assert_array_equal, assert_array_almost_equal
 import unittest
 
@@ -173,17 +174,23 @@
                 for theta in [0.0, 0.6, -0.4]:
                     module_angle = numpy.pi * wrap + theta
 
-                    # We have redefined the angle to be the sin of the angle.
+                    # We have redefined the angle to be the softened sin of the angle.
                     # That way, when the module flips directions, the slip angle also flips
                     # directions to keep it stable.
                     computed_angle = self.slip_angle[i](
                         utils.state_vector(velocity=velocity,
-                                           module_angle=numpy.pi * wrap +
-                                           theta),
+                                           module_angle=module_angle),
                         self.I,
                     )[0, 0]
 
-                    expected = numpy.sin(numpy.pi * wrap + theta)
+                    # Compute out the expected value directly to confirm we got it right.
+                    loggain = 20.0
+                    vy = 1.5 * numpy.sin(-module_angle)
+                    vx = 1.5 * numpy.cos(-module_angle)
+                    expected = numpy.sin(-numpy.arctan2(
+                        vy,
+                        scipy.special.logsumexp([1.0, abs(vx) * loggain]) /
+                        loggain))
 
                     self.assertAlmostEqual(
                         expected,