Add join operator to vector and matrix

fsprojects · Sep 19, 2024 · 29f98d6 · 29f98d6
1 parent f2e49ba
commit 29f98d6
Showing 1 changed file with 16 additions and 2 deletions.
diff --git a/src/FSharpPlus.TypeLevel/Data/Matrix.fs b/src/FSharpPlus.TypeLevel/Data/Matrix.fs
@@ -251,6 +251,10 @@ module Vector =
 
   let inline apply (f: Vector<'a -> 'b, 'n>) (v: Vector<'a, 'n>) : Vector<'b, 'n> = map2 id f v
 
+  [<MethodImpl(MethodImplOptions.AggressiveInlining)>]
+  let join (vv: Vector<Vector<'a, 'n>, 'n>): Vector<'a, 'n> =
+    { Items = Array.init (Array.length vv.Items) (fun i -> vv.Items.[i].Items.[i]) }
+
   let inline norm (v: Vector< ^a, ^n >) : ^a =
     v |> toArray |> Array.sumBy (fun x -> x * x) |> sqrt
   let inline maximumNorm (v: Vector< ^a, ^n >) : ^a =
@@ -327,6 +331,14 @@ module Matrix =
       for j = 0 to Array2D.length2 m1.Items - 1 do
         f i j m1.Items.[i, j] m2.Items.[i, j]
 
+  let inline apply (f: Matrix<'a -> 'b, 'm, 'n>) (m: Matrix<'a, 'm, 'n>) : Matrix<'b, 'm, 'n> = map2 id f m
+
+  [<MethodImpl(MethodImplOptions.AggressiveInlining)>]
+  let join (m: Matrix<Matrix<'a, 'm, 'n>, 'm, 'n>) : Matrix<'a, 'm, 'n> =
+    { Items =
+        Array2D.init (Array2D.length1 m.Items) (Array2D.length2 m.Items)
+          (fun i j -> m.Items.[i, j].Items.[i, j] ) }
+
   let inline rowLength (_: Matrix<'a, 'm, 'n>) : 'm = Singleton<'m>
   let inline colLength (_: Matrix<'a, 'm, 'n>) : 'n = Singleton<'n>
   let inline rowLength' (_: Matrix<'a, ^m, 'n>) : int = RuntimeValue (Singleton< ^m >)
@@ -571,8 +583,9 @@ type Matrix<'Item, 'Row, 'Column> with
 
   static member inline Return (x: 'x) : Matrix<'x, 'm, 'n> = Matrix.replicate Singleton Singleton x
   static member inline Pure   (x: 'x) : Matrix<'x, 'm, 'n> = Matrix.replicate Singleton Singleton x
-  static member inline ( <*> ) (f: Matrix<'x -> 'y, 'm, 'n>, x: Matrix<'x, 'm, 'n>) = Matrix.map2 id f x
-  static member inline ( <.> ) (f: Matrix<'x -> 'y, 'm, 'n>, x: Matrix<'x, 'm, 'n>) = Matrix.map2 id f x
+  static member inline ( <*> ) (f: Matrix<'x -> 'y, 'm, 'n>, x: Matrix<'x, 'm, 'n>) = Matrix.apply f x
+  static member inline ( <.> ) (f: Matrix<'x -> 'y, 'm, 'n>, x: Matrix<'x, 'm, 'n>) = Matrix.apply f x
+  static member inline Join (x: Matrix<Matrix<'x, 'm, 'n>, 'm, 'n>) = Matrix.join x
   static member inline get_Zero () : Matrix<'a, 'm, 'n> = Matrix.zero
   static member inline ( + ) (m1, m2) = Matrix.map2 (+) m1 m2
   static member inline ( - ) (m1, m2) = Matrix.map2 (-) m1 m2
@@ -607,6 +620,7 @@ type Vector<'Item, 'Length> with
   static member inline Pure   (x: 'x) : Vector<'x, 'n> = Vector.replicate Singleton x
   static member inline ( <*> ) (f: Vector<'x -> 'y, 'n>, x: Vector<'x, 'n>) : Vector<'y, 'n> = Vector.apply f x
   static member inline ( <.> ) (f: Vector<'x -> 'y, 'n>, x: Vector<'x, 'n>) : Vector<'y, 'n> = Vector.apply f x
+  static member inline Join (x: Vector<Vector<'x, 'n>, 'n>) : Vector<'x, 'n> = Vector.join x
 
   [<EditorBrowsable(EditorBrowsableState.Never)>]
   static member inline Zip (x, y) = Vector.zip x y