Skip to content
GitLab
Projects
Groups
Snippets
/
Help
Help
Support
Community forum
Keyboard shortcuts
?
Submit feedback
Contribute to GitLab
Sign in
Toggle navigation
Menu
Open sidebar
Connor Hainje
pidanalysis
Commits
c2a6aa22
Commit
c2a6aa22
authored
Dec 07, 2021
by
connor.hainje@pnnl.gov
Browse files
Completely package and refactor contribution plotting methods
parent
236dc4f2
Changes
2
Hide whitespace changes
Inline
Side-by-side
packages/pidplots/pidplots/contribution.py
0 → 100644
View file @
c2a6aa22
import
numpy
as
np
import
matplotlib.pyplot
as
plt
from
.
import
const
,
colors
from
.
import
distributions
as
dist
from
.utils
import
(
_ax_xlabel
,
connected_region_borders
,
imshow_2D_axis_labels
,
outer_legend
,
better_xlabel
,
text_array_values
,
xticks_radians_to_degrees
,
)
def
subplot_contribution_distribution
(
ax
,
df
,
color_list
=
None
,
det
=
"max"
):
color_list
=
colors
.
fill_color_list
(
color_list
)
contrib_cols
=
[
f
"contrib_
{
det
}
"
for
det
in
const
.
DETECTORS
]
if
det
==
"all"
:
dist
.
all
(
ax
,
df
[
contrib_cols
].
values
,
labels
=
const
.
DETECTORS
,
colors
=
color_list
,
bin_lims
=
(
-
1
,
1
),
)
elif
det
==
"min"
:
dist
.
smallest
(
ax
,
df
[
contrib_cols
].
values
,
labels
=
const
.
DETECTORS
,
colors
=
color_list
,
bin_lims
=
(
-
1
,
1
),
)
elif
det
==
"max"
:
dist
.
largest
(
ax
,
df
[
contrib_cols
].
values
,
labels
=
const
.
DETECTORS
,
colors
=
color_list
,
bin_lims
=
(
-
1
,
1
),
)
elif
det
in
const
.
DETECTORS
:
dist
.
single
(
ax
,
df
[
f
"contrib_
{
det
}
"
].
values
,
label
=
det
,
color
=
color_list
[
const
.
DETECTORS
.
index
(
det
)],
bin_lims
=
(
-
1
,
1
),
)
else
:
raise
ValueError
(
"`det` not understood."
)
def
plot_contribution_distribution
(
df
,
det
=
"max"
,
color_list
=
None
,
yscale
=
"log"
,
figsize
=
None
):
fig
,
ax
=
plt
.
subplots
(
figsize
=
figsize
)
subplot_contribution_distribution
(
ax
,
df
,
det
=
det
,
color_list
=
color_list
)
ax
.
set_yscale
(
yscale
)
if
det
==
"all"
:
xlabel
=
"Detector contribution"
elif
det
==
"min"
:
xlabel
=
"Smallest detector contribution"
elif
det
==
"max"
:
xlabel
=
"Largest detector contribution"
elif
det
in
const
.
DETECTORS
:
xlabel
=
f
"
{
det
}
contribution"
ax
.
set_xlabel
(
xlabel
)
fig
.
tight_layout
()
if
det
not
in
const
.
DETECTORS
:
outer_legend
(
fig
,
loc
=
"bottom"
,
y
=
0.02
)
return
fig
def
subplot_avg_contribution
(
ax
,
df
,
color_list
=
None
):
ctrbs
=
df
[[
f
"contrib_
{
det
}
"
for
det
in
const
.
DETECTORS
]].
values
ctrbs
=
np
.
mean
(
ctrbs
,
axis
=
0
)
x_vals
=
np
.
arange
(
len
(
const
.
DETECTORS
))
bars
=
ax
.
bar
(
x_vals
,
ctrbs
,
width
=
0.8
)
color_list
=
colors
.
fill_color_list
(
color_list
)
for
i
in
range
(
len
(
const
.
DETECTORS
)):
bars
[
i
].
set_facecolor
(
color_list
[
i
])
ax
.
set_xticks
(
x_vals
)
ax
.
set_xticklabels
(
const
.
DETECTORS
)
ax
.
set_xlim
(
x_vals
[
0
]
-
0.5
,
x_vals
[
-
1
]
+
0.5
)
def
plot_avg_contribution
(
df
,
title
=
None
,
color_list
=
None
,
figsize
=
None
):
fig
,
ax
=
plt
.
subplots
(
figsize
=
figsize
)
subplot_avg_contribution
(
ax
,
df
,
color_list
=
color_list
)
ax
.
set_xlabel
(
"Detector"
)
ax
.
set_ylabel
(
"Average contribution value"
)
ax
.
set_title
(
title
)
fig
.
tight_layout
()
return
fig
def
subplot_avg_contribution_1D
(
ax
,
df
,
bin_by
=
"p"
,
ls
=
"-"
,
color_list
=
None
):
if
bin_by
==
"p"
:
ctrbs
=
df
.
groupby
([
"p_bin"
]).
mean
()
bin_centers
=
const
.
P_BIN_CENTERS
elif
bin_by
==
"theta"
:
ctrbs
=
df
.
groupby
([
"theta_bin"
]).
mean
()
bin_centers
=
const
.
THETA_BIN_CENTERS
color_list
=
colors
.
fill_color_list
(
color_list
)
ctrbs
=
ctrbs
[[
f
"contrib_
{
det
}
"
for
det
in
const
.
DETECTORS
]]
for
i
,
det
in
enumerate
(
const
.
DETECTORS
):
vals
=
np
.
zeros
(
len
(
bin_centers
))
vals
[
ctrbs
.
index
]
=
ctrbs
[
f
"contrib_
{
det
}
"
]
ax
.
plot
(
bin_centers
,
vals
,
ls
=
ls
,
c
=
color_list
[
i
],
label
=
det
)
def
plot_avg_contribution_1D
(
df
,
bin_by
=
"p"
,
title
=
None
,
figsize
=
None
,
color_list
=
None
):
fig
,
ax
=
plt
.
subplots
(
figsize
=
figsize
)
subplot_avg_contribution_1D
(
ax
,
df
,
bin_by
=
bin_by
,
color_list
=
color_list
)
fig
.
tight_layout
()
_ax_xlabel
(
ax
,
bin_by
)
if
bin_by
==
"theta"
:
xticks_radians_to_degrees
(
ax
)
ax
.
set_ylabel
(
"Average contribution value"
)
outer_legend
(
fig
,
loc
=
"top"
,
ncol
=
len
(
const
.
DETECTORS
),
y
=
0.98
)
better_xlabel
(
fig
,
title
,
loc
=
"top"
,
y
=
1.08
)
return
fig
def
subplot_avg_contribution_2D
(
ax
,
df
,
det
=
"max"
,
color_list
=
None
,
cell_fontsize
=
10
):
color_list
=
colors
.
fill_color_list
(
color_list
)
gb
=
df
.
groupby
([
"p_bin"
,
"theta_bin"
])
# handle single detector avg contribution vs p, theta
if
det
in
const
.
DETECTORS
:
cmap
=
colors
.
make_linear_colormap
(
color_list
[
const
.
DETECTORS
.
index
(
det
)])
val
=
np
.
zeros
((
const
.
N_P_BINS
,
const
.
N_THETA_BINS
))
for
(
i
,
j
),
g
in
gb
:
val
[
i
,
j
]
=
g
.
mean
()[
f
"contrib_
{
det
}
"
]
im
=
ax
.
imshow
(
val
,
cmap
=
cmap
,
aspect
=
"auto"
)
# add cell values
if
cell_fontsize
>
0
:
text_array_values
(
ax
,
val
,
fontsize
=
cell_fontsize
)
# handle min/max avg contributions vs p, theta
elif
det
==
"min"
or
det
==
"max"
:
cmap
=
colors
.
make_colormap
(
color_list
)
all_vals
=
np
.
zeros
((
const
.
N_P_BINS
,
const
.
N_THETA_BINS
,
len
(
const
.
DETECTORS
)))
for
(
i
,
j
),
g
in
gb
:
all_vals
[
i
,
j
]
=
g
.
mean
()[
[
f
"contrib_
{
det
}
"
for
det
in
const
.
DETECTORS
]
].
values
if
det
==
"min"
:
idx
=
np
.
argmin
(
all_vals
,
axis
=
2
)
else
:
idx
=
np
.
argmax
(
all_vals
,
axis
=
2
)
# TODO: make sure vlims are correct
im
=
ax
.
imshow
(
idx
,
cmap
=
cmap
,
vmin
=-
0.5
,
vmax
=
len
(
const
.
DETECTORS
)
-
0.5
,
aspect
=
"auto"
)
# add cell values
val
=
np
.
take_along_axis
(
all_vals
,
np
.
expand_dims
(
idx
,
2
),
axis
=
2
).
squeeze
()
if
cell_fontsize
>
0
:
text_array_values
(
ax
,
val
,
fontsize
=
cell_fontsize
)
connected_region_borders
(
ax
,
idx
)
else
:
raise
ValueError
(
"Value given for argument 'max' not understood."
)
imshow_2D_axis_labels
(
ax
)
return
im
def
plot_avg_contribution_2D
(
df
,
det
=
"max"
,
figsize
=
None
,
color_list
=
None
,
cell_fontsize
=
10
):
fig
,
ax
=
plt
.
subplots
(
figsize
=
figsize
)
color_list
=
colors
.
fill_color_list
(
color_list
)
im
=
subplot_avg_contribution_2D
(
ax
,
df
,
color_list
=
color_list
,
det
=
det
,
cell_fontsize
=
cell_fontsize
)
fig
.
tight_layout
()
if
det
in
const
.
DETECTORS
:
cb
=
fig
.
colorbar
(
im
)
cb
.
set_label
(
f
"
{
det
}
average contribution value"
)
else
:
from
matplotlib.patches
import
Patch
handles
=
[
Patch
(
color
=
color_list
[
i
])
for
i
in
range
(
len
(
const
.
DETECTORS
))]
title
=
(
"Largest"
if
det
==
"max"
else
"Smallest"
)
+
"
\n
contributor"
outer_legend
(
fig
,
handles
,
const
.
DETECTORS
,
loc
=
"right"
,
ncol
=
1
,
title
=
title
,
x
=
0.98
,
)
return
fig
packages/pidplots/pidplots/distributions.py
0 → 100644
View file @
c2a6aa22
import
numpy
as
np
# There are a few different ways that we can plot the distributions
# We can plot histograms of...
# - all distributions
# - a single distribution
# - the largest/smallest (which we'll split by distribution of origin)
def
all
(
ax
,
data
,
bin_lims
=
(
0
,
1
),
n_bins
=
30
,
labels
=
None
,
colors
=
None
):
bins
=
np
.
linspace
(
*
bin_lims
,
num
=
n_bins
+
1
)
for
i
in
range
(
data
.
shape
[
1
]):
label
=
labels
[
i
]
if
labels
else
""
color
=
colors
[
i
]
if
colors
else
None
ax
.
hist
(
data
[:,
i
],
bins
=
bins
,
histtype
=
"step"
,
label
=
label
,
color
=
color
)
# single(ax, data, i, bin_lims=bin_lims, n_bins=n_bins, label=label)
def
single
(
ax
,
data
,
bin_lims
=
(
0
,
1
),
n_bins
=
30
,
label
=
""
,
color
=
None
):
bins
=
np
.
linspace
(
*
bin_lims
,
num
=
n_bins
+
1
)
ax
.
hist
(
data
,
bins
=
bins
,
histtype
=
"step"
,
label
=
label
,
color
=
color
)
def
_min_or_max
(
ax
,
data
,
idx
,
bin_lims
=
(
0
,
1
),
n_bins
=
30
,
labels
=
None
,
colors
=
None
):
val
=
data
[
np
.
arange
(
len
(
idx
)),
idx
]
bins
=
np
.
linspace
(
*
bin_lims
,
num
=
n_bins
+
1
)
for
i
in
range
(
data
.
shape
[
1
]):
label
=
labels
[
i
]
if
labels
else
""
color
=
colors
[
i
]
if
colors
else
None
ax
.
hist
(
val
[
idx
==
i
],
bins
=
bins
,
histtype
=
"step"
,
label
=
label
,
color
=
color
)
def
largest
(
ax
,
data
,
bin_lims
=
(
0
,
1
),
n_bins
=
30
,
labels
=
None
,
colors
=
None
):
idx
=
np
.
argmax
(
data
,
axis
=
1
)
_min_or_max
(
ax
,
data
,
idx
,
bin_lims
=
bin_lims
,
n_bins
=
n_bins
,
labels
=
labels
,
colors
=
colors
)
def
smallest
(
ax
,
data
,
bin_lims
=
(
0
,
1
),
n_bins
=
30
,
labels
=
None
,
colors
=
None
):
idx
=
np
.
argmin
(
data
,
axis
=
1
)
_min_or_max
(
ax
,
data
,
idx
,
bin_lims
=
bin_lims
,
n_bins
=
n_bins
,
labels
=
labels
,
colors
=
colors
)
Write
Preview
Supports
Markdown
0%
Try again
or
attach a new file
.
Cancel
You are about to add
0
people
to the discussion. Proceed with caution.
Finish editing this message first!
Cancel
Please
register
or
sign in
to comment