sdss-almanac 0.2.1__py3-none-any.whl

almanac/data_models/utils.py

@@ -0,0 +1,185 @@
+import numpy as np
+from astropy.table import Table, hstack
+from scipy.spatial.distance import cdist
+from subprocess import check_output
+from typing import List, Tuple
+
+from almanac import utils
+from almanac.config import config
+from almanac.logger import logger
+
+def sanitise_twomass_designation(v):
+    # The target_ids seem to be styled '2MASS-J...'
+    v = str(v).strip()
+    v = v[5:] if v.startswith("2MASS") else v
+    v = str(v.lstrip("-Jdb_"))
+    if v.lower() == "na" or v == "None":
+        return ""
+    return v
+
+def match_planned_to_plugged(planned, plugged, tol=1e-5):
+
+    is_apogee = (
+        (planned["holetype"] == "APOGEE")
+    |   (planned["holetype"] == "APOGEE_SHARED")
+    |   (planned["holetype"] == "APOGEE_SOUTH")
+    )
+    if not any(is_apogee):
+        return []
+
+    planned = planned[is_apogee]
+    plugged = plugged[plugged["spectrographId"] == 2]
+
+    ra_dist = cdist(
+        planned["target_ra"].reshape((-1, 1)),
+        plugged["ra"].reshape((-1, 1)),
+    )
+    dec_dist = cdist(
+        planned["target_dec"].reshape((-1, 1)),
+        plugged["dec"].reshape((-1, 1)),
+    )
+
+    meets_tolerance = (ra_dist < tol) & (dec_dist < tol)
+    n_matches_to_plugged_holes = np.sum(meets_tolerance, axis=0)
+
+    N = np.sum(n_matches_to_plugged_holes > 1)
+    if N > 0:
+        raise RuntimeError("Cannot uniquely match plugged holes to planned holes!")
+
+    dist = np.sqrt(ra_dist**2 + dec_dist**2)
+    has_match = (n_matches_to_plugged_holes == 1)
+    planned_hole_indices = np.argmin(dist[:, has_match], axis=0)
+
+    return hstack(
+        [
+            plugged[has_match],
+            planned[planned_hole_indices]
+        ],
+        metadata_conflicts="silent",
+        uniq_col_name="{table_name}{col_name}",
+        table_names=("", "planned_")
+    )
+
+
+def get_headers(path, head=20_000):
+    keys = (
+        "FIELDID", "DESIGNID", "CONFIGID", "SEEING", "EXPTYPE",
+        "NREAD", "IMAGETYP", "LAMPQRTZ", "LAMPTHAR", "LAMPUNE", "FOCUS",
+        "NAME", "PLATEID", "CARTID", "MAPID", "PLATETYP", "OBSCMNT",
+        "COLLPIST", "COLPITCH", "DITHPIX", "TCAMMID", "TLSDETB",
+    )
+    keys_str = "|".join(keys)
+
+    commands = " | ".join(
+        ['hexdump -n {head} -e \'80/1 "%_p" "\\n"\' {path}', 'egrep "{keys_str}"']
+    ).format(head=head, path=path, keys_str=keys_str)
+    outputs = check_output(commands, shell=True, text=True)
+    outputs = outputs.strip().split("\n")
+
+    values = _parse_headers(outputs, keys)
+    return dict(zip(map(str.lower, keys), values))
+
+
+
+def _parse_headers(output: List[str], keys: Tuple[str, ...], default=None) -> List[str]:
+    """
+    Parse hexdump output to extract header key-value pairs.
+
+    :param output:
+        List of strings from hexdump output containing header information.
+    :param keys:
+        Tuple of header keys to look for in the output.
+    :param default:
+        Default value to use when a key is not found.
+
+    :returns:
+        List of header values corresponding to the input keys, with defaults for missing keys.
+    """
+    meta = [default] * len(keys)
+    for line in output:
+        try:
+            key, value = line.split("=", 2)
+        except ValueError:  # grep'd something in the data
+            continue
+
+        key = key.strip()
+        if key in keys:
+            index = keys.index(key)
+            if "/" in value:
+                # could be comment
+                *parts, comment = value.split("/")
+                value = "/".join(parts)
+
+            value = value.strip("' ")
+            meta[index] = value.strip()
+    return meta
+
+
+def get_exposure_path(observatory, mjd, prefix, exposure, chip):
+    return (
+        f"{config.apogee_dir}/"
+        f"{observatory}/"
+        f"{mjd}/"
+        f"{prefix}-{chip}-{get_exposure_string(mjd, exposure)}.apz"
+    )
+
+def mjd_to_exposure_prefix(mjd: int) -> int:
+    """Convert MJD to exposure prefix.
+
+    The exposure prefix is calculated as (MJD - 55562) * 10000, with a minimum of 0.
+
+    :param mjd:
+        Modified Julian Date (MJD) as an integer.
+
+    :returns:
+        Exposure prefix as an integer.
+    """
+    return max(0, (int(mjd) - 55_562) * 10_000)
+
+
+def get_exposure_string(mjd, exposure):
+    if isinstance(exposure, str):
+        return exposure
+    else:
+        return f"{mjd_to_exposure_prefix(mjd) + exposure:08d}"
+
+
+
+def input_id_to_designation(input_id: str) -> Tuple[str, str]:
+    """
+    Convert an input ID to a standardized designation format.
+
+    The input identifier might be a 2MASS-style designation (in many different
+    formats), or a Gaia DR2-style designation, or an input catalog identifier.
+
+    :param input_id:
+        The input ID string.
+
+    :returns:
+        A two-length tuple containing the designation type, and the cleaned
+        designation identifier.
+    """
+    cleaned = str(input_id).strip().lower()
+    if cleaned == "na":
+        return ("", "")
+    is_gaia = cleaned.startswith("gaia")
+    if is_gaia:
+        dr, source_id = cleaned.split(" ")
+        dr = dr.split("_")[1].lstrip("dr")
+        return (f"Gaia_DR{dr}", source_id)
+
+    is_twomass = cleaned.startswith("2m") or input_id.startswith("j")
+    if is_twomass:
+        if cleaned.startswith("2mass"):
+            cleaned = cleaned[5:]
+        if cleaned.startswith("2m"):
+            cleaned = cleaned[2:]
+        designation = str(cleaned.lstrip("-jdb_"))
+        return ("2MASS", designation)
+    else:
+        try:
+            catalogid = np.int64(cleaned)
+        except:
+            return ("Unknown", input_id)
+        else:
+            return ("catalog", cleaned)

