com.simiacryptus.mindseye.lang.Tensor Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of mindseye-core Show documentation
Core Neural Networks Framework
There is a newer version: 2.1.0
/*
 * Copyright (c) 2018 by Andrew Charneski.
 *
 * The author licenses this file to you under the
 * Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance
 * with the License.  You may obtain a copy
 * of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing,
 * software distributed under the License is distributed on an
 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
 * KIND, either express or implied.  See the License for the
 * specific language governing permissions and limitations
 * under the License.
 */

package com.simiacryptus.mindseye.lang;

import com.google.gson.JsonArray;
import com.google.gson.JsonElement;
import com.google.gson.JsonObject;
import com.google.gson.JsonPrimitive;
import org.apache.commons.lang3.ArrayUtils;

import javax.annotation.Nonnull;
import javax.annotation.Nullable;
import java.awt.image.BufferedImage;
import java.io.Serializable;
import java.util.*;
import java.util.function.*;
import java.util.stream.*;

/**
 * A multi-dimensional array of data. Represented internally as a single double[] array. This class is central to data
 * handling in MindsEye, and may have some odd-looking or suprising optimizations.
 */
@SuppressWarnings("serial")
public final class Tensor extends ReferenceCountingBase implements Serializable {

  /**
   * The constant json_precision.
   */
  @Nonnull
  public static DataSerializer json_precision = SerialPrecision.Float;
  /**
   * The Dimensions.
   */
  @Nullable
  protected final int[] dimensions;
  /**
   * The Strides.
   */
  @Nullable
  protected final int[] strides;
  /**
   * The Data.
   */
  @Nullable
  protected volatile double[] data;

  @Nullable
  protected volatile UUID id;

  /**
   * Instantiates a new Tensor.
   */
  private Tensor() {
    super();
    data = null;
    strides = null;
    dimensions = null;
  }

  /**
   * Instantiates a new Tensor.
   *
   * @param ds the ds
   */
  public Tensor(@Nonnull final double... ds) {
    this(ds, ds.length);
  }

  /**
   * Instantiates a new Tensor.
   *
   * @param data the data
   * @param dims the dims
   */
  public Tensor(@Nullable final double[] data, @Nonnull final int... dims) {
    if (Tensor.length(dims) > Integer.MAX_VALUE) throw new IllegalArgumentException();
    if (null != data && Tensor.length(dims) != data.length)
      throw new IllegalArgumentException(Arrays.toString(dims) + " != " + data.length);
    dimensions = (null == dims || 0 == dims.length) ? new int[]{} : Arrays.copyOf(dims, dims.length);
    strides = Tensor.getSkips(dims);
    //this.data = data;// Arrays.copyOf(data, data.length);
    if (null != data) {
      this.data = RecycleBin.DOUBLES.copyOf(data, data.length);
    }
    assert isValid();
    //assert (null == data || Tensor.length(dims) == data.length);
  }

  private Tensor(int[] dims, @Nullable double[] data) {
    this(dims, Tensor.getSkips(dims), data);
  }

  private Tensor(int[] dimensions, int[] strides, @Nullable double[] data) {
    if (Tensor.length(dimensions) >= Integer.MAX_VALUE) throw new IllegalArgumentException();
    assert null == data || data.length == Tensor.length(dimensions);
    this.dimensions = dimensions;
    this.strides = strides;
    this.data = data;
    assert isValid();
  }

  /**
   * Instantiates a new Tensor.
   *
   * @param data the data
   * @param dims the dims
   */
  public Tensor(@Nullable final float[] data, @Nonnull final int... dims) {
    if (Tensor.length(dims) >= Integer.MAX_VALUE) throw new IllegalArgumentException();
    dimensions = Arrays.copyOf(dims, dims.length);
    strides = Tensor.getSkips(dims);
    if (null != data) {
      this.data = RecycleBin.DOUBLES.obtain(data.length);// Arrays.copyOf(data, data.length);
      Arrays.parallelSetAll(this.data, i -> {
        final double v = data[i];
        return Double.isFinite(v) ? v : 0;
      });
      assert Arrays.stream(this.data).allMatch(v -> Double.isFinite(v));
    }
    assert isValid();
    //assert (null == data || Tensor.length(dims) == data.length);
  }

  /**
   * Instantiates a new Tensor.
   *
   * @param dims the dims
   */
  public Tensor(@Nonnull final int... dims) {
    this((double[]) null, dims);
    assert dims.length > 0;
  }

