scitex 2.0.0__py2.py3-none-any.whl

scitex/plt/ax/_style/_set_meta.py

@@ -0,0 +1,261 @@
+#!./env/bin/python3
+# -*- coding: utf-8 -*-
+# Time-stamp: "2025-06-04 11:35:00 (ywatanabe)"
+# Author: Yusuke Watanabe (ywata1989@gmail.com)
+
+"""
+Scientific metadata management for figures with YAML export.
+"""
+
+# Imports
+import yaml
+from typing import Optional, List, Dict, Any
+
+
+# Functions
+def set_meta(ax, caption=None, methods=None, stats=None, keywords=None, 
+             experimental_details=None, journal_style=None, 
+             significance=None, **kwargs):
+    """Set comprehensive scientific metadata for figures with YAML export
+    
+    Parameters
+    ----------
+    ax : matplotlib.axes.Axes or scitex AxisWrapper
+        The axes to modify
+    caption : str, optional
+        Figure caption text
+    methods : str, optional
+        Experimental methods description
+    stats : str, optional
+        Statistical analysis details
+    keywords : List[str], optional
+        Keywords for categorization and search
+    experimental_details : Dict[str, Any], optional
+        Structured experimental parameters (n_samples, temperature, etc.)
+    journal_style : str, optional
+        Target journal style ('nature', 'science', 'ieee', 'cell', etc.)
+    significance : str, optional
+        Significance statement or implications
+    **kwargs : additional metadata
+        Any additional metadata fields
+        
+    Returns
+    -------
+    ax : matplotlib.axes.Axes or scitex AxisWrapper
+        The modified axes
+        
+    Examples
+    --------
+    >>> fig, ax = scitex.plt.subplots()
+    >>> ax.plot(x, y, id='neural_data')
+    >>> ax.set_xyt(x='Time (ms)', y='Voltage (mV)', t='Neural Recording')
+    >>> ax.set_meta(
+    ...     caption='Intracellular recording showing action potentials.',
+    ...     methods='Whole-cell patch-clamp in acute brain slices.',
+    ...     stats='Statistical analysis using paired t-test (p<0.05).',
+    ...     keywords=['electrophysiology', 'neural_recording', 'patch_clamp'],
+    ...     experimental_details={
+    ...         'n_samples': 15,
+    ...         'temperature': 32,
+    ...         'recording_duration': 600,
+    ...         'electrode_resistance': '3-5 MΩ'
+    ...     },
+    ...     journal_style='nature',
+    ...     significance='Demonstrates novel neural dynamics in layer 2/3 pyramidal cells.'
+    ... )
+    >>> scitex.io.save(fig, 'neural_recording.png')  # YAML metadata auto-saved
+    """
+    
+    # Build comprehensive metadata dictionary
+    metadata = {}
+    
+    if caption is not None:
+        metadata['caption'] = caption
+    if methods is not None:
+        metadata['methods'] = methods
+    if stats is not None:
+        metadata['statistical_analysis'] = stats
+    if keywords is not None:
+        metadata['keywords'] = keywords if isinstance(keywords, list) else [keywords]
+    if experimental_details is not None:
+        metadata['experimental_details'] = experimental_details
+    if journal_style is not None:
+        metadata['journal_style'] = journal_style
+    if significance is not None:
+        metadata['significance'] = significance
+    
+    # Add any additional metadata
+    for key, value in kwargs.items():
+        if value is not None:
+            metadata[key] = value
+    
+    # Add automatic metadata
+    import datetime
+    metadata['created_timestamp'] = datetime.datetime.now().isoformat()
+    metadata['scitex_version'] = '1.11.0'  # TODO: Get from __version__
+    
+    # Store metadata in figure for automatic saving
+    fig = ax.get_figure()
+    if not hasattr(fig, '_scitex_metadata'):
+        fig._scitex_metadata = {}
+    
+    # Use axis as key for panel-specific metadata
+    fig._scitex_metadata[ax] = metadata
+    
+    # Also store as YAML-ready structure
+    if not hasattr(fig, '_scitex_yaml_metadata'):
+        fig._scitex_yaml_metadata = {}
+    fig._scitex_yaml_metadata[ax] = metadata
+    
+    # Backward compatibility - store simple caption
+    if caption is not None:
+        if not hasattr(fig, '_scitex_captions'):
+            fig._scitex_captions = {}
+        fig._scitex_captions[ax] = caption
+    
+    return ax
+
+
+def set_figure_meta(ax, caption=None, methods=None, stats=None, significance=None,
+                   funding=None, conflicts=None, data_availability=None, **kwargs):
+    """Set figure-level metadata for multi-panel figures
+    
+    Parameters
+    ----------
+    ax : matplotlib.axes.Axes or scitex AxisWrapper
+        Any axis in the figure (figure accessed via ax.get_figure())
+    caption : str, optional
+        Figure-level caption
+    methods : str, optional
+        Overall experimental methods
+    stats : str, optional
+        Overall statistical approach
+    significance : str, optional
+        Significance and implications
+    funding : str, optional
+        Funding acknowledgments
+    conflicts : str, optional
+        Conflict of interest statement
+    data_availability : str, optional
+        Data availability statement
+    **kwargs : additional metadata
+        Any additional figure-level metadata
+        
+    Returns
+    -------
+    ax : matplotlib.axes.Axes or scitex AxisWrapper
+        The modified axes
+        
+    Examples
+    --------
+    >>> fig, ((ax1, ax2), (ax3, ax4)) = scitex.plt.subplots(2, 2)
+    >>> # Set individual panel metadata...
+    >>> ax1.set_meta(caption='Panel A analysis...')
+    >>> ax2.set_meta(caption='Panel B comparison...')
+    >>> 
+    >>> # Set figure-level metadata
+    >>> ax1.set_figure_meta(
+    ...     caption='Comprehensive analysis of neural dynamics...',
+    ...     significance='This work demonstrates novel therapeutic targets.',
+    ...     funding='Supported by NIH grant R01-NS123456.',
+    ...     data_availability='Data available at doi:10.5061/dryad.example'
+    ... )
+    """
+    
+    # Build figure-level metadata
+    figure_metadata = {}
+    
+    if caption is not None:
+        figure_metadata['main_caption'] = caption
+    if methods is not None:
+        figure_metadata['overall_methods'] = methods
+    if stats is not None:
+        figure_metadata['overall_statistics'] = stats
+    if significance is not None:
+        figure_metadata['significance'] = significance
+    if funding is not None:
+        figure_metadata['funding'] = funding
+    if conflicts is not None:
+        figure_metadata['conflicts_of_interest'] = conflicts
+    if data_availability is not None:
+        figure_metadata['data_availability'] = data_availability
+    
+    # Add any additional metadata
+    for key, value in kwargs.items():
+        if value is not None:
+            figure_metadata[key] = value
+    
+    # Add automatic metadata
+    import datetime
+    figure_metadata['created_timestamp'] = datetime.datetime.now().isoformat()
+    
+    # Store in figure
+    fig = ax.get_figure()
+    fig._scitex_figure_metadata = figure_metadata
+    
+    # Backward compatibility
+    if caption is not None:
+        fig._scitex_main_caption = caption
+    
+    return ax
+
+
+def export_metadata_yaml(fig, filepath):
+    """Export all figure metadata to YAML file
+    
+    Parameters
+    ----------
+    fig : matplotlib.figure.Figure
+        Figure with metadata
+    filepath : str
+        Output YAML file path
+    """
+    import datetime
+    
+    # Collect all metadata
+    export_data = {
+        'figure_metadata': {},
+        'panel_metadata': {},
+        'export_info': {
+            'timestamp': datetime.datetime.now().isoformat(),
+            'scitex_version': '1.11.0'
+        }
+    }
+    
+    # Figure-level metadata
+    if hasattr(fig, '_scitex_figure_metadata'):
+        export_data['figure_metadata'] = fig._scitex_figure_metadata
+    
+    # Panel-level metadata
+    if hasattr(fig, '_scitex_yaml_metadata'):
+        for i, (ax, metadata) in enumerate(fig._scitex_yaml_metadata.items()):
+            panel_key = f'panel_{i+1}'
+            export_data['panel_metadata'][panel_key] = metadata
+    
+    # Write YAML file
+    with open(filepath, 'w') as f:
+        yaml.dump(export_data, f, default_flow_style=False, sort_keys=False, indent=2)
+
+
+if __name__ == "__main__":
+    # Start
+    import sys
+    import matplotlib.pyplot as plt
+    CONFIG, sys.stdout, sys.stderr, plt, CC = scitex.gen.start(sys, plt)
+
+    # Example usage
+    fig, ax = plt.subplots()
+    ax.plot([1, 2, 3], [1, 4, 2])
+    
+    set_meta(ax, 
+             caption='Example figure showing data trends.',
+             methods='Synthetic data generated for demonstration.',
+             keywords=['example', 'demo', 'synthetic'],
+             experimental_details={'n_samples': 3, 'data_type': 'synthetic'})
+    
+    export_metadata_yaml(fig, 'example_metadata.yaml')
+
+    # Close
+    scitex.gen.close(CONFIG)
+
+# EOF

