starplot 0.13.0__py2.py3-none-any.whl → 0.15.0__py2.py3-none-any.whl

starplot/plotters/experimental.py

@@ -0,0 +1,171 @@
+import geopandas as gpd
+
+from shapely import MultiPolygon
+from shapely import (
+    MultiPoint,
+    intersection,
+    delaunay_triangles,
+    distance,
+)
+
+from starplot.data import DataFiles
+from starplot.data.constellations import (
+    CONSTELLATIONS_FULL_NAMES,
+    CONSTELLATION_HIP_IDS,
+)
+from starplot.styles import PathStyle
+
+
+class ExperimentalPlotterMixin:
+    def _constellation_borders(self):
+        from shapely import LineString, MultiLineString
+        from shapely.ops import unary_union
+
+        constellation_borders = gpd.read_file(
+            DataFiles.CONSTELLATIONS,
+            engine="pyogrio",
+            use_arrow=True,
+            bbox=self._extent_mask(),
+        )
+
+        if constellation_borders.empty:
+            return
+
+        geometries = []
+
+        for i, constellation in constellation_borders.iterrows():
+            geometry_types = constellation.geometry.geom_type
+
+            # equinox = LineString([[0, 90], [0, -90]])
+            """
+            Problems:
+                - Need to handle multipolygon borders too (SER)
+                - Shapely's union doesn't handle geodesy (e.g. TRI + AND)
+                - ^^ TRI is plotted with ra < 360, but AND has ra > 360
+                - ^^ idea: create union first and then remove duplicate lines?
+            
+                TODO: create new static data file of constellation border lines
+            """
+
+            if "Polygon" in geometry_types and "MultiPolygon" not in geometry_types:
+                polygons = [constellation.geometry]
+
+            elif "MultiPolygon" in geometry_types:
+                polygons = constellation.geometry.geoms
+
+            for p in polygons:
+                coords = list(zip(*p.exterior.coords.xy))
+                # coords = [(ra * -1, dec) for ra, dec in coords]
+
+                new_coords = []
+
+                for i, c in enumerate(coords):
+                    ra, dec = c
+                    if i > 0:
+                        if new_coords[i - 1][0] - ra > 60:
+                            ra += 360
+
+                        elif ra - new_coords[i - 1][0] > 60:
+                            new_coords[i - 1][0] += 360
+
+                    new_coords.append([ra, dec])
+
+                ls = LineString(new_coords)
+                geometries.append(ls)
+
+        mls = MultiLineString(geometries)
+        geometries = unary_union(mls)
+
+        for ls in list(geometries.geoms):
+            x, y = ls.xy
+
+            self.line(zip(x, y), self.style.constellation_borders)
+            # x, y = ls.xy
+            # newx = [xx * -1 for xx in list(x)]
+            # self.ax.plot(
+            #     # list(x),
+            #     newx,
+            #     list(y),
+            #     # **self._plot_kwargs(),
+            #     # transform=self._geodetic,
+            #     transform=self._plate_carree,
+            #     **style_kwargs,
+            # )
+
+    def _plot_constellation_labels_experimental(
+        self,
+        style: PathStyle = None,
+        labels: dict[str, str] = CONSTELLATIONS_FULL_NAMES,
+    ):
+        def sorter(g):
+            # higher score is better
+            d = distance(g.centroid, points_line.centroid)
+            # d = distance(g.centroid, constellation.boundary.centroid)
+            extent = abs(g.bounds[2] - g.bounds[0])
+            area = g.area / constellation.boundary.area
+            # return ((extent**3)) * area**2
+            # return ((extent**2) - (d/2)) * area**2
+            # print(str(extent) + " " + str(area) + " " + str(d))
+            return d**2 * -1
+            return (extent / 2 + area) - (d / 5)
+
+        for constellation in self.objects.constellations:
+            constellation_stars = [
+                s
+                for s in self.objects.stars
+                if s.constellation_id == constellation.iau_id and s.magnitude < 5
+            ]
+            constellation_line_stars = [
+                s
+                for s in self.objects.stars
+                if s.constellation_id == constellation.iau_id
+                and s.hip in CONSTELLATION_HIP_IDS[constellation.iau_id]
+            ]
+            points = MultiPoint([(s.ra, s.dec) for s in constellation_stars])
+            points_line = MultiPoint([(s.ra, s.dec) for s in constellation_line_stars])
+
+            triangles = delaunay_triangles(
+                geometry=points,
+                tolerance=2,
+            )
+            print(constellation.name + " " + str(len(triangles.geoms)))
+
+            polygons = []
+            for t in triangles.geoms:
+                try:
+                    inter = intersection(t, constellation.boundary)
+                except Exception:
+                    continue
+                if (
+                    inter.geom_type == "Polygon"
+                    and len(list(zip(*inter.exterior.coords.xy))) > 2
+                ):
+                    polygons.append(inter)
+
+            p_by_area = {pg.area: pg for pg in polygons}
+            polygons_sorted = [
+                p_by_area[k] for k in sorted(p_by_area.keys(), reverse=True)
+            ]
+
+            # sort by combination of horizontal extent and area
+            polygons_sorted = sorted(polygons_sorted, key=sorter, reverse=True)
+
+            if len(polygons_sorted) > 0:
+                i = 0
+                ra, dec = polygons_sorted[i].centroid.x, polygons_sorted[i].centroid.y
+            else:
+                ra, dec = constellation.ra, constellation.dec
+
+            text = labels.get(constellation.iau_id)
+            style = style or self.style.constellation.label
+            style.anchor_point = "center"
+            self.text(
+                text,
+                ra,
+                dec,
+                style,
+                hide_on_collision=self.hide_colliding_labels,
+                area=MultiPolygon(polygons_sorted[:3])
+                if len(polygons_sorted)
+                else constellation.boundary,
+            )

