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}