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 0 : 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 : 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 : 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 : distances(i) = sqrt((lu_teff(closest_indices(i)) - teff)**2 + &
81 : (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 :
118 : integer :: i, n, j
119 0 : real(dp) :: distance, norm_teff, norm_logg, norm_meta
120 0 : real(dp), dimension(:), allocatable :: scaled_lu_teff, scaled_lu_logg, scaled_lu_meta
121 0 : real(dp), dimension(4) :: min_distances
122 : integer, dimension(4) :: indices
123 0 : real(dp) :: teff_min, teff_max, logg_min, logg_max, meta_min, meta_max
124 0 : real(dp) :: teff_dist, logg_dist, meta_dist
125 :
126 0 : n = size(lu_teff)
127 0 : min_distances = huge(1.0)
128 0 : indices = -1
129 :
130 : ! Find min and max for normalization
131 0 : teff_min = minval(lu_teff)
132 0 : teff_max = maxval(lu_teff)
133 0 : logg_min = minval(lu_logg)
134 0 : logg_max = maxval(lu_logg)
135 0 : meta_min = minval(lu_meta)
136 0 : meta_max = maxval(lu_meta)
137 :
138 : ! Allocate and scale lookup table values
139 0 : allocate (scaled_lu_teff(n), scaled_lu_logg(n), scaled_lu_meta(n))
140 :
141 0 : if (teff_max - teff_min > 0.00) then
142 0 : scaled_lu_teff = (lu_teff - teff_min)/(teff_max - teff_min)
143 : end if
144 :
145 0 : if (logg_max - logg_min > 0.00) then
146 0 : scaled_lu_logg = (lu_logg - logg_min)/(logg_max - logg_min)
147 : end if
148 :
149 0 : if (meta_max - meta_min > 0.00) then
150 0 : scaled_lu_meta = (lu_meta - meta_min)/(meta_max - meta_min)
151 : end if
152 :
153 : ! Normalize input parameters
154 0 : norm_teff = (teff - teff_min)/(teff_max - teff_min)
155 0 : norm_logg = (log_g - logg_min)/(logg_max - logg_min)
156 0 : norm_meta = (metallicity - meta_min)/(meta_max - meta_min)
157 :
158 : ! Find closest models
159 0 : do i = 1, n
160 0 : teff_dist = 0.0
161 0 : logg_dist = 0.0
162 0 : meta_dist = 0.0
163 :
164 0 : if (teff_max - teff_min > 0.00) then
165 0 : teff_dist = scaled_lu_teff(i) - norm_teff
166 : end if
167 :
168 0 : if (logg_max - logg_min > 0.00) then
169 0 : logg_dist = scaled_lu_logg(i) - norm_logg
170 : end if
171 :
172 0 : if (meta_max - meta_min > 0.00) then
173 0 : meta_dist = scaled_lu_meta(i) - norm_meta
174 : end if
175 :
176 : ! Using squared distance without sqrt (monotonic transform)
177 0 : distance = teff_dist**2 + logg_dist**2 + meta_dist**2
178 :
179 0 : do j = 1, 4
180 0 : if (distance < min_distances(j)) then
181 : ! Shift larger distances down
182 0 : if (j < 4) then
183 0 : min_distances(j + 1:4) = min_distances(j:3)
184 0 : indices(j + 1:4) = indices(j:3)
185 : end if
186 0 : min_distances(j) = distance
187 0 : indices(j) = i
188 0 : exit
189 : end if
190 : end do
191 : end do
192 :
193 0 : closest_indices = indices
194 0 : end subroutine get_closest_stellar_models
195 :
196 : !---------------------------------------------------------------------------
197 : ! Linear interpolation (binary search version for efficiency)
198 : !---------------------------------------------------------------------------
199 0 : subroutine linear_interpolate(x, y, x_val, y_val)
200 : real(dp), intent(in) :: x(:), y(:), x_val
201 : real(dp), intent(out) :: y_val
202 : integer :: low, high, mid
203 :
204 : ! Validate input sizes
205 0 : if (size(x) < 2) then
206 0 : print *, "Error: x array has fewer than 2 points."
207 0 : y_val = 0.0_dp
208 : return
209 : end if
210 :
211 0 : if (size(x) /= size(y)) then
212 0 : print *, "Error: x and y arrays have different sizes."
213 0 : y_val = 0.0_dp
214 : return
215 : end if
216 :
217 : ! Handle out-of-bounds cases
218 0 : if (x_val <= x(1)) then
219 0 : y_val = y(1)
220 0 : return
221 0 : else if (x_val >= x(size(x))) then
222 0 : y_val = y(size(y))
223 : return
224 : end if
225 :
226 : ! Binary search to find the proper interval [x(low), x(low+1)]
227 0 : low = 1
228 0 : high = size(x)
229 0 : do while (high - low > 1)
230 0 : mid = (low + high)/2
231 0 : if (x(mid) <= x_val) then
232 : low = mid
233 : else
234 0 : high = mid
235 : end if
236 : end do
237 :
238 : ! Linear interpolation between x(low) and x(low+1)
239 0 : y_val = y(low) + (y(low + 1) - y(low))/(x(low + 1) - x(low))*(x_val - x(low))
240 : end subroutine linear_interpolate
241 :
242 : !---------------------------------------------------------------------------
243 : ! Array interpolation for SED construction
244 : !---------------------------------------------------------------------------
245 0 : subroutine interpolate_array(x_in, y_in, x_out, y_out)
246 : real(dp), intent(in) :: x_in(:), y_in(:), x_out(:)
247 : real(dp), intent(out) :: y_out(:)
248 : integer :: i
249 :
250 : ! Validate input sizes
251 0 : if (size(x_in) < 2 .or. size(y_in) < 2) then
252 0 : print *, "Error: x_in or y_in arrays have fewer than 2 points."
253 0 : stop
254 : end if
255 :
256 0 : if (size(x_in) /= size(y_in)) then
257 0 : print *, "Error: x_in and y_in arrays have different sizes."
258 0 : stop
259 : end if
260 :
261 0 : if (size(x_out) <= 0) then
262 0 : print *, "Error: x_out array is empty."
263 0 : stop
264 : end if
265 :
266 0 : do i = 1, size(x_out)
267 0 : call linear_interpolate(x_in, y_in, x_out(i), y_out(i))
268 : end do
269 0 : end subroutine interpolate_array
270 :
271 : end module knn_interp
|
