Create a scatter plot with dynamic opacity.

dynamic_opacity.py

"""Scatter plots with dynamic opacities.

Inspired by:

https://observablehq.com/@rreusser/selecting-the-right-opacity-for-2d-point-clouds.

Copyright (C) 2024 by Paul Brodersen.

This program is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, either version 3 of the License, or (at
your option) any later version. This program is distributed in the
hope that it will be useful, but WITHOUT ANY WARRANTY; without even
the implied warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR
PURPOSE.  See the GNU General Public License for more details. You
should have received a copy of the GNU General Public License along
with this program. If not, see <https://www.gnu.org/licenses/.
"""

import numpy as np
import matplotlib.pyplot as plt

from matplotlib.transforms import Affine2D
from matplotlib.colors import to_rgba_array
from matplotlib.collections import CircleCollection


def rgba_to_grayscale(r : float, g : float, b : float, a : float = 1) -> float:
    # Adapted from: https://stackoverflow.com/a/689547/2912349
    return (0.299 * r + 0.587 * g + 0.114 * b) * a


class ScatterPlot2D:
    """Draw a scatter plot in which the opacity of the markers is adjusted depending on the number of markers in view.

    Parameters
    ----------
    x, y : NDArray
        The x and y origins of each marker.
    sizes : Optional[int | float | NDArray]
        The marker sizes.
    alpha : Optional[float]
        The initial opacity.
    ax: Optional[plt.Axes]
        The matplotlib axis instance.
    *args, **kwargs:
        Passed through to matplotlib.collections.CircleCollection

    Attributes
    ----------
    circles : matplotlib.collections.CircleCollection
        The circle collection.

    Notes
    -----
    As this class uses matplotlib.collections.CircleCollection to draw the data points,
    other marker shapes are not supported.
    """

    def __init__(self, x, y, sizes=10, alpha=0.01, ax=None, *args, **kwargs):

        # Initialize the axis if none given.
        if ax is None:
            ax = plt.gca()
        fig = ax.get_figure()

        # Draw markers.
        sizes = np.full_like(x, sizes) if isinstance(sizes, (float, int)) else sizes
        self.circles = CircleCollection(
            sizes,
            offsets=np.c_[x, y],
            offset_transform=ax.transData,
            alpha=alpha,
            *args, **kwargs
        )
        self.circles.set_transform(Affine2D().scale(fig.dpi / 72.0))  # the points to pixels transform
        ax.add_collection(self.circles, autolim=True)
        ax.autoscale_view()

        # Computed the weighted intensity of each point: size * b/w color intensity.
        colors = self.circles.get_facecolor()
        colors = colors if colors.shape[0] > 1 else np.ones_like(x)[:, np.newaxis] * colors
        intensities = np.array([rgba_to_grayscale(*rgba) * size for rgba, size in zip(colors, sizes)])
        total_intensity = np.sum(intensities)

        def on_zoom(event):
            # Determine points remaining in view.
            xmin, xmax = ax.get_xlim()
            ymin, ymax = ax.get_ylim()
            mask = np.logical_and(
                np.logical_and(x >= xmin, x < xmax),
                np.logical_and(y >= ymin, y < ymax),
            )
            # Rescale alpha by dividing it by the remaining fraction of the total intensity.
            # For large quotients, ensures a smooth approach to one by applying the hyperbolic tangent.
            remaining_intensity = np.sum(intensities[mask]) / total_intensity
            self.circles.set_alpha(np.tanh(alpha / remaining_intensity))
            fig.canvas.draw()

        # Most zoom events trigger changes in both, x and y.
        # Connecting to both would result in evaluating the code twice on each zoom event.
        # ax.callbacks.connect('xlim_changed', on_zoom)
        ax.callbacks.connect('ylim_changed', on_zoom)


if __name__ == "__main__":

    def generate_roessler_attractor(total_points, a=0.2, b=0.2, c=5.7, dt=0.006):

        xn = 2.644838333129883
        yn = 4.060488700866699
        zn = 2.8982460498809814

        for ii in range(total_points):
            dx = -yn - zn
            dy = xn + a * yn
            dz = b + zn * (xn - c)

            xh = xn + 0.5 * dt * dx
            yh = yn + 0.5 * dt * dy
            zh = zn + 0.5 * dt * dz

            dx = -yh - zh
            dy = xh + a * yh
            dz = b + zh * (xh - c)

            xn1 = xn + dt * dx
            yn1 = yn + dt * dy
            zn1 = zn + dt * dz

            xn = xn1
            yn = yn1
            zn = zn1

            yield xn1, yn1, zn1

    x, y, z = np.array(list(generate_roessler_attractor(100_000))).T

    fig, ax = plt.subplots()
    sp = ScatterPlot2D(x, y, alpha=0.01, ax=ax)
    ax.axis("off")
    plt.show()