Fixed radius near neighbours in elm

FixedRadius.elm

module FixedRadius exposing (Buckets, Point, fromList, search, searchReflexive, map)

{-| A module for fixed-radious near neighbours

@docs Point, Buckets, fromList, search, searchReflexive, map

-}

import Dict exposing (Dict)


{-| A data point
-}
type alias Point a =
    { a | x : Float, y : Float }


type alias Bucket =
    ( Int, Int )


{-| Opaque type
-}
type Buckets a
    = Buckets { radius : Float, buckets : Dict Bucket (List (Point a)) }


{-| Construct `Buckets` from a radius and a list of points
-}
fromList : Float -> List (Point a) -> Buckets a
fromList radius points =
    let
        buckets =
            List.foldl (insert radius) Dict.empty points
    in
        Buckets { radius = radius, buckets = buckets }


{-| Search for unique pairs of close vertices: if `(A, B)` is in the output, then `(B, A)` never is.
-}
search : Buckets a -> List ( Point a, Point a )
search ((Buckets { buckets }) as wrapped) =
    List.concatMap (\b -> pair b wrapped) (Dict.keys buckets)


{-| Search for pairs of close vertices: if `(A, B)` is in the output, then `(B, A)` is too.
-}
searchReflexive : Buckets a -> List ( Point a, Point a )
searchReflexive ((Buckets { buckets }) as wrapped) =
    List.concatMap (\b -> pairReflexive b wrapped) (Dict.keys buckets)


insert : Float -> Point a -> Dict Bucket (List (Point a)) -> Dict Bucket (List (Point a))
insert radius point buckets =
    let
        updater old =
            case old of
                Nothing ->
                    Just [ point ]

                Just currentPoints ->
                    Just (point :: currentPoints)
    in
        Dict.update (toKey radius point) updater buckets


{-| Transform/update the points

The Buckets are automatically recalculated
-}
map : (Point a -> Point b) -> Buckets a -> Buckets b
map tagger (Buckets { radius, buckets }) =
    let
        newBuckets =
            buckets
                |> Dict.foldr (\k v accum -> List.foldl (insert radius) accum (List.map tagger v)) Dict.empty
    in
        Buckets { buckets = newBuckets, radius = radius }


toKey : Float -> Point a -> ( Int, Int )
toKey radius { x, y } =
    let
        i1 =
            floor (x / radius)

        i2 =
            floor (y / radius)
    in
        ( i1, i2 )


distanceSq : Point a -> Point b -> Float
distanceSq p1 p2 =
    (p1.x - p2.x) * (p1.x - p2.x) + (p1.y - p2.y) * (p1.y - p2.y)


getForwardBuckets : Bucket -> Dict Bucket a -> List Bucket
getForwardBuckets ( i1, i2 ) buckets =
    let
        topRight =
            ( i1 + 1, i2 - 1 )

        right =
            ( i1 + 1, i2 )

        bottomRight =
            ( i1 + 1, i2 + 1 )

        bottom =
            ( i1, i2 + 1 )

        potentials =
            [ topRight, right, bottomRight, bottom ]
    in
        [ topRight, right, bottomRight, bottom ]


{-| Keep duplicate pairs, so the result will contain `(A, B)` and `(B, A)`.
-}
pairReflexive : Bucket -> Buckets a -> List ( Point a, Point a )
pairReflexive bucket (Buckets { buckets, radius }) =
    let
        neighbours : List (Point a)
        neighbours =
            getForwardBuckets bucket buckets
                |> List.filterMap (\b -> Dict.get b buckets)
                |> List.concat

        predicate p1 p2 =
            if p1 == p2 then
                Nothing
            else if distanceSq p1 p2 < radius then
                Just ( p1, p2 )
            else
                Nothing
    in
        case Dict.get bucket buckets of
            Nothing ->
                []

            Just points ->
                points
                    |> List.concatMap (\point -> List.map (predicate point) neighbours)
                    |> List.filterMap identity


withRemainder : List a -> List ( a, List a )
withRemainder list =
    case list of
        [] ->
            []

        x :: xs ->
            ( x, xs ) :: withRemainder xs


{-| Get only unique pairs: if `(A, B)` is in the output, then `(B, A)` never is.
-}
pair : Bucket -> Buckets a -> List ( Point a, Point a )
pair bucket (Buckets { buckets, radius }) =
    let
        neighbours : List (Point a)
        neighbours =
            getForwardBuckets bucket buckets
                |> List.filterMap (\b -> Dict.get b buckets)
                |> List.concat

        predicate p1 p2 =
            if p1 == p2 then
                Nothing
            else if distanceSq p1 p2 < radius then
                Just ( p1, p2 )
            else
                Nothing
    in
        case Dict.get bucket buckets of
            Nothing ->
                []

            Just points ->
                points
                    |> withRemainder
                    |> List.concatMap (\( point, remainder ) -> List.map (predicate point) (neighbours ++ remainder))
                    |> List.filterMap identity