No commit message

paulxshen · Jan 5, 2025 · 454bca6 · 454bca6
1 parent 1e00506
commit 454bca6
Show file tree

Hide file tree

Showing 3 changed files with 15 additions and 6 deletions.
diff --git a/src/ArrayPadding.jl b/src/ArrayPadding.jl
@@ -1,5 +1,4 @@
 module ArrayPadding
 include("main.jl")
-export pad, pad!, Ramp, ReplicateRamp, PaddedArray, left, right
-export place, place!, array
+export pad, pad!, TanhRamp
 end # module Pad
diff --git a/src/alg.jl b/src/alg.jl
@@ -1,11 +1,12 @@
-function lr(a::S, v, i, l, r, ol=0, or=0, dofill=false) where {S}
+function lr(a::S, v0, i, l, r, ol=0, or=0, dofill=false) where {S}
     N = ndims(a)
     sz = size(a)
     T = eltype(a)
     ax = axes(a, i)
     sel = (1:N) .== i
     al = ar = nothing
     if l > 0
+        v = _fv(v0, l)
         if v == :periodic
             I = ifelse.(sel, (ax[end-l+1:end],), :)
             al = a[I...]
@@ -27,7 +28,8 @@ function lr(a::S, v, i, l, r, ol=0, or=0, dofill=false) where {S}
             #     al = 2a[I...] - a[I2...]
             # end
         elseif isa(v, Function)
-            al = repeat(v.(constructor(S)(1:l)), (sel .+ .!(sel) .* sz)...)
+            b = reshape(v.(constructor(S)(1:l)), (l .* sel .+ .!(sel))...)
+            al = repeat(b, (sel .+ .!(sel) .* sz)...)
         else
             if dofill
                 f = fillfunc(S)
@@ -39,6 +41,7 @@ function lr(a::S, v, i, l, r, ol=0, or=0, dofill=false) where {S}
         end
     end
     if r > 0
+        v = _fv(v0, r)
         if v == :periodic
             I = ifelse.(sel, (1:r,), :)
             ar = a[I...]
@@ -60,7 +63,8 @@ function lr(a::S, v, i, l, r, ol=0, or=0, dofill=false) where {S}
             #     ar = 2a[I...] - a[I2...]
             # end
         elseif isa(v, Function)
-            ar = repeat(v.(constructor(S)(1:r)), (sel .+ .!(sel) .* sz)...)
+            b = reshape(v.(constructor(S)(1:r)), (r .* sel .+ .!(sel))...)
+            ar = repeat(b, (sel .+ .!(sel) .* sz)...)
         else
             if dofill
                 f = fillfunc(S)

diff --git a/src/utils.jl b/src/utils.jl
@@ -7,4 +7,10 @@ constructor(::Type{<:Array}) = Array
 
 # (f::Function)(::Type{Buffer{T,S}}) where {T,S} = f(S)
 fillfunc(::Type{<:Buffer{T,S}}) where {T,S} = fillfunc(S)
-constructor(::Type{<:Buffer{T,S}}) where {T,S} = constructor(S)
+constructor(::Type{<:Buffer{T,S}}) where {T,S} = constructor(S)
+
+struct TanhRamp
+    v
+end
+_fv(v::TanhRamp, n) = i -> v.v * tanh(2i / n)
+_fv(v, n) = v