-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathplot.py
More file actions
75 lines (59 loc) · 2.3 KB
/
plot.py
File metadata and controls
75 lines (59 loc) · 2.3 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
"""Functions for plotting data files."""
from __future__ import print_function
import numpy as np
import matplotlib.pyplot as plt
def rivet_paths(file_name):
"""Return all :py:mod:`rivet` paths found at file_name."""
from . import yodaplot
return yodaplot.data_object_names(file_name)
def gridplot(file_name, uses_rivet_plot_info=True):
"""Convenience function to plot all :py:mod:`yoda` data objects
from a :py:mod:`yoda` file into a subplots grid.
:param str file_name: The path to the :py:mod:`yoda` file.
:return: fig, axes_list
"""
all_rivet_paths = rivet_paths(file_name)
# setup axes
if len(all_rivet_paths) == 1:
fig, axes_list = plt.subplots()
else:
ncols = 2
nrows = (len(all_rivet_paths) - 1) / ncols + 1
fig, axes_list = plt.subplots(nrows, ncols, squeeze=False)
# plot into axes
for rivet_path, axes in zip(all_rivet_paths, np.ravel(axes_list)):
plt.sca(axes)
plot(file_name, rivet_path, uses_rivet_plot_info=uses_rivet_plot_info)
return fig, axes_list
def plot(filename_or_data_object, rivet_path,
uses_rivet_plot_info=True, errors_enabled=None,
**kwargs):
"""Plot a :py:mod:`yoda` data object, potentially from a :py:mod:`yoda` file."""
from . import yodaplot
print("Plotting", rivet_path, end="")
if isinstance(filename_or_data_object, str):
print(" from", filename_or_data_object, "...")
else:
print()
if uses_rivet_plot_info and errors_enabled is None:
from . import rivetplot
errors_enabled = rivetplot.errors_enabled(rivet_path)
else:
errors_enabled = True if errors_enabled is None else errors_enabled
if "rebin_count" in kwargs:
rebin_count = kwargs["rebin_count"]
del kwargs["rebin_count"]
else:
if uses_rivet_plot_info:
from . import rivetplot
rebin_count = rivetplot.rebin_count(rivet_path)
else:
rebin_count = 1
result = yodaplot.plot(filename_or_data_object, rivet_path,
errors_enabled=errors_enabled,
rebin_count=rebin_count,
**kwargs)
if uses_rivet_plot_info:
from . import rivetplot
rivetplot.apply_plot_info(rivet_path)
return result