<commit_before>/* Copyright (C) 2008-2010 University of Massachusetts Amherst,
   Department of Computer Science.
   This file is part of "FACTORIE" (Factor graphs, Imperative, Extensible)
   http://factorie.cs.umass.edu, http://code.google.com/p/factorie/
   Licensed 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 cc.factorie

import cc.factorie.la._
import cc.factorie.optimize._
import scala.collection.mutable.HashMap

class SuffStats extends HashMap[DotFamily, Vector] {
  override def default(template: DotFamily) = {
    template.freezeDomains
    val vector: Vector = template.weights match {
      case w: SparseVector => new SparseVector(w.length)
      case w: DenseVector => new DenseVector(w.length)
      case w: SparseOuter1DenseVector1 => new SparseOuter1DenseVector1(w.length1, w.length2)
      case w: SparseOuter2DenseVector1 => new SparseOuter2DenseVector1(w.length1, w.length2, w.length3)
    }
    this(template) = vector
    vector
  }
  // To help make sure sort order of vectors matches
  def sortedKeys = keys.toSeq.sortWith(_.hashCode > _.hashCode)
}

class LogLinearOptimizable[V <: DiscreteVarWithTarget with NoVariableCoordination](val templates: Seq[DotFamily],
                                                                                   val model: Model,
                                                                                   val variableSets: Seq[Seq[V]],
                                                                                   val gaussianPriorVariance: Double,
                                                                                   val familiesToUpdate: Seq[DotFamily],
                                                                                   val constraints: SuffStats)
          extends OptimizableFamilies(templates) with OptimizableByValueAndGradient {
  // Cached values
  private var oValue = Double.NaN
  private var oGradient: Array[Double] = new Array[Double](numOptimizableParameters)
  // Flush cache when parameters change
  override def setOptimizableParameters(a: Array[Double]): Unit = {
    oValue = Double.NaN
    model.families.foreach(_.clearCachedStatistics) // Parameter changing, so cache no longer valid
    super.setOptimizableParameters(a)
  }

  override def optimizableParameter_=(index: Int, d: Double): Unit = {oValue = Double.NaN; super.optimizableParameter_=(index, d)}
  // Calculation of value and gradient
  def setOptimizableValueAndGradient: Unit = {
    if (variableSets.forall(_.size == 1)) setOptimizableValueAndGradientIID
    else setOptimizableValueAndGradientBP
  }

  def setOptimizableValueAndGradientBP: Unit = {
    val expectations = new SuffStats
    oValue = 0.0
    java.util.Arrays.fill(oGradient, 0.0)
    variableSets.foreach(variables => {
      if (variables.size > 0) {
        val lattice = new BPLattice(variables, model)
        // Do inference on the tree
        lattice.updateTreewise(expectations)
        // For all factors // TODO Here skip factors that would have been left out in the TRP spanning tree of a loopy graph
        // TODO Note that this will only work for variables with TrueSetting.  Where to enforce this?
        variables.foreach(_.asInstanceOf[VarWithTargetValue].setToTarget(null))
        // oValue += model.factors(variables).foldLeft(0.0)(_+_.cachedStatistics.score) - logZ
        for (bpfactor <- lattice.bpFactors.values) oValue += bpfactor.factor.cachedStatistics.score
        oValue -= lattice.sumLogZ
      }
    })
    val invVariance = -1.0 / gaussianPriorVariance
    familiesToUpdate.foreach {
      t =>
        oValue += 0.5 * t.weights.dot(t.weights) * invVariance
        // sum positive constraints into (previously negated) expectations
        expectations(t) += constraints(t)
        //vecPlusEq(expectations(t), constraints(t), 1.0)
        // subtract weights due to regularization
        expectations(t) += t.weights * invVariance
    }
    // constraints.keys.foreach(t => expectations(t) += constraints(t))
    oGradient = (new ArrayFromVectors(expectations.sortedKeys.map(expectations(_)))).getVectorsInArray(oGradient)
  }

  def setOptimizableValueAndGradientIID: Unit = {
    val expectations = new SuffStats
    oValue = 0.0
    java.util.Arrays.fill(oGradient, 0.0)
    var distribution = Array.fill(0)(0.0)
    variableSets.foreach(_.foreach(v => {
      if (v.domain.size != distribution.length) distribution = new Array[Double](v.domain.size)
      forIndex(distribution.length)(i => {
        v.set(i)(null)
        distribution(i) = model.score(Seq(v))
      })

      maths.expNormalize(distribution)

      forIndex(distribution.length)(i => {
        v.set(i)(null)
        // put negative expectations into 'expectations' StatMap
        model.factorsOfFamilies(Seq(v), familiesToUpdate).foreach(f => expectations(f.family) +=  f.statistics.vector *(-distribution(i)))
      })

      oValue += math.log(distribution(v.targetIntValue))
    }))
    val invVariance = -1.0 / gaussianPriorVariance
    familiesToUpdate.foreach {
      t =>
        oValue += 0.5 * t.weights.dot(t.weights) * invVariance
        // sum positive constraints into (previously negated) expectations
        expectations(t) += constraints(t)
        // subtract weights due to regularization
        expectations(t) += t.weights * invVariance
    }
    // constraints.keys.foreach(t => expectations(t) += constraints(t))
    oGradient = (new ArrayFromVectors(expectations.sortedKeys.map(expectations(_)))).getVectorsInArray(oGradient)
  }

  def optimizableValue: Double = {
    if (oValue.isNaN) setOptimizableValueAndGradient
    oValue
  }

  def getOptimizableGradient(a: Array[Double]) = {
    if (oValue.isNaN) setOptimizableValueAndGradient
    Array.copy(oGradient, 0, a, 0, oGradient.length)
  }
}

// TODO preliminary: currently only supports trivial inference on IID discrete variables
// and inference on tree-shaped graphs of discrete variables.
// In the future there will be a choice of different inference methods over arbitrary graphical model structures

/**
 * Maximum likelihood parameter estimation for the weights of DotTemplate.
 * @author <NAME>, <NAME>, <NAME>
 */
