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
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
open Ast
(** Helper: get values from a VVector or VList *)
let to_value_array label = function
| VVector arr -> Ok arr
| VList items -> Ok (Array.of_list (List.map snd items))
| VNA _ -> Error (Error.na_value_error label)
| _ -> Error (Error.type_error (Printf.sprintf "Function `%s` expects a Vector or List." label))
let register env =
(*
--# Lag values
--#
--# Shifts a vector forward by n positions, filling with NA.
--#
--# @name lag
--# @param x :: Vector The input vector.
--# @param n :: Int (Optional) Number of positions to shift. Default is 1.
--# @return :: Vector The shifted vector.
--# @example
--# lag([1, 2, 3])
--# -- Returns = [NA, 1, 2]
--# @family colcraft
--# @seealso lead
--# @export
*)
(* lag(x) or lag(x, n): shift values forward, filling with NA *)
let env = Env.add "lag"
(make_builtin ~name:"lag" ~variadic:true 1 (fun args _env ->
match args with
| [arg] | [arg; VInt 1] ->
(match to_value_array "lag" arg with
| Error e -> e
| Ok arr ->
let n = Array.length arr in
if n = 0 then VVector [||]
else
let result = Array.make n ((VNA NAGeneric)) in
for i = 1 to n - 1 do
result.(i) <- arr.(i - 1)
done;
VVector result)
| [arg; VInt offset] when offset >= 0 ->
(match to_value_array "lag" arg with
| Error e -> e
| Ok arr ->
let n = Array.length arr in
if n = 0 then VVector [||]
else
let result = Array.make n ((VNA NAGeneric)) in
for i = offset to n - 1 do
result.(i) <- arr.(i - offset)
done;
VVector result)
| [_; VInt _] -> Error.value_error "Function `lag` offset must be non-negative."
| [_; _] -> Error.type_error "Function `lag` expects an integer offset."
| _ -> Error.make_error ArityError "Function `lag` takes 1 or 2 arguments."
))
env
in
(*
--# Lead values
--#
--# Shifts a vector backward by n positions, filling with NA.
--#
--# @name lead
--# @param x :: Vector The input vector.
--# @param n :: Int (Optional) Number of positions to shift. Default is 1.
--# @return :: Vector The shifted vector.
--# @example
--# lead([1, 2, 3])
--# -- Returns = [2, 3, NA]
--# @family colcraft
--# @seealso lag
--# @export
*)
(* lead(x) or lead(x, n): shift values backward, filling with NA *)
let env = Env.add "lead"
(make_builtin ~name:"lead" ~variadic:true 1 (fun args _env ->
match args with
| [arg] | [arg; VInt 1] ->
(match to_value_array "lead" arg with
| Error e -> e
| Ok arr ->
let n = Array.length arr in
if n = 0 then VVector [||]
else
let result = Array.make n ((VNA NAGeneric)) in
for i = 0 to n - 2 do
result.(i) <- arr.(i + 1)
done;
VVector result)
| [arg; VInt offset] when offset >= 0 ->
(match to_value_array "lead" arg with
| Error e -> e
| Ok arr ->
let n = Array.length arr in
if n = 0 then VVector [||]
else
let result = Array.make n ((VNA NAGeneric)) in
for i = 0 to n - 1 - offset do
result.(i) <- arr.(i + offset)
done;
VVector result)
| [_; VInt _] -> Error.value_error "Function `lead` offset must be non-negative."
| [_; _] -> Error.type_error "Function `lead` expects an integer offset."
| _ -> Error.make_error ArityError "Function `lead` takes 1 or 2 arguments."
))
env
in
env