texture_proc_real_1step_not_aligned_2step_aligned.py

texture_proc_real_1step_not_aligned_2step_aligned.py · 36 KiB · Python Raw

import json
import numpy as np
import os
import scipy
from scipy.optimize import minimize
from scipy.optimize import curve_fit
from scipy.signal import argrelextrema
from pystog import Pre_Proc
import matplotlib.pyplot as plt
import random

rebin = Pre_Proc.rebin

# ================
# Input parameters
# vvvvvvvvvvvvvvvv

# ----------
# Some flags
# ----------
testing = False     # flag for testing

# ------------------------------------------------------------------
# Directory and config files definition
#
# run_dir           => directory for storing the reduced data files
#                      according to the the two theta and azimuthal
#                      angle bins.
# out_dir           => directory for storing the outputs.
# det_map_file      => detector group information for all the two
#                      theta and azimuthal bins.
# ttheta_group_file => config file defining groups for second stage
#                      processing.
# ------------------------------------------------------------------
run_dir = os.path.join(
    ".",
    "texture_proc"
)
out_dir = os.path.join(
    run_dir,
    "texture_proc_output"
)
det_map_file = os.path.join(
    "/SNS/NOM/shared/scripts/texture",
    "nom_texture_pointer.json"
)
ttheta_group_file = os.path.join(
    ".",
    "ttheta_group_params_new_new_new.json"
)

# ---------------------
# Processing parameters
# ---------------------
stem_name = "NOM_Si_640e_"
ttheta_binw = 5.    # in degrees
azm_binw = 10.      # in degrees
header_num = 2
sph_l_max = 1       # for first stage
tt_sph_l_max = 2    # for second stage
amp_thresh = .9     # Threshold for peak intensity -- in a certain two theta
                    # bins, all patterns with the fitted peak intensity
                    # different from the median intensity by more than this
                    # percentage threshold will be discarded.

# ^^^^^^^^^^^^^^^^
# Input parameters
# ================

if testing:
    print("\n!!!!!!!!!!!!!!!!!!!!!!!!!!!!!")
    print("[Warning] Testing mode is on.")
    print("!!!!!!!!!!!!!!!!!!!!!!!!!!!!!\n")

def gaussian(x, a, b, c):
    """Gaussian function.
    Parameters
    ----------
    x : ndarray Independent variable.
    a : float Amplitude.
    b : float Center.
    c : flaot Width.

Returns
    -------
    float Gaussian function value.
    """
    return a * np.exp(-(x - b)**2 / (2 * c**2))

def bkg_expo(x, a, b, c):
    """Background exponential function.
    Parameters
    ----------
    x : ndarray Independent variable.
    a : float offset.
    b : float coefficient.
    c : flaot exponent.

Returns
    -------
    float Exponential background value.
    """
    return a - b * np.exp(-c * x**2.)

def sph_harm_real(m, l, theta, phi):
    """Real spherical harmonics. Refer to the Wikipedia,
    https://en.wikipedia.org/wiki/Table_of_spherical_harmonics
    We could follow the standard definition of the spherical harmonics to
    define the function here, but for the purpose of this script, we only
    need the first few terms so we just implemented the function directly.

Parameters
    ----------
    m : order of spherical harmonics
    l : degree of spherical harmonics
    theta : polar angle in radian
    phi : azimuthal angle in radian

Returns
    -------
    float Spherical harmonics real value
    """
    if l == 0 and m == 0:
        return np.sqrt(1. / (4. * np.pi))
    if l == 1 and m == -1:
        return np.sqrt(3. / (4. * np.pi)) * np.sin(phi) * np.sin(theta)
    if l == 1 and m == 0:
        return np.sqrt(3. / (4. * np.pi)) * np.cos(theta)
    if l == 1 and m == 1:
        return np.sqrt(3. / (4. * np.pi)) * np.cos(phi) * np.sin(theta)
    if l == 2 and m == -2:
        part = np.sin(phi) * np.cos(phi) * (np.sin(theta))**2.
        return np.sqrt(15. / (4. * np.pi)) * part
    if l == 2 and m == -1:
        part = np.sin(phi) * np.sin(theta) * np.cos(theta)
        return np.sqrt(15. / (4. * np.pi)) * part
    if l == 2 and m == 0:
        part = 3. * (np.cos(theta))**2. - 1.
        return np.sqrt(5. / (16. * np.pi)) * part
    if l == 2 and m == 1:
        part = np.cos(phi) * np.sin(theta) * np.cos(theta)
        return np.sqrt(15. / (4. * np.pi)) * part
    if l == 2 and m == 2:
        part = ((np.cos(phi))**2. - (np.sin(phi))**2.) * (np.sin(theta))**2.
        return np.sqrt(15. / (16. * np.pi)) * part