starplot/plotters/milkyway.py

@@ -0,0 +1,39 @@
+from shapely.ops import unary_union
+
+from starplot.data import db
+from starplot.styles import PolygonStyle
+from starplot.styles.helpers import use_style
+from starplot.geometry import unwrap_polygon_360
+from starplot.profile import profile
+
+
+class MilkyWayPlotterMixin:
+    @profile
+    @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
+        """
+        con = db.connect()
+        mw = con.table("milky_way")
+
+        extent = self._extent_mask()
+        result = mw.filter(mw.geometry.intersects(extent)).to_pandas()
+
+        mw_union = unary_union(
+            [unwrap_polygon_360(row.geometry) for row in result.itertuples()]
+        )
+
+        if mw_union.geom_type == "MultiPolygon":
+            polygons = mw_union.geoms
+        else:
+            polygons = [mw_union]
+
+        for p in polygons:
+            self.polygon(
+                geometry=p,
+                style=style,
+            )

starplot/plotters/stars.py

@@ -1,51 +1,26 @@
 from typing import Callable, Mapping
-from functools import cache
 
-from skyfield.api import Star as SkyfieldStar
-
-# import numpy as np
+import rtree
+import numpy as np
+from skyfield.api import Star as SkyfieldStar, wgs84
 
 from starplot import callables
-from starplot.data import bayer, stars, flamsteed
-from starplot.data.stars import StarCatalog, STAR_NAMES
+from starplot.data import stars
+from starplot.data.stars import StarCatalog
 from starplot.models.star import Star, from_tuple
 from starplot.styles import ObjectStyle, use_style
+from starplot.profile import profile
 
 
 class StarPlotterMixin:
-    @cache
-    def _load_stars(self, catalog, limiting_magnitude=None):
-        stardata = stars.load(catalog)
-
-        ra_buffer = (self.ra_max - self.ra_min) / 4
-        dec_buffer = (self.dec_max - self.dec_min) / 4
-
-        if limiting_magnitude is not None:
-            stardata = stardata[(stardata["magnitude"] <= limiting_magnitude)]
-
-        stardata = stardata[
-            (stardata["dec_degrees"] < self.dec_max + dec_buffer)
-            & (stardata["dec_degrees"] > self.dec_min - dec_buffer)
-        ]
-
-        if self.ra_max <= 24 and self.ra_min >= 0:
-            stardata = stardata[
-                (stardata["ra_hours"] < self.ra_max + ra_buffer)
-                & (stardata["ra_hours"] > self.ra_min - ra_buffer)
-            ]
-        elif self.ra_max > 24:
-            # handle wrapping
-            stardata = stardata[
-                (stardata["ra_hours"] > self.ra_min - ra_buffer)
-                | (stardata["ra_hours"] < self.ra_max - 24 + ra_buffer)
-            ]
-        elif self.ra_min < 0:
-            stardata = stardata[
-                (stardata["ra_hours"] > 24 + self.ra_min - ra_buffer)
-                | (stardata["ra_hours"] < self.ra_max + ra_buffer)
-            ]
-
-        return stardata
+    def _load_stars(self, catalog, filters=None):
+        extent = self._extent_mask()
+
+        return stars.load(
+            extent=extent,
+            catalog=catalog,
+            filters=filters,
+        )
 
     def _scatter_stars(self, ras, decs, sizes, alphas, colors, style=None, **kwargs):
         style = style or self.style.star
@@ -81,7 +56,7 @@ class StarPlotterMixin:
         self,
         star_objects: list[Star],
         star_sizes: list[float],
-        where_labels: list,
+        label_row_ids: list,
         style: ObjectStyle,
         labels: Mapping[str, str],
         bayer_labels: bool,
@@ -93,7 +68,7 @@ class StarPlotterMixin:
 
         # 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]):
+            if s._row_id not in label_row_ids:
                 continue
 
             if (
@@ -108,17 +83,21 @@ class StarPlotterMixin:
             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_fn is not None:
+                label = label_fn(s)
+            elif s.hip in labels:
+                label = labels.get(s.hip)
+            else:
+                label = s.name
+
+            bayer_desig = s.bayer
+            flamsteed_num = s.flamsteed
 
             if label:
                 self.text(
                     label,
                     s.ra,
                     s.dec,
-                    # style,
-                    # _offset(style, star_sizes[i]),
                     style=style.label.offset_from_marker(
                         marker_symbol=style.marker.symbol,
                         marker_size=star_sizes[i],
@@ -129,13 +108,13 @@ class StarPlotterMixin:
                 )
 
             if bayer_labels and bayer_desig:
-                _bayer.append((bayer_desig, s.ra, s.dec, star_sizes[i], s.hip))
+                _bayer.append((bayer_desig, s.ra, s.dec, star_sizes[i]))
 
-            if flamsteed_labels and flamsteed_num:
-                _flamsteed.append((flamsteed_num, s.ra, s.dec, star_sizes[i], s.hip))
+            if flamsteed_labels and flamsteed_num and not bayer_desig:
+                _flamsteed.append((flamsteed_num, s.ra, s.dec, star_sizes[i]))
 
         # Plot bayer/flamsteed
-        for bayer_desig, ra, dec, star_size, _ in _bayer:
+        for bayer_desig, ra, dec, star_size in _bayer:
             self.text(
                 bayer_desig,
                 ra,
@@ -149,9 +128,7 @@ class StarPlotterMixin:
                 gid="stars-label-bayer",
             )
 
-        for flamsteed_num, ra, dec, star_size, hip in _flamsteed:
-            if hip in bayer.hip:
-                continue
+        for flamsteed_num, ra, dec, star_size in _flamsteed:
             self.text(
                 flamsteed_num,
                 ra,
@@ -165,27 +142,26 @@ class StarPlotterMixin:
                 gid="stars-label-flamsteed",
             )
 
-    def _prepare_star_coords(self, df):
+    def _prepare_star_coords(self, df, limit_by_altaz=False):
         df["x"], df["y"] = (
             df["ra"],
             df["dec"],
         )
         return df
 
+    @profile
     @use_style(ObjectStyle, "star")
     def stars(
         self,
-        mag: float = 6.0,
-        catalog: StarCatalog = StarCatalog.HIPPARCOS,
+        where: list = None,
+        where_labels: list = None,
+        catalog: StarCatalog = StarCatalog.BIG_SKY_MAG11,
         style: ObjectStyle = None,
-        rasterize: bool = False,
         size_fn: Callable[[Star], float] = callables.size_by_magnitude,
-        alpha_fn: Callable[[Star], float] = callables.alpha_by_magnitude,
+        alpha_fn: Callable[[Star], float] = None,
         color_fn: Callable[[Star], str] = None,
         label_fn: Callable[[Star], str] = None,
-        where: list = None,
-        where_labels: list = None,
-        labels: Mapping[int, str] = STAR_NAMES,
+        labels: Mapping[int, str] = None,
         legend_label: str = "Star",
         bayer_labels: bool = False,
         flamsteed_labels: bool = False,
@@ -195,23 +171,26 @@ class StarPlotterMixin:
         """
         Plots stars
 
