stripe / rainier

Bayesian inference in Scala.
https://rainier.fit
Apache License 2.0
433 stars 51 forks source link

Proposal: new type approach for Generator, RandomVariable etc #391

Closed avi-stripe closed 5 years ago

avi-stripe commented 5 years ago

I'll add comments to explain this in more detail soon.

trait RandomVariable[F[_]] {
    type G[_]

    def map[T, U](rv: F[T])(f: T => U): F[U]
    def zip[T, U](rvt: F[T], rvu: F[U]): F[(T,U)]
    def zip[T, U](rv: F[T], g: Generator[U]): G[(T,U)]
    def flatMap[T, U](rv: F[T])(f: T => Generator[U]): G[U]
    def sample[T, X](rv: F[T], k: Int)(implicit dec: Decoder[T,X], rng: RNG): List[X]
}

trait Generator[T]
object Generator {
    implicit val rv: RandomVariable[Generator] = new RandomVariable[Generator] {
        type G[_] = Generator[_]
        //...
    }
}

trait Model[T] {   
    def flatMap[U](f: T => Model[U]): Model[U]
    def sample(sampler: Sampler)(implicit rng: RNG): Sample[T]
    def optimize[X]()(implicit dec: Decoder[T,X]): X
}
object Model {
    implicit val rv: RandomVariable[Model] = new RandomVariable[Model] {
        type G[_] = Prediction[_]
        //...
    }
}

trait Prediction[T] {        
    def sample(sampler: Sampler)(implicit rng: RNG): Sample[T]
}
object Prediction {
    implicit val rv: RandomVariable[Prediction] = new RandomVariable[Prediction] {
        type G[_] = Prediction[_]
        //...
    }
}

trait Sample[T] {    
    def toList[X](implicit dec: Decoder[T,X]): List[X]
}
object Sample {
    implicit val rv: RandomVariable[Sample] = new RandomVariable[Sample] {
        type G[_] = Sample[_]
        //...
    }
}