def target_function(params, sph_mat, txt_int):
    """Target function for the optimization of the spherical harmonic
    coefficients.

Parameters
    ----------
    params : ndarray Spherical harmonic coefficients.
    sph_mat : ndarray Spherical harmonic matrix.
    txt_int : ndarray Textured S(Q) intensity.

Returns
    -------
    float : Norm of the residuals.
    """
    result = sph_mat @ params
    residuals = result - txt_int

return np.linalg.norm(residuals)

if not os.path.exists(out_dir):
    os.makedirs(out_dir)

# ================
# Read in Q values
# ================
print("[Info] Read in Q values...")
bank = 1
while True:
    tr_file = os.path.join(
        run_dir,
        stem_name + f"bank{bank}.dat"
    )
    if os.path.exists(tr_file):
        data = np.loadtxt(tr_file, skiprows=header_num, dtype=float)
        q_vals = data[:, 0]
        break
    else:
        bank += 1

# ===================
# Read in S(Q) values
# ===================
print("[Info] Read in all S(Q)'s...")
num_groups = int(180. / ttheta_binw) * int(360. / azm_binw)
all_sofq = dict()
for i in range(num_groups):
    group_tmp = i + 1
    tr_file = os.path.join(
        run_dir,
        stem_name + "bank%d.dat" % group_tmp
    )
    if os.path.exists(tr_file):
        data = np.loadtxt(tr_file, skiprows=header_num, dtype=float)
        sofq_tmp = data[:, 1]
        all_sofq[group_tmp] = sofq_tmp

# ===========================================================================
# Read in the detector map, telling the center two theta and azimuthal angle,
# and the ID for each group of detectors.
# ===========================================================================
det_pointer = json.load(open(det_map_file, "r"))

# =============================================================================
# Determine the valid azimuthal bins for each two theta bin. Here follows are
# the criteria:
#   1. The 2theta bin is in use
#   2. The peak fitting succeeded
#   3. The fitted peak intensity is close to the median intensity
# =============================================================================
with open(ttheta_group_file, "r") as f:
    ttheta_groups = json.load(f)

print(
    "[Info] Pre-processing the data for all "
    "the two theta and azimuthal bins..."
)
ttheta_bins = np.arange(0, 181, ttheta_binw)
for i in range(len(ttheta_bins) - 1):
    ttheta_cv = (ttheta_bins[i] + ttheta_bins[i + 1]) / 2.
    ttheta_cvk = "{:d}".format(int(np.floor(ttheta_cv)))
    if ttheta_cvk in det_pointer:
        ttheta_valid = False
        for tkey, item in ttheta_groups.items():
            if ttheta_cvk in item["QRange"].keys():
                ttheta_valid = True
                ttheta_group = tkey
                break

# A valid two theta value means it belongs to one of the defined
        # groups in the `ttheta_group_file` file.
        if ttheta_valid:
            pa_lb = ttheta_groups[ttheta_group]["PAParams"]["LeftBound"]
            pa_rb = ttheta_groups[ttheta_group]["PAParams"]["RightBound"]
            mask = (q_vals >= pa_lb) & (q_vals <= pa_rb)
            q_vals_fit = q_vals[mask]

b0 = (pa_lb + pa_rb) / 2.
            c0 = 1.

amp_tmp = dict()
            center_tmp = dict()

for key, item in det_pointer[ttheta_cvk]["AZMDict"].items():
                # Check whether the data exist for each azimuthal bin. If not,
                # we just discard the bin.
                if int(item["GroupID"]) in all_sofq.keys():
                    sofq_fit = all_sofq[int(item["GroupID"])][mask]
                    a0 = np.max(sofq_fit)
                    p_init = [a0, b0, c0]
                    try:
                        popt, _ = curve_fit(
                            gaussian,
                            q_vals_fit,
                            sofq_fit,
                            p0=p_init
                        )
                        amp_tmp[key] = popt[0]
                        center_tmp[key] = popt[1]
                    except RuntimeError:
                        # We want to handle the peak fitting failure. For
                        # testing mode, we just set the amplitude to 1 and the
                        # center to 0. For real data, we just discard the data.
                        if testing:
                            amp_tmp[key] = 1.
                            center_tmp[key] = 0.
                        else:
                            det_pointer[ttheta_cvk][
                                "AZMDict"
                            ][key]["Use"] = False
                            det_pointer[ttheta_cvk][
                                "AZMDict"
                            ][key]["Offset"] = 0.
                else:
                    det_pointer[ttheta_cvk][
                        "AZMDict"
                    ][key]["Use"] = False
                    det_pointer[ttheta_cvk][
                        "AZMDict"
                    ][key]["Offset"] = 0.

