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}