starplot 0.20.5__py3-none-any.whl

starplot/plotters/gradients.py

@@ -0,0 +1,153 @@
+import numpy as np
+from matplotlib.colors import LinearSegmentedColormap
+from starplot.profile import profile
+from starplot.styles import GradientDirection
+
+
+class GradientBackgroundMixin:
+    """
+    Mixin class to handle adding gradients to plots.
+
+    Handles a variety of projections and can be inherited by HorizonPlot, OpticPlot,
+    and MapPlot for vertical, radial, and mollweide gradients. However, some more
+    obscure projections may throw errors if attempting to plot a gradient with them.
+    """
+
+    @profile
+    def _plot_gradient_background(
+        self, gradient_preset: list[tuple[float, str]]
+    ) -> None:
+        """
+        Adds a gradient background to the plot, beneath the GeoAxes.
+        The background_color of the map must be set as a RGBA value with full
+        transparency (e.g. #ffffff00) for this function to render the desired
+        result.
+
+        Args:
+            gradient_preset: A list of tuples (e.g. [(0.0, '#000000'), (1.0, '#000080')])
+                    where each tuple contains a position value [0-1] and a color
+                    value to describe the range of colors in the gradient.
+        """
+        direction = self._gradient_direction
+        reverse = True if direction == GradientDirection.RADIAL else False
+        cmap = self._create_colormap(gradient_preset, reverse=reverse)
+        background_ax = self._create_background_ax()
+        background_ax.set_axis_off()
+        self._background_ax = background_ax
+
+        X, Y, gradient = self._create_gradient_arrays()
+
+        # Radial specific axes adjustments
+        if self._gradient_direction == GradientDirection.RADIAL:
+            background_ax.set_ylim(Y.min(), Y.max() * 1.06)
+
+        # if getattr(self, "optic", None):
+        #     self._camera_optic_transform(background_ax)
+        # Camera specific axes adjustments
+        # if gradient_shape == "camera":
+        #     self._camera_optic_transform(background_ax)
+
+        # Render gradient
+        background_ax.pcolormesh(
+            X,
+            Y,
+            gradient,
+            cmap=cmap,
+            shading="gouraud",
+            rasterized=True,
+            zorder=0,
+            clip_path=self._background_clip_path,
+        )
+
+        # Set plot in self.ax's zorder to 1 so it appears above the gradient
+        self.ax.zorder = 1
+
+        # Event driven function so background_ax matches GeoAxes when plotted
+        self.ax.figure.canvas.mpl_connect(
+            "draw_event",
+            lambda event: background_ax.set_position(self.ax.get_position()),
+        )
+
+    def _create_colormap(
+        self, gradient_preset: list[tuple[float, str]], reverse: bool = False
+    ) -> LinearSegmentedColormap:
+        """Creates a matplotlib colormap from a gradient preset."""
+        positions, colors = zip(*gradient_preset)
+
+        if self._gradient_direction == GradientDirection.RADIAL:
+            positions = [p / 2 for p in positions]
+            positions[-1] = 1
+
+        from pydantic.color import Color
+
+        colors = [Color(c).as_hex() for c in colors]
+
+        cmap = LinearSegmentedColormap.from_list(
+            "custom_gradient", list(zip(positions, colors)), N=750
+        )
+        return cmap.reversed() if reverse else cmap
+
+    def _create_background_ax(self):
+        """Adds a set of axes to take the gradient image."""
+        bbox = self.ax.get_position()
+        projection = None
+
+        if self._gradient_direction == GradientDirection.RADIAL:
+            projection = "polar"
+        elif self._gradient_direction == GradientDirection.MOLLWEIDE:
+            projection = "mollweide"
+
+        return self.ax.figure.add_axes(bbox, zorder=0, projection=projection)
+
+    def _create_gradient_arrays(self):
+        """Creates arrays for the gradient placement and the gradient meshgrid."""
+        # Radial gradient
+        if self._gradient_direction == GradientDirection.RADIAL:
+            rad = np.linspace(0, 1, 50)
+            azm = np.linspace(0, 2 * np.pi, 100)
+            Y, X = np.meshgrid(rad, azm)
+            gradient = Y**2.0
+            return X, Y, gradient
+
+        # Mollweide gradient
+        if self._gradient_direction == GradientDirection.MOLLWEIDE:
+            return self._create_mollweide_gradient()
+
+        # Default Vertical Gradient
+        x_array = np.linspace(0, 1, 2)
+        y_array = np.linspace(0, 1, 750)
+        X, Y = np.meshgrid(x_array, y_array)
+        gradient = np.linspace(0, 1, 750).reshape(-1, 1)
+        gradient = np.repeat(gradient, 2, axis=1)
+        return X, Y, gradient
+
+    def _create_mollweide_gradient(self):
+        """Generate meshgrid and gradient for a mollweide projection."""
+        x = np.linspace(-np.pi, np.pi, 250)
+        y = np.linspace(-np.pi / 2, np.pi / 2, 250)
+        X, Y = np.meshgrid(x, y)
+        # Rotation matrix (ICRS → Galactic)
+        R = np.array(
+            [
+                [-0.0548755604162154, -0.8734370902348850, -0.4838350155487132],
+                [0.4941094278755837, -0.4448296299600112, 0.7469822444972189],
+                [-0.8676661490190047, -0.1980763734312015, 0.4559837761750669],
+            ]
+        )
+        # Equatorial unit vectors
+        cos_y = np.cos(Y)
+        eq = np.stack(
+            [cos_y * np.cos(X), cos_y * np.sin(X) * -1, np.sin(Y) * -1], axis=-1
+        )
+        # Rotate into Galactic coords
+        gal = eq @ R.T
+        # Gradient follows galactic latitude
+        gradient = np.arcsin(gal[..., 2])
+        return X, Y, gradient
+
+    def _camera_optic_transform(self, background_ax) -> None:
+        """Apply camera-specific axes transformations for the gradient."""
+        background_ax.set_xlim(-0.11, 1.11)
+        background_ax.set_ylim(-0.07, 1.07)
+        if self.optic.rotation == 0:
+            return