if (len(center_tmp) > 0):
                # center_ref = next(iter(center_tmp.items()))[1]  # first item
                center_ref = 5.06

# Here we check the peak intensity. If the intensity is too
                # far from the median intensity, we discard the data. If not,
                # we calculate the offset for each azimuthal bin, relative to
                # the center of the first azimuthal bin, within the same two
                # theta bin.
                median_amp = np.median(list(amp_tmp.values()))
                used_keys = list()
                for key, item in amp_tmp.items():
                    if abs((item - median_amp) / median_amp) < amp_thresh:
                        used_keys.append(key)
                for key in det_pointer[ttheta_cvk]["AZMDict"].keys():
                    if key in used_keys:
                        det_pointer[ttheta_cvk]["AZMDict"][key]["Use"] = True
                        offset = center_ref - center_tmp[key]
                        det_pointer[ttheta_cvk][
                            "AZMDict"
                        ][key]["Offset"] = offset
                    else:
                        det_pointer[ttheta_cvk][
                            "AZMDict"
                        ][key]["Use"] = False
                        det_pointer[ttheta_cvk][
                            "AZMDict"
                        ][key]["Offset"] = 0.
            else:
                # If none of the detector groups in the two theta bin is
                # valid (either file not existing, or fitting failed),
                # we mask out the whole two theta bin.
                for key in det_pointer[ttheta_cvk]["AZMDict"].keys():
                    det_pointer[ttheta_cvk]["AZMDict"][key]["Use"] = False
                    det_pointer[ttheta_cvk]["AZMDict"][key]["Offset"] = 0.
        else:
            # If the two theta bin is invalid, we just discard the bin.
            # By 'invalid' two theta bin, we mean that the two theta value
            # is not included in the `ttheta_group_file`.
            for key in det_pointer[ttheta_cvk]["AZMDict"].keys():
                det_pointer[ttheta_cvk]["AZMDict"][key]["Use"] = False
                det_pointer[ttheta_cvk]["AZMDict"][key]["Offset"] = 0.

# =============================================================================
# Loop over the two theta bins and optimize the spherical harmonic coefficients
# =============================================================================
print("[Info] Working on the isotropic correction for each two theta bin...")
ttheta_bins = np.arange(0, 181, ttheta_binw)
total_items = len(ttheta_bins)
sofq_corrected_first = list()
ttheta_valid = list()
print("[Info] Processing 2Theta array (0, 180)...")
for i in range(len(ttheta_bins) - 1):
    ttheta_cv = (ttheta_bins[i] + ttheta_bins[i + 1]) / 2.
    ttheta_cvk = "{:d}".format(int(np.floor(ttheta_cv)))
    print(f"[Info] Processing 2Theta: {ttheta_cvk}...")
    if ttheta_cvk in det_pointer:
        sofq_corrected = list()
        s_coefficients = list()
        valid_group_found = True
        for j in range(len(q_vals)):
            sofq_phi = list()
            sph_matrix = list()
            for _, item in det_pointer[ttheta_cvk]["AZMDict"].items():
                # According to the `Use` flag created for each azimuthal bin
                # in previous step, we are going to check whether there
                # exist valid azimuthal bins in a two theta bin.
                # NOTE: We may need to add in further check for whether the
                # number of valid azimuthal bins is greater than the number
                # of parameters to fit.
                if int(item["GroupID"]) in all_sofq.keys() and item["Use"]:
                    azm_v = item["CentrAMZ"]
                    sph_row = list()
                    for l in range(sph_l_max + 1):
                        for m in range(-l, l + 1):
                            if l > 0 and m == 0:
                                continue
                            sph_row.append(
                                sph_harm_real(
                                    m,
                                    l,
                                    ttheta_cv * np.pi / 180.,
                                    azm_v * np.pi / 180.
                                )
                            )
                    sph_matrix.append(sph_row)
                    sofq_phi.append(all_sofq[int(item["GroupID"])][j])
                else:
                    pass

