org.nd4j.linalg.api.iter.NdIndexIterator Maven / Gradle / Ivy
package org.nd4j.linalg.api.iter;
import org.apache.commons.math3.util.Pair;
import org.nd4j.linalg.api.shape.Shape;
import org.nd4j.linalg.util.ArrayUtil;
import java.util.HashMap;
import java.util.Iterator;
import java.util.Map;
/**
* Iterates and returns int arrays
* over a particular shape.
*
* This iterator starts at zero and increments
* the shape until each item in the "position"
* hits the current shape
*
* @author Adam Gibson
*/
public class NdIndexIterator implements Iterator {
private int length = -1;
private int i = 0;
private int[] shape;
private char order = 'c';
private boolean cache = false;
private static Map, LinearIndexLookup> lookupMap = new HashMap<>();
private LinearIndexLookup lookup;
/**
* Pass in the shape to iterate over.
* Defaults to c ordering
* @param shape the shape to iterate over
*/
public NdIndexIterator(int... shape) {
this('c', shape);
this.cache = false;
}
/**
* Pass in the shape to iterate over.
* Defaults to c ordering
* @param shape the shape to iterate over
*/
public NdIndexIterator(char order, boolean cache, int... shape) {
this.shape = ArrayUtil.copy(shape);
this.length = ArrayUtil.prod(shape);
this.order = order;
this.cache = cache;
if (this.cache) {
LinearIndexLookup lookup = lookupMap.get(new Pair<>(shape, order));
if (lookup == null) {
lookup = new LinearIndexLookup(shape, order);
//warm up the cache
for (int i = 0; i < length; i++) {
lookup.lookup(i);
}
lookupMap.put(new Pair<>(shape, order), lookup);
this.lookup = lookup;
} else {
this.lookup = lookupMap.get(new Pair<>(shape, order));
}
}
}
/**
* Pass in the shape to iterate over
* @param shape the shape to iterate over
*/
public NdIndexIterator(char order, int... shape) {
this(order, false, shape);
}
@Override
public boolean hasNext() {
return i < length;
}
@Override
public int[] next() {
if (lookup != null)
return lookup.lookup(i++);
switch (order) {
case 'c':
return Shape.ind2subC(shape, i++);
case 'f':
return Shape.ind2sub(shape, i++);
default:
throw new IllegalArgumentException("Illegal ordering " + order);
}
}
@Override
public void remove() {
}
}