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 final int rank = shape == null ? 0 : 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 final int rank = shape == null ? 0 : shape.length; 083 zero = shape == null; 084 endrank = rank - 1; 085 pos = new int[rank]; 086 delta = new int[rank]; 087 this.element = element; 088 if (strides == null) { 089 offset = 0; 090 stride = new int[rank]; 091 int s = isize; 092 for (int j = endrank; j >= 0; j--) { 093 stride[j] = s; 094 s *= shape[j]; 095 } 096 } else { 097 stride = strides.clone(); 098 } 099 100 for (int j = endrank; j >= 0; j--) { 101 int t = stride[j]; 102 offset += t * start[j]; 103 t *= step[j]; 104 stride[j] = t; 105 int s = sshape[j]; 106 if (!zero) { 107 zero = s == 0; 108 } 109 delta[j] = s * t; 110 } 111 112 nstart = offset; 113 } 114 115 @Override 116 void calcGap() { 117 // do nothing 118 } 119 120 @Override 121 public boolean hasNext() { 122 if (zero) { 123 return false; 124 } 125 126 // now move on one position 127 int j = endrank; 128 if (j < 0) { 129 index += istep; 130 return index < istep; 131 } 132 for (; j >= 0; j--) { 133 index += stride[j]; 134 final int s = step[j]; 135 final int p = pos[j] + s; 136 if ((s > 0 && p < stop[j]) || (s < 0 && p > stop[j])) { 137 pos[j] = p; 138 break; 139 } 140 pos[j] = start[j]; 141 index -= delta[j]; // reset this dimension 142 } 143 return j >= 0; 144 } 145 146 @Override 147 public int[] getPos() { 148 return pos; 149 } 150 151 @Override 152 public void reset() { 153 System.arraycopy(start, 0, pos, 0, start.length); 154 if (endrank >= 0) { 155 pos[endrank] -= step[endrank]; 156 index = nstart - stride[endrank]; 157 } else { 158 index = -istep; 159 } 160 index += element; 161 } 162}