scitex/plt/ax/_style/_set_n_ticks.py

@@ -0,0 +1,37 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+# Timestamp: "2025-04-29 12:02:14 (ywatanabe)"
+# File: /home/ywatanabe/proj/scitex_repo/src/scitex/plt/ax/_set_n_ticks.py
+# ----------------------------------------
+import os
+
+__FILE__ = "./src/scitex/plt/ax/_set_n_ticks.py"
+__DIR__ = os.path.dirname(__FILE__)
+# ----------------------------------------
+
+import matplotlib
+
+
+def set_n_ticks(
+    ax,
+    n_xticks=4,
+    n_yticks=4,
+):
+    """
+    Example:
+        ax = set_n_ticks(ax)
+    """
+
+    if n_xticks is not None:
+        ax.xaxis.set_major_locator(matplotlib.ticker.MaxNLocator(n_xticks))
+
+    if n_yticks is not None:
+        ax.yaxis.set_major_locator(matplotlib.ticker.MaxNLocator(n_yticks))
+
+    # Force the figure to redraw to reflect changes
+    ax.figure.canvas.draw()
+
+    return ax
+
+
+# EOF

scitex/plt/ax/_style/_set_size.py

@@ -0,0 +1,16 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+# Time-stamp: "2022-12-09 13:38:11 (ywatanabe)"
+
+
+def set_size(ax, w, h):
+    """w, h: width, height in inches"""
+    # if not ax: ax=plt.gca()
+    l = ax.figure.subplotpars.left
+    r = ax.figure.subplotpars.right
+    t = ax.figure.subplotpars.top
+    b = ax.figure.subplotpars.bottom
+    figw = float(w) / (r - l)
+    figh = float(h) / (t - b)
+    ax.figure.set_size_inches(figw, figh)
+    return ax

