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/**
016 * <p>Class to provide iteration through a dataset</p>
017 * <p>Instantiate an iterator and use it in a while loop:
018 * <pre>
019 * Dataset ds = DatasetFactory.createLinearSpace(DoubleDataset.class, 0, 10, 0.25);
020 * PositionIterator iter = ds.getPositionIterator();
021 * int[] pos = iter.getPos()
022 *
023 * while (iter.hasNext()) {
024 *     ds.set(1.2, pos);
025 * }
026 * </pre>
027 *
028 */
029public class PositionIterator extends IndexIterator {
030        private int offset;
031        final private int[] shape;
032        final private int[] start;
033        final private int[] stop;
034        final private int[] step;
035        final private int endrank;
036
037        final private boolean[] omit; // axes to miss out
038
039        /**
040         * position in dataset
041         */
042        final private int[] pos;
043        private boolean once;
044        final private boolean zero; // no iterations allowed
045
046        /**
047         * Constructor for an iterator over elements of a dataset that are within
048         * the shape
049         *
050         * @param shape
051         */
052        public PositionIterator(int[] shape) {
053                this(new SliceND(shape), null);
054        }
055
056        /**
057         * Constructor for an iterator over a single item broadcasted to given shape
058         *
059         * @param offset offset to single item
060         * @param shape
061         */
062        public PositionIterator(int offset, int[] shape) {
063                this(offset, new SliceND(shape), null);
064        }
065
066        /**
067         * Constructor for an iterator that misses out several axes
068         * @param shape
069         * @param axes missing axes, can be null for full dataset
070         */
071        public PositionIterator(int[] shape, int... axes) {
072                this(new SliceND(shape), axes);
073        }
074
075        /**
076         * Constructor for an iterator that misses out several axes
077         * @param shape
078         * @param slice
079         * @param axes missing axes
080         */
081        public PositionIterator(int[] shape, Slice[] slice, int[] axes) {
082                this(new SliceND(shape, slice), axes);
083        }
084
085        /**
086         * Constructor for an iterator that misses out several axes
087         * @param shape
088         * @param start
089         * @param stop
090         * @param step
091         * @param axes missing axes
092         */
093        public PositionIterator(int[] shape, int[] start, int[] stop, int[] step, int[] axes) {
094                this(new SliceND(shape, start, stop, step), axes);
095        }
096
097        /**
098         * Constructor for an iterator that misses out several axes
099         * @param slice
100         * @param axes missing axes
101         */
102        public PositionIterator(SliceND slice, int... axes) {
103                this(0, slice, axes);
104        }
105
106        /**
107         * Constructor for an iterator that misses out several axes
108         * 
109         * @param offset offset to start with
110         * @param slice
111         * @param axes missing axes
112         */
113        public PositionIterator(int offset, SliceND slice, int... axes) {
114                this.offset = offset;
115                int[] oshape = slice.getShape();
116                start = slice.getStart();
117                stop  = slice.getStop();
118                step  = slice.getStep();
119                for (int s : step) {
120                        if (s < 0) {
121                                throw new UnsupportedOperationException("Negative steps not implemented");
122                        }
123                }
124                int rank = oshape.length;
125                endrank = rank - 1;
126
127                omit = new boolean[rank];
128                shape = oshape.clone();
129                if (axes != null) {
130                        for (int a : axes) {
131                                a = ShapeUtils.checkAxis(rank, a);
132                                if (a >= 0 && a <= endrank) {
133                                        omit[a] = true;
134                                        shape[a] = 1;
135                                } else if (a > endrank) {
136                                        throw new IllegalArgumentException("Specified axis exceeds dataset rank");
137                                }
138                        }
139                }
140
141                pos = new int[rank];
142
143                zero = ShapeUtils.calcSize(shape) == 0;
144
145                reset();
146        }
147
148        @Override
149        public boolean hasNext() {
150                // now move on one position
151                if (zero) {
152                        return false;
153                }
154                if (once) {
155                        once = false;
156                        return true;
157                }
158                for (int j = endrank; j >= 0; j--) {
159                        if (omit[j]) {
160                                continue;
161                        }
162                        pos[j] += step[j];
163                        if (pos[j] >= stop[j]) {
164                                pos[j] = start[j];
165                        } else {
166                                return true;
167                        }
168                }
169                return false;
170        }
171
172        @Override
173        public int[] getPos() {
174                return pos;
175        }
176
177        /**
178         * @return omit array - array where true means miss out
179         */
180        public boolean[] getOmit() {
181                return omit;
182        }
183
184        @Override
185        public void reset() {
186                for (int i = 0; i <= endrank; i++) {
187                        pos[i] = start[i];
188                }
189                if (zero) {
190                        return;
191                }
192
193                int j = 0;
194                for (; j <= endrank; j++) {
195                        if (!omit[j]) {
196                                break;
197                        }
198                }
199                if (j > endrank) {
200                        once = true;
201                        return;
202                }
203
204                if (omit[endrank]) {
205                        pos[endrank] = start[endrank];
206                        
207                        for (int i = endrank - 1; i >= 0; i--) {
208                                if (!omit[i]) {
209                                        pos[i] -= step[i];
210                                        break;
211                                }
212                        }
213                } else {
214                        pos[endrank] -= step[endrank];
215                }
216
217                index = offset;
218        }
219
220        @Override
221        public int[] getShape() {
222                return shape;
223        }
224
225        public int[] getStop() {
226                return stop;
227        }
228}