    001/*-
    002 *******************************************************************************
    003 * Copyright (c) 2011, 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
    016import java.util.Arrays;
    017
    018/**
    019 * Class to run over a slice of a dataset
    020 */
    021public class SliceIterator extends IndexIterator {
    022        int[] shape;
    023        int isize;
    024        int endrank; // last shape index
    025        int[] gap; // gaps in dataset
    026        int imax; // maximum index in array
    027        int[] start;
    028        int[] stop;
    029        int[] step;
    030        int[] sshape; // slice shape
    031        int[] pos; // position in dataset
    032        int istep; // step in last index
    033
    034        SliceIterator() {
    035        }
    036
    037        /**
    038         * Constructor for an iterator over the elements of a sliced dataset
    039         *
    040         * @param shape or dataShape
    041         * @param length of entire data array
    042         * @param sshape shape of new dataset, i.e. slice
    043         */
    044        public SliceIterator(final int[] shape, final int length, final int[] sshape) {
    045                this(shape, length, null, null, sshape, 1);
    046        }
    047
    048        /**
    049         * Constructor for an iterator over the elements of a sliced dataset
    050         *
    051         * @param shape or dataShape
    052         * @param length of entire data array
    053         * @param start (if null then equivalent to the origin)
    054         * @param sshape shape of new dataset, i.e. slice
    055         */
    056        public SliceIterator(final int[] shape, final int length, final int[] start, final int[] sshape) {
    057                this(shape, length, start, null, sshape, 1);
    058        }
    059
    060        /**
    061         * Constructor for an iterator over the elements of a sliced dataset
    062         *
    063         * @param shape or dataShape
    064         * @param length of entire data array
    065         * @param sshape shape of new dataset, i.e. slice
    066         * @param isize number of elements in an item
    067         */
    068        public SliceIterator(final int[] shape, final int length, final int[] sshape, final int isize) {
    069                this(shape, length, null, null, sshape, isize);
    070        }
    071
    072        /**
    073         * Constructor for an iterator over the elements of a sliced dataset
    074         *
    075         * @param shape or dataShape
    076         * @param length of entire data array
    077         * @param start (if null then equivalent to the origin)
    078         * @param sshape shape of new dataset, i.e. slice
    079         * @param isize number of elements in an item
    080         */
    081        public SliceIterator(final int[] shape, final int length, final int[] start, final int[] sshape, final int isize) {
    082                this(shape, length, start, null, sshape, isize);
    083        }
    084
    085        /**
    086         * Constructor for an iterator over the elements of a sliced dataset
    087         *
    088         * @param shape or dataShape
    089         * @param length of entire data array
    090         * @param start (if null then equivalent to the origin)
    091         * @param step (cannot contain zeros, if null then equivalent to ones)
    092         * @param sshape shape of new dataset, i.e. slice
    093         */
    094        public SliceIterator(final int[] shape, final int length, final int[] start, final int[] step, final int[] sshape) {
    095                this(shape, length, start, step, sshape, 1);
    096        }
    097
    098        /**
    099         * Constructor for an iterator over the elements of a sliced dataset
    100         *
    101         * @param shape or dataShape
    102         * @param length of entire data array
    103         * @param slice
    104         */
    105        public SliceIterator(final int[] shape, final int length, final SliceND slice) {
    106                this(shape, length, slice.getStart(), slice.getStep(), slice.getShape(), 1);
    107        }
    108
    109        /**
    110         * Constructor for an iterator over the elements of a sliced dataset
    111         *
    112         * @param shape or dataShape
    113         * @param length of entire data array
    114         * @param isize number of elements in an item
    115         * @param slice
    116         */
    117        public SliceIterator(final int[] shape, final int length, final int isize, final SliceND slice) {
    118                this(shape, length, slice.getStart(), slice.getStep(), slice.getShape(), isize);
    119        }
    120
    121        /**
    122         * Constructor for an iterator over the elements of a sliced dataset
    123         *
    124         * @param shape or dataShape
    125         * @param length of entire data array
    126         * @param start (if null then equivalent to the origin)
    127         * @param step (cannot contain zeros, if null then equivalent to ones)
    128         * @param sshape shape of new dataset, i.e. slice
    129         * @param isize number of elements in an item
    130         */
    131        public SliceIterator(final int[] shape, final int length, final int[] start, final int[] step, final int[] sshape, final int isize) {
    132                this.isize = isize;
    133                int rank = shape.length;
    134                endrank = rank - 1;
    135                this.shape = shape;
    136                this.start = new int[rank];
    137                this.sshape = sshape;
    138                if (step == null) {
    139                        this.step = new int[rank];
    140                        Arrays.fill(this.step, 1);
    141                } else {
    142                        this.step = step;
    143                }
    144
    145                if (rank == 0) {
    146                        istep = isize;
    147                        imax = length*isize;
    148                        stop = new int[0];
    149                        pos = new int[0];
    150                        gap = null;
    151                } else {
    152                        istep = this.step[endrank] * isize;
    153                        imax = length * isize;
    154                        stop = new int[rank];
    155                        gap = new int[endrank + 1];
    156                        pos = new int[rank];
    157                        calcGap();
    158                }
    159
    160                setStart(start);
    161        }
    162
    163        void calcGap() {
    164                int chunk = isize;
    165                for (int i = endrank; i >= 0; i--) {
    166                        stop[i] = start[i] + sshape[i] * step[i];
    167
    168                        if (step[i] < 0)
    169                                stop[i]++; // adjust for -ve steps so later code has more succinct test
    170
    171                        if (i > 0) {
    172                                gap[i] = (shape[i] * step[i - 1] - sshape[i] * step[i]) * chunk;
    173                                chunk *= shape[i];
    174                        }
    175                }
    176        }
    177
    178        /**
    179         * Set start (prefix with zeros if necessary)
    180         * @param newStart if null, then treat as origin
    181         */
    182        public void setStart(int... newStart) {
    183                final int rank = shape.length;
    184                if (rank == 0) {
    185                        index = -istep;
    186                        return;
    187                }
    188
    189                if (newStart == null) {
    190                        for (int i = 0; i < rank; i++) {
    191                                start[i] = 0;
    192                        }
    193                } else if (newStart.length > rank) {
    194                        throw new IllegalArgumentException("Length of start array greater than rank");
    195                } else {
    196                        int extra = rank - newStart.length;
    197                        for (int i = 0; i < extra; i++) {
    198                                start[i] = 0;
    199                        }
    200                        for (int i = 0; i < newStart.length; i++) {
    201                                start[i+extra] = newStart[i];
    202                        }
    203                }
    204
    205                reset();
    206                calcGap();
    207        }
    208
    209        @Override
    210        public void reset() {
    211                if (shape.length == 0) {
    212                        index = -istep;
    213                } else {
    214                        // work out index of first position
    215                        for (int i = 0; i < shape.length; i++)
    216                                pos[i] = start[i];
    217                        pos[endrank] -= step[endrank];
    218
    219                        index = pos[0];
    220                        for (int j = 1; j <= endrank; j++)
    221                                index = index * shape[j] + pos[j];
    222                        index *= isize;
    223                }
    224        }
    225
    226        @Override
    227        public boolean hasNext() {
    228                // now move on one position in slice
    229                int j = endrank;
    230                for (; j >= 0; j--) {
    231                        pos[j] += step[j];
    232
    233                        if ((pos[j] >= stop[j]) == (step[j] > 0)) { // stop index has been adjusted in code for -ve steps
    234                                pos[j] = start[j];
    235                                index += gap[j];
    236                        } else {
    237                                break;
    238                        }
    239                }
    240                if (j == -1 && endrank >= 0) {
    241                        index = imax;
    242                        return false;
    243                }
    244
    245                index += istep;
    246                return index < imax;
    247        }
    248
    249        public int[] getStart() {
    250                return start;
    251        }
    252
    253        @Override
    254        public int[] getPos() {
    255                return pos;
    256        }
    257
    258        public int[] getStep() {
    259                return step;
    260        }
    261
    262        @Override
    263        public int[] getShape() {
    264                return sshape;
    265        }
    266}