scitex/plt/ax/_style/_set_supxyt.py

@@ -0,0 +1,116 @@
+#!./env/bin/python3
+# -*- coding: utf-8 -*-
+# Time-stamp: "2024-07-13 07:56:46 (ywatanabe)"
+# Author: Yusuke Watanabe (ywata1989@gmail.com)
+
+"""
+This script does XYZ.
+"""
+
+# Imports
+import matplotlib.pyplot as plt
+
+from ._format_label import format_label
+
+
+# Functions
+def set_supxyt(ax, xlabel=False, ylabel=False, title=False, format_labels=True):
+    """Sets xlabel, ylabel and title"""
+    fig = ax.get_figure()
+
+    # if xlabel is not False:
+    #     fig.supxlabel(xlabel)
+
+    # if ylabel is not False:
+    #     fig.supylabel(ylabel)
+
+    # if title is not False:
+    #     fig.suptitle(title)
+    if xlabel is not False:
+        xlabel = format_label(xlabel) if format_labels else xlabel
+        fig.supxlabel(xlabel)
+
+    if ylabel is not False:
+        ylabel = format_label(ylabel) if format_labels else ylabel
+        fig.supylabel(ylabel)
+
+    if title is not False:
+        title = format_label(title) if format_labels else title
+        fig.suptitle(title)
+
+    return ax
+
+
+def set_supxytc(ax, xlabel=False, ylabel=False, title=False, caption=False, methods=False, stats=False, significance=False, format_labels=True):
+    """Sets figure-level xlabel, ylabel, title, and caption with SciTeX-Paper integration
+    
+    Parameters
+    ----------
+    ax : matplotlib.axes.Axes or scitex AxisWrapper
+        The axes to modify (figure accessed via ax.get_figure())
+    xlabel : str or False, optional
+        Figure-level X-axis label, by default False
+    ylabel : str or False, optional 
+        Figure-level Y-axis label, by default False
+    title : str or False, optional
+        Figure-level title (suptitle), by default False
+    caption : str or False, optional
+        Figure-level caption to store for later use with scitex.io.save(), by default False
+    methods : str or False, optional
+        Overall methods description for SciTeX-Paper integration, by default False
+    stats : str or False, optional
+        Overall statistical analysis details for SciTeX-Paper integration, by default False
+    significance : str or False, optional
+        Significance statement for SciTeX-Paper integration, by default False
+    format_labels : bool, optional
+        Whether to apply automatic formatting, by default True
+        
+    Returns
+    -------
+    ax : matplotlib.axes.Axes or scitex AxisWrapper
+        The modified axes
+        
+    Examples
+    --------
+    >>> fig, ((ax1, ax2), (ax3, ax4)) = scitex.plt.subplots(2, 2)
+    >>> # Add plots to each panel...
+    >>> ax1.set_supxytc(xlabel='Time (s)', ylabel='Signal Amplitude', 
+    ...                 title='Multi-Panel Analysis',
+    ...                 caption='Comprehensive analysis showing (A) raw data, (B) filtered signal, (C) power spectrum, and (D) phase analysis.',
+    ...                 methods='All experiments performed using standardized protocols.',
+    ...                 significance='This work demonstrates novel therapeutic targets.')
+    >>> scitex.io.save(fig, 'multi_panel.png')  # Caption automatically saved
+    """
+    # Set labels and title using existing function
+    set_supxyt(ax, xlabel=xlabel, ylabel=ylabel, title=title, format_labels=format_labels)
+    
+    # Store figure-level caption and extended metadata
+    if caption is not False or methods is not False or stats is not False or significance is not False:
+        fig = ax.get_figure()
+        # Store comprehensive figure-level metadata
+        fig_metadata = {
+            'main_caption': caption if caption is not False else None,
+            'methods': methods if methods is not False else None,
+            'stats': stats if stats is not False else None,
+            'significance': significance if significance is not False else None
+        }
+        
+        fig._scitex_figure_metadata = fig_metadata
+        
+        # Backward compatibility - also store simple caption
+        if caption is not False:
+            fig._scitex_main_caption = caption
+    
+    return ax
+
+
+if __name__ == "__main__":
+    # Start
+    CONFIG, sys.stdout, sys.stderr, plt, CC = scitex.gen.start(sys, plt)
+
+    # (YOUR AWESOME CODE)
+
+    # Close
+    scitex.gen.close(CONFIG)
+
+# EOF

