Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions .github/workflows/build_docs.yml
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,7 @@ on:

jobs:
deploy:
runs-on: ubuntu-20.04
runs-on: ubuntu-latest
permissions:
contents: write
concurrency:
Expand All @@ -33,7 +33,7 @@ jobs:
poetry run make github
- name: Deploy
uses: peaceiris/actions-gh-pages@v3
if: ${{ github.ref == 'refs/heads/master' }}
if: ${{ github.ref == 'refs/heads/main' }}
with:
github_token: ${{ secrets.GITHUB_TOKEN }}
publish_dir: ./docs
2 changes: 1 addition & 1 deletion ParticlePhaseSpace/DataLoaders.py
Original file line number Diff line number Diff line change
Expand Up @@ -437,7 +437,7 @@ def _import_data(self):
self.data[self._columns['particle type']] = particle_types_pdg
self.data[self._columns['particle id']] = np.arange(
len(self.data)) # may want to replace with track ID if available?
self.data[self._columns['time']] = pd.Series(0 * np.ones(self.data.shape[0]), dtype="category") # may want to replace with time feature if available?
self.data[self._columns['time']] = pd.Series(0 * np.ones(self.data.shape[0]))
# figure out the momentums:
DirCosineX = data['Cosine X']
DirCosineY = data['Cosine Y']
Expand Down
240 changes: 240 additions & 0 deletions ParticlePhaseSpace/_ParticlePhaseSpace.py
Original file line number Diff line number Diff line change
Expand Up @@ -540,6 +540,193 @@ def n_particles_v_time(self, n_bins: int = 100, grid: bool = False): # pragma:
plt.grid()
plt.tight_layout()

def radial_energy_boxplot(self, n_bins=5, beam_direction='z', rlim=None):
"""
Plots boxplots of kinetic energy distribution per radial bin

:param n_bins: Number of radial bins (boxplots)
:type n_bins: int
:param beam_direction: Direction of beam ('z', 'x', or 'y'); default 'z'
:type beam_direction: str
:param rlim: Tuple (r_min, r_max); if None, uses full range
:type rlim: list or None
"""
box_alpha = 0.7
box_color = 'dodgerblue'

ps_data = self._PS._ps_data

# Compute radial coordinate
if beam_direction == 'z':
x = ps_data[self._PS.columns['x']]
y = ps_data[self._PS.columns['y']]
r = np.sqrt(x ** 2 + y ** 2)
elif beam_direction == 'x':
y = ps_data[self._PS.columns['y']]
z = ps_data[self._PS.columns['z']]
r = np.sqrt(y ** 2 + z ** 2)
elif beam_direction == 'y':
x = ps_data[self._PS.columns['x']]
z = ps_data[self._PS.columns['z']]
r = np.sqrt(x ** 2 + z ** 2)
else:
raise ValueError("beam_direction must be 'x', 'y', or 'z'")

# Set r limits
if rlim is None:
r_min, r_max = r.min(), r.max()
else:
r_min, r_max = rlim

# Bin edges and bin centers
r_edges = np.linspace(r_min, r_max, n_bins + 1)
r_centers = 0.5 * (r_edges[:-1] + r_edges[1:])

# Get kinetic energy
if not self._PS.columns['Ek'] in ps_data.columns:
self._PS.fill.kinetic_E()
ps_data = self._PS._ps_data
ek = ps_data[self._PS.columns['Ek']]

# Collect energy per bin
energy_bins = []
for i in range(n_bins):
in_bin = (r >= r_edges[i]) & (r < r_edges[i + 1])
energy_bins.append(ek[in_bin])

# Prepare boxplot data (remove empty bins)
boxplot_data = [eb.values for eb in energy_bins if len(eb) > 0]
r_centers_plot = [r_centers[i] for i, eb in enumerate(energy_bins) if len(eb) > 0]

fig, ax = plt.subplots(figsize=(10, 6))
bp = ax.boxplot(
boxplot_data,
positions=r_centers_plot,
widths=(r_edges[1] - r_edges[0]) * 0.7,
patch_artist=True,
showfliers=False, # Hide outliers for cleaner look
boxprops=dict(facecolor=box_color, color='navy', alpha=box_alpha, linewidth=1.5),
whiskerprops=dict(color='navy', linewidth=1.5),
capprops=dict(color='navy', linewidth=1.5),
medianprops=dict(color='orange', linewidth=2)
)