sph_matrix = np.array(sph_matrix)
            sofq_phi = np.array(sofq_phi)

if len(sph_matrix) == 0 or len(sofq_phi) == 0:
                valid_group_found = False
                break

num_terms = (sph_l_max + 1)**2 - sph_l_max
            params = np.array(
                [1 for _ in range(num_terms)]
            )

result = minimize(
                target_function,
                params,
                args=(sph_matrix, sofq_phi),
                method='SLSQP'
            )

optimal_params = result.x

# Calculate the S(Q) corrected, taking into account only the
            # isotropic term.
            y00 = sph_harm_real(0, 0, 0, 0)
            s00 = optimal_params[0]
            sofq_corrected.append(y00 * s00)

s_coefficients.append(optimal_params)

# Output the corrected S(Q) values and the coefficients for
        # each two theta bin.
        # For a certain two theta bin, if no valid azimuthal bin
        # found for any Q value, the `valid_group_found` would be
        # False and thus the corresponding two theta bin will be
        # marked as invalid, i.e., not included in the valid two
        # theta bin list `ttheta_valid`.
        if valid_group_found:
            # The valid two theta bin array and the isotropic
            # corrected S(Q) array are synchronized, i.e., they
            # share the same index, making it convenient for next
            # step implementation.
            sofq_corrected_first.append(sofq_corrected)
            ttheta_valid.append(ttheta_cv)
            with open(
                os.path.join(
                    out_dir,
                    f"{stem_name}2theta_{ttheta_cvk}.dat"
                ),
                "w"
            ) as f:
                f.write(f"{len(q_vals)}\n")
                f.write("\n")
                for j in range(len(q_vals)):
                    f.write(
                        f"{q_vals[j]:.4f} {sofq_corrected[j]:.4f}\n"
                    )

with open(
                os.path.join(
                    out_dir,
                    f"{stem_name}2theta_{ttheta_cvk}_coeffs.dat"
                ),
                "w"
            ) as f:
                for j in range(len(q_vals)):
                    f.write(f"{q_vals[j]:.4f} ")
                    for k in range(len(s_coefficients[j])):
                        f.write(
                            f"{s_coefficients[j][k]:.4f} "
                        )
                    f.write("\n")

ttheta_valid = np.array(ttheta_valid)
sofq_corrected_first = np.array(sofq_corrected_first)

# =============================================================================
# Pre-process the isotropic corrected S(Q) data corresponding to all the two
# theta bins. Here is the thing to do:
#   1. For each two theta bin, grab the data in a certain Q range according to
#      read-in parameters from `ttheta_group_file`.
#   2. For each two theta bin, fill the data in the low Q range with the data
#      in the same low Q range that comes from the first two theta bin in the
#      same group.
# =============================================================================
print(
    "[Info] Pre-processing (step-1) the isotropic corrected data for the "
    "second stage correction..."
)
ttheta_k_valid = [
    "{:d}".format(int(np.floor(ttheta))) for ttheta in ttheta_valid
]
for i, ttheta in enumerate(ttheta_valid):
    ttheta_k = ttheta_k_valid[i]
    for key, item in ttheta_groups.items():
        if ttheta_k in item["QRange"].keys():
            group = key
            break

qrange_tmp = ttheta_groups[group]["QRange"]
    left_b = qrange_tmp[ttheta_k]["LeftBound"]
    leftmost_b = next(iter(qrange_tmp.items()))[1]["LeftBound"]
    right_b_fill = left_b

# Fill the low Q range with the data from the first two theta bin
    # in the same group.
    ref_key_ig = next(iter(qrange_tmp))
    ref_key_id = ttheta_k_valid.index(ref_key_ig)
    mask_fill = (q_vals >= leftmost_b) & (q_vals < right_b_fill)
    sofq_corrected_first[i][mask_fill] = sofq_corrected_first[
        ref_key_id
    ][mask_fill]

