-
Notifications
You must be signed in to change notification settings - Fork 4
/
Categorical.scala
64 lines (58 loc) · 2.2 KB
/
Categorical.scala
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
/*
* Copyright 2021 Arman Bilge
*
* 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 schrodinger.kernel
import algebra.Priority
import algebra.ring.AdditiveMonoid
import cats.Functor
import cats.Invariant
import cats.NonEmptyTraverse
import cats.Reducible
import cats.collections.HashMap
import cats.data.Chain
import cats.data.NonEmptyVector
import cats.kernel.Hash
import cats.kernel.Order
import cats.syntax.all.*
import schrodinger.math.syntax.*
trait Categorical[F[_], V, I]:
def categorical(probabilites: V): F[I]
object Categorical:
inline def apply[F[_], V, I](probabilites: V)(using
c: Categorical[F, V, I],
): F[I] = c.categorical(probabilites)
def apply[F[_], P, G[_]: NonEmptyTraverse, A](support: G[(A, P)])(using
F: Priority[Functor[F], InvariantAndHash[F, A]],
P: AdditiveMonoid[P],
c: Categorical[F, G[P], Long],
): F[A] =
val probabilities = support.map(_._2)
F match
case Priority.Preferred(given Functor[f]) =>
apply(probabilities).map(support.get(_).get._1)
case Priority.Fallback(InvariantAndHash(given Invariant[f], given Hash[a])) =>
val inv = HashMap.fromIterableOnce(support.toIterable.view.map(_._1).zipWithIndex)
apply(probabilities).imap(support.get(_).get._1)(inv.getOrElse(_, -1))
given [F[_]: Functor, G[_]: Reducible, P: Order](using
P: AdditiveMonoid[P],
u: Uniform[F, P],
): Categorical[F, G[P], Long] with
def categorical(probabilities: G[P]) =
val cumSum =
probabilities.reduceLeftTo(Chain.one(_))((c, p) => c :+ (c.lastOption.get + p)).toVector
val sum = cumSum.last
Uniform(P.zero, sum).map(
cumSum.search(_)(using Order[P].toOrdering).insertionPoint,
)