1import os
2import os.path as op
3import logging
5from pyrocko import guts, util
6from pyrocko.guts import Dict, List, Tuple, Float, Unicode, Object, String
8from grond.meta import StringID
9from grond.plot.config import PlotFormat
12guts_prefix = 'grond'
14logger = logging.getLogger('grond.plot.collection')
17class PlotItem(Object):
18 name = StringID.T()
19 attributes = Dict.T(
20 StringID.T(), List.T(String.T()))
21 title = Unicode.T(
22 optional=True,
23 help='item\'s description')
24 description = Unicode.T(
25 optional=True,
26 help='item\'s description')
29class PlotGroup(Object):
30 name = StringID.T(
31 help='group name')
32 section = StringID.T(
33 optional=True,
34 help='group\'s section path, e.g. results.waveforms')
35 title = Unicode.T(
36 optional=True,
37 help='group\'s title')
38 description = Unicode.T(
39 optional=True,
40 help='group description')
41 formats = List.T(
42 PlotFormat.T(),
43 help='plot format')
44 variant = StringID.T(
45 help='variant of the group')
46 feather_icon = String.T(
47 default='bar-chart-2',
48 help='Feather icon for the HTML report.')
49 size_cm = Tuple.T(2, Float.T())
50 items = List.T(PlotItem.T())
51 attributes = Dict.T(StringID.T(), List.T(String.T()))
53 def filename_image(self, item, format):
54 return '%s.%s.%s.%s' % (
55 self.name,
56 self.variant,
57 item.name,
58 format.extension)
61class PlotCollection(Object):
62 group_refs = List.T(Tuple.T(2, StringID.T()))
65class PlotCollectionManager(object):
67 def __init__(self, path, show=False):
68 self._path = path
69 self.load_collection()
70 self._show = show
72 def load_collection(self):
73 path = self.path_collection()
74 if op.exists(path):
75 self._collection = guts.load(filename=self.path_collection())
76 else:
77 self._collection = PlotCollection()
79 def dump_collection(self):
80 path = self.path_collection()
81 util.ensuredirs(path)
82 guts.dump(self._collection, filename=path)
84 def path_collection(self):
85 return op.join(self._path, 'plot_collection.yaml')
87 def path_image(self, group, item, format):
88 return op.join(
89 self._path, group.name, group.variant,
90 group.filename_image(item, format))
92 def path_group(self, group_ref=None, group=None):
93 if group_ref is not None:
94 group_name, group_variant = group_ref
95 else:
96 group_name = group.name
97 group_variant = group.variant
99 return op.join(
100 self._path, group_name, group_variant, '%s.%s.%s' % (
101 group_name, group_variant, 'plot_group.yaml'))
103 def paths_group_dirs(self, group_ref=None, group=None):
104 if group_ref is not None:
105 group_name, group_variant = group_ref
106 else:
107 group_name = group.name
108 group_variant = group.variant
110 return [
111 op.join(self._path, group_name, group_variant),
112 op.join(self._path, group_name)]
114 def create_group_mpl(self, config, iter_item_figure, **kwargs):
115 from matplotlib import pyplot as plt
116 group = PlotGroup(
117 formats=guts.clone(config.formats),
118 size_cm=config.size_cm,
119 name=config.name,
120 variant=config.variant,
121 **kwargs)
123 path_group = self.path_group(group=group)
124 if os.path.exists(path_group):
125 self.remove_group_files(path_group)
127 group_ref = (group.name, group.variant)
128 if group_ref in self._collection.group_refs:
129 self._collection.group_refs.remove(group_ref)
131 self.dump_collection()
133 figs_to_close = []
134 for item, fig in iter_item_figure:
135 group.items.append(item)
136 for format in group.formats:
137 path = self.path_image(group, item, format)
138 util.ensuredirs(path)
139 format.render_mpl(
140 fig,
141 path=path,
142 dpi=format.get_dpi(group.size_cm))
144 logger.info('Figure saved: %s' % path)
146 if not self._show:
147 plt.close(fig)
148 else:
149 figs_to_close.append(fig)
151 util.ensuredirs(path_group)
152 group.validate()
153 group.dump(filename=path_group)
154 self._collection.group_refs.append(group_ref)
155 self.dump_collection()
157 if self._show:
158 plt.show()
160 for fig in figs_to_close:
161 plt.close(fig)
163 def create_group_automap(self, config, iter_item_figure, **kwargs):
164 group = PlotGroup(
165 formats=guts.clone(config.formats),
166 size_cm=config.size_cm,
167 name=config.name,
168 variant=config.variant,
169 **kwargs)
171 path_group = self.path_group(group=group)
172 if os.path.exists(path_group):
173 self.remove_group_files(path_group)
175 group_ref = (group.name, group.variant)
176 if group_ref in self._collection.group_refs:
177 self._collection.group_refs.remove(group_ref)
179 self.dump_collection()
181 for item, automap in iter_item_figure:
182 group.items.append(item)
183 for format in group.formats:
184 path = self.path_image(group, item, format)
185 util.ensuredirs(path)
186 format.render_automap(
187 automap,
188 path=path,
189 resolution=format.get_dpi(group.size_cm))
191 logger.info('Figure saved: %s' % path)
193 util.ensuredirs(path_group)
194 group.dump(filename=path_group)
195 self._collection.group_refs.append(group_ref)
196 self.dump_collection()
198 def create_group_gmtpy(self, config, iter_item_figure):
199 pass
201 def remove_group_files(self, path_group):
202 group = guts.load(filename=path_group)
203 for item in group.items:
204 for format in group.formats:
205 path = self.path_image(group, item, format)
206 try:
207 os.unlink(path)
208 except OSError:
209 pass
211 os.unlink(path_group)
212 for path in self.paths_group_dirs(group=group):
213 try:
214 os.rmdir(path)
215 except OSError:
216 pass
219__all__ = [
220 'PlotItem',
221 'PlotGroup',
222 'PlotCollection',
223 'PlotCollectionManager',
224]