ax.set_xlabel(f"r [{self._PS._units.length.label}]", fontsize=14)
ax.set_ylabel(f"Kinetic Energy [{self._PS._units.energy.label}]", fontsize=14)
ax.set_title("Energy Distribution In Radial Bins", fontsize=16)
ax.set_xlim(r_min, r_max)
ax.grid(True, linestyle='--', alpha=0.5)
plt.tight_layout()
plt.show()

def fluence_map_2D(self, beam_direction: str = 'z', at: float = 0.0, quantity: str = 'energy',
bins: int = 100, normalize: bool = False,
xlim=None, ylim=None):
"""
Plot fluence (particle or energy per unit area) as a 2D map in the plane orthogonal to beam_direction.

The phase space data is firt projected to the "at" coordinate. any data which would never make it there is discarded.

:param beam_direction: Direction of beam travel ('x', 'y', or 'z')
:type beam_direction: str, optional
:param at: Position of plane along beam direction
:type at: float, optional
:param quantity: 'particle' for particle fluence, 'energy' for energy fluence
:type quantity: str, optional
:param bins: Number of bins per axis
:type bins: int, optional
:param normalize: If True, normalize fluence map to max value
:type normalize: bool, optional
:param xlim: x-axis limits (tuple or list)
:type xlim: tuple or list, optional
:param ylim: y-axis limits (tuple or list)
:type ylim: tuple or list, optional
:return: fluence_map (2D numpy array), (xedges, yedges)
"""

# Project all particles to the plane
plot_ps = self._PS.transform.project_to_plane(beam_direction=beam_direction, at=at)
ps_data = plot_ps._ps_data
# filter any particles not at the requested plane (travelling in wrong direction)
if beam_direction == 'x':
coord = ps_data[plot_ps.columns['x']]
elif beam_direction == 'y':
coord = ps_data[plot_ps.columns['y']]
elif beam_direction == 'z':
coord = ps_data[plot_ps.columns['z']]
else:
raise ValueError(f"beam direction: {beam_direction} is not supported.")

tol = 1e-6
mask = np.abs(coord - at) < tol
# Filter for particles actually "on" the plane
ps_data = ps_data[mask]
columns = plot_ps.columns
units = plot_ps._units

# Determine orthogonal axes
if beam_direction == 'x':
axis1, axis2 = columns['y'], columns['z']
elif beam_direction == 'y':
axis1, axis2 = columns['x'], columns['z']
elif beam_direction == 'z':
axis1, axis2 = columns['x'], columns['y']
else:
raise ValueError("beam_direction must be one of 'x', 'y', or 'z'")

x = ps_data[axis1]
y = ps_data[axis2]

# Set limits if not specified
if xlim is None:
# Default to percentiles to avoid extreme outliers
xlim = (np.percentile(x, 1), np.percentile(x, 99))
if ylim is None:
ylim = (np.percentile(y, 1), np.percentile(y, 99))

xedges = np.linspace(xlim[0], xlim[1], bins + 1)
yedges = np.linspace(ylim[0], ylim[1], bins + 1)

# Mask out particles outside bounds to avoid weird bins
mask = (x >= xlim[0]) & (x <= xlim[1]) & (y >= ylim[0]) & (y <= ylim[1])
x = x[mask]
y = y[mask]
weights = ps_data['weight'][mask]

if quantity == 'energy':
plot_ps.fill.kinetic_E() # Fill on projected phase space
ps_data = plot_ps._ps_data
if isinstance(weights.dtype, pd.CategoricalDtype):
weights = weights.astype(float)
weights = weights * ps_data[columns['Ek']][mask]
elif quantity != 'particle':
raise ValueError("quantity must be 'particle' or 'energy'")

H, _, _ = np.histogram2d(x, y, bins=[xedges, yedges], weights=weights)
bin_area = (xedges[1] - xedges[0]) * (yedges[1] - yedges[0])
fluence_map = H / bin_area