class LogLinearMaximumLikelihood(model: Model) {
  //type TemplatesToUpdate = DotTemplate
  var gaussianPriorVariance = 10.0
  def familiesToUpdate: Seq[DotFamily] = model.familiesOfClass(classOf[DotFamily])

  //def process[V <: DiscreteVariableWithTrueSetting with NoVariableCoordination](variables: Seq[V], numIterations: Int): Unit = process(List(variables), numIterations)
  // TODO Figure out how to reinstate something like this.
  //def process[V <: DiscreteVariableWithTrueSetting with NoVariableCoordination](variables: Seq[V]): Unit = process(List(variables), Int.MaxValue)

  def process[V <: DiscreteVarWithTarget with NoVariableCoordination](variableSet: Seq[V], numIterations: Int = Int.MaxValue): Unit = 
    processAll(List(variableSet), numIterations)

  /**First argument is a collection of collections-of-variables.  The former are considered iid.  The later may have dependencies.  */
  def processAll[V <: DiscreteVarWithTarget with NoVariableCoordination](variableSets: Seq[Seq[V]], numIterations: Int = Int.MaxValue): Unit = {
    // Data structure for holding per-template constraints and expectations

    val constraints = new SuffStats
    // Add all model dot templates to constraints
    //familiesToUpdate.foreach(t => constraints(t) = constraints.default(t)) // TODO Why is this line necessary? Delete it? -akm
    // Gather constraints
    variableSets.foreach(_.foreach(_.setToTarget(null)))
    variableSets.foreach(vars => model.factorsOfFamilies(vars, familiesToUpdate).foreach(f => { val vector = f.cachedStatistics.vector ; constraints(f.family) += vector }))

    def templates = constraints.sortedKeys

    // Currently only supports iid single DiscreteVariables
    val optimizable = new LogLinearOptimizable(templates, model, variableSets, gaussianPriorVariance, familiesToUpdate, constraints)

    // Do the gradient-climbing optimization!
    try {
      val optimizer = new LimitedMemoryBFGS(optimizable)
      optimizer.optimize(numIterations)
    }
    catch {
      case e : Error => e.printStackTrace
    }
    // Resetting and running again sometimes improves results
    try {
      val optimizer = new LimitedMemoryBFGS(optimizable)
      optimizer.optimize(numIterations)
    }
    catch {
      case e : Error => e.printStackTrace
    }
  }
}

