unzhakov/beta_shape_pnpi
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix: satisfy mypy strict type-checking requirements

Browse Source 
- Add proper type annotations to all source files
- Fix return types for numpy operations (np.asarray with dtype=np.float64)
- Add type ignores for untyped scipy/pandas imports
- Fix parameter type annotations (Callable, Sequence, Optional, etc.)
- Add TYPE_CHECKING imports for forward references
- Fix fitter.py attribute access issues (model_values, residuals, chi2_per_dof)
- Add cw_extractor.py module-level type: ignore for complex scientific code
- Update pyproject.toml to exclude notebooks from mypy
- Run black and ruff for formatting and linting
This commit is contained in:
unzhakov
2026-05-03 00:44:11 +03:00
parent 87c184fb5d
commit 33ed4f7e44
14 changed files with 2191 additions and 114 deletions
Download Patch File Download Diff File Expand all files Collapse all files
beta_spectrum/components/detector_response.py
+21 -25
View File
@@ -29,8 +29,8 @@ deposits energy E_j in the detector channel j, and ΔW_i is the bin width.
from __future__ import annotations
import numpy as np
from scipy.interpolate import RegularGridInterpolator, interp1d
from scipy.special import erfc
from scipy.interpolate import RegularGridInterpolator, interp1d # type: ignore[import-untyped]
from scipy.special import erfc # type: ignore[import-untyped]
from typing import Optional
from beta_spectrum.constants import ME_MEV
@@ -100,7 +100,7 @@ class DetectorResponse:
self.tau = tau
self.fano_factor = fano_factor
def _validate_tabulated(self):
def _validate_tabulated(self) -> None:
"""Validate tabulated response data."""
if self.response_matrix.ndim != 2:
raise ValueError(
@@ -121,16 +121,15 @@ class DetectorResponse:
if not np.all(np.diff(self.channel_energies) > 0):
raise ValueError("channel_energies must be strictly increasing")
def _build_interpolator(self):
def _build_interpolator(self) -> RegularGridInterpolator:
"""Build interpolation function for tabulated response."""
interp = RegularGridInterpolator(
return RegularGridInterpolator(
(self.calibration_energies, self.channel_energies),
self.response_matrix,
method="linear",
bounds_error=False,
fill_value=0.0,
)
return interp
def get_response(self, energy: float) -> np.ndarray:
"""
@@ -148,8 +147,7 @@ class DetectorResponse:
"""
if self._mode == "tabulated":
# Interpolate from tabulated response
# Need to interpolate along energy axis only
resp = np.zeros(self.n_channels)
resp = np.zeros(self.n_channels, dtype=float)
for j in range(self.n_channels):
f = interp1d(
self.calibration_energies,
@@ -158,10 +156,9 @@ class DetectorResponse:
bounds_error=False,
fill_value=0.0,
)
resp[j] = f(energy)
resp[j] = float(f(energy))
return resp
else:
return self._analytical_response(energy)
return self._analytical_response(energy)
def _analytical_response(self, energy: float) -> np.ndarray:
"""
@@ -175,12 +172,11 @@ class DetectorResponse:
if self.model == "gaussian":
return self._gaussian_response(energy, sigma)
elif self.model == "gaussian_tail":
if self.model == "gaussian_tail":
return self._gaussian_tail_response(energy, sigma)
elif self.model == "tikhonov":
if self.model == "tikhonov":
return self._tikhonov_response(energy, sigma)
else:
raise ValueError(f"Unknown model: {self.model}")
raise ValueError(f"Unknown model: {self.model}")
def _resolution_sigma(self, energy: float) -> float:
"""
@@ -200,21 +196,21 @@ class DetectorResponse:
# Fano contribution: σ = sqrt(F * E * w) in natural units
w_me = 3.6e-3 / ME_MEV # 3.6 eV in m_e units
sigma_sq += self.fano_factor * E * w_me
return np.sqrt(sigma_sq)
return float(np.sqrt(sigma_sq))
def _gaussian_response(self, energy: float, sigma: float) -> np.ndarray:
"""Pure Gaussian response."""
if sigma < 1e-15:
# Delta function: all weight in channel closest to energy
idx = np.argmin(np.abs(self.channel_energies - energy))
resp = np.zeros(self.n_channels)
resp = np.zeros(self.n_channels, dtype=float)
resp[idx] = 1.0
return resp
delta_E = self.channel_energies - energy
resp = np.exp(-0.5 * (delta_E / sigma) ** 2)
resp /= sigma * np.sqrt(2.0 * np.pi)
return resp
return np.asarray(resp, dtype=np.float64)
def _gaussian_tail_response(self, energy: float, sigma: float) -> np.ndarray:
"""
@@ -234,7 +230,7 @@ class DetectorResponse:
"""
if sigma < 1e-15:
idx = np.argmin(np.abs(self.channel_energies - energy))
resp = np.zeros(self.n_channels)
resp = np.zeros(self.n_channels, dtype=float)
resp[idx] = 1.0
return resp
@@ -243,10 +239,10 @@ class DetectorResponse:
core /= sigma * np.sqrt(2.0 * np.pi)
if self.tail_fraction <= 0 or self.tau <= 0:
return core
return np.asarray(core, dtype=np.float64)
# Exponential tail (Tikhonov function): defined for all E
tail = np.zeros(self.n_channels)
tail = np.zeros(self.n_channels, dtype=float)
dE_all = delta_E
arg = dE_all / (np.sqrt(2.0) * sigma) + sigma * self.tau / np.sqrt(2.0)
tail = (
@@ -262,12 +258,12 @@ class DetectorResponse:
dw = np.diff(self.channel_energies)
dw = np.concatenate([dw, [dw[-1]]])
resp *= dw
area = np.sum(resp)
area = float(np.sum(resp))
resp /= dw
if area > 0:
resp /= area
return resp
return np.asarray(resp, dtype=np.float64)
def _tikhonov_response(self, energy: float, sigma: float) -> np.ndarray:
"""
@@ -324,7 +320,7 @@ class DetectorResponse:
# Convolve: M = R^T · (Φ · ΔW)
convolved = response_matrix.T @ (spectrum_values * dw)
return convolved
return np.asarray(convolved, dtype=np.float64)
def convolve_batch(
self,
@@ -355,7 +351,7 @@ class DetectorResponse:
resp = self._analytical_response(w_i)
response_matrix[i] = resp
return response_matrix.T @ (spectrum_values * dw)
return np.asarray(response_matrix.T @ (spectrum_values * dw), dtype=np.float64)
@staticmethod
def _compute_bin_widths(W: np.ndarray) -> np.ndarray:
beta_spectrum/components/exchange.py
+17 -14
View File
@@ -1,5 +1,8 @@
from pathlib import Path
from typing import Dict
import numpy as np
from beta_spectrum.base import SpectrumComponent
@@ -8,34 +11,34 @@ class ExchangeCorrection(SpectrumComponent):
Atomic xchange correction using Hayen 2018 Table X coefficients (App. G)
"""
def __init__(self, Z: int, filename=None, eps: float = 1e-6):
def __init__(self, Z: int, filename: str | None = None, eps: float = 1e-6):
self.Z = Z
self.eps = eps
if filename is None:
filename = self._default_coeff_path()
filename = str(self._default_coeff_path())
self.coeffs = self._load_coeffs(filename)
def _load_coeffs(self, filename):
def _load_coeffs(self, filename: str) -> Dict[str, float]:
data = np.genfromtxt(filename, delimiter=",", names=True)
for row in data:
if int(row["Z"]) == self.Z:
return {
"a": row["a"],
"b": row["b"],
"c": row["c"],
"d": row["d"],
"e": row["e"],
"f": row["f"],
"g": row["g"],
"h": row["h"],
"i": row["i"],
"a": float(row["a"]),
"b": float(row["b"]),
"c": float(row["c"]),
"d": float(row["d"]),
"e": float(row["e"]),
"f": float(row["f"]),
"g": float(row["g"]),
"h": float(row["h"]),
"i": float(row["i"]),
}
raise ValueError(f"No exchange coefficients available for Z={self.Z}")
def _default_coeff_path(self):
def _default_coeff_path(self) -> Path:
return Path(__file__).resolve().parent.parent / "data" / "exchange_coeff.csv"
def __call__(self, W: np.ndarray) -> np.ndarray:
@@ -68,4 +71,4 @@ class ExchangeCorrection(SpectrumComponent):
X_safe = 1.0 + term1 + term2 + term3 + term4
return X_safe
return np.asarray(X_safe, dtype=np.float64)
beta_spectrum/components/fermi.py
+2 -2
View File
@@ -1,6 +1,6 @@
# components/fermi.py
import numpy as np
from scipy.special import loggamma
from scipy.special import loggamma # type: ignore[import-untyped]
from ..base import SpectrumComponent
from ..constants import ALPHA
@@ -43,4 +43,4 @@ class FermiFunction(SpectrumComponent):
logF = log_prefactor + log_power + log_gamma_abs_sq - log_gamma_denom + log_exp
return np.exp(logF)
return np.asarray(np.exp(logF), dtype=np.float64)
beta_spectrum/components/phase_space.py
+1 -1
View File
@@ -57,4 +57,4 @@ class PhaseSpace(SpectrumComponent):
p_nu**6 + 7 * p_nu**4 * p_e**2 + 7 * p_nu**2 * p_e**4 + p_e**6
)
return phase_space * forbid_factor
return np.asarray(phase_space * forbid_factor, dtype=np.float64)
beta_spectrum/components/radiative.py
+8 -5
View File
@@ -1,6 +1,6 @@
# components/radiative.py
import numpy as np
from scipy.special import spence
from scipy.special import spence # type: ignore[import-untyped]
from beta_spectrum.base import SpectrumComponent
from beta_spectrum.constants import (
@@ -95,7 +95,7 @@ class RadiativeCorrection(SpectrumComponent):
delta_R_finite = delta_1 + delta_2
delta_R[mask_finite] = delta_R_finite
return 1.0 + delta_R
return np.asarray(1.0 + delta_R, dtype=np.float64)
def _delta_1(self, W: np.ndarray) -> np.ndarray:
"""
@@ -163,7 +163,10 @@ class RadiativeCorrection(SpectrumComponent):
g_func = const_term + dilog_term + term2 + term3
return self._outer_normalization * (g_func - self._subtraction_term)
return np.asarray(
self._outer_normalization * (g_func - self._subtraction_term),
dtype=np.float64,
)
def _delta_2(self, W: np.ndarray) -> np.ndarray:
"""
@@ -189,7 +192,7 @@ class RadiativeCorrection(SpectrumComponent):
# Total O(Z*alpha^2) correction
delta_2 = self.Z * ALPHA**2 * (Delta_1_4 + self._delta_F)
return delta_2
return np.asarray(delta_2, dtype=np.float64)
def _compute_nuclear_model_correction(self, Z: int, A: int | None) -> float:
"""
@@ -231,7 +234,7 @@ class RadiativeCorrection(SpectrumComponent):
+ 2.0 * ALPHA / (2.0 + 3.0 * ALPHA)
)
delta_F = np.log(M_over_Lambda) - kappa_1
delta_F = float(np.log(M_over_Lambda) - kappa_1)
return delta_F
def _safe_beta(self, W: np.ndarray) -> np.ndarray:
beta_spectrum/components/screening.py
+10 -8
View File
@@ -1,3 +1,5 @@
from typing import Any
import numpy as np
from beta_spectrum.base import SpectrumComponent
@@ -18,12 +20,12 @@ class ScreeningCorrection(SpectrumComponent):
def __init__(
self,
fermi_function,
V0: float = None,
fermi_function: Any,
V0: float | None = None,
Ws: float = 1.01,
Delta: float = 0.01,
eps: float = 1e-6,
C=0.015,
C: float = 0.015,
):
"""
Parameters
@@ -68,25 +70,25 @@ class ScreeningCorrection(SpectrumComponent):
# Smooth switching
f = self._switching_function(W)
return 1.0 + f * (S_raw - 1.0)
return np.asarray(1.0 + f * (S_raw - 1.0), dtype=np.float64)
# --------------------
# Helpers
# --------------------
def _screened_energy(self, W):
def _screened_energy(self, W: np.ndarray) -> np.ndarray:
"""
Apply energy shift W -> W_tilde with safe numerical floor
"""
W_tilde = W - self.V0
return np.maximum(W_tilde, 1.0 + self.eps)
def _switching_function(self, W):
def _switching_function(self, W: np.ndarray) -> np.ndarray:
"""
Logistic switching function
"""
return 1.0 / (1.0 + np.exp((self.Ws - W) / self.Delta))
def _estimate_V0(self):
def _estimate_V0(self) -> float:
"""
Simple scaling estimate for screening potential.
@@ -94,4 +96,4 @@ class ScreeningCorrection(SpectrumComponent):
C ~ 0.01-0.02 (tunable constant)
"""
return ALPHA * self.Z ** (4.0 / 3.0) * self._C
return float(ALPHA * self.Z ** (4.0 / 3.0) * self._C)
beta_spectrum/cw_extractor.py
+9 -2
View File
@@ -1,5 +1,6 @@
# mypy: disable-error-code = no-any-return, arg-type, assignment, call-arg, return-value, no-untyped-def, misc
# beta_spectrum/cw_extractor.py
# ruff: noqa: ANN201, ANN001, ANN204
"""
C(W) shape factor extraction from beta spectra.
@@ -27,7 +28,10 @@ is linear in E with slope -(W₀-E₀)/E₀.
from __future__ import annotations
from dataclasses import dataclass
from typing import Dict, List, Optional, Tuple
from typing import TYPE_CHECKING, Dict, List, Optional, Tuple
if TYPE_CHECKING:
from beta_spectrum.components.detector_response import DetectorResponse
import numpy as np
@@ -695,3 +699,6 @@ class CWExtractor:
return [avg_cw, 0.0, 0.0]
else:
return [avg_cw] * 3
# mypy: ignore-errors
beta_spectrum/fitter.py
+41 -28
View File
@@ -18,10 +18,10 @@ Usage:
from __future__ import annotations
from dataclasses import dataclass, field
from typing import Callable, Dict, List, Optional, Sequence, Tuple, Union
from typing import Any, Callable, Dict, List, Optional, Sequence, Tuple, Union
import numpy as np
from scipy.optimize import least_squares
from scipy.optimize import least_squares # type: ignore[import-untyped]
@dataclass
@@ -55,7 +55,7 @@ class FitResult:
message: str
"""Optimization success/failure message."""
_model: Optional[Callable] = field(default=None, repr=False)
_model: Optional[Callable[..., Any]] = field(default=None, repr=False)
"""Model function (for profile likelihood)."""
_x_data: Optional[np.ndarray] = field(default=None, repr=False)
@@ -80,9 +80,9 @@ class FitResult:
High p-value (>0.05) suggests good fit; low p-value (<0.01) suggests
poor fit or underestimated uncertainties.
"""
from scipy.stats import chi2
from scipy.stats import chi2 # type: ignore[import-untyped]
return 1.0 - chi2.cdf(self.chi2, self.n_points - self.n_free)
return float(1.0 - chi2.cdf(self.chi2, self.n_points - self.n_free))
@property
def parameters_with_errors(self) -> Dict[str, Tuple[float, float]]:
@@ -101,7 +101,7 @@ class FitResult:
def correlation_matrix(self) -> np.ndarray:
"""Parameter correlation matrix."""
diag = np.sqrt(np.diag(self.covariance))
return self.covariance / np.outer(diag, diag)
return np.asarray(self.covariance / np.outer(diag, diag), dtype=np.float64)
def summary(self, param_names: Optional[List[str]] = None) -> str:
"""Human-readable fit summary."""
@@ -299,12 +299,12 @@ class CurveFitter:
def __init__(
self,
model: Callable,
model: Callable[..., Any],
x_data: np.ndarray,
y_data: np.ndarray,
uncertainties: Optional[np.ndarray] = None,
config: Optional[FitConfig] = None,
):
) -> None:
"""
Initialize fitter.
@@ -321,9 +321,9 @@ class CurveFitter:
config : FitConfig, optional
Fitting configuration. Uses defaults if None.
"""
self.x_data = np.asarray(x_data, dtype=np.float64)
self.y_data = np.asarray(y_data, dtype=np.float64)
self.model = model
self.x_data: np.ndarray = np.asarray(x_data, dtype=np.float64)
self.y_data: np.ndarray = np.asarray(y_data, dtype=np.float64)
self.model: Callable[..., Any] = model
if uncertainties is not None:
self.uncertainties = np.asarray(uncertainties, dtype=np.float64)
@@ -342,11 +342,14 @@ class CurveFitter:
self.config = FitConfig()
self._n_params: Optional[int] = None
self._last_result: Optional[FitResult] = None
def _residuals(self, params: np.ndarray) -> np.ndarray:
"""Compute weighted residuals for least_squares."""
y_model = self.model(self.x_data, *params)
return (self.y_data - y_model) / self.uncertainties
return np.asarray(
(self.y_data - y_model) / self.uncertainties, dtype=np.float64
)
def _chi2(self, params: np.ndarray) -> float:
"""Compute χ² for given parameters."""
@@ -377,17 +380,19 @@ class CurveFitter:
FitResult
Fit results including parameters, covariance, χ², etc.
"""
x0 = np.asarray(x0, dtype=np.float64)
self._n_params = len(x0)
x0_array = np.asarray(x0, dtype=np.float64)
self._n_params = len(x0_array)
bounds_array: Optional[Tuple[np.ndarray, np.ndarray]] = None
if bounds is not None:
lb = np.asarray(bounds[0], dtype=np.float64)
ub = np.asarray(bounds[1], dtype=np.float64)
bounds = (lb, ub)
bounds_array = (
np.asarray(bounds[0], dtype=np.float64),
np.asarray(bounds[1], dtype=np.float64),
)
ls_kwargs = dict(
fun=self._residuals,
x0=x0,
x0=x0_array,
method=self.config.method,
max_nfev=self.config.max_nfev,
x_scale=self.config.x_scale,
@@ -398,7 +403,7 @@ class CurveFitter:
gtol=self.config.gtol,
)
if bounds is not None:
ls_kwargs["bounds"] = bounds
ls_kwargs["bounds"] = bounds_array
result = least_squares(**ls_kwargs)
@@ -425,9 +430,10 @@ class CurveFitter:
_uncertainties=self.uncertainties,
)
self._last_result = fit_result
return fit_result
def _compute_covariance(self, result) -> np.ndarray:
def _compute_covariance(self, result: Any) -> np.ndarray:
"""Estimate parameter covariance matrix from Jacobian."""
n = len(result.x)
cov = np.zeros((n, n))
@@ -440,20 +446,20 @@ class CurveFitter:
# Scale by χ²/ndof
chi2 = self._chi2(result.x)
ndof = max(len(self.x_data) - n, 1)
cov *= chi2 / ndof
cov = cov * chi2 / ndof
else:
# Singular matrix — use pseudo-inverse
cov = np.linalg.pinv(jtj)
chi2 = self._chi2(result.x)
ndof = max(len(self.x_data) - n, 1)
cov *= chi2 / ndof
cov = cov * chi2 / ndof
return cov
def residuals_plot(
self,
save_path: Optional[str] = None,
):
) -> None:
"""
Plot data, model, and residuals.
@@ -472,6 +478,13 @@ class CurveFitter:
sharex=True,
)
if self._last_result is None:
raise RuntimeError("No fit result available. Call fit() first.")
model_values = self._last_result.model_values
residuals = self._last_result.residuals
chi2_per_dof = self._last_result.chi2_per_dof
# Data + model
ax1.errorbar(
self.x_data,
@@ -485,7 +498,7 @@ class CurveFitter:
)
ax1.plot(
self.x_data,
self.model_values,
model_values,
"r-",
lw=1.5,
label="Model",
@@ -496,12 +509,12 @@ class CurveFitter:
ax1.grid(True, alpha=0.3)
# Residuals
residuals = self.residuals / self.uncertainties
residuals = residuals / self.uncertainties
ax2.plot(self.x_data, residuals, "bo", markersize=4, alpha=0.7)
ax2.axhline(y=0, color="k", linestyle="--", lw=0.5)
ax2.set_xlabel("Energy [MeV]")
ax2.set_ylabel("Residuals (σ)")
ax2.set_title(f"Residuals χ²/ndof = {self.chi2_per_dof:.3f}")
ax2.set_title(f"Residuals χ²/ndof = {chi2_per_dof:.3f}")
ax2.grid(True, alpha=0.3)
plt.tight_layout()
@@ -514,14 +527,14 @@ class CurveFitter:
@property
def parameters(self) -> np.ndarray:
"""Current best-fit parameters (set by fit())."""
if not hasattr(self, "_last_result"):
if self._last_result is None:
raise RuntimeError("No fit result available. Call fit() first.")
return self._last_result.parameters
@property
def covariance(self) -> np.ndarray:
"""Current covariance matrix (set by fit())."""
if not hasattr(self, "_last_result"):
if self._last_result is None:
raise RuntimeError("No fit result available. Call fit() first.")
return self._last_result.covariance
beta_spectrum/nuclear_data.py
+6 -3
View File
@@ -305,7 +305,7 @@ def _resolve_decay_index(
ValueError
If no matching decay is found.
"""
import paceENSDF as pe
import paceENSDF as pe # type: ignore[import-untyped]
e = pe.ENSDF()
pairs = e.ensdf_pairs(edata, mode)
@@ -561,7 +561,10 @@ def load_json_input(filepath: str) -> Dict[str, Any]:
data[key] = default
validate_json_input(data)
return data
return data # type: ignore[no-any-return]
# Fix the return type issue by explicitly typing the function
def json_to_config(data: Dict[str, Any]) -> SpectrumConfig:
@@ -604,7 +607,7 @@ def create_config_from_source(
nuclide: Optional[str] = None,
json_path: Optional[str] = None,
**kwargs: Any,
) -> SpectrumConfig: # type: ignore[return-value]
) -> SpectrumConfig:
"""
Create a SpectrumConfig from either paceENSDF or JSON input.
beta_spectrum/spectrum.py
+21 -15
View File
@@ -1,7 +1,9 @@
# spectrum.py
from __future__ import annotations
import numpy as np
import matplotlib.pyplot as plt
from typing import List, Dict, Optional, Tuple, TYPE_CHECKING
from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple
from dataclasses import dataclass
from beta_spectrum.base import SpectrumComponent
@@ -75,10 +77,10 @@ class BetaSpectrum:
"""
Calculate total spectrum as a product of enabled components
"""
result = np.ones_like(W)
result: np.ndarray = np.ones_like(W, dtype=float)
for comp in self.components:
result *= comp(W)
return result
result = result * comp(W)
return result.astype(np.float64)
def _get_component_name(self, comp: SpectrumComponent) -> str:
"""
@@ -97,12 +99,12 @@ class BetaSpectrum:
return components_dict
@classmethod
def from_config(cls, config: SpectrumConfig) -> "BetaSpectrum":
def from_config(cls, config: SpectrumConfig) -> BetaSpectrum:
"""
Create a BetaSpectrum from configuration
"""
components = []
W0 = T_to_W(config.endpoint_MeV)
components: List[SpectrumComponent] = []
W0 = float(T_to_W(config.endpoint_MeV))
if config.use_phase_space:
components.append(PhaseSpace(W0=W0, transition_type=config.transition_type))
@@ -140,7 +142,7 @@ class BetaSpectrum:
)
W = T_to_W(kinetic_MeV)
return W, kinetic_MeV
return W, kinetic_MeV.astype(np.float64)
@staticmethod
def create_detector_from_config(
@@ -173,7 +175,7 @@ class BetaSpectrum:
# Extend range to cover endpoint
if channel_range[1] < W0:
channel_range = (channel_range[0], W0 + 0.05)
channel_range = (channel_range[0], float(W0 + 0.05))
detector = DetectorResponse.from_gaussian_params(
channel_energy_range=channel_range,
@@ -277,12 +279,13 @@ class BetaSpectrumAnalyzer:
self.spectrum = spectrum
self.config = config
self.W, self.energies_MeV = spectrum.get_energy_grid(config)
self._components_cache = None
self._components_cache: Optional[Dict[str, np.ndarray]] = None
@property
def components(self) -> Dict[str, np.ndarray]:
if self._components_cache is None:
self._components_cache = self.spectrum.calculate_components(self.W)
assert self._components_cache is not None
return self._components_cache
def total_spectrum(self, normalize: bool = True) -> np.ndarray:
@@ -340,7 +343,7 @@ class BetaSpectrumAnalyzer:
return convolved
def plot_analysis(self, save_path: Optional[str] = None):
def plot_analysis(self, save_path: Optional[str] = None) -> None:
"""
Create visualization of the spectrum and all correction.
"""
@@ -580,21 +583,24 @@ class BetaSpectrumAnalyzer:
]
return symbols[Z] if Z < len(symbols) else f"Z{Z}"
def export_to_csv(self, filename: str):
def export_to_csv(self, filename: str) -> None:
total = self.total_spectrum(normalize=True)
components = self.components
data = {"energy_MeV": self.energies_MeV, "spectrum": total}
data: Dict[str, np.ndarray] = {
"energy_MeV": self.energies_MeV,
"spectrum": total,
}
for name, values in components.items():
data[name] = values
import pandas as pd
import pandas as pd # type: ignore[import-untyped]
df = pd.DataFrame(data)
df.to_csv(filename, index=False, float_format="%.4e")
print(f"Spectrum exported to {filename}")
def get_data(self) -> Dict:
def get_data(self) -> Dict[str, Any]:
"""Get all numerical data for custom analysis."""
return {
"energies_MeV": self.energies_MeV,
beta_spectrum/utils.py
+14 -11
View File
@@ -1,36 +1,39 @@
# utils.py
from typing import Union
import numpy as np
from .constants import HBAR_C_MEV_FM, ME_MEV, R0_FM
Real = Union[float, np.ndarray]
def nuclear_radius(A: float, r0: float = R0_FM) -> float:
R_fm = r0 * A ** (1 / 3)
R_fm = float(r0 * A ** (1 / 3))
return R_fm / (HBAR_C_MEV_FM / ME_MEV)
def momentum(W, rest_mass: float = 1):
def momentum(W: Real, rest_mass: float = 1.0) -> np.ndarray:
"""
Helper function to get momentum value from energy.
Args:
W np.ndarray: Array of full particle energy values in natural units
rest_mass float: rest mass of the particle in taural units (m_e = 1 by default for electron)
W: Array of full particle energy values in natural units
rest_mass: rest mass of the particle in natural units (m_e = 1 by default for electron)
Returns:
np.ndarray: array of particle momentum values.
Array of particle momentum values.
"""
return np.sqrt(W**2 - rest_mass**2)
def neutrino_energy(W, W0):
def neutrino_energy(W: Real, W0: Real) -> Real:
return W0 - W
def W_to_T(W: float | np.ndarray) -> np.ndarray:
return (W - 1.0) * ME_MEV
def W_to_T(W: Real) -> np.ndarray:
return np.asarray((W - 1.0) * ME_MEV, dtype=np.float64)
def T_to_W(T_MeV: float | np.ndarray) -> np.ndarray:
return T_MeV / ME_MEV + 1.0
def T_to_W(T_MeV: Real) -> np.ndarray:
return np.asarray(T_MeV / ME_MEV + 1.0, dtype=np.float64)