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}