ENH: Fixups to multi-field assignment helpers

ahaldane · ahaldane · commit c89273320a96 · 2018-11-22T18:15:51.000-05:00
diff --git a/doc/release/1.16.0-notes.rst b/doc/release/1.16.0-notes.rst
@@ -19,6 +19,11 @@ Highlights
 New functions
 =============
 
+ * New functions in the `numpy.lib.recfunctions` module to ease the structured
+   assignment changes: `assign_fields_by_name`, `structured_to_unstructured`, 
+   `unstructured_to_structured`, `apply_along_fields`, and `require_fields`.
+   See the user guide at <https://docs.scipy.org/doc/numpy/user/basics.rec.html>
+   for more info.
 
 Deprecations
 ============
diff --git a/numpy/lib/recfunctions.py b/numpy/lib/recfunctions.py
@@ -888,7 +888,7 @@ def _get_fields_and_offsets(dt, offset=0):
             fields.extend(_get_fields_and_offsets(field[0], field[1] + offset))
     return fields
 
-def structured_to_unstructured(arr, dtype=None):
+def structured_to_unstructured(arr, dtype=None, copy=False, casting='unsafe'):
     """
     Converts and n-D structured array into an (n+1)-D unstructured array.
 
@@ -903,9 +903,15 @@ def structured_to_unstructured(arr, dtype=None):
     Parameters
     ----------
     arr : ndarray
-       Structured array or dtype to convert.
+       Structured array or dtype to convert. Cannot contain object datatype.
     dtype : dtype, optional
        The dtype of the output unstructured array
+    copy : bool, optional
+        See copy argument to `ndarray.astype`. If true, always return a copy.
+        If false, and `dtype` requirements are satisfied, a view is returned.
+    casting : {'no', 'equiv', 'safe', 'same_kind', 'unsafe'}, optional
+        See casting argument of `ndarray.astype`. Controls what kind of data
+        casting may occur.
 
     Returns
     -------
@@ -932,39 +938,46 @@ def structured_to_unstructured(arr, dtype=None):
     array([ 3. ,  5.5,  9. , 11. ])
 
     """
-    if not arr.dtype.names:
+    if arr.dtype.names is None:
         raise ValueError('arr must be a structured array')
 
     fields = _get_fields_and_offsets(arr.dtype)
-    n_elem = sum(f[2] for f in fields)
+    names, dts, counts, offsets = zip(*fields)
+    n_fields = len(names)
 
     if dtype is None:
-        out_dtype = np.result_type(*[f[1].base for f in fields])
+        out_dtype = np.result_type(*[dt.base for dt in dts])
     else:
         out_dtype = dtype
 
-    out = np.empty(arr.shape + (n_elem,), dtype=out_dtype)
+    # Use a series of views and casts to convert to an unstructured array:
 
-    index = 0
-    for name, dt, count, offset in fields:
-        if count == 1:
-            out[...,index] = arr.getfield(dt, offset)
-            index += 1
-        else:
-            out[...,index:index+count] = arr.getfield(dt, offset)
-            index += count
+    # first view using flattened fields (doesn't work for object arrays)
+    # Note: dts may include a shape for subarrays
+    flattened_fields = np.dtype({'names': names,
+                                 'formats': dts,
+                                 'offsets': offsets,
+                                 'itemsize': arr.dtype.itemsize})
+    arr = arr.view(flattened_fields)
 
-    return out
+    # next cast to a packed format with all fields converted to new dtype
+    packed_fields = np.dtype({'names': names,
+                              'formats': [(out_dtype, c) for c in counts]})
+    arr = arr.astype(packed_fields, copy=copy, casting=casting)
 
