From 745b5fc2714e80e8aeefe8c0a7b0b0cd7dad2315 Mon Sep 17 00:00:00 2001 From: David Fang Date: Tue, 28 Jul 2026 11:37:44 +0100 Subject: [PATCH 1/3] feat: rotate option for heatmap (spacetime) --- idefix2python/axes.py | 23 ++++++--- idefix2python/quantities.py | 49 +++++++++++++------ idefix2python/renderer.py | 97 +++++++++++++++++++++++-------------- idefix2python/tools.py | 2 +- 4 files changed, 111 insertions(+), 60 deletions(-) diff --git a/idefix2python/axes.py b/idefix2python/axes.py index 3637336..b49b1aa 100644 --- a/idefix2python/axes.py +++ b/idefix2python/axes.py @@ -14,15 +14,21 @@ # No data should appear in Fig, Ax: they are sent by Renderer. DPI = 300 -COLUMN_WIDTH = 6 -ROW_HEIGHT = 6 -# ROW_HEIGHT = 2 +plt.rcParams.update({"font.size": 16}) class Fig: counter = -1 - def __init__(self, quantities, suptitle=None, suptitle_kwargs=None, **kwargs): + def __init__( + self, + quantities, + suptitle=None, + suptitle_kwargs=None, + column_width=6, + row_height=6, + **kwargs, + ): Fig.counter += 1 self.name = f"fig{Fig.counter}" self.quantities = quantities @@ -33,6 +39,8 @@ def __init__(self, quantities, suptitle=None, suptitle_kwargs=None, **kwargs): self.axesMovie = [] self.axesTimeline = [] self.movie = False + self.column_width = column_width + self.row_height = row_height self.rows = 1 self.columns = 1 for qtyInfo in quantities: @@ -65,14 +73,15 @@ def init(self): self.axes[*qtyInfo.plot_coords].add_quantity(qtyInfo) def generate_figure(self, custom_suptitle=None): - fig_width = max(6, COLUMN_WIDTH * self.columns) # minimum width of 6 - fig_height = max(4, ROW_HEIGHT * self.rows) # minimum height of 4 + fig_width = max(6, self.column_width * self.columns) # minimum width of 6 + fig_height = max(4, self.row_height * self.rows) # minimum height of 4 fig, axs = plt.subplots( self.rows, self.columns, figsize=(fig_width, fig_height), squeeze=False, - layout="constrained", + # layout="constrained", + layout="compressed", **self.kwargs, ) self.fig = fig diff --git a/idefix2python/quantities.py b/idefix2python/quantities.py index 8b56d5a..250bbca 100644 --- a/idefix2python/quantities.py +++ b/idefix2python/quantities.py @@ -1,5 +1,30 @@ from itertools import count +default_pmesh_kwargs = { + "rasterized": True, + "edgecolors": None, + "antialiased": True, +} + +default_streamline_kwargs = { + "linewidth": 0.2, + "arrowstyle": "->", + "color": "#d3d3d3", + # "color": (1, 1, 1, 0.5), + "density": 2, +} + +default_parts_kwargs = {"marker": "x", "markersize": 0.2} + +default_ref_plot_kwargs = { + "zorder": 3, + "ls": "--", + "lw": 1, + "alpha": 0.8, + "color": "limegreen", + "label": "Reference", +} + class Data: """ @@ -65,7 +90,6 @@ def __init__(self, key, symbol="", plot_coords=None, bounds=None, **kwargs): self.style_kwargs = kwargs.get("style_kwargs", {}) - default_parts_kwargs = {"marker": "x", "markersize": 0.2} self.parts_kwargs = merge_default_to_dict( default_parts_kwargs, kwargs.get("parts_kwargs", {}) ) @@ -77,14 +101,7 @@ def __init__(self, key, symbol="", plot_coords=None, bounds=None, **kwargs): self.ref_function = kwargs.get("ref_function", None) self.pointsRef = [] self.valuesRef = [] - default_ref_plot_kwargs = { - "zorder": 3, - "ls": "--", - "lw": 1, - "alpha": 0.8, - "color": "limegreen", - "label": "Reference", - } + if self.ref_function is not None: if not hasattr(self.ref_function, "plot_kwargs"): self.ref_function.plot_kwargs = {} @@ -172,13 +189,7 @@ def __init__( raise Exception( f"Invalid streamline configuration: {streamlines}. Expected a list/tuple of length 2." ) - default_streamline_kwargs = { - "linewidth": 0.2, - "arrowstyle": "->", - "color": "#d3d3d3", - # "color": (1, 1, 1, 0.5), - "density": 2, - } + self.streamline_kwargs = merge_default_to_dict( default_streamline_kwargs, kwargs.get("streamline_kwargs", {}) ) @@ -255,8 +266,14 @@ def __init__( bounds=None, norm="linear", uids=None, + rotate=False, **kwargs, ): + self.rotate = rotate + self.style_kwargs = merge_default_to_dict( + default_pmesh_kwargs, + kwargs.get("style_kwargs", {}), + ) super().__init__(key, symbol, plot_coords, bounds, **kwargs) self.set_norm(norm) self.uids = uids diff --git a/idefix2python/renderer.py b/idefix2python/renderer.py index 8aabd50..ad15894 100644 --- a/idefix2python/renderer.py +++ b/idefix2python/renderer.py @@ -161,26 +161,26 @@ def _pre_render(self): qtyInfo.points = ( self.gridInfo.X1Line ) # TODO to change if user wants custom xqty - qtyInfo.xmin = ( - qtyInfo.xmin - if qtyInfo.xmin is not None - else np.min(self.processor.years) - ) - qtyInfo.xmax = ( - qtyInfo.xmax - if qtyInfo.xmax is not None - else np.max(self.processor.years) - ) - qtyInfo.ymin = ( - qtyInfo.ymin - if qtyInfo.ymin is not None - else np.nanmin(qtyInfo.points) - ) - qtyInfo.ymax = ( - qtyInfo.ymax - if qtyInfo.ymax is not None - else np.nanmax(qtyInfo.points) - ) + # qtyInfo.xmin = ( + # qtyInfo.xmin + # if qtyInfo.xmin is not None + # else np.min(self.processor.years) + # ) + # qtyInfo.xmax = ( + # qtyInfo.xmax + # if qtyInfo.xmax is not None + # else np.max(self.processor.years) + # ) + # qtyInfo.ymin = ( + # qtyInfo.ymin + # if qtyInfo.ymin is not None + # else np.nanmin(qtyInfo.points) + # ) + # qtyInfo.ymax = ( + # qtyInfo.ymax + # if qtyInfo.ymax is not None + # else np.nanmax(qtyInfo.points) + # ) qtyInfo.set_default_xlabel(r"$t$ [yr]") qtyInfo.set_default_ylabel(self.gridInfo.axis_name_1) @@ -265,6 +265,7 @@ def render_Frame(self, frame_nb=None, vtkPath=None, partPath=None): partvtk = None if partPath is None else readVTK(partPath) commonvtk = self.processor.process(datavtk=datavtk, partvtk=partvtk) custom_suptitle = f"{self.context.runName}\n{Path(*vtkPath.parts[-4:])}\n$t={VTK.t[0]:.1e}$" + custom_suptitle = rf"$t={VTK.t[0] / (2 * np.pi):.1e}\,\mathrm{{yr}}$" else: custom_suptitle = None @@ -471,10 +472,16 @@ def do_timeline_stuff(self, figure, timeline, frame_nb=-1): ax = figure.axes[*timeline.plot_coords].ax if getattr(ax, "show_time_indicator", True): if frame_nb > 0: - ax.axvline( - x=self.processor.years[frame_nb], - **TIMEINDICATOR_KWARGS, - ) + if getattr(timeline, "rotate", None): + ax.axhline( + y=self.processor.years[frame_nb], + **TIMEINDICATOR_KWARGS, + ) + else: + ax.axvline( + x=self.processor.years[frame_nb], + **TIMEINDICATOR_KWARGS, + ) ax.show_time_indicator = False def _render_SpaceTimeHeatmap(self, figure, sptime, commonvtk, frame_nb=-1): @@ -626,8 +633,6 @@ def _draw_pcolormesh(self, figure, qtyInfo, data=None): np.asarray(self.processor.years), np.asarray(self.gridInfo.X1Line), ) - print(np.shape(np.transpose(qtyInfo.values))) - print(np.shape(self.gridInfo.mask1)) data_mesh = np.transpose(qtyInfo.values)[self.gridInfo.mask1] vmin, vmax = qtyInfo.bounds if vmin is None or self.userArgs.noBounds: @@ -654,16 +659,13 @@ def _draw_pcolormesh(self, figure, qtyInfo, data=None): ax = figure.axes[*qtyInfo.plot_coords].ax + is_rotated = getattr(qtyInfo, "rotate", False) cmesh = ax.pcolormesh( - grid1, - grid2, - data_mesh, + grid2.T if is_rotated else grid1, + grid1.T if is_rotated else grid2, + data_mesh.T if is_rotated else data_mesh, norm=norm, **qtyInfo.style_kwargs, - rasterized=True, - shading="gouraud", - edgecolors="none", - antialiased=True, ) cbar = None @@ -681,14 +683,37 @@ def _draw_pcolormesh(self, figure, qtyInfo, data=None): def colorbar(mappable, cbformat): + loc = "bottom" last_axes = plt.gca() ax = mappable.axes fig = ax.figure + + cbar = fig.colorbar(mappable, ax=ax, location=loc, format=cbformat) + + return cbar # loc = "left" - loc = "bottom" divider = make_axes_locatable(ax) - cax = divider.append_axes(loc, size="2%", pad=0.75) + # cax = divider.append_axes( + # loc, + # size="2%", + # ) + pos = ax.get_position() + + # Define how far below the plot you want it (y-offset) + # Increase this value to push the colorbar lower down + gap = 0.2 + + # Define colorbar dimensions relative to the main plot + cb_height = 0.02 + cb_width = pos.width # Make it 80% of the plot's width + + # Center it horizontally under the main plot + cb_left = pos.x0 + (pos.width - 0.57) / 2 + cb_bottom = pos.y0 - gap + + # Create the isolated axis + cax = fig.add_axes([cb_left, cb_bottom, cb_width, cb_height]) # cax = divider.append_axes(loc, size="4%") - cbar = fig.colorbar(mappable, cax=cax, location=loc, format=cbformat) + cbar = fig.colorbar(mappable, cax=cax, location=loc, format=cbformat, pad="1000%") plt.sca(last_axes) return cbar diff --git a/idefix2python/tools.py b/idefix2python/tools.py index 85e8077..d7ec02d 100644 --- a/idefix2python/tools.py +++ b/idefix2python/tools.py @@ -61,7 +61,7 @@ def movie(pattern_png, movie_path, fps=10): print(movie_path, "from", pattern_png) ffmpeg.input(pattern_png, pattern_type="glob", framerate=fps).filter( "scale", - 3840, + 1920, "-2", # TODO More flexible ).output( str(movie_path), From ea7ccc635f944b14256f4d5681b0bdadebef023c Mon Sep 17 00:00:00 2001 From: David Fang Date: Thu, 30 Jul 2026 20:38:10 +0100 Subject: [PATCH 2/3] flip labels --- idefix2python/quantities.py | 3 +++ idefix2python/renderer.py | 2 ++ 2 files changed, 5 insertions(+) diff --git a/idefix2python/quantities.py b/idefix2python/quantities.py index 250bbca..6b65e0c 100644 --- a/idefix2python/quantities.py +++ b/idefix2python/quantities.py @@ -141,6 +141,9 @@ def set_default_ylabel(self, ylabel): if self.ylabel is None: self.ylabel = ylabel + def switch_labels(self): + self.xlabel, self.ylabel = self.ylabel, self.xlabel + def __str__(self): return self.key diff --git a/idefix2python/renderer.py b/idefix2python/renderer.py index ad15894..ab7c768 100644 --- a/idefix2python/renderer.py +++ b/idefix2python/renderer.py @@ -203,6 +203,8 @@ def _pre_render(self): LOG( f"Warning: Failed to compute ref_function for {qtyInfo.key}. Error: {e}" ) + if getattr(qtyInfo, "rotate", None): + qtyInfo.switch_labels() if qtyInfo.is_timeline: qtyInfo.points = self.processor.years From 22254c2c28939bd27a1a164b514fc1166aa5e5fa Mon Sep 17 00:00:00 2001 From: David Fang Date: Thu, 30 Jul 2026 20:38:31 +0100 Subject: [PATCH 3/3] default labels for OneComponentOneVariable --- idefix2python/renderer.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/idefix2python/renderer.py b/idefix2python/renderer.py index ab7c768..b70607f 100644 --- a/idefix2python/renderer.py +++ b/idefix2python/renderer.py @@ -184,7 +184,9 @@ def _pre_render(self): qtyInfo.set_default_xlabel(r"$t$ [yr]") qtyInfo.set_default_ylabel(self.gridInfo.axis_name_1) - elif isinstance(qtyInfo, PartQuantity): + elif isinstance(qtyInfo, PartQuantity) or isinstance( + qtyInfo, OneComponentOneVariable + ): qtyInfo.set_default_xlabel(r"$t$ [yr]") qtyInfo.set_default_ylabel(qtyInfo.symbol)