scitex/plt/ax/_style/_set_ticks.py

@@ -0,0 +1,276 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+# Timestamp: "2025-04-27 20:04:55 (ywatanabe)"
+# File: /ssh:sp:/home/ywatanabe/proj/scitex_repo/src/scitex/plt/ax/_set_ticks.py
+# ----------------------------------------
+import os
+
+__FILE__ = "./src/scitex/plt/ax/_style/_set_ticks.py"
+__DIR__ = os.path.dirname(__FILE__)
+# ----------------------------------------
+
+import matplotlib.pyplot as plt
+import numpy as np
+
+from ....dict._to_str import to_str
+from ....types import is_listed_X
+
+
+def set_ticks(ax, xvals=None, xticks=None, yvals=None, yticks=None):
+    """Set custom tick labels on both x and y axes.
+
+    Convenience function to set tick positions and labels for both axes
+    at once. Automatically handles canvas updates for interactive backends.
+
+    Parameters
+    ----------
+    ax : matplotlib.axes.Axes
+        The axes object to modify.
+    xvals : array-like, optional
+        Values corresponding to x-axis data points.
+    xticks : list, optional
+        Desired tick labels for the x-axis.
+    yvals : array-like, optional
+        Values corresponding to y-axis data points.
+    yticks : list, optional
+        Desired tick labels for the y-axis.
+
+    Returns
+    -------
+    matplotlib.axes.Axes
+        The modified axes object.
+
+    Examples
+    --------
+    >>> fig, ax = plt.subplots()
+    >>> x = np.linspace(0, 10, 100)
+    >>> ax.plot(x, np.sin(x))
+    >>> ax = set_ticks(ax, xvals=x, xticks=[0, 5, 10],
+    ...                yvals=[-1, 0, 1], yticks=['-1', '0', '1'])
+
+    See Also
+    --------
+    set_x_ticks : Set ticks for x-axis only
+    set_y_ticks : Set ticks for y-axis only
+    """
+    ax = set_x_ticks(ax, x_vals=xvals, x_ticks=xticks)
+    ax = set_y_ticks(ax, y_vals=yvals, y_ticks=yticks)
+    canvas_type = type(ax.figure.canvas).__name__
+    if "TkAgg" in canvas_type:
+        ax.get_figure().canvas.draw()  # Redraw the canvas once after making all updates
+    return ax
+
+
+def set_x_ticks(ax, x_vals=None, x_ticks=None):
+    """
+    Set custom tick labels on the x and y axes based on specified values and desired ticks.
+
+    Parameters:
+    - ax: The axis object to modify.
+    - x_vals: Array of x-axis values.
+    - x_ticks: List of desired tick labels on the x-axis.
+    - y_vals: Array of y-axis values.
+    - y_ticks: List of desired tick labels on the y-axis.
+
+    Example:
+        import matplotlib.pyplot as plt
+        import numpy as np
+
+        fig, axes = plt.subplots(nrows=4)
+        x = np.linspace(0, 10, 100)
+        y = np.sin(x)
+        for ax in axes:
+            ax.plot(x, y)  # Plot a sine wave
+
+        set_ticks(axes[0])  # Do nothing # OK
+        set_ticks(axes[1], x_vals=x+3) # OK
+        set_ticks(axes[2], x_ticks=[1,2])  # OK
+        set_ticks(axes[3], x_vals=x+3, x_ticks=[4,5])  # Auto-generate ticks across the range
+        fig.tight_layout()
+        plt.show()
+    """
+
+    def _avoid_overlaps(values):
+        values = np.array(values)
+        if ("int" in str(values.dtype)) or ("float" in str(values.dtype)):
+            values = values.astype(float) + np.arange(len(values)) * 1e-5
+        return values
+
+    def _set_x_vals(ax, x_vals):
+        x_vals = _avoid_overlaps(x_vals)
+        new_x_axis = np.linspace(*ax.get_xlim(), len(x_vals))
+        ax.set_xticks(new_x_axis)
+        ax.set_xticklabels([f"{xv}" for xv in x_vals])
+        return ax
+
+    def _set_x_ticks(ax, x_ticks):
+        x_ticks = np.array(x_ticks)
+        if x_ticks.dtype.kind in ["U", "S", "O"]:  # If x_ticks are strings
+            ax.set_xticks(range(len(x_ticks)))
+            ax.set_xticklabels(x_ticks)
+        else:
+            x_vals = np.array(
+                [label.get_text().replace("−", "-") for label in ax.get_xticklabels()]
+            )
+            x_vals = x_vals.astype(float)
+            x_indi = np.argmin(
+                np.array(np.abs(x_vals[:, np.newaxis] - x_ticks[np.newaxis, :])),
+                axis=0,
+            )
+            ax.set_xticks(ax.get_xticks()[x_indi])
+            ax.set_xticklabels([f"{xt}" for xt in x_ticks])
+        return ax
+
+    x_vals_passed = x_vals is not None
+    x_ticks_passed = x_ticks is not None
+
+    if is_listed_X(x_ticks, dict):
+        x_ticks = [to_str(xt, delimiter="\n") for xt in x_ticks]
+
+    if (not x_vals_passed) and (not x_ticks_passed):
+        # Do nothing
+        pass
+
+    elif x_vals_passed and (not x_ticks_passed):
+        # Replaces the x axis to x_vals
+        x_ticks = np.linspace(x_vals[0], x_vals[-1], 4)
+        ax = _set_x_vals(ax, x_ticks)
+
+    elif (not x_vals_passed) and x_ticks_passed:
+        # Locates 'x_ticks' on the original x axis
+        ax.set_xticks(x_ticks)
+
+    elif x_vals_passed and x_ticks_passed:
+        if isinstance(x_vals, str):
+            if x_vals == "auto":
+                x_vals = np.arange(len(x_ticks))
+
+        # Replaces the original x axis to 'x_vals' and locates the 'x_ticks' on the new axis
+        ax = _set_x_vals(ax, x_vals)
+        ax = _set_x_ticks(ax, x_ticks)
+
+    return ax
+
+
+def set_y_ticks(ax, y_vals=None, y_ticks=None):
+    """
+    Set custom tick labels on the y-axis based on specified values and desired ticks.
+
+    Parameters:
+    - ax: The axis object to modify.
+    - y_vals: Array of y-axis values where ticks should be placed.
+    - y_ticks: List of labels for ticks on the y-axis.
+
+    Example:
+        import matplotlib.pyplot as plt
+        import numpy as np
+
+        fig, ax = plt.subplots()
+        x = np.linspace(0, 10, 100)
+        y = np.sin(x)
+        ax.plot(x, y)  # Plot a sine wave
+
+        set_y_ticks(ax, y_vals=y, y_ticks=['Low', 'High'])  # Set custom y-axis ticks
+        plt.show()
+    """
+
+    def _avoid_overlaps(values):
+        values = np.array(values)
+        if ("int" in str(values.dtype)) or ("float" in str(values.dtype)):
+            values = values.astype(float) + np.arange(len(values)) * 1e-5
+        return values
+
+    def _set_y_vals(ax, y_vals):
+        y_vals = _avoid_overlaps(y_vals)
+        new_y_axis = np.linspace(*ax.get_ylim(), len(y_vals))
+        ax.set_yticks(new_y_axis)
+        ax.set_yticklabels([f"{yv:.2f}" for yv in y_vals])
+        return ax
+
+    # def _set_y_ticks(ax, y_ticks):
+    #     y_ticks = np.array(y_ticks)
+    #     y_vals = np.array(
+    #         [
+    #             label.get_text().replace("−", "-")
+    #             for label in ax.get_yticklabels()
+    #         ]
+    #     )
+    #     y_vals = y_vals.astype(float)
+    #     y_indi = np.argmin(
+    #         np.array(np.abs(y_vals[:, np.newaxis] - y_ticks[np.newaxis, :])),
+    #         axis=0,
+    #     )
+
+    #     # y_indi = [np.argmin(np.abs(y_vals - yt)) for yt in y_ticks]
+    #     ax.set_yticks(ax.get_yticks()[y_indi])
+    #     ax.set_yticklabels([f"{yt}" for yt in y_ticks])
+    #     return ax
+    def _set_y_ticks(ax, y_ticks):
+        y_ticks = np.array(y_ticks)
+        if y_ticks.dtype.kind in ["U", "S", "O"]:  # If y_ticks are strings
+            ax.set_yticks(range(len(y_ticks)))
+            ax.set_yticklabels(y_ticks)
+        else:
+            y_vals = np.array(
+                [label.get_text().replace("−", "-") for label in ax.get_yticklabels()]
+            )
+            y_vals = y_vals.astype(float)
+            y_indi = np.argmin(
+                np.array(np.abs(y_vals[:, np.newaxis] - y_ticks[np.newaxis, :])),
+                axis=0,
+            )
+            ax.set_yticks(ax.get_yticks()[y_indi])
+            ax.set_yticklabels([f"{yt}" for yt in y_ticks])
+        return ax
+
+    y_vals_passed = y_vals is not None
+    y_ticks_passed = y_ticks is not None
+
+    if is_listed_X(y_ticks, dict):
+        y_ticks = [to_str(yt, delimiter="\n") for yt in y_ticks]
+
+    if (not y_vals_passed) and (not y_ticks_passed):
+        # Do nothing
+        pass
+
+    elif y_vals_passed and (not y_ticks_passed):
+        # Replaces the y axis to y_vals
+        ax = _set_y_vals(ax, y_vals)
+
+    elif (not y_vals_passed) and y_ticks_passed:
+        # Locates 'y_ticks' on the original y axis
+        ax.set_yticks(y_ticks)
+
+    elif y_vals_passed and y_ticks_passed:
+        # Replaces the original y axis to 'y_vals' and locates the 'y_ticks' on the new axis
+        if y_vals == "auto":
+            y_vals = np.arange(len(y_ticks))
+
+        ax = _set_y_vals(ax, y_vals)
+        ax = _set_y_ticks(ax, y_ticks)
+    return ax
+
+
+if __name__ == "__main__":
+    import scitex
+
+    xx, tt, fs = scitex.dsp.demo_sig()
+    pha, amp, freqs = scitex.dsp.wavelet(xx, fs)
+
+    i_batch, i_ch = 0, 0
+    ff = freqs[i_batch, i_ch]
+    fig, ax = scitex.plt.subplots()
+
+    ax.image2d(amp[i_batch, i_ch])
+
+    ax = set_ticks(
+        ax,
+        x_vals=tt,
+        x_ticks=[0, 1, 2, 3, 4],
+        y_vals=ff,
+        y_ticks=[0, 128, 256],
+    )
+
+    plt.show()
+
+# EOF