starplot/plotters/legend.py

@@ -0,0 +1,253 @@
+from typing import Callable
+
+import numpy as np
+from matplotlib.legend import Legend
+from matplotlib.lines import Line2D
+
+from starplot import callables
+from starplot.data.translations import translate
+from starplot.models.star import Star
+from starplot.styles import (
+    MarkerStyle,
+    LegendLocationEnum,
+    LegendStyle,
+)
+from starplot.styles.helpers import use_style
+
+
+class LegendPlotterMixin:
+    def _create_legend(self, handles, labels, title, style, set_anchor: bool = False):
+        style_kwargs = style.matplot_kwargs(self.scale)
+
+        target = self.ax
+
+        if style.location.startswith("outside"):
+            target = self.fig
+
+        style_kwargs["borderaxespad"] = -1 * style.padding
+
+        legend = Legend(
+            target,
+            handles=handles,
+            labels=labels,
+            title=title,
+            **style_kwargs,
+        )
+
+        legend.set_zorder(
+            # zorder is not a valid kwarg to legend(), so we have to set it afterwards
+            style.zorder
+        )
+        legend.get_title().set_color(style.font_color.as_hex())
+
+        if not set_anchor:
+            return legend
+
+        display_to_axes_transform = self.ax.transAxes.inverted()
+        origin_x, origin_y = self.ax.transAxes.transform((0, 0))
+        padding_x, padding_y = display_to_axes_transform.transform(
+            (origin_x + style.padding_x / 2, origin_y + style.padding_y / 2)
+        )
+
+        if style.location.startswith("outside"):
+            extent = legend.get_window_extent(renderer=self.fig.canvas.get_renderer())
+            min_x, min_y = display_to_axes_transform.transform(extent.min)
+            max_x, max_y = display_to_axes_transform.transform(extent.max)
+
+            baseline, _ = display_to_axes_transform.transform(
+                (origin_x + 200, origin_y + style.padding_y / 2)
+            )
+
+            padding_x += baseline
+            width = max_x - min_x + padding_x
+            # height = max_y - min_y
+            # top_x, top_y = display_to_figure_transform.transform(self.ax.transAxes.transform((1, 1)))
+
+            bbox = {
+                LegendLocationEnum.OUTSIDE_TOP_RIGHT: (1 + width, 1 - padding_y),
+                LegendLocationEnum.OUTSIDE_TOP_LEFT: (-1 * width, 1 - padding_y),
+                LegendLocationEnum.OUTSIDE_BOTTOM_RIGHT: (1 + width, padding_y),
+                LegendLocationEnum.OUTSIDE_BOTTOM_LEFT: (-1 * width, padding_y),
+            }.get(style.location)
+
+        else:
+            bbox = {
+                LegendLocationEnum.INSIDE_TOP_LEFT: (padding_x, 1 - padding_y),
+                LegendLocationEnum.INSIDE_TOP_RIGHT: (1 - padding_x, 1 - padding_y),
+                LegendLocationEnum.INSIDE_TOP: (0.5, 1 - padding_y),
+                LegendLocationEnum.INSIDE_BOTTOM_LEFT: (padding_x, padding_y),
+                LegendLocationEnum.INSIDE_BOTTOM_RIGHT: (1 - padding_x, padding_y),
+                LegendLocationEnum.INSIDE_BOTTOM: (0.5, padding_y),
+            }.get(style.location)
+
+        legend.set_bbox_to_anchor(
+            bbox=bbox,
+            transform=self.ax.transAxes,
+        )
+
+        return legend
+
+    @use_style(LegendStyle, "legend")
+    def legend(self, title: str = "Legend", style: LegendStyle = None):
+        """
+        Plots the legend.
+
+        If the legend is already plotted, then it'll be removed first and then plotted again. So, it's safe to call this function multiple times if you need to 'refresh' the legend.
+
+        Args:
+            title: Title of the legend, which will be plotted at the top
+            style: Styling of the legend. If None, then the plot's style (specified when creating the plot) will be used
+        """
+        if not self._legend_handles:
+            return
+
+        if self._legend:
+            self._legend.remove()
+
+        target = self.ax
+
+        title = translate(title, self.language)
+
+        if style.location.startswith("outside"):
+            target = self.fig
+
+        legend = self._create_legend(
+            handles=self._legend_handles.values(),
+            labels=self._legend_handles.keys(),
+            title=title,
+            style=style,
+            set_anchor=True,
+        )
+
+        target.add_artist(legend)
+
+        self._legend = legend
+        self._legend_target = target
+
+    def _add_to_legend(self, legend):
+        if not self._legend:
+            self.legend()
+
+        target = self._legend_target
+
+        legend_base = self._legend.get_children()[0]
+        legend_2 = legend.get_children()[0]
+
+        # empty legend for padding
+        empty = Legend(
+            target,
+            handles=[],
+            labels=[],
+            title="",
+            # **style_kwargs,
+        )
+
+        # add empty legend for padding between legend and scale
+        legend_base.get_children().extend(empty.get_children()[0].get_children())
+
+        legend_base.get_children().extend(legend_2.get_children())
+
+        legend_base.get_children().extend(empty.get_children()[0].get_children())
+
+        # target.add_artist(self._legend)
+
+    @use_style(LegendStyle, "legend")
+    def star_magnitude_scale(
+        self,
+        title: str = "Star Magnitude",
+        style: LegendStyle = None,
+        size_fn: Callable[[Star], float] = callables.size_by_magnitude,
+        label_fn: Callable[float, str] = lambda m: str(m),
+        start: float = -1,
+        stop: float = 9,
+        step: float = 1,
+        add_to_legend: bool = False,
+    ):
+        """
+        Plots a magnitude scale for stars.
+
+        !!! example "Experimental"
+
+            This is currently an "experimental" feature, which means it's likely to be changed and improved in upcoming versions of Starplot.
+            It also means the feature likely has limitations.
+
+            **Help us improve this feature by submitting feedback on [GitHub (open an issue)](https://github.com/steveberardi/starplot/issues) or chat with us on [Discord](https://discord.gg/WewJJjshFu). Thanks!**
+
+        !!! note "Current Limitations"
+            - Only supports size functions that determine size based on magnitude only
+            - Only supports default marker for stars (point)
+            - Labels can only be plotted to the right of the marker
+            - Does not automatically determine the magnitude range of the stars you already plotted
+
+        Args:
+            title: Title of the legend, which will be plotted at the top
+            style: Styling of the magnitude scale. If None, then the plot's `legend` style will be used
+            size_fn: Size function for the star markers
+            label_fn: Function to determine the label for each magnitude
+            start: Starting magnitude
+            stop: Stop point (exclusive)
+            step: Step-size of each scale entry (i.e. how much to increment each step)
+            add_to_legend: If True, the scale will be added to the bottom of the legend (and if the legend isn't already plotted, then it'll plot the legend)
+        """
+        target = self.ax
+
+        if style.location.startswith("outside"):
+            target = self.fig
+
+        title = translate(title, self.language)
+
+        def scale(
+            size_fn,
+            style: MarkerStyle,
+            label_fn,
+            start: float,
+            stop: float,
+            step: float = 1,
+        ):
+            for mag in np.arange(start, stop, step):
+                s = style.matplot_kwargs()
+                s["markersize"] = (
+                    size_fn(Star(pk=1, ra=0, dec=0, magnitude=mag, geometry=None))
+                    ** 0.5
+                ) * self.scale
+                label = label_fn(mag)
+                yield Line2D(
+                    [],
+                    [],
+                    **s,
+                    linestyle="None",
+                    label=label,
+                )
+
+        handles = [
+            h
+            for h in scale(
+                size_fn=size_fn,
+                style=self.style.star.marker,
+                label_fn=label_fn,
+                start=start,
+                stop=stop,
+                step=step,
+            )
+        ]
+        labels = [str(m) for m in np.arange(start, stop, step)]
+
+        scale = self._create_legend(
+            handles=handles,
+            labels=labels,
+            title=title,
+            style=style,
+            set_anchor=True if not add_to_legend else False,
+        )
+
+        if add_to_legend:
+            self._add_to_legend(scale)
+        else:
+            target.add_artist(scale)
+
+        # for text in magnitude_scale.get_texts():
+        #     text.set_ha("center") # horizontal alignment of text item
+        #     text.set_x(-85) # x-position
+        #     text.set_y(-90) # y-position
+
+        # self.ax.add_artist(magnitude_scale)