  /**
   * From json tensor.
   *
   * @param json      the json
   * @param resources the resources
   * @return the tensor
   */
  @Nullable
  public static Tensor fromJson(@Nullable final JsonElement json, @Nullable Map resources) {
    if (null == json) return null;
    if (json.isJsonArray()) {
      final JsonArray array = json.getAsJsonArray();
      final int size = array.size();
      if (array.get(0).isJsonPrimitive()) {
        final double[] doubles = IntStream.range(0, size).mapToObj(i -> {
          return array.get(i);
        }).mapToDouble(element -> {
          return element.getAsDouble();
        }).toArray();
        @Nonnull Tensor tensor = new Tensor(doubles);
        assert tensor.isValid();
        return tensor;
      } else {
        final List elements = IntStream.range(0, size).mapToObj(i -> {
          return array.get(i);
        }).map(element -> {
          return Tensor.fromJson(element, resources);
        }).collect(Collectors.toList());
        @Nonnull final int[] dimensions = elements.get(0).getDimensions();
        if (!elements.stream().allMatch(t -> Arrays.equals(dimensions, t.getDimensions()))) {
          throw new IllegalArgumentException();
        }
        @Nonnull final int[] newDdimensions = Arrays.copyOf(dimensions, dimensions.length + 1);
        newDdimensions[dimensions.length] = size;
        @Nonnull final Tensor tensor = new Tensor(newDdimensions);
        @Nullable final double[] data = tensor.getData();
        for (int i = 0; i < size; i++) {
          @Nullable final double[] e = elements.get(i).getData();
          System.arraycopy(e, 0, data, i * e.length, e.length);
        }
        for (@Nonnull Tensor t : elements) {
          t.freeRef();
        }
        assert tensor.isValid();
        return tensor;
      }
    } else if (json.isJsonObject()) {
      JsonObject jsonObject = json.getAsJsonObject();
      @Nonnull int[] dims = fromJsonArray(jsonObject.getAsJsonArray("length"));
      @Nonnull Tensor tensor = new Tensor(dims);
      SerialPrecision precision = SerialPrecision.valueOf(jsonObject.getAsJsonPrimitive("precision").getAsString());
      JsonElement base64 = jsonObject.get("base64");
      if (null == base64) {
        if (null == resources) throw new IllegalArgumentException("No Data Resources");
        CharSequence resourceId = jsonObject.getAsJsonPrimitive("resource").getAsString();
        tensor.setBytes(resources.get(resourceId), precision);
      } else {
        tensor.setBytes(Base64.getDecoder().decode(base64.getAsString()), precision);
      }
      assert tensor.isValid();
      JsonElement id = jsonObject.get("id");
      if(null != id) {
        tensor.setId(UUID.fromString(id.getAsString()));
      }
      return tensor;
    } else {
      @Nonnull Tensor tensor = new Tensor(json.getAsJsonPrimitive().getAsDouble());
      assert tensor.isValid();
      return tensor;
    }
  }

  private static double bound8bit(final double value) {
    final int max = 0xFF;
    final int min = 0;
    return value < min ? min : value > max ? max : value;
  }

  private static int bound8bit(final int value) {
    final int max = 0xFF;
    final int min = 0;
    return value < min ? min : value > max ? max : value;
  }

  /**
   * Dim l long.
   *
   * @param dims the dims
   * @return the long
   */
  public static int length(@Nonnull int... dims) {
    long total = 1;
    for (final int dim : dims) {
      total *= dim;
    }
    return (int) total;
  }

  /**
   * From rgb tensor.
   *
   * @param img the img
   * @return the tensor
   */
  @Nonnull
  public static Tensor fromRGB(@Nonnull final BufferedImage img) {
    final int width = img.getWidth();
    final int height = img.getHeight();
    @Nonnull final Tensor a = new Tensor(width, height, 3);
    IntStream.range(0, width).parallel().forEach(x -> {
      @Nonnull final int[] coords = {0, 0, 0};
      IntStream.range(0, height).forEach(y -> {
        coords[0] = x;
        coords[1] = y;
        coords[2] = 0;
        a.set(coords, img.getRGB(x, y) & 0xFF);
        coords[2] = 1;
        a.set(coords, img.getRGB(x, y) >> 8 & 0xFF);
        coords[2] = 2;
        a.set(coords, img.getRGB(x, y) >> 16 & 0x0FF);
      });
    });
    return a;
  }

  /**
   * Get doubles double [ ].
   *
   * @param stream the stream
   * @param dim    the length
   * @return the double [ ]
   */
  public static double[] getDoubles(@Nonnull final DoubleStream stream, final int dim) {
    final double[] doubles = RecycleBin.DOUBLES.obtain(dim);
    stream.forEach(new DoubleConsumer() {
      int j = 0;

      @Override
      public void accept(final double value) {
        doubles[j++] = value;
      }
    });
    return doubles;
  }