# =============================================================================
# Process each of the two theta groups separately. Here is the thing to do:
#   1. For each group, align all the patterns according to the peak fitting.
#   2. Within each group, we,
#       - for each two theta, figure out and remove the background, using an
#         exponential function
#       - integrate intensities in those pre-defined Q chunks
#       - for each Q chunk, fit against the spherical harmonics across two
#         theta angles, just like what we did with each single Q point when
#         correcting the phi part.
#       - extract the isotropic term for each Q chunk, calculate the ratio
#         of the original integrated intensity over the isotropic intensity
#         for each two theta value. Then rescale each Q chunk for each two
#         theta.
#       - add back the exponential background
# =============================================================================
print(
    "[Info] Pre-processing (step-2) the isotropic corrected data for the "
    "second stage correction..."
)
sofq_bkg_removed = dict()
bkg_off = np.inf

for tg_key, item in ttheta_groups.items():
    print(f"[Info] Processing 2Theta group {tg_key}...")
    pa_lb = item["PAParams"]["LeftBound"]
    pa_rb = item["PAParams"]["RightBound"]
    center_tmp = dict()
    for tt_key in item["QRange"].keys():
        if tt_key not in ttheta_k_valid:
            continue
        tt_index = ttheta_k_valid.index(tt_key)
        sofq_tmp = sofq_corrected_first[tt_index]
        mask = (q_vals >= pa_lb) & (q_vals <= pa_rb)
        q_vals_fit = q_vals[mask]
        sofq_fit = sofq_tmp[mask]

p_init = [np.max(sofq_fit), (pa_lb + pa_rb) / 2., 1.]

if testing:
            center_tmp[tt_key] = 0.
        else:
            try:
                popt, _ = curve_fit(
                    gaussian,
                    q_vals_fit,
                    sofq_fit,
                    p0=p_init
                )
                center_tmp[tt_key] = popt[1]
            except RuntimeError:
                center_tmp[tt_key] = 0.

# center_ref = next(iter(center_tmp.items()))[1]
    center_ref = 5.06
    for tt_key, c_item in center_tmp.items():
        # According to the peak fitting failure handling just above, the
        # corresponding peak center would be set to 0.
        if c_item == 0.:
            offset = 0.
        else:
            offset = center_ref - c_item

tt_index = ttheta_k_valid.index(tt_key)
        q_vals_tmp = q_vals + offset
        sofq_tmp = sofq_corrected_first[tt_index].copy()
        q_bin_s = q_vals[1] - q_vals[0]
        if offset < 0:
            while q_vals_tmp[-1] < q_vals[-1]:
                q_vals_tmp = np.append(
                    q_vals_tmp,
                    q_vals_tmp[-1] + q_bin_s
                )
                sofq_tmp = np.append(
                    sofq_tmp,
                    sofq_corrected_first[tt_index][-1]
                )
        else:
            while q_vals_tmp[0] > q_vals[0]:
                q_vals_tmp = np.insert(
                    q_vals_tmp,
                    0,
                    q_vals_tmp[0] - q_bin_s
                )
                sofq_tmp = np.insert(
                    sofq_tmp,
                    0,
                    sofq_corrected_first[tt_index][0]
                )
        _, sofq_rebin = rebin(
            q_vals_tmp,
            sofq_tmp,
            q_vals[0],
            q_bin_s,
            q_vals[-1]
        )
        sofq_corrected_first[tt_index] = sofq_rebin

for tt_key in item["QRange"].keys():
        if tt_key not in ttheta_k_valid:
            continue
        tt_index = ttheta_k_valid.index(tt_key)
        ttheta = ttheta_valid[tt_index]
        sofq_ttheta = sofq_corrected_first[tt_index]

bkg_lb = item["QRange"][tt_key]["LeftBound"]
        bkg_rb = item["QRange"][tt_key]["RightBound"]

# Here we output the aligned patterns.
        # N.B. The alignment only happens internally in each of the defined
        # groups. Therefore, a final stage alignment may still be needed to
        # align the patterns after all steps.
        with open(
            os.path.join(
                out_dir,
                f"{stem_name}ttheta_{tt_key}_corrected_aligned.dat"
            ),
            "w"
        ) as f:
            f.write(f"{len(q_vals)}\n")
            f.write("\n")
            for i, q_val in enumerate(q_vals):
                f.write("{0:10.3F}{1:20.5F}\n".format(q_val, sofq_ttheta[i]))

# We try to use the `argrelextrema` method in `scipy` to figure out
        # local valley points on the pattern and use those valley points for
        # fitting an exponential background. 
        mask_tmp = q_vals >= bkg_lb
        sofq_ttheta_m = sofq_ttheta[mask_tmp]
        q_vals_m = q_vals[mask_tmp]