<commit_msg>Refactoring LogLinearMaximumLikelihood to make it more easily profileable<commit_after>/* Copyright (C) 2008-2010 University of Massachusetts Amherst,
   Department of Computer Science.
   This file is part of "FACTORIE" (Factor graphs, Imperative, Extensible)
   http://factorie.cs.umass.edu, http://code.google.com/p/factorie/
   Licensed 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 cc.factorie

import cc.factorie.la._
import cc.factorie.optimize._
import scala.collection.mutable.HashMap

class SuffStats extends HashMap[DotFamily, Vector] {
  override def default(template: DotFamily) = {
    template.freezeDomains
    val vector: Vector = template.weights match {
      case w: SparseVector => new SparseVector(w.length)
      case w: DenseVector => new DenseVector(w.length)
      case w: SparseOuter1DenseVector1 => new SparseOuter1DenseVector1(w.length1, w.length2)
      case w: SparseOuter2DenseVector1 => new SparseOuter2DenseVector1(w.length1, w.length2, w.length3)
    }
    this(template) = vector
    vector
  }
  // To help make sure sort order of vectors matches
  def sortedKeys = keys.toSeq.sortWith(_.hashCode > _.hashCode)
}

class LogLinearOptimizable[V <: DiscreteVarWithTarget with NoVariableCoordination](val templates: Seq[DotFamily],
                                                                                   val model: Model,
                                                                                   val variableSets: Seq[Seq[V]],
                                                                                   val gaussianPriorVariance: Double,
                                                                                   val familiesToUpdate: Seq[DotFamily],
                                                                                   val constraints: SuffStats)
          extends OptimizableFamilies(templates) with OptimizableByValueAndGradient {
  // Cached values
  private var oValue = Double.NaN
  private var oGradient: Array[Double] = new Array[Double](numOptimizableParameters)
  // Flush cache when parameters change
  override def setOptimizableParameters(a: Array[Double]): Unit = {
    oValue = Double.NaN
    model.families.foreach(_.clearCachedStatistics) // Parameter changing, so cache no longer valid
    super.setOptimizableParameters(a)
  }

  override def optimizableParameter_=(index: Int, d: Double): Unit = {oValue = Double.NaN; super.optimizableParameter_=(index, d)}
  // Calculation of value and gradient

  val doSetOptimizableValueAndGradient = if (variableSets.forall(_.size == 1)) { () => setOptimizableValueAndGradientIID  } else { () => setOptimizableValueAndGradientBP}

  def setOptimizableValueAndGradient: Unit = doSetOptimizableValueAndGradient()

  def updateOValueBP(expectations: SuffStats) {
    variableSets.foreach(variables => {
      if (variables.size > 0) {
        val lattice = new BPLattice(variables, model)
        // Do inference on the tree
        lattice.updateTreewise(expectations)
        // For all factors // TODO Here skip factors that would have been left out in the TRP spanning tree of a loopy graph
        // TODO Note that this will only work for variables with TrueSetting.  Where to enforce this?
        variables.foreach(_.asInstanceOf[VarWithTargetValue].setToTarget(null))
        // oValue += model.factors(variables).foldLeft(0.0)(_+_.cachedStatistics.score) - logZ
        for (bpfactor <- lattice.bpFactors.values) oValue += bpfactor.factor.cachedStatistics.score
        oValue -= lattice.sumLogZ
      }
    })
  }

  val invVariance = -1.0 / gaussianPriorVariance
  def updateExpectations(expectations: SuffStats, constraints: SuffStats) {
    familiesToUpdate.foreach {
      t =>
        oValue += 0.5 * t.weights.dot(t.weights) * invVariance
        // sum positive constraints into (previously negated) expectations
        expectations(t) += constraints(t)
        //vecPlusEq(expectations(t), constraints(t), 1.0)
        // subtract weights due to regularization
        expectations(t) += t.weights * invVariance
    }
  }

  def setOptimizableValueAndGradientBP: Unit = {
    val expectations = new SuffStats
    oValue = 0.0
    //java.util.Arrays.fill(oGradient, 0.0)
    updateOValueBP(expectations)
    updateExpectations(expectations, constraints)
    // constraints.keys.foreach(t => expectations(t) += constraints(t))
    oGradient = (new ArrayFromVectors(expectations.sortedKeys.map(expectations(_)))).getVectorsInArray(oGradient)
  }

  var distribution = Array.fill(0)(0.0)
  def doOValueIID(expectations: SuffStats) {
    variableSets.foreach(_.foreach(v => {
      if (v.domain.size != distribution.length) distribution = new Array[Double](v.domain.size)
      var i = 0
      while (i < distribution.length) {
        v.set(i)(null)
        distribution(i) = model.score(Seq(v))
        i += 1
      }

      maths.expNormalize(distribution)
      i = 0
      while (i < distribution.length) {
        v.set(i)(null)
        // put negative expectations into 'expectations' StatMap
        model.factorsOfFamilies(Seq(v), familiesToUpdate).foreach(f => expectations(f.family) +=  f.statistics.vector *(-distribution(i)))
        i += 1
      }

      oValue += math.log(distribution(v.targetIntValue))
    }))
  }

  def setOptimizableValueAndGradientIID: Unit = {
    val expectations = new SuffStats
    oValue = 0.0
    java.util.Arrays.fill(oGradient, 0.0)
    doOValueIID(expectations)
    updateExpectations(expectations, constraints)
    // constraints.keys.foreach(t => expectations(t) += constraints(t))
    oGradient = (new ArrayFromVectors(expectations.sortedKeys.map(expectations(_)))).getVectorsInArray(oGradient)
  }

  def optimizableValue: Double = {
    if (oValue.isNaN) setOptimizableValueAndGradient
    oValue
  }

  def getOptimizableGradient(a: Array[Double]) = {
    if (oValue.isNaN) setOptimizableValueAndGradient
    Array.copy(oGradient, 0, a, 0, oGradient.length)
  }
}

// TODO preliminary: currently only supports trivial inference on IID discrete variables
// and inference on tree-shaped graphs of discrete variables.
// In the future there will be a choice of different inference methods over arbitrary graphical model structures

/**
 * Maximum likelihood parameter estimation for the weights of DotTemplate.
 * @author <NAME>, <NAME>, <NAME>
 */
