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
045        /**
046         * Constructor for an iterator over elements of a dataset that are within
047         * the shape
048         *
049         * @param shape
050         */
051        public PositionIterator(int[] shape) {
052                this(new SliceND(shape), null);
053        }
054
055        /**
056         * Constructor for an iterator over a single item broadcasted to given shape
057         *
058         * @param offset offset to single item
059         * @param shape
060         */
061        public PositionIterator(int offset, int[] shape) {
062                this(offset, new SliceND(shape), null);
063        }
064
065        /**
066         * Constructor for an iterator that misses out several axes
067         * @param shape
068         * @param axes missing axes, can be null for full dataset
069         */
070        public PositionIterator(int[] shape, int... axes) {
071                this(new SliceND(shape), axes);
072        }
073
074        /**
075         * Constructor for an iterator that misses out several axes
076         * @param shape
077         * @param slice
078         * @param axes missing axes
079         */
080        public PositionIterator(int[] shape, Slice[] slice, int[] axes) {
081                this(new SliceND(shape, slice), axes);
082        }
083
084        /**
085         * Constructor for an iterator that misses out several axes
086         * @param shape
087         * @param start
088         * @param stop
089         * @param step
090         * @param axes missing axes
091         */
092        public PositionIterator(int[] shape, int[] start, int[] stop, int[] step, int[] axes) {
093                this(new SliceND(shape, start, stop, step), axes);
094        }
095
096        /**
097         * Constructor for an iterator that misses out several axes
098         * @param slice
099         * @param axes missing axes
100         */
101        public PositionIterator(SliceND slice, int... axes) {
102                this(0, slice, axes);
103        }
104
105        /**
106         * Constructor for an iterator that misses out several axes
107         * 
108         * @param offset offset to start with
109         * @param slice
110         * @param axes missing axes
111         */
112        public PositionIterator(int offset, SliceND slice, int... axes) {
113                this.offset = offset;
114                int[] oshape = slice.getShape();
115                start = slice.getStart();
116                stop  = slice.getStop();
117                step  = slice.getStep();
118                for (int s : step) {
119                        if (s < 0) {
120                                throw new UnsupportedOperationException("Negative steps not implemented");
121                        }
122                }
123                int rank = oshape.length;
124                endrank = rank - 1;
125
126                omit = new boolean[rank];
127                shape = oshape.clone();
128                if (axes != null) {
129                        for (int a : axes) {
130                                if (a < 0) {
131                                        a += rank;
132                                }
133                                if (a >= 0 && a <= endrank) {
134                                        omit[a] = true;
135                                        shape[a] = 1;
136                                } else if (a > endrank) {
137                                        throw new IllegalArgumentException("Specified axis exceeds dataset rank");
138                                }
139                        }
140                }
141
142                pos = new int[rank];
143
144                reset();
145        }
146
147        @Override
148        public boolean hasNext() {
149                // now move on one position
150                if (once) {
151                        once = false;
152                        return true;
153                }
154                for (int j = endrank; j >= 0; j--) {
155                        if (omit[j]) {
156                                continue;
157                        }
158                        pos[j] += step[j];
159                        if (pos[j] >= stop[j]) {
160                                pos[j] = start[j];
161                        } else {
162                                return true;
163                        }
164                }
165                return false;
166        }
167
168        @Override
169        public int[] getPos() {
170                return pos;
171        }
172
173        /**
174         * @return omit array - array where true means miss out
175         */
176        public boolean[] getOmit() {
177                return omit;
178        }
179
180        @Override
181        public void reset() {
182                for (int i = 0; i <= endrank; i++)
183                        pos[i] = start[i];
184
185                int j = 0;
186                for (; j <= endrank; j++) {
187                        if (!omit[j])
188                                break;
189                }
190                if (j > endrank) {
191                        once = true;
192                        return;
193                }
194
195                if (omit[endrank]) {
196                        pos[endrank] = start[endrank];
197                        
198                        for (int i = endrank - 1; i >= 0; i--) {
199                                if (!omit[i]) {
200                                        pos[i] -= step[i];
201                                        break;
202                                }
203                        }
204                } else {
205                        pos[endrank] -= step[endrank];
206                }
207
208                index = offset;
209        }
210
211        @Override
212        public int[] getShape() {
213                return shape;
214        }
215
216        public int[] getStop() {
217                return stop;
218        }
219}