    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                int rank = 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                int rank = shape.length;
    083                endrank = rank - 1;
    084                pos = new int[rank];
    085                delta = new int[rank];
    086                this.element = element;
    087                if (strides == null) {
    088                        offset = 0;
    089                        stride = new int[rank];
    090                        int s = isize;
    091                        for (int j = endrank; j >= 0; j--) {
    092                                stride[j] = s;
    093                                s *= shape[j];
    094                        }
    095                } else {
    096                        stride = strides.clone();
    097                }
    098
    099                for (int j = endrank; j >= 0; j--) {
    100                        int t = stride[j];
    101                        offset += t * start[j];
    102                        t *= step[j];
    103                        stride[j] = t;
    104                        int s = sshape[j];
    105                        if (!zero) {
    106                                zero = s == 0;
    107                        }
    108                        delta[j] = s * t;
    109                }
    110
    111                nstart = offset;
    112        }
    113
    114        @Override
    115        void calcGap() {
    116                // do nothing
    117        }
    118
    119        @Override
    120        public boolean hasNext() {
    121                if (zero) {
    122                        return false;
    123                }
    124
    125                // now move on one position
    126                int j = endrank;
    127                if (j < 0) {
    128                        index += istep;
    129                        return index < istep;
    130                }
    131                for (; j >= 0; j--) {
    132                        index += stride[j];
    133                        final int s = step[j];
    134                        final int p = pos[j] + s;
    135                        if ((s > 0 && p < stop[j]) || (s < 0 && p > stop[j])) {
    136                                pos[j] = p;
    137                                break;
    138                        }
    139                        pos[j] = start[j];
    140                        index -= delta[j]; // reset this dimension
    141                }
    142                return j >= 0;
    143        }
    144
    145        @Override
    146        public int[] getPos() {
    147                return pos;
    148        }
    149
    150        @Override
    151        public void reset() {
    152                System.arraycopy(start, 0, pos, 0, start.length);
    153                if (endrank >= 0) {
    154                        pos[endrank] -= step[endrank];
    155                        index = nstart - stride[endrank];
    156                } else {
    157                        index = -istep;
    158                }
    159                index += element;
    160        }
    161}