ENH: add multi-field assignment helpers in np.lib.recfunctions

Adds helper functions for the copy->view transition for multi-field
indexes. Adds `structured_to_unstructured`, `apply_along_fields`,
`assign_fields_by_name`, `require_fields`.

2 files changed, 332 insertions, 2 deletions

diff --git a/numpy/lib/recfunctions.py b/numpy/lib/recfunctions.py
index 53a586f56..11c04f03d 100644
--- a/numpy/lib/recfunctions.py
+++ b/numpy/lib/recfunctions.py
@@ -102,7 +102,7 @@ def get_fieldspec(dtype):
         fields = ((name, dtype.fields[name]) for name in dtype.names)
         # keep any titles, if present
         return [
-            (name if len(f) == 2 else (f[2], name), f[0]) 
+            (name if len(f) == 2 else (f[2], name), f[0])
             for name, f in fields
         ]
 
@@ -870,6 +870,285 @@ def repack_fields(a, align=False, recurse=False):
     dt = np.dtype(fieldinfo, align=align)
     return np.dtype((a.type, dt))
 
+def _get_fields_and_offsets(dt, offset=0):
+    """
+    Returns a flat list of (name, dtype, count, offset) tuples of all the
+    scalar fields in the dtype "dt", including nested fields, in left
+    to right order.
+    """
+    fields = []
+    for name in dt.names:
+        field = dt.fields[name]
+        if field[0].names is None:
+            count = 1
+            for size in field[0].shape:
+                count *= size
+            fields.append((name, field[0], count, field[1] + offset))
+        else:
+            fields.extend(_get_fields_and_offsets(field[0], field[1] + offset))
+    return fields
+
+def structured_to_unstructured(arr, dtype=None):
+    """
+    Converts and n-D structured array into an (n+1)-D unstructured array.
+
+    The new array will have a new last dimension equal in size to the
+    number of field-elements of the input array. If not supplied, the output
+    datatype is determined from the numpy type promotion rules applied to all
+    the field datatypes.
+
+    Nested fields, as well as each element of any subarray fields, all count
+    as a single field-elements.
+
+    Parameters
+    ----------
+    arr : ndarray
+       Structured array or dtype to convert.
+    dtype : dtype, optional
+       The dtype of the output unstructured array
+
+    Returns
+    -------
+    unstructured : ndarray
+       Unstructured array with one more dimension.
+
+    Examples
+    --------
+
+    >>> a = np.zeros(4, dtype=[('a', 'i4'), ('b', 'f4,u2'), ('c', 'f4', 2)])
+    >>> a
+    array([(0, (0., 0), [0., 0.]), (0, (0., 0), [0., 0.]),
+           (0, (0., 0), [0., 0.]), (0, (0., 0), [0., 0.])],
+          dtype=[('a', '<i4'), ('b', [('f0', '<f4'), ('f1', '<u2')]), ('c', '<f4', (2,))])
+    >>> structured_to_unstructured(arr)
+    array([[0., 0., 0., 0., 0.],
+           [0., 0., 0., 0., 0.],
+           [0., 0., 0., 0., 0.],
+           [0., 0., 0., 0., 0.]])
+
+    >>> b = np.array([(1, 2, 5), (4, 5, 7), (7, 8 ,11), (10, 11, 12)],
+    ...              dtype=[('x', 'i4'), ('y', 'f4'), ('z', 'f8')])
+    >>> np.mean(structured_to_unstructured(b[['x', 'z']]), axis=-1)
+    array([ 3. ,  5.5,  9. , 11. ])
+
+    """
+    if not arr.dtype.names:
+        raise ValueError('arr must be a structured array')
+
+    fields = _get_fields_and_offsets(arr.dtype)
+    n_elem = sum(f[2] for f in fields)
+
+    if dtype is None:
+        out_dtype = np.result_type(*[f[1].base for f in fields])
+    else:
+        out_dtype = dtype
+
+    out = np.empty(arr.shape + (n_elem,), dtype=out_dtype)
+
+    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
+
+    return out
+
+def unstructured_to_structured(arr, dtype=None, names=None, align=False):
+    """
+    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
+    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.
+
+    Nested fields, as well as each element of any subarray fields, all count
+    towards the number of field-elements.
+
+    Parameters
+    ----------
+    arr : ndarray
+       Unstructured array or dtype to convert.
+    dtype : dtype, optional
+       The structured dtype of the output array
+    names : list of strings, optional
+       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.
+
+    Returns
+    -------
+    structured : ndarray
+       Structured array with fewer dimensions.
+
+    Examples
+    --------
+
+    >>> dt = np.dtype([('a', 'i4'), ('b', 'f4,u2'), ('c', 'f4', 2)])
+    >>> a = np.arange(20).reshape((4,5))
+    >>> a
+    array([[ 0,  1,  2,  3,  4],
+           [ 5,  6,  7,  8,  9],
+           [10, 11, 12, 13, 14],
+           [15, 16, 17, 18, 19]])
+    >>> unstructured_to_structured(a, dt)
+    array([( 0, ( 1.,  2), [ 3.,  4.]), ( 5, ( 6.,  7), [ 8.,  9.]),
+           (10, (11., 12), [13., 14.]), (15, (16., 17), [18., 19.])],
+          dtype=[('a', '<i4'), ('b', [('f0', '<f4'), ('f1', '<u2')]), ('c', '<f4', (2,))])
+
+    """
+    if arr.shape == ():
+        raise ValueError('arr must have at least one dimension')
+    n_elem = arr.shape[-1]
+
+    if dtype is None:
+        if names is None:
+            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)
+    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:
+            raise ValueError('The length of the last dimension of arr must '
+                             'be equal to the number of fields in dtype')
+        out_dtype = dtype
+
+    out = np.empty(arr.shape[:-1], dtype=out_dtype)
+
+    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
+
+    return out
+
+def apply_along_fields(func, arr):
+    """
+    Apply function 'func' as a reduction across fields of a structured array.
+
+    This is similar to `apply_along_axis`, but treats the fields of a
+    structured array as an extra axis.
+
+    Parameters
+    ----------
+    func : function
+       Function to apply on the "field" dimension. This function must
+       support an `axis` argument, like np.mean, np.sum, etc.
+    arr : ndarray
+       Structured array for which to apply func.
+
+    Returns
+    -------
+    out : ndarray
+       Result of the recution operation
+
+    Examples
+    --------
+
+    >>> b = np.array([(1, 2, 5), (4, 5, 7), (7, 8 ,11), (10, 11, 12)],
+    ...              dtype=[('x', 'i4'), ('y', 'f4'), ('z', 'f8')])
+    >>> apply_along_fields(np.mean, b)
+    array([ 2.66666667,  5.33333333,  8.66666667, 11.        ])
+    >>> apply_along_fields(np.mean, b[['x', 'z']])
+    array([ 3. ,  5.5,  9. , 11. ])
+
+    """
+    if not arr.dtype.names:
+        raise ValueError('arr must be a structured array')
+
+    uarr = structured_to_unstructured(arr)
+    return func(uarr, axis=-1)
+    # works and avoids axis requirement, but very, very slow:
+    #return np.apply_along_axis(func, -1, uarr)
+
+def assign_fields_by_name(dst, src, zero_unassigned=True):
+    """
+    Assigns values from one structured array to another by field name.
+
+    Normally in numpy >= 1.14, assignment of one structured array to another
+    copies fields "by position", meaning that the first field from the src is
+    copied to the first field of the dst, and so on, regardless of field name.
+
+    This function instead copies "by field name", such that fields in the dst
+    are assigned from the identically named field in the src. This applies
+    recursively for nested structures. This is how structure assignment worked
+    in numpy >= 1.6 to <= 1.13.
+
+    Parameters
+    ----------
+    dst : ndarray
+    src : ndarray
+        The source and destination arrays during assignment.
+    zero_unassigned : bool, optional
+        If True, fields in the dst for which there was no matching
+        field in the src are filled with the value 0 (zero). This
+        was the behavior of numpy <= 1.13. If False, those fields
+        are not modified.
+    """
+
+    if dst.dtype.names is None:
+        dst[:] = src
+        return
+
+    for name in dst.dtype.names:
+        if name not in src.dtype.names:
+            if zero_unassigned:
+                dst[name] = 0
+        else:
+            assign_fields_by_name(dst[name], src[name],
+                                  zero_unassigned)
+
+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.
+
+    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.
+
+    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.
+
+    Parameters
+    ----------
+    a : ndarray
+       array to cast
+    required_dtype : dtype
+       datatype for output array
+
+    Returns
+    -------
+    out : ndarray
+        array with the new dtype, with field values copied from the fields in
+        the input array with the same name
+
+    Examples
+    --------
+
+    >>> a = np.ones(4, dtype=[('a', 'i4'), ('b', 'f8'), ('c', 'u1')])
+    >>> require_fields(a, [('b', 'f4'), ('c', 'u1')])
+    """
+    out = np.empty(array.shape, dtype=required_dtype)
+    assign_fields_by_name(out, array)
+    return out
+
 
 def _stack_arrays_dispatcher(arrays, defaults=None, usemask=None,
                              asrecarray=None, autoconvert=None):
diff --git a/numpy/lib/tests/test_recfunctions.py b/numpy/lib/tests/test_recfunctions.py
index 5585a95f9..8b61ba0bf 100644
--- a/numpy/lib/tests/test_recfunctions.py
+++ b/numpy/lib/tests/test_recfunctions.py
@@ -10,7 +10,8 @@ from numpy.testing import assert_, assert_raises
 from numpy.lib.recfunctions import (
     drop_fields, rename_fields, get_fieldstructure, recursive_fill_fields,
     find_duplicates, merge_arrays, append_fields, stack_arrays, join_by,
-    repack_fields)
+    repack_fields, unstructured_to_structured, structured_to_unstructured,
+    apply_along_fields, require_fields, assign_fields_by_name)
 get_names = np.lib.recfunctions.get_names
 get_names_flat = np.lib.recfunctions.get_names_flat
 zip_descr = np.lib.recfunctions.zip_descr
@@ -204,6 +205,56 @@ class TestRecFunctions(object):
         dt = np.dtype((np.record, dt))
         assert_(repack_fields(dt).type is np.record)
 
+    def test_structured_to_unstructured(self):
+        a = np.zeros(4, dtype=[('a', 'i4'), ('b', 'f4,u2'), ('c', 'f4', 2)])
+        out = structured_to_unstructured(a)
+        assert_equal(out, np.zeros((4,5), dtype='f8'))
+
+        b = np.array([(1, 2, 5), (4, 5, 7), (7, 8 ,11), (10, 11, 12)],
+                     dtype=[('x', 'i4'), ('y', 'f4'), ('z', 'f8')])
+        out = np.mean(structured_to_unstructured(b[['x', 'z']]), axis=-1)
+        assert_equal(out, np.array([ 3. ,  5.5,  9. , 11. ]))
+
+        c = np.arange(20).reshape((4,5))
+        out = unstructured_to_structured(c, a.dtype)
+        want = np.array([( 0, ( 1.,  2), [ 3.,  4.]),
+                         ( 5, ( 6.,  7), [ 8.,  9.]),
+                         (10, (11., 12), [13., 14.]),
+                         (15, (16., 17), [18., 19.])],
+                     dtype=[('a', '<i4'),
+                            ('b', [('f0', '<f4'), ('f1', '<u2')]),
+                            ('c', '<f4', (2,))])
+        assert_equal(out, want)
+
+        d = np.array([(1, 2, 5), (4, 5, 7), (7, 8 ,11), (10, 11, 12)],
+                     dtype=[('x', 'i4'), ('y', 'f4'), ('z', 'f8')])
+        assert_equal(apply_along_fields(np.mean, d),
+                     np.array([ 8.0/3,  16.0/3,  26.0/3, 11. ]))
+        assert_equal(apply_along_fields(np.mean, d[['x', 'z']]),
+                     np.array([ 3. ,  5.5,  9. , 11. ]))
+
+    def test_field_assignment_by_name(self):
+        a = np.ones(2, dtype=[('a', 'i4'), ('b', 'f8'), ('c', 'u1')])
+        newdt = [('b', 'f4'), ('c', 'u1')]
+
+        assert_equal(require_fields(a, newdt), np.ones(2, newdt))
+
+        b = np.array([(1,2), (3,4)], dtype=newdt)
+        assign_fields_by_name(a, b, zero_unassigned=False)
+        assert_equal(a, np.array([(1,1,2),(1,3,4)], dtype=a.dtype))
+        assign_fields_by_name(a, b)
+        assert_equal(a, np.array([(0,1,2),(0,3,4)], dtype=a.dtype))
+
+        # test nested fields
+        a = np.ones(2, dtype=[('a', [('b', 'f8'), ('c', 'u1')])])
+        newdt = [('a', [('c', 'u1')])]
+        assert_equal(require_fields(a, newdt), np.ones(2, newdt))
+        b = np.array([((2,),), ((3,),)], dtype=newdt)
+        assign_fields_by_name(a, b, zero_unassigned=False)
+        assert_equal(a, np.array([((1,2),), ((1,3),)], dtype=a.dtype))
+        assign_fields_by_name(a, b)
+        assert_equal(a, np.array([((0,2),), ((0,3),)], dtype=a.dtype))
+
 
 class TestRecursiveFillFields(object):
     # Test recursive_fill_fields.