001/*- 002 ******************************************************************************* 003 * Copyright (c) 2015, 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 * Class to provide slice iteration through a dataset 017 * <p>It allows a number of axes to be omitted and iterates over 018 * the axes left over. The omitted axes define an inner shape and 019 * the remaining axes define an outer shape which is iterated over. 020 */ 021public class SliceNDIterator extends IndexIterator { 022 final private int[] shape; 023 final private int[] start; 024 final private int[] stop; 025 final private int[] step; 026 final private int endrank; 027 028 final private boolean[] omit; // axes to miss out 029 030 /** 031 * position in source dataset 032 */ 033 final private int[] pos; 034 final private int[] end; 035 private boolean once; 036 037 private SliceND cSlice; // current slice 038 039 private int sRank; // number of dimensions used (i.e. not missing) 040 final private SliceND iSlice; // omitted source or inner slice 041 042 final private SliceND sSlice; // shortened slice 043 final private int[] sStart; // shortened position 044 final private int[] sStop; // shortened end 045 046 private SliceND dSlice; // destination slice 047 private int[] dStart; 048 private int[] dStop; 049 050 /** 051 * Constructor for an iterator that misses out several axes 052 * @param slice 053 * @param axes missing axes 054 */ 055 public SliceNDIterator(SliceND slice, int... axes) { 056 cSlice = slice.clone(); 057 int[] sShape = cSlice.getSourceShape(); 058 shape = cSlice.getShape().clone(); 059 start = cSlice.getStart(); 060 stop = cSlice.getStop(); 061 step = cSlice.getStep(); 062 for (int s : step) { 063 if (s < 0) { 064 throw new UnsupportedOperationException("Negative steps not implemented"); 065 } 066 } 067 int rank = shape.length; 068 endrank = rank - 1; 069 070 omit = new boolean[rank]; 071 dSlice = new SliceND(shape); 072 dStart = dSlice.getStart(); 073 dStop = dSlice.getStop(); 074 sRank = rank; 075 if (axes != null) { 076 axes = ShapeUtils.checkAxes(rank, axes); 077 for (int a : axes) { 078 if (a >= 0 && a <= endrank) { 079 sRank--; 080 omit[a] = true; 081 shape[a] = 1; 082 } else if (a > endrank) { 083 throw new IllegalArgumentException("Specified axis exceeds dataset rank"); 084 } 085 } 086 } 087 088 cSlice = cSlice.clone(); 089 pos = cSlice.getStart(); 090 end = cSlice.getStop(); 091 if (sRank == rank) { 092 sStart = pos; 093 sStop = null; 094 iSlice = null; 095 sSlice = cSlice; 096 } else { 097 int[] dShape = dSlice.getShape(); 098 int[] oShape = new int[sRank]; // outer shape 099 int[] iShape = new int[rank - sRank]; // inner shape 100 for (int i = 0, j = 0, k = 0; i < rank; i++) { 101 if (omit[i]) { 102 iShape[j++] = sShape[i]; 103 } else { 104 oShape[k++] = sShape[i]; 105 dShape[i] = 1; 106 } 107 } 108 sSlice = new SliceND(oShape); 109 sStart = sSlice.getStart(); 110 sStop = sSlice.getStop(); 111 112 iSlice = new SliceND(iShape); 113 for (int i = 0, j = 0, k = 0; i < rank; i++) { 114 if (omit[i]) { 115 iSlice.setSlice(j++, start[i], stop[i], step[i]); 116 } else { 117 sSlice.setSlice(k++, start[i], stop[i], step[i]); 118 } 119 } 120 } 121 122 reset(); 123 } 124 125 @Override 126 public boolean hasNext() { 127 // now move on one position 128 if (once) { 129 once = false; 130 return true; 131 } 132 int k = sRank - 1; 133 for (int j = endrank; j >= 0; j--) { 134 if (omit[j]) { 135 continue; 136 } 137 pos[j] += step[j]; 138 end[j] = pos[j] + step[j]; 139 dStart[j]++; 140 dStop[j]++; 141 if (pos[j] >= stop[j]) { 142 pos[j] = start[j]; 143 end[j] = pos[j] + step[j]; 144 dStart[j] = 0; 145 dStop[j] = 1; 146 if (sStop != null) { 147 sStart[k] = pos[j]; 148 sStop[k] = end[j]; 149 k--; 150 } 151 } else { 152 if (sStop != null) { 153 sStart[k] = pos[j]; 154 sStop[k] = end[j]; 155 k--; 156 } 157 return true; 158 } 159 } 160 return false; 161 } 162 163 @Override 164 public int[] getPos() { 165 return pos; 166 } 167 168 /** 169 * Get omitted part of source slice which never changes. I.e. the inner slice 170 * @return slice (can be null) 171 */ 172 public SliceND getOmittedSlice() { 173 return iSlice; 174 } 175 176 /** 177 * Get output or destination slice 178 * @return slice 179 */ 180 public SliceND getOutputSlice() { 181 return dSlice; 182 } 183 184 /** 185 * Get current slice 186 * @return slice 187 */ 188 public SliceND getCurrentSlice() { 189 return cSlice; 190 } 191 192 /** 193 * Shortened position where axes are omitted 194 * @return used position 195 */ 196 public int[] getUsedPos() { 197 return sStart; 198 } 199 200 /** 201 * Shortened slice where axes are omitted 202 * @return used (or outer) slice 203 */ 204 public SliceND getUsedSlice() { 205 return sSlice; 206 } 207 208 /** 209 * @return omit array - array where true means miss out 210 */ 211 public boolean[] getOmit() { 212 return omit; 213 } 214 215 @Override 216 public void reset() { 217 for (int i = 0, k = 0; i <= endrank; i++) { 218 int b = start[i]; 219 int d = step[i]; 220 if (!omit[i]) { 221 cSlice.setSlice(i, b, b + d, d); 222 dStart[i] = 0; 223 dStop[i] = 1; 224 if (sStop != null) { 225 sSlice.setSlice(k++, b, b + d, d); 226 } 227 } else { 228 cSlice.setSlice(i, b, end[i], d); 229 } 230 } 231 232 int j = 0; 233 for (; j <= endrank; j++) { 234 if (!omit[j]) 235 break; 236 } 237 if (j > endrank) { 238 once = true; 239 return; 240 } 241 242 if (omit[endrank]) { 243 pos[endrank] = start[endrank]; 244 for (int i = endrank - 1; i >= 0; i--) { 245 if (!omit[i]) { 246 end[i] = pos[i]; 247 pos[i] -= step[i]; 248 dStart[i]--; 249 dStop[i]--; 250 break; 251 } 252 } 253 } else { 254 end[endrank] = pos[endrank]; 255 pos[endrank] -= step[endrank]; 256 dStart[endrank]--; 257 dStop[endrank]--; 258 } 259 260 if (sStart != pos) { 261 for (int i = 0, k = 0; i <= endrank; i++) { 262 if (!omit[i]) { 263 sStart[k++] = pos[i]; 264 } 265 } 266 } 267 } 268 269 @Override 270 public int[] getShape() { 271 return shape; 272 } 273}