Like 0
Fork 0

Last active 1742233591

yuanpeng's Avatar yuanpeng revised this gist 1742233591. Go to revision

1 file changed, 1010 insertions

texture_proc_real_1step_not_aligned_2step_aligned.py(file created)

@@ -0,0 +1,1010 @@
1 + import json
2 + import numpy as np
3 + import os
4 + import scipy
5 + from scipy.optimize import minimize
6 + from scipy.optimize import curve_fit
7 + from scipy.signal import argrelextrema
8 + from pystog import Pre_Proc
9 + import matplotlib.pyplot as plt
10 + import random
11 +
12 + rebin = Pre_Proc.rebin
13 +
14 + # ================
15 + # Input parameters
16 + # vvvvvvvvvvvvvvvv
17 +
18 + # ----------
19 + # Some flags
20 + # ----------
21 + testing = False # flag for testing
22 +
23 + # ------------------------------------------------------------------
24 + # Directory and config files definition
25 + #
26 + # run_dir => directory for storing the reduced data files
27 + # according to the the two theta and azimuthal
28 + # angle bins.
29 + # out_dir => directory for storing the outputs.
30 + # det_map_file => detector group information for all the two
31 + # theta and azimuthal bins.
32 + # ttheta_group_file => config file defining groups for second stage
33 + # processing.
34 + # ------------------------------------------------------------------
35 + run_dir = os.path.join(
36 + ".",
37 + "texture_proc"
38 + )
39 + out_dir = os.path.join(
40 + run_dir,
41 + "texture_proc_output"
42 + )
43 + det_map_file = os.path.join(
44 + "/SNS/NOM/shared/scripts/texture",
45 + "nom_texture_pointer.json"
46 + )
47 + ttheta_group_file = os.path.join(
48 + ".",
49 + "ttheta_group_params_new_new_new.json"
50 + )
51 +
52 + # ---------------------
53 + # Processing parameters
54 + # ---------------------
55 + stem_name = "NOM_Si_640e_"
56 + ttheta_binw = 5. # in degrees
57 + azm_binw = 10. # in degrees
58 + header_num = 2
59 + sph_l_max = 1 # for first stage
60 + tt_sph_l_max = 2 # for second stage
61 + amp_thresh = .9 # Threshold for peak intensity -- in a certain two theta
62 + # bins, all patterns with the fitted peak intensity
63 + # different from the median intensity by more than this
64 + # percentage threshold will be discarded.
65 +
66 + # ^^^^^^^^^^^^^^^^
67 + # Input parameters
68 + # ================
69 +
70 + if testing:
71 + print("\n!!!!!!!!!!!!!!!!!!!!!!!!!!!!!")
72 + print("[Warning] Testing mode is on.")
73 + print("!!!!!!!!!!!!!!!!!!!!!!!!!!!!!\n")
74 +
75 +
76 + def gaussian(x, a, b, c):
77 + """Gaussian function.
78 + Parameters
79 + ----------
80 + x : ndarray Independent variable.
81 + a : float Amplitude.
82 + b : float Center.
83 + c : flaot Width.
84 +
85 + Returns
86 + -------
87 + float Gaussian function value.
88 + """
89 + return a * np.exp(-(x - b)**2 / (2 * c**2))
90 +
91 +
92 + def bkg_expo(x, a, b, c):
93 + """Background exponential function.
94 + Parameters
95 + ----------
96 + x : ndarray Independent variable.
97 + a : float offset.
98 + b : float coefficient.
99 + c : flaot exponent.
100 +
101 + Returns
102 + -------
103 + float Exponential background value.
104 + """
105 + return a - b * np.exp(-c * x**2.)
106 +
107 +
108 + def sph_harm_real(m, l, theta, phi):
109 + """Real spherical harmonics. Refer to the Wikipedia,
110 + https://en.wikipedia.org/wiki/Table_of_spherical_harmonics
111 + We could follow the standard definition of the spherical harmonics to
112 + define the function here, but for the purpose of this script, we only
113 + need the first few terms so we just implemented the function directly.
114 +
115 + Parameters
116 + ----------
117 + m : order of spherical harmonics
118 + l : degree of spherical harmonics
119 + theta : polar angle in radian
120 + phi : azimuthal angle in radian
121 +
122 + Returns
123 + -------
124 + float Spherical harmonics real value
125 + """
126 + if l == 0 and m == 0:
127 + return np.sqrt(1. / (4. * np.pi))
128 + if l == 1 and m == -1:
129 + return np.sqrt(3. / (4. * np.pi)) * np.sin(phi) * np.sin(theta)
130 + if l == 1 and m == 0:
131 + return np.sqrt(3. / (4. * np.pi)) * np.cos(theta)
132 + if l == 1 and m == 1:
133 + return np.sqrt(3. / (4. * np.pi)) * np.cos(phi) * np.sin(theta)
134 + if l == 2 and m == -2:
135 + part = np.sin(phi) * np.cos(phi) * (np.sin(theta))**2.
136 + return np.sqrt(15. / (4. * np.pi)) * part
137 + if l == 2 and m == -1:
138 + part = np.sin(phi) * np.sin(theta) * np.cos(theta)
139 + return np.sqrt(15. / (4. * np.pi)) * part
140 + if l == 2 and m == 0:
141 + part = 3. * (np.cos(theta))**2. - 1.
142 + return np.sqrt(5. / (16. * np.pi)) * part
143 + if l == 2 and m == 1:
144 + part = np.cos(phi) * np.sin(theta) * np.cos(theta)
145 + return np.sqrt(15. / (4. * np.pi)) * part
146 + if l == 2 and m == 2:
147 + part = ((np.cos(phi))**2. - (np.sin(phi))**2.) * (np.sin(theta))**2.
148 + return np.sqrt(15. / (16. * np.pi)) * part
149 +
150 +
151 + def target_function(params, sph_mat, txt_int):
152 + """Target function for the optimization of the spherical harmonic
153 + coefficients.
154 +
155 + Parameters
156 + ----------
157 + params : ndarray Spherical harmonic coefficients.
158 + sph_mat : ndarray Spherical harmonic matrix.
159 + txt_int : ndarray Textured S(Q) intensity.
160 +
161 + Returns
162 + -------
163 + float : Norm of the residuals.
164 + """
165 + result = sph_mat @ params
166 + residuals = result - txt_int
167 +
168 + return np.linalg.norm(residuals)
169 +
170 +
171 + if not os.path.exists(out_dir):
172 + os.makedirs(out_dir)
173 +
174 + # ================
175 + # Read in Q values
176 + # ================
177 + print("[Info] Read in Q values...")
178 + bank = 1
179 + while True:
180 + tr_file = os.path.join(
181 + run_dir,
182 + stem_name + f"bank{bank}.dat"
183 + )
184 + if os.path.exists(tr_file):
185 + data = np.loadtxt(tr_file, skiprows=header_num, dtype=float)
186 + q_vals = data[:, 0]
187 + break
188 + else:
189 + bank += 1
190 +
191 + # ===================
192 + # Read in S(Q) values
193 + # ===================
194 + print("[Info] Read in all S(Q)'s...")
195 + num_groups = int(180. / ttheta_binw) * int(360. / azm_binw)
196 + all_sofq = dict()
197 + for i in range(num_groups):
198 + group_tmp = i + 1
199 + tr_file = os.path.join(
200 + run_dir,
201 + stem_name + "bank%d.dat" % group_tmp
202 + )
203 + if os.path.exists(tr_file):
204 + data = np.loadtxt(tr_file, skiprows=header_num, dtype=float)
205 + sofq_tmp = data[:, 1]
206 + all_sofq[group_tmp] = sofq_tmp
207 +
208 + # ===========================================================================
209 + # Read in the detector map, telling the center two theta and azimuthal angle,
210 + # and the ID for each group of detectors.
211 + # ===========================================================================
212 + det_pointer = json.load(open(det_map_file, "r"))
213 +
214 + # =============================================================================
215 + # Determine the valid azimuthal bins for each two theta bin. Here follows are
216 + # the criteria:
217 + # 1. The 2theta bin is in use
218 + # 2. The peak fitting succeeded
219 + # 3. The fitted peak intensity is close to the median intensity
220 + # =============================================================================
221 + with open(ttheta_group_file, "r") as f:
222 + ttheta_groups = json.load(f)
223 +
224 + print(
225 + "[Info] Pre-processing the data for all "
226 + "the two theta and azimuthal bins..."
227 + )
228 + ttheta_bins = np.arange(0, 181, ttheta_binw)
229 + for i in range(len(ttheta_bins) - 1):
230 + ttheta_cv = (ttheta_bins[i] + ttheta_bins[i + 1]) / 2.
231 + ttheta_cvk = "{:d}".format(int(np.floor(ttheta_cv)))
232 + if ttheta_cvk in det_pointer:
233 + ttheta_valid = False
234 + for tkey, item in ttheta_groups.items():
235 + if ttheta_cvk in item["QRange"].keys():
236 + ttheta_valid = True
237 + ttheta_group = tkey
238 + break
239 +
240 + # A valid two theta value means it belongs to one of the defined
241 + # groups in the `ttheta_group_file` file.
242 + if ttheta_valid:
243 + pa_lb = ttheta_groups[ttheta_group]["PAParams"]["LeftBound"]
244 + pa_rb = ttheta_groups[ttheta_group]["PAParams"]["RightBound"]
245 + mask = (q_vals >= pa_lb) & (q_vals <= pa_rb)
246 + q_vals_fit = q_vals[mask]
247 +
248 + b0 = (pa_lb + pa_rb) / 2.
249 + c0 = 1.
250 +
251 + amp_tmp = dict()
252 + center_tmp = dict()
253 +
254 + for key, item in det_pointer[ttheta_cvk]["AZMDict"].items():
255 + # Check whether the data exist for each azimuthal bin. If not,
256 + # we just discard the bin.
257 + if int(item["GroupID"]) in all_sofq.keys():
258 + sofq_fit = all_sofq[int(item["GroupID"])][mask]
259 + a0 = np.max(sofq_fit)
260 + p_init = [a0, b0, c0]
261 + try:
262 + popt, _ = curve_fit(
263 + gaussian,
264 + q_vals_fit,
265 + sofq_fit,
266 + p0=p_init
267 + )
268 + amp_tmp[key] = popt[0]
269 + center_tmp[key] = popt[1]
270 + except RuntimeError:
271 + # We want to handle the peak fitting failure. For
272 + # testing mode, we just set the amplitude to 1 and the
273 + # center to 0. For real data, we just discard the data.
274 + if testing:
275 + amp_tmp[key] = 1.
276 + center_tmp[key] = 0.
277 + else:
278 + det_pointer[ttheta_cvk][
279 + "AZMDict"
280 + ][key]["Use"] = False
281 + det_pointer[ttheta_cvk][
282 + "AZMDict"
283 + ][key]["Offset"] = 0.
284 + else:
285 + det_pointer[ttheta_cvk][
286 + "AZMDict"
287 + ][key]["Use"] = False
288 + det_pointer[ttheta_cvk][
289 + "AZMDict"
290 + ][key]["Offset"] = 0.
291 +
292 + if (len(center_tmp) > 0):
293 + # center_ref = next(iter(center_tmp.items()))[1] # first item
294 + center_ref = 5.06
295 +
296 + # Here we check the peak intensity. If the intensity is too
297 + # far from the median intensity, we discard the data. If not,
298 + # we calculate the offset for each azimuthal bin, relative to
299 + # the center of the first azimuthal bin, within the same two
300 + # theta bin.
301 + median_amp = np.median(list(amp_tmp.values()))
302 + used_keys = list()
303 + for key, item in amp_tmp.items():
304 + if abs((item - median_amp) / median_amp) < amp_thresh:
305 + used_keys.append(key)
306 + for key in det_pointer[ttheta_cvk]["AZMDict"].keys():
307 + if key in used_keys:
308 + det_pointer[ttheta_cvk]["AZMDict"][key]["Use"] = True
309 + offset = center_ref - center_tmp[key]
310 + det_pointer[ttheta_cvk][
311 + "AZMDict"
312 + ][key]["Offset"] = offset
313 + else:
314 + det_pointer[ttheta_cvk][
315 + "AZMDict"
316 + ][key]["Use"] = False
317 + det_pointer[ttheta_cvk][
318 + "AZMDict"
319 + ][key]["Offset"] = 0.
320 + else:
321 + # If none of the detector groups in the two theta bin is
322 + # valid (either file not existing, or fitting failed),
323 + # we mask out the whole two theta bin.
324 + for key in det_pointer[ttheta_cvk]["AZMDict"].keys():
325 + det_pointer[ttheta_cvk]["AZMDict"][key]["Use"] = False
326 + det_pointer[ttheta_cvk]["AZMDict"][key]["Offset"] = 0.
327 + else:
328 + # If the two theta bin is invalid, we just discard the bin.
329 + # By 'invalid' two theta bin, we mean that the two theta value
330 + # is not included in the `ttheta_group_file`.
331 + for key in det_pointer[ttheta_cvk]["AZMDict"].keys():
332 + det_pointer[ttheta_cvk]["AZMDict"][key]["Use"] = False
333 + det_pointer[ttheta_cvk]["AZMDict"][key]["Offset"] = 0.
334 +
335 + # =============================================================================
336 + # Loop over the two theta bins and optimize the spherical harmonic coefficients
337 + # =============================================================================
338 + print("[Info] Working on the isotropic correction for each two theta bin...")
339 + ttheta_bins = np.arange(0, 181, ttheta_binw)
340 + total_items = len(ttheta_bins)
341 + sofq_corrected_first = list()
342 + ttheta_valid = list()
343 + print("[Info] Processing 2Theta array (0, 180)...")
344 + for i in range(len(ttheta_bins) - 1):
345 + ttheta_cv = (ttheta_bins[i] + ttheta_bins[i + 1]) / 2.
346 + ttheta_cvk = "{:d}".format(int(np.floor(ttheta_cv)))
347 + print(f"[Info] Processing 2Theta: {ttheta_cvk}...")
348 + if ttheta_cvk in det_pointer:
349 + sofq_corrected = list()
350 + s_coefficients = list()
351 + valid_group_found = True
352 + for j in range(len(q_vals)):
353 + sofq_phi = list()
354 + sph_matrix = list()
355 + for _, item in det_pointer[ttheta_cvk]["AZMDict"].items():
356 + # According to the `Use` flag created for each azimuthal bin
357 + # in previous step, we are going to check whether there
358 + # exist valid azimuthal bins in a two theta bin.
359 + # NOTE: We may need to add in further check for whether the
360 + # number of valid azimuthal bins is greater than the number
361 + # of parameters to fit.
362 + if int(item["GroupID"]) in all_sofq.keys() and item["Use"]:
363 + azm_v = item["CentrAMZ"]
364 + sph_row = list()
365 + for l in range(sph_l_max + 1):
366 + for m in range(-l, l + 1):
367 + if l > 0 and m == 0:
368 + continue
369 + sph_row.append(
370 + sph_harm_real(
371 + m,
372 + l,
373 + ttheta_cv * np.pi / 180.,
374 + azm_v * np.pi / 180.
375 + )
376 + )
377 + sph_matrix.append(sph_row)
378 + sofq_phi.append(all_sofq[int(item["GroupID"])][j])
379 + else:
380 + pass
381 +
382 + sph_matrix = np.array(sph_matrix)
383 + sofq_phi = np.array(sofq_phi)
384 +
385 + if len(sph_matrix) == 0 or len(sofq_phi) == 0:
386 + valid_group_found = False
387 + break
388 +
389 + num_terms = (sph_l_max + 1)**2 - sph_l_max
390 + params = np.array(
391 + [1 for _ in range(num_terms)]
392 + )
393 +
394 + result = minimize(
395 + target_function,
396 + params,
397 + args=(sph_matrix, sofq_phi),
398 + method='SLSQP'
399 + )
400 +
401 + optimal_params = result.x
402 +
403 + # Calculate the S(Q) corrected, taking into account only the
404 + # isotropic term.
405 + y00 = sph_harm_real(0, 0, 0, 0)
406 + s00 = optimal_params[0]
407 + sofq_corrected.append(y00 * s00)
408 +
409 + s_coefficients.append(optimal_params)
410 +
411 + # Output the corrected S(Q) values and the coefficients for
412 + # each two theta bin.
413 + # For a certain two theta bin, if no valid azimuthal bin
414 + # found for any Q value, the `valid_group_found` would be
415 + # False and thus the corresponding two theta bin will be
416 + # marked as invalid, i.e., not included in the valid two
417 + # theta bin list `ttheta_valid`.
418 + if valid_group_found:
419 + # The valid two theta bin array and the isotropic
420 + # corrected S(Q) array are synchronized, i.e., they
421 + # share the same index, making it convenient for next
422 + # step implementation.
423 + sofq_corrected_first.append(sofq_corrected)
424 + ttheta_valid.append(ttheta_cv)
425 + with open(
426 + os.path.join(
427 + out_dir,
428 + f"{stem_name}2theta_{ttheta_cvk}.dat"
429 + ),
430 + "w"
431 + ) as f:
432 + f.write(f"{len(q_vals)}\n")
433 + f.write("\n")
434 + for j in range(len(q_vals)):
435 + f.write(
436 + f"{q_vals[j]:.4f} {sofq_corrected[j]:.4f}\n"
437 + )
438 +
439 + with open(
440 + os.path.join(
441 + out_dir,
442 + f"{stem_name}2theta_{ttheta_cvk}_coeffs.dat"
443 + ),
444 + "w"
445 + ) as f:
446 + for j in range(len(q_vals)):
447 + f.write(f"{q_vals[j]:.4f} ")
448 + for k in range(len(s_coefficients[j])):
449 + f.write(
450 + f"{s_coefficients[j][k]:.4f} "
451 + )
452 + f.write("\n")
453 +
454 + ttheta_valid = np.array(ttheta_valid)
455 + sofq_corrected_first = np.array(sofq_corrected_first)
456 +
457 + # =============================================================================
458 + # Pre-process the isotropic corrected S(Q) data corresponding to all the two
459 + # theta bins. Here is the thing to do:
460 + # 1. For each two theta bin, grab the data in a certain Q range according to
461 + # read-in parameters from `ttheta_group_file`.
462 + # 2. For each two theta bin, fill the data in the low Q range with the data
463 + # in the same low Q range that comes from the first two theta bin in the
464 + # same group.
465 + # =============================================================================
466 + print(
467 + "[Info] Pre-processing (step-1) the isotropic corrected data for the "
468 + "second stage correction..."
469 + )
470 + ttheta_k_valid = [
471 + "{:d}".format(int(np.floor(ttheta))) for ttheta in ttheta_valid
472 + ]
473 + for i, ttheta in enumerate(ttheta_valid):
474 + ttheta_k = ttheta_k_valid[i]
475 + for key, item in ttheta_groups.items():
476 + if ttheta_k in item["QRange"].keys():
477 + group = key
478 + break
479 +
480 + qrange_tmp = ttheta_groups[group]["QRange"]
481 + left_b = qrange_tmp[ttheta_k]["LeftBound"]
482 + leftmost_b = next(iter(qrange_tmp.items()))[1]["LeftBound"]
483 + right_b_fill = left_b
484 +
485 + # Fill the low Q range with the data from the first two theta bin
486 + # in the same group.
487 + ref_key_ig = next(iter(qrange_tmp))
488 + ref_key_id = ttheta_k_valid.index(ref_key_ig)
489 + mask_fill = (q_vals >= leftmost_b) & (q_vals < right_b_fill)
490 + sofq_corrected_first[i][mask_fill] = sofq_corrected_first[
491 + ref_key_id
492 + ][mask_fill]
493 +
494 + # =============================================================================
495 + # Process each of the two theta groups separately. Here is the thing to do:
496 + # 1. For each group, align all the patterns according to the peak fitting.
497 + # 2. Within each group, we,
498 + # - for each two theta, figure out and remove the background, using an
499 + # exponential function
500 + # - integrate intensities in those pre-defined Q chunks
501 + # - for each Q chunk, fit against the spherical harmonics across two
502 + # theta angles, just like what we did with each single Q point when
503 + # correcting the phi part.
504 + # - extract the isotropic term for each Q chunk, calculate the ratio
505 + # of the original integrated intensity over the isotropic intensity
506 + # for each two theta value. Then rescale each Q chunk for each two
507 + # theta.
508 + # - add back the exponential background
509 + # =============================================================================
510 + print(
511 + "[Info] Pre-processing (step-2) the isotropic corrected data for the "
512 + "second stage correction..."
513 + )
514 + sofq_bkg_removed = dict()
515 + bkg_off = np.inf
516 +
517 + for tg_key, item in ttheta_groups.items():
518 + print(f"[Info] Processing 2Theta group {tg_key}...")
519 + pa_lb = item["PAParams"]["LeftBound"]
520 + pa_rb = item["PAParams"]["RightBound"]
521 + center_tmp = dict()
522 + for tt_key in item["QRange"].keys():
523 + if tt_key not in ttheta_k_valid:
524 + continue
525 + tt_index = ttheta_k_valid.index(tt_key)
526 + sofq_tmp = sofq_corrected_first[tt_index]
527 + mask = (q_vals >= pa_lb) & (q_vals <= pa_rb)
528 + q_vals_fit = q_vals[mask]
529 + sofq_fit = sofq_tmp[mask]
530 +
531 + p_init = [np.max(sofq_fit), (pa_lb + pa_rb) / 2., 1.]
532 +
533 + if testing:
534 + center_tmp[tt_key] = 0.
535 + else:
536 + try:
537 + popt, _ = curve_fit(
538 + gaussian,
539 + q_vals_fit,
540 + sofq_fit,
541 + p0=p_init
542 + )
543 + center_tmp[tt_key] = popt[1]
544 + except RuntimeError:
545 + center_tmp[tt_key] = 0.
546 +
547 + # center_ref = next(iter(center_tmp.items()))[1]
548 + center_ref = 5.06
549 + for tt_key, c_item in center_tmp.items():
550 + # According to the peak fitting failure handling just above, the
551 + # corresponding peak center would be set to 0.
552 + if c_item == 0.:
553 + offset = 0.
554 + else:
555 + offset = center_ref - c_item
556 +
557 + tt_index = ttheta_k_valid.index(tt_key)
558 + q_vals_tmp = q_vals + offset
559 + sofq_tmp = sofq_corrected_first[tt_index].copy()
560 + q_bin_s = q_vals[1] - q_vals[0]
561 + if offset < 0:
562 + while q_vals_tmp[-1] < q_vals[-1]:
563 + q_vals_tmp = np.append(
564 + q_vals_tmp,
565 + q_vals_tmp[-1] + q_bin_s
566 + )
567 + sofq_tmp = np.append(
568 + sofq_tmp,
569 + sofq_corrected_first[tt_index][-1]
570 + )
571 + else:
572 + while q_vals_tmp[0] > q_vals[0]:
573 + q_vals_tmp = np.insert(
574 + q_vals_tmp,
575 + 0,
576 + q_vals_tmp[0] - q_bin_s
577 + )
578 + sofq_tmp = np.insert(
579 + sofq_tmp,
580 + 0,
581 + sofq_corrected_first[tt_index][0]
582 + )
583 + _, sofq_rebin = rebin(
584 + q_vals_tmp,
585 + sofq_tmp,
586 + q_vals[0],
587 + q_bin_s,
588 + q_vals[-1]
589 + )
590 + sofq_corrected_first[tt_index] = sofq_rebin
591 +
592 + for tt_key in item["QRange"].keys():
593 + if tt_key not in ttheta_k_valid:
594 + continue
595 + tt_index = ttheta_k_valid.index(tt_key)
596 + ttheta = ttheta_valid[tt_index]
597 + sofq_ttheta = sofq_corrected_first[tt_index]
598 +
599 + bkg_lb = item["QRange"][tt_key]["LeftBound"]
600 + bkg_rb = item["QRange"][tt_key]["RightBound"]
601 +
602 + # Here we output the aligned patterns.
603 + # N.B. The alignment only happens internally in each of the defined
604 + # groups. Therefore, a final stage alignment may still be needed to
605 + # align the patterns after all steps.
606 + with open(
607 + os.path.join(
608 + out_dir,
609 + f"{stem_name}ttheta_{tt_key}_corrected_aligned.dat"
610 + ),
611 + "w"
612 + ) as f:
613 + f.write(f"{len(q_vals)}\n")
614 + f.write("\n")
615 + for i, q_val in enumerate(q_vals):
616 + f.write("{0:10.3F}{1:20.5F}\n".format(q_val, sofq_ttheta[i]))
617 +
618 + # We try to use the `argrelextrema` method in `scipy` to figure out
619 + # local valley points on the pattern and use those valley points for
620 + # fitting an exponential background.
621 + mask_tmp = q_vals >= bkg_lb
622 + sofq_ttheta_m = sofq_ttheta[mask_tmp]
623 + q_vals_m = q_vals[mask_tmp]
624 +
625 + # The `order` parameter here controls how 'local' the valley will be.
626 + # If set to a very small value, it will reproduce the overall pattern
627 + # itself. But setting the value too high may induce too few points.
628 + q_min = argrelextrema(sofq_ttheta_m, np.less, order=20)
629 + sofq_min = np.array([sofq_ttheta_m[item_i] for item_i in q_min[0]])
630 + q_min = np.array([q_vals_m[item_i] for item_i in q_min[0]])
631 +
632 + # For the background fitting purpose, we set all S(Q) values above
633 + # the specified threshold to the value just below (inclusive) the
634 + # boundary. This will get rid of the impact of those strong noise
635 + # in the high Q region.
636 + indices = np.where(q_min <= bkg_rb)[0]
637 + max_tmp = np.max(q_min[indices])
638 + up_index = np.where(q_min == max_tmp)[0][0]
639 + for qm_i, qm_val in enumerate(q_min):
640 + if qm_val > bkg_rb:
641 + sofq_min[qm_i] = sofq_min[up_index]
642 +
643 + # Specifying the bounding parameter for `curve_fit` helps a lot with
644 + # fitting convergence.
645 + p_init = [0, 0.1, 0.1]
646 + bounds_l = [-np.inf, 0., 0.]
647 + bounds_u = [np.inf, 1., 1.]
648 + popt, pconv = curve_fit(
649 + bkg_expo,
650 + q_min,
651 + sofq_min,
652 + p0=p_init,
653 + bounds=(bounds_l, bounds_u)
654 + )
655 +
656 + is_infinite = False
657 + for row in pconv:
658 + for element in row:
659 + if np.isinf(element):
660 + is_infinite = True
661 + break
662 +
663 + if is_infinite:
664 + print("[Warning] Infinite covariance in the curve fitting.")
665 +
666 + with open(
667 + os.path.join(
668 + out_dir,
669 + f"{stem_name}ttheta_{tt_key}_corrected_bkg_pts.dat"
670 + ),
671 + "w"
672 + ) as f:
673 + f.write(f"{len(q_min)}\n")
674 + f.write("\n")
675 + for i, q_val in enumerate(q_min):
676 + f.write(
677 + "{0:10.3F}{1:20.5F}\n".format(
678 + q_val, sofq_min[i]
679 + )
680 + )
681 +
682 + sofq_ttheta_bkg = bkg_expo(q_vals, popt[0], popt[1], popt[2])
683 + sofq_ttheta_mb = sofq_ttheta - sofq_ttheta_bkg
684 + sofq_bkg_removed[tt_key] = sofq_ttheta_mb
685 +
686 + # We want to keep a record of the background that most closely resemble
687 + # 1 at high Q.
688 + bkg_hq_val = bkg_expo(q_vals[-1], popt[0], popt[1], popt[2])
689 + bkg_off_tmp = abs(1. - bkg_hq_val)
690 + if bkg_off_tmp < bkg_off:
691 + bkg_off = bkg_off_tmp
692 + bkg_popt = popt
693 +
694 + with open(
695 + os.path.join(
696 + out_dir,
697 + f"{stem_name}ttheta_{tt_key}_corrected_bkg_rm_fittedbkg.dat"
698 + ),
699 + "w"
700 + ) as f:
701 + f.write(f"{len(q_vals)}\n")
702 + f.write("\n")
703 + for i, q_val in enumerate(q_vals):
704 + f.write(
705 + "{0:10.3F}{1:20.5F}{2:20.5F}\n".format(
706 + q_val, sofq_ttheta_mb[i], sofq_ttheta_bkg[i]
707 + )
708 + )
709 +
710 + # Integrate intensities inside Q chunks.
711 + ttheta_groups[tg_key]["QRange"][tt_key]["QChunkInt"] = list()
712 + for qv_i, q_val in enumerate(q_vals):
713 + v_t = sofq_ttheta_mb[qv_i]
714 + for qc_i, qc_val in enumerate(item["QChunkDef"]):
715 + if qc_i == 0:
716 + if q_val <= qc_val:
717 + if len(item["QRange"][tt_key]["QChunkInt"]) == 0:
718 + ttheta_groups[
719 + tg_key
720 + ]["QRange"][tt_key]["QChunkInt"].append(v_t)
721 + else:
722 + ttheta_groups[
723 + tg_key
724 + ]["QRange"][tt_key]["QChunkInt"][0] += v_t
725 + break
726 + else:
727 + if item["QChunkDef"][qc_i - 1] < q_val <= qc_val:
728 + if len(item["QRange"][tt_key]["QChunkInt"]) == qc_i:
729 + in_range1 = bkg_lb <= item["QChunkDef"][qc_i - 1]
730 + in_range2 = bkg_rb >= qc_val
731 + if in_range1 and in_range2:
732 + ttheta_groups[
733 + tg_key
734 + ]["QRange"][tt_key]["QChunkInt"].append(v_t)
735 + else:
736 + ttheta_groups[
737 + tg_key
738 + ]["QRange"][tt_key]["QChunkInt"].append(np.inf)
739 + else:
740 + ttheta_groups[
741 + tg_key
742 + ]["QRange"][tt_key]["QChunkInt"][-1] += v_t
743 + break
744 +
745 + # Treat each Q chunk as a single 'point' and perform the isotropic correction
746 + # just as before.
747 + sph00 = sph_harm_real(0, 0, 0, 0)
748 + for tg_key, item in ttheta_groups.items():
749 + params_all = list()
750 + qchunks_len = len(item["QChunkDef"])
751 + for i in range(qchunks_len):
752 + sph_matrix = list()
753 + sofq_ttheta = list()
754 + lin_limit = item["LinLimit"][i]
755 +
756 + num_eqs = 0
757 + for tt_key in item["QRange"].keys():
758 + if tt_key in ttheta_k_valid:
759 + num_eqs += 1
760 +
761 + tt_sph_l_max = 0
762 + while (tt_sph_l_max + 1 + 1)**2 < num_eqs:
763 + tt_sph_l_max += 1
764 +
765 + for tt_key in item["QRange"].keys():
766 + if "Factor" not in item["QRange"][tt_key].keys():
767 + ttheta_groups[tg_key]["QRange"][tt_key]["Factor"] = list()
768 +
769 + if tt_key not in ttheta_k_valid:
770 + continue
771 +
772 + tt_index = ttheta_k_valid.index(tt_key)
773 + ttheta_tmp = ttheta_valid[tt_index]
774 + v_t = item["QRange"][tt_key]["QChunkInt"][i]
775 + if v_t == np.inf:
776 + continue
777 + sofq_ttheta.append(v_t)
778 + sph_row = list()
779 + for l in range(tt_sph_l_max + 1):
780 + for m in range(-l, l + 1):
781 + sph_row.append(
782 + sph_harm_real(
783 + m,
784 + l,
785 + ttheta_tmp * np.pi / 180.,
786 + 0.
787 + )
788 + )
789 +
790 + if "Infinity" not in lin_limit:
791 + sph_row.append(ttheta_tmp)
792 +
793 + sph_matrix.append(sph_row)
794 +
795 + sph_matrix = np.array(sph_matrix)
796 + sofq_ttheta = np.array(sofq_ttheta)
797 +
798 + if len(sph_matrix) > 0 or len(sofq_ttheta) > 0:
799 + if "Infinity" in lin_limit:
800 + num_terms = (tt_sph_l_max + 1)**2
801 + else:
802 + num_terms = (tt_sph_l_max + 1)**2 + 1
803 +
804 + params = np.array(
805 + [1 for _ in range(num_terms)]
806 + )
807 +
808 + min_b = min(sofq_ttheta) / sph00
809 + max_b = max(sofq_ttheta) / sph00
810 + bounds = [(min_b, max_b)]
811 + if "Infinity" not in lin_limit:
812 + num_b_add = num_terms - 2
813 + else:
814 + num_b_add = num_terms - 1
815 +
816 + for _ in range(num_b_add):
817 + bounds.append(
818 + (-float('inf'), float('inf'))
819 + )
820 +
821 + if "Infinity" not in lin_limit:
822 + bounds.append((lin_limit[0], lin_limit[1]))
823 +
824 + result = minimize(
825 + target_function,
826 + params,
827 + args=(sph_matrix, sofq_ttheta),
828 + method='SLSQP',
829 + bounds=bounds
830 + )
831 +
832 + optimal_params = result.x
833 +
834 + y00 = sph_harm_real(0, 0, 0, 0)
835 + s00 = optimal_params[0]
836 + if "Infinity" not in lin_limit:
837 + ctt = optimal_params[-1]
838 + else:
839 + y00 = 0.
840 + s00 = 0.
841 +
842 + # Calculate the ratio of the original integrated intensity over
843 + # the isotropic intensity.
844 + for tt_key in item["QRange"].keys():
845 + if tt_key not in ttheta_k_valid:
846 + continue
847 +
848 + if i == 0:
849 + factor_v = 1.
850 + else:
851 + tt_index = ttheta_k_valid.index(tt_key)
852 + ttheta_tmp = ttheta_valid[tt_index]
853 + if "Infinity" in lin_limit:
854 + top_val = y00 * s00
855 + else:
856 + top_val = y00 * s00 + ctt * ttheta_tmp
857 +
858 + bottom_val = item["QRange"][tt_key]["QChunkInt"][i]
859 + factor_v = top_val / bottom_val
860 +
861 + # if tg_key == "Group-1":
862 + # if i == 2:
863 + # print("Debugging -> ")
864 + # print(
865 + # f"fit val: {top_val}, init val: {bottom_val}",
866 + # f"y00: {y00}, s00: {s00}",
867 + # f"tt: {ttheta_tmp}"
868 + # )
869 + # print(
870 + # f"fit val: {top_val}, init val: {bottom_val}",
871 + # f"y00: {y00}, s00: {s00}, ctt: {ctt}",
872 + # f"tt: {ttheta_tmp}"
873 + # )
874 +
875 + ttheta_groups[tg_key]["QRange"][tt_key]["Factor"].append(
876 + factor_v
877 + )
878 +
879 + # Rescale the data by Q chunks according to the ratio obtained above.
880 + for tt_key in item["QRange"].keys():
881 + if tt_key not in ttheta_k_valid:
882 + continue
883 +
884 + for i, q_val in enumerate(q_vals):
885 + for ii, qc_val in enumerate(item["QChunkDef"]):
886 + if ii > 0:
887 + # if q_vals[ii - 1] < q_val <= qc_val and q_val > 4.:
888 + if q_vals[ii - 1] < q_val <= qc_val:
889 + sofq_bkg_removed[tt_key][i] *= item[
890 + "QRange"
891 + ][tt_key]["Factor"][ii]
892 +
893 + break
894 +
895 + # Add back in the exponential background.
896 + bkg_add = bkg_expo(q_vals, bkg_popt[0], bkg_popt[1], bkg_popt[2])
897 + sofq_final = sofq_bkg_removed[tt_key] + bkg_add
898 +
899 + with open(
900 + os.path.join(
901 + out_dir,
902 + f"{stem_name}ttheta_{tt_key}_corrected_wo_bkg.dat"
903 + ),
904 + "w"
905 + ) as f:
906 + f.write(f"{len(q_vals)}\n")
907 + f.write("\n")
908 + for i, q_val in enumerate(q_vals):
909 + f.write(
910 + "{0:10.3F}{1:20.5F}\n".format(
911 + q_val, sofq_bkg_removed[tt_key][i]
912 + )
913 + )
914 +
915 + with open(
916 + os.path.join(
917 + out_dir,
918 + f"{stem_name}ttheta_{tt_key}_corrected_final.dat"
919 + ),
920 + "w"
921 + ) as f:
922 + f.write(f"{len(q_vals)}\n")
923 + f.write("\n")
924 + for i, q_val in enumerate(q_vals):
925 + f.write(
926 + "{0:10.3F}{1:20.5F}\n".format(
927 + q_val, sofq_final[i]
928 + )
929 + )
930 +
931 + with open(os.path.join(out_dir, f"{stem_name}_ttheta_1na_2a.json"), "w") as f:
932 + json.dump(ttheta_groups, f, indent=4)
933 +
934 + print("\n=========================================================")
935 + print("[Info] Done with the two-stage isotropic S(Q) correction.")
936 + print("[Info] Proceed to merge the corrected S(Q) data by bank.")
937 + print("=========================================================\n")
938 +
939 + with open("output_group.json", "r") as f:
940 + group_dict = json.load(f)
941 +
942 + bank_dict = dict()
943 + for key in ttheta_k_valid:
944 + for keykey in ttheta_groups.keys():
945 + if key in ttheta_groups[keykey]["QRange"].keys():
946 + key_group = keykey
947 + break
948 +
949 + file_t = os.path.join(
950 + out_dir,
951 + f"{stem_name}ttheta_{key}_corrected_final.dat"
952 + )
953 + if os.path.exists(file_t):
954 + y_tmp = np.loadtxt(file_t, skiprows=2, dtype=float)[:, 1]
955 + for gkey, item in group_dict.items():
956 + if item[0] <= float(key) < item[1]:
957 + bank = gkey
958 + break
959 +
960 + if bank not in bank_dict:
961 + bank_dict[bank] = {
962 + "y": y_tmp,
963 + "num_bands": 1,
964 + "q_counts": [0 for _ in range(len(q_vals))]
965 + }
966 + else:
967 + bank_dict[bank]["y"] += y_tmp
968 + bank_dict[bank]["num_bands"] += 1
969 +
970 + q_chunk_tmp = ttheta_groups[key_group]["QChunkDef"]
971 + for i, q_val in enumerate(q_vals):
972 + for ii, qc_val in enumerate(q_chunk_tmp):
973 + if ii == 0:
974 + if q_val <= qc_val:
975 + q_chunk_id = ii
976 + break
977 + else:
978 + if q_chunk_tmp[ii - 1] < q_val <= qc_val:
979 + q_chunk_id = ii
980 + break
981 +
982 + if q_chunk_id == 0:
983 + factor = 1.
984 + else:
985 + factor = ttheta_groups[
986 + key_group
987 + ]["QRange"][key]["Factor"][q_chunk_id]
988 +
989 + if factor > 0.:
990 + bank_dict[bank]["q_counts"][i] += 1
991 +
992 + for key, item in bank_dict.items():
993 + for i in range(len(q_vals)):
994 + if item["q_counts"][i] > 0:
995 + item["y"][i] /= item["q_counts"][i]
996 + else:
997 + item["y"][i] /= item["num_bands"]
998 +
999 + with open(
1000 + os.path.join(out_dir, f"{stem_name}{key}_merged_1na_2a_ave.dat"),
1001 + "w"
1002 + ) as f:
1003 + f.write(f"{len(q_vals)}\n")
1004 + f.write("\n")
1005 + for i, q_val in enumerate(q_vals):
1006 + f.write("{0:10.3F}{1:20.5F}\n".format(q_val, item["y"][i]))
1007 +
1008 + print("=========================================================")
1009 + print("[Info] Done with merging the corrected S(Q) by banks.")
1010 + print("=========================================================\n")
Newer Older

Powered by Opengist ⋅ Load: 61ms⋅

English