    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
    027        public StrideIterator(final int[] shape) {
    028                this(shape, null, 0);
    029        }
    030
    031        public StrideIterator(final int isize, final int[] shape) {
    032                this(isize, shape, null, 0);
    033        }
    034
    035        public StrideIterator(final int[] shape, final int[] strides) {
    036                this(shape, strides, 0);
    037        }
    038
    039        public StrideIterator(final int[] shape, final int[] strides, final int offset) {
    040                this(1, shape, strides, offset);
    041        }
    042
    043        public StrideIterator(final int isize, final int[] shape, final int[] strides, final int offset) {
    044                this(isize, shape, strides, offset, 0);
    045        }
    046
    047        public StrideIterator(final int isize, final int[] shape, final int[] strides, final int offset, final int element) {
    048                init(isize, shape, strides, offset, element);
    049                reset();
    050        }
    051
    052        public StrideIterator(final int isize, final int[] shape, final int[] start, final int[] stop, final int[] step) {
    053                this(isize, shape, null, 0, start, stop, step);
    054        }
    055
    056        public StrideIterator(final int isize, final int[] shape, final int[] oStrides, final int oOffset, final int[] start, final int[] stop, final int[] step) {
    057                this(isize, shape, oStrides, oOffset, new SliceND(shape, start, stop, step));
    058        }
    059
    060        public StrideIterator(final int isize, final int[] shape, final int[] oStrides, final int oOffset, final SliceND slice) {
    061                int rank = shape.length;
    062                int[] strides = new int[rank];
    063                int[] offset = new int[1];
    064                int[] newShape = AbstractDataset.createStrides(slice, isize, shape, oStrides, oOffset, strides, offset);
    065
    066                init(isize, newShape, strides, offset[0], 0);
    067                reset();
    068        }
    069
    070        private void init(final int isize, final int[] shape, final int[] strides, final int offset, final int element) {
    071                this.isize = isize;
    072                istep = isize;
    073                this.shape = shape;
    074                int rank = shape.length;
    075                endrank = rank - 1;
    076                pos = new int[rank];
    077                delta = new int[rank];
    078                this.element = element;
    079                if (strides != null) {
    080                        stride = strides;
    081                        for (int j = endrank; j >= 0; j--) {
    082                                delta[j] = stride[j] * shape[j];
    083                        }
    084                } else {
    085                        stride = new int[rank];
    086                        int s = isize;
    087                        for (int j = endrank; j >= 0; j--) {
    088                                stride[j] = s;
    089                                s *= shape[j];
    090                                delta[j] = s;
    091                        }
    092                        imax = s;
    093                }
    094                nstart = offset;
    095        }
    096
    097        @Override
    098        void calcGap() {
    099                // do nothing
    100        }
    101
    102        @Override
    103        public int[] getShape() {
    104                return shape;
    105        }
    106
    107        @Override
    108        public boolean hasNext() {
    109                // now move on one position
    110                int j = endrank;
    111                if (j < 0) {
    112                        index += istep;
    113                        return index < istep;
    114                }
    115                for (; j >= 0; j--) {
    116                        index += stride[j];
    117                        final int p = pos[j] + 1;
    118                        if (p < shape[j]) {
    119                                pos[j] = p;
    120                                break;
    121                        }
    122                        pos[j] = 0;
    123                        index -= delta[j]; // reset this dimension
    124                }
    125                return j >= 0;
    126        }
    127
    128        @Override
    129        public int[] getPos() {
    130                return pos;
    131        }
    132
    133        @Override
    134        public void reset() {
    135                Arrays.fill(pos, 0);
    136                if (endrank >= 0) {
    137                        pos[endrank] = -1;
    138                        index = nstart - stride[endrank];
    139                } else {
    140                        index = -istep;
    141                }
    142                index += element;
    143        }
    144}