starplot/plotters/milkyway.py

@@ -0,0 +1,51 @@
+from shapely.ops import unary_union
+
+from ibis import _
+
+from starplot.data import db
+from starplot.data.catalogs import Catalog, MILKY_WAY
+from starplot.styles import PolygonStyle
+from starplot.styles.helpers import use_style
+from starplot.geometry import split_polygon_at_zero
+from starplot.profile import profile
+from starplot.models.milky_way import from_tuple
+
+
+class MilkyWayPlotterMixin:
+    @profile
+    @use_style(PolygonStyle, "milky_way")
+    def milky_way(self, style: PolygonStyle = None, catalog: Catalog = MILKY_WAY):
+        """
+        Plots the Milky Way
+
+        Args:
+            style: Styling of the Milky Way. If None, then the plot's style (specified when creating the plot) will be used
+            catalog: Catalog to use for Milky Way polygons
+        """
+        con = db.connect()
+        mw = catalog._load(connection=con, table_name="milky_way")
+        mw = mw.mutate(
+            geometry=_.geometry.cast("geometry"),  # cast WKB to geometry type
+        )
+
+        extent = self._extent_mask()
+        df = mw.filter(_.geometry.intersects(extent)).to_pandas()
+
+        milky_ways = [from_tuple(m) for m in df.itertuples()]
+
+        polygons = []
+        for milky_way in milky_ways:
+            polygons.extend(split_polygon_at_zero(milky_way.geometry))
+
+        mw_union = unary_union(polygons)
+
+        if mw_union.geom_type == "MultiPolygon":
+            polygons = mw_union.geoms
+        else:
+            polygons = [mw_union]
+
+        for p in polygons:
+            self.polygon(
+                geometry=p.buffer(0.001),
+                style=style,
+            )