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