Merge pull request #1133 from WardBrian/complex-containers

Complex containers

Author: WardBrian

src/analysis_and_optimization/Debug_data_generation.ml

@@ -232,6 +232,22 @@ let gen_matrix mm m n t =
       Expr.Helpers.matrix_from_rows
         (repeat_th m (fun () -> gen_row_vector mm n t))
 
+let gen_complex_unwrapped () =
+  ( gen_num_real Map.Poly.empty Transformation.Identity
+  , gen_num_real Map.Poly.empty Transformation.Identity )
+
+let gen_complex () = Expr.Helpers.complex (gen_complex_unwrapped ())
+
+let gen_complex_row_vector n =
+  Expr.Helpers.complex_row_vector (repeat_th n gen_complex_unwrapped)
+
+let gen_complex_vector n =
+  Expr.Helpers.complex_vector (repeat_th n gen_complex_unwrapped)
+
+let gen_complex_matrix m n =
+  Expr.Helpers.complex_matrix_from_rows
+    (repeat_th m (fun () -> gen_complex_row_vector n))
+
 let rec gen_array m st n t =
   let elt () = generate_value m st t in
   match (t : Expr.Typed.t Transformation.t) with
@@ -252,11 +268,15 @@ and generate_value m st t =
   | SReal -> Expr.Helpers.float (gen_num_real m t)
   | SComplex ->
       (* when serialzied, a complex number looks just like a 2-array of reals *)
-      generate_value m (SArray (SReal, Expr.Helpers.int 2)) t
+      gen_complex ()
   | SVector (_, e) -> gen_vector m (unwrap_int_exn m e) t
   | SRowVector (_, e) -> gen_row_vector m (unwrap_int_exn m e) t
   | SMatrix (_, e1, e2) ->
       gen_matrix m (unwrap_int_exn m e1) (unwrap_int_exn m e2) t
+  | SComplexVector e -> gen_complex_vector (unwrap_int_exn m e)
+  | SComplexRowVector e -> gen_complex_row_vector (unwrap_int_exn m e)
+  | SComplexMatrix (e1, e2) ->
+      gen_complex_matrix (unwrap_int_exn m e1) (unwrap_int_exn m e2)
   | SArray (st, e) -> gen_array m st (unwrap_int_exn m e) t
 
 let rec pp_value_json ppf e =
@@ -267,6 +287,8 @@ let rec pp_value_json ppf e =
   | FunApp (StanLib (transpose, _, _), [e])
     when String.equal transpose (Operator.to_string Transpose) ->
       pp_value_json ppf e
+  | FunApp (StanLib ("to_complex", _, _), [r; i]) ->
+      Fmt.(pf ppf "[@[<hov 1>%a, %a@]]" pp_value_json r pp_value_json i)
   | _ ->
       Common.FatalError.fatal_error_msg
         [%message "Could not evaluate expression " (e : Expr.Typed.t)]

src/analysis_and_optimization/Mem_pattern.ml

@@ -425,6 +425,9 @@ let rec query_initial_demotable_stmt (in_loop : bool) (acc : string Set.Poly.t)
       Set.Poly.union_list
         [ acc; query_expr acc predicate
         ; query_initial_demotable_stmt true acc body ]