# The `order` parameter here controls how 'local' the valley will be.
        # If set to a very small value, it will reproduce the overall pattern
        # itself. But setting the value too high may induce too few points.
        q_min = argrelextrema(sofq_ttheta_m, np.less, order=20)
        sofq_min = np.array([sofq_ttheta_m[item_i] for item_i in q_min[0]])
        q_min = np.array([q_vals_m[item_i] for item_i in q_min[0]])

# For the background fitting purpose, we set all S(Q) values above
        # the specified threshold to the value just below (inclusive) the
        # boundary. This will get rid of the impact of those strong noise
        # in the high Q region.
        indices = np.where(q_min <= bkg_rb)[0]
        max_tmp = np.max(q_min[indices])
        up_index = np.where(q_min == max_tmp)[0][0]
        for qm_i, qm_val in enumerate(q_min):
            if qm_val > bkg_rb:
                sofq_min[qm_i] = sofq_min[up_index]

# Specifying the bounding parameter for `curve_fit` helps a lot with
        # fitting convergence. 
        p_init = [0, 0.1, 0.1]
        bounds_l = [-np.inf, 0., 0.]
        bounds_u = [np.inf, 1., 1.]
        popt, pconv = curve_fit(
            bkg_expo,
            q_min,
            sofq_min,
            p0=p_init,
            bounds=(bounds_l, bounds_u)
        )

is_infinite = False
        for row in pconv:
            for element in row:
                if np.isinf(element):
                    is_infinite = True
                    break

if is_infinite:
            print("[Warning] Infinite covariance in the curve fitting.")

with open(
            os.path.join(
                out_dir,
                f"{stem_name}ttheta_{tt_key}_corrected_bkg_pts.dat"
            ),
            "w"
        ) as f:
            f.write(f"{len(q_min)}\n")
            f.write("\n")
            for i, q_val in enumerate(q_min):
                f.write(
                    "{0:10.3F}{1:20.5F}\n".format(
                        q_val, sofq_min[i]
                    )
                )

sofq_ttheta_bkg = bkg_expo(q_vals, popt[0], popt[1], popt[2])
        sofq_ttheta_mb = sofq_ttheta - sofq_ttheta_bkg
        sofq_bkg_removed[tt_key] = sofq_ttheta_mb

# We want to keep a record of the background that most closely resemble
        # 1 at high Q.
        bkg_hq_val = bkg_expo(q_vals[-1], popt[0], popt[1], popt[2])
        bkg_off_tmp = abs(1. - bkg_hq_val)
        if bkg_off_tmp < bkg_off:
            bkg_off = bkg_off_tmp
            bkg_popt = popt

with open(
            os.path.join(
                out_dir,
                f"{stem_name}ttheta_{tt_key}_corrected_bkg_rm_fittedbkg.dat"
            ),
            "w"
        ) as f:
            f.write(f"{len(q_vals)}\n")
            f.write("\n")
            for i, q_val in enumerate(q_vals):
                f.write(
                    "{0:10.3F}{1:20.5F}{2:20.5F}\n".format(
                        q_val, sofq_ttheta_mb[i], sofq_ttheta_bkg[i]
                    )
                )

# Integrate intensities inside Q chunks.
        ttheta_groups[tg_key]["QRange"][tt_key]["QChunkInt"] = list()
        for qv_i, q_val in enumerate(q_vals):
            v_t = sofq_ttheta_mb[qv_i]
            for qc_i, qc_val in enumerate(item["QChunkDef"]):
                if qc_i == 0:
                    if q_val <= qc_val:
                        if len(item["QRange"][tt_key]["QChunkInt"]) == 0:
                            ttheta_groups[
                                tg_key
                            ]["QRange"][tt_key]["QChunkInt"].append(v_t)
                        else:
                            ttheta_groups[
                                tg_key
                            ]["QRange"][tt_key]["QChunkInt"][0] += v_t
                        break
                else:
                    if item["QChunkDef"][qc_i - 1] < q_val <= qc_val:
                        if len(item["QRange"][tt_key]["QChunkInt"]) == qc_i:
                            in_range1 = bkg_lb <= item["QChunkDef"][qc_i - 1]
                            in_range2 = bkg_rb >= qc_val
                            if in_range1 and in_range2:
                                ttheta_groups[
                                    tg_key
                                ]["QRange"][tt_key]["QChunkInt"].append(v_t)
                            else:
                                ttheta_groups[
                                    tg_key
                                ]["QRange"][tt_key]["QChunkInt"].append(np.inf)
                        else:
                            ttheta_groups[
                                tg_key
                            ]["QRange"][tt_key]["QChunkInt"][-1] += v_t
                        break

