pipeline/session.py

@@ -0,0 +1,353 @@
+"""What a session contains, decided by reading the frames rather than the names.
+
+The first version of this pipeline parsed filenames: it expected
+`calibrated-T32-...-Luminance-BIN2-W-300-001.fit` and broke on anything else.
+That works for exactly one telescope, one binning and one exposure length. Every
+other session in the archive - narrowband from 2024, an Ha-only run from 2021,
+an RGB set with no luminance at all - fails that pattern.
+
+So this module opens each frame and asks the header. FITS carries FILTER,
+EXPTIME, XBINNING, OBJECT and DATE-OBS precisely so that software does not have
+to guess from a filename, and iTelescope fills them in properly.
+
+The unit of work is a *session*: one target, one telescope, one night, in one
+directory. Everything else in the pipeline takes a Session and asks it what
+filters exist and which frames belong to each.
+"""
+import glob
+import json
+import os
+import zipfile
+from collections import defaultdict
+
+from astropy.io import fits
+
+import layout
+
+# Filters that carry broadband colour, in the order a colour image wants them.
+BROADBAND = ("Red", "Green", "Blue")
+LUMINANCE = ("Luminance", "Clear", "Lum", "L")
+# Narrowband, with the aliases different telescopes write into FILTER.
+NARROWBAND = {
+    "Ha": ("Ha", "H-alpha", "Halpha", "HA", "H_Alpha"),
+    "OIII": ("OIII", "O3", "Oiii", "O-III"),
+    "SII": ("SII", "S2", "Sii", "S-II"),
+}
+
+
+def canonical_filter(raw):
+    """Map whatever the header says onto a name the pipeline knows."""
+    if raw is None:
+        return "Unknown"
+    name = str(raw).strip()
+    for canon, aliases in NARROWBAND.items():
+        if name in aliases or name.lower() == canon.lower():
+            return canon
+    for canon in BROADBAND:
+        if name.lower() == canon.lower() or name.lower() == canon[0].lower():
+            return canon
+    if name.lower() in [a.lower() for a in LUMINANCE]:
+        return "Luminance"
+    return name or "Unknown"
+
+
+class Frame:
+    """One light frame and the header facts the pipeline needs."""
+
+    __slots__ = ("path", "filter", "exptime", "binning", "object", "date_obs",
+                 "telescope", "instrument", "shape", "scale", "calibrated")
+
+    def __init__(self, path, header, shape, calibrated=True):
+        self.calibrated = calibrated
+        self.path = path
+        self.filter = canonical_filter(header.get("FILTER"))
+        self.exptime = float(header.get("EXPTIME", header.get("EXPOSURE", 0)))
+        self.binning = int(header.get("XBINNING", 1))
+        self.object = str(header.get("OBJECT", "")).strip()
+        self.date_obs = str(header.get("DATE-OBS", ""))
+        self.telescope = str(header.get("TELESCOP", "")).strip()
+        self.instrument = str(header.get("INSTRUME", "")).strip()
+        self.shape = shape
+        # iTelescope writes the plate scale in; without it, solve.py estimates.
+        scale = header.get("HIERARCH iTelescopePlateScaleH")
+        try:
+            self.scale = float(scale) if scale is not None else None
+        except (TypeError, ValueError):
+            self.scale = None
+
+    @property
+    def name(self):
+        return os.path.basename(self.path)
+
+    def __repr__(self):
+        return (f"<{self.filter} {self.exptime:.0f}s bin{self.binning} "
+                f"{self.name}>")
+
+
+class Session:
+    """One target, one telescope, one night."""
+
+    def __init__(self, root):
+        self.root = os.path.abspath(root)
+        self.frames = []
+        self._by_filter = {}
+
+    # ---------------------------------------------------------------- ingest
+    @staticmethod
+    def _is_mirror_junk(path):
+        """True for artefacts a broken segmented download leaves behind.
+
+        Cyberduck writes a DIRECTORY named after the file it is fetching, and
+        leaves `.cyberducksegment` parts inside when interrupted. Those
+        directories are named `something.fit`, so anything that trusts the
+        extension will try to open a directory as a frame. A re-mirror over an
+        already-organised tree produces dozens of them.
+        """
+        name = os.path.basename(path).lower()
+        if ".cyberducksegment" in name or name.endswith((".part", ".partial")):
+            return True
+        parent = os.path.basename(os.path.dirname(path)).lower()
+        return parent.endswith((".fit", ".fits", ".tif", ".tiff", ".zip"))
+
+    def ingest(self, verbose=True):
+        """Put the session into the standard layout and uncompress the lights.
+
+        Handles the three states a downloaded session can be in: still zipped,
+        already extracted, or a mixture of both (which is what happens when
+        someone unzips a few by hand).
+        """
+        raw = os.path.join(self.root, layout.RAW)
+        cal = os.path.join(self.root, layout.CALIBRATED)
+        os.makedirs(raw, exist_ok=True)
+        os.makedirs(cal, exist_ok=True)
+
+        moved = extracted = skipped = duplicates = 0
+        # Loose files at the session root get filed first.
+        for path in sorted(glob.glob(os.path.join(self.root, "*"))):
+            if os.path.isdir(path):
+                continue
+            name = os.path.basename(path)
+            low = name.lower()
+            if self._is_mirror_junk(path):
+                continue
+            if low.endswith(".zip") or (low.startswith("jpeg-") and
+                                        low.endswith(".jpg")):
+                dest = os.path.join(raw, name)
+            elif low.endswith((".fit", ".fits", ".tif", ".tiff")):
+                dest = os.path.join(cal, name)
+            else:
+                continue
+            if not os.path.exists(dest):
+                os.replace(path, dest)
+                moved += 1
+            elif os.path.getsize(dest) == os.path.getsize(path):
+                # Already filed, and the mirror has put an identical copy back.
+                os.remove(path)
+                duplicates += 1
+
+        # iTelescope ships each frame twice: 'calibrated-*' with bias, dark
+        # and flat applied, and 'raw-*' without. Extracting both put two copies
+        # of every frame into the stack. Only the calibrated set is extracted;
+        # the raw archives stay zipped in raw/ as the untouched originals. If a
+        # session has no calibrated archives at all, the raw ones are used and
+        # the fact is recorded.
+        archives = sorted(glob.glob(os.path.join(raw, "*.zip")))
+        calibrated_archives = [z for z in archives
+                               if os.path.basename(z).lower().startswith(
+                                   "calibrated")]
+        self.used_raw = not calibrated_archives and bool(archives)
+        if self.used_raw:
+            print("    NOTE: no calibrated archives in this session - "
+                  "extracting the raw frames instead. They have NOT had bias, "
+                  "dark or flat correction applied.")
+        else:
+            archives = calibrated_archives
+
+        for zpath in archives:
+            try:
+                with zipfile.ZipFile(zpath) as zf:
+                    for info in zf.infolist():
+                        base = os.path.basename(info.filename)
+                        if not base.lower().endswith((".fit", ".fits", ".tif",
+                                                      ".tiff")):
+                            continue
+                        target = os.path.join(cal, base)
+                        if os.path.exists(target) and \
+                                os.path.getsize(target) == info.file_size:
+                            skipped += 1
+                            continue
+                        with zf.open(info) as src, open(target, "wb") as dst:
+                            while chunk := src.read(1 << 20):
+                                dst.write(chunk)
+                        extracted += 1
+            except zipfile.BadZipFile:
+                print(f"    WARNING: {os.path.basename(zpath)} is not a "
+                      f"readable zip - skipped")
+        if verbose:
+            print(f"    filed {moved}, extracted {extracted}, "
+                  f"already present {skipped}, "
+                  f"duplicate copies removed {duplicates}")
+        return dict(moved=moved, extracted=extracted, skipped=skipped,
+                    duplicates=duplicates)
+
+    # ----------------------------------------------------------------- scan
+    def scan(self, verbose=True):
+        """Read every calibrated light frame's header."""
+        cal = os.path.join(self.root, layout.CALIBRATED)
+        paths = sorted(p for p in glob.glob(os.path.join(cal, "*.fit")) +
+                       glob.glob(os.path.join(cal, "*.fits"))
+                       if os.path.isfile(p) and not self._is_mirror_junk(p))
+        self.frames = []
+        for path in paths:
+            try:
+                with fits.open(path, memmap=True) as hd:
+                    hdr = hd[0].header
+                    shape = hd[0].shape
+            except Exception as exc:                          # noqa: BLE001
+                print(f"    WARNING: cannot read {os.path.basename(path)}: "
+                      f"{exc}")
+                continue
+            # Skip anything that is not a light frame.
+            if str(hdr.get("IMAGETYP", "Light")).lower().startswith(
+                    ("dark", "bias", "flat")):
+                continue
+            self.frames.append(Frame(path, hdr,
+                                     shape, calibrated=bool(
+                                         str(hdr.get("CALSTAT", "")).strip())))
+
+        # If both calibrated and uncalibrated frames are present - which
+        # happens when a session was extracted before this rule existed - keep
+        # only the calibrated ones rather than stacking each frame twice.
+        if any(f.calibrated for f in self.frames):
+            dropped = [f for f in self.frames if not f.calibrated]
+            if dropped and verbose:
+                print(f"    ignoring {len(dropped)} uncalibrated duplicate "
+                      f"frames (no CALSTAT)")
+            self.frames = [f for f in self.frames if f.calibrated]
+
+        self._by_filter = defaultdict(list)
+        for f in self.frames:
+            self._by_filter[f.filter].append(f)
+        if verbose:
+            self.describe()
+        return self.frames
+
+    # ------------------------------------------------------------ properties
+    @property
+    def filters(self):
+        """Filters present, deepest first - the deepest becomes the reference."""
+        return sorted(self._by_filter,
+                      key=lambda f: -sum(x.exptime for x in self._by_filter[f]))
+
+    def by_filter(self, name):
+        return list(self._by_filter.get(name, []))
+
+    @property
+    def target(self):
+        names = [f.object for f in self.frames if f.object]
+        return max(set(names), key=names.count) if names else \
+            os.path.basename(os.path.dirname(self.root))
+
+    @property
+    def telescope(self):
+        names = [f.telescope for f in self.frames if f.telescope]
+        return max(set(names), key=names.count) if names else "unknown"
+
+    @property
+    def scale(self):
+        vals = [f.scale for f in self.frames if f.scale]
+        return sum(vals) / len(vals) if vals else None
+
+    @property
+    def palette(self):
+        """Which colour scheme this session's filters support."""
+        present = set(self._by_filter)
+        has_lum = "Luminance" in present
+        broad = [c for c in BROADBAND if c in present]
+        narrow = [n for n in ("SII", "Ha", "OIII") if n in present]
+        if len(broad) == 3 and has_lum:
+            return "LRGB"
+        if len(broad) == 3:
+            return "RGB"
+        if len(narrow) >= 3:
+            return "SHO"
+        if "Ha" in narrow and "OIII" in narrow:
+            return "HOO"
+        if len(present) == 1:
+            return "MONO"
+        if broad or narrow:
+            return "PARTIAL"
+        return "UNKNOWN"
+
+    def describe(self):
+        print(f"    target      {self.target}")
+        print(f"    telescope   {self.telescope}")
+        if self.scale:
+            print(f"    plate scale {self.scale:.4f} arcsec/px")
+        print(f"    palette     {self.palette}")
+        for name in self.filters:
+            fs = self._by_filter[name]
+            exps = sorted({f.exptime for f in fs})
+            bins = sorted({f.binning for f in fs})
+            total = sum(f.exptime for f in fs) / 60.0
+            shapes = sorted({f.shape for f in fs})
+            shape_txt = "/".join(f"{s[1]}x{s[0]}" for s in shapes)
+            print(f"      {name:10s} {len(fs):3d} frames  "
+                  f"{'/'.join(f'{e:.0f}' for e in exps):>12s} s  "
+                  f"bin{'/'.join(str(b) for b in bins)}  "
+                  f"{shape_txt:>11s}  {total:6.1f} min")
+
+    # ------------------------------------------------------------- manifest
+    def manifest(self):
+        return {
+            "root": self.root,
+            "target": self.target,
+            "telescope": self.telescope,
+            "plate_scale": self.scale,
+            "palette": self.palette,
+            "filters": {
+                name: {
+                    "frames": len(self._by_filter[name]),
+                    "total_exposure_s": sum(f.exptime
+                                            for f in self._by_filter[name]),
+                    "exposures": sorted({f.exptime
+                                         for f in self._by_filter[name]}),
+                    "binning": sorted({f.binning
+                                       for f in self._by_filter[name]}),
+                } for name in self.filters
+            },
+        }
+
+    def save_manifest(self):
+        path = os.path.join(self.root, "manifest.json")
+        with open(path, "w", encoding="utf-8") as fh:
+            json.dump(self.manifest(), fh, indent=2)
+        return path
+
+
+def discover(root):
+    """Find every session under a directory tree.
+
+    A session is any directory that directly contains light frames or their
+    archives, which matches how iTelescope's downloads are organised
+    (telescope / object / date).
+    """
+    found = []
+    for dirpath, dirnames, filenames in os.walk(root):
+        base = os.path.basename(dirpath)
+        if base.lower().endswith((".fit", ".fits", ".tif", ".tiff", ".zip"))                 or ".cyberducksegment" in base.lower():
+            dirnames[:] = []
+            continue
+        if base in (layout.RAW, layout.CALIBRATED, "intermediates",
+                    "final", "renderings", "science", "stacks"):
+            dirnames[:] = []
+            continue
+        has_data = any(f.lower().endswith((".zip", ".fit", ".fits"))
+                       for f in filenames)
+        has_layout = any(d in dirnames for d in (layout.RAW,
+                                                 layout.CALIBRATED))
+        if has_data or has_layout:
+            found.append(dirpath)
+            dirnames[:] = [d for d in dirnames
+                           if d in (layout.RAW, layout.CALIBRATED)]
+    return sorted(found)