starplot 0.12.5__py2.py3-none-any.whl → 0.14.0__py2.py3-none-any.whl

starplot/plotters/milkyway.py

@@ -0,0 +1,41 @@
+from starplot.data import DataFiles
+from starplot.styles import PolygonStyle
+from starplot.styles.helpers import use_style
+from starplot.utils import lon_to_ra
+
+
+class MilkyWayPlotterMixin:
+    @use_style(PolygonStyle, "milky_way")
+    def milky_way(self, style: PolygonStyle = None):
+        """
+        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
+        """
+
+        mw = self._read_geo_package(DataFiles.MILKY_WAY.value)
+
+        if mw.empty:
+            return
+
+        def _prepare_polygon(p):
+            points = list(zip(*p.boundary.coords.xy))
+            # convert lon to RA and reverse so the coordinates are counterclockwise order
+            return [(lon_to_ra(lon) * 15, dec) for lon, dec in reversed(points)]
+
+        # create union of all Milky Way patches
+        gs = mw.geometry.to_crs(self._plate_carree)
+        mw_union = gs.buffer(0.1).unary_union.buffer(-0.1)
+        polygons = []
+
+        if mw_union.geom_type == "MultiPolygon":
+            polygons.extend([_prepare_polygon(polygon) for polygon in mw_union.geoms])
+        else:
+            polygons.append(_prepare_polygon(mw_union))
+
+        for polygon_points in polygons:
+            self._polygon(
+                points=polygon_points,
+                style=style,
+            )

starplot/plotters/stars.py

@@ -1,13 +1,17 @@
 from typing import Callable, Mapping
 from functools import cache
 
+import rtree
 from skyfield.api import Star as SkyfieldStar
 
+import numpy as np
+
 from starplot import callables
-from starplot.data import bayer, stars
+from starplot.coordinates import CoordinateSystem
+from starplot.data import bayer, stars, flamsteed
 from starplot.data.stars import StarCatalog, STAR_NAMES
 from starplot.models.star import Star, from_tuple
-from starplot.styles import ObjectStyle, LabelStyle, use_style
+from starplot.styles import ObjectStyle, use_style
 
 
 class StarPlotterMixin:
@@ -64,6 +68,7 @@ class StarPlotterMixin:
             zorder=kwargs.pop("zorder", None) or style.marker.zorder,
             edgecolors=edge_colors,
             alpha=alphas,
+            gid="stars",
             **self._plot_kwargs(),
             **kwargs,
         )
@@ -77,24 +82,88 @@ class StarPlotterMixin:
     def _star_labels(
         self,
         star_objects: list[Star],
+        star_sizes: list[float],
         where_labels: list,
-        style: LabelStyle,
+        style: ObjectStyle,
         labels: Mapping[str, str],
         bayer_labels: bool,
+        flamsteed_labels: bool,
         label_fn: Callable[[Star], str],
     ):
-        for s in star_objects:
+        _bayer = []
+        _flamsteed = []
+
+        # Plot all star common names first
+        for i, s in enumerate(star_objects):
             if where_labels and not all([e.evaluate(s) for e in where_labels]):
                 continue
 
+            if (
+                s.hip
+                and s.hip in self._labeled_stars
+                or s.tyc
+                and s.tyc in self._labeled_stars
+            ):
+                continue
+            elif s.hip:
+                self._labeled_stars.append(s.hip)
+            elif s.tyc:
+                self._labeled_stars.append(s.tyc)
+
             label = labels.get(s.hip) if label_fn is None else label_fn(s)
             bayer_desig = bayer.hip.get(s.hip)
+            flamsteed_num = flamsteed.hip.get(s.hip)
 
             if label:
-                self.text(label, s.ra, s.dec, style)
+                self.text(
+                    label,
+                    s.ra,
+                    s.dec,
+                    style=style.label.offset_from_marker(
+                        marker_symbol=style.marker.symbol,
+                        marker_size=star_sizes[i],
+                        scale=self.scale,
+                    ),
+                    hide_on_collision=self.hide_colliding_labels,
+                    gid="stars-label-name",
+                )
 
             if bayer_labels and bayer_desig:
-                self.text(bayer_desig, s.ra, s.dec, self.style.bayer_labels)
+                _bayer.append((bayer_desig, s.ra, s.dec, star_sizes[i], s.hip))
+
+            if flamsteed_labels and flamsteed_num:
+                _flamsteed.append((flamsteed_num, s.ra, s.dec, star_sizes[i], s.hip))
+
+        # Plot bayer/flamsteed
+        for bayer_desig, ra, dec, star_size, _ in _bayer:
+            self.text(
+                bayer_desig,
+                ra,
+                dec,
+                style=self.style.bayer_labels.offset_from_marker(
+                    marker_symbol=style.marker.symbol,
+                    marker_size=star_size,
+                    scale=self.scale,
+                ),
+                hide_on_collision=self.hide_colliding_labels,
+                gid="stars-label-bayer",
+            )
+
+        for flamsteed_num, ra, dec, star_size, hip in _flamsteed:
+            if hip in bayer.hip:
+                continue
+            self.text(
+                flamsteed_num,
+                ra,
+                dec,
+                style=self.style.flamsteed_labels.offset_from_marker(
+                    marker_symbol=style.marker.symbol,
+                    marker_size=star_size,
+                    scale=self.scale,
+                ),
+                hide_on_collision=self.hide_colliding_labels,
+                gid="stars-label-flamsteed",
+            )
 
     def _prepare_star_coords(self, df):
         df["x"], df["y"] = (
@@ -119,6 +188,7 @@ class StarPlotterMixin:
         labels: Mapping[int, str] = STAR_NAMES,
         legend_label: str = "Star",
         bayer_labels: bool = False,
+        flamsteed_labels: bool = False,
         *args,
         **kwargs,
     ):
