From e077da39c81110ef213f897b3e2c9e169976d474 Mon Sep 17 00:00:00 2001
From: Alexander Nozik <altavir@gmail.com>
Date: Sun, 9 Jun 2019 17:17:13 +0300
Subject: [PATCH] Sampler space

---
 .../kotlin/scientifik/kmath/chains/Chain.kt   | 13 +++++++-
 .../scientifik/kmath/prob/Distribution.kt     |  4 +--
 .../scientifik/kmath/prob/SamplerAlgebra.kt   | 31 +++++++++++++++++++
 3 files changed, 45 insertions(+), 3 deletions(-)
 create mode 100644 kmath-prob/src/commonMain/kotlin/scientifik/kmath/prob/SamplerAlgebra.kt

diff --git a/kmath-coroutines/src/commonMain/kotlin/scientifik/kmath/chains/Chain.kt b/kmath-coroutines/src/commonMain/kotlin/scientifik/kmath/chains/Chain.kt
index 58841e7b8..317cdcea0 100644
--- a/kmath-coroutines/src/commonMain/kotlin/scientifik/kmath/chains/Chain.kt
+++ b/kmath-coroutines/src/commonMain/kotlin/scientifik/kmath/chains/Chain.kt
@@ -37,6 +37,8 @@ interface Chain<out R> {
      */
     fun fork(): Chain<R>
 
+    companion object
+
 }
 
 /**
@@ -137,4 +139,13 @@ fun <T, S, R> Chain<T>.mapWithState(state: S, stateFork: (S) -> S, mapper: suspe
     object : Chain<R> {
         override suspend fun next(): R = state.mapper(this@mapWithState)
         override fun fork(): Chain<R> = this@mapWithState.fork().mapWithState(stateFork(state), stateFork, mapper)
-    }
\ No newline at end of file
+    }
+
+/**
+ * Zip two chains together using given transformation
+ */
+fun <T, U, R> Chain<T>.zip(other: Chain<U>, block: suspend (T, U) -> R): Chain<R> = object : Chain<R> {
+    override suspend fun next(): R = block(this@zip.next(), other.next())
+
+    override fun fork(): Chain<R> = this@zip.fork().zip(other.fork(), block)
+}
\ No newline at end of file
diff --git a/kmath-prob/src/commonMain/kotlin/scientifik/kmath/prob/Distribution.kt b/kmath-prob/src/commonMain/kotlin/scientifik/kmath/prob/Distribution.kt
index d3cac9163..487e13876 100644
--- a/kmath-prob/src/commonMain/kotlin/scientifik/kmath/prob/Distribution.kt
+++ b/kmath-prob/src/commonMain/kotlin/scientifik/kmath/prob/Distribution.kt
@@ -5,7 +5,7 @@ import scientifik.kmath.chains.map
 import kotlin.jvm.JvmName
 
 interface Sampler<T : Any> {
-    fun sample(generator: RandomGenerator): RandomChain<T>
+    fun sample(generator: RandomGenerator): Chain<T>
 }
 
 /**
@@ -22,7 +22,7 @@ interface Distribution<T : Any> : Sampler<T> {
      * Create a chain of samples from this distribution.
      * The chain is not guaranteed to be stateless.
      */
-    override fun sample(generator: RandomGenerator): RandomChain<T>
+    override fun sample(generator: RandomGenerator): Chain<T>
 
     /**
      * An empty companion. Distribution factories should be written as its extensions
diff --git a/kmath-prob/src/commonMain/kotlin/scientifik/kmath/prob/SamplerAlgebra.kt b/kmath-prob/src/commonMain/kotlin/scientifik/kmath/prob/SamplerAlgebra.kt
new file mode 100644
index 000000000..ca9d74aab
--- /dev/null
+++ b/kmath-prob/src/commonMain/kotlin/scientifik/kmath/prob/SamplerAlgebra.kt
@@ -0,0 +1,31 @@
+package scientifik.kmath.prob
+
+import scientifik.kmath.chains.Chain
+import scientifik.kmath.chains.ConstantChain
+import scientifik.kmath.chains.pipe
+import scientifik.kmath.chains.zip
+import scientifik.kmath.operations.Space
+
+class BasicSampler<T : Any>(val chainBuilder: (RandomGenerator) -> Chain<T>) : Sampler<T> {
+    override fun sample(generator: RandomGenerator): Chain<T> = chainBuilder(generator)
+}
+
+class ConstantSampler<T : Any>(val value: T) : Sampler<T> {
+    override fun sample(generator: RandomGenerator): Chain<T> = ConstantChain(value)
+}
+
+/**
+ * A space for samplers. Allows to perform simple operations on distributions
+ */
+class SamplerSpace<T : Any>(val space: Space<T>) : Space<Sampler<T>> {
+
+    override val zero: Sampler<T> = ConstantSampler(space.zero)
+
+    override fun add(a: Sampler<T>, b: Sampler<T>): Sampler<T> = BasicSampler { generator ->
+        a.sample(generator).zip(b.sample(generator)) { aValue, bValue -> space.run { aValue + bValue } }
+    }
+
+    override fun multiply(a: Sampler<T>, k: Number): Sampler<T> = BasicSampler { generator ->
+        a.sample(generator).pipe { space.run { it * k.toDouble() } }
+    }
+}
\ No newline at end of file