starplot 0.20.5__py3-none-any.whl

starplot/plotters/arrow.py

@@ -0,0 +1,174 @@
+import random
+
+import numpy as np
+
+from matplotlib import patches
+from shapely import LineString, Polygon, affinity, ops
+
+from starplot.geometry import circle
+from starplot.profile import profile
+from starplot.coordinates import CoordinateSystem
+from starplot.styles import ArrowStyle
+from starplot.styles.helpers import use_style
+from starplot.utils import points_on_line
+
+
+class ArrowPlotterMixin:
+    def _to_axes(self, points):
+        ax_points = []
+
+        for ra, dec in points:
+            x, y = self._proj.transform_point(ra, dec, self._crs)
+            data_to_axes = self.ax.transData + self.ax.transAxes.inverted()
+            x_axes, y_axes = data_to_axes.transform((x, y))
+            ax_points.append([x_axes, y_axes])
+        return ax_points
+
+    def _to_display(self, points):
+        display_points = []
+
+        for ra, dec in points:
+            if self._coordinate_system == CoordinateSystem.RA_DEC:
+                ra *= -1
+
+            data_x, data_y = self._proj.transform_point(ra, dec, self._geodetic)
+            if np.isnan(data_x) or np.isnan(data_y):
+                continue
+
+            display_x, display_y = self.ax.transData.transform((data_x, data_y))
+            display_points.append((display_x, display_y))
+
+        return display_points
+
+    @profile
+    @use_style(ArrowStyle, "arrow")
+    def arrow(
+        self,
+        origin: tuple[float, float] = None,
+        target: tuple[float, float] = None,
+        style: ArrowStyle = None,
+        scale: float = 0.99,
+        length: float = 5,
+        max_attempts: int = 100,
+    ):
+        """
+        Plots an arrow from one point to another if you specify both `origin` and `target`. If you only specify a `target`, then the arrow will serve as a "pointer" to that target.
+
+        Args:
+            origin: Starting point (ra, dec)
+            target: Target of the arrow (ra, dec)
+            style: Style of the arrow
+            scale: Scaling factor for the arrow, to make it offset from the origin/target
+            length: If you only specify a target, then this will be the length of the arrow (in degrees). This value is ignored if you're plotting an arrow from one point to another.
+            max_attempts: If you only specify a target, then this will be the max number of attempts for plotting the arrow without colliding with labels. Arrow will be plotted on the final attempt.
+
+        """
+
+        def create_arrow(p0, p1):
+            ax_coords = self._to_axes([origin, target])
+            x, y = ax_coords[0]
+            target_x, target_y = ax_coords[1]
+
+            # angle to rotate head
+            angle_radians = np.atan2(target_y - y, target_x - x)
+            angle_degrees = np.degrees(angle_radians)
+
+            line_x = [
+                (1 - scale) * ax_coords[1][0] + scale * ax_coords[0][0],
+                (1 - scale) * ax_coords[0][0] + scale * ax_coords[1][0],
+            ]
+            line_y = [
+                (1 - scale) * ax_coords[1][1] + scale * ax_coords[0][1],
+                (1 - scale) * ax_coords[0][1] + scale * ax_coords[1][1],
+            ]
+
+            arrow_x = line_x[-1]
+            arrow_y = line_y[-1]
+
+            body_width = style.body_width / self.resolution * 0.5
+            head_width = style.head_width / self.resolution * 0.5
+            head_height = style.head_height / self.resolution * 0.5
+
+            arrow_line = LineString(list(zip(line_x, line_y)))
+            arrow_body = arrow_line.buffer(body_width, **style.shapely_kwargs())
+
+            arrow_head = Polygon(
+                [
+                    (arrow_x, arrow_y),
+                    (arrow_x + head_width, arrow_y - head_height),
+                    (arrow_x - head_width, arrow_y - head_height),
+                    (arrow_x, arrow_y),
+                ]
+            )
+
+            arrow_head = affinity.rotate(
+                arrow_head,
+                angle_degrees + 270,
+                origin=(arrow_x, arrow_y),
+            )
+
+            result = ops.split(arrow_line, arrow_head)
+            arrow_body = result.geoms[0].buffer(body_width, **style.shapely_kwargs())
+            return arrow_body.union(arrow_head.buffer(0.0001, **style.shapely_kwargs()))
+
+        if origin and target:
+            origin = self._prepare_coords(*origin)
+            target = self._prepare_coords(*target)
+            arrow_polygon = create_arrow(origin, target)
+
+        elif target:
+            # we only need to prepare the target coords because the origin will be
+            # based on the circle around the prepared target coords (so the origin will be automatically in the expected coord system)
+            target = self._prepare_coords(*target)
+            arrow_polygon = None
+            attempts = 0
+            padding = 8
+            polygon = circle(
+                center=target,
+                diameter_degrees=length * 2,
+                num_pts=200,
+            )
+
+            origins = list(zip(*polygon.exterior.coords.xy))
+
+            while arrow_polygon is None and attempts < max_attempts:
+                attempts += 1
+                origin = random.choice(origins)
+                display_points = self._to_display([origin, target])
+
+                if len(display_points) < 2:
+                    continue
+
+                points_arrow = points_on_line(display_points[0], display_points[1], 25)
+
+                # check if arrow body collides with any labels
+                collides_with_label = False
+                for x, y in points_arrow:
+                    bbox = (
+                        x - padding,
+                        y - padding,
+                        x + padding,
+                        y + padding,
+                    )
+                    if self._is_label_collision(bbox):
+                        collides_with_label = True
+                        break
+
+                # if arrow body does not collide with labels, we can create polygon and exit loop
+                if not collides_with_label or attempts == max_attempts:
+                    arrow_polygon = create_arrow(origin, target)
+
+        else:
+            raise ValueError(
+                "To plot an arrow you must specify a target or a target and origin."
+            )
+
+        patch = patches.Polygon(
+            list(zip(*arrow_polygon.exterior.coords.xy)),
+            **style.matplot_kwargs(self.scale),
+            transform=self.ax.transAxes,
+        )
+        self.ax.add_patch(patch)
+        # Need to set clip path AFTER adding patch
+        patch.set_clip_on(True)
+        patch.set_clip_path(self._background_clip_path)