+  | Decl {decl_type= Type.Sized st; decl_id; _}
+    when SizedType.is_complex_type st ->
+      Set.Poly.add acc decl_id
   | Skip | Break | Continue | Decl _ -> acc
 
 (** Look through a statement to see whether the objects used in it need to be

src/frontend/Ast.ml

@@ -329,8 +329,13 @@ let get_loc_dt (t : untyped_expression Type.t) =
   match t with
   | Type.Unsized _ | Sized (SInt | SReal | SComplex) -> None
   | Sized
-      (SVector (_, e) | SRowVector (_, e) | SMatrix (_, e, _) | SArray (_, e))
-    ->
+      ( SVector (_, e)
+      | SRowVector (_, e)
+      | SMatrix (_, e, _)
+      | SComplexVector e
+      | SComplexRowVector e
+      | SComplexMatrix (e, _)
+      | SArray (_, e) ) ->
       Some e.emeta.loc.begin_loc
 
 let get_loc_tf (t : untyped_expression Transformation.t) =

src/frontend/Ast_to_Mir.ml

@@ -231,7 +231,8 @@ let param_size transform sizedtype =
     | SizedType.SArray (t, d) -> SizedType.SArray (shrink_eigen f t, d)
     | SVector (mem_pattern, d) | SMatrix (mem_pattern, d, _) ->
         SVector (mem_pattern, f d)
-    | SInt | SReal | SComplex | SRowVector _ ->
+    | SInt | SReal | SComplex | SRowVector _ | SComplexRowVector _
+     |SComplexVector _ | SComplexMatrix _ ->
         Common.FatalError.fatal_error_msg
           [%message
             "Expecting SVector or SMatrix, got " (st : Expr.Typed.t SizedType.t)]
@@ -240,7 +241,8 @@ let param_size transform sizedtype =
     match st with
     | SizedType.SArray (t, d) -> SizedType.SArray (shrink_eigen_mat f t, d)
     | SMatrix (mem_pattern, d1, d2) -> SVector (mem_pattern, f d1 d2)
-    | SInt | SReal | SComplex | SRowVector _ | SVector _ ->
+    | SInt | SReal | SComplex | SRowVector _ | SVector _ | SComplexRowVector _
+     |SComplexVector _ | SComplexMatrix _ ->
         Common.FatalError.fatal_error_msg
           [%message "Expecting SMatrix, got " (st : Expr.Typed.t SizedType.t)]
   in
@@ -308,6 +310,16 @@ let check_sizedtype name =
         let er = trans_expr r in
         let ec = trans_expr c in
         (check r er @ check c ec, SizedType.SMatrix (mem_pattern, er, ec))
+    | SComplexVector s ->
+        let e = trans_expr s in
+        (check s e, SizedType.SComplexVector e)
+    | SComplexRowVector s ->
+        let e = trans_expr s in
+        (check s e, SizedType.SComplexRowVector e)
+    | SComplexMatrix (r, c) ->
+        let er = trans_expr r in
+        let ec = trans_expr c in
+        (check r er @ check c ec, SizedType.SComplexMatrix (er, ec))
     | SArray (t, s) ->
         let e = trans_expr s in
         let ll, t = sizedtype t in
@@ -581,6 +593,12 @@ let trans_sizedtype_decl declc tr name =
     | SRowVector (mem_pattern, s) ->
         let l, s = grab_size FnValidateSize n s in
         (l, SizedType.SRowVector (mem_pattern, s))
+    | SComplexRowVector s ->
+        let l, s = grab_size FnValidateSize n s in
+        (l, SizedType.SComplexRowVector s)
+    | SComplexVector s ->
+        let l, s = grab_size FnValidateSize n s in
+        (l, SizedType.SComplexVector s)
     | SMatrix (mem_pattern, r, c) ->
         let l1, r = grab_size FnValidateSize n r in
         let l2, c = grab_size FnValidateSize (n + 1) c in
@@ -598,6 +616,10 @@ let trans_sizedtype_decl declc tr name =
                 ; meta= r.Expr.Fixed.meta.Expr.Typed.Meta.loc } ]
           | _ -> [] in
         (l1 @ l2 @ cf_cov, SizedType.SMatrix (mem_pattern, r, c))
+    | SComplexMatrix (r, c) ->
+        let l1, r = grab_size FnValidateSize n r in
+        let l2, c = grab_size FnValidateSize (n + 1) c in
+        (l1 @ l2, SizedType.SComplexMatrix (r, c))
     | SArray (t, s) ->
         let l, s = grab_size FnValidateSize n s in
         let ll, t = go (n + 1) t in

src/frontend/Info.ml

@@ -8,7 +8,9 @@ let rec sized_basetype_dims t =
   | SReal -> ("real", 0)
   | SComplex -> ("complex", 0)
   | SVector _ | SRowVector _ -> ("real", 1)
+  | SComplexVector _ | SComplexRowVector _ -> ("complex", 1)
   | SMatrix _ -> ("real", 2)
+  | SComplexMatrix _ -> ("complex", 2)
   | SArray (t, _) ->
       let bt, n = sized_basetype_dims t in
       (bt, n + 1)
@@ -19,7 +21,9 @@ let rec unsized_basetype_dims t =
   | UReal -> ("real", 0)
   | UComplex -> ("complex", 0)
   | UVector | URowVector -> ("real", 1)
+  | UComplexVector | UComplexRowVector -> ("complex", 1)
   | UMatrix -> ("real", 2)
