comparison: src/Gradient.jl
src/Gradient.jl
- changeset 1
- 5762e68842c9
- parent 0
- 888dfd34d24a
- child 2
- fc7430da1d7a
equal
deleted
inserted
replaced
| 0:888dfd34d24a | 1:5762e68842c9 |
|---|---|
| 1 ######################## | |
| 2 # Discretised gradients | |
| 3 ######################## | |
| 4 | |
| 5 module Gradient | |
| 6 | |
| 7 ############## | |
| 8 # Our exports | |
| 9 ############## | |
| 10 | |
| 11 export ∇₂!, ∇₂ᵀ!, ∇₂fold!, | |
| 12 ∇₂_norm₂₂_est, ∇₂_norm₂₂_est², | |
| 13 ∇₂_norm₂∞_est, ∇₂_norm₂∞_est², | |
| 14 ∇₂c!, | |
| 15 ∇₃!, ∇₃ᵀ!, | |
| 16 vec∇₃!, vec∇₃ᵀ! | |
| 17 | |
| 18 ################## | |
| 19 # Helper routines | |
| 20 ################## | |
| 21 | |
| 22 @inline function imfold₂′!(f_aa!, f_a0!, f_ab!, | |
| 23 f_0a!, f_00!, f_0b!, | |
| 24 f_ba!, f_b0!, f_bb!, | |
| 25 n, m, state) | |
| 26 # First row | |
| 27 state = f_aa!(state, (1, 1)) | |
| 28 for j = 2:m-1 | |
| 29 state = f_a0!(state, (1, j)) | |
| 30 end | |
| 31 state = f_ab!(state, (1, m)) | |
| 32 | |
| 33 # Middle rows | |
| 34 for i=2:n-1 | |
| 35 state = f_0a!(state, (i, 1)) | |
| 36 for j = 2:m-1 | |
| 37 state = f_00!(state, (i, j)) | |
| 38 end | |
| 39 state = f_0b!(state, (i, m)) | |
| 40 end | |
| 41 | |
| 42 # Last row | |
| 43 state = f_ba!(state, (n, 1)) | |
| 44 for j =2:m-1 | |
| 45 state = f_b0!(state, (n, j)) | |
| 46 end | |
| 47 return f_bb!(state, (n, m)) | |
| 48 end | |
| 49 | |
| 50 ######################### | |
| 51 # 2D forward differences | |
| 52 ######################### | |
| 53 | |
| 54 ∇₂_norm₂₂_est² = 8 | |
| 55 ∇₂_norm₂₂_est = √∇₂_norm₂₂_est² | |
| 56 ∇₂_norm₂∞_est² = 2 | |
| 57 ∇₂_norm₂∞_est = √∇₂_norm₂∞_est² | |
| 58 | |
| 59 function ∇₂!(u₁, u₂, u) | |
| 60 @. @views begin | |
| 61 u₁[1:(end-1), :] = u[2:end, :] - u[1:(end-1), :] | |
| 62 u₁[end, :, :] = 0 | |
| 63 | |
| 64 u₂[:, 1:(end-1)] = u[:, 2:end] - u[:, 1:(end-1)] | |
| 65 u₂[:, end] = 0 | |
| 66 end | |
| 67 return u₁, u₂ | |
| 68 end | |
| 69 | |
| 70 function ∇₂!(v, u) | |
| 71 ∇₂!(@view(v[1, :, :]), @view(v[2, :, :]), u) | |
| 72 end | |
| 73 | |
| 74 @inline function ∇₂fold!(f!::Function, u, state) | |
| 75 g! = (state, pt) -> begin | |
| 76 (i, j) = pt | |
| 77 g = @inbounds [u[i+1, j]-u[i, j], u[i, j+1]-u[i, j]] | |
| 78 return f!(g, state, pt) | |
| 79 end | |
| 80 gr! = (state, pt) -> begin | |
| 81 (i, j) = pt | |
| 82 g = @inbounds [u[i+1, j]-u[i, j], 0.0] | |
| 83 return f!(g, state, pt) | |
| 84 end | |
| 85 gb! = (state, pt) -> begin | |
| 86 (i, j) = pt | |
| 87 g = @inbounds [0.0, u[i, j+1]-u[i, j]] | |
| 88 return f!(g, state, pt) | |
| 89 end | |
| 90 g0! = (state, pt) -> begin | |
| 91 return f!([0.0, 0.0], state, pt) | |
| 92 end | |
| 93 return imfold₂′!(g!, g!, gr!, | |
| 94 g!, g!, gr!, | |
| 95 gb!, gb!, g0!, | |
| 96 size(u, 1), size(u, 2), state) | |
| 97 end | |
| 98 | |
| 99 function ∇₂ᵀ!(v, v₁, v₂) | |
| 100 @. @views begin | |
| 101 v[2:(end-1), :] = v₁[1:(end-2), :] - v₁[2:(end-1), :] | |
| 102 v[1, :] = -v₁[1, :] | |
| 103 v[end, :] = v₁[end-1, :] | |
| 104 | |
| 105 v[:, 2:(end-1)] += v₂[:, 1:(end-2)] - v₂[:, 2:(end-1)] | |
| 106 v[:, 1] += -v₂[:, 1] | |
| 107 v[:, end] += v₂[:, end-1] | |
| 108 end | |
| 109 return v | |
| 110 end | |
| 111 | |
| 112 function ∇₂ᵀ!(u, v) | |
| 113 ∇₂ᵀ!(u, @view(v[1, :, :]), @view(v[2, :, :])) | |
| 114 end | |
| 115 | |
| 116 ################################################## | |
| 117 # 2D central differences (partial implementation) | |
| 118 ################################################## | |
| 119 | |
| 120 function ∇₂c!(v, u) | |
| 121 @. @views begin | |
| 122 v[1, 2:(end-1), :] = (u[3:end, :] - u[1:(end-2), :])/2 | |
| 123 v[1, end, :] = (u[end, :] - u[end-1, :])/2 | |
| 124 v[1, 1, :] = (u[2, :] - u[1, :])/2 | |
| 125 | |
| 126 v[2, :, 2:(end-1)] = (u[:, 3:end] - u[:, 1:(end-2)])/2 | |
| 127 v[2, :, end] = (u[:, end] - u[:, end-1])/2 | |
| 128 v[2, :, 1] = (u[:, 2] - u[:, 1])/2 | |
| 129 end | |
| 130 end | |
| 131 | |
| 132 ######################### | |
| 133 # 3D forward differences | |
| 134 ######################### | |
| 135 | |
| 136 function ∇₃!(u₁,u₂,u₃,u) | |
| 137 @. @views begin | |
| 138 u₁[1:(end-1), :, :] = u[2:end, :, :] - u[1:(end-1), :, :] | |
| 139 u₁[end, :, :] = 0 | |
| 140 | |
| 141 u₂[:, 1:(end-1), :] = u[:, 2:end, :] - u[:, 1:(end-1), :] | |
| 142 u₂[:, end, :] = 0 | |
| 143 | |
| 144 u₃[:, :, 1:(end-1)] = u[:, :, 2:end] - u[:, :, 1:(end-1)] | |
| 145 u₃[:, :, end] = 0 | |
| 146 end | |
| 147 return u₁, u₂, u₃ | |
| 148 end | |
| 149 | |
| 150 function ∇₃ᵀ!(v,v₁,v₂,v₃) | |
| 151 @. @views begin | |
| 152 v[2:(end-1), :, :] = v₁[1:(end-2), :, :] - v₁[2:(end-1), :, :] | |
| 153 v[1, :, :] = -v₁[1, :, :] | |
| 154 v[end, :, :] = v₁[end-1, :, :] | |
| 155 | |
| 156 v[:, 2:(end-1), :] += v₂[:, 1:(end-2), :] - v₂[:, 2:(end-1), :] | |
| 157 v[:, 1, :] += -v₂[:, 1, :] | |
| 158 v[:, end, :] += v₂[:, end-1, :] | |
| 159 | |
| 160 v[:, :, 2:(end-1)] += v₃[:, :, 1:(end-2)] - v₃[:, :, 2:(end-1)] | |
| 161 v[:, :, 1] += -v₃[:, :, 1] | |
| 162 v[:, :, end] += v₃[:, :, end-1] | |
| 163 end | |
| 164 return v | |
| 165 end | |
| 166 | |
| 167 ########################################### | |
| 168 # 3D forward differences for vector fields | |
| 169 ########################################### | |
| 170 | |
| 171 function vec∇₃!(u₁,u₂,u₃,u) | |
| 172 @. @views for j=1:size(u, 1) | |
| 173 ∇₃!(u₁[j, :, :, :],u₂[j, :, :, :],u₃[j, :, :, :],u[j, :, :, :]) | |
| 174 end | |
| 175 return u₁, u₂, u₃ | |
| 176 end | |
| 177 | |
| 178 function vec∇₃ᵀ!(u,v₁,v₂,v₃) | |
| 179 @. @views for j=1:size(u, 1) | |
| 180 ∇₃ᵀ!(u[j, :, :, :],v₁[j, :, :, :],v₂[j, :, :, :],v₃[j, :, :, :]) | |
| 181 end | |
| 182 return u | |
| 183 end | |
| 184 | |
| 185 end # Module |