  @Nonnull
  private static int[] getSkips(@Nonnull final int[] dims) {
    @Nonnull final int[] skips = new int[dims.length];
    for (int i = 0; i < skips.length; i++) {
      if (i == 0) {
        skips[0] = 1;
      } else {
        skips[i] = skips[i - 1] * dims[i - 1];
      }
    }
    return skips;
  }

  /**
   * Product tensor.
   *
   * @param left  the left
   * @param right the right
   * @return the tensor
   */
  @Nonnull
  public static Tensor product(@Nonnull final Tensor left, @Nonnull final Tensor right) {
    if (left.length() == 1 && right.length() != 1) return Tensor.product(right, left);
    assert left.length() == right.length() || 1 == right.length();
    @Nonnull final Tensor result = new Tensor(left.getDimensions());
    @Nullable final double[] resultData = result.getData();
    @Nullable final double[] leftData = left.getData();
    @Nullable final double[] rightData = right.getData();
    for (int i = 0; i < resultData.length; i++) {
      final double l = leftData[i];
      final double r = rightData[1 == rightData.length ? 0 : i];
      resultData[i] = l * r;
    }
    return result;
  }

  /**
   * To doubles double [ ].
   *
   * @param data the data
   * @return the double [ ]
   */
  public static double[] toDoubles(@Nonnull final float[] data) {
    final double[] buffer = RecycleBin.DOUBLES.obtain(data.length);
    for (int i = 0; i < data.length; i++) {
      buffer[i] = data[i];
    }
    return buffer;
  }

  /**
   * To floats float [ ].
   *
   * @param data the data
   * @return the float [ ]
   */
  @Nonnull
  public static float[] toFloats(@Nonnull final double[] data) {
    @Nonnull final float[] buffer = new float[data.length];
    for (int i = 0; i < data.length; i++) {
      buffer[i] = (float) data[i];
    }
    return buffer;
  }

  /**
   * To json array json array.
   *
   * @param ints the ints
   * @return the json array
   */
  @Nonnull
  public static JsonArray toJsonArray(@Nonnull int[] ints) {
    @Nonnull JsonArray dim = new JsonArray();
    for (int i = 0; i < ints.length; i++) {
      dim.add(new JsonPrimitive(ints[i]));
    }
    return dim;
  }

  /**
   * From json array int [ ].
   *
   * @param ints the ints
   * @return the int [ ]
   */
  @Nonnull
  public static int[] fromJsonArray(@Nonnull JsonArray ints) {
    @Nonnull int[] array = new int[ints.size()];
    for (int i = 0; i < ints.size(); i++) {
      array[i] = ints.get(i).getAsInt();
    }
    return array;
  }

  /**
   * Reverse dimensions tensor.
   *
   * @param tensor the tensor
   * @return the tensor
   */
  @Nonnull
  public static Tensor reverseDimensions(@Nonnull Tensor tensor) {
    return tensor.rearrange(Tensor::reverse);
  }

  /**
   * Permute int [ ].
   *
   * @param key        the key
   * @param data       the data
   * @param dimensions the dimensions
   * @return the int [ ]
   */
  @Nonnull
  public static int[] permute(@Nonnull int[] key, int[] data, final int[] dimensions) {
    @Nonnull int[] copy = new int[key.length];
    for (int i = 0; i < key.length; i++) {
      int k = key[i];
      if (k == Integer.MAX_VALUE) {
        copy[i] = dimensions[0] - data[0] - 1;
      } else if (k < 0) {
        copy[i] = dimensions[-k] - data[-k] - 1;
      } else {
        copy[i] = data[k];
      }
    }
    return copy;
  }

  /**
   * Reverse int [ ].
   *
   * @param dimensions the dimensions
   * @return the int [ ]
   */
  @Nonnull
  public static int[] reverse(@Nonnull int[] dimensions) {
    @Nonnull int[] copy = Arrays.copyOf(dimensions, dimensions.length);
    ArrayUtils.reverse(copy);
    return copy;
  }

  /**
   * Pretty print string.
   *
   * @param doubles the doubles
   * @return the string
   */
  public static CharSequence prettyPrint(double[] doubles) {
    @Nonnull Tensor t = new Tensor(doubles);
    String prettyPrint = t.prettyPrint();
    t.freeRef();
    return prettyPrint;
  }

  /**
   * Get pixel double [ ].
   *
   * @param tensor the tensor
   * @param x      the x
   * @param y      the y
   * @param bands  the bands
   * @return the double [ ]
   */
  public static double[] getPixel(final Tensor tensor, final int x, final int y, final int bands) {
    return IntStream.range(0, bands).mapToDouble(band -> tensor.get(x, y, band)).toArray();
  }

