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