+  | UComplexMatrix -> ("complex", 2)
   | UArray t ->
       let bt, n = unsized_basetype_dims t in
       (bt, n + 1)

src/frontend/Pretty_printing.ml

@@ -180,7 +180,7 @@ let pp_autodifftype = Middle.UnsizedType.pp_autodifftype
 let rec unwind_sized_array_type st =
   match st with
   | Middle.SizedType.SInt | SReal | SComplex | SVector _ | SRowVector _
-   |SMatrix _ ->
+   |SMatrix _ | SComplexMatrix _ | SComplexVector _ | SComplexRowVector _ ->
       (st, [])
   | SArray (st, dim) ->
       let st', dims = unwind_sized_array_type st in
@@ -319,6 +319,10 @@ let rec pp_sizedtype ppf = function
   | SRowVector (_, e) -> pf ppf "row_vector[%a]" pp_expression e
   | SMatrix (_, e1, e2) ->
       pf ppf "matrix[%a, %a]" pp_expression e1 pp_expression e2
+  | SComplexVector e -> pf ppf "complex_vector[%a]" pp_expression e
+  | SComplexRowVector e -> pf ppf "complex_row_vector[%a]" pp_expression e
+  | SComplexMatrix (e1, e2) ->
+      pf ppf "complex_matrix[%a, %a]" pp_expression e1 pp_expression e2
   | SArray _ as arr ->
       let ty, ixs = unwind_sized_array_type arr in
       pf ppf "array[@[%a@]]@ %a"
@@ -348,9 +352,12 @@ let pp_transformed_type ppf (pst, trans) =
       let pp_possibly_transformed_type ppf (st, trans) =
         let sizes_fmt =
           match st with
-          | SizedType.SVector (_, e) | SRowVector (_, e) ->
+          | SizedType.SVector (_, e)
+           |SRowVector (_, e)
+           |SComplexVector e
+           |SComplexRowVector e ->
               const (fun ppf -> pf ppf "[%a]" pp_expression) e
-          | SMatrix (_, e1, e2) ->
+          | SMatrix (_, e1, e2) | SComplexMatrix (e1, e2) ->
               const
                 (fun ppf -> pf ppf "[%a, %a]" pp_expression e1 pp_expression)
                 e2

src/frontend/Promotion.ml

@@ -21,25 +21,27 @@ let promote_inner (exp : Ast.typed_expression) prom =
         { expr= Ast.Promotion (exp, UReal, AutoDiffable)
         ; emeta=
             { emeta with
-              type_= UnsizedType.promote_array emeta.type_ UReal
+              type_= UnsizedType.promote_container emeta.type_ UReal
             ; ad_level= AutoDiffable } }
   | ToComplexVar ->
       Ast.
         { expr= Ast.Promotion (exp, UComplex, AutoDiffable)
         ; emeta=
             { emeta with
-              type_= UnsizedType.promote_array emeta.type_ UComplex
+              type_= UnsizedType.promote_container emeta.type_ UComplex
             ; ad_level= AutoDiffable } }
   | IntToReal when UnsizedType.is_int_type emeta.type_ ->
       Ast.
         { expr= Ast.Promotion (exp, UReal, emeta.ad_level)
-        ; emeta= {emeta with type_= UnsizedType.promote_array emeta.type_ UReal}
+        ; emeta=
+            {emeta with type_= UnsizedType.promote_container emeta.type_ UReal}
         }
   | (IntToComplex | RealToComplex)
     when not (UnsizedType.is_complex_type emeta.type_) ->
       (* these two promotions are separated for cost, but are actually the same promotion *)
       { expr= Promotion (exp, UComplex, emeta.ad_level)
-      ; emeta= {emeta with type_= UnsizedType.promote_array emeta.type_ UComplex}
+      ; emeta=
+          {emeta with type_= UnsizedType.promote_container emeta.type_ UComplex}
       }
   | _ -> exp
 
@@ -54,37 +56,54 @@ let rec promote (exp : Ast.typed_expression) prom =
       { expr= ArrayExpr pes
       ; emeta=
           { exp.emeta with
-            type_= UnsizedType.promote_array exp.emeta.type_ type_
+            type_= UnsizedType.promote_container exp.emeta.type_ type_
           ; ad_level } }
   | RowVectorExpr (_ :: _ as es) ->
       let pes = List.map ~f:(fun e -> promote e prom) es in
       let fst = List.hd_exn pes in
       let ad_level = fst.emeta.ad_level in
