    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}