    001/*-
    002 *******************************************************************************
    003 * Copyright (c) 2015, 2016 Diamond Light Source Ltd.
    004 * All rights reserved. This program and the accompanying materials
    005 * are made available under the terms of the Eclipse Public License v1.0
    006 * which accompanies this distribution, and is available at
    007 * http://www.eclipse.org/legal/epl-v10.html
    008 *
    009 * Contributors:
    010 *    Peter Chang - initial API and implementation and/or initial documentation
    011 *******************************************************************************/
    012
    013package org.eclipse.january.dataset;
    014
    015/**
    016 * Class to provide slice iteration through a dataset
    017 * <p>It allows a number of axes to be omitted and iterates over
    018 * the axes left over. The omitted axes define an inner shape and
    019 * the remaining axes define an outer shape which is iterated over.
    020 */
    021public class SliceNDIterator extends IndexIterator {
    022        final private int[] shape;
    023        final private int[] start;
    024        final private int[] stop;
    025        final private int[] step;
    026        final private int endrank;
    027
    028        final private boolean[] omit; // axes to miss out
    029
    030        /**
    031         * position in source dataset
    032         */
    033        final private int[] pos;
    034        final private int[] end;
    035        private boolean once;
    036
    037        private SliceND cSlice; // current slice
    038
    039        private int sRank; // number of dimensions used (i.e. not missing)
    040        final private SliceND iSlice; // omitted source or inner slice
    041
    042        final private SliceND sSlice; // shortened slice
    043        final private int[] sStart; // shortened position
    044        final private int[] sStop; // shortened end
    045
    046        private SliceND dSlice; // destination slice
    047        private int[] dStart;
    048        private int[] dStop;
    049
    050        /**
    051         * Constructor for an iterator that misses out several axes
    052         * @param slice
    053         * @param axes missing axes
    054         */
    055        public SliceNDIterator(SliceND slice, int... axes) {
    056                cSlice = slice.clone();
    057                int[] sShape = cSlice.getSourceShape();
    058                shape = cSlice.getShape().clone();
    059                start = cSlice.getStart();
    060                stop  = cSlice.getStop();
    061                step  = cSlice.getStep();
    062                for (int s : step) {
    063                        if (s < 0) {
    064                                throw new UnsupportedOperationException("Negative steps not implemented");
    065                        }
    066                }
    067                int rank = shape.length;
    068                endrank = rank - 1;
    069
    070                omit = new boolean[rank];
    071                dSlice = new SliceND(shape);
    072                dStart = dSlice.getStart();
    073                dStop  = dSlice.getStop();
    074                sRank = rank;
    075                if (axes != null) {
    076                        axes = ShapeUtils.checkAxes(rank, axes);
    077                        for (int a : axes) {
    078                                if (a >= 0 && a <= endrank) {
    079                                        sRank--;
    080                                        omit[a] = true;
    081                                        shape[a] = 1;
    082                                } else if (a > endrank) {
    083                                        throw new IllegalArgumentException("Specified axis exceeds dataset rank");
    084                                }
    085                        }
    086                }
    087
    088                cSlice = cSlice.clone();
    089                pos = cSlice.getStart();
    090                end = cSlice.getStop();
    091                if (sRank == rank) {
    092                        sStart = pos;
    093                        sStop = null;
    094                        iSlice = null;
    095                        sSlice = cSlice;
    096                } else {
    097                        int[] dShape = dSlice.getShape();
    098                        int[] oShape = new int[sRank]; // outer shape
    099                        int[] iShape = new int[rank - sRank]; // inner shape
    100                        for (int i = 0, j = 0, k = 0; i < rank; i++) {
    101                                if (omit[i]) {
    102                                        iShape[j++] = sShape[i];
    103                                } else {
    104                                        oShape[k++] = sShape[i];
    105                                        dShape[i] = 1;
    106                                }
    107                        }
    108                        sSlice = new SliceND(oShape);
    109                        sStart = sSlice.getStart();
    110                        sStop = sSlice.getStop();
    111
    112                        iSlice = new SliceND(iShape);
    113                        for (int i = 0, j = 0, k = 0; i < rank; i++) {
    114                                if (omit[i]) {
    115                                        iSlice.setSlice(j++, start[i], stop[i], step[i]);
    116                                } else {
    117                                        sSlice.setSlice(k++, start[i], stop[i], step[i]);
    118                                }
    119                        }
    120                }
    121
    122                reset();
    123        }
    124
    125        @Override
    126        public boolean hasNext() {
    127                // now move on one position
    128                if (once) {
    129                        once = false;
    130                        return true;
    131                }
    132                int k = sRank - 1;
    133                for (int j = endrank; j >= 0; j--) {
    134                        if (omit[j]) {
    135                                continue;
    136                        }
    137                        pos[j] += step[j];
    138                        end[j] = pos[j] + step[j];
    139                        dStart[j]++;
    140                        dStop[j]++;
    141                        if (pos[j] >= stop[j]) {
    142                                pos[j] = start[j];
    143                                end[j] = pos[j] + step[j];
    144                                dStart[j] = 0;
    145                                dStop[j] = 1;
    146                                if (sStop != null) {
    147                                        sStart[k] = pos[j];
    148                                        sStop[k] = end[j];
    149                                        k--;
    150                                }
    151                        } else {
    152                                if (sStop != null) {
    153                                        sStart[k] = pos[j];
    154                                        sStop[k] = end[j];
    155                                        k--;
    156                                }
    157                                return true;
    158                        }
    159                }
    160                return false;
    161        }
    162
    163        @Override
    164        public int[] getPos() {
    165                return pos;
    166        }
    167
    168        /**
    169         * Get omitted part of source slice which never changes. I.e. the inner slice
    170         * @return slice (can be null)
    171         */
    172        public SliceND getOmittedSlice() {
    173                return iSlice;
    174        }
    175
    176        /**
    177         * Get output or destination slice
    178         * @return slice
    179         */
    180        public SliceND getOutputSlice() {
    181                return dSlice;
    182        }
    183
    184        /**
    185         * Get current slice
    186         * @return slice
    187         */
    188        public SliceND getCurrentSlice() {
    189                return cSlice;
    190        }
    191
    192        /**
    193         * Shortened position where axes are omitted
    194         * @return used position
    195         */
    196        public int[] getUsedPos() {
    197                return sStart;
    198        }
    199
    200        /**
    201         * Shortened slice where axes are omitted
    202         * @return used (or outer) slice
    203         */
    204        public SliceND getUsedSlice() {
    205                return sSlice;
    206        }
    207
    208        /**
    209         * @return omit array - array where true means miss out
    210         */
    211        public boolean[] getOmit() {
    212                return omit;
    213        }
    214
    215        @Override
    216        public void reset() {
    217                for (int i = 0, k = 0; i <= endrank; i++) {
    218                        int b = start[i];
    219                        int d = step[i];
    220                        if (!omit[i]) {
    221                                cSlice.setSlice(i, b, b + d, d);
    222                                dStart[i] = 0;
    223                                dStop[i] = 1;
    224                                if (sStop != null) {
    225                                        sSlice.setSlice(k++, b, b + d, d);
    226                                }
    227                        } else {
    228                                cSlice.setSlice(i, b, end[i], d);
    229                        }
    230                }
    231
    232                int j = 0;
    233                for (; j <= endrank; j++) {
    234                        if (!omit[j])
    235                                break;
    236                }
    237                if (j > endrank) {
    238                        once = true;
    239                        return;
    240                }
    241
    242                if (omit[endrank]) {
    243                        pos[endrank] = start[endrank];
    244                        for (int i = endrank - 1; i >= 0; i--) {
    245                                if (!omit[i]) {
    246                                        end[i] = pos[i];
    247                                        pos[i] -= step[i];
    248                                        dStart[i]--;
    249                                        dStop[i]--;
    250                                        break;
    251                                }
    252                        }
    253                } else {
    254                        end[endrank] = pos[endrank];
    255                        pos[endrank] -= step[endrank];
    256                        dStart[endrank]--;
    257                        dStop[endrank]--;
    258                }
    259
    260                if (sStart != pos) {
    261                        for (int i = 0, k = 0; i <= endrank; i++) {
    262                                if (!omit[i]) {
    263                                        sStart[k++] = pos[i];
    264                                }
    265                        }
    266                }
    267        }
    268
    269        @Override
    270        public int[] getShape() {
    271                return shape;
    272        }
    273}