Improve iris.pandas cube -> data.frame

lib/iris/pandas.py

@@ -124,65 +124,7 @@ def _as_pandas_coord(coord):
     return index
 
 
-def _assert_shared(np_obj, pandas_obj):
-    """Ensure the pandas object shares memory."""
-    values = pandas_obj.values
-
-    def _get_base(array):
-        # Chase the stack of NumPy `base` references back to the original array
-        while array.base is not None:
-            array = array.base
-        return array
-
-    base = _get_base(values)
-    np_base = _get_base(np_obj)
-    if base is not np_base:
-        msg = "Pandas {} does not share memory".format(
-            type(pandas_obj).__name__
-        )
-        raise AssertionError(msg)
-
-
-def as_series(cube, copy=True):
-    """
-    Convert a 1D cube to a Pandas Series.
-
-    Args:
-
-        * cube - The cube to convert to a Pandas Series.
-
-    Kwargs:
-
-        * copy - Whether to make a copy of the data.
-                 Defaults to True. Must be True for masked data.
-
-    .. note::
-
-        This function will copy your data by default.
-        If you have a large array that cannot be copied,
-        make sure it is not masked and use copy=False.
-
-    """
-    data = cube.data
-    if ma.isMaskedArray(data):
-        if not copy:
-            raise ValueError("Masked arrays must always be copied.")
-        data = data.astype("f").filled(np.nan)
-    elif copy:
-        data = data.copy()
-
-    index = None
-    if cube.dim_coords:
-        index = _as_pandas_coord(cube.dim_coords[0])
-
-    series = pandas.Series(data, index)
-    if not copy:
-        _assert_shared(data, series)
-
-    return series
-
-
-def as_data_frame(cube, copy=True):
+def as_data_frame(cube, dropna=True, asmultiindex=False, add_aux_coord=None):
     """
     Convert a 2D cube to a Pandas DataFrame.
 
@@ -192,39 +134,55 @@ def as_data_frame(cube, copy=True):
 
     Kwargs:
 
-        * copy - Whether to make a copy of the data.
-                 Defaults to True. Must be True for masked data
-                 and some data types (see notes below).
+        * dropna - Remove missing values from returned dataframe.
+                    Defaults to True.
 
     .. note::
 
-        This function will copy your data by default.
-        If you have a large array that cannot be copied,
-        make sure it is not masked and use copy=False.
-
-    .. note::
-
-        Pandas will sometimes make a copy of the array,
-        for example when creating from an int32 array.
-        Iris will detect this and raise an exception if copy=False.
+        TBC
 
     """
     data = cube.data
     if ma.isMaskedArray(data):
-        if not copy:
-            raise ValueError("Masked arrays must always be copied.")
         data = data.astype("f").filled(np.nan)
     elif copy:
         data = data.copy()
 
-    index = columns = None
-    if cube.coords(dimensions=[0]):
-        index = _as_pandas_coord(cube.coord(dimensions=[0]))
-    if cube.coords(dimensions=[1]):
-        columns = _as_pandas_coord(cube.coord(dimensions=[1]))
-
-    data_frame = pandas.DataFrame(data, index, columns)
-    if not copy:
-        _assert_shared(data, data_frame)
+    # Extract dim coord information
+    if cube.ndim != len(cube.dim_coords):
+        # Create dummy dim coord information if dim coords not defined
+        coord_names = ["dim" + str(n) for n in range(cube.ndim)]
+        coords = [range(dim) for dim in cube.shape]
+        for c in cube.dim_coords:
+            for i, dummyc in enumerate(coords):
+                if len(dummyc) == len(c.points):
+                    coords[i] = _as_pandas_coord(c)
+                    coord_names[i] = c.name()
+                else:
+                    pass
+    else:
+        coord_names = list(map(lambda x: x.name(), cube.dim_coords))
+        coords = list(map(lambda x: _as_pandas_coord(x), cube.dim_coords))
+
+    index = pandas.MultiIndex.from_product(coords, names=coord_names)
+    data_frame = pandas.DataFrame({cube.name(): data.flatten()}, index=index)
+
+    # Add aux coord information
+    if add_aux_coord:
+        aux_coord_names = list(map(lambda x: x.name(), cube.aux_coords))
+        for acoord in add_aux_coord:
+            assert acoord in aux_coord_names, f'\"{acoord}\" not in cube' # Check aux coord exists
+            aux_coord = cube.coord(acoord)
+            coord_bool = np.array(cube.shape) == aux_coord.shape[0] # Which dim coords match aux coord length
+            aux_coord_index = np.array(coords)[coord_bool][0] # Get corresponding dim coord
+            # Build aux coord dataframe
+            acoord_df = pd.DataFrame({acoord: aux_coord.points}, index = pd.Index(data=aux_coord_index, name=np.array(coord_names)[coord_bool][0]))
+            # Join to main data frame
+            data_frame = data_frame.join(acoord_df, on=np.array(coord_names)[coord_bool][0])
+
+    if dropna:
+        data_frame.dropna(inplace=True)
+    if not asmultiindex:
+        data_frame.reset_index(inplace=True)
 
     return data_frame