Stateless optimisers and HMCMC

.scalafmt.conf

@@ -11,6 +11,11 @@ align.tokens = [{code = "="}, {code = "<-"}, {code = ":="}, {code = "->"}, {code
 continuationIndent.callSite = 2
 continuationIndent.defnSite = 2
 optIn.breakChainOnFirstMethodDot = true
+rewrite.rules = [SortModifiers]
+rewrite.sortModifiers.order = [
+  "implicit", "final", "sealed", "abstract",
+  "override", "private", "protected", "lazy"
+]
 
 spaces {
   inImportCurlyBraces = true

build.sbt

@@ -1,4 +1,4 @@
-ThisBuild / baseVersion := "0.8.0"
+ThisBuild / baseVersion := "0.9.0"
 
 ThisBuild / organization      := "ai.entrolution"
 ThisBuild / organizationName  := "Greg von Nessi"

project/Dependencies.scala

@@ -3,8 +3,8 @@ import sbt.*
 object DependencyVersions {
   val scala2p13Version = "2.13.11"
 
-  val bengalStmVersion           = "0.9.3"
-  val bigMathVersion             = "2.3.0"
+  val bengalStmVersion           = "0.9.4"
+  val bigMathVersion             = "2.3.2"
   val breezeVersion              = "2.1.0"
   val catsEffectVersion          = "3.4.8"
   val catsEffectTestingVersion   = "1.4.0"

src/main/scala/thylacine/config/ConjugateGradientConfig.scala

@@ -21,5 +21,5 @@ case class ConjugateGradientConfig(
   convergenceThreshold: Double,
   goldenSectionTolerance: Double,
   lineProbeExpansionFactor: Double,
-  numberOfResultsToRetain: Int
+  minimumNumberOfIterations: Int
 )

src/main/scala/thylacine/config/HmcmcConfig.scala

@@ -21,6 +21,5 @@ case class HmcmcConfig(
   stepsBetweenSamples: Int,
   stepsInDynamicsSimulation: Int,
   warmupStepCount: Int,
-  dynamicsSimulationStepSize: Double,
-  sampleParallelism: Option[Int]
+  dynamicsSimulationStepSize: Double
 )

src/main/scala/thylacine/model/components/forwardmodel/ForwardModel.scala

@@ -31,7 +31,7 @@ private[thylacine] trait ForwardModel[F[_]] extends GenericMapping with CanValid
     input: IndexedVectorCollection
   ): F[VectorContainer]
 
-  private[thylacine] final def evalAt(input: Map[String, Vector[Double]]): F[Vector[Double]] =
+  final private[thylacine] def evalAt(input: Map[String, Vector[Double]]): F[Vector[Double]] =
     evalAt(IndexedVectorCollection(input)).map(_.scalaVector)
 
   // Note that input validation should be done within
@@ -40,7 +40,7 @@ private[thylacine] trait ForwardModel[F[_]] extends GenericMapping with CanValid
     input: IndexedVectorCollection
   ): F[IndexedMatrixCollection]
 
-  private[thylacine] final def jacobianAt(input: Map[String, Vector[Double]]): F[Map[String, Vector[Vector[Double]]]] =
+  final private[thylacine] def jacobianAt(input: Map[String, Vector[Double]]): F[Map[String, Vector[Vector[Double]]]] =
     jacobianAt(IndexedVectorCollection(input))
       .map(_.genericScalaRepresentation)
 }

src/main/scala/thylacine/model/components/forwardmodel/InMemoryMemoizedForwardModel.scala

@@ -28,12 +28,12 @@ private[thylacine] trait InMemoryMemoizedForwardModel[F[_]] extends ForwardModel
   protected val evalCache: CachedComputation[F, VectorContainer]
   protected val jacobianCache: CachedComputation[F, IndexedMatrixCollection]
 
-  private[thylacine] override final def evalAt(
+  final override private[thylacine] def evalAt(
     input: ModelParameterCollection
   ): F[VectorContainer] =
     evalCache.performComputation(input)
 
-  private[thylacine] override def jacobianAt(
+  override private[thylacine] def jacobianAt(
     input: ModelParameterCollection
   ): F[IndexedMatrixCollection] =
     jacobianCache.performComputation(input)

src/main/scala/thylacine/model/components/forwardmodel/LinearForwardModel.scala

@@ -33,13 +33,13 @@ import scala.annotation.unused
 // A linear forward model may work across more than
 // one model parameter generator
 case class LinearForwardModel[F[_]: STM: Async](
-  protected override val evalCache: CachedComputation[F, VectorContainer],
-  protected override val jacobianCache: CachedComputation[F, IndexedMatrixCollection],
+  override protected val evalCache: CachedComputation[F, VectorContainer],
+  override protected val jacobianCache: CachedComputation[F, IndexedMatrixCollection],
   private[thylacine] val transform: IndexedMatrixCollection,
   private[thylacine] val vectorOffset: Option[VectorContainer],
   override val domainDimension: Int,
   override val rangeDimension: Int,
-  private[thylacine] override val validated: Boolean = false
+  override private[thylacine] val validated: Boolean = false
 ) extends StmImplicits[F]
     with InMemoryMemoizedForwardModel[F] {
   if (!validated) {
@@ -49,7 +49,7 @@ case class LinearForwardModel[F[_]: STM: Async](
     )
   }
 
-  private[thylacine] override lazy val getValidated: LinearForwardModel[F] =
+  override private[thylacine] lazy val getValidated: LinearForwardModel[F] =
     if (validated) {
       this
     } else {
@@ -65,7 +65,7 @@ case class LinearForwardModel[F[_]: STM: Async](
   private[thylacine] val getJacobian: IndexedMatrixCollection =
     transform
 
-  private[thylacine] override def jacobianAt(
+  override private[thylacine] def jacobianAt(
     input: IndexedVectorCollection
   ): F[IndexedMatrixCollection] =
     Async[F].pure(getJacobian)

src/main/scala/thylacine/model/components/forwardmodel/NonLinearForwardModel.scala

@@ -26,15 +26,15 @@ import cats.effect.kernel.Async
 import cats.syntax.all._
 
 case class NonLinearForwardModel[F[_]: STM: Async](
-  protected override val evalCache: CachedComputation[F, VectorContainer],
-  protected override val jacobianCache: CachedComputation[F, IndexedMatrixCollection],
+  override protected val evalCache: CachedComputation[F, VectorContainer],
+  override protected val jacobianCache: CachedComputation[F, IndexedMatrixCollection],
   private val domainDimensions: Map[String, Int],
   override val rangeDimension: Int,
-  private[thylacine] override val validated: Boolean = false
+  override private[thylacine] val validated: Boolean = false
 ) extends StmImplicits[F]
     with InMemoryMemoizedForwardModel[F] {
 
-  private[thylacine] override val getValidated = this
+  override private[thylacine] val getValidated = this
 
   override val domainDimension: Int = domainDimensions.values.sum
 }

src/main/scala/thylacine/model/components/likelihood/CauchyLikelihood.scala

@@ -29,24 +29,24 @@ import java.util.UUID
 import scala.annotation.unused
 
 case class CauchyLikelihood[F[_]: Async](
-  private[thylacine] override val posteriorTermIdentifier: TermIdentifier,
+  override private[thylacine] val posteriorTermIdentifier: TermIdentifier,
   private[thylacine] val observations: RecordedData,
-  private[thylacine] override val forwardModel: ForwardModel[F],
-  private[thylacine] override val validated: Boolean = false
+  override private[thylacine] val forwardModel: ForwardModel[F],
+  override private[thylacine] val validated: Boolean = false
 ) extends AsyncImplicits[F]
     with Likelihood[F, ForwardModel[F], CauchyDistribution] {
   if (!validated) {
     assert(forwardModel.rangeDimension == observations.data.dimension)
   }
 
-  private[thylacine] override lazy val getValidated: CauchyLikelihood[F] =
+  override private[thylacine] lazy val getValidated: CauchyLikelihood[F] =
     if (validated) {
       this
     } else {
       this.copy(observations = observations.getValidated, forwardModel = forwardModel.getValidated, validated = true)
     }
 
-  private[thylacine] override lazy val observationDistribution: CauchyDistribution =
+  override private[thylacine] lazy val observationDistribution: CauchyDistribution =
     CauchyDistribution(observations)
 
 }

src/main/scala/thylacine/model/components/likelihood/GaussianLikelihood.scala

@@ -28,24 +28,24 @@ import cats.effect.kernel.Async
 import java.util.UUID
 
 case class GaussianLikelihood[F[_]: Async, T <: ForwardModel[F]](
-  private[thylacine] override val posteriorTermIdentifier: TermIdentifier,
+  override private[thylacine] val posteriorTermIdentifier: TermIdentifier,
   private[thylacine] val observations: RecordedData,
-  private[thylacine] override val forwardModel: T,
-  private[thylacine] override val validated: Boolean = false
+  override private[thylacine] val forwardModel: T,
+  override private[thylacine] val validated: Boolean = false
 ) extends AsyncImplicits[F]
     with Likelihood[F, T, GaussianDistribution] {
   if (!validated) {
     assert(forwardModel.rangeDimension == observations.data.dimension)
   }
 
-  private[thylacine] override lazy val getValidated: GaussianLikelihood[F, T] =
+  override private[thylacine] lazy val getValidated: GaussianLikelihood[F, T] =
     if (validated) {
       this
     } else {
       this.copy(observations = observations.getValidated, validated = true)
     }
 
-  private[thylacine] override lazy val observationDistribution: GaussianDistribution =
+  override private[thylacine] lazy val observationDistribution: GaussianDistribution =
     GaussianDistribution(observations)
 
 }

src/main/scala/thylacine/model/components/likelihood/GaussianLinearLikelihood.scala

@@ -34,24 +34,24 @@ import scala.annotation.unused
 // Thus, it gets a dedicated case class that is leveraged to do an analytic check in the posterior
 // construction.
 case class GaussianLinearLikelihood[F[_]: Async](
-  private[thylacine] override val posteriorTermIdentifier: TermIdentifier,
+  override private[thylacine] val posteriorTermIdentifier: TermIdentifier,
   private[thylacine] val observations: RecordedData,
-  private[thylacine] override val forwardModel: LinearForwardModel[F],
-  private[thylacine] override val validated: Boolean = false
+  override private[thylacine] val forwardModel: LinearForwardModel[F],
+  override private[thylacine] val validated: Boolean = false
 ) extends AsyncImplicits[F]
     with Likelihood[F, LinearForwardModel[F], GaussianDistribution] {
   if (!validated) {
     assert(forwardModel.rangeDimension == observations.data.dimension)
   }
 
-  private[thylacine] override lazy val getValidated: GaussianLinearLikelihood[F] =
+  override private[thylacine] lazy val getValidated: GaussianLinearLikelihood[F] =
     if (validated) {
       this
     } else {
       this.copy(observations = observations.getValidated, validated = true)
     }
 
-  private[thylacine] override lazy val observationDistribution: GaussianDistribution =
+  override private[thylacine] lazy val observationDistribution: GaussianDistribution =
     GaussianDistribution(observations)
 
 }

src/main/scala/thylacine/model/components/likelihood/Likelihood.scala

@@ -37,18 +37,18 @@ private[thylacine] trait Likelihood[F[_], +FM <: ForwardModel[F], +D <: Distribu
   private[thylacine] def observationDistribution: D
   private[thylacine] def forwardModel: FM
 
-  override final val domainDimension: Int =
+  final override val domainDimension: Int =
     forwardModel.domainDimension
 
-  private[thylacine] override final def logPdfAt(
+  final override private[thylacine] def logPdfAt(
     input: ModelParameterCollection
   ): F[Double] =
     for {
       mappedVec <- forwardModel.evalAt(input)
       res       <- Async[F].delay(observationDistribution.logPdfAt(mappedVec))
     } yield res
 
-  private[thylacine] override final def logPdfGradientAt(
+  final override private[thylacine] def logPdfGradientAt(
     input: ModelParameterCollection
   ): F[ModelParameterCollection] =
     for {

src/main/scala/thylacine/model/components/likelihood/UniformLikelihood.scala

@@ -30,18 +30,18 @@ import java.util.UUID
 import scala.annotation.unused
 
 case class UniformLikelihood[F[_]: Async, T <: ForwardModel[F]](
-  private[thylacine] override val posteriorTermIdentifier: TermIdentifier,
+  override private[thylacine] val posteriorTermIdentifier: TermIdentifier,
   private[thylacine] val upperBounds: VectorContainer,
   private[thylacine] val lowerBounds: VectorContainer,
-  private[thylacine] override val forwardModel: T,
-  private[thylacine] override val validated: Boolean = false
+  override private[thylacine] val forwardModel: T,
+  override private[thylacine] val validated: Boolean = false
 ) extends AsyncImplicits[F]
     with Likelihood[F, T, UniformDistribution] {
   if (!validated) {
     assert(lowerBounds.dimension == upperBounds.dimension)
   }
 
-  private[thylacine] override lazy val getValidated: UniformLikelihood[F, T] =
+  override private[thylacine] lazy val getValidated: UniformLikelihood[F, T] =
     if (validated) {
       this
     } else {
@@ -52,7 +52,7 @@ case class UniformLikelihood[F[_]: Async, T <: ForwardModel[F]](
       )
     }
 
-  private[thylacine] override lazy val observationDistribution: UniformDistribution =
+  override private[thylacine] lazy val observationDistribution: UniformDistribution =
     distributions.UniformDistribution(upperBounds = upperBounds, lowerBounds = lowerBounds)
 
 }

src/main/scala/thylacine/model/components/posterior/ConjugateGradientOptimisedPosterior.scala

@@ -21,46 +21,47 @@ import thylacine.config.ConjugateGradientConfig
 import thylacine.model.components.likelihood.Likelihood
 import thylacine.model.components.prior.Prior
 import thylacine.model.core.AsyncImplicits
+import thylacine.model.core.telemetry.OptimisationTelemetryUpdate
 import thylacine.model.optimization.gradientdescent.ConjugateGradientEngine
 
 import cats.effect.kernel.Async
 
 case class ConjugateGradientOptimisedPosterior[F[_]: Async](
   private[thylacine] val gradientDescentConfig: ConjugateGradientConfig,
-  protected override val newMaximumCallback: Double => F[Unit],
-  protected override val isConvergedCallback: Unit => F[Unit],
-  private[thylacine] override val priors: Set[Prior[F, _]],
-  private[thylacine] override val likelihoods: Set[Likelihood[F, _, _]]
+  override protected val iterationUpdateCallback: OptimisationTelemetryUpdate => F[Unit],
+  override protected val isConvergedCallback: Unit => F[Unit],
+  override private[thylacine] val priors: Set[Prior[F, _]],
+  override private[thylacine] val likelihoods: Set[Likelihood[F, _, _]]
 ) extends AsyncImplicits[F]
     with Posterior[F, Prior[F, _], Likelihood[F, _, _]]
     with ConjugateGradientEngine[F] {
 
   override protected val convergenceThreshold: Double =
     gradientDescentConfig.convergenceThreshold
 
-  override protected val numberOfResultsToRetain: Int =
-    gradientDescentConfig.numberOfResultsToRetain
-
   override protected val goldenSectionTolerance: Double =
     gradientDescentConfig.goldenSectionTolerance
 
   override protected val lineProbeExpansionFactor: Double =
     gradientDescentConfig.lineProbeExpansionFactor
+
+  override protected def minimumNumberOfIterations: Int =
+    gradientDescentConfig.minimumNumberOfIterations
 }
 
 object ConjugateGradientOptimisedPosterior {
 
   def from[F[_]: Async](
     conjugateGradientConfig: ConjugateGradientConfig,
     posterior: Posterior[F, Prior[F, _], Likelihood[F, _, _]],
-    newMaximumCallback: Double => F[Unit],
+    iterationUpdateCallback: OptimisationTelemetryUpdate => F[Unit],
     isConvergedCallback: Unit => F[Unit]
   ): ConjugateGradientOptimisedPosterior[F] =
     ConjugateGradientOptimisedPosterior(
-      gradientDescentConfig = conjugateGradientConfig,
-      newMaximumCallback    = newMaximumCallback,
-      isConvergedCallback   = isConvergedCallback,
-      priors                = posterior.priors,
-      likelihoods           = posterior.likelihoods
+      gradientDescentConfig   = conjugateGradientConfig,
+      iterationUpdateCallback = iterationUpdateCallback,
+      isConvergedCallback     = isConvergedCallback,
+      priors                  = posterior.priors,
+      likelihoods             = posterior.likelihoods
     )
 }