    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
    015import java.util.Arrays;
    016
    017
    018/**
    019 * Class to run over a contiguous dataset using strides
    020 */
    021public class StrideIterator extends SliceIterator {
    022        private int[] stride;
    023        private int[] delta;  // reset values
    024        private int nstart;
    025        private int element;
    026        private boolean zero;
    027
    028        public StrideIterator(final int[] shape) {
    029                this(shape, null, 0);
    030        }
    031
    032        public StrideIterator(final int isize, final int[] shape) {
    033                this(isize, shape, null, 0);
    034        }
    035
    036        public StrideIterator(final int[] shape, final int[] strides) {
    037                this(shape, strides, 0);
    038        }
    039
    040        public StrideIterator(final int[] shape, final int[] strides, final int offset) {
    041                this(1, shape, strides, offset);
    042        }
    043
    044        public StrideIterator(final int isize, final int[] shape, final int[] strides, final int offset) {
    045                this(isize, shape, strides, offset, 0);
    046        }
    047
    048        public StrideIterator(final int isize, final int[] shape, final int[] strides, final int offset, final int element) {
    049                final int rank = shape == null ? 0 : shape.length;
    050                start = new int[rank];
    051                stop = shape;
    052                step = new int[rank];
    053                Arrays.fill(step, 1);
    054                this.sshape = shape;
    055
    056                init(isize, shape, strides, offset, element);
    057                reset();
    058        }
    059
    060        public StrideIterator(final int isize, final int[] shape, final int[] start, final int[] stop, final int[] step) {
    061                this(isize, shape, null, 0, start, stop, step);
    062        }
    063
    064        public StrideIterator(final int isize, final int[] shape, final int[] oStrides, final int oOffset, final int[] start, final int[] stop, final int[] step) {
    065                this(isize, shape, oStrides, oOffset, new SliceND(shape, start, stop, step));
    066        }
    067
    068        public StrideIterator(final int isize, final int[] shape, final int[] strides, final int offset, final SliceND slice) {
    069                start = slice.getStart();
    070                stop = slice.getStop();
    071                step = slice.getStep();
    072                this.sshape = slice.getShape();
    073
    074                init(isize, shape, strides, offset, 0);
    075                reset();
    076        }
    077
    078        private void init(final int isize, final int[] shape, final int[] strides, int offset, final int element) {
    079                this.isize = isize;
    080                istep = isize;
    081                this.shape = shape;
    082                final int rank = shape == null ? 0 : shape.length;
    083                zero = shape == null;
    084                endrank = rank - 1;
    085                pos = new int[rank];
    086                delta = new int[rank];
    087                this.element = element;
    088                if (strides == null) {
    089                        offset = 0;
    090                        stride = new int[rank];
    091                        int s = isize;
    092                        for (int j = endrank; j >= 0; j--) {
    093                                stride[j] = s;
    094                                s *= shape[j];
    095                        }
    096                } else {
    097                        stride = strides.clone();
    098                }
    099
    100                for (int j = endrank; j >= 0; j--) {
    101                        int t = stride[j];
    102                        offset += t * start[j];
    103                        t *= step[j];
    104                        stride[j] = t;
    105                        int s = sshape[j];
    106                        if (!zero) {
    107                                zero = s == 0;
    108                        }
    109                        delta[j] = s * t;
    110                }
    111
    112                nstart = offset;
    113        }
    114
    115        @Override
    116        void calcGap() {
    117                // do nothing
    118        }
    119
    120        @Override
    121        public boolean hasNext() {
    122                if (zero) {
    123                        return false;
    124                }
    125
    126                // now move on one position
    127                int j = endrank;
    128                if (j < 0) {
    129                        index += istep;
    130                        return index < istep;
    131                }
    132                for (; j >= 0; j--) {
    133                        index += stride[j];
    134                        final int s = step[j];
    135                        final int p = pos[j] + s;
    136                        if ((s > 0 && p < stop[j]) || (s < 0 && p > stop[j])) {
    137                                pos[j] = p;
    138                                break;
    139                        }
    140                        pos[j] = start[j];
    141                        index -= delta[j]; // reset this dimension
    142                }
    143                return j >= 0;
    144        }
    145
    146        @Override
    147        public int[] getPos() {
    148                return pos;
    149        }
    150
    151        @Override
    152        public void reset() {
    153                System.arraycopy(start, 0, pos, 0, start.length);
    154                if (endrank >= 0) {
    155                        pos[endrank] -= step[endrank];
    156                        index = nstart - stride[endrank];
    157                } else {
    158                        index = -istep;
    159                }
    160                index += element;
    161        }
    162}