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.Arrays; 016 017 018/** 019 * Class to run over a contiguous dataset using strides 020 */ 021public class StrideIterator extends SliceIterator { 022 private int[] stride; 023 private int[] delta; // reset values 024 private int nstart; 025 private int element; 026 private boolean zero; 027 028 public StrideIterator(final int[] shape) { 029 this(shape, null, 0); 030 } 031 032 public StrideIterator(final int isize, final int[] shape) { 033 this(isize, shape, null, 0); 034 } 035 036 public StrideIterator(final int[] shape, final int[] strides) { 037 this(shape, strides, 0); 038 } 039 040 public StrideIterator(final int[] shape, final int[] strides, final int offset) { 041 this(1, shape, strides, offset); 042 } 043 044 public StrideIterator(final int isize, final int[] shape, final int[] strides, final int offset) { 045 this(isize, shape, strides, offset, 0); 046 } 047 048 public StrideIterator(final int isize, final int[] shape, final int[] strides, final int offset, final int element) { 049 int rank = shape.length; 050 start = new int[rank]; 051 stop = shape; 052 step = new int[rank]; 053 Arrays.fill(step, 1); 054 this.sshape = shape; 055 056 init(isize, shape, strides, offset, element); 057 reset(); 058 } 059 060 public StrideIterator(final int isize, final int[] shape, final int[] start, final int[] stop, final int[] step) { 061 this(isize, shape, null, 0, start, stop, step); 062 } 063 064 public StrideIterator(final int isize, final int[] shape, final int[] oStrides, final int oOffset, final int[] start, final int[] stop, final int[] step) { 065 this(isize, shape, oStrides, oOffset, new SliceND(shape, start, stop, step)); 066 } 067 068 public StrideIterator(final int isize, final int[] shape, final int[] strides, final int offset, final SliceND slice) { 069 start = slice.getStart(); 070 stop = slice.getStop(); 071 step = slice.getStep(); 072 this.sshape = slice.getShape(); 073 074 init(isize, shape, strides, offset, 0); 075 reset(); 076 } 077 078 private void init(final int isize, final int[] shape, final int[] strides, int offset, final int element) { 079 this.isize = isize; 080 istep = isize; 081 this.shape = shape; 082 int rank = shape.length; 083 endrank = rank - 1; 084 pos = new int[rank]; 085 delta = new int[rank]; 086 this.element = element; 087 if (strides == null) { 088 offset = 0; 089 stride = new int[rank]; 090 int s = isize; 091 for (int j = endrank; j >= 0; j--) { 092 stride[j] = s; 093 s *= shape[j]; 094 } 095 } else { 096 stride = strides.clone(); 097 } 098 099 for (int j = endrank; j >= 0; j--) { 100 int t = stride[j]; 101 offset += t * start[j]; 102 t *= step[j]; 103 stride[j] = t; 104 int s = sshape[j]; 105 if (!zero) { 106 zero = s == 0; 107 } 108 delta[j] = s * t; 109 } 110 111 nstart = offset; 112 } 113 114 @Override 115 void calcGap() { 116 // do nothing 117 } 118 119 @Override 120 public boolean hasNext() { 121 if (zero) { 122 return false; 123 } 124 125 // now move on one position 126 int j = endrank; 127 if (j < 0) { 128 index += istep; 129 return index < istep; 130 } 131 for (; j >= 0; j--) { 132 index += stride[j]; 133 final int s = step[j]; 134 final int p = pos[j] + s; 135 if ((s > 0 && p < stop[j]) || (s < 0 && p > stop[j])) { 136 pos[j] = p; 137 break; 138 } 139 pos[j] = start[j]; 140 index -= delta[j]; // reset this dimension 141 } 142 return j >= 0; 143 } 144 145 @Override 146 public int[] getPos() { 147 return pos; 148 } 149 150 @Override 151 public void reset() { 152 System.arraycopy(start, 0, pos, 0, start.length); 153 if (endrank >= 0) { 154 pos[endrank] -= step[endrank]; 155 index = nstart - stride[endrank]; 156 } else { 157 index = -istep; 158 } 159 index += element; 160 } 161}