if normalize and fluence_map.max() > 0:
fluence_map = fluence_map / fluence_map.max()

plt.figure(figsize=(8, 6))
plt.imshow(fluence_map.T, origin='lower', aspect='auto',
extent=[xedges[0], xedges[-1], yedges[0], yedges[-1]],
cmap='inferno')
if quantity == "energy":
plt.colorbar(label=f"Energy fluence [{units.energy.label}/{units.length.label}$^2$]")
else:
plt.colorbar(label=f"Particle fluence [particles/{units.length.label}$^2$]")
plt.xlabel(axis1)
plt.ylabel(axis2)
plt.title(
f"{quantity.capitalize()} fluence map in plane orthogonal to {beam_direction} @ {at}")
plt.xlim(xlim)
plt.ylim(ylim)
plt.tight_layout()
plt.show()



class _Transform(_PhaseSpace_MethodHolder):
"""
Expand Down Expand Up @@ -687,6 +874,59 @@ def project(self, direction: str = 'z', distance: float = 100, in_place: bool =
new_instance = self._return_position_update(new_x, new_y, new_z, in_place)
return new_instance # nb: None if in_place = True

def project_to_plane(self, beam_direction: str = 'z', at: float = 0.0, in_place: bool = False):
"""
Projects all particles to the specified plane (e.g., z=at, y=at, x=at) along their trajectories.
Assumes straight-line motion.
If a particle would NEVER reach the requested plane, it will keep it's oroginal positions.

:param beam_direction: Direction normal to the plane ('x', 'y', or 'z')
:type beam_direction: str
:param at: Position of plane along beam_direction
:type at: float
:param in_place: If True, modify in place. If False, return new PhaseSpace
:type in_place: bool
:return: new PhaseSpace object with all particles on the plane
"""
# Get required columns
col = self._PS.columns
psd = self._PS._ps_data

for axis in [col['x'], col['y'], col['z'], col['px'], col['py'], col['pz']]:
if isinstance(psd[axis].dtype, pd.CategoricalDtype):
psd[axis] = psd[axis].cat.codes.astype(float)

# Calculate displacement needed to bring each particle to the plane
if beam_direction == 'x':
d = at - psd[col['x']]
# Avoid divide by zero
px = psd[col['px']].replace(0, np.nan)
new_y = psd[col['y']] + d * psd[col['py']] / px
new_z = psd[col['z']] + d * psd[col['pz']] / px
new_x = np.full_like(psd[col['x']], at)
elif beam_direction == 'y':
d = at - psd[col['y']]
py = psd[col['py']].replace(0, np.nan)
new_x = psd[col['x']] + d * psd[col['px']] / py
new_z = psd[col['z']] + d * psd[col['pz']] / py
new_y = np.full_like(psd[col['y']], at)
elif beam_direction == 'z':
d = at - psd[col['z']] # projection amount
pz = psd[col['pz']].replace(0, np.nan)
new_x = psd[col['x']] + d * psd[col['px']] / pz
new_y = psd[col['y']] + d * psd[col['py']] / pz
new_z = np.full_like(psd[col['z']], at)
else:
raise ValueError("beam_direction must be one of 'x', 'y', or 'z'")

# Replace any nan (from divide by zero) with current position
new_x = np.where(np.isnan(new_x), psd[col['x']], new_x)
new_y = np.where(np.isnan(new_y), psd[col['y']], new_y)
new_z = np.where(np.isnan(new_z), psd[col['z']], new_z)

# Use existing _return_position_update for in_place or return new
return self._return_position_update(new_x, new_y, new_z, in_place) # nb: None if in_place = True

def regrid(self, quantities: (list, None) = None, n_bins: (int, list) = 10, in_place=False):
"""
this re-grids each quantity in quantities onto a new grid. The new grid is defined by
Expand Down
1 change: 1 addition & 0 deletions docsrc/examples.rst
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@ Examples
/new_data_loader
/IAEA
/new_data_exporter
/plotting
/transform
/units
/grid_merge
Expand Down
388 changes: 388 additions & 0 deletions docsrc/plotting.ipynb

Large diffs are not rendered by default.