Sat, 28 Dec 2019 10:14:02 +0200
@background
0 | 1 | ######################### |
2 | # Some utility functions | |
3 | ######################### | |
4 | ||
4
59fd17a3cea0
Add __precompile__() for what it is worth
Tuomo Valkonen <tuomov@iki.fi>
parents:
0
diff
changeset
|
5 | __precompile__() |
59fd17a3cea0
Add __precompile__() for what it is worth
Tuomo Valkonen <tuomov@iki.fi>
parents:
0
diff
changeset
|
6 | |
0 | 7 | module Util |
8 | ||
9 | ############## | |
10 | # Our exports | |
11 | ############## | |
12 | ||
13 | export map_first_slice!, | |
14 | reduce_first_slice, | |
15 | norm₂, | |
16 | γnorm₂, | |
17 | norm₂w, | |
18 | norm₂², | |
19 | norm₂w², | |
20 | norm₂₁, | |
21 | γnorm₂₁, | |
22 | dot, | |
23 | mean, | |
24 | proj_norm₂₁ball!, | |
25 | curry, | |
8 | 26 | ⬿, |
9 | 27 | @threadsif, |
28 | @background | |
8 | 29 | |
30 | ||
31 | ########## | |
32 | # Threads | |
33 | ########## | |
34 | ||
35 | macro threadsif(threads, loop) | |
36 | return esc(:(if $threads | |
37 | Threads.@threads $loop | |
38 | else | |
39 | $loop | |
40 | end)) | |
41 | end | |
0 | 42 | |
9 | 43 | macro background(bgtask, fgtask) |
44 | return :(t = Threads.@spawn $(esc(bgtask)); | |
45 | $(esc(fgtask)); | |
46 | wait(t)) | |
47 | end | |
48 | ||
0 | 49 | ######################## |
50 | # Functional programming | |
51 | ######################### | |
52 | ||
53 | curry = (f::Function,y...)->(z...)->f(y...,z...) | |
54 | ||
55 | ############################### | |
56 | # For working with NamedTuples | |
57 | ############################### | |
58 | ||
59 | ⬿ = merge | |
60 | ||
61 | ###### | |
62 | # map | |
63 | ###### | |
64 | ||
65 | @inline function map_first_slice!(f!, y) | |
66 | for i in CartesianIndices(size(y)[2:end]) | |
67 | @inbounds f!(@view(y[:, i])) | |
68 | end | |
69 | end | |
70 | ||
71 | @inline function map_first_slice!(x, f!, y) | |
72 | for i in CartesianIndices(size(y)[2:end]) | |
73 | @inbounds f!(@view(x[:, i]), @view(y[:, i])) | |
74 | end | |
75 | end | |
76 | ||
77 | @inline function reduce_first_slice(f, y; init=0.0) | |
78 | accum=init | |
79 | for i in CartesianIndices(size(y)[2:end]) | |
80 | @inbounds accum=f(accum, @view(y[:, i])) | |
81 | end | |
82 | return accum | |
83 | end | |
84 | ||
85 | ########################### | |
86 | # Norms and inner products | |
87 | ########################### | |
88 | ||
89 | @inline function dot(x, y) | |
90 | @assert(length(x)==length(y)) | |
91 | ||
92 | accum=0 | |
93 | for i=1:length(y) | |
94 | @inbounds accum += x[i]*y[i] | |
95 | end | |
96 | return accum | |
97 | end | |
98 | ||
99 | @inline function norm₂w²(y, w) | |
100 | #Insane memory allocs | |
101 | #return @inbounds sum(i -> y[i]*y[i]*w[i], 1:length(y)) | |
102 | accum=0 | |
103 | for i=1:length(y) | |
104 | @inbounds accum=accum+y[i]*y[i]*w[i] | |
105 | end | |
106 | return accum | |
107 | end | |
108 | ||
109 | @inline function norm₂w(y, w) | |
110 | return √(norm₂w²(y, w)) | |
111 | end | |
112 | ||
113 | @inline function norm₂²(y) | |
114 | #Insane memory allocs | |
115 | #return @inbounds sum(i -> y[i]*y[i], 1:length(y)) | |
116 | accum=0 | |
117 | for i=1:length(y) | |
118 | @inbounds accum=accum+y[i]*y[i] | |
119 | end | |
120 | return accum | |
121 | end | |
122 | ||
123 | @inline function norm₂(y) | |
124 | return √(norm₂²(y)) | |
125 | end | |
126 | ||
127 | @inline function γnorm₂(y, γ) | |
128 | hubersq = xsq -> begin | |
129 | x=√xsq | |
130 | return if x > γ | |
131 | x-γ/2 | |
132 | elseif x<-γ | |
133 | -x-γ/2 | |
134 | else | |
135 | xsq/(2γ) | |
136 | end | |
137 | end | |
138 | ||
139 | if γ==0 | |
140 | return norm₂(y) | |
141 | else | |
142 | return hubersq(norm₂²(y)) | |
143 | end | |
144 | end | |
145 | ||
146 | function norm₂₁(y) | |
147 | return reduce_first_slice((s, x) -> s+norm₂(x), y) | |
148 | end | |
149 | ||
150 | function γnorm₂₁(y,γ) | |
151 | return reduce_first_slice((s, x) -> s+γnorm₂(x, γ), y) | |
152 | end | |
153 | ||
154 | function mean(v) | |
155 | return sum(v)/prod(size(v)) | |
156 | end | |
157 | ||
158 | @inline function proj_norm₂₁ball!(y, α) | |
159 | α²=α*α | |
160 | ||
7 | 161 | if ndims(y)==3 && size(y, 1)==2 |
162 | @inbounds for i=1:size(y, 2) | |
163 | @simd for j=1:size(y, 3) | |
164 | n² = y[1,i,j]*y[1,i,j]+y[2,i,j]*y[2,i,j] | |
165 | if n²>α² | |
166 | v = α/√n² | |
167 | y[1, i, j] *= v | |
168 | y[2, i, j] *= v | |
169 | end | |
170 | end | |
171 | end | |
172 | else | |
173 | y′=reshape(y, (size(y, 1), prod(size(y)[2:end]))) | |
174 | ||
175 | @inbounds @simd for i=1:size(y′, 2)# in CartesianIndices(size(y)[2:end]) | |
176 | n² = norm₂²(@view(y′[:, i])) | |
177 | if n²>α² | |
178 | y′[:, i] .*= (α/√n²) | |
179 | end | |
0 | 180 | end |
181 | end | |
182 | end | |
183 | ||
184 | end # Module | |
185 |