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 a = ShapeUtils.checkAxis(rank, a); 077 if (a >= 0 && a <= endrank) { 078 sRank--; 079 omit[a] = true; 080 shape[a] = 1; 081 } else if (a > endrank) { 082 throw new IllegalArgumentException("Specified axis exceeds dataset rank"); 083 } 084 } 085 } 086 087 cSlice = cSlice.clone(); 088 pos = cSlice.getStart(); 089 end = cSlice.getStop(); 090 if (sRank == rank) { 091 sStart = pos; 092 sStop = null; 093 oSlice = null; 094 sSlice = cSlice; 095 } else { 096 int[] dShape = dSlice.getShape(); 097 int[] lShape = new int[sRank]; // inner shape 098 int[] oShape = new int[rank - sRank]; // output shape 099 for (int i = 0, j = 0, k = 0; i < rank; i++) { 100 if (omit[i]) { 101 oShape[j++] = sShape[i]; 102 } else { 103 lShape[k++] = sShape[i]; 104 dShape[i] = 1; 105 } 106 } 107 sSlice = new SliceND(lShape); 108 sStart = sSlice.getStart(); 109 sStop = sSlice.getStop(); 110 111 oSlice = new SliceND(oShape); 112 for (int i = 0, j = 0, k = 0; i < rank; i++) { 113 if (omit[i]) { 114 oSlice.setSlice(j++, start[i], stop[i], step[i]); 115 } else { 116 sSlice.setSlice(k++, start[i], stop[i], step[i]); 117 } 118 } 119 } 120 121 reset(); 122 } 123 124 @Override 125 public boolean hasNext() { 126 // now move on one position 127 if (once) { 128 once = false; 129 return true; 130 } 131 int k = sRank - 1; 132 for (int j = endrank; j >= 0; j--) { 133 if (omit[j]) { 134 continue; 135 } 136 pos[j] += step[j]; 137 end[j] = pos[j] + step[j]; 138 dStart[j]++; 139 dStop[j]++; 140 if (pos[j] >= stop[j]) { 141 pos[j] = start[j]; 142 end[j] = pos[j] + step[j]; 143 dStart[j] = 0; 144 dStop[j] = 1; 145 if (sStop != null) { 146 sStart[k] = pos[j]; 147 sStop[k] = end[j]; 148 k--; 149 } 150 } else { 151 if (sStop != null) { 152 sStart[k] = pos[j]; 153 sStop[k] = end[j]; 154 k--; 155 } 156 return true; 157 } 158 } 159 return false; 160 } 161 162 @Override 163 public int[] getPos() { 164 return pos; 165 } 166 167 /** 168 * Get omitted part of source slice which never changes 169 * @return slice (can be null) 170 */ 171 public SliceND getOmittedSlice() { 172 return oSlice; 173 } 174 175 /** 176 * Get output or destination slice 177 * @return slice 178 */ 179 public SliceND getOutputSlice() { 180 return dSlice; 181 } 182 183 /** 184 * Get current slice 185 * @return slice 186 */ 187 public SliceND getCurrentSlice() { 188 return cSlice; 189 } 190 191 /** 192 * Shortened position where axes are omitted 193 * @return used position 194 */ 195 public int[] getUsedPos() { 196 return sStart; 197 } 198 199 /** 200 * Shortened slice where axes are omitted 201 * @return used slice 202 */ 203 public SliceND getUsedSlice() { 204 return sSlice; 205 } 206 207 /** 208 * @return omit array - array where true means miss out 209 */ 210 public boolean[] getOmit() { 211 return omit; 212 } 213 214 @Override 215 public void reset() { 216 for (int i = 0, k = 0; i <= endrank; i++) { 217 int b = start[i]; 218 int d = step[i]; 219 if (!omit[i]) { 220 cSlice.setSlice(i, b, b + d, d); 221 dStart[i] = 0; 222 dStop[i] = 1; 223 if (sStop != null) { 224 sSlice.setSlice(k++, b, b + d, d); 225 } 226 } else { 227 cSlice.setSlice(i, b, end[i], d); 228 } 229 } 230 231 int j = 0; 232 for (; j <= endrank; j++) { 233 if (!omit[j]) 234 break; 235 } 236 if (j > endrank) { 237 once = true; 238 return; 239 } 240 241 if (omit[endrank]) { 242 pos[endrank] = start[endrank]; 243 for (int i = endrank - 1; i >= 0; i--) { 244 if (!omit[i]) { 245 end[i] = pos[i]; 246 pos[i] -= step[i]; 247 dStart[i]--; 248 dStop[i]--; 249 break; 250 } 251 } 252 } else { 253 end[endrank] = pos[endrank]; 254 pos[endrank] -= step[endrank]; 255 dStart[endrank]--; 256 dStop[endrank]--; 257 } 258 259 if (sStart != pos) { 260 for (int i = 0, k = 0; i <= endrank; i++) { 261 if (!omit[i]) { 262 sStart[k++] = pos[i]; 263 } 264 } 265 } 266 } 267 268 @Override 269 public int[] getShape() { 270 return shape; 271 } 272}