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}