-      {expr= RowVectorExpr pes; emeta= {exp.emeta with ad_level}}
+      let type_ =
+        (* "RowVectorExpr" can also be a matrix expr, depends on what is inside *)
+        match fst.emeta.type_ with
+        | UComplexRowVector -> UnsizedType.UComplexMatrix
+        | URowVector -> UMatrix
+        | UComplex -> UComplexRowVector
+        | _ -> URowVector in
+      {expr= RowVectorExpr pes; emeta= {exp.emeta with type_; ad_level}}
   | _ -> promote_inner exp prom
 
 let promote_list es promotions = List.map2_exn es promotions ~f:promote
 
 (** Get the promotion needed to make the second type into the first.
   Types NEED to have previously been checked to be promotable
 *)
-let rec get_type_promotion_exn (ad, ty) (ad2, ty2) =
-  match (ty, ty2) with
+let rec get_type_promotion_exn (ad_orig, ty_orig) (ad_expect, ty_expect) =
+  match (ty_orig, ty_expect) with
   | UnsizedType.(UReal, (UReal | UInt) | UVector, UVector | UMatrix, UMatrix)
-    when ad <> ad2 ->
+    when ad_orig <> ad_expect ->
       ToVar
-  | UComplex, (UReal | UInt | UComplex) when ad <> ad2 -> ToComplexVar
+  | UComplex, (UReal | UInt | UComplex)
+   |UComplexMatrix, (UMatrix | UComplexMatrix)
+   |UComplexVector, (UVector | UComplexVector)
+   |UComplexRowVector, (URowVector | UComplexRowVector)
+    when ad_orig <> ad_expect ->
+      ToComplexVar
   | UReal, UInt -> IntToReal
   | UComplex, UInt -> IntToComplex
-  | UComplex, UReal -> RealToComplex
-  | UArray nt1, UArray nt2 -> get_type_promotion_exn (ad, nt1) (ad2, nt2)
+  | UComplex, UReal
+   |UComplexMatrix, UMatrix
+   |UComplexVector, UVector
+   |UComplexRowVector, URowVector ->
+      RealToComplex
+  | UArray nt1, UArray nt2 ->
+      get_type_promotion_exn (ad_orig, nt1) (ad_expect, nt2)
   | t1, t2 when t1 = t2 -> NoPromotion
   | _, _ ->
       Common.FatalError.fatal_error_msg
         [%message
           "Tried to get promotion of mismatched types!"
-            (ty : UnsizedType.t)
-            (ty2 : UnsizedType.t)]
+            (ty_orig : UnsizedType.t)
+            (ty_expect : UnsizedType.t)]
 
 (** Calculate the "cost"/number of promotions performed.
     Used to disambiguate function signatures

src/frontend/SignatureMismatch.ml

@@ -17,7 +17,8 @@ let pp_unsized_type ctx ppf =
   let rec pp ppf ty =
     match ty with
     | UnsizedType.UInt | UReal | UVector | URowVector | UMatrix | UComplex
-     |UMathLibraryFunction ->
+     |UComplexRowVector | UComplexVector | UComplexMatrix | UMathLibraryFunction
+      ->
         UnsizedType.pp ppf ty
     | UArray ut ->
         let ut2, d = UnsizedType.unwind_array_type ut in
@@ -140,8 +141,13 @@ let rec check_same_type depth t1 t2 =
   match (t1, t2) with
   | t1, t2 when t1 = t2 -> Ok Promotion.NoPromotion
   | UnsizedType.(UReal, UInt) when depth < 1 -> Ok IntToReal
-  | UnsizedType.(UComplex, UInt) when depth < 1 -> Ok IntToComplex
-  | UnsizedType.(UComplex, UReal) when depth < 1 -> Ok RealToComplex
+  | UComplex, UInt when depth < 1 -> Ok IntToComplex
+  | UComplex, UReal
+   |UComplexMatrix, UMatrix
+   |UComplexVector, UVector
+   |UComplexRowVector, URowVector
+    when depth < 1 ->
+      Ok RealToComplex
   (* Arrays: Try to recursively promote, but make sure the error is for these types,
      not the recursive call *)
   | UArray nt1, UArray nt2 ->