    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.calcSize(shape) == 0;
    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}