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 015 016import java.util.Arrays; 017 018/** 019 * Class to run over a slice of a dataset 020 */ 021public class SliceIterator extends IndexIterator { 022 int[] shape; 023 int isize; 024 int endrank; // last shape index 025 int[] gap; // gaps in dataset 026 int imax; // maximum index in array 027 int[] start; 028 int[] stop; 029 int[] step; 030 int[] sshape; // slice shape 031 int[] pos; // position in dataset 032 int istep; // step in last index 033 034 SliceIterator() { 035 } 036 037 /** 038 * Constructor for an iterator over the elements of a sliced dataset 039 * 040 * @param shape or dataShape 041 * @param length of entire data array 042 * @param sshape shape of new dataset, i.e. slice 043 */ 044 public SliceIterator(final int[] shape, final int length, final int[] sshape) { 045 this(shape, length, null, null, sshape, 1); 046 } 047 048 /** 049 * Constructor for an iterator over the elements of a sliced dataset 050 * 051 * @param shape or dataShape 052 * @param length of entire data array 053 * @param start (if null then equivalent to the origin) 054 * @param sshape shape of new dataset, i.e. slice 055 */ 056 public SliceIterator(final int[] shape, final int length, final int[] start, final int[] sshape) { 057 this(shape, length, start, null, sshape, 1); 058 } 059 060 /** 061 * Constructor for an iterator over the elements of a sliced dataset 062 * 063 * @param shape or dataShape 064 * @param length of entire data array 065 * @param sshape shape of new dataset, i.e. slice 066 * @param isize number of elements in an item 067 */ 068 public SliceIterator(final int[] shape, final int length, final int[] sshape, final int isize) { 069 this(shape, length, null, null, sshape, isize); 070 } 071 072 /** 073 * Constructor for an iterator over the elements of a sliced dataset 074 * 075 * @param shape or dataShape 076 * @param length of entire data array 077 * @param start (if null then equivalent to the origin) 078 * @param sshape shape of new dataset, i.e. slice 079 * @param isize number of elements in an item 080 */ 081 public SliceIterator(final int[] shape, final int length, final int[] start, final int[] sshape, final int isize) { 082 this(shape, length, start, null, sshape, isize); 083 } 084 085 /** 086 * Constructor for an iterator over the elements of a sliced dataset 087 * 088 * @param shape or dataShape 089 * @param length of entire data array 090 * @param start (if null then equivalent to the origin) 091 * @param step (cannot contain zeros, if null then equivalent to ones) 092 * @param sshape shape of new dataset, i.e. slice 093 */ 094 public SliceIterator(final int[] shape, final int length, final int[] start, final int[] step, final int[] sshape) { 095 this(shape, length, start, step, sshape, 1); 096 } 097 098 /** 099 * Constructor for an iterator over the elements of a sliced dataset 100 * 101 * @param shape or dataShape 102 * @param length of entire data array 103 * @param slice 104 */ 105 public SliceIterator(final int[] shape, final int length, final SliceND slice) { 106 this(shape, length, slice.getStart(), slice.getStep(), slice.getShape(), 1); 107 } 108 109 /** 110 * Constructor for an iterator over the elements of a sliced dataset 111 * 112 * @param shape or dataShape 113 * @param length of entire data array 114 * @param isize number of elements in an item 115 * @param slice 116 */ 117 public SliceIterator(final int[] shape, final int length, final int isize, final SliceND slice) { 118 this(shape, length, slice.getStart(), slice.getStep(), slice.getShape(), isize); 119 } 120 121 /** 122 * Constructor for an iterator over the elements of a sliced dataset 123 * 124 * @param shape or dataShape 125 * @param length of entire data array 126 * @param start (if null then equivalent to the origin) 127 * @param step (cannot contain zeros, if null then equivalent to ones) 128 * @param sshape shape of new dataset, i.e. slice 129 * @param isize number of elements in an item 130 */ 131 public SliceIterator(final int[] shape, final int length, final int[] start, final int[] step, final int[] sshape, final int isize) { 132 this.isize = isize; 133 int rank = shape.length; 134 endrank = rank - 1; 135 this.shape = shape; 136 this.start = new int[rank]; 137 this.sshape = sshape; 138 if (step == null) { 139 this.step = new int[rank]; 140 Arrays.fill(this.step, 1); 141 } else { 142 this.step = step; 143 } 144 145 if (rank == 0) { 146 istep = isize; 147 imax = length*isize; 148 stop = new int[0]; 149 pos = new int[0]; 150 gap = null; 151 } else { 152 istep = this.step[endrank] * isize; 153 imax = length * isize; 154 stop = new int[rank]; 155 gap = new int[endrank + 1]; 156 pos = new int[rank]; 157 calcGap(); 158 } 159 160 setStart(start); 161 } 162 163 void calcGap() { 164 int chunk = isize; 165 for (int i = endrank; i >= 0; i--) { 166 stop[i] = start[i] + sshape[i] * step[i]; 167 168 if (step[i] < 0) 169 stop[i]++; // adjust for -ve steps so later code has more succinct test 170 171 if (i > 0) { 172 gap[i] = (shape[i] * step[i - 1] - sshape[i] * step[i]) * chunk; 173 chunk *= shape[i]; 174 } 175 } 176 } 177 178 /** 179 * Set start (prefix with zeros if necessary) 180 * @param newStart if null, then treat as origin 181 */ 182 public void setStart(int... newStart) { 183 final int rank = shape.length; 184 if (rank == 0) { 185 index = -istep; 186 return; 187 } 188 189 if (newStart == null) { 190 for (int i = 0; i < rank; i++) { 191 start[i] = 0; 192 } 193 } else if (newStart.length > rank) { 194 throw new IllegalArgumentException("Length of start array greater than rank"); 195 } else { 196 int extra = rank - newStart.length; 197 for (int i = 0; i < extra; i++) { 198 start[i] = 0; 199 } 200 for (int i = 0; i < newStart.length; i++) { 201 start[i+extra] = newStart[i]; 202 } 203 } 204 205 reset(); 206 calcGap(); 207 } 208 209 @Override 210 public void reset() { 211 if (shape.length == 0) { 212 index = -istep; 213 } else { 214 // work out index of first position 215 for (int i = 0; i < shape.length; i++) 216 pos[i] = start[i]; 217 pos[endrank] -= step[endrank]; 218 219 index = pos[0]; 220 for (int j = 1; j <= endrank; j++) 221 index = index * shape[j] + pos[j]; 222 index *= isize; 223 } 224 } 225 226 @Override 227 public boolean hasNext() { 228 // now move on one position in slice 229 int j = endrank; 230 for (; j >= 0; j--) { 231 pos[j] += step[j]; 232 233 if ((pos[j] >= stop[j]) == (step[j] > 0)) { // stop index has been adjusted in code for -ve steps 234 pos[j] = start[j]; 235 index += gap[j]; 236 } else { 237 break; 238 } 239 } 240 if (j == -1 && endrank >= 0) { 241 index = imax; 242 return false; 243 } 244 245 index += istep; 246 return index < imax; 247 } 248 249 public int[] getStart() { 250 return start; 251 } 252 253 @Override 254 public int[] getPos() { 255 return pos; 256 } 257 258 public int[] getStep() { 259 return step; 260 } 261 262 @Override 263 public int[] getShape() { 264 return sshape; 265 } 266}