-def unstructured_to_structured(arr, dtype=None, names=None, align=False):
+    # finally is it safe to view the packed fields as the unstructured type
+    return arr.view((out_dtype, sum(counts)))
+
+def unstructured_to_structured(arr, dtype=None, names=None, align=False,
+                               copy=False, casting='unsafe'):
     """
     Converts and n-D unstructured array into an (n-1)-D structured array.
 
-    The last dimension of the array is converted into a structure, with
+    The last dimension of the input array is converted into a structure, with
     number of field-elements equal to the size of the last dimension of the
-    input array. By default all fields will have the same dtype as the
-    original array, but you may supply a custom dtype with the right
-    number of fields-elements.
+    input array. By default all output fields have the input array's dtype, but
+    an output structured dtype with an equal number of fields-elements can be
+    supplied instead.
 
     Nested fields, as well as each element of any subarray fields, all count
     towards the number of field-elements.
@@ -979,7 +992,13 @@ def unstructured_to_structured(arr, dtype=None, names=None, align=False):
        If dtype is not supplied, this specifies the field names for the output
        dtype, in order. The field dtypes will be the same as the input array.
     align : boolean, optional
-       If dtype is not supplied, whether to create an aligned memory layout.
+       Whether to create an aligned memory layout.
+    copy : bool, optional
+        See copy argument to `ndarray.astype`. If true, always return a copy.
+        If false, and `dtype` requirements are satisfied, a view is returned.
+    casting : {'no', 'equiv', 'safe', 'same_kind', 'unsafe'}, optional
+        See casting argument of `ndarray.astype`. Controls what kind of data
+        casting may occur.
 
     Returns
     -------
@@ -1011,29 +1030,36 @@ def unstructured_to_structured(arr, dtype=None, names=None, align=False):
             names = ['f{}'.format(n) for n in range(n_elem)]
         out_dtype = np.dtype([(n, arr.dtype) for n in names], align=align)
         fields = _get_fields_and_offsets(out_dtype)
+        names, dts, counts, offsets = zip(*fields)
     else:
         if names is not None:
             raise ValueError("don't supply both dtype and names")
         # sanity check of the input dtype
         fields = _get_fields_and_offsets(dtype)
-        n_fields = sum(f[2] for f in fields)
-        if n_fields != n_elem:
+        names, dts, counts, offsets = zip(*fields)
+        if n_elem != sum(counts):
             raise ValueError('The length of the last dimension of arr must '
                              'be equal to the number of fields in dtype')
         out_dtype = dtype
+        if align and not out_dtype.isalignedstruct:
+            raise ValueError("align was True but dtype is not aligned")
 
-    out = np.empty(arr.shape[:-1], dtype=out_dtype)
+    # Use a series of views and casts to convert to a structured array:
 
-    n = 0
-    for name, dt, count, offset in fields:
-        if count == 1:
-            out.setfield(arr[...,n], dt, offset)
-            n += 1
-        else:
-            out.setfield(arr[...,n:n+count], dt, offset)
-            n += count
+    # first view as a packed structured array of one dtype
+    packed_fields = np.dtype({'names': names,
+                              'formats': [(arr.dtype, c) for c in counts]})
+    arr = np.ascontiguousarray(arr).view(packed_fields)
 
-    return out
+    # next cast to an unpacked but flattened format with varied dtypes
+    flattened_fields = np.dtype({'names': names,
+                                 'formats': dts,
+                                 'offsets': offsets,
+                                 'itemsize': out_dtype.itemsize})
+    arr = arr.astype(flattened_fields, copy=copy, casting=casting)
+
+    # finally view as the final nested dtype and remove the last axis
+    return arr.view(out_dtype)[..., 0]
 
 def apply_along_fields(func, arr):
     """
@@ -1066,7 +1092,7 @@ def apply_along_fields(func, arr):
     array([ 3. ,  5.5,  9. , 11. ])
 
     """
-    if not arr.dtype.names:
+    if arr.dtype.names is None:
         raise ValueError('arr must be a structured array')
 
     uarr = structured_to_unstructured(arr)
@@ -1113,15 +1139,11 @@ def assign_fields_by_name(dst, src, zero_unassigned=True):
 
 def require_fields(array, required_dtype):
     """
-    Casts the array to the required dtype using assignment by field-name.
-
-    Normal structured array casting/assignment works "by position" in numpy
-    1.14+, meaning that the first field from the source's dtype is copied to
-    the first field of the destination's dtype, and so on.
+    Casts a structured array to a new dtype using assignment by field-name.
 
-    This function assigns by name instead, so the value of a field in the
-    output array is the value of the field with the same name in the source
-    array.
+    This function assigns to from the old to the new array by name, so the
+    value of a field in the output array is the value of the field with the
+    same name in the source array.
 
     If a field name in the required_dtype does not exist in the
     input array, that field is set to 0 in the output array.
diff --git a/numpy/lib/tests/test_recfunctions.py b/numpy/lib/tests/test_recfunctions.py
@@ -233,6 +233,14 @@ def test_structured_to_unstructured(self):
         assert_equal(apply_along_fields(np.mean, d[['x', 'z']]),
                      np.array([ 3. ,  5.5,  9. , 11. ]))
 
+        # check that for uniform field dtypes we get a view, not a copy:
+        d = np.array([(1, 2, 5), (4, 5, 7), (7, 8 ,11), (10, 11, 12)],
+                     dtype=[('x', 'i4'), ('y', 'i4'), ('z', 'i4')])
+        dd = structured_to_unstructured(d)
+        ddd = unstructured_to_structured(dd, d.dtype)
+        assert_(dd.base is d)
+        assert_(ddd.base is d)
+
     def test_field_assignment_by_name(self):
         a = np.ones(2, dtype=[('a', 'i4'), ('b', 'f8'), ('c', 'u1')])
         newdt = [('b', 'f4'), ('c', 'u1')]
