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