CCPP SciDoc v7.0.0  v7.0.0
Common Community Physics Package Developed at DTC
 
Loading...
Searching...
No Matches
module_mp_tempo_ml.F90
1! TEMPO neural network module designed for cloud droplet number concentration prediction
2! A big thanks to David John Gagne (NCAR) for very useful discussions
3! Please see https://github.com/NCAR/mlmicrophysics for a more detailed way to connect
4! ML models in python to Fortran microphysics code
5
6module module_mp_tempo_ml
7
8 implicit none
9 private
10 public :: mldata, tempo_save_or_read_ml_data, predict_number, predict_number_sub
11
12 type mldata
13 integer :: input_size
14 integer :: output_size
15 integer :: node_size
16 real, allocatable, dimension(:) :: transform_mean
17 real, allocatable, dimension(:) :: transform_var
18 real, allocatable, dimension(:,:) :: weights00
19 real, allocatable, dimension(:) :: bias00
20 real, allocatable, dimension(:,:) :: weights01
21 real, allocatable, dimension(:) :: bias01
22 end type mldata
23
24contains
25
26!----------------------------------------------------------------------------------
27 ! Initializes or returns neural network information.
28 ! Saves nueral network data information if a neural network is not
29 ! initialized. If output data is requested, the initialized and saved
30 ! neural network will be used.
31 subroutine tempo_save_or_read_ml_data(ml_data_in, ml_data_out)
32
33 logical, save :: not_initialized = .true.
34 type(mldata), dimension(2), intent(in), optional :: ml_data_in
35 type(mldata), dimension(2), intent(out), optional :: ml_data_out
36 type(mldata), dimension(2), save :: tempo_ml_data_save
37
38 if (not_initialized) then
39 if (present(ml_data_in)) tempo_ml_data_save = ml_data_in
40 not_initialized = .false.
41 endif
42
43 if (present(ml_data_out)) then
44 ml_data_out = tempo_ml_data_save
45 endif
46
47 end subroutine tempo_save_or_read_ml_data
48
49!------------------------------------------------------------------------------------------------------
50 ! Predicts number concentration
51 real function predict_number(qc, qr, qi, qs, pres, temp, w, predict_nc)
52
53 real, intent(in) :: qc, qr, qi, qs, pres, temp, w
54 logical, intent(in) :: predict_nc
55 type(mldata), dimension(2) :: get_ml_data
56 type(mldata) :: ml_data
57 integer, parameter :: input_rows = 1
58 real, allocatable, dimension(:,:) :: input, input_transformed
59 real, allocatable, dimension(:,:) :: output00, output00activ
60 real, allocatable, dimension(:,:) :: output01, output01activ
61 real, parameter :: logmin = -6.0 ! R2
62 real, parameter :: logmax = 9.3010299957 ! 2000 cm^-3
63 double precision :: predictexp
64
65 ! Get neural network data
66 call tempo_save_or_read_ml_data(ml_data_out=get_ml_data)
67
68 if (predict_nc) then
69 ml_data = get_ml_data(1)
70 else
71 ml_data = get_ml_data(2)
72 endif
73
74 ! Allocate arrays for input data and transformation
75 if (.not. allocated(input)) then
76 allocate(input(input_rows, ml_data%input_size))
77 endif
78 if (.not. allocated(input_transformed)) then
79 allocate(input_transformed(input_rows, ml_data%input_size))
80 endif
81
82 ! Collect input data
83 input(1,1) = qc
84 input(1,2) = qr
85 input(1,3) = qi
86 input(1,4) = qs
87 input(1,5) = pres
88 input(1,6) = temp
89 input(1,7) = w
90
91 ! Transform input data
92 call standard_scaler_transform(mean=ml_data%transform_mean, var=ml_data%transform_var, &
93 raw_data=input, transformed_data=input_transformed)
94
95 ! Allocate arrays
96 if (.not. allocated(output00)) then
97 allocate(output00(ml_data%node_size, ml_data%output_size))
98 endif
99 if (.not. allocated(output00activ)) then
100 allocate(output00activ(ml_data%node_size, ml_data%output_size))
101 endif
102 if (.not. allocated(output01)) then
103 allocate(output01(ml_data%output_size, ml_data%output_size))
104 endif
105 if (.not. allocated(output01activ)) then
106 allocate(output01activ(ml_data%output_size, ml_data%output_size))
107 endif
108
109 ! Reconstruct neural network
110 ! First layer
111 output00 = matmul(transpose(ml_data%weights00), transpose(input_transformed)) + &
112 reshape(ml_data%bias00, (/size(ml_data%bias00), ml_data%output_size/))
113 call relu_activation(input=output00, output=output00activ)
114
115 ! Second layer
116 output01 = matmul(transpose(ml_data%weights01), output00activ) + &
117 reshape(ml_data%bias01, (/size(ml_data%bias01), ml_data%output_size/))
118 call relu_activation(input=output01, output=output01activ)
119
120 ! Prediction
121 predictexp = min(logmax, max(logmin, output01activ(1,1)))
122 predict_number = 10.**predictexp
123
124 end function predict_number
125
126!------------------------------------------------------------------------------------------------------
127 ! Predicts number concentration
128 subroutine predict_number_sub(kts, kte, qc, qr, qi, qs, pres, temp, w, predict_number, predict_nc)
129
130 integer, intent(in) :: kts, kte
131 real, dimension(kts:kte), intent(in) :: qc, qr, qi, qs, pres, temp, w
132 double precision, dimension(kts:kte), intent(inout) :: predict_number
133 logical, intent(in) :: predict_nc
134 type(mldata), dimension(2) :: get_ml_data
135 type(mldata) :: ml_data
136 integer, parameter :: input_rows = 1
137 real, allocatable, dimension(:,:) :: input, input_transformed
138 real, allocatable, dimension(:,:) :: output00, output00activ, reshaped_bias00
139 real, allocatable, dimension(:,:) :: output01, output01activ, reshaped_bias01
140 real, parameter :: logmin = -6.0 ! R2
141 real, parameter :: logmax = 9.3010299957 ! 2000 cm^-3
142 double precision :: predictexp
143 integer :: k
144
145 ! Get neural network data
146 call tempo_save_or_read_ml_data(ml_data_out=get_ml_data)
147
148 if (predict_nc) then
149 ml_data = get_ml_data(1)
150 else
151 ml_data = get_ml_data(2)
152 endif
153
154 ! Allocate arrays for input data and transformation
155 if (.not. allocated(input)) then
156 allocate(input(ml_data%input_size, kte))
157 endif
158 if (.not. allocated(input_transformed)) then
159 allocate(input_transformed(ml_data%input_size, kte))
160 endif
161
162 ! Collect input data
163 input(1,:) = qc
164 input(2,:) = qr
165 input(3,:) = qi
166 input(4,:) = qs
167 input(5,:) = pres
168 input(6,:) = temp
169 input(7,:) = w
170
171 ! Transform input data
172 call standard_scaler_transform(mean=ml_data%transform_mean, var=ml_data%transform_var, &
173 raw_data=input, transformed_data=input_transformed)
174
175 ! Allocate arrays
176 if (.not. allocated(output00)) then
177 allocate(output00(ml_data%node_size, kte))
178 endif
179 if (.not. allocated(output00activ)) then
180 allocate(output00activ(ml_data%node_size, kte))
181 endif
182 if (.not. allocated(reshaped_bias00)) then
183 allocate(reshaped_bias00(ml_data%node_size, kte))
184 endif
185
186 if (.not. allocated(output01)) then
187 allocate(output01(ml_data%output_size, kte))
188 endif
189 if (.not. allocated(output01activ)) then
190 allocate(output01activ(ml_data%output_size, kte))
191 endif
192 if (.not. allocated(reshaped_bias01)) then
193 allocate(reshaped_bias01(ml_data%output_size, kte))
194 endif
195
196 do k = kts, kte
197 reshaped_bias00(:,k) = ml_data%bias00
198 reshaped_bias01(1,k) = ml_data%bias01(1)
199 enddo
200
201 ! Reconstruct neural network
202 ! First layer
203 output00 = matmul(ml_data%weights00, input_transformed) + reshaped_bias00
204 call relu_activation(input=output00, output=output00activ)
205
206 ! Second layer
207 output01 = matmul(ml_data%weights01, output00activ) + reshaped_bias01
208 call relu_activation(input=output01, output=output01activ)
209
210 do k = kts, kte
211 ! Prediction
212 predictexp = min(logmax, max(logmin, output01activ(1,k)))
213 predict_number(k) = 10.**predictexp
214 enddo
215
216 end subroutine predict_number_sub
217
218!------------------------------------------------------------------------------------------------------
219 ! Standard transformer
220 subroutine standard_scaler_transform(mean, var, raw_data, transformed_data)
221 real, dimension(:,:), intent(in) :: raw_data
222 real, dimension(:), intent(in) :: mean, var
223 real, intent(out) :: transformed_data(size(raw_data, 1), size(raw_data, 2))
224 integer :: i
225
226 do i = 1, size(raw_data, 1)
227 transformed_data(i,:) = (raw_data(i,:) - mean(i)) / sqrt(var(i))
228 end do
229
230 end subroutine standard_scaler_transform
231
232!----------------------------------------------------------------------------------
233 ! Activation function
234 subroutine relu_activation(input, output)
235
236 real, dimension(:,:), intent(in) :: input
237 real, dimension(size(input, 1), size(input, 2)), intent(out) :: output
238 integer :: i, j
239
240 do i = 1, size(input, 1)
241 do j = 1, size(input,2)
242 output(i, j) = max(input(i,j), 0.)
243 end do
244 end do
245
246 end subroutine relu_activation
247!----------------------------------------------------------------------------------
248
249end module module_mp_tempo_ml
module_mp_tempo_ml
Definition module_mp_tempo_ml.F90:6
module_mp_tempo_ml::mldata
Definition module_mp_tempo_ml.F90:12