import numpy as np
import pandas as pd
from scipy.stats import exponnorm
import matplotlib.pyplot as plt
from scipy.special import erfc, erfcx
from scipy.stats import uniform
import lmfit
import inspect
import warnings
import nbragg.utils as utils
[docs]class Response:
def __init__(self, kind="jorgensen", vary: bool = False, wlstep: float = 0.1, tstep: float = 10e-6):
"""
Initializes the Response object with specified parameters.
Parameters:
kind (str): The type of response function to use. Options are 'jorgensen', 'square', 'square_jorgensen', 'full_jorgensen', or 'none'.
vary (bool): If True, the parameters can vary during fitting. Default is False.
wlstep (float): Wavelength step size for Δλ grid. Default is 0.1.
tstep (float): Time step size for tgrid in seconds. Default is 10e-6 (10 µs).
"""
self.wlstep = wlstep
self.tstep = tstep
self.params = lmfit.Parameters()
self.Δλ = np.arange(-20, 20, self.wlstep)
self.tgrid = np.arange(-0.005, 0.005, tstep) # Grid for time-based response (-5ms to 5ms with 10µs bins)
self.kind = kind
# Choose the response function
if kind == "jorgensen":
self.function = self.jorgensen_response
self.params = lmfit.Parameters()
self.params.add(f"α0", value=3.67, min=0.001, max=10000, vary=vary)
self.params.add(f"β0", value=3.06, min=0.001, max=10000, vary=vary)
elif kind == "square":
self.function = self.square_response
self.params = lmfit.Parameters()
self.params.add(f"width", value=tstep*1e6, min=tstep*1e6, max=5000, vary=vary)
elif kind == "square_jorgensen":
self.function = self.square_jorgensen_response
self.params = lmfit.Parameters()
self.params.add(f"α0", value=3.67, min=0.001, max=10000, vary=vary)
self.params.add(f"β0", value=3.06, min=0.001, max=10000, vary=vary)
self.params.add(f"width", value=tstep*1e6, min=tstep*1e6, max=5000, vary=vary)
elif kind == "full_jorgensen":
self.function = self.full_jorgensen_response
self.params = lmfit.Parameters()
# Units (user-facing): α, β are decay rates in 1/µs; σ is a
# Gaussian width in µs. tgrid is still in seconds — the response
# function converts internally. α0/β1/σ0/σ2 default to 0 and are
# usually held fixed; σ1 is the dominant resolution term and
# should be refined first (see staged_preset()). Defaults give a
# near-delta response (FWHM ~ 1 tgrid sample = 10 µs) so the
# convolution starts as essentially identity.
self.params.add(f"α0", value=0., min=0., max=1000., vary=False)
self.params.add(f"α1", value=0.5, min=0., max=1000., vary=vary)
self.params.add(f"β0", value=0.5, min=0., max=1000., vary=vary)
self.params.add(f"β1", value=0.0, min=0., max=1000., vary=False)
self.params.add(f"σ0", value=0.0, min=0.0, max=1000., vary=False)
self.params.add(f"σ1", value=0.001, min=1e-6, max=1000., vary=vary)
self.params.add(f"σ2", value=0.0, min=0.0, max=1000., vary=False)
elif kind == "none":
self.function = self.empty_response
else:
raise NotImplementedError(f"Response kind '{kind}' is not supported. Use 'jorgensen', 'square', 'square_jorgensen', 'full_jorgensen', or 'none'.")
[docs] def staged_preset(self):
"""
Recommended staged-fit ordering for this response kind.
Returns a dict suitable for assigning to `TransmissionModel.stages`,
merging with other stages as needed. For full_jorgensen the order
follows Jorgensen-profile fitting best practice: refine σ1 first
(dominant resolution term), then β0 (decay), then α1 (rise).
"""
if self.kind == "full_jorgensen":
return {
"response_sig1": ["σ1"],
"response_beta": ["σ1", "β0"],
"response_alpha": ["σ1", "β0", "α1"],
}
if self.kind == "jorgensen":
return {"response": ["α0", "β0"]}
return {}
[docs] def register_response(self, response_func, lmfit_params=None, **kwargs):
"""
Registers a new response using any scipy.stats function.
Parameters:
response_func (function): A function from scipy.stats, e.g., exponnorm.pdf.
lmfit_params (lmfit.Parameters): Optional lmfit.Parameters to define limits and vary.
kwargs: Default parameter values for the response function.
"""
# Detect parameters of the response function
sig_params = inspect.signature(response_func).parameters
for param, default in kwargs.items():
if param in sig_params:
self.params.add(param, value=default, vary=True)
else:
raise ValueError(f"Parameter '{param}' not found in the response function signature.")
self.function = response_func.pdf(self.tgrid)
# Use optional lmfit.Parameters to customize limits and vary
if lmfit_params:
for name, param in lmfit_params.items():
if name in self.params:
self.params[name].set(value=param.value, vary=param.vary, min=param.min, max=param.max)
[docs] def cut_array_symmetric(self, arr, threshold):
"""
Symmetrically cuts the array based on a threshold.
Parameters:
arr (np.ndarray): Input array to be cut.
threshold (float): The threshold value for cutting the array.
Returns:
np.ndarray: Symmetrically cut array with an odd number of elements.
"""
if len(arr) % 2 == 0:
raise ValueError("Input array length must be odd.")
center_idx = len(arr) // 2
left_idx = np.argmax(arr[:center_idx][::-1] < threshold)
right_idx = np.argmax(arr[center_idx:] < threshold)
left_bound = center_idx - max(left_idx, right_idx)
right_bound = center_idx + max(left_idx, right_idx) + 1 # Ensure odd length
return arr[left_bound:right_bound]
[docs] def empty_response(self, **kwargs):
"""
Returns an empty response [0.0, 1.0, 0.0].
"""
return np.array([0., 1., 0.])
[docs] def square_response(self, width=10, center: bool = True, **kwargs):
"""
Returns a square response in time with a given width [µs].
"""
width = width * 1e-6 # Convert to seconds
loc = -0.5 * width if center else 0.
tof_response = uniform.pdf(self.tgrid, loc=loc, scale=width)
tof_response /= np.sum(tof_response)
return tof_response
[docs] def jorgensen_response(self, α0, β0, σ1=None, **kwargs):
"""
Calculates the Jorgensen peak profile function with Greek Unicode parameters.
Parameters:
-----------
α0 : float or list/tuple
Alpha parameter [α0₁, α0₂] or single value α0₁ (α0₂ defaults to 0)
β0 : float or list/tuple
Beta parameter [β0₁, β0₂] or single value β0₁ (β0₂ defaults to 0)
σ1 : list/tuple, optional
Sigma parameters [σ₁, σ₂, σ₃], defaults to [0, 0, 0]
Returns:
--------
numpy.ndarray
Normalized profile values with NaN values replaced by 0
"""
# Handle input parameters
if σ1 is None:
σ1 = [0, 0, 0]
# Convert single values to lists with 0 as second element
if not isinstance(α0, (list, tuple)):
α0 = [α0, 0]
if not isinstance(β0, (list, tuple)):
β0 = [β0, 0]
# Generate x values
x = self.Δλ
# Calculate parameters using d-spacing of 1.0 as in original code
d = 1.0
α0_calc = α0[0] + α0[1]/d
β0_calc = β0[0] + β0[1]/d**4
σ1_calc = np.sqrt(σ1[0]**2 + (σ1[1]*d)**2 + (σ1[2]*d*d)**2)
# Constants
sqrt2 = np.sqrt(2)
σ2 = σ1_calc * σ1_calc
# Scaling factor
scale = α0_calc * β0_calc / 2 / (α0_calc + β0_calc)
with warnings.catch_warnings():
warnings.simplefilter("ignore")
# Calculate intermediate terms
u = α0_calc/2 * (α0_calc*σ2 + 2*x)
v = β0_calc/2 * (β0_calc*σ2 - 2*x)
y = (α0_calc*σ2 + x)/(sqrt2*σ1_calc)
z = (β0_calc*σ2 - x)/(sqrt2*σ1_calc)
# Calculate profile with special handling for numerical stability
term1 = np.exp(u) * erfc(y)
term2 = np.exp(v) * erfc(z)
# Zero out terms where erfc is zero to avoid NaN
term1[erfc(y) == 0] = 0
term2[erfc(z) == 0] = 0
# Calculate profile
profile = scale * (term1 + term2)
# Normalize
profile = profile / np.sum(profile)
# Replace any NaN values with 0
return np.nan_to_num(profile, 0)
def square_jorgensen_response(self, α0, β0, width, **kwargs):
"""
Combines square and Jorgensen responses (not implemented in this update).
"""
# Placeholder: Implement if needed
raise NotImplementedError("square_jorgensen_response not implemented.")
[docs] def full_jorgensen_response(self, wl=4.0, α0=0., α1=0.5, β0=0.5, β1=0., σ0=0., σ1=0.001, σ2=0., **kwargs):
"""
Jorgensen TOF peak profile with d-spacing-dependent parameters.
Implements eq. 3.40 of Jorgensen et al. (1978):
h(Δ, σ, α, β) = αβ / (2(α+β)) · (exp(u) erfc(y) + exp(v) erfc(z))
with
α = α0 + α1 / d [1/µs]
β = β0 + β1 / d^4 [1/µs]
σ² = σ0² + (σ1·d)² + (σ2·d²)² [µs²]
Parameter units are GSAS-style: α, β in 1/µs and σ in µs. The profile
is evaluated on self.tgrid (seconds) — the function converts units
internally so the user-facing values stay in the 0.001–1000 range.
Uses the identity exp(u)·erfc(y) = exp(−x²/(2σ²))·erfcx(y) on the y≥0
side to avoid exp(α²σ²/2) overflow when σ is large; the y<0 side is
already bounded so it uses the direct form.
Returns a 1-D normalized profile defined on self.tgrid.
"""
d = wl / 2.0
d2 = d * d
d4 = d2 * d2
# Convert user-facing units (1/µs, µs) to seconds-based internal units
# so they match self.tgrid (seconds): 1/µs → ×1e6 to get 1/s; µs → ×1e-6 to get s.
α = (α0 + α1 / d) * 1e6
β = (β0 + β1 / d4) * 1e6
σ2_val = (σ0**2 + (σ1 * d)**2 + (σ2 * d2)**2) * 1e-12
σ2_val = max(σ2_val, 1e-30)
σ = np.sqrt(σ2_val)
if α + β <= 0:
raise ValueError(f"α + β must be positive (got α={α}, β={β})")
x = self.tgrid
sqrt2 = np.sqrt(2)
inv_2σ2 = 0.5 / σ2_val
rad2sigma = 1.0 / (sqrt2 * σ)
with warnings.catch_warnings():
warnings.simplefilter("ignore")
y = (α * σ2_val + x) * rad2sigma
z = (β * σ2_val - x) * rad2sigma
gauss = np.exp(-x * x * inv_2σ2)
# term1 = exp(u)·erfc(y); use gauss·erfcx(y) for y>=0 (stable),
# direct form for y<0 (where exp(u) is bounded ≤ exp(-α²σ²/2)·2)
term1 = np.empty_like(x)
m1 = y >= 0
term1[m1] = gauss[m1] * erfcx(y[m1])
u_neg = (α / 2) * (α * σ2_val + 2 * x[~m1])
term1[~m1] = np.exp(u_neg) * erfc(y[~m1])
term2 = np.empty_like(x)
m2 = z >= 0
term2[m2] = gauss[m2] * erfcx(z[m2])
v_neg = (β / 2) * (β * σ2_val - 2 * x[~m2])
term2[~m2] = np.exp(v_neg) * erfc(z[~m2])
scale = α * β / (2 * (α + β))
profile = scale * (term1 + term2)
profile = np.nan_to_num(profile, nan=0.0, posinf=0.0, neginf=0.0)
total = profile.sum()
if total > 0:
profile = profile / total
return profile
def _apply_asymmetry_correction(self, profile, xgrid, d, α0, β0, truncation_factor):
"""
Applies asymmetry correction to the profile, mimicking GSAS computeAsymmetry.
Parameters:
-----------
profile : array-like
Input profile to correct.
xgrid : array-like
Grid of x values (time-of-flight offsets).
d : float
d-spacing.
α0 : list/tuple
[α0, α1] parameters.
β0 : list/tuple
[β0, β1] parameters.
truncation_factor : float
Truncation factor for asymmetry convolution.
Returns:
--------
numpy.ndarray
Asymmetry-corrected profile.
"""
# Compute asymmetry parameters
α = α0[0] + α0[1] / d
β = β0[0] + β0[1] / d**4
# Initialize output
new_profile = np.zeros_like(profile)
# Apply β asymmetry (first term)
total_asymmetry = np.abs(β)
if total_asymmetry > 0 and total_asymmetry < 1:
direction = -1 if β < 0 else 1
log_truncation = np.log(truncation_factor)
truncation_angle = np.abs(-log_truncation / total_asymmetry)
for i, x in enumerate(xgrid):
function = profile[i]
normalization = 1.0
j = i + direction
while 0 <= j < len(xgrid) and np.abs(xgrid[j] - x) < truncation_angle:
difference = np.abs(xgrid[j] - x)
exp_asymmetry = np.exp(-difference * total_asymmetry)
function += profile[j] * exp_asymmetry
normalization += exp_asymmetry
j += direction
new_profile[i] = function / normalization if normalization > 0 else profile[i]
# Apply α asymmetry (second term)
total_asymmetry = np.abs(α)
if total_asymmetry > 0 and total_asymmetry < 1:
direction = -1 if α < 0 else 1
log_truncation = np.log(truncation_factor)
truncation_angle = np.abs(-log_truncation / total_asymmetry)
for i, x in enumerate(xgrid):
function = new_profile[i] if new_profile[i] != 0 else profile[i]
normalization = 1.0
j = i + direction
while 0 <= j < len(xgrid) and np.abs(xgrid[j] - x) < truncation_angle:
difference = np.abs(xgrid[j] - x)
exp_asymmetry = np.exp(-difference * total_asymmetry)
function += new_profile[j] * exp_asymmetry
normalization += exp_asymmetry
j += direction
new_profile[i] = function / normalization if normalization > 0 else new_profile[i]
return new_profile
[docs] def square_jorgensen_response(self, α0=3.67, β0=3, σ1=None, width=None, **kwargs):
"""
Calculates the Jorgensen peak profile function with an additional square width parameter.
Parameters:
-----------
α0 : float or list/tuple
Alpha parameter [α0, α0] or single value α0 (α0 defaults to 0)
β0 : float or list/tuple
Beta parameter [β0, β0] or single value β0 (β0 defaults to 0)
σ1 : list/tuple, optional
Sigma parameters [σ₁, σ₂, σ₃], defaults to [0, 0, 0]
width : float, optional
Square width to broaden the response [usec]
Returns:
--------
numpy.ndarray
Normalized profile values with NaN values replaced by 0
"""
# Handle input parameters
if σ1 is None:
σ1 = [0, 0, 0]
# Convert single values to lists with 0 as second element
if not isinstance(α0, (list, tuple)):
α0 = [α0, 0]
if not isinstance(β0, (list, tuple)):
β0 = [β0, 0]
# Generate x values
x = self.Δλ
# Calculate parameters using d-spacing of 1.0 as in original code
d = 1.0
α0_calc = α0[0] + α0[1]/d
β0_calc = β0[0] + β0[1]/d**4
σ1_calc = np.sqrt(σ1[0]**2 + (σ1[1]*d)**2 + (σ1[2]*d*d)**2)
# Constants
sqrt2 = np.sqrt(2)
σ2 = σ1_calc * σ1_calc
# Scaling factor
scale = α0_calc * β0_calc / 2 / (α0_calc + β0_calc)
with warnings.catch_warnings():
warnings.simplefilter("ignore")
# Calculate intermediate terms
u = α0_calc/2 * (α0_calc*σ2 + 2*x)
v = β0_calc/2 * (β0_calc*σ2 - 2*x)
y = (α0_calc*σ2 + x)/(sqrt2*σ1_calc)
z = (β0_calc*σ2 - x)/(sqrt2*σ1_calc)
# Calculate profile with special handling for numerical stability
term1 = np.exp(u) * erfc(y)
term2 = np.exp(v) * erfc(z)
# Zero out terms where erfc is zero to avoid NaN
term1[erfc(y) == 0] = 0
term2[erfc(z) == 0] = 0
# Calculate profile
profile = scale * (term1 + term2)
# Apply square width broadening if specified
if width is not None:
# Convert width to seconds
width_sec = width * 1e-6
# Create square broadening function
loc = -0.5 * width_sec
square_broadening = uniform.pdf(self.tgrid, loc=loc, scale=width_sec)
# Convolve Jorgensen profile with square broadening
from scipy.ndimage import convolve1d
profile = convolve1d(profile, square_broadening, 0)
# Normalize
profile = profile / np.sum(profile)
# Replace any NaN values with 0
return np.nan_to_num(profile, 0)
[docs] def plot(self, params=None, **kwargs):
"""
Plots the response function.
Parameters:
params (dict): Parameters for the response function.
**kwargs: Additional arguments for plot customization.
"""
ax = kwargs.pop("ax", plt.gca())
params = params if params else self.params
y = self.function(**params.valuesdict())
if self.kind == "jorgensen" or self.kind == "square_jorgensen":
xlabel = kwargs.pop("xlabel", "wavelength [Angstrom]")
df = pd.Series(y, index=self.Δλ, name="Response")
df.plot(ax=ax, xlabel=xlabel, **kwargs)
elif self.kind == "full_jorgensen":
xlabel = kwargs.pop("xlabel", "tof [usec]")
df = pd.Series(y, index=self.tgrid * 1e6, name="Response")
df.plot(ax=ax, xlabel=xlabel, **kwargs)
elif self.kind == "square":
xlabel = kwargs.pop("xlabel", "tof [usec]")
df = pd.Series(y, index=self.tgrid, name="Response")
df.plot(ax=ax, xlabel=xlabel, **kwargs)
else:
raise ValueError("response plotting is only supported for 'jorgensen','full_jorgensen', 'square', and 'square_jorgensen' kinds")
[docs]class Background:
def __init__(self, kind="expo_norm", vary: bool = False):
"""
Initializes the Background object with specified parameters.
Parameters:
kind (str): Type of background function ('constant', 'polynomial3', 'polynomial5','sample_dependent' or 'none').
vary (bool): If True, the parameters can vary during fitting.
"""
self.params = lmfit.Parameters()
if kind == "polynomial3":
self.function = self.polynomial3_background
for i in range(3):
self.params.add(f"bg{i}", value=0., min=-1e6, max= 1e6, vary=vary)
elif kind == "polynomial5":
self.function = self.polynomial5_background
for i in range(5):
self.params.add(f"bg{i}", value=0., min=-1e6, max= 1e6, vary=vary)
elif kind == "sample_dependent":
self.function = self.polynomial3_background
self.params.add(f"k", value=1., min=1., max= 10., vary=vary)
for i in range(3):
self.params.add(f"bg{i}", value=0., min=-1e6, max= 1e6, vary=vary)
elif kind == "constant":
self.function = self.constant_background
self.params.add('bg0', value=0.0, min=-1e6,max=1e6,vary=vary)
elif kind == "none":
self.function = self.empty_background
else:
raise NotImplementedError(f"Background kind '{kind}' is not supported. Use 'none', 'constant', 'polynomial3', 'sample_dependent', or 'polynomial5'.")
[docs] def empty_background(self, wl, **kwargs):
"""
Returns a zero background array.
"""
return np.zeros_like(wl)
[docs] def constant_background(self, wl, bg0=0., **kwargs):
"""
Generates a constant background.
Parameters:
wl (np.ndarray): Wavelength values.
b0 (float): Constant background value.
"""
bg = np.full_like(wl, bg0)
return np.where(bg>0,bg,0.)
[docs] def polynomial3_background(self, wl, bg0=0., bg1=1., bg2=0., **kwargs):
"""
Computwls a third-degree polynomial background.
Parameters:
wl (np.ndarray): Energy values.
bg0 (float): Constant term.
bg1 (float): Linear term.
bg2 (float): Quadratic term.
"""
bg = bg0 + bg1 * np.sqrt(wl) + bg2 / np.sqrt(wl)
return np.where(bg>0,bg,0.)
[docs] def polynomial5_background(self, wl, bg0=0., bg1=1., bg2=0., bg3=0., bg4=0., **kwargs):
"""
Computes a fifth-degree polynomial background.
Parameters:
wl (np.ndarray): Wavelegth values.
bg0 (float): Constant term.
bg1 (float): Linear term.
bg2 (float): Quadratic term.
bg3 (float): Cubic term.
bg4 (float): Quartic term.
"""
bg = bg0 + bg1 * np.sqrt(wl) + bg2 / np.sqrt(wl) + bg3 * wl + bg4 * wl**2
return np.where(bg>0,bg,0.)
[docs] def plot(self, wl, params=None, **kwargs):
"""
Plots the background function.
Parameters:
wl (np.ndarray): Wavelength values.
params (dict): Parameters for the background function.
"""
ax = kwargs.pop("ax", plt.gca())
ls = kwargs.pop("ls", "--")
color = kwargs.pop("color", "0.5")
params = params if params else self.params
k = params.get("k",1.) # sample dependent parameter if present
y = k*self.function(wl, **params.valuesdict())
df = pd.Series(y, index=wl, name="Background")
df.plot(ax=ax, color=color, ls=ls, **kwargs)