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