"""
Query module for weight window data retrieval.
Provides the QueryResult class for structured access to weight window query results.
"""
from dataclasses import dataclass
from typing import List, Tuple
import numpy as np
import pandas as pd
from wwinpy.header import Header
[docs]
@dataclass
class QueryResult:
"""Store the results of a weight window query.
:param header: Weight window file header information
:type header: Header
:param particle_types: List of particle type identifiers
:type particle_types: List[int]
:param ww_values: List of weight window values arrays, one per particle type
:type ww_values: List[np.ndarray]
:param energy_intervals: List of energy interval pairs (starts, ends) for each particle type
:type energy_intervals: List[Tuple[np.ndarray, np.ndarray]]
:param time_intervals: List of time interval pairs (starts, ends) for each particle type
:type time_intervals: List[Tuple[np.ndarray, np.ndarray]]
:param x_intervals: Spatial interval pairs (starts, ends) for x-direction
:type x_intervals: Tuple[np.ndarray, np.ndarray]
:param y_intervals: Spatial interval pairs (starts, ends) for y-direction
:type y_intervals: Tuple[np.ndarray, np.ndarray]
:param z_intervals: Spatial interval pairs (starts, ends) for z-direction
:type z_intervals: Tuple[np.ndarray, np.ndarray]
"""
header: Header
particle_types: List[int]
ww_values: List[np.ndarray]
energy_intervals: List[Tuple[np.ndarray, np.ndarray]]
time_intervals: List[Tuple[np.ndarray, np.ndarray]]
x_intervals: Tuple[np.ndarray, np.ndarray]
y_intervals: Tuple[np.ndarray, np.ndarray]
z_intervals: Tuple[np.ndarray, np.ndarray]
[docs]
def to_dataframe(self) -> pd.DataFrame:
"""Convert query results to a pandas DataFrame.
:return: DataFrame containing weight window data with particle type, spatial,
energy, and time information
:rtype: pd.DataFrame
"""
data_rows = []
for p_idx, p_type in enumerate(self.particle_types):
e_starts, e_ends = self.energy_intervals[p_idx]
t_starts, t_ends = self.time_intervals[p_idx]
x_starts, x_ends = self.x_intervals
y_starts, y_ends = self.y_intervals
z_starts, z_ends = self.z_intervals
ww_vals = self.ww_values[p_idx]
if len(t_starts) == 0 and len(t_ends) == 0:
t_starts = np.array([0.0])
t_ends = np.array([np.inf])
# Get array shape
time_dim = len(t_starts)
for t_idx, (t_start, t_end) in enumerate(zip(t_starts, t_ends)):
for e_idx, (e_start, e_end) in enumerate(zip(e_starts, e_ends)):
for z_idx, (z_start, z_end) in enumerate(zip(z_starts, z_ends)):
for y_idx, (y_start, y_end) in enumerate(zip(y_starts, y_ends)):
for x_idx, (x_start, x_end) in enumerate(zip(x_starts, x_ends)):
t_index = t_idx if time_dim > 1 else 0
# Determine the appropriate energy index
if ww_vals.shape[1] == 1:
# Only one energy bin available, use e_idx=0
effective_e_idx = 0
else:
# Multiple energy bins available
effective_e_idx = e_idx
# Additional safety check
if e_idx >= ww_vals.shape[1]:
raise IndexError(
f"Energy index {e_idx} out of bounds for ww_values with shape {ww_vals.shape}"
)
try:
ww_value = float(ww_vals[t_index, effective_e_idx, z_idx, y_idx, x_idx])
except IndexError as ie:
raise IndexError(
f"Error accessing ww_values at indices "
f"(t={t_index}, e={effective_e_idx}, z={z_idx}, y={y_idx}, x={x_idx}): {ie}"
)
data_rows.append({
'particle_type': p_type,
'time_start': t_start,
'time_end': t_end,
'energy_start': e_start,
'energy_end': e_end,
'x_start': x_start,
'x_end': x_end,
'y_start': y_start,
'y_end': y_end,
'z_start': z_start,
'z_end': z_end,
'ww_value': ww_value
})
df = pd.DataFrame(data_rows)
# Format energy columns to scientific notation with 5 decimal places
df['energy_start'] = df['energy_start'].apply(lambda x: '{:.5e}'.format(x))
df['energy_end'] = df['energy_end'].apply(lambda x: '{:.5e}'.format(x))
df['ww_value'] = df['ww_value'].apply(lambda x: '{:.5e}'.format(x))
return df