Line data Source code
1 : ! ***********************************************************************
2 : !
3 : ! Copyright (C) 2025 Niall Miller & The MESA Team
4 : !
5 : ! This program is free software: you can redistribute it and/or modify
6 : ! it under the terms of the GNU Lesser General Public License
7 : ! as published by the Free Software Foundation,
8 : ! either version 3 of the License, or (at your option) any later version.
9 : !
10 : ! This program is distributed in the hope that it will be useful,
11 : ! but WITHOUT ANY WARRANTY; without even the implied warranty of
12 : ! MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
13 : ! See the GNU Lesser General Public License for more details.
14 : !
15 : ! You should have received a copy of the GNU Lesser General Public License
16 : ! along with this program. If not, see <https://www.gnu.org/licenses/>.
17 : !
18 : ! ***********************************************************************
19 :
20 : ! ***********************************************************************
21 : ! K-Nearest Neighbors interpolation module for spectral energy distributions (SEDs)
22 : ! ***********************************************************************
23 :
24 : module knn_interp
25 : use const_def, only: dp
26 : use colors_utils, only: dilute_flux, load_sed
27 : implicit none
28 :
29 : private
30 : public :: construct_sed_knn, load_sed, interpolate_array, dilute_flux
31 :
32 : contains
33 :
34 : !---------------------------------------------------------------------------
35 : ! Main entry point: Construct a SED using KNN interpolation
36 : !---------------------------------------------------------------------------
37 0 : subroutine construct_sed_knn(teff, log_g, metallicity, R, d, file_names, &
38 0 : lu_teff, lu_logg, lu_meta, stellar_model_dir, &
39 : wavelengths, fluxes)
40 : real(dp), intent(in) :: teff, log_g, metallicity, R, d
41 : real(dp), intent(in) :: lu_teff(:), lu_logg(:), lu_meta(:)
42 : character(len=*), intent(in) :: stellar_model_dir
43 : character(len=100), intent(in) :: file_names(:)
44 : real(dp), dimension(:), allocatable, intent(out) :: wavelengths, fluxes
45 :
46 : integer, dimension(4) :: closest_indices
47 0 : real(dp), dimension(:), allocatable :: temp_wavelengths, temp_flux, common_wavelengths
48 0 : real(dp), dimension(:, :), allocatable :: model_fluxes
49 : real(dp), dimension(4) :: weights, distances
50 : integer :: i, n_points
51 : real(dp) :: sum_weights
52 0 : real(dp), dimension(:), allocatable :: diluted_flux
53 :
54 : ! Get the four closest stellar models
55 : call get_closest_stellar_models(teff, log_g, metallicity, lu_teff, &
56 0 : lu_logg, lu_meta, closest_indices)
57 :
58 : ! Load the first SED to define the wavelength grid
59 0 : call load_sed(trim(stellar_model_dir)//trim(file_names(closest_indices(1))), &
60 0 : closest_indices(1), temp_wavelengths, temp_flux)
61 :
62 0 : n_points = size(temp_wavelengths)
63 0 : allocate (common_wavelengths(n_points))
64 0 : common_wavelengths = temp_wavelengths
65 :
66 : ! Allocate flux array for the models (4 models, n_points each)
67 0 : allocate (model_fluxes(4, n_points))
68 0 : call interpolate_array(temp_wavelengths, temp_flux, common_wavelengths, model_fluxes(1, :))
69 :
70 : ! Load and interpolate remaining SEDs
71 0 : do i = 2, 4
72 0 : call load_sed(trim(stellar_model_dir)//trim(file_names(closest_indices(i))), &
73 0 : closest_indices(i), temp_wavelengths, temp_flux)
74 :
75 0 : call interpolate_array(temp_wavelengths, temp_flux, common_wavelengths, model_fluxes(i, :))
76 : end do
77 :
78 : ! Compute distances and weights for the four models
79 0 : do i = 1, 4
80 0 : distances(i) = sqrt((lu_teff(closest_indices(i)) - teff)**2 + &
81 0 : (lu_logg(closest_indices(i)) - log_g)**2 + &
82 0 : (lu_meta(closest_indices(i)) - metallicity)**2)
83 0 : if (distances(i) == 0.0) distances(i) = 1.0d-10 ! Prevent division by zero
84 0 : weights(i) = 1.0/distances(i)
85 : end do
86 :
87 : ! Normalize weights
88 0 : sum_weights = sum(weights)
89 0 : weights = weights/sum_weights
90 :
91 : ! Allocate output arrays
92 0 : allocate (wavelengths(n_points), fluxes(n_points))
93 0 : wavelengths = common_wavelengths
94 0 : fluxes = 0.0
95 :
96 : ! Perform weighted combination of the model fluxes (still at the stellar surface)
97 0 : do i = 1, 4
98 0 : fluxes = fluxes + weights(i)*model_fluxes(i, :)
99 : end do
100 :
101 : ! Now, apply the dilution factor (R/d)^2 to convert the surface flux density
102 : ! into the observed flux density at Earth.
103 0 : allocate (diluted_flux(n_points))
104 0 : call dilute_flux(fluxes, R, d, diluted_flux)
105 0 : fluxes = diluted_flux
106 :
107 0 : end subroutine construct_sed_knn
108 :
109 : !---------------------------------------------------------------------------
110 : ! Identify the four closest stellar models
111 : !---------------------------------------------------------------------------
112 0 : subroutine get_closest_stellar_models(teff, log_g, metallicity, lu_teff, &
113 0 : lu_logg, lu_meta, closest_indices)
114 : real(dp), intent(in) :: teff, log_g, metallicity
115 : real(dp), intent(in) :: lu_teff(:), lu_logg(:), lu_meta(:)
116 : integer, dimension(4), intent(out) :: closest_indices
117 : logical :: use_teff_dim, use_logg_dim, use_meta_dim
118 :
119 : integer :: i, n, j
120 : real(dp) :: distance, norm_teff, norm_logg, norm_meta
121 0 : real(dp), dimension(:), allocatable :: scaled_lu_teff, scaled_lu_logg, scaled_lu_meta
122 : real(dp), dimension(4) :: min_distances
123 : integer, dimension(4) :: indices
124 : real(dp) :: teff_min, teff_max, logg_min, logg_max, meta_min, meta_max
125 : real(dp) :: teff_dist, logg_dist, meta_dist
126 :
127 0 : n = size(lu_teff)
128 0 : min_distances = huge(1.0)
129 0 : indices = -1
130 :
131 : ! Find min and max for normalization
132 0 : teff_min = minval(lu_teff)
133 0 : teff_max = maxval(lu_teff)
134 0 : logg_min = minval(lu_logg)
135 0 : logg_max = maxval(lu_logg)
136 0 : meta_min = minval(lu_meta)
137 0 : meta_max = maxval(lu_meta)
138 :
139 : ! Allocate and scale lookup table values
140 0 : allocate (scaled_lu_teff(n), scaled_lu_logg(n), scaled_lu_meta(n))
141 :
142 0 : if (teff_max - teff_min > 0.00) then
143 0 : scaled_lu_teff = (lu_teff - teff_min)/(teff_max - teff_min)
144 : end if
145 :
146 0 : if (logg_max - logg_min > 0.00) then
147 0 : scaled_lu_logg = (lu_logg - logg_min)/(logg_max - logg_min)
148 : end if
149 :
150 0 : if (meta_max - meta_min > 0.00) then
151 0 : scaled_lu_meta = (lu_meta - meta_min)/(meta_max - meta_min)
152 : end if
153 :
154 : ! Normalize input parameters
155 0 : norm_teff = (teff - teff_min)/(teff_max - teff_min)
156 0 : norm_logg = (log_g - logg_min)/(logg_max - logg_min)
157 0 : norm_meta = (metallicity - meta_min)/(meta_max - meta_min)
158 :
159 : ! Find closest models
160 0 : do i = 1, n
161 0 : teff_dist = 0.0
162 0 : logg_dist = 0.0
163 0 : meta_dist = 0.0
164 :
165 0 : if (teff_max - teff_min > 0.00) then
166 0 : teff_dist = scaled_lu_teff(i) - norm_teff
167 : end if
168 :
169 0 : if (logg_max - logg_min > 0.00) then
170 0 : logg_dist = scaled_lu_logg(i) - norm_logg
171 : end if
172 :
173 0 : if (meta_max - meta_min > 0.00) then
174 0 : meta_dist = scaled_lu_meta(i) - norm_meta
175 : end if
176 :
177 :
178 : ! Detect dummy axes once
179 :
180 : use_teff_dim = .not.(all(lu_teff == 0.0_dp) .or. all(lu_teff == 999.0_dp) .or. all(lu_teff == -999.0_dp))
181 : use_logg_dim = .not.(all(lu_logg == 0.0_dp) .or. all(lu_logg == 999.0_dp) .or. all(lu_logg == -999.0_dp))
182 : use_meta_dim = .not.(all(lu_meta == 0.0_dp) .or. all(lu_meta == 999.0_dp) .or. all(lu_meta == -999.0_dp))
183 :
184 :
185 :
186 : ! Inside the loop:
187 : distance = 0.0_dp
188 0 : if (use_teff_dim) distance = distance + teff_dist**2
189 0 : if (use_logg_dim) distance = distance + logg_dist**2
190 0 : if (use_meta_dim) distance = distance + meta_dist**2
191 :
192 :
193 :
194 0 : distance = teff_dist**2 + logg_dist**2 + meta_dist**2
195 :
196 0 : do j = 1, 4
197 0 : if (distance < min_distances(j)) then
198 : ! Shift larger distances down
199 0 : if (j < 4) then
200 0 : min_distances(j + 1:4) = min_distances(j:3)
201 0 : indices(j + 1:4) = indices(j:3)
202 : end if
203 0 : min_distances(j) = distance
204 0 : indices(j) = i
205 0 : exit
206 : end if
207 : end do
208 : end do
209 :
210 0 : closest_indices = indices
211 0 : end subroutine get_closest_stellar_models
212 :
213 : !---------------------------------------------------------------------------
214 : ! Linear interpolation (binary search version for efficiency)
215 : !---------------------------------------------------------------------------
216 0 : subroutine linear_interpolate(x, y, x_val, y_val)
217 : real(dp), intent(in) :: x(:), y(:), x_val
218 : real(dp), intent(out) :: y_val
219 : integer :: low, high, mid
220 :
221 : ! Validate input sizes
222 0 : if (size(x) < 2) then
223 0 : print *, "Error: x array has fewer than 2 points."
224 0 : y_val = 0.0_dp
225 : return
226 : end if
227 :
228 0 : if (size(x) /= size(y)) then
229 0 : print *, "Error: x and y arrays have different sizes."
230 0 : y_val = 0.0_dp
231 : return
232 : end if
233 :
234 : ! Handle out-of-bounds cases
235 0 : if (x_val <= x(1)) then
236 0 : y_val = y(1)
237 0 : return
238 0 : else if (x_val >= x(size(x))) then
239 0 : y_val = y(size(y))
240 : return
241 : end if
242 :
243 : ! Binary search to find the proper interval [x(low), x(low+1)]
244 0 : low = 1
245 0 : high = size(x)
246 0 : do while (high - low > 1)
247 0 : mid = (low + high)/2
248 0 : if (x(mid) <= x_val) then
249 : low = mid
250 : else
251 0 : high = mid
252 : end if
253 : end do
254 :
255 : ! Linear interpolation between x(low) and x(low+1)
256 0 : y_val = y(low) + (y(low + 1) - y(low))/(x(low + 1) - x(low))*(x_val - x(low))
257 : end subroutine linear_interpolate
258 :
259 : !---------------------------------------------------------------------------
260 : ! Array interpolation for SED construction
261 : !---------------------------------------------------------------------------
262 0 : subroutine interpolate_array(x_in, y_in, x_out, y_out)
263 : real(dp), intent(in) :: x_in(:), y_in(:), x_out(:)
264 : real(dp), intent(out) :: y_out(:)
265 : integer :: i
266 :
267 : ! Validate input sizes
268 0 : if (size(x_in) < 2 .or. size(y_in) < 2) then
269 0 : print *, "Error: x_in or y_in arrays have fewer than 2 points."
270 0 : stop
271 : end if
272 :
273 0 : if (size(x_in) /= size(y_in)) then
274 0 : print *, "Error: x_in and y_in arrays have different sizes."
275 0 : stop
276 : end if
277 :
278 0 : if (size(x_out) <= 0) then
279 0 : print *, "Error: x_out array is empty."
280 0 : stop
281 : end if
282 :
283 0 : do i = 1, size(x_out)
284 0 : call linear_interpolate(x_in, y_in, x_out(i), y_out(i))
285 : end do
286 0 : end subroutine interpolate_array
287 :
288 : end module knn_interp
|
