Repository: dynamicslab/pysindy

Issue:
[BUG] AxesArray `ravel()` needed for some plotting functions. 
Currently, `AxesArray` tries to implement the most commonly used numpy functions that preserve certain axis information.  For others, the numpy function executes and then calls `AxesArray.__array_finalize__(self, obj)`, which tries to guess how to assign axes.  This causes an error in `ravel()`, called by some plotting functions.

### Reproducing code example:


```python
import pysindy
import pysindy as ps
import numpy as np
from matplotlib import pyplot as plt

arr = ps.AxesArray(np.arange(5), {"ax_1": 0})
plt.scatter(arr, arr)
```

### Error message:
```
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/home/jake/github/pysindy/env/lib/python3.10/site-packages/matplotlib/pyplot.py", line 3684, in scatter
    __ret = gca().scatter(
  File "/home/jake/github/pysindy/env/lib/python3.10/site-packages/matplotlib/__init__.py", line 1465, in inner
    return func(ax, *map(sanitize_sequence, args), **kwargs)
  File "/home/jake/github/pysindy/env/lib/python3.10/site-packages/matplotlib/axes/_axes.py", line 4649, in scatter
    x = np.ma.ravel(x)
  File "/home/jake/github/pysindy/env/lib/python3.10/site-packages/numpy/ma/core.py", line 6882, in __call__
    return method(marr, *args, **params)
  File "/home/jake/github/pysindy/env/lib/python3.10/site-packages/numpy/ma/core.py", line 4677, in ravel
    r = ndarray.ravel(self._data, order=order).view(type(self))
  File "/home/jake/github/pysindy/pysindy/utils/axes.py", line 353, in __array_finalize__
    self._ax_map = _AxisMapping(obj.axes, obj.ndim)
  File "/home/jake/github/pysindy/pysindy/utils/axes.py", line 305, in __getattr__
    raise AttributeError(f"'{type(self)}' object has no attribute '{name}'")
AttributeError: '<class 'pysindy.utils.axes.AxesArray'>' object has no attribute 'axes'
```

### PySINDy/Python version information:

current master, a0e36ddb08080ebbf2aeff47b7bae3e51ff8bc2a


Discussion:

Patch:
diff --git a/pysindy/utils/axes.py b/pysindy/utils/axes.py
index ed27957..25c42db 100644
--- a/pysindy/utils/axes.py
+++ b/pysindy/utils/axes.py
@@ -346,11 +346,23 @@ class AxesArray(np.lib.mixins.NDArrayOperatorsMixin, np.ndarray):
         elif all(
             (
                 isinstance(obj, AxesArray),
+                hasattr(obj, "_ax_map"),
                 not hasattr(self, "_ax_map"),
                 self.shape == obj.shape,
             )
         ):
             self._ax_map = _AxisMapping(obj.axes, obj.ndim)
+        # Using a poorly-initialized AxesArray
+        # Occurs in MaskedArray.ravel, used in some plotting.  MaskedArray views
+        # of AxesArray lose the axes attributes, and then the _ax_map attributes.
+        # See numpy.ma.core:asanyarray
+        elif all(
+            (
+                isinstance(obj, AxesArray),
+                not hasattr(obj, "_ax_map"),
+            )
+        ):
+            self._ax_map = _AxisMapping({"ax_unk": 0}, in_ndim=1)
         # maybe add errors for incompatible views?
 
     def __array_ufunc__(
@@ -418,6 +430,16 @@ def _implements(numpy_function):
     return decorator
 
 
+@_implements(np.ravel)
+def ravel(a, order="C"):
+    out = np.ravel(np.asarray(a), order=order)
+    is_1d_already = len(a.shape) == 1
+    if is_1d_already:
+        return AxesArray(out, a.axes)
+    else:
+        return AxesArray(out, {"ax_unk": 0})
+
+
 @_implements(np.ix_)
 def ix_(*args: AxesArray):
     calc = np.ix_(*(np.asarray(arg) for arg in args))


Test patch:
diff --git a/test/utils/test_axes.py b/test/utils/test_axes.py
index b26a738..6e85cd3 100644
--- a/test/utils/test_axes.py
+++ b/test/utils/test_axes.py
@@ -629,6 +629,27 @@ def test_tensordot_list_axes():
     assert_array_equal(result, super_result)
 
 
+def test_ravel_1d():
+    arr = AxesArray(np.array([1, 2]), axes={"ax_a": 0})
+    result = np.ravel(arr)
+    assert_array_equal(result, arr)
+    assert result.axes == arr.axes
+
+
+def test_ravel_nd():
+    arr = AxesArray(np.array([[1, 2], [3, 4]]), axes={"ax_a": 0, "ax_b": 1})
+    result = np.ravel(arr)
+    expected = np.ravel(np.asarray(arr))
+    assert_array_equal(result, expected)
+    assert result.axes == {"ax_unk": 0}
+
+
+def test_ma_ravel():
+    arr = AxesArray(np.array([1, 2]), axes={"ax_a": 0})
+    marr = np.ma.MaskedArray(arr)
+    np.ma.ravel(marr)
+
+
 @pytest.mark.skip
 def test_einsum_implicit():
     ...