-
Notifications
You must be signed in to change notification settings - Fork 36
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
* extract tweaks * add rasters.sample * get rid of some toy code * just write kw... in docs Co-authored-by: Rafael Schouten <[email protected]> * use Type{G} Co-authored-by: Rafael Schouten <[email protected]> * add a missing where G * neater formatting * better reuse _rowtypes code * more code style Co-authored-by: Rafael Schouten <[email protected]> * extract tweaks * add rasters.sample * get rid of some toy code * neater formatting * better reuse _rowtypes code * more code style Co-authored-by: Rafael Schouten <[email protected]> * merge _rowtype changes from other branch * optimization for extract points with skipmissing * streamline _rowtype * correctly use _rowtype * improve type stability * switch to using dimindices internally * let users specify geometry type * just use a broadcast * don't need to be compatible with julia < 1.9 Co-authored-by: Rafael Schouten <[email protected]> * deal with type stability better * allow not providing `n` * tweak the docstring --------- Co-authored-by: Rafael Schouten <[email protected]>
- Loading branch information
1 parent
3202707
commit 1f00143
Showing
7 changed files
with
255 additions
and
38 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,13 @@ | ||
| module RastersStatsBaseExt | ||
|
|
||
| using Rasters, StatsBase | ||
| using StatsBase.Random | ||
|
|
||
| const RA = Rasters | ||
|
|
||
| import Rasters: _True, _False, _booltype, istrue | ||
| import Rasters.DimensionalData as DD | ||
|
|
||
| include("sample.jl") | ||
|
|
||
| end # Module |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,94 @@ | ||
| Rasters.sample(x::RA.RasterStackOrArray, args...; kw...) = Rasters.sample(Random.GLOBAL_RNG, x, args...; kw...) | ||
| @inline function Rasters.sample( | ||
| rng::Random.AbstractRNG, x::RA.RasterStackOrArray, args...; | ||
| geometry=(X,Y), | ||
| index = false, | ||
| names=RA._names(x), | ||
| name=names, | ||
| skipmissing=false, | ||
| weights=nothing, | ||
| weightstype::Type{<:StatsBase.AbstractWeights}=StatsBase.Weights, | ||
| kw... | ||
| ) | ||
| na = DD._astuple(name) | ||
| geometry, geometrytype, dims = _geometrytype(x, geometry) | ||
|
|
||
| _sample(rng, x, args...; | ||
| dims, | ||
| names = NamedTuple{na}(na), | ||
| geometry, | ||
| geometrytype, | ||
| # These keywords are converted to _True/_False for type stability later on | ||
| index = _booltype(index), | ||
| skipmissing = _booltype(skipmissing), | ||
| weights, | ||
| weightstype, | ||
| kw... # passed to StatsBase.sample, could be replace or ordered | ||
| ) | ||
| end | ||
| function _sample( | ||
| rng, x, n::Integer; | ||
| dims, names::NamedTuple{K}, geometry, geometrytype::Type{G}, index, skipmissing, weights, weightstype, kw..., | ||
| ) where {K, G} | ||
| indices = sample_indices(rng, x, skipmissing, weights, weightstype, n; kw...) | ||
| T = RA._rowtype(x, G; geometry, index, skipmissing, skipinvalid = _True(), names) | ||
| x2 = x isa AbstractRasterStack ? x[K] : RasterStack(NamedTuple{K}((x,))) | ||
| return _getindices(T, x2, dims, indices) | ||
| end | ||
| function _sample( | ||
| rng, x; | ||
| dims, names::NamedTuple{K}, geometry, geometrytype::Type{G}, index, skipmissing, weights, weightstype, kw..., | ||
| ) where {K, G} | ||
| indices = sample_indices(rng, x, skipmissing, weights, weightstype) | ||
| T = RA._rowtype(x, G; geometry, index, skipmissing, skipinvalid = _True(), names) | ||
| x2 = x isa AbstractRasterStack ? x[K] : RasterStack(NamedTuple{K}((x,))) | ||
| return _getindex(T, x2, dims, indices) | ||
| end | ||
|
|
||
| _getindices(::Type{T}, x, dims, indices) where {T} = | ||
| broadcast(I -> _getindex(T, x, dims, I), indices) | ||
|
|
||
| function _getindex(::Type{T}, x::AbstractRasterStack{<:Any, NT}, dims, idx) where {T, NT} | ||
| RA._maybe_add_fields( | ||
| T, | ||
| NT(x[RA.commondims(idx, x)]), | ||
| DimPoints(dims)[RA.commondims(idx, dims)], | ||
| val(idx) | ||
| ) | ||
| end | ||
|
|
||
| # args may be an integer or nothing | ||
| function sample_indices(rng, x, skipmissing, weights::Nothing, weightstype, args...; kw...) | ||
| if istrue(skipmissing) | ||
| wts = StatsBase.Weights(vec(boolmask(x))) | ||
| StatsBase.sample(rng, RA.DimIndices(x), wts, args...; kw...) | ||
| else | ||
| StatsBase.sample(rng, RA.DimIndices(x), args...; kw...) | ||
| end | ||
| end | ||
| function sample_indices(rng, x, skipmissing, weights::AbstractDimArray, ::Type{W}, args...; kw...) where W | ||
| wts = if istrue(skipmissing) | ||
| @d boolmask(x) .* weights | ||
| elseif dims(weights) == dims(x) | ||
| weights | ||
| else | ||
| @d ones(eltype(weights), dims(x)) .* weights | ||
| end |> vec |> W | ||
| StatsBase.sample(rng, RA.DimIndices(x), wts, args...; kw...) | ||
| end | ||
| function _geometrytype(x, geometry::Bool) | ||
| if geometry | ||
| error("Specify a geometry type by setting `geometry` to a Tuple or NamedTuple of Dimensions. E.g. `geometry = (X, Y)`") | ||
| else | ||
| return _False(), Nothing, dims(x) | ||
| end | ||
| end | ||
|
|
||
| function _geometrytype(x, geometry::Tuple) | ||
| dims = DD.commondims(DD.dims(x), geometry) | ||
| return _True(), Tuple{map(eltype, dims)...}, dims | ||
| end | ||
| function _geometrytype(x, geometry::NamedTuple{K}) where K | ||
| dims = DD.commondims(DD.dims(x), values(geometry)) | ||
| return _True(), NamedTuple{K, Tuple{map(eltype, dims)...}}, dims | ||
| end |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.