@@ -138,7 +208,8 @@ class StarPlotterMixin:
             where_labels: A list of expressions that determine which stars are labeled on the plot. See [Selecting Objects](/reference-selecting-objects/) for details.
             labels: A dictionary that maps a star's HIP id to the label that'll be plotted for that star. If you want to hide name labels, then set this arg to `None`.
             legend_label: Label for stars in the legend. If `None`, then they will not be in the legend.
-            bayer_labels: If True, then Bayer labels for stars will be plotted. Set this to False if you want to hide Bayer labels.
+            bayer_labels: If True, then Bayer labels for stars will be plotted.
+            flamsteed_labels: If True, then Flamsteed number labels for stars will be plotted.
         """
         self.logger.debug("Plotting stars...")
 
@@ -151,6 +222,7 @@ class StarPlotterMixin:
         color_fn = color_fn or (lambda d: color_hex)
 
         where = where or []
+        stars_to_index = []
 
         if where:
             mag = None
@@ -160,8 +232,6 @@ class StarPlotterMixin:
         else:
             labels = {**STAR_NAMES, **labels}
 
-        star_size_multiplier = self._size_multiplier * style.marker.size / 5
-
         nearby_stars_df = self._load_stars(catalog, mag)
         nearby_stars = SkyfieldStar.from_dataframe(nearby_stars_df)
         astrometric = self.ephemeris["earth"].at(self.timescale).observe(nearby_stars)
@@ -173,17 +243,66 @@ class StarPlotterMixin:
         self._prepare_star_coords(nearby_stars_df)
 
         starz = []
+        ctr = 0
 
         for star in nearby_stars_df.itertuples():
             obj = from_tuple(star, catalog)
+            ctr += 1
 
             if not all([e.evaluate(obj) for e in where]):
                 continue
 
-            size = size_fn(obj) * star_size_multiplier
+            if self._coordinate_system == CoordinateSystem.RA_DEC:
+                data_xy = self._proj.transform_point(
+                    star.ra * -1, star.dec, self._geodetic
+                )
+            elif self._coordinate_system == CoordinateSystem.AZ_ALT:
+                data_xy = self._proj.transform_point(star.x, star.y, self._crs)
+            else:
+                raise ValueError("Unrecognized coordinate system")
+
+            display_x, display_y = self.ax.transData.transform(data_xy)
+
+            if (
+                display_x < 0
+                or display_y < 0
+                or np.isnan(display_x)
+                or np.isnan(display_y)
+                or self._is_clipped([(display_x, display_y)])
+            ):
+                continue
+
+            size = size_fn(obj) * self.scale**2
             alpha = alpha_fn(obj)
             color = color_fn(obj) or style.marker.color.as_hex()
 
+            if obj.magnitude < 5:
+                # radius = ((size**0.5 / 2) / self.scale) #/ 3.14
+                radius = size**0.5 / 5
+
+                bbox = np.array(
+                    (
+                        display_x - radius,
+                        display_y - radius,
+                        display_x + radius,
+                        display_y + radius,
+                    )
+                )
+                # if obj.name == "Sirius":
+                #     print(bbox)
+                # if obj.name == "Electra":
+                #     print(bbox)
+
+                if self._stars_rtree.get_size() > 0:
+                    self._stars_rtree.insert(
+                        0,
+                        bbox,
+                        None,
+                    )
+                else:
+                    # if the index has no stars yet, then wait until end to load for better performance
+                    stars_to_index.append((ctr, bbox, None))
+
             starz.append((star.x, star.y, size, alpha, color, obj))
 
         starz.sort(key=lambda s: s[2], reverse=True)  # sort by descending size
@@ -215,6 +334,17 @@ class StarPlotterMixin:
 
         self._add_legend_handle_marker(legend_label, style.marker)
 
-        self._star_labels(
-            star_objects, where_labels, style.label, labels, bayer_labels, label_fn
-        )
+        if stars_to_index:
+            self._stars_rtree = rtree.index.Index(stars_to_index)
+
+        if labels:
+            self._star_labels(
+                star_objects,
+                sizes,
+                where_labels,
+                style,
+                labels,
+                bayer_labels,
+                flamsteed_labels,
+                label_fn,
+            )