From 00908b624de3f0ccf39592ca8e270e829f475935 Mon Sep 17 00:00:00 2001
From: Peter Sharpe <pds@mit.edu>
Date: Thu, 4 Jan 2024 01:04:18 +0100
Subject: [PATCH] fixes to modulo function, which had discrepancies with np.mod
 for negative numbers. Added centered_mod, which is a version of mod that is
 centered on zero (i.e., a half-phase-shifted mod)

---
 aerosandbox/numpy/arithmetic_dyadic.py | 27 +++++++++++-
 aerosandbox/numpy/calculus.py          | 57 +++++++++++++++-----------
 2 files changed, 59 insertions(+), 25 deletions(-)

diff --git a/aerosandbox/numpy/arithmetic_dyadic.py b/aerosandbox/numpy/arithmetic_dyadic.py
index 1e2b9523c..2fad0ece2 100644
--- a/aerosandbox/numpy/arithmetic_dyadic.py
+++ b/aerosandbox/numpy/arithmetic_dyadic.py
@@ -1,6 +1,7 @@
 import numpy as _onp
 import casadi as _cas
 from typing import Tuple, Iterable, Union
+from aerosandbox.numpy.conditionals import where
 
 from aerosandbox.numpy.determine_type import is_casadi_type
 
@@ -65,4 +66,28 @@ def mod(x1, x2):
         return _onp.mod(x1, x2)
 
     else:
-        return _cas.fmod(x1, x2)
+        out = _cas.fmod(x1, x2)
+        out = where(
+            x1 < 0,
+            out + x2,
+            out
+        )
+        return out
+
+
+def centered_mod(x1, x2):
+    """
+    Return element-wise remainder of division, centered on zero.
+
+    See syntax here: https://numpy.org/doc/stable/reference/generated/numpy.mod.html
+    """
+    if not is_casadi_type(x1) and not is_casadi_type(x2):
+        remainder = _onp.mod(x1, x2)
+        return where(
+            remainder > x2 / 2,
+            remainder - x2,
+            remainder
+        )
+
+    else:
+        return _cas.remainder(x1, x2)
diff --git a/aerosandbox/numpy/calculus.py b/aerosandbox/numpy/calculus.py
index 50f4f13c7..7f3ad0867 100644
--- a/aerosandbox/numpy/calculus.py
+++ b/aerosandbox/numpy/calculus.py
@@ -1,8 +1,7 @@
 import numpy as _onp
 import casadi as _cas
 from aerosandbox.numpy.determine_type import is_casadi_type
-from aerosandbox.numpy.conditionals import where as _where
-from aerosandbox.numpy.arithmetic_dyadic import mod as _mod
+from aerosandbox.numpy.arithmetic_dyadic import centered_mod as _centered_mod
 from aerosandbox.numpy.array import array, concatenate, reshape
 
 
@@ -13,12 +12,9 @@ def diff(a, n=1, axis=-1, period=None):
     See syntax here: https://numpy.org/doc/stable/reference/generated/numpy.diff.html
     """
     if period is not None:
-        original_result = diff(a, n=n, axis=axis)
-        remainder = _mod(original_result, period)
-        return _where(
-            remainder > period / 2,
-            remainder - period,
-            remainder
+        return _centered_mod(
+            diff(a, n=n, axis=axis),
+            period
         )
 
     if not is_casadi_type(a):
@@ -239,22 +235,35 @@ def trapz(x, modify_endpoints=False):  # TODO unify with NumPy trapz, this is di
 
 
 if __name__ == '__main__':
-    import numpy as np
+    import aerosandbox as asb
+    import aerosandbox.numpy as np
 
-    # a = np.linspace(-500, 500, 21) % 360 - 180
-    # print(diff(a, period=360))
-
-    x = np.cumsum(np.arange(10))
-    y = x ** 2
-
-    print(gradient(y, x, edge_order=1))
-    print(gradient(y, x, edge_order=1))
-    print(gradient(y, x, edge_order=1, n=2))
-
-    from aerosandbox.tools.code_benchmarking import Timer
     import casadi as cas
 
-    # with Timer("np"):
-    #     print(gradient(np.eye(50), 2, 1))
-    # with Timer("custom"):
-    #     print(gradient(cas.MX_eye(50), 2, 1))
+    print(diff(cas.DM([355, 5]), period=360))
+    #
+    # # a = np.linspace(-500, 500, 21) % 360 - 180
+    # # print(diff(a, period=360))
+    #
+    # x = np.cumsum(np.arange(10))
+    # y = x ** 2
+    #
+    # print(gradient(y, x, edge_order=1))
+    # print(gradient(y, x, edge_order=1))
+    # print(gradient(y, x, edge_order=1, n=2))
+    #
+    # opti = asb.Opti()
+    # x = opti.variable(init_guess=[355, 5])
+    # d = diff(x, period=360)
+    # opti.subject_to([
+    #     # x[0] == 3,
+    #     x[0] > 180,
+    #     x[1] < 180,
+    #     d < 20,
+    #     d > -20
+    # ])
+    # opti.maximize(np.sum(np.cosd(x)))
+    # sol = opti.solve(
+    #     behavior_on_failure="return_last"
+    # )
+    # print(sol.value(x))
\ No newline at end of file