    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
    015import java.util.ArrayList;
    016import java.util.Arrays;
    017import java.util.List;
    018
    019/**
    020 * Class to run over an array of integer datasets and return its items
    021 */
    022public class IntegersIterator extends IndexIterator {
    023        final private int[] ishape; // shape of input
    024        final private int irank; // rank of input shape
    025        final private int[] oshape; // shape of output
    026        final private int orank; // rank of output shape
    027        private int offset; // offset of index subspace in new position
    028        private int srank; // rank of subspace
    029
    030        final private IndexIterator it;
    031
    032        /**
    033         * position in input shape
    034         */
    035        final private int[] ipos;
    036        /**
    037         * position in output shape
    038         */
    039        final private int[] opos;
    040        final private List<Object> indexes;
    041
    042        /**
    043         * Constructor for an iterator over the items of an array of objects
    044         * @param shape of entire data array
    045         * @param index an array of integer dataset, boolean dataset, slices or null entries (same as full slices)
    046         */
    047        public IntegersIterator(final int[] shape, final Object... index) {
    048                this(false, shape, index);
    049        }
    050
    051        /**
    052         * Constructor for an iterator over the items of an array of objects
    053         * @param restrict1D if true, allow only one 1D integer datasets otherwise they must match shape
    054         * @param shape of entire data array
    055         * @param index an array of integer dataset, boolean dataset, slices or null entries (same as full slices)
    056         */
    057        public IntegersIterator(final boolean restrict1D, final int[] shape, final Object... index) {
    058                if (shape == null) {
    059                        ishape = null;
    060                        irank = 0;
    061                }  else {
    062                        ishape = shape.clone();
    063                        irank = shape.length;
    064                }
    065                if (irank < index.length) {
    066                        throw new IllegalArgumentException("Number of index datasets is greater than rank of dataset");
    067                }
    068                indexes = new ArrayList<Object>();
    069                for (Object i : index) {
    070                        if (i instanceof BooleanDataset) { // turn boolean datasets into integer ones
    071                                for (IntegerDataset id : Comparisons.nonZero((Dataset) i)) {
    072                                        indexes.add(id);
    073                                }
    074                        } else if (i == null || i instanceof Slice) {
    075                                indexes.add(i);
    076                        } else if (i instanceof IntegerDataset) {
    077                                Dataset id = (Dataset) i;
    078                                int r = id.getRank();
    079                                if (restrict1D && r > 1) {
    080                                        throw new IllegalArgumentException("Integer datasets were restricted to zero or one dimensions");
    081                                }
    082                                if (r == 0) { // avoid zero-rank datasets
    083                                        i = id.reshape(1);
    084                                }
    085                                indexes.add(i);
    086                        } else {
    087                                throw new IllegalArgumentException("Unsupported object for indexing");
    088                        }
    089                }
    090                if (indexes.size() < irank) { // pad out index list
    091                        for (int i = indexes.size(); i < irank; i++) {
    092                                indexes.add(null);
    093                        }
    094                } else if (indexes.size() > irank) {
    095                        throw new IllegalArgumentException("Too many indices (a boolean dataset may have too many dimensions)");
    096                }
    097
    098                int ilength = -1;
    099                int[] cshape = null;
    100                int first = -1; // index of first null or slice after non-null index
    101                boolean intact = true;
    102                srank = 0;
    103                for (int i = 0; i < irank; i++) { // see if shapes are consistent and subspace is intact
    104                        Object obj = indexes.get(i);
    105                        if (obj instanceof IntegerDataset && !restrict1D) {
    106                                IntegerDataset ind = (IntegerDataset) obj;
    107                                if (first > 0) {
    108                                        intact = false;
    109                                }
    110
    111                                int l = ind.size;
    112                                if (ilength < l) {
    113                                        ilength = l;
    114                                        cshape = null;
    115                                } else if (l != 1 && l != ilength) {
    116                                        throw new IllegalArgumentException("Index datasets do not have same size");
    117                                }
    118                                if (cshape == null) {
    119                                        cshape = ind.shape;
    120                                        srank = cshape.length;
    121                                        offset = i;
    122                                } else if (l > 1 && !Arrays.equals(ind.shape, cshape)) { // broadcast
    123                                        throw new IllegalArgumentException("Index datasets do not have same shape");
    124                                }
    125                        } else {
    126                                if (cshape != null) {
    127                                        if (first < 0)
    128                                                first = i;
    129                                }
    130                        }
    131                }
    132
    133                List<Integer> oShape = new ArrayList<Integer>(irank);
    134
    135                if (intact) { // get new output shape list
    136                        boolean used = false;
    137                        for (int i = 0; i < irank; i++) {
    138                                Object obj = indexes.get(i);
    139                                if (obj instanceof IntegerDataset) {
    140                                        if (restrict1D || !used) {
    141                                                used = true;
    142                                                int[] lshape = restrict1D ? ((IntegerDataset) obj).shape : cshape;
    143                                                for (int j : lshape) {
    144                                                        oShape.add(j);
    145                                                }
    146                                        }
    147                                } else if (obj instanceof Slice) {
    148                                        Slice s = (Slice) obj;
    149                                        int l = ishape[i];
    150                                        s.setLength(l);
    151                                        oShape.add(s.getNumSteps());
    152                                } else {
    153                                        oShape.add(ishape[i]);
    154                                }
    155                        }
    156                } else {
    157                        assert cshape != null;
    158                        for (int j : cshape) {
    159                                oShape.add(j);
    160                        }
    161                        for (int i = 0; i < irank; i++) {
    162                                Object obj = indexes.get(i);
    163                                if (obj == null) {
    164                                        oShape.add(ishape[i]);
    165                                } else if (obj instanceof Slice) {
    166                                        Slice s = (Slice) obj;
    167                                        int l = ishape[i];
    168                                        s.setLength(l);
    169                                        oShape.add(s.getNumSteps());
    170                                }
    171                        }
    172                }
    173                orank = oShape.size();
    174                oshape = orank == 0 && ishape == null ? null : new int[orank];
    175                for (int i = 0; i < orank; i++) {
    176                        oshape[i] = oShape.get(i);
    177                }
    178
    179                for (int i = 0; i < irank; i++) { // check input indexes for out of bounds
    180                        Object obj = indexes.get(i);
    181                        if (obj instanceof IntegerDataset) {
    182                                IntegerDataset ind = (IntegerDataset) obj;
    183                                if (ind.getSize() > 0) {
    184                                        int l = ishape[i];
    185                                        if (ind.min().intValue() < -l || ind.max().intValue() >= l) {
    186                                                throw new IllegalArgumentException("A value in index datasets is outside permitted range");
    187                                        }
    188                                }
    189                        }
    190                }
    191
    192                ipos = new int[irank];
    193                it = new PositionIterator(oshape);
    194                opos = it.getPos();
    195        }
    196
    197        @Override
    198        public int[] getShape() {
    199                return oshape;
    200        }
    201
    202        @Override
    203        public boolean hasNext() {
    204                if (it.hasNext()) {
    205                        int i = 0;
    206                        for (; i < offset; i++) {
    207                                Object obj = indexes.get(i);
    208                                if (obj == null) {
    209                                        ipos[i] = opos[i];
    210                                } else if (obj instanceof Slice) {
    211                                        Slice s = (Slice) obj;
    212                                        ipos[i] = s.getPosition(opos[i]); // overwrite position
    213                                } else {
    214                                        throw new IllegalStateException("Bad state: index dataset before offset");
    215                                }
    216                        }
    217                        int[] spos = srank > 0 ? Arrays.copyOfRange(opos, i, i+srank) : opos;
    218                        if (spos == opos) {
    219                                for (; i < irank; i++) {
    220                                        Object obj = indexes.get(i);
    221                                        if (obj == null) {
    222                                                ipos[i] = opos[i];
    223                                        } else if (obj instanceof Slice) {
    224                                                Slice s = (Slice) obj;
    225                                                ipos[i] = s.getPosition(opos[i]); // overwrite position
    226                                        } else if (obj instanceof IntegerDataset) { // allowed when restricted to 1D
    227                                                int p = ((Dataset) obj).getInt(opos[i]);
    228                                                ipos[i] = p < 0 ? p + ishape[i] : p;
    229                                        } else {
    230                                                throw new IllegalStateException("Bad state: index dataset after subspace");
    231                                        }
    232                                }
    233                        } else {
    234                                for (int j = 0; j < irank; j++) {
    235                                        Object obj = indexes.get(j);
    236                                        if (obj instanceof IntegerDataset) {
    237                                                IntegerDataset ind = (IntegerDataset) obj;
    238                                                int p = ind.size > 1 ? ind.get(spos) : ind.getAbs(0); // broadcasting
    239                                                ipos[i] = p < 0 ? p + ishape[i] : p;
    240                                                i++;
    241                                        }
    242                                }
    243                                int o = orank - irank;
    244                                for (; i < irank; i++) {
    245                                        Object obj = indexes.get(i);
    246                                        if (obj == null) {
    247                                                ipos[i] = opos[i+o];
    248                                        } else if (obj instanceof Slice) {
    249                                                Slice s = (Slice) obj;
    250                                                ipos[i] = s.getPosition(opos[i+o]); // overwrite position
    251                                        } else {
    252                                                throw new IllegalStateException("Bad state: index dataset after subspace");
    253                                        }
    254                                }
    255                        }
    256//                      System.err.println(Arrays.toString(opos) + ", " + Arrays.toString(spos) + ", " + Arrays.toString(ipos));
    257                        return true;
    258                }
    259                return false;
    260        }
    261
    262        @Override
    263        public int[] getPos() {
    264                return ipos;
    265        }
    266
    267        @Override
    268        public void reset() {
    269                it.reset();
    270                Arrays.fill(ipos, 0);
    271                index = 0;
    272        }
    273}