001/*-
002 *******************************************************************************
003 * Copyright (c) 2015, 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/**
016 * Class to provide slice iteration through a dataset
017 * <p>It allows a number of axes to be omitted and iterates over
018 * the axes left over. The omitted axes define an inner shape and
019 * the remaining axes define an outer shape which is iterated over.
020 */
021public class SliceNDIterator extends IndexIterator {
022        final private int[] shape;
023        final private int[] start;
024        final private int[] stop;
025        final private int[] step;
026        final private int endrank;
027
028        final private boolean[] omit; // axes to miss out
029
030        /**
031         * position in source dataset
032         */
033        final private int[] pos;
034        final private int[] end;
035        private boolean once;
036
037        private SliceND cSlice; // current slice
038        
039        private int sRank; // number of dimensions used (i.e. not missing)
040        final private SliceND iSlice; // omitted source or inner slice
041
042        final private SliceND sSlice; // shortened slice
043        final private int[] sStart; // shortened position
044        final private int[] sStop; // shortened end
045
046        private SliceND dSlice; // destination slice
047        private int[] dStart;
048        private int[] dStop;
049
050        /**
051         * Constructor for an iterator that misses out several axes
052         * @param slice
053         * @param axes missing axes
054         */
055        public SliceNDIterator(SliceND slice, int... axes) {
056                cSlice = slice.clone();
057                int[] sShape = cSlice.getSourceShape();
058                shape = cSlice.getShape().clone();
059                start = cSlice.getStart();
060                stop  = cSlice.getStop();
061                step  = cSlice.getStep();
062                for (int s : step) {
063                        if (s < 0) {
064                                throw new UnsupportedOperationException("Negative steps not implemented");
065                        }
066                }
067                int rank = shape.length;
068                endrank = rank - 1;
069
070                omit = new boolean[rank];
071                dSlice = new SliceND(shape);
072                dStart = dSlice.getStart();
073                dStop  = dSlice.getStop();
074                sRank = rank;
075                if (axes != null) {
076                        axes = ShapeUtils.checkAxes(rank, axes);
077                        for (int a : axes) {
078                                if (a >= 0 && a <= endrank) {
079                                        sRank--;
080                                        omit[a] = true;
081                                        shape[a] = 1;
082                                } else if (a > endrank) {
083                                        throw new IllegalArgumentException("Specified axis exceeds dataset rank");
084                                }
085                        }
086                }
087
088                cSlice = cSlice.clone();
089                pos = cSlice.getStart();
090                end = cSlice.getStop();
091                if (sRank == rank) {
092                        sStart = pos;
093                        sStop = null;
094                        iSlice = null;
095                        sSlice = cSlice;
096                } else {
097                        int[] dShape = dSlice.getShape();
098                        int[] oShape = new int[sRank]; // outer shape
099                        int[] iShape = new int[rank - sRank]; // inner shape
100                        for (int i = 0, j = 0, k = 0; i < rank; i++) {
101                                if (omit[i]) {
102                                        iShape[j++] = sShape[i];
103                                } else {
104                                        oShape[k++] = sShape[i];
105                                        dShape[i] = 1;
106                                }
107                        }
108                        sSlice = new SliceND(oShape);
109                        sStart = sSlice.getStart();
110                        sStop = sSlice.getStop();
111
112                        iSlice = new SliceND(iShape);
113                        for (int i = 0, j = 0, k = 0; i < rank; i++) {
114                                if (omit[i]) {
115                                        iSlice.setSlice(j++, start[i], stop[i], step[i]);
116                                } else {
117                                        sSlice.setSlice(k++, start[i], stop[i], step[i]);
118                                }
119                        }
120                }
121                
122                reset();
123        }
124
125        @Override
126        public boolean hasNext() {
127                // now move on one position
128                if (once) {
129                        once = false;
130                        return true;
131                }
132                int k = sRank - 1;
133                for (int j = endrank; j >= 0; j--) {
134                        if (omit[j]) {
135                                continue;
136                        }
137                        pos[j] += step[j];
138                        end[j] = pos[j] + step[j];
139                        dStart[j]++;
140                        dStop[j]++;
141                        if (pos[j] >= stop[j]) {
142                                pos[j] = start[j];
143                                end[j] = pos[j] + step[j];
144                                dStart[j] = 0;
145                                dStop[j] = 1;
146                                if (sStop != null) {
147                                        sStart[k] = pos[j];
148                                        sStop[k] = end[j];
149                                        k--;
150                                }
151                        } else {
152                                if (sStop != null) {
153                                        sStart[k] = pos[j];
154                                        sStop[k] = end[j];
155                                        k--;
156                                }
157                                return true;
158                        }
159                }
160                return false;
161        }
162
163        @Override
164        public int[] getPos() {
165                return pos;
166        }
167
168        /**
169         * Get omitted part of source slice which never changes. I.e. the inner slice
170         * @return slice (can be null)
171         */
172        public SliceND getOmittedSlice() {
173                return iSlice;
174        }
175
176        /**
177         * Get output or destination slice
178         * @return slice 
179         */
180        public SliceND getOutputSlice() {
181                return dSlice;
182        }
183
184        /**
185         * Get current slice
186         * @return slice
187         */
188        public SliceND getCurrentSlice() {
189                return cSlice;
190        }
191
192        /**
193         * Shortened position where axes are omitted
194         * @return used position
195         */
196        public int[] getUsedPos() {
197                return sStart;
198        }
199
200        /**
201         * Shortened slice where axes are omitted
202         * @return used (or outer) slice
203         */
204        public SliceND getUsedSlice() {
205                return sSlice;
206        }
207
208        /**
209         * @return omit array - array where true means miss out
210         */
211        public boolean[] getOmit() {
212                return omit;
213        }
214
215        @Override
216        public void reset() {
217                for (int i = 0, k = 0; i <= endrank; i++) {
218                        int b = start[i];
219                        int d = step[i];
220                        if (!omit[i]) {
221                                cSlice.setSlice(i, b, b + d, d);
222                                dStart[i] = 0;
223                                dStop[i] = 1;
224                                if (sStop != null) {
225                                        sSlice.setSlice(k++, b, b + d, d);
226                                }
227                        } else {
228                                cSlice.setSlice(i, b, end[i], d);
229                        }
230                }
231
232                int j = 0;
233                for (; j <= endrank; j++) {
234                        if (!omit[j])
235                                break;
236                }
237                if (j > endrank) {
238                        once = true;
239                        return;
240                }
241
242                if (omit[endrank]) {
243                        pos[endrank] = start[endrank];
244                        for (int i = endrank - 1; i >= 0; i--) {
245                                if (!omit[i]) {
246                                        end[i] = pos[i];
247                                        pos[i] -= step[i];
248                                        dStart[i]--;
249                                        dStop[i]--;
250                                        break;
251                                }
252                        }
253                } else {
254                        end[endrank] = pos[endrank];
255                        pos[endrank] -= step[endrank];
256                        dStart[endrank]--;
257                        dStop[endrank]--;
258                }
259
260                if (sStart != pos) {
261                        for (int i = 0, k = 0; i <= endrank; i++) {
262                                if (!omit[i]) {
263                                        sStart[k++] = pos[i];
264                                }
265                        }
266                }
267        }
268
269        @Override
270        public int[] getShape() {
271                return shape;
272        }
273}