# Treat each Q chunk as a single 'point' and perform the isotropic correction
# just as before.
sph00 = sph_harm_real(0, 0, 0, 0)
for tg_key, item in ttheta_groups.items():
    params_all = list()
    qchunks_len = len(item["QChunkDef"])
    for i in range(qchunks_len):
        sph_matrix = list()
        sofq_ttheta = list()
        lin_limit = item["LinLimit"][i]

num_eqs = 0
        for tt_key in item["QRange"].keys():
            if tt_key in ttheta_k_valid:
                num_eqs += 1

tt_sph_l_max = 0
        while (tt_sph_l_max + 1 + 1)**2 < num_eqs:
            tt_sph_l_max += 1

for tt_key in item["QRange"].keys():
            if "Factor" not in item["QRange"][tt_key].keys():
                ttheta_groups[tg_key]["QRange"][tt_key]["Factor"] = list()

if tt_key not in ttheta_k_valid:
                continue

tt_index = ttheta_k_valid.index(tt_key)
            ttheta_tmp = ttheta_valid[tt_index]
            v_t = item["QRange"][tt_key]["QChunkInt"][i]
            if v_t == np.inf:
                continue
            sofq_ttheta.append(v_t)
            sph_row = list()
            for l in range(tt_sph_l_max + 1):
                for m in range(-l, l + 1):
                    sph_row.append(
                        sph_harm_real(
                            m,
                            l,
                            ttheta_tmp * np.pi / 180.,
                            0.
                        )
                    )

if "Infinity" not in lin_limit:
                sph_row.append(ttheta_tmp)

sph_matrix.append(sph_row)

sph_matrix = np.array(sph_matrix)
        sofq_ttheta = np.array(sofq_ttheta)

if len(sph_matrix) > 0 or len(sofq_ttheta) > 0:
            if "Infinity" in lin_limit:
                num_terms = (tt_sph_l_max + 1)**2
            else:
                num_terms = (tt_sph_l_max + 1)**2 + 1

params = np.array(
                [1 for _ in range(num_terms)]
            )

min_b = min(sofq_ttheta) / sph00
            max_b = max(sofq_ttheta) / sph00
            bounds = [(min_b, max_b)]
            if "Infinity" not in lin_limit:
                num_b_add = num_terms - 2
            else:
                num_b_add = num_terms - 1

for _ in range(num_b_add):
                bounds.append(
                    (-float('inf'), float('inf'))
                )

if "Infinity" not in lin_limit:
                bounds.append((lin_limit[0], lin_limit[1]))

result = minimize(
                target_function,
                params,
                args=(sph_matrix, sofq_ttheta),
                method='SLSQP',
                bounds=bounds
            )

optimal_params = result.x

y00 = sph_harm_real(0, 0, 0, 0)
            s00 = optimal_params[0]
            if "Infinity" not in lin_limit:
                ctt = optimal_params[-1]
        else:
            y00 = 0.
            s00 = 0.

# Calculate the ratio of the original integrated intensity over
        # the isotropic intensity.
        for tt_key in item["QRange"].keys():
            if tt_key not in ttheta_k_valid:
                continue

if i == 0:
                factor_v = 1.
            else:
                tt_index = ttheta_k_valid.index(tt_key)
                ttheta_tmp = ttheta_valid[tt_index]
                if "Infinity" in lin_limit:
                    top_val = y00 * s00
                else:
                    top_val = y00 * s00 + ctt * ttheta_tmp

bottom_val = item["QRange"][tt_key]["QChunkInt"][i]
                factor_v = top_val / bottom_val

# if tg_key == "Group-1":
                #     if i == 2:
                #         print("Debugging -> ")
                #         print(
                #             f"fit val: {top_val}, init val: {bottom_val}",
                #             f"y00: {y00}, s00: {s00}",
                #             f"tt: {ttheta_tmp}"
                #         )
                #         print(
                #             f"fit val: {top_val}, init val: {bottom_val}",
                #             f"y00: {y00}, s00: {s00}, ctt: {ctt}",
                #             f"tt: {ttheta_tmp}"
                #         )

ttheta_groups[tg_key]["QRange"][tt_key]["Factor"].append(
                factor_v
            )

