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 an n-D 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 = sShape == null ? null : 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                final int rank = shape == null ? 0 : 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                        int[] dShape = dSlice.getShape();
097                        for (int i = 0; i < rank; i++) {
098                                dShape[i] = 1;
099                        }
100                } else {
101                        int[] dShape = dSlice.getShape();
102                        int[] oShape = new int[sRank]; // outer shape
103                        int[] iShape = new int[rank - sRank]; // inner shape
104                        for (int i = 0, j = 0, k = 0; i < rank; i++) {
105                                if (omit[i]) {
106                                        iShape[j++] = sShape[i];
107                                } else {
108                                        oShape[k++] = sShape[i];
109                                        dShape[i] = 1;
110                                }
111                        }
112                        sSlice = new SliceND(oShape);
113                        sStart = sSlice.getStart();
114                        sStop = sSlice.getStop();
115
116                        iSlice = new SliceND(iShape);
117                        for (int i = 0, j = 0, k = 0; i < rank; i++) {
118                                if (omit[i]) {
119                                        iSlice.setSlice(j++, start[i], stop[i], step[i]);
120                                } else {
121                                        sSlice.setSlice(k++, start[i], stop[i], step[i]);
122                                }
123                        }
124                }
125                
126                reset();
127        }
128
129        @Override
130        public boolean hasNext() {
131                // now move on one position
132                if (once) {
133                        once = false;
134                        return true;
135                }
136                int k = sRank - 1;
137                for (int j = endrank; j >= 0; j--) {
138                        if (omit[j]) {
139                                continue;
140                        }
141                        pos[j] += step[j];
142                        end[j] = pos[j] + step[j];
143                        dStart[j]++;
144                        dStop[j]++;
145                        if (pos[j] >= stop[j]) {
146                                pos[j] = start[j];
147                                end[j] = pos[j] + step[j];
148                                dStart[j] = 0;
149                                dStop[j] = 1;
150                                if (sStop != null) {
151                                        sStart[k] = pos[j];
152                                        sStop[k] = end[j];
153                                        k--;
154                                }
155                        } else {
156                                if (sStop != null) {
157                                        sStart[k] = pos[j];
158                                        sStop[k] = end[j];
159                                        k--;
160                                }
161                                return true;
162                        }
163                }
164                return false;
165        }
166
167        @Override
168        public int[] getPos() {
169                return pos;
170        }
171
172        /**
173         * Get omitted part of source slice which never changes. I.e. the inner slice
174         * @return slice (can be null)
175         */
176        public SliceND getOmittedSlice() {
177                return iSlice;
178        }
179
180        /**
181         * Get output or destination slice
182         * @return slice 
183         */
184        public SliceND getOutputSlice() {
185                return dSlice;
186        }
187
188        /**
189         * Get current slice
190         * @return slice
191         */
192        public SliceND getCurrentSlice() {
193                return cSlice;
194        }
195
196        /**
197         * Shortened position where axes are omitted
198         * @return used position
199         */
200        public int[] getUsedPos() {
201                return sStart;
202        }
203
204        /**
205         * Shortened slice where axes are omitted
206         * @return used (or outer) slice
207         */
208        public SliceND getUsedSlice() {
209                return sSlice;
210        }
211
212        /**
213         * @return omit array - array where true means miss out
214         */
215        public boolean[] getOmit() {
216                return omit;
217        }
218
219        @Override
220        public void reset() {
221                once = false;
222                for (int i = 0, k = 0; i <= endrank; i++) {
223                        int b = start[i];
224                        int d = step[i];
225                        if (!omit[i]) {
226                                cSlice.setSlice(i, b, b + d, d);
227                                dStart[i] = 0;
228                                dStop[i] = 1;
229                                if (sStop != null) {
230                                        sSlice.setSlice(k++, b, b + d, d);
231                                }
232                        } else {
233                                cSlice.setSlice(i, b, end[i], d);
234                        }
235                }
236
237                int j = 0;
238                for (; j <= endrank; j++) {
239                        if (!omit[j])
240                                break;
241                }
242                if (j > endrank) {
243                        once = shape != null;
244                        return;
245                }
246
247                if (omit[endrank]) {
248                        pos[endrank] = start[endrank];
249                        for (int i = endrank - 1; i >= 0; i--) {
250                                if (!omit[i]) {
251                                        end[i] = pos[i];
252                                        pos[i] -= step[i];
253                                        dStart[i]--;
254                                        dStop[i]--;
255                                        break;
256                                }
257                        }
258                } else {
259                        end[endrank] = pos[endrank];
260                        pos[endrank] -= step[endrank];
261                        dStart[endrank]--;
262                        dStop[endrank]--;
263                }
264
265                if (sStart != pos) {
266                        for (int i = 0, k = 0; i <= endrank; i++) {
267                                if (!omit[i]) {
268                                        sStart[k++] = pos[i];
269                                }
270                        }
271                }
272        }
273
274        @Override
275        public int[] getShape() {
276                return shape;
277        }
278}