class LogLinearMaximumLikelihood(model: Model) {
  //type TemplatesToUpdate = DotTemplate
  var gaussianPriorVariance = 10.0
  def familiesToUpdate: Seq[DotFamily] = model.familiesOfClass(classOf[DotFamily])

  //def process[V <: DiscreteVariableWithTrueSetting with NoVariableCoordination](variables: Seq[V], numIterations: Int): Unit = process(List(variables), numIterations)
  // TODO Figure out how to reinstate something like this.
  //def process[V <: DiscreteVariableWithTrueSetting with NoVariableCoordination](variables: Seq[V]): Unit = process(List(variables), Int.MaxValue)

  def process[V <: DiscreteVarWithTarget with NoVariableCoordination](variableSet: Seq[V], numIterations: Int = Int.MaxValue): Unit = 
    processAll(List(variableSet), numIterations)

  /**First argument is a collection of collections-of-variables.  The former are considered iid.  The later may have dependencies.  */
  def processAll[V <: DiscreteVarWithTarget with NoVariableCoordination](variableSets: Seq[Seq[V]], numIterations: Int = Int.MaxValue): Unit = {
    // Data structure for holding per-template constraints and expectations

    val constraints = new SuffStats
    // Add all model dot templates to constraints
    //familiesToUpdate.foreach(t => constraints(t) = constraints.default(t)) // TODO Why is this line necessary? Delete it? -akm
    // Gather constraints
    variableSets.foreach(_.foreach(_.setToTarget(null)))
    variableSets.foreach(vars => model.factorsOfFamilies(vars, familiesToUpdate).foreach(f => { val vector = f.cachedStatistics.vector ; constraints(f.family) += vector }))

    def templates = constraints.sortedKeys

    // Currently only supports iid single DiscreteVariables
    val optimizable = new LogLinearOptimizable(templates, model, variableSets, gaussianPriorVariance, familiesToUpdate, constraints)

    // Do the gradient-climbing optimization!
    try {
      val optimizer = new LimitedMemoryBFGS(optimizable)
      optimizer.optimize(numIterations)
    }
    catch {
      case e : Error => e.printStackTrace
    }
    // Resetting and running again sometimes improves results
    try {
      val optimizer = new LimitedMemoryBFGS(optimizable)
      optimizer.optimize(numIterations)
    }
    catch {
      case e : Error => e.printStackTrace
    }
  }
}


{'fraction_non_alphanumeric': 0.0925068870523416, 'fraction_numerical': 0.006060606060606061, 'mean_word_length': 3.105632209907261, 'pattern_counts': {'":': 0, '<': 23, '<?xml version=': 0, '>': 48, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 42, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1789885', 'n_tokens_mistral': 5075, 'n_tokens_neox': 4874, 'n_words': 1359}