  /**
   * Reduce tensor.
   *
   * @return the tensor
   */
  @Nonnull
  public Tensor sumChannels() {
    int[] dimensions = getDimensions();
    Tensor self = this;
    return new Tensor(dimensions[0], dimensions[1], 1).setByCoord(c -> {
      int[] coords = c.getCoords();
      return IntStream.range(0, dimensions[2]).mapToDouble(j -> self.get(coords[0], coords[1], j)).sum();
    });
  }

  /**
   * Gets pixel stream.
   *
   * @return the pixel stream
   */
  @Nonnull
  public Stream getPixelStream() {
    int[] dimensions = getDimensions();
    int width = dimensions[0];
    int height = dimensions[1];
    int bands = dimensions[2];
    return IntStream.range(0, width).mapToObj(x -> x).parallel().flatMap(x -> {
      return IntStream.range(0, height).mapToObj(y -> y).map(y -> {
        return getPixel(this, x, y, bands);
      });
    });
  }

  /**
   * Rescale rms tensor.
   *
   * @param rms the rms
   * @return the tensor
   */
  public Tensor rescaleRms(final double rms) {
    return scale(rms / rms());
  }

  /**
   * Normalize distribution tensor.
   *
   * @return the tensor
   */
  public Tensor normalizeDistribution() {
    double[] sortedValues = Arrays.stream(getData()).sorted().toArray();
    Tensor result = map(v -> Math.abs(((double) Arrays.binarySearch(sortedValues, v)) / ((double) sortedValues.length)));
    return result;
  }

  /**
   * Reorder dimensions tensor.
   *
   * @param fn the fn
   * @return the tensor
   */
  @Nonnull
  public Tensor rearrange(@Nonnull UnaryOperator fn) {
    return rearrange(fn, fn.apply(getDimensions()));
  }

  /**
   * Reorder dimensions tensor.
   *
   * @param fn         the fn
   * @param outputDims the output dims
   * @return the tensor
   */
  @Nonnull
  public Tensor rearrange(@Nonnull UnaryOperator fn, int[] outputDims) {
    @Nonnull Tensor result = new Tensor(outputDims);
    coordStream(false).forEach(c -> {
      int[] inCoords = c.getCoords();
      int[] outCoords = fn.apply(inCoords);
      result.set(outCoords, get(c));
    });
    return result;
  }

  /**
   * Is valid boolean.
   *
   * @return the boolean
   */
  public boolean isValid() {
    return !isFinalized() && (null == this.data || this.data.length == Tensor.length(dimensions));
  }

  /**
   * Accum.
   *
   * @param tensor the tensor
   */
  public void addInPlace(@Nonnull final Tensor tensor) {
    assert Arrays.equals(getDimensions(), tensor.getDimensions()) : Arrays.toString(getDimensions()) + " != " + Arrays.toString(tensor.getDimensions());
    setParallelByIndex(c -> get(c) + tensor.get(c));
  }

  /**
   * Add.
   *
   * @param coords the coords
   * @param value  the value
   */
  public void add(@Nonnull final Coordinate coords, final double value) {
    add(coords.getIndex(), value);
  }

  /**
   * Add tensor.
   *
   * @param index the index
   * @param value the value
   * @return the tensor
   */
  @Nonnull
  public final Tensor add(final int index, final double value) {
    getData()[index] += value;
    return this;
  }

  /**
   * Add.
   *
   * @param coords the coords
   * @param value  the value
   */
  public void add(@Nonnull final int[] coords, final double value) {
    add(index(coords), value);
  }

  /**
   * Add right.
   *
   * @param right the right
   * @return the right
   */
  @Nullable
  public Tensor add(@Nonnull final Tensor right) {
    assert Arrays.equals(getDimensions(), right.getDimensions());
    return mapCoords((c) -> get(c) + right.get(c));
  }

  /**
   * Add and free tensor.
   *
   * @param right the right
   * @return the tensor
   */
  @Nullable
  public Tensor addAndFree(@Nonnull final Tensor right) {
    assertAlive();
    right.assertAlive();
    if (1 == currentRefCount()) {
      addInPlace(right);
      return this;
    } else {
      assert Arrays.equals(getDimensions(), right.getDimensions());
      return mapCoordsAndFree((c) -> get(c) + right.get(c));
    }
  }

  /**
   * Coord stream stream.
   *
   * @param parallel the safe
   * @return the stream
   */
  @Nonnull
  public Stream coordStream(boolean parallel) {
    //ConcurrentHashSet