Skip to content

Commit

Permalink
findall: support custom function as argument (#488)
Browse files Browse the repository at this point in the history
* findall: support custom function as arg

```
findall(>(100), cl)
```
  • Loading branch information
iblislin authored Mar 18, 2021
1 parent 80d3705 commit ecee85a
Show file tree
Hide file tree
Showing 3 changed files with 22 additions and 0 deletions.
12 changes: 12 additions & 0 deletions docs/src/split.md
Original file line number Diff line number Diff line change
Expand Up @@ -84,6 +84,18 @@ typeof(red)
ohlc[red]
```

The following example won't create a temporary `Bool` vector, and gains better
performance.

```@setup findall
using TimeSeries
using MarketData
```

```@repl findall
findall(>(100), cl)
```

## Splitting by head and tail

### `head`
Expand Down
5 changes: 5 additions & 0 deletions src/split.jl
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,11 @@ to(ta::TimeArray{T, N, D}, d::D) where {T, N, D} =
###### findall ##################

Base.findall(ta::TimeArray{Bool,1}) = findall(values(ta))
Base.findall(f::Function, ta::TimeArray{T,1}) where {T} = findall(f, values(ta))
function Base.findall(f::Function, ta::TimeArray{T,2}) where {T}
A = values(ta)
collect(i for i in axes(A, 1) if f(view(A, i, :)))
end

###### findwhen #################

Expand Down
5 changes: 5 additions & 0 deletions test/split.jl
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,11 @@ using TimeSeries
@test findwhen(cl .> op)[2] == Date(2000, 1, 5)
@test length(findwhen(cl .> op)) == 244
end

@testset "findall(f::Function, ta)" begin
@test findall(cl .> 100) == findall(x -> x > 100, cl)
@test findall(cl .> 100) == findall(x -> x[4] > 100, ohlc)
end
end


Expand Down

0 comments on commit ecee85a

Please sign in to comment.