feat: rotate option for heatmap (spacetime)

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

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 = {}
@@ -124,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
 
@@ -172,13 +192,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 +269,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

idefix2python/renderer.py

@@ -161,30 +161,32 @@ 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)
 
-                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)
 
@@ -203,6 +205,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
@@ -265,6 +269,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 +476,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 +637,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 +663,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 +687,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

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),