    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 an n-D 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 = sShape == null ? null : 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                final int rank = shape == null ? 0 : 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                        int[] dShape = dSlice.getShape();
    097                        for (int i = 0; i < rank; i++) {
    098                                dShape[i] = 1;
    099                        }
    100                } else {
    101                        int[] dShape = dSlice.getShape();
    102                        int[] oShape = new int[sRank]; // outer shape
    103                        int[] iShape = new int[rank - sRank]; // inner shape
    104                        for (int i = 0, j = 0, k = 0; i < rank; i++) {
    105                                if (omit[i]) {
    106                                        iShape[j++] = sShape[i];
    107                                } else {
    108                                        oShape[k++] = sShape[i];
    109                                        dShape[i] = 1;
    110                                }
    111                        }
    112                        sSlice = new SliceND(oShape);
    113                        sStart = sSlice.getStart();
    114                        sStop = sSlice.getStop();
    115
    116                        iSlice = new SliceND(iShape);
    117                        for (int i = 0, j = 0, k = 0; i < rank; i++) {
    118                                if (omit[i]) {
    119                                        iSlice.setSlice(j++, start[i], stop[i], step[i]);
    120                                } else {
    121                                        sSlice.setSlice(k++, start[i], stop[i], step[i]);
    122                                }
    123                        }
    124                }
    125
    126                reset();
    127        }
    128
    129        @Override
    130        public boolean hasNext() {
    131                // now move on one position
    132                if (once) {
    133                        once = false;
    134                        return true;
    135                }
    136                int k = sRank - 1;
    137                for (int j = endrank; j >= 0; j--) {
    138                        if (omit[j]) {
    139                                continue;
    140                        }
    141                        pos[j] += step[j];
    142                        end[j] = pos[j] + step[j];
    143                        dStart[j]++;
    144                        dStop[j]++;
    145                        if (pos[j] >= stop[j]) {
    146                                pos[j] = start[j];
    147                                end[j] = pos[j] + step[j];
    148                                dStart[j] = 0;
    149                                dStop[j] = 1;
    150                                if (sStop != null) {
    151                                        sStart[k] = pos[j];
    152                                        sStop[k] = end[j];
    153                                        k--;
    154                                }
    155                        } else {
    156                                if (sStop != null) {
    157                                        sStart[k] = pos[j];
    158                                        sStop[k] = end[j];
    159                                        k--;
    160                                }
    161                                return true;
    162                        }
    163                }
    164                return false;
    165        }
    166
    167        @Override
    168        public int[] getPos() {
    169                return pos;
    170        }
    171
    172        /**
    173         * Get omitted part of source slice which never changes. I.e. the inner slice
    174         * @return slice (can be null)
    175         */
    176        public SliceND getOmittedSlice() {
    177                return iSlice;
    178        }
    179
    180        /**
    181         * Get output or destination slice
    182         * @return slice
    183         */
    184        public SliceND getOutputSlice() {
    185                return dSlice;
    186        }
    187
    188        /**
    189         * Get current slice
    190         * @return slice
    191         */
    192        public SliceND getCurrentSlice() {
    193                return cSlice;
    194        }
    195
    196        /**
    197         * Shortened position where axes are omitted
    198         * @return used position
    199         */
    200        public int[] getUsedPos() {
    201                return sStart;
    202        }
    203
    204        /**
    205         * Shortened slice where axes are omitted
    206         * @return used (or outer) slice
    207         */
    208        public SliceND getUsedSlice() {
    209                return sSlice;
    210        }
    211
    212        /**
    213         * @return omit array - array where true means miss out
    214         */
    215        public boolean[] getOmit() {
    216                return omit;
    217        }
    218
    219        @Override
    220        public void reset() {
    221                once = false;
    222                for (int i = 0, k = 0; i <= endrank; i++) {
    223                        int b = start[i];
    224                        int d = step[i];
    225                        if (!omit[i]) {
    226                                cSlice.setSlice(i, b, b + d, d);
    227                                dStart[i] = 0;
    228                                dStop[i] = 1;
    229                                if (sStop != null) {
    230                                        sSlice.setSlice(k++, b, b + d, d);
    231                                }
    232                        } else {
    233                                cSlice.setSlice(i, b, end[i], d);
    234                        }
    235                }
    236
    237                int j = 0;
    238                for (; j <= endrank; j++) {
    239                        if (!omit[j])
    240                                break;
    241                }
    242                if (j > endrank) {
    243                        once = shape != null;
    244                        return;
    245                }
    246
    247                if (omit[endrank]) {
    248                        pos[endrank] = start[endrank];
    249                        for (int i = endrank - 1; i >= 0; i--) {
    250                                if (!omit[i]) {
    251                                        end[i] = pos[i];
    252                                        pos[i] -= step[i];
    253                                        dStart[i]--;
    254                                        dStop[i]--;
    255                                        break;
    256                                }
    257                        }
    258                } else {
    259                        end[endrank] = pos[endrank];
    260                        pos[endrank] -= step[endrank];
    261                        dStart[endrank]--;
    262                        dStop[endrank]--;
    263                }
    264
    265                if (sStart != pos) {
    266                        for (int i = 0, k = 0; i <= endrank; i++) {
    267                                if (!omit[i]) {
    268                                        sStart[k++] = pos[i];
    269                                }
    270                        }
    271                }
    272        }
    273
    274        @Override
    275        public int[] getShape() {
    276                return shape;
    277        }
    278}