Skip to main content
`Ctrl`+`K`

# Source code for ndcube.utils.collection

    import numbers

    import numpy as np

    __all__ = ['assert_aligned_axes_compatible']

    def _sanitize_aligned_axes(keys, data, aligned_axes):
        if aligned_axes is None:
            return None
        # If aligned_axes set to "all", assume all axes are aligned in order.
        if isinstance(aligned_axes, str) and aligned_axes.lower() == "all":
            # Check all cubes are of same shape
            cube0_dims = data[0].shape
            cubes_same_shape = all(all(d.shape[i] == dim for i, dim in enumerate(cube0_dims))
                                    for d in data)
            if cubes_same_shape is not True:
                raise ValueError(
                    "All cubes in data not of same shape. Please set aligned_axes kwarg.")
            sanitized_axes = tuple([tuple(range(len(cube0_dims)))] * len(data))
        else:
            # Else, sanitize user-supplied aligned axes.
            sanitized_axes = _sanitize_user_aligned_axes(data, aligned_axes)

        return dict(zip(keys, sanitized_axes))

    def _sanitize_user_aligned_axes(data, aligned_axes):
        """
        Converts input aligned_axes to standard format.
        aligned_axes can be supplied by the user in a few ways:
        *. A tuple of tuples of ints, where each tuple corresponds to a cube
        in the collection, and each int designates the an aligned axis in numpy order.
        In this case, the axis represented by the 0th int in the 0th tuple is aligned
        with the 0th int in the 1st tuple and so on.
        *. A single tuple of ints if all aligned axes are in the same order.
        *. A single int if only one axis is aligned and if the aligned axis in each cube
        is in the same order.
        """
        aligned_axes_error_message = ("aligned_axes must contain ints or "
                                      "a tuple of ints for each element in data.")
        if isinstance(data[0].shape, tuple):
            cube0_dims = np.array(data[0].shape, dtype=object)[np.array(aligned_axes[0])]
        else:
            cube0_dims = data[0].shape[np.array(aligned_axes[0])]
        # If user entered a single int or string, convert to length 1 tuple of int.
        if isinstance(aligned_axes, numbers.Integral):
            aligned_axes = (aligned_axes,)
        if not isinstance(aligned_axes, tuple):
            raise ValueError(aligned_axes_error_message)
        # Check type of each element.
        axes_all_ints = all(isinstance(axis, numbers.Integral) for axis in aligned_axes)
        axes_all_tuples = all(isinstance(axis, tuple) for axis in aligned_axes)
        # If all elements are int, duplicate tuple so there is one for each cube.
        n_cubes = len(data)
        if axes_all_ints:
            n_aligned_axes = len(aligned_axes)
            aligned_axes = tuple([aligned_axes for i in range(n_cubes)])

        # If all elements are tuple, ensure there is a tuple for each cube and
        # all elements of each sub-tuple are ints.
        elif axes_all_tuples:
            if len(aligned_axes) != n_cubes:
                raise ValueError("aligned_axes must have a tuple for each element in data.")

            n_aligned_axes = len(aligned_axes[0])

            # Ensure all elements of sub-tuples are ints,
            # each tuple has the same number of aligned axes,
            # number of aligned axes are <= number of cube dimensions,
            # and the dimensions of the aligned axes in each cube are the same.
            subtuples_are_ints = [False] * n_cubes
            aligned_axes_same_lengths = [False] * n_cubes
            if not all(len(axes) == n_aligned_axes for axes in aligned_axes):
                raise ValueError("Each element in aligned_axes must have same length.")
            for i in range(n_cubes):
                # Check each cube has at least as many dimensions as there are aligned axes
                # and that all cubes have enough dimensions to accommodate aligned axes.
                n_cube_dims = len(data[i].shape)
                max_aligned_axis = max(aligned_axes[i])
                if n_cube_dims < max([max_aligned_axis, n_aligned_axes]):
                    raise ValueError(
                        "Each cube in data must have at least as many axes as aligned axes "
                        "and aligned axis indices must be less than number of cube axes.\n"
                        f"Cube number: {i};\n"
                        f"Number of cube dimensions: {n_cube_dims};\n"
                        f"No. aligned axes: {n_aligned_axes};\n"
                        f"Highest aligned axis: {max_aligned_axis}")
                subtuple_types = [False] * n_aligned_axes
                cube_lengths_equal = [False] * n_aligned_axes
                for j, axis in enumerate(aligned_axes[i]):
                    subtuple_types[j] = isinstance(axis, numbers.Integral)
                    cube_lengths_equal[j] = data[i].shape[axis] == cube0_dims[j]
                subtuples_are_ints[i] = all(subtuple_types)
                aligned_axes_same_lengths[i] = all(cube_lengths_equal)
            if not all(subtuples_are_ints):
                raise ValueError(aligned_axes_error_message)
            if not all(aligned_axes_same_lengths):
                raise ValueError("Aligned cube/sequence axes must be of same length.")
        else:
            raise ValueError(aligned_axes_error_message)

        # Ensure all aligned axes are of same length.
        check_dimensions = {len({cube.shape[cube_aligned_axes[j]]
                                         for cube, cube_aligned_axes in zip(data, aligned_axes)})
                                for j in range(n_aligned_axes)}
        if check_dimensions != {1}:
            raise ValueError("Aligned axes are not all of same length.")

        return aligned_axes

    def _update_aligned_axes(drop_aligned_axes_indices, aligned_axes, first_key):
        # Remove dropped axes from aligned_axes. MUST BE A BETTER WAY TO DO THIS.
        if len(drop_aligned_axes_indices) <= 0:
            new_aligned_axes = tuple(aligned_axes.values())
        elif len(drop_aligned_axes_indices) == len(aligned_axes[first_key]):
            new_aligned_axes = None
        else:
            new_aligned_axes = []
            for key in aligned_axes.keys():
                cube_aligned_axes = np.array(aligned_axes[key])
                for drop_axis_index in drop_aligned_axes_indices:
                    drop_axis = cube_aligned_axes[drop_axis_index]
                    cube_aligned_axes = np.delete(cube_aligned_axes, drop_axis_index)
                    w = np.where(cube_aligned_axes > drop_axis)
                    cube_aligned_axes[w] -= 1
                    w = np.where(drop_aligned_axes_indices > drop_axis_index)
                    drop_aligned_axes_indices[w] -= 1
                new_aligned_axes.append(tuple(cube_aligned_axes))
            new_aligned_axes = tuple(new_aligned_axes)

        return new_aligned_axes

    def assert_aligned_axes_compatible(data_dimensions1, data_dimensions2, data_axes1, data_axes2):
        """
        Checks whether two sets of aligned axes are compatible.

        Parameters
        ----------
        data_dimensions1: `tuple` of `int`
            The dimension lengths of data cube 1.
        data_dimensions2: `tuple` of `int`
            The dimension lengths of data cube 2.
        data_axes1: `tuple` of `int`
            The aligned axes of data cube 1.
        data_axes2: `tuple` of `int`
            The aligned axes of data cube 2.
        """
        # If one set of aligned axes is None and the other isn't, they are defined as not compatible.
        if (data_axes1 is None and data_axes2 is not None) or (data_axes1 is not None and data_axes2 is None):
            raise ValueError("Both collections must use aligned_axes or both axes must not use aligned_axes."
                             f"Currently {data_axes1} != {data_axes2}")

        # Aligned_axes are being used for both collections
        if data_axes1 is not None:
            # Confirm same number of aligned axes.
            if len(data_axes1) != len(data_axes2):
                raise ValueError(f"Number of aligned axes must be equal: {len(data_axes1)} != {len(data_axes2)}")
            # Confirm dimension lengths of each aligned axis is the same.
            if not all(np.array(data_dimensions1)[np.array(data_axes1)] == np.array(data_dimensions2)[np.array(data_axes2)]):
                raise ValueError("All corresponding aligned axes between cubes must be of same length.")