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
open Ast
(*
--# Square root
--#
--# Calculates the square root of x.
--#
--# @name sqrt
--# @param x :: Number | Vector | NDArray The input value (must be non-negative).
--# @param na_ignore :: Bool Whether to preserve NA values in inputs. Default is false.
--# @return :: Float | Vector | NDArray The square root.
--# @example
--# sqrt(16)
--# -- Returns = 4.0
--# @family math
--# @seealso pow
--# @export
*)
let register env =
Env.add "sqrt"
(make_builtin_named ~name:"sqrt" ~variadic:true 1 (fun named_args _env ->
match Math_common.get_bool_flag "na_ignore" false named_args with
| Error e -> e
| Ok na_ignore ->
let args = Math_common.positional_args_without [ "na_ignore" ] named_args in
match args with
| [VInt n] ->
if n < 0 then Error.value_error "Function `sqrt` is undefined for negative numbers."
else VFloat (Float.sqrt (float_of_int n))
| [VFloat f] ->
if f < 0.0 then Error.value_error "Function `sqrt` is undefined for negative numbers."
else VFloat (Float.sqrt f)
| [VVector arr] ->
let result = Array.make (Array.length arr) (VNA NAGeneric) in
let had_error = ref None in
Array.iteri (fun i v ->
if !had_error = None then
match v with
| VInt n ->
if n < 0 then had_error := Some (Error.value_error "Function `sqrt` is undefined for negative numbers.")
else result.(i) <- VFloat (Float.sqrt (float_of_int n))
| VFloat f ->
if f < 0.0 then had_error := Some (Error.value_error "Function `sqrt` is undefined for negative numbers.")
else result.(i) <- VFloat (Float.sqrt f)
| VNA na_t when na_ignore -> result.(i) <- VNA na_t
| VNA _ -> had_error := Some (Error.na_value_error "sqrt")
| _ -> had_error := Some (Error.type_error "Function `sqrt` requires numeric values.")
) arr;
(match !had_error with Some e -> e | None -> VVector result)
| [VNDArray arr] ->
let result = Array.map (fun f ->
if f < 0.0 then nan (* Use NaN for negative sqrt in NDArray to match typical vectorized behavior, or error? *)
else Float.sqrt f
) arr.data in
if Array.exists Float.is_nan result then
Error.value_error "Function `sqrt` encountered negative values in NDArray."
else
VNDArray { shape = arr.shape; data = result }
| [VNA na_t] when na_ignore -> VNA na_t
| [VNA _] -> Error.na_value_error "sqrt"
| [_] -> Error.type_error "Function `sqrt` expects a number, numeric Vector, or NDArray."
| _ -> Error.arity_error_named "sqrt" 1 (List.length args)
))
env