+        Labels for stars are determined in this order:
+
+        1. Return value from `label_fn`
+        2. Value for star's HIP id in `labels`
+        3. IAU-designated name, as listed in the [data reference](/data/star-designations/)
+
         Args:
-            mag: Limiting magnitude of stars to plot. For more control of what stars to plot, use the `where` kwarg. **Note:** if you pass `mag` and `where` then `mag` will be ignored
-            catalog: The catalog of stars to use: "hipparcos", "big-sky-mag11", or "big-sky" -- see [`StarCatalog`](/reference-data/#starplot.data.stars.StarCatalog) for details
+            where: A list of expressions that determine which stars to plot. See [Selecting Objects](/reference-selecting-objects/) for details.
+            where_labels: A list of expressions that determine which stars are labeled on the plot (this includes all labels: name, Bayer, and Flamsteed). If you want to hide **all** labels, then set this arg to `[False]`. See [Selecting Objects](/reference-selecting-objects/) for details.
+            catalog: The catalog of stars to use: "big-sky-mag11", or "big-sky" -- see [`StarCatalog`](/reference-data/#starplot.data.stars.StarCatalog) for details
             style: If `None`, then the plot's style for stars will be used
-            rasterize: If True, then the stars will be rasterized when plotted, which can speed up exporting to SVG and reduce the file size but with a loss of image quality
             size_fn: Callable for calculating the marker size of each star. If `None`, then the marker style's size will be used.
             alpha_fn: Callable for calculating the alpha value (aka "opacity") of each star. If `None`, then the marker style's alpha will be used.
             color_fn: Callable for calculating the color of each star. If `None`, then the marker style's color will be used.
             label_fn: Callable for determining the label of each star. If `None`, then the names in the `labels` kwarg will be used.
-            where: A list of expressions that determine which stars to plot. See [Selecting Objects](/reference-selecting-objects/) for details.
-            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`.
+            labels: A dictionary that maps a star's HIP id to the label that'll be plotted for that star. If `None`, then the star's IAU-designated name will be used.
             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.
             flamsteed_labels: If True, then Flamsteed number labels for stars will be plotted.
         """
-        self.logger.debug("Plotting stars...")
 
         # fallback to style if callables are None
         color_hex = (
@@ -222,62 +201,86 @@ class StarPlotterMixin:
         color_fn = color_fn or (lambda d: color_hex)
 
         where = where or []
+        where_labels = where_labels or []
+        stars_to_index = []
+        labels = labels or {}
+
+        star_results = self._load_stars(catalog, filters=where)
 
-        if where:
-            mag = None
+        star_results_labeled = star_results
+        for f in where_labels:
+            star_results_labeled = star_results_labeled.filter(f)
 
-        if labels is None:
-            labels = {}
+        label_row_ids = star_results_labeled.to_pandas()["rowid"].tolist()
+
+        stars_df = star_results.to_pandas()
+
+        if getattr(self, "projection", None) == "zenith":
+            # filter stars for zenith plots to only include those above horizon
+            self.location = self.earth + wgs84.latlon(self.lat, self.lon)
+            stars_apparent = (
+                self.location.at(self.timescale)
+                .observe(SkyfieldStar.from_dataframe(stars_df))
+                .apparent()
+            )
+            # we only need altitude
+            stars_alt, _, _ = stars_apparent.altaz()
+            stars_df["alt"] = stars_alt.degrees
+            stars_df = stars_df[stars_df["alt"] > 0]
         else:
-            labels = {**STAR_NAMES, **labels}
-
-        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)
-        stars_ra, stars_dec, _ = astrometric.radec()
-        nearby_stars_df["ra"], nearby_stars_df["dec"] = (
-            stars_ra.hours * 15,
-            stars_dec.degrees,
-        )
-        self._prepare_star_coords(nearby_stars_df)
+            nearby_stars = SkyfieldStar.from_dataframe(stars_df)
+            astrometric = self.earth.at(self.timescale).observe(nearby_stars)
+            stars_ra, stars_dec, _ = astrometric.radec()
+            stars_df["ra"], stars_df["dec"] = (
+                stars_ra.hours * 15,
+                stars_dec.degrees,
+            )
 
+        stars_df = self._prepare_star_coords(stars_df)
         starz = []
+        rtree_id = 1
 
-        for star in nearby_stars_df.itertuples():
-            obj = from_tuple(star, catalog)
+        for star in stars_df.itertuples():
+            data_xy = self._proj.transform_point(star.x, star.y, self._crs)
+            display_x, display_y = self.ax.transData.transform(data_xy)
 
-            if not all([e.evaluate(obj) for e in where]):
+            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
 
+            obj = from_tuple(star)
             size = size_fn(obj) * self.scale**2
             alpha = alpha_fn(obj)
             color = color_fn(obj) or style.marker.color.as_hex()
 
-            starz.append((star.x, star.y, size, alpha, color, obj))
+            if obj.magnitude < 5:
+                rtree_id += 1
+                # 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 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((rtree_id, bbox, None))
 
-            # Experimental code for keeping spatial index of plotted stars (for better label placement)
-            # if getattr(self, "_geodetic", None):
-            #     # TODO : clean up!
-            #     x, y = self._proj.transform_point(
-            #         star.ra * -1, star.dec, self._geodetic
-            #     )
-            #     x0, y0 = self.ax.transData.transform((x, y))
-
-            #     if (
-            #         x0 < 0
-            #         or y0 < 0
-            #         or obj.magnitude > 5
-            #         or np.isnan(x0)
-            #         or np.isnan(y0)
-            #     ):
-            #         continue
-            #     radius = 1 + (5 - obj.magnitude)
-            #     # radius = max(((size**0.5 / 2) / self.scale)/1.44 - 6, 0) #size / self.scale**2 / 200
-            #     self._stars_rtree.insert(
-            #         0,
-            #         np.array((x0 - radius, y0 - radius, x0 + radius, y0 + radius)),
-            #         obj=star.x,
-            #     )
+            starz.append((star.x, star.y, size, alpha, color, obj))
 
         starz.sort(key=lambda s: s[2], reverse=True)  # sort by descending size
 
@@ -303,19 +306,20 @@ class StarPlotterMixin:
             edgecolors=style.marker.edge_color.as_hex()
             if style.marker.edge_color
             else "none",
-            rasterized=rasterize,
         )
 
         self._add_legend_handle_marker(legend_label, style.marker)
 
-        if labels:
-            self._star_labels(
-                star_objects,
-                sizes,
-                where_labels,
-                style,
-                labels,
-                bayer_labels,
-                flamsteed_labels,
-                label_fn,
-            )
+        if stars_to_index:
+            self._stars_rtree = rtree.index.Index(stars_to_index)
+
+        self._star_labels(
+            star_objects,
+            sizes,
+            label_row_ids,
+            style,
+            labels,
+            bayer_labels,
+            flamsteed_labels,
+            label_fn,
+        )

starplot/profile.py

@@ -0,0 +1,16 @@
+import time
+
+
+def profile(func):
+    def wrapper(*args, **kwargs):
+        start = time.time()
+
+        result = func(*args, **kwargs)
+
+        duration = round(time.time() - start, 4)
+
+        args[0].logger.debug(f"{func.__name__} = {str(duration)} sec")
+
+        return result
+
+    return wrapper

starplot/settings.py

@@ -0,0 +1,26 @@
+import os
+
+from pathlib import Path
+
+
+def env(name, default):
+    return os.environ.get(name) or default
+
+
+STARPLOT_PATH = Path(__file__).resolve().parent
+"""Path of starplot source"""
+
+DATA_PATH = STARPLOT_PATH / "data" / "library"
+"""Path of starplot data"""
+
+DOWNLOAD_PATH = Path(env("STARPLOT_DOWNLOAD_PATH", str(DATA_PATH / "downloads")))
+"""Path for downloaded data"""
+
+DUCKDB_EXTENSION_PATH = Path(
+    env("STARPLOT_DUCKDB_EXTENSIONS_PATH", str(DATA_PATH / "duckdb-extensions"))
+)
+"""Path for DuckDB extensions"""
+
+
+RAW_DATA_PATH = STARPLOT_PATH.parent.parent / "raw"
+BUILD_PATH = STARPLOT_PATH.parent.parent / "build"

starplot/styles/__init__.py

@@ -2,3 +2,5 @@ from .base import *  # noqa: F401,F403
 from .helpers import *  # noqa: F401,F403
 
 # from .extensions import *  # noqa: F401
+
+import starplot.styles.extensions as style_extensions  # noqa: F401