Lead/lag with respect to a time variable

Project.toml

@@ -10,6 +10,7 @@ julia = "1"
 
 [extras]
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+Dates = "ade2ca70-3891-5945-98fb-dc099432e06a"
 
 [targets]
-test = ["Test"]
+test = ["Test", "Dates"]

src/lag.jl

@@ -97,3 +97,70 @@ julia> s = lead(v, (0, 2))
 ```
 """
 lead(v::AbstractArray, n = 1; default = missing) = ShiftedArray(v, map(-, n); default = default)
+
+
+
+
+
+
+
+
+"""
+    lag(v::AbstractVector, times::AbstractVector, period = oneunit(eltype(times)); default = missing) -> Vector
+
+Shifts with respect to a times given in the vector `times`.  The third variable `period` gives the period by which to shift.
+`default` specifies a default value when the shifted time is not in `times`.
+Elements in `times` must all be distinct.
+
+# Examples
+
+```jldoctest lead
+julia> v = [1, 3, 5, 4];
+julia> times = [1990, 1992, 1993];
+julia> lag(v, times)
+3-element Array{Union{Missing, Int64},1}:
+  missing
+  missing
+ 3
+julia> using Dates
+julia> times = [Date(1990, 1, 1), Date(1990, 1, 3), Date(1990, 1, 4)]
+julia> lag(v, times, Day(1))
+3-element Array{Union{Missing, Int64},1}:
+ missing
+ missing
+3
+"""
+function lag(v::AbstractVector, times::AbstractVector, period = oneunit(eltype(times)); default = missing)
+   	# Code follows the function indexin
+    inds = keys(times)
+    timesdict = Dict{eltype(times),eltype(inds)}()
+    for (val, ind) in zip(times, inds)
+         out = get!(timesdict, val, ind)
+         out != ind && error("Times must be distinct")
+    end
+    return Union{eltype(v), typeof(default)}[
+         (i = get(timesdict, x - period, nothing); i !== nothing ? v[i] : default) for x in times
+    ]
+end
+
+"""
+    lead(v::AbstractVector, times::AbstractVector, period = oneunit(eltype(times)); default = missing) -> Vector
+
+Shifts with respect to a vector of times `times`. The third variable `period` gives the period by which to shift.
+`default` specifies a default value when the shifted time is not in `times`.
+Elements in `times` must all be distinct.
+
+# Examples
+
+```jldoctest lead
+julia> v = [1, 3, 5, 4];
+julia> times = [1990, 1992, 1993];
+julia> lead(v, times, 1)
+3-element Array{Union{Missing, Int64},1}:
+  missing
+  5
+  missing
+"""
+function lead(v::AbstractVector, times::AbstractVector, period = oneunit(eltype(times)); default = missing)
+    lag(v, times, -period; default = default)
+end

test/runtests.jl

@@ -1,4 +1,4 @@
-using ShiftedArrays, Test
+using ShiftedArrays, Test, Dates
 
 @testset "ShiftedVector" begin
     v = [1, 3, 5, 4]
@@ -96,4 +96,25 @@ end
     @test isequal(diff2, [-7, -9, missing, missing])
 
     @test all(lead(v, 2, default = -100) .== coalesce.(lead(v, 2), -100))
+
+
+    v = [4, 5, 6]
+    times = [1989, 1991, 1992]
+    @test all(lag(v, times) .=== [missing, missing, 5])
+    @test all(lag(v, times; default = 0) .=== [0, 0, 5])
+    @test all(lead(v, times) .=== [missing, 6, missing])
+    @test all(lead(v, times; default = 0) .=== [0, 6, 0])
+    times = [Date(1989, 1, 1), Date(1989, 1, 3), Date(1989, 1, 4)]
+    @test all(lag(v, times, Day(1)) .=== [missing, missing, 5])
+    @test all(lag(v, times, Day(2)) .=== [missing, 4, missing])
+    @test all(lag(v, times, Day(5)) .=== [missing, missing, missing])
+    @test all(lead(v, times, Day(1)) .=== [missing, 6, missing])
+    @test all(lead(v, times, Day(1)) .=== lag(v, times, -Day(1)))
+    @test all(lead(v, times, Day(2)) .=== [5, missing, missing])
+    @test all(lead(v, times, Day(5)) .=== [missing, missing, missing])
+    times = [DateTime(1989, 1, 1), DateTime(1989, 1, 3), DateTime(1989, 1, 4)]
+    @test all(lag(v, times, Millisecond(1)) .=== [missing, missing, missing])
+    @test all(lag(v, times, Day(1)) .=== [missing, missing, 5])
+    times = [Date(1989, 1, 1), Date(1989, 1, 3), Date(1989, 1, 3)]
+    @test_throws ErrorException lag(v, times, Day(1))
 end