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