# Rescale the data by Q chunks according to the ratio obtained above.
    for tt_key in item["QRange"].keys():
        if tt_key not in ttheta_k_valid:
            continue

for i, q_val in enumerate(q_vals):
            for ii, qc_val in enumerate(item["QChunkDef"]):
                if ii > 0:
                    # if q_vals[ii - 1] < q_val <= qc_val and q_val > 4.:
                    if q_vals[ii - 1] < q_val <= qc_val:
                        sofq_bkg_removed[tt_key][i] *= item[
                            "QRange"
                        ][tt_key]["Factor"][ii]

break

# Add back in the exponential background.
        bkg_add = bkg_expo(q_vals, bkg_popt[0], bkg_popt[1], bkg_popt[2])
        sofq_final = sofq_bkg_removed[tt_key] + bkg_add

with open(
            os.path.join(
                out_dir,
                f"{stem_name}ttheta_{tt_key}_corrected_wo_bkg.dat"
            ),
            "w"
        ) as f:
            f.write(f"{len(q_vals)}\n")
            f.write("\n")
            for i, q_val in enumerate(q_vals):
                f.write(
                    "{0:10.3F}{1:20.5F}\n".format(
                        q_val, sofq_bkg_removed[tt_key][i]
                    )
                )

with open(
            os.path.join(
                out_dir,
                f"{stem_name}ttheta_{tt_key}_corrected_final.dat"
            ),
            "w"
        ) as f:
            f.write(f"{len(q_vals)}\n")
            f.write("\n")
            for i, q_val in enumerate(q_vals):
                f.write(
                    "{0:10.3F}{1:20.5F}\n".format(
                        q_val, sofq_final[i]
                    )
                )

with open(os.path.join(out_dir, f"{stem_name}_ttheta_1na_2a.json"), "w") as f:
    json.dump(ttheta_groups, f, indent=4)

print("\n=========================================================")
print("[Info] Done with the two-stage isotropic S(Q) correction.")
print("[Info] Proceed to merge the corrected S(Q) data by bank.")
print("=========================================================\n")

with open("output_group.json", "r") as f:
    group_dict = json.load(f)

bank_dict = dict()
for key in ttheta_k_valid:
    for keykey in ttheta_groups.keys():
        if key in ttheta_groups[keykey]["QRange"].keys():
            key_group = keykey
            break

file_t = os.path.join(
        out_dir,
        f"{stem_name}ttheta_{key}_corrected_final.dat"
    )
    if os.path.exists(file_t):
        y_tmp = np.loadtxt(file_t, skiprows=2, dtype=float)[:, 1]
        for gkey, item in group_dict.items():
            if item[0] <= float(key) < item[1]:
                bank = gkey
                break

if bank not in bank_dict:
            bank_dict[bank] = {
                "y": y_tmp,
                "num_bands": 1,
                "q_counts": [0 for _ in range(len(q_vals))]
            }
        else:
            bank_dict[bank]["y"] += y_tmp
            bank_dict[bank]["num_bands"] += 1

q_chunk_tmp = ttheta_groups[key_group]["QChunkDef"]
        for i, q_val in enumerate(q_vals):
            for ii, qc_val in enumerate(q_chunk_tmp):
                if ii == 0:
                    if q_val <= qc_val:
                        q_chunk_id = ii
                        break
                else:
                    if q_chunk_tmp[ii - 1] < q_val <= qc_val:
                        q_chunk_id = ii
                        break

if q_chunk_id == 0:
                factor = 1.
            else:
                factor = ttheta_groups[
                    key_group
                ]["QRange"][key]["Factor"][q_chunk_id]

if factor > 0.:
                bank_dict[bank]["q_counts"][i] += 1

for key, item in bank_dict.items():
    for i in range(len(q_vals)):
        if item["q_counts"][i] > 0:
            item["y"][i] /= item["q_counts"][i]
        else:
            item["y"][i] /= item["num_bands"]

with open(
        os.path.join(out_dir, f"{stem_name}{key}_merged_1na_2a_ave.dat"),
        "w"
    ) as f:
        f.write(f"{len(q_vals)}\n")
        f.write("\n")
        for i, q_val in enumerate(q_vals):
            f.write("{0:10.3F}{1:20.5F}\n".format(q_val, item["y"][i]))

print("=========================================================")
print("[Info] Done with merging the corrected S(Q) by banks.")
print("=========================================================\n")

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")