    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}