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 to use
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 to use
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 to use
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 to use
078         * @param slice for each dimension
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 to use
088         * @param start
089         *            start point, may be {@code null}
090         * @param stop
091         *            stop point, may be {@code null}
092         * @param step
093         *            step, may be {@code null}
094         * @param axes missing axes
095         */
096        public PositionIterator(int[] shape, int[] start, int[] stop, int[] step, int[] axes) {
097                this(new SliceND(shape, start, stop, step), axes);
098        }
099
100        /**
101         * Constructor for an iterator that misses out several axes
102         * @param slice an n-D slice
103         * @param axes missing axes
104         */
105        public PositionIterator(SliceND slice, int... axes) {
106                this(0, slice, axes);
107        }
108
109        /**
110         * Constructor for an iterator that misses out several axes
111         * 
112         * @param offset offset to start with
113         * @param slice an n-D slice
114         * @param axes missing axes
115         */
116        public PositionIterator(int offset, SliceND slice, int... axes) {
117                this.offset = offset;
118                int[] oshape = slice.getShape();
119                start = slice.getStart();
120                stop  = slice.getStop();
121                step  = slice.getStep();
122                for (int s : step) {
123                        if (s < 0) {
124                                throw new UnsupportedOperationException("Negative steps not implemented");
125                        }
126                }
127                int rank;
128                if (oshape == null) {
129                        rank = 0;
130                        shape = null;
131                } else {
132                        rank = oshape.length;
133                        shape = oshape.clone();
134                }
135                endrank = rank - 1;
136
137                omit = new boolean[rank];
138                if (axes != null) {
139                        for (int a : axes) {
140                                a = ShapeUtils.checkAxis(rank, a);
141                                if (a >= 0 && a <= endrank) {
142                                        omit[a] = true;
143                                        shape[a] = 1;
144                                } else if (a > endrank) {
145                                        throw new IllegalArgumentException("Specified axis exceeds dataset rank");
146                                }
147                        }
148                }
149
150                pos = new int[rank];
151
152                zero = ShapeUtils.isZeroSize(shape);
153
154                reset();
155        }
156
157        @Override
158        public boolean hasNext() {
159                // now move on one position
160                if (zero) {
161                        return false;
162                }
163                if (once) {
164                        once = false;
165                        return true;
166                }
167                for (int j = endrank; j >= 0; j--) {
168                        if (omit[j]) {
169                                continue;
170                        }
171                        pos[j] += step[j];
172                        if (pos[j] >= stop[j]) {
173                                pos[j] = start[j];
174                        } else {
175                                return true;
176                        }
177                }
178                return false;
179        }
180
181        @Override
182        public int[] getPos() {
183                return pos;
184        }
185
186        /**
187         * @return omit array - array where true means miss out
188         */
189        public boolean[] getOmit() {
190                return omit;
191        }
192
193        @Override
194        public void reset() {
195                for (int i = 0; i <= endrank; i++) {
196                        pos[i] = start[i];
197                }
198                if (zero) {
199                        return;
200                }
201
202                int j = 0;
203                for (; j <= endrank; j++) {
204                        if (!omit[j]) {
205                                break;
206                        }
207                }
208                if (j > endrank) {
209                        once = true;
210                        return;
211                }
212
213                if (omit[endrank]) {
214                        pos[endrank] = start[endrank];
215                        
216                        for (int i = endrank - 1; i >= 0; i--) {
217                                if (!omit[i]) {
218                                        pos[i] -= step[i];
219                                        break;
220                                }
221                        }
222                } else {
223                        pos[endrank] -= step[endrank];
224                }
225
226                index = offset;
227        }
228
229        @Override
230        public int[] getShape() {
231                return shape;
232        }
233
234        public int[] getStop() {
235                return stop;
236        }
237}