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
015
016import java.util.Arrays;
017
018/**
019 * Class to run over a slice of a dataset
020 */
021public class SliceIterator extends IndexIterator {
022        int[] shape;
023        int isize;
024        int endrank; // last shape index
025        int[] gap; // gaps in dataset
026        int imax; // maximum index in array
027        int[] start;
028        int[] stop;
029        int[] step;
030        int[] sshape; // slice shape
031        int[] pos; // position in dataset
032        int istep; // step in last index
033
034        SliceIterator() {
035        }
036
037        /**
038         * Constructor for an iterator over the elements of a sliced dataset
039         * 
040         * @param shape or dataShape
041         * @param length of entire data array
042         * @param sshape shape of new dataset, i.e. slice
043         */
044        public SliceIterator(final int[] shape, final int length, final int[] sshape) {
045                this(shape, length, null, null, sshape, 1);
046        }
047
048        /**
049         * Constructor for an iterator over the elements of a sliced dataset
050         * 
051         * @param shape or dataShape
052         * @param length of entire data array
053         * @param start (if null then equivalent to the origin)
054         * @param sshape shape of new dataset, i.e. slice
055         */
056        public SliceIterator(final int[] shape, final int length, final int[] start, final int[] sshape) {
057                this(shape, length, start, null, sshape, 1);
058        }
059
060        /**
061         * Constructor for an iterator over the elements of a sliced dataset
062         * 
063         * @param shape or dataShape
064         * @param length of entire data array
065         * @param sshape shape of new dataset, i.e. slice
066         * @param isize number of elements in an item
067         */
068        public SliceIterator(final int[] shape, final int length, final int[] sshape, final int isize) {
069                this(shape, length, null, null, sshape, isize);
070        }
071
072        /**
073         * Constructor for an iterator over the elements of a sliced dataset
074         * 
075         * @param shape or dataShape
076         * @param length of entire data array
077         * @param start (if null then equivalent to the origin)
078         * @param sshape shape of new dataset, i.e. slice
079         * @param isize number of elements in an item
080         */
081        public SliceIterator(final int[] shape, final int length, final int[] start, final int[] sshape, final int isize) {
082                this(shape, length, start, null, sshape, isize);
083        }
084
085        /**
086         * Constructor for an iterator over the elements of a sliced dataset
087         * 
088         * @param shape or dataShape
089         * @param length of entire data array
090         * @param start (if null then equivalent to the origin)
091         * @param step (cannot contain zeros, if null then equivalent to ones)
092         * @param sshape shape of new dataset, i.e. slice
093         */
094        public SliceIterator(final int[] shape, final int length, final int[] start, final int[] step, final int[] sshape) {
095                this(shape, length, start, step, sshape, 1);
096        }
097
098        /**
099         * Constructor for an iterator over the elements of a sliced dataset
100         * 
101         * @param shape or dataShape
102         * @param length of entire data array
103         * @param slice
104         */
105        public SliceIterator(final int[] shape, final int length, final SliceND slice) {
106                this(shape, length, slice.getStart(), slice.getStep(), slice.getShape(), 1);
107        }
108
109        /**
110         * Constructor for an iterator over the elements of a sliced dataset
111         * 
112         * @param shape or dataShape
113         * @param length of entire data array
114         * @param isize number of elements in an item
115         * @param slice
116         */
117        public SliceIterator(final int[] shape, final int length, final int isize, final SliceND slice) {
118                this(shape, length, slice.getStart(), slice.getStep(), slice.getShape(), isize);
119        }
120
121        /**
122         * Constructor for an iterator over the elements of a sliced dataset
123         * 
124         * @param shape or dataShape
125         * @param length of entire data array
126         * @param start (if null then equivalent to the origin)
127         * @param step (cannot contain zeros, if null then equivalent to ones)
128         * @param sshape shape of new dataset, i.e. slice
129         * @param isize number of elements in an item
130         */
131        public SliceIterator(final int[] shape, final int length, final int[] start, final int[] step, final int[] sshape, final int isize) {
132                this.isize = isize;
133                int rank = shape.length;
134                endrank = rank - 1;
135                this.shape = shape;
136                this.start = new int[rank];
137                this.sshape = sshape;
138                if (step == null) {
139                        this.step = new int[rank];
140                        Arrays.fill(this.step, 1);
141                } else {
142                        this.step = step;
143                }
144
145                if (rank == 0) {
146                        istep = isize;
147                        imax = length*isize;
148                        stop = new int[0];
149                        pos = new int[0];
150                        gap = null;
151                } else {
152                        istep = this.step[endrank] * isize;
153                        imax = length * isize;
154                        stop = new int[rank];
155                        gap = new int[endrank + 1];
156                        pos = new int[rank];
157                        calcGap();
158                }
159
160                setStart(start);
161        }
162
163        void calcGap() {
164                int chunk = isize;
165                for (int i = endrank; i >= 0; i--) {
166                        stop[i] = start[i] + sshape[i] * step[i];
167        
168                        if (step[i] < 0)
169                                stop[i]++; // adjust for -ve steps so later code has more succinct test
170        
171                        if (i > 0) {
172                                gap[i] = (shape[i] * step[i - 1] - sshape[i] * step[i]) * chunk;
173                                chunk *= shape[i];
174                        }
175                }
176        }
177
178        /**
179         * Set start (prefix with zeros if necessary)
180         * @param newStart if null, then treat as origin
181         */
182        public void setStart(int... newStart) {
183                final int rank = shape.length;
184                if (rank == 0) {
185                        index = -istep;
186                        return;
187                }
188
189                if (newStart == null) {
190                        for (int i = 0; i < rank; i++) {
191                                start[i] = 0;
192                        }
193                } else if (newStart.length > rank) {
194                        throw new IllegalArgumentException("Length of start array greater than rank");
195                } else {
196                        int extra = rank - newStart.length;
197                        for (int i = 0; i < extra; i++) {
198                                start[i] = 0;
199                        }
200                        for (int i = 0; i < newStart.length; i++) {
201                                start[i+extra] = newStart[i];
202                        }
203                }
204
205                reset();
206                calcGap();
207        }
208
209        @Override
210        public void reset() {
211                if (shape.length == 0) {
212                        index = -istep;
213                } else {
214                        // work out index of first position
215                        for (int i = 0; i < shape.length; i++)
216                                pos[i] = start[i];
217                        pos[endrank] -= step[endrank];
218        
219                        index = pos[0];
220                        for (int j = 1; j <= endrank; j++)
221                                index = index * shape[j] + pos[j];
222                        index *= isize;
223                }
224        }
225
226        @Override
227        public boolean hasNext() {
228                // now move on one position in slice
229                int j = endrank;
230                for (; j >= 0; j--) {
231                        pos[j] += step[j];
232
233                        if ((pos[j] >= stop[j]) == (step[j] > 0)) { // stop index has been adjusted in code for -ve steps
234                                pos[j] = start[j];
235                                index += gap[j];
236                        } else {
237                                break;
238                        }
239                }
240                if (j == -1 && endrank >= 0) {
241                        index = imax;
242                        return false;
243                }
244
245                index += istep;
246                return index < imax;
247        }
248
249        public int[] getStart() {
250                return start;
251        }
252
253        @Override
254        public int[] getPos() {
255                return pos;
256        }
257
258        public int[] getStep() {
259                return step;
260        }
261
262        @Override
263        public int[] getShape() {
264                return sshape;
265        }
266}