From 7af9ae020ab7fbc78f9cb940699194654055026a Mon Sep 17 00:00:00 2001 From: Noe Brucy Date: Tue, 2 Nov 2021 14:06:22 +0100 Subject: [PATCH] [plotter] added LIC + vector overlay overhaul --- plotter.py | 181 +++++++++++++++++++++++++++++++---------------------- 1 file changed, 106 insertions(+), 75 deletions(-) diff --git a/plotter.py b/plotter.py index 09ee976..600812d 100644 --- a/plotter.py +++ b/plotter.py @@ -29,7 +29,10 @@ if os.environ.get("DISPLAY", "") == "": import datetime import matplotlib.pyplot as plt -from moviepy.video.io.ImageSequenceClip import ImageSequenceClip +try: + from moviepy.video.io.ImageSequenceClip import ImageSequenceClip +except ModuleNotFoundError: + print("WARNING: no movie support (missing module moviepy)") from mpl_toolkits.axes_grid1.inset_locator import inset_axes import pspec_read @@ -39,6 +42,13 @@ from studyprocessor import StudyProcessor from run_selector import RunSelector from units import U, unit_str, convert_exp +try: + import lic +except ModuleNotFoundError: + print("WARNING: no LIC support (missing module lic)") + +from matplotlib.cm import ScalarMappable + from astrophysix.simdm.experiment import ( ParameterSetting, @@ -57,6 +67,81 @@ def not_array_error(err): return str(err)[-len(epy2) :] == epy2 or str(err)[-len(epy3) :] == epy3 + +def gethv(map_h, map_v, extent): + # Number of selected vectors + nh = map_h.shape[0] + nv = map_h.shape[1] + + # Creates vectors position grid + size_h = extent[1] - extent[0] # size of the map + dh = size_h / map_h.shape[0] # size of cell + seph = size_h / nh # separation between vectors + h = extent[0] + dh + np.arange(nh) * seph + + size_v = extent[3] - extent[2] + dv = size_v / map_h.shape[1] + sepv = size_v / nv + v = extent[2] + dv + np.arange(nv) * sepv + + return np.meshgrid(h, v) + +def streamplot(ax, map_h, map_v, extent, **kwargs): + """ + Add an overlay : streamlines + """ + hh, vv = gethv(map_h, map_v, extent) + ax.streamplot(hh, vv, map_h, map_v, **kwargs) + + +def quiver(ax, map_h, map_v, extent, key_v, label, **kwargs): + + + hh, vv = gethv(map_h, map_v, extent) + + # plot vector field + vec_field = ax.quiver(hh, vv, map_h, map_v, units="width", **kwargs) + + # get norm information + norm_v = np.sqrt(map_h ** 2 + map_v ** 2) + max_v = np.max(norm_v) + min_v = np.min(norm_v) + # add vector key + if key_v is None: + key_v = (max_v + min_v) / 2.0 + ax.quiverkey( + vec_field, + 0.6, + 0.98, + key_v, + f"${key_v:g}$ {label}", + labelpos="E", + coordinates="figure", + ) + +def line_integral_convolution(ax, map_h, map_v, extent, **kwargs): + """ + from Adnan Ali Ahmad + """ + + lic_res = lic.lic(map_v, map_h,length=20) #compute line integral convolution + + # Amplify contrast on lic + lim=(.1,.9) + lic_data_clip = np.clip(lic_res,lim[0],lim[1]) + lic_data_rgba = ScalarMappable(norm=None, cmap="binary").to_rgba(lic_data_clip) + lic_data_clip_rescale = (lic_data_clip-lim[0])/(lim[1]-lim[0]) + lic_data_rgba[...,3] = lic_data_clip_rescale * 1 + + args = [lic_data_rgba] + plot_args = {**kwargs} + plot_args["cmap"] = "binary" + plot_args["extent"] = extent + plot_args["origin"] = "lower" + ax.imshow(*args, **plot_args) + + + class PlotRule(Rule): """ The rule class, speficic to plot. @@ -626,7 +711,6 @@ class Plotter(Aggregator, BaseProcessor): pad=1, color="white", frameon=False, - size_vertical=1, ) plt.gca().add_artist(scalebar) @@ -858,90 +942,37 @@ class Plotter(Aggregator, BaseProcessor): # Scatter plot plt.scatter(part_h, part_v, s=mass[mask] / 5e3, **kwargs) - def _overlay_speed( - self, ax_los, im_extent, unit=U.km_s, unit_coeff=1.0, key_v=None, **kwargs + def _overlay_vector( + self, name, ax_los, extent, unit=U.km_s, unit_coeff=1.0, reduce_res=1, kind="quiver", **kwargs ): """ Add an overlay : velocity vector field """ - dmap_vh = self.current_processor.get_value("/maps/speed_h_{}".format(ax_los)) - dmap_vv = self.current_processor.get_value("/maps/speed_v_{}".format(ax_los)) - label, unit_old, unit = self._ax_label_unit( - f"/maps/speed_h_{ax_los}", "", unit, unit_coeff - ) + self.current_processor.process(f"{name}_h", ax_los) + self.current_processor.process(f"{name}_v", ax_los) + + map_h = self.current_processor.get_value("/maps/speed_h_{}".format(ax_los)) + map_v = self.current_processor.get_value("/maps/speed_v_{}".format(ax_los)) + label, unit_old, unit = self._ax_label_unit(f"/maps/speed_h_{ax_los}", "", unit, unit_coeff) - vel_red = self.params.plot.vel_red # take only a subset of velocities - map_vh_red = dmap_vh[::vel_red, ::vel_red] * unit_old.express(unit) - map_vv_red = dmap_vv[::vel_red, ::vel_red] * unit_old.express(unit) + map_h = map_h[::reduce_res, ::reduce_res] * unit_old.express(unit) + map_v = map_v[::reduce_res, ::reduce_res] * unit_old.express(unit) - # get norm information - norm_v = np.sqrt(map_vh_red ** 2 + map_vv_red ** 2) - max_v = np.max(norm_v) - min_v = np.min(norm_v) + if kind == "quiver": + quiver(plt.gca(), map_h, map_v, extent=extent, label=label, **kwargs) + elif kind == "streamplot": + streamplot(plt.gca(), map_h, map_v, extent=extent, **kwargs) + elif kind == "lic": + line_integral_convolution(plt.gca(), map_h, map_v, extent=extent, **kwargs) - # Number of selected vectors - nh = map_vh_red.shape[0] - nv = map_vh_red.shape[1] + def _overlay_speed(self, ax_los, **kwargs): + self._overlay_vector("speed", ax_los, **kwargs) - # Creates vectors position grid - size_h = im_extent[1] - im_extent[0] # size of the map - dh = size_h / dmap_vh.shape[0] # size of cell - seph = size_h / nh # separation between vectors - h = im_extent[0] + dh + np.arange(nh) * seph - - size_v = im_extent[3] - im_extent[2] - dv = size_v / dmap_vh.shape[1] - sepv = size_v / nv - v = im_extent[2] + dv + np.arange(nv) * sepv - - hh, vv = np.meshgrid(h, v) - - # plot vector field - vec_field = plt.quiver(hh, vv, map_vh_red, map_vv_red, units="width", **kwargs) - - # add vector key - if key_v is None: - key_v = (max_v + min_v) / 2.0 - plt.quiverkey( - vec_field, - 0.6, - 0.98, - key_v, - r"${:g}$".format(key_v) + label, - labelpos="E", - coordinates="figure", - ) - - def _overlay_B(self, ax_los, im_extent, **kwargs): - """ - Add an overlay : magnetic streamlines - """ - dmap_Bh = self.current_processor.get_value(f"/maps/B_h_{ax_los}") - dmap_Bv = self.current_processor.get_value(f"/maps/B_v_{ax_los}") - - # TODO : redo this with im_extent - - vel_red = self.params.plot.vel_red - radius = self.current_processor.attribute("/maps", "radius") - center = self.current_processor.attribute("/maps", "center") - lbox = self.current_processor.lbox - - map_Bh_red = dmap_Bh[::vel_red, ::vel_red] # take only a subset of velocities - map_Bv_red = dmap_Bv[::vel_red, ::vel_red] - nh = map_Bh_red.shape[0] - nv = map_Bv_red.shape[1] - vec_h = ( - np.arange(nh) * 2.0 / nh * radius - radius + center[0] + radius / nh - ) * lbox - vec_v = ( - np.arange(nv) * 2.0 / nv * radius - radius + center[1] + radius / nv - ) * lbox - hh, vv = np.meshgrid(vec_h, vec_v) - - plt.streamplot(hh, vv, map_Bh_red, map_Bv_red, **kwargs) + def _overlay_B(self, ax_los, **kwargs): + self._overlay_vector("B", ax_los, **kwargs) def _plot_hist( self,