gt:       for (int i = 0; i < result.neighborIndices.length; i++) {
context: package edu.stanford.nlp.loglinear.inference;
  
  import edu.stanford.nlp.loglinear.model.ConcatVector;
  import edu.stanford.nlp.loglinear.model.GraphicalModel;
  import edu.stanford.nlp.loglinear.model.NDArrayDoubles;
  
  import java.util.ArrayList;
  import java.util.Iterator;
  import java.util.List;
  import java.util.function.BiFunction;
  
  /**
   * Created by keenon on 8/11/15.
   * <p>
   * Holds a factor populated by doubles that knows how to do all the important operations for PGM inference. Internally,
   * these are just different flavors of two basic data-flow operations:
   * <p>
   * - Factor product
   * - Factor marginalization
   * <p>
   * The output here is different ways to grow and shrink factors that turn out to be useful for downstream uses in PGMs.
   * Basically, we care about message passing, as that will be the primary operation.
   * <p>
   * Everything is represented as log-linear, because the primary use for TableFactor is in CliqueTree, and that is
   * intended for use with log-linear models.
   */
  public class TableFactor extends NDArrayDoubles {
    public int[] neighborIndices;
  
    /**
     * Construct a TableFactor for inference within a model. This just copies the important bits from the model factor,
     * and replaces the ConcatVectorTable with an internal datastructure that has done all the dotproducts with the
     * weights out, and so stores only doubles.
     * <p>
     * Each element of the table is given by: t_i = exp(f_i*w)
     *
     * @param weights the vector to dot product with every element of the factor table
     * @param factor  the feature factor to be multiplied in
     */
    public TableFactor(ConcatVector weights, GraphicalModel.Factor factor) {
      super(factor.featuresTable.getDimensions());
      this.neighborIndices = factor.neigborIndices;
  
      // Calculate the factor residents by dot product with the weights
  
      // OPTIMIZATION:
      // Rather than use the standard iterator, which creates lots of int[] arrays on the heap, which need to be GC'd,
      // we use the fast version that just mutates one array. Since this is read once for us here, this is ideal.
      Iterator<int[]> fastPassByReferenceIterator = factor.featuresTable.fastPassByReferenceIterator();
      int[] assignment = fastPassByReferenceIterator.next();
      while (true) {
        setAssignmentLogValue(assignment, factor.featuresTable.getAssignmentValue(assignment).get().dotProduct(weights));
        // This mutates the assignment[] array, rather than creating a new one
        if (fastPassByReferenceIterator.hasNext()) fastPassByReferenceIterator.next();
        else break;
      }
    }
  
    /**
     * Construct a TableFactor for inference within a model. This is the same as the other constructor, except that the
     * table is observed out before any unnecessary dot products are done out, so hopefully we dramatically reduce the
     * number of computations required to calculate the resulting table.
     * <p>
     * Each element of the table is given by: t_i = exp(f_i*w)
     *
     * @param weights the vector to dot product with every element of the factor table
     * @param factor  the feature factor to be multiplied in
     */
    public TableFactor(ConcatVector weights, GraphicalModel.Factor factor, int[] observations) {
      super();
      assert (observations.length == factor.neigborIndices.length);
  
      int size = 0;
      for (int observation : observations) if (observation == -1) size++;
  
      neighborIndices = new int[size];
      dimensions = new int[size];
      int[] forwardPointers = new int[size];
      int[] factorAssignment = new int[factor.neigborIndices.length];
  
      int cursor = 0;
      for (int i = 0; i < factor.neigborIndices.length; i++) {
        if (observations[i] == -1) {
          neighborIndices[cursor] = factor.neigborIndices[i];
          dimensions[cursor] = factor.featuresTable.getDimensions()[i];
          forwardPointers[cursor] = i;
          cursor++;
        } else factorAssignment[i] = observations[i];
      }
      assert (cursor == size);
  
      values = new double[combinatorialNeighborStatesCount()];
  
      for (int[] assn : this) {
        for (int i = 0; i < assn.length; i++) {
          factorAssignment[forwardPointers[i]] = assn[i];
        }
        setAssignmentLogValue(assn, factor.featuresTable.getAssignmentValue(factorAssignment).get().dotProduct(weights));
      }
    }
  
    /**
     * FOR PRIVATE USE AND TESTING ONLY
     */
    TableFactor(int[] neighborIndices, int[] dimensions) {
      super(dimensions);
      this.neighborIndices = neighborIndices;
      for (int i = 0; i < values.length; i++) {
        values[i] = Double.NEGATIVE_INFINITY;
      }
    }
  
    /**
     * Remove a variable by observing it at a certain value, return a new factor without that variable.
     *
     * @param variable the variable to be observed
     * @param value    the value the variable takes when observed
     * @return a new factor with 'variable' in it
     */
    public TableFactor observe(int variable, final int value) {
      return marginalize(variable, 0, (marginalizedVariableValue, assignment) -> {
        if (marginalizedVariableValue == value) {
          return (old, n) -> {
            // This would mean that we're observing something with 0 probability, which will wonk up downstream
            // stuff
            // assert(n != 0);
            return n;
          };
        } else {
          return (old, n) -> old;
        }
      });
    }
  
    /**
     * Returns the summed marginals for each element in the factor. These are represented in log space, and are summed
     * using the numerically stable variant, even though it's slightly slower.
     *
     * @return an array of doubles one-to-one with variable states for each variable
     */
    public double[][] getSummedMarginals() {
      double[][] results = new double[neighborIndices.length][];
      for (int i = 0; i < neighborIndices.length; i++) {
        results[i] = new double[getDimensions()[i]];
      }
  
      double[][] maxValues = new double[neighborIndices.length][];
      for (int i = 0; i < neighborIndices.length; i++) {
        maxValues[i] = new double[getDimensions()[i]];
        for (int j = 0; j < maxValues[i].length; j++) maxValues[i][j] = Double.NEGATIVE_INFINITY;
      }
  
      // Get max values
  
      // OPTIMIZATION:
      // Rather than use the standard iterator, which creates lots of int[] arrays on the heap, which need to be GC'd,
      // we use the fast version that just mutates one array. Since this is read once for us here, this is ideal.
  
      Iterator<int[]> fastPassByReferenceIterator = fastPassByReferenceIterator();
      int[] assignment = fastPassByReferenceIterator.next();
      while (true) {
        double v = getAssignmentLogValue(assignment);
        for (int i = 0; i < neighborIndices.length; i++) {
          if (maxValues[i][assignment[i]] < v) maxValues[i][assignment[i]] = v;
        }
        // This mutates the resultAssignment[] array, rather than creating a new one
        if (fastPassByReferenceIterator.hasNext()) {
          fastPassByReferenceIterator.next();
        } else break;
      }
  
      // Do the summation
  
      // OPTIMIZATION:
      // Rather than use the standard iterator, which creates lots of int[] arrays on the heap, which need to be GC'd,
      // we use the fast version that just mutates one array. Since this is read once for us here, this is ideal.
  
      Iterator<int[]> secondFastPassByReferenceIterator = fastPassByReferenceIterator();
      assignment = secondFastPassByReferenceIterator.next();
      while (true) {
        double v = getAssignmentLogValue(assignment);
        for (int i = 0; i < neighborIndices.length; i++) {
          results[i][assignment[i]] += Math.exp(v - maxValues[i][assignment[i]]);
        }
        // This mutates the resultAssignment[] array, rather than creating a new one
        if (secondFastPassByReferenceIterator.hasNext()) {
          secondFastPassByReferenceIterator.next();
        } else break;
      }
  
      // normalize results, and move to linear space
  
      for (int i = 0; i < neighborIndices.length; i++) {
        double sum = 0.0;
        for (int j = 0; j < results[i].length; j++) {
          results[i][j] = Math.exp(maxValues[i][j]) * results[i][j];
          sum += results[i][j];
        }
        if (Double.isInfinite(sum)) {
          for (int j = 0; j < results[i].length; j++) {
            results[i][j] = 1.0 / results[i].length;
          }
        } else {
          for (int j = 0; j < results[i].length; j++) {
            results[i][j] /= sum;
          }
        }
      }
  
      return results;
    }
  
    /**
     * Convenience function to max out all but one variable, and return the marginal array.
     *
     * @return an array of doubles one-to-one with variable states for each variable
     */
    public double[][] getMaxedMarginals() {
      double[][] maxValues = new double[neighborIndices.length][];
      for (int i = 0; i < neighborIndices.length; i++) {
        maxValues[i] = new double[getDimensions()[i]];
        for (int j = 0; j < maxValues[i].length; j++) maxValues[i][j] = Double.NEGATIVE_INFINITY;
      }
  
      // Get max values
  
      // OPTIMIZATION:
      // Rather than use the standard iterator, which creates lots of int[] arrays on the heap, which need to be GC'd,
      // we use the fast version that just mutates one array. Since this is read once for us here, this is ideal.
  
      Iterator<int[]> fastPassByReferenceIterator = fastPassByReferenceIterator();
      int[] assignment = fastPassByReferenceIterator.next();
      while (true) {
        double v = getAssignmentLogValue(assignment);
        for (int i = 0; i < neighborIndices.length; i++) {
          if (maxValues[i][assignment[i]] < v) maxValues[i][assignment[i]] = v;
        }
        // This mutates the resultAssignment[] array, rather than creating a new one
        if (fastPassByReferenceIterator.hasNext()) {
          fastPassByReferenceIterator.next();
        } else break;
      }
  
      for (int i = 0; i < neighborIndices.length; i++) {
        normalizeLogArr(maxValues[i]);
      }
  
      return maxValues;
    }
  
    /**
     * Marginalize out a variable by taking the max value.
     *
     * @param variable the variable to be maxed out.
     * @return a table factor that will contain the largest value of the variable being marginalized out.
     */
    public TableFactor maxOut(int variable) {
      return marginalize(variable, Double.NEGATIVE_INFINITY, (marginalizedVariableValue, assignment) -> Math::max);
    }
  
    /**
     * Marginalize out a variable by taking a sum.
     *
     * @param variable the variable to be summed out
     * @return a factor with variable removed
     */
    public TableFactor sumOut(int variable) {
  
      // OPTIMIZATION: This is by far the most common case, for linear chain inference, and is worth making fast
      // We can use closed loops, and not bother with using the basic iterator to loop through indices.
      // If this special case doesn't trip, we fall back to the standard (but slower) algorithm for the general case
  
      if (getDimensions().length == 2) {
        if (neighborIndices[0] == variable) {
          TableFactor marginalized = new TableFactor(new int[]{neighborIndices[1]}, new int[]{getDimensions()[1]});
  
          for (int i = 0; i < marginalized.values.length; i++) marginalized.values[i] = 0;
  
          // We use the stable log-sum-exp trick here, so first we calculate the max
  
          double[] max = new double[getDimensions()[1]];
          for (int j = 0; j < getDimensions()[1]; j++) {
            max[j] = Double.NEGATIVE_INFINITY;
          }
  
          for (int i = 0; i < getDimensions()[0]; i++) {
            int k = i * getDimensions()[1];
            for (int j = 0; j < getDimensions()[1]; j++) {
              int index = k + j;
              if (values[index] > max[j]) {
                max[j] = values[index];
              }
            }
          }
  
          // Then we take the sum, minus the max
  
          for (int i = 0; i < getDimensions()[0]; i++) {
            int k = i * getDimensions()[1];
            for (int j = 0; j < getDimensions()[1]; j++) {
              int index = k + j;
              if (Double.isFinite(max[j])) {
                marginalized.values[j] += Math.exp(values[index] - max[j]);
              }
            }
          }
  
          // And now we exponentiate, and add back in the values
  
          for (int j = 0; j < getDimensions()[1]; j++) {
            if (Double.isFinite(max[j])) {
              marginalized.values[j] = max[j] + Math.log(marginalized.values[j]);
            } else {
              marginalized.values[j] = max[j];
            }
          }
  
          return marginalized;
        } else {
          assert (neighborIndices[1] == variable);
          TableFactor marginalized = new TableFactor(new int[]{neighborIndices[0]}, new int[]{getDimensions()[0]});
  
          for (int i = 0; i < marginalized.values.length; i++) marginalized.values[i] = 0;
  
          // We use the stable log-sum-exp trick here, so first we calculate the max
  
          double[] max = new double[getDimensions()[0]];
          for (int i = 0; i < getDimensions()[0]; i++) {
            max[i] = Double.NEGATIVE_INFINITY;
          }
  
          for (int i = 0; i < getDimensions()[0]; i++) {
            int k = i * getDimensions()[1];
            for (int j = 0; j < getDimensions()[1]; j++) {
              int index = k + j;
              if (values[index] > max[i]) {
                max[i] = values[index];
              }
            }
          }
  
          // Then we take the sum, minus the max
  
          for (int i = 0; i < getDimensions()[0]; i++) {
            int k = i * getDimensions()[1];
            for (int j = 0; j < getDimensions()[1]; j++) {
              int index = k + j;
              if (Double.isFinite(max[i])) {
                marginalized.values[i] += Math.exp(values[index] - max[i]);
              }
            }
          }
  
          // And now we exponentiate, and add back in the values
  
          for (int i = 0; i < getDimensions()[0]; i++) {
            if (Double.isFinite(max[i])) {
              marginalized.values[i] = max[i] + Math.log(marginalized.values[i]);
            } else {
              marginalized.values[i] = max[i];
            }
          }
  
          return marginalized;
        }
      } else {
        // This is a little tricky because we need to use the stable log-sum-exp trick on top of our marginalize
        // dataflow operation.
  
        // First we calculate all the max values to use as pivots to prevent overflow
        TableFactor maxValues = maxOut(variable);
  
        // Then we do the sum against an offset from the pivots
        TableFactor marginalized = marginalize(variable, 0, (marginalizedVariableValue, assignment) -> (a, b) -> a + Math.exp(b - maxValues.getAssignmentLogValue(assignment)));
  
        // Then we factor the max values back in, and
        for (int[] assignment : marginalized) {
          marginalized.setAssignmentLogValue(assignment, maxValues.getAssignmentLogValue(assignment) + Math.log(marginalized.getAssignmentLogValue(assignment)));
        }
  
        return marginalized;
      }
    }
  
    /**
     * Product two factors, taking the multiplication at the intersections.
     *
     * @param other the other factor to be multiplied
     * @return a factor containing the union of both variable sets
     */
    public TableFactor multiply(TableFactor other) {
  
      // Calculate the result domain
  
      List<Integer> domain = new ArrayList<>();
      List<Integer> otherDomain = new ArrayList<>();
      List<Integer> resultDomain = new ArrayList<>();
  
      for (int n : neighborIndices) {
        domain.add(n);
        resultDomain.add(n);
      }
      for (int n : other.neighborIndices) {
        otherDomain.add(n);
        if (!resultDomain.contains(n)) resultDomain.add(n);
      }
  
      // Create result TableFactor
  
      int[] resultNeighborIndices = new int[resultDomain.size()];
      int[] resultDimensions = new int[resultNeighborIndices.length];
      for (int i = 0; i < resultDomain.size(); i++) {
        int var = resultDomain.get(i);
        resultNeighborIndices[i] = var;
        // assert consistency about variable size, we can't have the same variable with two different sizes
        assert ((getVariableSize(var) == 0 && other.getVariableSize(var) > 0) ||
            (getVariableSize(var) > 0 && other.getVariableSize(var) == 0) ||
            (getVariableSize(var) == other.getVariableSize(var)));
        resultDimensions[i] = Math.max(getVariableSize(resultDomain.get(i)), other.getVariableSize(resultDomain.get(i)));
      }
      TableFactor result = new TableFactor(resultNeighborIndices, resultDimensions);
  
      // OPTIMIZATION:
      // If we're a factor of size 2 receiving a message of size 1, then we can optimize that pretty heavily
      // We could just use the general algorithm at the end of this set of special cases, but this is the fastest way
      if (otherDomain.size() == 1 && (resultDomain.size() == domain.size()) && domain.size() == 2) {
        int msgVar = otherDomain.get(0);
        int msgIndex = resultDomain.indexOf(msgVar);
  
        if (msgIndex == 0) {
          for (int i = 0; i < resultDimensions[0]; i++) {
            double d = other.values[i];
            int k = i * resultDimensions[1];
            for (int j = 0; j < resultDimensions[1]; j++) {
              int index = k + j;
              result.values[index] = values[index] + d;
            }
          }
        } else if (msgIndex == 1) {
          for (int i = 0; i < resultDimensions[0]; i++) {
            int k = i * resultDimensions[1];
            for (int j = 0; j < resultDimensions[1]; j++) {
              int index = k + j;
              result.values[index] = values[index] + other.values[j];
            }
          }
        }
      }
      // OPTIMIZATION:
      // The special case where we're a message of size 1, and the other factor is receiving the message, and of size 2
      else if (domain.size() == 1 && (resultDomain.size() == otherDomain.size()) && resultDomain.size() == 2) {
        return other.multiply(this);
      }
      // Otherwise we follow the big comprehensive, slow general purpose algorithm
      else {
  
        // Calculate back-pointers from the result domain indices to original indices
  
        int[] mapping = new int[result.neighborIndices.length];
        int[] otherMapping = new int[result.neighborIndices.length];
  