almanac/database.py

@@ -0,0 +1,22 @@
+import sdssdb
+from sdssdb.peewee.sdss5db import database
+
+from time import time
+from almanac import config, logger
+from dataclasses import asdict
+
+# Create a temporary sdssdb profile based on almanac settings
+t = -time()
+sdssdb.config.update(almanac=asdict(config.sdssdb))
+is_database_available = database.set_profile("almanac", reuse_if_open=True)
+t += time()
+
+if not is_database_available:
+    logger.warning(f"Unable to connect to SDSS database after {t:.1f} seconds")
+elif t > float(config.database_connect_time_warning):
+    logger.warning(
+        f"Took {t:.1f} s to connect to SDSS database.\n"
+        f"You can suppress this warning with the `database_connect_time_warning` configuration."
+    )
+
+from sdssdb.peewee.sdss5db import catalogdb, opsdb

almanac/display.py

@@ -0,0 +1,422 @@
+import logging
+import numpy as np
+from itertools import cycle
+from datetime import datetime, timedelta
+from rich.console import Console
+from rich.table import Table
+from rich.live import Live
+from rich.text import Text
+from rich.align import Align
+from rich.table import Table as RichTable
+
+from almanac import config
+from almanac.data_models import Exposure
+from typing import Optional, List, Tuple, Dict, Any
+
+def mjd_to_datetime(mjd):
+    """Convert MJD to datetime - mock implementation"""
+    # MJD epoch is November 17, 1858
+    epoch = datetime(1858, 11, 17)
+    return epoch + timedelta(days=mjd)
+
+class BufferedHandler(logging.Handler):
+    """Custom logging handler that buffers log records"""
+    def __init__(self):
+        super().__init__()
+        self.buffer = []
+
+    def emit(self, record):
+        self.buffer.append(record)
+
+    def flush_to_console(self, console=None):
+        """Flush buffered records to console"""
+        if console is None:
+            console = Console()
+
+        for record in self.buffer:
+            log_message = self.format(record)
+
+            # Color code based on log level
+            if record.levelno >= logging.ERROR:
+                style = "red"
+            elif record.levelno >= logging.WARNING:
+                style = "yellow"
+            elif record.levelno >= logging.INFO:
+                style = "blue"
+            else:
+                style = "dim"
+
+            console.print(log_message, style=style)
+
+        self.buffer.clear()
+
+class ObservationsDisplay:
+    color_outside_range = "black"
+    color_unknown = "white"
+    color_no_data = "bright_black"
+    color_apo = "dodger_blue3"
+    color_lco = "green4"
+    color_both = "purple4"
+    color_missing = "red"
+
+    def __init__(self, mjd_min, mjd_max, observatories=("apo", "lco")):
+        self.console = Console()
+        self.start_date = mjd_to_datetime(mjd_min) if isinstance(mjd_min, (int, float)) else mjd_min
+        self.end_date = mjd_to_datetime(mjd_max) if isinstance(mjd_max, (int, float)) else mjd_max
+        self.days_per_week = 7
+
+        # Track completion status for each day
+        self.completed = dict(apo=set(), lco=set())
+        self.no_data = dict(apo=set(), lco=set())
+        self.missing = set()
+        self.offset = 0
+        # Setup logging buffer
+        self.log_buffer = BufferedHandler()
+        self.log_buffer.setFormatter(logging.Formatter(
+            '%(asctime)s - %(name)s - %(levelname)s - %(message)s',
+            datefmt='%H:%M:%S'
+        ))
+
+        # Calculate the grid dimensions based on the date range
+        self._setup_grid()
+        self.observatories = observatories
+
+    def _setup_grid(self):
+        """Setup the grid based on start and end dates, organized by years"""
+        # Calculate which years we need to display
+        self.start_year = self.start_date.year
+        self.end_year = self.end_date.year
+        self.year_rows = list(range(self.start_year, self.end_year + 1))
+
+        # For each year, we'll create a full year grid (Jan 1 to Dec 31)
+        # but pad the beginning if our data doesn't start on Jan 1
+        self.year_grids = {}
+        self.dates = []
+
+        self.offset = 0
+        for i, year in enumerate(self.year_rows):
+            year_start = datetime(year, 1, 1)
+            year_end = datetime(year, 12, 31)
+
+            # Find Monday of the week containing Jan 1
+            jan1_weekday = year_start.weekday()  # 0=Monday, 6=Sunday
+            grid_start = year_start - timedelta(days=jan1_weekday)
+            if i == 0:
+                self.offset = (self.start_date - grid_start).days
+
+            # Find Sunday of the week containing Dec 31
+            dec31_weekday = year_end.weekday()  # 0=Monday, 6=Sunday
+            days_to_sunday = 6 - dec31_weekday
+            grid_end = year_end + timedelta(days=days_to_sunday)
+
+            # Generate dates for this year's grid
+            year_dates = []
+            current_date = grid_start
+            while current_date <= grid_end:
+                year_dates.append(current_date)
+                current_date += timedelta(days=1)
+
+            # Calculate weeks for this year
+            total_days = len(year_dates)
+            weeks = total_days // 7
+
+            self.year_grids[year] = {
+                'dates': year_dates,
+                'weeks': weeks,
+                'grid_start': grid_start,
+                'grid_end': grid_end,
+                'year_start': year_start,
+                'year_end': year_end
+            }
+
+            # Add to master dates list with offset
+            for i, date in enumerate(year_dates):
+                self.dates.append(date)
+
+
+    def get_day_color(self, day_index):
+        """Return the color for a given day based on completion status"""
+        if day_index >= len(self.dates):
+            return self.color_outside_range
+
+        date = self.dates[day_index]
+
+        # Only show colored squares for dates within our actual range
+        if date < self.start_date or date > self.end_date:
+            return self.color_outside_range
+
+        if day_index in self.missing:
+            return self.color_missing
+        elif day_index in self.completed["apo"] and day_index in self.completed["lco"]:
+            return self.color_both
+        elif day_index in self.completed["apo"]:
+            return self.color_apo
+        elif day_index in self.completed["lco"]:
+            return self.color_lco
+        elif day_index in self.no_data["apo"] and day_index in self.no_data["lco"]:
+            return self.color_no_data
+        else:
+            return self.color_unknown
+
+    def _get_month_headers_for_year(self, year):
+        """Generate month headers for a specific year"""
+        year_data = self.year_grids[year]
+        weeks = year_data['weeks']
+        year_dates = year_data['dates']
+
+        headers = ["   "]  # Space for day labels
+        current_month = None
+        text_to_add = ""
+
+        for week in range(weeks):
+            week_start_index = week * 7
+
+            if week_start_index < len(year_dates):
+                week_date = year_dates[week_start_index]
+
+                # Only show month headers for dates within the actual year
+                if week_date.year == year:
+                    month_abbr = week_date.strftime("%b")
+
+                    # Only show month if it's different from previous week
+                    if current_month != month_abbr:
+                        headers.append(Text(f"{month_abbr[:1]}", style="dim"))
+                        text_to_add = month_abbr[1:]
+                        current_month = month_abbr
+                    else:
+                        if len(text_to_add) > 0:
+                            headers.append(Text(text_to_add[:1], style="dim"))
+                            text_to_add = text_to_add[1:]
+                        else:
+                            headers.append(Text(" "))
+                else:
+                    # This is padding (before Jan 1 or after Dec 31)
+                    headers.append(Text(" "))
+            else:
+                headers.append(Text(" "))
+
+        return headers
+
+    def create_contributions_grid_for_year(self, year):
+        """Create the contributions grid for a specific year"""
+        year_data = self.year_grids[year]
+        weeks = year_data['weeks']
+        year_dates = year_data['dates']
+
+        table = Table.grid(padding=0)
+
+        # Add columns for day labels and each week in this year
+        table.add_column()  # For day labels
+        for _ in range(weeks):
+            table.add_column()
+
+        # Add month headers
+        month_headers = self._get_month_headers_for_year(year)
+        table.add_row(*month_headers)
+
+        # Create rows for each day of the week
+        day_names = ["S", "M", "T", "W", "T", "F", "S"]
+
+        for day_of_week in range(self.days_per_week):
+            row = [Text(day_names[day_of_week].ljust(3), style="dim")]
+
+            for week in range(weeks):
+                day_index_in_year = week * 7 + day_of_week
+
+                if day_index_in_year < len(year_dates):
+                    date = year_dates[day_index_in_year]
+
+                    # Find this date in our master dates list to get the right index
+                    master_day_index = None
+                    for i, master_date in enumerate(self.dates):
+                        if master_date == date:
+                            master_day_index = i
+                            break
+
+                    if master_day_index is not None:
+                        color = self.get_day_color(master_day_index)
+
+                        # Show square only for dates within the actual year and our date range
+                        if (date.year == year and
+                            self.start_date <= date <= self.end_date):
+                            square = Text("■", style=color)
+                        elif date.year == year:
+                            # Within the year but outside our date range
+                            square = Text("■", style=self.color_no_data)
+                        else:
+                            # Padding dates (before Jan 1 or after Dec 31)
+                            square = Text(" ")
+                    else:
+                        square = Text(" ")
+                else:
+                    square = Text(" ")
+
+                row.append(square)
+
+            table.add_row(*row)
+
+        return table
+
+    def create_display(self):
+        """Create the complete display with title and yearly grids"""
+        date_range = f"{self.start_date.strftime('%b %d, %Y')} - {self.end_date.strftime('%b %d, %Y')}"
+        title = Text("SDSS/APOGEE Observations", style="bold white")
+        subtitle = Text(date_range, style="dim")
+
+        # Create legend
+        legend = Table.grid(padding=(0, 1))
+        legend.add_column()
+        legend.add_column()
+        legend.add_column()
+        legend.add_column()
+        legend.add_column()
+
+        items = [
+            Text("■", style=self.color_no_data),
+            Text("None", style="dim"),
+        ]
+        if "apo" in self.observatories:
+            items.extend([
+                Text("■", style=self.color_apo),
+                Text("APO", style="dim"),
+            ])
+        if "lco" in self.observatories:
+            items.extend([
+                Text("■", style=self.color_lco),
+                Text("LCO", style="dim"),
+            ])
+
+        if "apo" in self.observatories and "lco" in self.observatories:
+            items.extend([
+                Text("■", style=self.color_both),
+                Text("Both", style="dim"),
+            ])
+        legend.add_row(*items)
+
+        # Combine everything
+        main_table = Table.grid()
+        main_table.add_column()
+        main_table.add_row(Align.center(title))
+        main_table.add_row(Align.center(subtitle))
+        main_table.add_row("")
+
+        # Add each year's grid with year header
+        for i, year in enumerate(self.year_rows):
+            # Add year header
+            year_header = Text(str(year), style="bold cyan")
+            main_table.add_row(Align.left(year_header))
+
+            # Add the grid for this year
+            year_grid = self.create_contributions_grid_for_year(year)
+            main_table.add_row(Align.center(year_grid))
+
+            # Add spacing between years except after the last one
+            if i < len(self.year_rows) - 1:
+                main_table.add_row("")
+
+        main_table.add_row("")
+        main_table.add_row(Align.center(legend))
+
+        return main_table
+
+    def add_observation(self, date, observatory):
+        """Add an observation for a specific date and observatory"""
+        # Find the day index for this date
+        for i, grid_date in enumerate(self.dates):
+            if grid_date.date() == date.date():
+                self.completed[observatory].add(i)
+                break
+
+
+def display_exposures(
+    exposures: List[Exposure],
+    sequences: Optional[Dict[str, List[Tuple[int, int]]]] = None,
+    console: Optional[Console] = None,
+    header_style: str = "bold cyan",
+    column_names: Optional[List[str]] = None,
+    sequence_styles: Tuple[str, ...] = ("green", "yellow"),
+    missing_style: str = "blink bold red",
+    title_style: str = "bold blue",
+) -> None:
+    """Display exposure information using Rich table formatting.
+
+    Args:
+        exposures: List of Exposure objects containing exposure data
+        sequences: Dictionary mapping sequence names to lists of (start, end) tuples (default: None)
+        console: Rich Console instance (default: None, creates new one)
+        header_style: Style for table headers (default: "bold cyan")
+        sequence_styles: Tuple of styles to cycle through for sequences (default: ("green", "yellow"))
+        missing_style: Style for missing/error entries (default: "red")
+        title_style: Style for the table title (default: "bold blue")
+    """
+    if console is None:
+        console = Console()
+
+    if len(exposures) == 0:
+        return
+
+    # Create the title
+    observatory, mjd = (exposures[0].observatory, exposures[0].mjd)
+    title = f"{len(exposures)} exposures from {observatory.upper()} on MJD {mjd}"
+
+    # Create Rich table
+    rich_table = RichTable(title=title, title_style=title_style, show_header=True, header_style=header_style)
+
+    field_names = config.display_field_names
+
+    for field_name in field_names:
+        rich_table.add_column(field_name, justify="center")
+
+    # Prepare sequence tracking
+    flattened_sequences = []
+    for k, v in (sequences or dict()).items():
+        flattened_sequences.extend(v)
+    flattened_sequences = np.array(flattened_sequences)
+
+    sequence_styles_cycle = cycle(sequence_styles)
+    in_sequence, current_sequence_style = (False, next(sequence_styles_cycle))
+
+    # Add rows to the table
+    for i, exposure in enumerate(exposures, start=1):
+        # Check if this row is part of a sequence
+        row_style = None
+        end_of_sequence = None
+        if len(flattened_sequences) > 0:
+            try:
+                j, k = np.where(flattened_sequences == i)
+            except:
+                pass
+            else:
+                # Could be start or end of sequence, and could be out of order
+                start_of_sequence = 0 in k
+                end_of_sequence = 1 in k
+
+                if start_of_sequence:
+                    in_sequence = True
+                    current_sequence_style = next(sequence_styles_cycle)
+                elif end_of_sequence:  # only end of sequence
+                    in_sequence = True
+
+        # Determine row style
+        if in_sequence:
+            row_style = current_sequence_style
+        else:
+            # Check if it's missing or has issues
+            if exposure.image_type == "missing":
+                row_style = missing_style
+
+        # Convert row data to strings and apply styling if needed
+        row_data = []
+        for field_name in field_names:
+            value = getattr(exposure, field_name)
+            if row_style:
+                row_data.append(Text(f"{value}", style=row_style))
+            else:
+                row_data.append(f"{value}")
+
+        rich_table.add_row(*row_data)
+        if end_of_sequence:
+            in_sequence = False
+
+    console.print(rich_table)
+    console.print()  # Add a blank line after the table