starplot/plotters/constellations.py

@@ -0,0 +1,298 @@
+from typing import Callable
+
+import rtree
+from shapely import (
+    MultiPoint,
+)
+from matplotlib.collections import LineCollection
+from ibis import _
+
+from starplot.coordinates import CoordinateSystem
+from starplot.data import db, constellations as condata
+from starplot.data.catalogs import (
+    Catalog,
+    CONSTELLATIONS_IAU,
+    CONSTELLATION_BORDERS,
+    BIG_SKY_MAG11,
+)
+from starplot.data.stars import load as load_stars
+from starplot.models import Star, Constellation
+from starplot.models.constellation import from_tuple
+from starplot.profile import profile
+from starplot.styles import LineStyle, LabelStyle
+from starplot.styles.helpers import use_style
+from starplot.geometry import is_wrapped_polygon, line_segment, split_line_at_meridian
+from starplot.plotters.text import CollisionHandler
+
+
+class ConstellationPlotterMixin:
+    @profile
+    def _prepare_constellation_stars(
+        self, constellations: list[Constellation]
+    ) -> dict[int, tuple[float, float]]:
+        """
+        Returns dictionary of stars and their position:
+
+        {hip: (x,y)}
+
+        Where (x, y) is the plotted coordinate system (RA/DEC or AZ/ALT)
+        """
+        hips = []
+        for c in constellations:
+            hips.extend(c.star_hip_ids)
+
+        results = load_stars(
+            catalog=BIG_SKY_MAG11,
+            filters=[_.hip.isin(hips)],
+        )
+        results = results.select(
+            "ra",
+            "dec",
+            "epoch_year",
+            "hip",
+            "parallax_mas",
+            "ra_mas_per_year",
+            "dec_mas_per_year",
+        )
+        df = results.to_pandas()
+        df["ra_hours"], df["dec_degrees"] = (df.ra / 15, df.dec)
+        df = self._prepare_star_coords(df, limit_by_altaz=False)
+
+        return {star.hip: (star.x, star.y) for star in df.itertuples()}
+
+    @profile
+    @use_style(LineStyle, "constellation_lines")
+    def constellations(
+        self,
+        style: LineStyle = None,
+        where: list = None,
+        sql: str = None,
+        catalog: Catalog = CONSTELLATIONS_IAU,
+    ):
+        """Plots the constellation lines **only**. To plot constellation borders and/or labels, see separate functions for them.
+
+        **Important:** If you're plotting the constellation lines, then it's good to plot them _first_, because Starplot will use the constellation lines to determine where to place labels that are plotted afterwards (labels will look better if they're not crossing a constellation line).
+
+        Args:
+            style: Styling of the constellations. If None, then the plot's style (specified when creating the plot) will be used
+            where: A list of expressions that determine which constellations to plot. See [Selecting Objects](/reference-selecting-objects/) for details.
+            sql: SQL query for selecting constellations (table name is `_`). This query will be applied _after_ any filters in the `where` kwarg.
+            catalog: The catalog of constellations to use -- see [catalogs overview](/data/overview/) for details
+        """
+        self.logger.debug("Plotting constellation lines...")
+
+        where = where or []
+        ctr = 0
+
+        extent = self._extent_mask()
+        results = condata.load(extent=extent, filters=where, sql=sql, catalog=catalog)
+        constellations_df = results.to_pandas()
+
+        if constellations_df.empty:
+            return
+
+        constellations = [from_tuple(c) for c in constellations_df.itertuples()]
+
+        constellation_points_to_index = []
+        lines = []
+        constars = self._prepare_constellation_stars(constellations)
+
+        for c in constellations:
+            hiplines = c.star_hip_lines
+            inbounds = False
+
+            for s1_hip, s2_hip in hiplines:
+                if not constars.get(s1_hip) or not constars.get(s2_hip):
+                    continue
+
+                x1, y1 = constars.get(s1_hip)
+                x2, y2 = constars.get(s2_hip)
+
+                if x1 == x2 and y1 == y2:
+                    continue
+
+                if x1 - x2 > 60:
+                    x2 += 360
+                elif x2 - x1 > 60:
+                    x1 += 360
+
+                if not inbounds and (
+                    self._in_bounds_xy(x1, y1) or self._in_bounds_xy(x2, y2)
+                ):
+                    inbounds = True
+                elif not inbounds:
+                    continue
+
+                xy_lines = []
+
+                if x2 > 360:
+                    xy_lines = [*split_line_at_meridian((x1, y1), (x2, y2))]
+                elif x1 > 360:
+                    xy_lines = [*split_line_at_meridian((x2, y2), (x1, y1))]
+                else:
+                    xy_lines = [[(x1, y1), (x2, y2)]]
+
+                data_lines = [
+                    [self._proj.transform_point(x, y, self._crs) for x, y in line]
+                    for line in xy_lines
+                ]
+                display_lines = [
+                    [self.ax.transData.transform(p) for p in line]
+                    for line in data_lines
+                ]
+
+                lines.extend(data_lines)
+
+                radius = style.width * self.scale if style.width else 1
+
+                for display_start, display_end in display_lines:
+                    for x, y in line_segment(display_start, display_end, radius * 4):
+                        if x < 0 or y < 0:
+                            continue
+                        bbox = (
+                            x - radius,
+                            y - radius,
+                            x + radius,
+                            y + radius,
+                        )
+                        if self.debug_text:
+                            self._debug_bbox(bbox, color="#39FF14", width=0.5)
+
+                        constellation_points_to_index.append((ctr, bbox, None))
+                        ctr += 1
+
+            if inbounds:
+                self._objects.constellations.append(c)
+
+        line_collection = LineCollection(
+            lines,
+            clip_on=True,
+            clip_path=self._background_clip_path,
+            gid="constellations-line",
+            **style.matplot_line_collection_kwargs(self.scale),
+        )
+        self.ax.add_collection(line_collection)
+
+        if self._constellations_rtree.get_size() == 0:
+            self._constellations_rtree = rtree.index.Index(
+                constellation_points_to_index
+            )
+        else:
+            for bbox in constellation_points_to_index:
+                self._constellations_rtree.insert(
+                    0,
+                    bbox,
+                    None,
+                )
+
+    @profile
+    @use_style(LineStyle, "constellation_borders")
+    def constellation_borders(
+        self, style: LineStyle = None, catalog: Catalog = CONSTELLATION_BORDERS
+    ):
+        """Plots the constellation borders
+
+        Args:
+            style: Styling of the constellation borders. If None, then the plot's style (specified when creating the plot) will be used
+            catalog: Catalog to use for constellation borders
+        """
+        con = db.connect()
+        borders = catalog._load(connection=con, table_name="constellation_borders")
+        borders = borders.mutate(
+            geometry=_.geometry.cast("geometry"),  # cast WKB to geometry type
+        )
+
+        extent = self._extent_mask()
+        borders_df = borders.filter(_.geometry.intersects(extent)).to_pandas()
+
+        if borders_df.empty:
+            return
+
+        border_lines = []
+        geometries = [line.geometry for line in borders_df.itertuples()]
+
+        for ls in geometries:
+            if ls.length < 360:
+                ls = ls.segmentize(1)
+
+            xy = [c for c in ls.coords]
+
+            if self._coordinate_system == CoordinateSystem.RA_DEC:
+                border_lines.append(xy)
+
+            elif self._coordinate_system == CoordinateSystem.AZ_ALT:
+                coords = [self._prepare_coords(*p) for p in xy]
+                border_lines.append(coords)
+
+            else:
+                raise ValueError("Unrecognized coordinate system")
+
+        line_collection = LineCollection(
+            border_lines,
+            **style.matplot_line_collection_kwargs(self.scale),
+            transform=self._crs,
+            clip_on=True,
+            clip_path=self._background_clip_path,
+            gid="constellations-border",
+        )
+        self.ax.add_collection(line_collection)
+
+    @profile
+    @use_style(LabelStyle, "constellation_labels")
+    def constellation_labels(
+        self,
+        style: LabelStyle = None,
+        label_fn: Callable[[Constellation], str] = Constellation.get_label,
+        collision_handler: CollisionHandler = None,
+    ):
+        """
+        Plots constellation labels for all constellations that have been plotted. This means you must plot the constellations before plotting their labels.
+
+        It's good to plot these last because they're area-based labels (vs point-based, like for star names), and area-based labels have more freedom to move around. If you plot area-based labels first, then it would limit the available space for point-based labels.
+
+        Args:
+            style: Styling of the constellation labels. If None, then the plot's style (specified when creating the plot) will be used
+            label_fn: Callable for determining the label for each constellation. The default function returns the constellation's name in uppercase.
+            collision_handler: An instance of [CollisionHandler][starplot.CollisionHandler] that describes what to do on collisions with other labels, markers, etc. If `None`, then `CollisionHandler(allow_constellation_line_collisions=True)` will be used (**Important: this function does NOT default to the plot's collision handler, since it's the only area-based label function and collisions should be handled differently**).
+        """
+
+        collision_handler = collision_handler or self.area_label_handler
+
+        hips = []
+        for c in self.objects.constellations:
+            hips.extend(c.star_hip_ids)
+
+        all_constellation_stars = Star.find(where=[_.hip.isin(hips)])
+
+        for constellation in self.objects.constellations:
+            constellation_line_stars = [
+                s
+                for s in all_constellation_stars
+                if s.hip in constellation.star_hip_ids
+            ]
+            if not constellation_line_stars:
+                continue
+
+            if is_wrapped_polygon(constellation.boundary):
+                starpoints = []
+                ra, dec = zip(*[(s.ra, s.dec) for s in constellation_line_stars])
+                new_ra = [r - 360 if r > 300 else r for r in ra]
+                starpoints = list(zip(new_ra, dec))
+
+            else:
+                ra, dec = zip(*[(s.ra, s.dec) for s in constellation_line_stars])
+                starpoints = list(zip(ra, dec))
+
+            points_line = MultiPoint(starpoints)
+            centroid = points_line.centroid
+            text = label_fn(constellation)
+
+            self.text(
+                text,
+                centroid.x,
+                centroid.y,
+                style,
+                area=constellation.boundary,  # TODO : make this intersection with clip path
+                collision_handler=collision_handler,
+                gid="constellations-label-name",
+            )

starplot/plotters/debug.py

@@ -0,0 +1,21 @@
+from matplotlib.patches import Rectangle
+from matplotlib.transforms import Bbox
+
+
+class DebugPlotterMixin:
+    def _debug_bbox(self, bbox, color, width=1):
+        bboxe = Bbox([[bbox[0], bbox[1]], [bbox[2], bbox[3]]])
+        bboxe = bboxe.transformed(self.ax.transAxes.inverted())
+        rect = Rectangle(
+            (bboxe.x0, bboxe.y0),  # (x, y) position in display pixels
+            width=bboxe.width,
+            height=bboxe.height,
+            transform=self.ax.transAxes,
+            fill=False,
+            facecolor="none",
+            edgecolor=color,
+            linewidth=width,
+            alpha=1,
+            zorder=100_000,
+        )
+        self.ax.add_patch(rect)