26598415173

Committed 28 May 2026 07:50PM UTC coverage: 79.865% (+0.2%) from 79.63%

Build # 26598415173

Build Type

Pull #146

github

Committed by

web-flow

Commit Message

Merge 40c542451 into 7038301bf

Pull Request Pull Request #146: Sigurd's smart-split

Coverage Stats

478 of 740 branches covered (64.59%)

Branch coverage included in aggregate %.

219 of 267 new or added lines in 2 files covered. (82.02%)

3203 of 3869 relevant lines covered (82.79%)

0.83 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

82.83

/src/socm/workflows/ml_null_tests/smart_split_null_test.py

import warnings
from math import ceil
from pathlib import Path
from typing import Dict, List, Optional, Union

import numpy as np
from scipy.ndimage import uniform_filter
from sotodlib.core import Context

from socm.workflows.ml_null_tests import NullTestWorkflow

# Simons Observatory site (Atacama, Chile) — consistent with the other null-test workflows
_SO_NAME = "San Pedro de Atacama"
_SO_COUNTRY = "Chile"
_SO_TIMEZONE = "America/Santiago"
_SO_LAT = -22.91    # degrees North
_SO_LON = -68.2     # degrees East
_SO_ALT = 5190.0    # metres above sea level

# Fields unique to this workflow that must be stripped before spawning child NullTestWorkflows
_SMART_FIELDS = frozenset({
    "nsplits", "block", "mode", "nopt", "opt_mode",
    "scanpat_tol", "constraint", "rad", "res", "weight", "prefix",
})


def _label_unique(patterns: np.ndarray, atol: float) -> np.ndarray:
    """Assign integer labels so rows within *atol* of an existing centroid share a label."""
    labels = np.full(len(patterns), -1, dtype=int)
    centroids: List[np.ndarray] = []
    for i, row in enumerate(patterns):
        for lid, cent in enumerate(centroids):
            if np.all(np.abs(row - cent) <= atol):
                labels[i] = lid
                break
        else:
            labels[i] = len(centroids)
            centroids.append(row.copy())
    return labels


class SmartSplitNullTestWorkflow(NullTestWorkflow):
    """
    A smart-split null-test workflow.

    Ports the algorithm from tenki/smartsplit.py for SO/sotodlib, replacing
    enlib/enact with astropy for sky-coordinate transforms and scipy for
    hitmap smoothing.  Observations are grouped into blocks, sorted largest-
    first, then assigned to splits with a greedy hitmap-score algorithm.  An
    optional sequence of relocate-based optimisation passes further improves
    balance.

    New parameters relative to the other null-test workflows
    ---------------------------------------------------------
    block : str
        Grouping mode.  ``"day"`` (default) or ``"tod"``, optionally with a
        colon-separated size multiplier, e.g. ``"day:2"`` for 2-day blocks.
    mode : str
        ``"plain"``, ``"crosslink"`` (default), or ``"scanpat"``.
        In crosslink mode rising and setting scans are balanced independently.
        In scanpat mode each unique (az, el, az_throw) scanning pattern is
        treated as a separate sub-population.
    nopt : int
        Number of block-relocate optimisation passes (default 2000).
    opt_mode : str
        ``"linear"`` (default) cycles through free blocks in order;
        ``"random"`` draws random block indices.
    scanpat_tol : float
        Tolerance in degrees for grouping scan patterns (only used when
        ``mode="scanpat"``).
    constraint : str or None
        Path to a directory that contains ``smart_split_N/query.txt`` files
        from a previous run.  Blocks whose observations appear in those files
        are pre-assigned and excluded from optimisation.
    rad : float
        Tophat smoothing radius in degrees applied to per-block hitmaps
        (default 0.7).
    res : float
        Sky-map pixel size in degrees (default 0.5).
    weight : str
        ``"plain"`` weights each scan by its duration; ``"det"`` weights by
        ``n_samples`` (proxy for detector count × duration).
    prefix : str or None
        Optional string prepended to virtual array names in the output.
    """

    nsplits: int = 4
    block: str = "day"
    mode: str = "crosslink"
    nopt: int = 2000
    opt_mode: str = "linear"
    scanpat_tol: float = 1.0
    constraint: Optional[str] = None
    rad: float = 0.7
    res: float = 0.5
    weight: str = "plain"
    prefix: Optional[str] = None
    name: str = "smart_split_null_test_workflow"

    # ------------------------------------------------------------------
    # Block construction
    # ------------------------------------------------------------------

    def _parse_block_params(self):
        """Return (block_mode, block_size) from the ``block`` field."""
        toks = self.block.split(":")
        return toks[0], float(toks[1]) if len(toks) > 1 else 1.0

    def _build_blocks(
        self, obs_info: Dict[str, Dict[str, Union[float, str]]]
    ) -> List[List[str]]:
        block_mode, block_size = self._parse_block_params()
        if block_mode == "tod":
            return [[oid] for oid in obs_info]
        elif block_mode == "day":
            groups: Dict[int, List[str]] = {}
            for oid, meta in obs_info.items():
                key = int(int(meta["start_time"]) // 86400 // block_size)
                groups.setdefault(key, []).append(oid)
            return [groups[k] for k in sorted(groups)]
        else:
            raise ValueError(f"Unknown block mode: {block_mode!r}. Expected 'day' or 'tod'.")

    # ------------------------------------------------------------------
    # Virtual array assignment
    # ------------------------------------------------------------------

    def _get_anames(
        self, obs_info: Dict[str, Dict[str, Union[float, str]]]
    ):
        """Return ``(aname_per_obs, unique_array_list, obs_to_array_index)``."""
        base = f"{self.prefix}_arr" if self.prefix else "arr"

        if self.mode == "crosslink":
            anames = {
                oid: base + ("r" if float(meta["az_center"]) % 360 < 180 else "s")
                for oid, meta in obs_info.items()
            }
        elif self.mode == "scanpat":
            obs_ids = list(obs_info.keys())
            patterns = np.array([
                [float(obs_info[oid]["az_center"]),
                 float(obs_info[oid]["el_center"]),
                 float(obs_info[oid]["az_throw"])]
                for oid in obs_ids
            ])
            pids = _label_unique(patterns, atol=self.scanpat_tol)
            anames = {oid: f"{base}p{pids[i]}" for i, oid in enumerate(obs_ids)}
        else:  # plain
            anames = {oid: base for oid in obs_info}

        arrays = sorted(set(anames.values()))
        ais = {oid: arrays.index(anames[oid]) for oid in obs_info}
        return anames, arrays, ais

    # ------------------------------------------------------------------
    # Hitmap construction
    # ------------------------------------------------------------------

    def _compute_scan_radec(
        self, obs_info: Dict[str, Dict[str, Union[float, str]]]
    ) -> Dict[str, tuple]:
        """Batch-convert all observation scan paths from AltAz to ICRS.

        Returns a dict mapping obs_id to ``(ra_deg_array, dec_deg_array)``.
        Requires an internet connection the first time (IERS table download);
        subsequent runs use the astropy cache.
        """
        import astropy.units as u
        from astral import LocationInfo
        from astropy.coordinates import AltAz, EarthLocation, SkyCoord
        from astropy.time import Time

        city = LocationInfo(_SO_NAME, _SO_COUNTRY, _SO_TIMEZONE, _SO_LAT, _SO_LON)
        site = EarthLocation(
            lat=city.latitude * u.deg, lon=city.longitude * u.deg, height=_SO_ALT * u.m
        )

        obs_ids: List[str] = list(obs_info.keys())
        all_az: List[np.ndarray] = []
        all_el: List[np.ndarray] = []
        all_t: List[np.ndarray] = []
        npts_per_obs: List[int] = []

        for oid in obs_ids:
            meta = obs_info[oid]
            az_c = float(meta["az_center"])
            el_c = float(meta["el_center"])
            az_throw = float(meta["az_throw"])
            dur = float(meta["duration"])
            t_mid = float(meta["start_time"]) + dur / 2.0
            n_az = max(2, int(abs(az_throw) / self.res + 1))
            azs = np.linspace(az_c - az_throw / 2.0, az_c + az_throw / 2.0, n_az)
            all_az.append(azs)
            all_el.append(np.full(n_az, el_c))
            all_t.append(np.full(n_az, t_mid))
            npts_per_obs.append(n_az)

        az_arr = np.concatenate(all_az)
        el_arr = np.concatenate(all_el)
        t_arr = np.concatenate(all_t)

        # Mirror elevation > 90° (same correction as the rest of the null-test
        # workflows): reflect through the zenith and rotate azimuth by 180°.
        high_el = el_arr > 90.0
        el_arr[high_el] = 180.0 - el_arr[high_el]
        az_arr[high_el] = az_arr[high_el] + 180.0

        times = Time(t_arr, format="unix", scale="utc")
        frame = AltAz(obstime=times, location=site)
        altaz_coords = SkyCoord(az=az_arr * u.deg, alt=el_arr * u.deg, frame=frame)
        icrs = altaz_coords.transform_to("icrs")
        ras: np.ndarray = np.asarray(icrs.ra.deg, dtype=np.float64)
        decs: np.ndarray = np.asarray(icrs.dec.deg, dtype=np.float64)

        result: Dict[str, tuple] = {}
        i = 0
        for oid, npts in zip(obs_ids, npts_per_obs):
            result[oid] = (ras[i:i + npts], decs[i:i + npts])
            i += npts
        return result

    def _build_grid(self, scan_radec: Dict[str, tuple]):
        """Determine the RA/Dec grid parameters from all scan paths.

        Returns ``(ra_min, dec_min, ny, nx, ra_median)`` where
        ``ra_min`` / ``dec_min`` are the lower-left corners in degrees.
        """
        all_ras = np.concatenate([r for r, _ in scan_radec.values()])
        all_decs = np.concatenate([d for _, d in scan_radec.values()])

        # Unwrap RA around the median to handle the 0°/360° boundary.
        ra_med = float(np.median(all_ras))
        ra_uw = all_ras.copy()
        ra_uw[ra_uw < ra_med - 180.0] += 360.0
        ra_uw[ra_uw > ra_med + 180.0] -= 360.0

        pad = max(self.rad * 2.0, self.res)
        ra_min = float(np.min(ra_uw)) - pad
        ra_max = float(np.max(ra_uw)) + pad
        dec_min = float(np.min(all_decs)) - pad
        dec_max = float(np.max(all_decs)) + pad

        ny = max(1, int(ceil((dec_max - dec_min) / self.res)))
        nx = max(1, int(ceil((ra_max - ra_min) / self.res)))
        return ra_min, dec_min, ny, nx, ra_med

    def _obs_hitmap(
        self,
        oid: str,
        obs_info: Dict[str, Dict[str, Union[float, str]]],
        scan_radec: Dict[str, tuple],
        ra_min: float,
        dec_min: float,
        ny: int,
        nx: int,
        ra_med: float,
    ) -> np.ndarray:
        """Rasterise one observation onto a ``(ny, nx)`` hitmap."""
        ras, decs = scan_radec[oid]
        meta = obs_info[oid]
        dur = float(meta["duration"])
        n_samples = int(meta["n_samples"])

        ra_uw = ras.copy()
        ra_uw[ra_uw < ra_med - 180.0] += 360.0
        ra_uw[ra_uw > ra_med + 180.0] -= 360.0

        xi = np.round((ra_uw - ra_min) / self.res).astype(int)
        yi = np.round((decs - dec_min) / self.res).astype(int)
        valid = (xi >= 0) & (xi < nx) & (yi >= 0) & (yi < ny)
        xi, yi = xi[valid], yi[valid]

        w = (n_samples if self.weight == "det" else dur) / max(1, len(ras))

        hm = np.zeros((ny, nx), dtype=np.float64)
        np.add.at(hm, (yi, xi), w)
        return hm

    def _build_block_hits(
        self,
        blocks: List[List[str]],
        obs_info: Dict[str, Dict[str, Union[float, str]]],
        ais: Dict[str, int],
        narray: int,
        scan_radec: Dict[str, tuple],
        ra_min: float,
        dec_min: float,
        ny: int,
        nx: int,
        ra_med: float,
    ) -> np.ndarray:
        """Build per-block, per-array hitmaps ``hits[nblock, narray, ny, nx]``."""
        nblock = len(blocks)
        hits = np.zeros((nblock, narray, ny, nx), dtype=np.float64)
        smooth_size = 2 * max(1, int(self.rad / self.res)) + 1

        for bi, block in enumerate(blocks):
            raw = np.zeros((narray, ny, nx), dtype=np.float64)
            for oid in block:
                raw[ais[oid]] += self._obs_hitmap(
                    oid, obs_info, scan_radec, ra_min, dec_min, ny, nx, ra_med
                )
            for ai in range(narray):
                if raw[ai].any():
                    raw[ai] = uniform_filter(raw[ai], size=smooth_size, mode="constant")
            hits[bi] = raw

        return hits

    def _build_mask(self, hits: np.ndarray, narray: int) -> np.ndarray:
        """Build coverage mask ``[narray, ny, nx]`` for well-sampled pixels."""
        ny, nx = hits.shape[-2], hits.shape[-1]
        mask = np.zeros((narray, ny, nx), dtype=bool)
        for ai in range(narray):
            ahits = hits[:, ai]
            nonzero = ahits[ahits > 0]
            if not nonzero.size:
                continue
            ref = float(np.median(nonzero))
            nblock_per_pix = np.sum(ahits > ref * 0.2, axis=0)
            pos = nblock_per_pix[nblock_per_pix > 0]
            nbref = float(np.median(pos)) if pos.size else 1.0
            nblock_lim = min(2 * self.nsplits, nbref * 0.2)
            mask[ai] = nblock_per_pix > nblock_lim
        return mask

    @staticmethod
    def _delta_score(
        split_hits: np.ndarray, bhits: np.ndarray, mask: np.ndarray
    ) -> np.ndarray:
        """Score each split by its fractional coverage gain from adding *bhits*.

        ``split_hits`` is ``[nsplit, narray, ny, nx]``;
        ``bhits`` and ``mask`` are ``[narray, ny, nx]``.
        Returns ``[nsplit]``.
        """
        with warnings.catch_warnings():
            warnings.simplefilter("ignore", RuntimeWarning)
            ratio = bhits / split_hits   # broadcasts to [nsplit, narray, ny, nx]
        # Replace only NaN (0/0: no hits on either side → no contribution).
        # Inf (something/0: empty split) is left for the min-cap so that empty
        # splits receive the maximum possible score — matching the original.
        ratio = np.where(np.isnan(ratio), 0.0, ratio)
        ratio = np.minimum(ratio, 1000.0)
        flat_mask = mask.ravel()
        ratio_flat = ratio.reshape(split_hits.shape[0], -1)
        return np.sum(ratio_flat[:, flat_mask], axis=-1)

    # ------------------------------------------------------------------
    # Constraint loading
    # ------------------------------------------------------------------

    def _load_constraint(
        self, blocks: List[List[str]]
    ) -> np.ndarray:
        """Return block ownership array (``-1`` = free) from *constraint* dir."""
        block_ownership = np.full(len(blocks), -1, dtype=int)
        if not self.constraint:
            return block_ownership

        constraint_dir = Path(self.constraint)
        for split_idx in range(self.nsplits):
            qfile = constraint_dir / f"smart_split_{split_idx + 1}" / "query.txt"
            if not qfile.exists():
                continue
            with open(qfile) as fh:
                existing = {line.strip() for line in fh if line.strip()}
            for bi, block in enumerate(blocks):
                overlap = set(block) & existing
                if not overlap:
                    continue
                if block_ownership[bi] >= 0 and block_ownership[bi] != split_idx:
                    raise ValueError(
                        f"Block {bi} has conflicting ownership in constraint directory."
                    )
                block_ownership[bi] = split_idx

        return block_ownership

    # ------------------------------------------------------------------
    # Core split algorithm
    # ------------------------------------------------------------------

    def _get_splits(
        self, ctx: Context, obs_info: Dict[str, Dict[str, Union[float, str]]]
    ) -> List[List[str]]:
        """Assign observations to *nsplits* balanced groups.

        Algorithm
        ---------
        1. Group observations into blocks; sort blocks largest-first.
        2. Compute RA/Dec scan paths and build per-block hitmaps.
        3. Pre-assign blocks that appear in a constraint directory.
        4. Greedily assign free blocks to the split with the highest
           fractional coverage gain.
        5. Run *nopt* relocate-optimisation passes.
        """
        if len(obs_info) < self.nsplits:
            obs_list = list(obs_info.keys())
            splits: List[List[str]] = [[] for _ in range(self.nsplits)]
            for i, oid in enumerate(obs_list):
                splits[i % self.nsplits].append(oid)
            return splits

        blocks = sorted(self._build_blocks(obs_info), key=len, reverse=True)

        _, arrays, ais = self._get_anames(obs_info)
        narray = len(arrays)

        block_ownership = self._load_constraint(blocks)
        fixed_blocks = np.where(block_ownership >= 0)[0]
        free_blocks = np.where(block_ownership < 0)[0]
        nfree = len(free_blocks)

        scan_radec = self._compute_scan_radec(obs_info)
        ra_min, dec_min, ny, nx, ra_med = self._build_grid(scan_radec)
        hits = self._build_block_hits(
            blocks, obs_info, ais, narray, scan_radec, ra_min, dec_min, ny, nx, ra_med
        )
        mask = self._build_mask(hits, narray)

        split_hits = np.zeros((self.nsplits, narray, ny, nx), dtype=np.float64)
        block_assignments = block_ownership.copy()

        # Pre-assign fixed blocks.
        for bi in fixed_blocks:
            split_hits[int(block_assignments[bi])] += hits[bi]

        # Greedy initial assignment for free blocks.
        for bi in free_blocks:
            score = self._delta_score(split_hits, hits[bi], mask)
            best = int(np.argmax(score))
            block_assignments[bi] = best
            split_hits[best] += hits[bi]

        # Relocate-based optimisation.
        nopt_actual = self.nopt if self.nsplits > 1 else 0
        rng = np.random.default_rng(seed=0)
        if self.opt_mode == "linear":
            opt_order = np.arange(nopt_actual) % max(1, nfree)
        else:  # "random"
            opt_order = rng.integers(0, max(1, nfree), nopt_actual)

        for step in range(nopt_actual):
            if nfree == 0:
                break
            bi = int(free_blocks[opt_order[step]])
            scur = int(block_assignments[bi])
            bhits = hits[bi]

            split_hits[scur] = np.maximum(0.0, split_hits[scur] - bhits)
            score = self._delta_score(split_hits, bhits, mask)
            best = int(np.argmax(score))
            split_hits[best] += bhits
            block_assignments[bi] = best

        result: List[List[str]] = [[] for _ in range(self.nsplits)]
        for bi, block in enumerate(blocks):
            result[int(block_assignments[bi])].extend(block)

        return result

    # ------------------------------------------------------------------
    # Workflow factory
    # ------------------------------------------------------------------

    @classmethod
    def get_workflows(cls, desc: Optional[dict] = None) -> List[NullTestWorkflow]:
        """Create one NullTestWorkflow per non-empty smart-split group."""
        smart_workflow = cls(**(desc or {}))

        # Strip smartsplit-specific fields so child NullTestWorkflows do not
        # receive unknown keyword arguments or spurious command-line flags.
        base_desc = {
            k: v
            for k, v in smart_workflow.model_dump(exclude_unset=True).items()
            if k not in _SMART_FIELDS
        }

        workflows: List[NullTestWorkflow] = []
        for split in smart_workflow._splits:
            if not split:
                continue
            desc_copy = dict(base_desc)
            split_idx = len(workflows) + 1
            desc_copy["name"] = f"smart_split_{split_idx}_null_test_workflow"
            desc_copy["output_dir"] = f"{smart_workflow.output_dir}/smart_split_{split_idx}"
            desc_copy["datasize"] = 0
            query_file = Path(desc_copy["output_dir"]) / "query.txt"
            query_file.parent.mkdir(parents=True, exist_ok=True)
            with open(query_file, "w") as fh:
                for oid in split:
                    fh.write(f"{oid}\n")
            desc_copy["query"] = f"file://{query_file.absolute()!s}"
            desc_copy["chunk_nobs"] = 1
            workflows.append(NullTestWorkflow(**dict(desc_copy)))

        return workflows

1	import warnings	1✔
2	from math import ceil	1✔
3	from pathlib import Path	1✔
4	from typing import Dict, List, Optional, Union	1✔
5
6	import numpy as np	1✔
7	from scipy.ndimage import uniform_filter	1✔
8	from sotodlib.core import Context	1✔
9
10	from socm.workflows.ml_null_tests import NullTestWorkflow	1✔
11
12	# Simons Observatory site (Atacama, Chile) — consistent with the other null-test workflows
13	_SO_NAME = "San Pedro de Atacama"	1✔
14	_SO_COUNTRY = "Chile"	1✔
15	_SO_TIMEZONE = "America/Santiago"	1✔
16	_SO_LAT = -22.91 # degrees North	1✔
17	_SO_LON = -68.2 # degrees East	1✔
18	_SO_ALT = 5190.0 # metres above sea level	1✔
19
20	# Fields unique to this workflow that must be stripped before spawning child NullTestWorkflows
21	_SMART_FIELDS = frozenset({	1✔
22	"nsplits", "block", "mode", "nopt", "opt_mode",
23	"scanpat_tol", "constraint", "rad", "res", "weight", "prefix",
24	})
25
26
27	def _label_unique(patterns: np.ndarray, atol: float) -> np.ndarray:	1✔
28	"""Assign integer labels so rows within atol of an existing centroid share a label."""
29	labels = np.full(len(patterns), -1, dtype=int)	1✔
30	centroids: List[np.ndarray] = []	1✔
31	for i, row in enumerate(patterns):	1✔
32	for lid, cent in enumerate(centroids):	1✔
33	if np.all(np.abs(row - cent) <= atol):	1✔
34	labels[i] = lid	1✔
35	break	1✔
36	else:
37	labels[i] = len(centroids)	1✔
38	centroids.append(row.copy())	1✔
39	return labels	1✔
40
41
42	class SmartSplitNullTestWorkflow(NullTestWorkflow):	1✔
43	"""
44	A smart-split null-test workflow.
45
46	Ports the algorithm from tenki/smartsplit.py for SO/sotodlib, replacing
47	enlib/enact with astropy for sky-coordinate transforms and scipy for
48	hitmap smoothing. Observations are grouped into blocks, sorted largest-
49	first, then assigned to splits with a greedy hitmap-score algorithm. An
50	optional sequence of relocate-based optimisation passes further improves
51	balance.
52
53	New parameters relative to the other null-test workflows
54	---------------------------------------------------------
55	block : str
56	Grouping mode. ``"day"`` (default) or ``"tod"``, optionally with a
57	colon-separated size multiplier, e.g. ``"day:2"`` for 2-day blocks.
58	mode : str
59	``"plain"``, ``"crosslink"`` (default), or ``"scanpat"``.
60	In crosslink mode rising and setting scans are balanced independently.
61	In scanpat mode each unique (az, el, az_throw) scanning pattern is
62	treated as a separate sub-population.
63	nopt : int
64	Number of block-relocate optimisation passes (default 2000).
65	opt_mode : str
66	``"linear"`` (default) cycles through free blocks in order;
67	``"random"`` draws random block indices.
68	scanpat_tol : float
69	Tolerance in degrees for grouping scan patterns (only used when
70	``mode="scanpat"``).
71	constraint : str or None
72	Path to a directory that contains ``smart_split_N/query.txt`` files
73	from a previous run. Blocks whose observations appear in those files
74	are pre-assigned and excluded from optimisation.
75	rad : float
76	Tophat smoothing radius in degrees applied to per-block hitmaps
77	(default 0.7).
78	res : float
79	Sky-map pixel size in degrees (default 0.5).
80	weight : str
81	``"plain"`` weights each scan by its duration; ``"det"`` weights by
82	``n_samples`` (proxy for detector count × duration).
83	prefix : str or None
84	Optional string prepended to virtual array names in the output.
85	"""
86
87	nsplits: int = 4	1✔
88	block: str = "day"	1✔
89	mode: str = "crosslink"	1✔
90	nopt: int = 2000	1✔
91	opt_mode: str = "linear"	1✔
92	scanpat_tol: float = 1.0	1✔
93	constraint: Optional[str] = None	1✔
94	rad: float = 0.7	1✔
95	res: float = 0.5	1✔
96	weight: str = "plain"	1✔
97	prefix: Optional[str] = None	1✔
98	name: str = "smart_split_null_test_workflow"	1✔
99
100	# ------------------------------------------------------------------
101	# Block construction
102	# ------------------------------------------------------------------
103
104	def _parse_block_params(self):	1✔
105	"""Return (block_mode, block_size) from the ``block`` field."""
106	toks = self.block.split(":")	1✔
107	return toks[0], float(toks[1]) if len(toks) > 1 else 1.0	1✔
108
109	def _build_blocks(	1✔
110	self, obs_info: Dict[str, Dict[str, Union[float, str]]]
111	) -> List[List[str]]:
112	block_mode, block_size = self._parse_block_params()	1✔
113	if block_mode == "tod":	1✔
114	return [[oid] for oid in obs_info]	1✔
115	elif block_mode == "day":	1✔
116	groups: Dict[int, List[str]] = {}	1✔
117	for oid, meta in obs_info.items():	1✔
118	key = int(int(meta["start_time"]) // 86400 // block_size)	1✔
119	groups.setdefault(key, []).append(oid)	1✔
120	return [groups[k] for k in sorted(groups)]	1✔
121	else:
122	raise ValueError(f"Unknown block mode: {block_mode!r}. Expected 'day' or 'tod'.")	1✔
123
124	# ------------------------------------------------------------------
125	# Virtual array assignment
126	# ------------------------------------------------------------------
127
128	def _get_anames(	1✔
129	self, obs_info: Dict[str, Dict[str, Union[float, str]]]
130	):
131	"""Return ``(aname_per_obs, unique_array_list, obs_to_array_index)``."""
132	base = f"{self.prefix}_arr" if self.prefix else "arr"	1✔
133
134	if self.mode == "crosslink":	1✔
135	anames = {	1✔
136	oid: base + ("r" if float(meta["az_center"]) % 360 < 180 else "s")
137	for oid, meta in obs_info.items()
138	}
139	elif self.mode == "scanpat":	1✔
140	obs_ids = list(obs_info.keys())	1✔
141	patterns = np.array([	1✔
142	[float(obs_info[oid]["az_center"]),
143	float(obs_info[oid]["el_center"]),
144	float(obs_info[oid]["az_throw"])]
145	for oid in obs_ids
146	])
147	pids = _label_unique(patterns, atol=self.scanpat_tol)	1✔
148	anames = {oid: f"{base}p{pids[i]}" for i, oid in enumerate(obs_ids)}	1✔
149	else: # plain
150	anames = {oid: base for oid in obs_info}	1✔
151
152	arrays = sorted(set(anames.values()))	1✔
153	ais = {oid: arrays.index(anames[oid]) for oid in obs_info}	1✔
154	return anames, arrays, ais	1✔
155
156	# ------------------------------------------------------------------
157	# Hitmap construction
158	# ------------------------------------------------------------------
159
160	def _compute_scan_radec(	1✔
161	self, obs_info: Dict[str, Dict[str, Union[float, str]]]
162	) -> Dict[str, tuple]:
163	"""Batch-convert all observation scan paths from AltAz to ICRS.
164
165	Returns a dict mapping obs_id to ``(ra_deg_array, dec_deg_array)``.
166	Requires an internet connection the first time (IERS table download);
167	subsequent runs use the astropy cache.
168	"""
NEW 169	import astropy.units as u	×
NEW 170	from astral import LocationInfo	×
NEW 171	from astropy.coordinates import AltAz, EarthLocation, SkyCoord	×
NEW 172	from astropy.time import Time	×
173
NEW 174	city = LocationInfo(_SO_NAME, _SO_COUNTRY, _SO_TIMEZONE, _SO_LAT, _SO_LON)	×
NEW 175	site = EarthLocation(	×
176	lat=city.latitude * u.deg, lon=city.longitude * u.deg, height=_SO_ALT * u.m
177	)
178
NEW 179	obs_ids: List[str] = list(obs_info.keys())	×
NEW 180	all_az: List[np.ndarray] = []	×
NEW 181	all_el: List[np.ndarray] = []	×
NEW 182	all_t: List[np.ndarray] = []	×
NEW 183	npts_per_obs: List[int] = []	×
184
NEW 185	for oid in obs_ids:	×
NEW 186	meta = obs_info[oid]	×
NEW 187	az_c = float(meta["az_center"])	×
NEW 188	el_c = float(meta["el_center"])	×
NEW 189	az_throw = float(meta["az_throw"])	×
NEW 190	dur = float(meta["duration"])	×
NEW 191	t_mid = float(meta["start_time"]) + dur / 2.0	×
NEW 192	n_az = max(2, int(abs(az_throw) / self.res + 1))	×
NEW 193	azs = np.linspace(az_c - az_throw / 2.0, az_c + az_throw / 2.0, n_az)	×
NEW 194	all_az.append(azs)	×
NEW 195	all_el.append(np.full(n_az, el_c))	×
NEW 196	all_t.append(np.full(n_az, t_mid))	×
NEW 197	npts_per_obs.append(n_az)	×
198
NEW 199	az_arr = np.concatenate(all_az)	×
NEW 200	el_arr = np.concatenate(all_el)	×
NEW 201	t_arr = np.concatenate(all_t)	×
202
203	# Mirror elevation > 90° (same correction as the rest of the null-test
204	# workflows): reflect through the zenith and rotate azimuth by 180°.
NEW 205	high_el = el_arr > 90.0	×
NEW 206	el_arr[high_el] = 180.0 - el_arr[high_el]	×
NEW 207	az_arr[high_el] = az_arr[high_el] + 180.0	×
208
NEW 209	times = Time(t_arr, format="unix", scale="utc")	×
NEW 210	frame = AltAz(obstime=times, location=site)	×
NEW 211	altaz_coords = SkyCoord(az=az_arr * u.deg, alt=el_arr * u.deg, frame=frame)	×
NEW 212	icrs = altaz_coords.transform_to("icrs")	×
NEW 213	ras: np.ndarray = np.asarray(icrs.ra.deg, dtype=np.float64)	×
NEW 214	decs: np.ndarray = np.asarray(icrs.dec.deg, dtype=np.float64)	×
215
NEW 216	result: Dict[str, tuple] = {}	×
NEW 217	i = 0	×
NEW 218	for oid, npts in zip(obs_ids, npts_per_obs):	×
NEW 219	result[oid] = (ras[i:i + npts], decs[i:i + npts])	×
NEW 220	i += npts	×
NEW 221	return result	×
222
223	def _build_grid(self, scan_radec: Dict[str, tuple]):	1✔
224	"""Determine the RA/Dec grid parameters from all scan paths.
225
226	Returns ``(ra_min, dec_min, ny, nx, ra_median)`` where
227	``ra_min`` / ``dec_min`` are the lower-left corners in degrees.
228	"""
229	all_ras = np.concatenate([r for r, _ in scan_radec.values()])	1✔
230	all_decs = np.concatenate([d for _, d in scan_radec.values()])	1✔
231
232	# Unwrap RA around the median to handle the 0°/360° boundary.
233	ra_med = float(np.median(all_ras))	1✔
234	ra_uw = all_ras.copy()	1✔
235	ra_uw[ra_uw < ra_med - 180.0] += 360.0	1✔
236	ra_uw[ra_uw > ra_med + 180.0] -= 360.0	1✔
237
238	pad = max(self.rad * 2.0, self.res)	1✔
239	ra_min = float(np.min(ra_uw)) - pad	1✔
240	ra_max = float(np.max(ra_uw)) + pad	1✔
241	dec_min = float(np.min(all_decs)) - pad	1✔
242	dec_max = float(np.max(all_decs)) + pad	1✔
243
244	ny = max(1, int(ceil((dec_max - dec_min) / self.res)))	1✔
245	nx = max(1, int(ceil((ra_max - ra_min) / self.res)))	1✔
246	return ra_min, dec_min, ny, nx, ra_med	1✔
247
248	def _obs_hitmap(	1✔
249	self,
250	oid: str,
251	obs_info: Dict[str, Dict[str, Union[float, str]]],
252	scan_radec: Dict[str, tuple],
253	ra_min: float,
254	dec_min: float,
255	ny: int,
256	nx: int,
257	ra_med: float,
258	) -> np.ndarray:
259	"""Rasterise one observation onto a ``(ny, nx)`` hitmap."""
260	ras, decs = scan_radec[oid]	1✔
261	meta = obs_info[oid]	1✔
262	dur = float(meta["duration"])	1✔
263	n_samples = int(meta["n_samples"])	1✔
264
265	ra_uw = ras.copy()	1✔
266	ra_uw[ra_uw < ra_med - 180.0] += 360.0	1✔
267	ra_uw[ra_uw > ra_med + 180.0] -= 360.0	1✔
268
269	xi = np.round((ra_uw - ra_min) / self.res).astype(int)	1✔
270	yi = np.round((decs - dec_min) / self.res).astype(int)	1✔
271	valid = (xi >= 0) & (xi < nx) & (yi >= 0) & (yi < ny)	1✔
272	xi, yi = xi[valid], yi[valid]	1✔
273
274	w = (n_samples if self.weight == "det" else dur) / max(1, len(ras))	1✔
275
276	hm = np.zeros((ny, nx), dtype=np.float64)	1✔
277	np.add.at(hm, (yi, xi), w)	1✔
278	return hm	1✔
279
280	def _build_block_hits(	1✔
281	self,
282	blocks: List[List[str]],
283	obs_info: Dict[str, Dict[str, Union[float, str]]],
284	ais: Dict[str, int],
285	narray: int,
286	scan_radec: Dict[str, tuple],
287	ra_min: float,
288	dec_min: float,
289	ny: int,
290	nx: int,
291	ra_med: float,
292	) -> np.ndarray:
293	"""Build per-block, per-array hitmaps ``hits[nblock, narray, ny, nx]``."""
294	nblock = len(blocks)	1✔
295	hits = np.zeros((nblock, narray, ny, nx), dtype=np.float64)	1✔
296	smooth_size = 2 * max(1, int(self.rad / self.res)) + 1	1✔
297
298	for bi, block in enumerate(blocks):	1✔
299	raw = np.zeros((narray, ny, nx), dtype=np.float64)	1✔
300	for oid in block:	1✔
301	raw[ais[oid]] += self._obs_hitmap(	1✔
302	oid, obs_info, scan_radec, ra_min, dec_min, ny, nx, ra_med
303	)
304	for ai in range(narray):	1✔
305	if raw[ai].any():	1!
306	raw[ai] = uniform_filter(raw[ai], size=smooth_size, mode="constant")	1✔
307	hits[bi] = raw	1✔
308
309	return hits	1✔
310
311	def _build_mask(self, hits: np.ndarray, narray: int) -> np.ndarray:	1✔
312	"""Build coverage mask ``[narray, ny, nx]`` for well-sampled pixels."""
313	ny, nx = hits.shape[-2], hits.shape[-1]	1✔
314	mask = np.zeros((narray, ny, nx), dtype=bool)	1✔
315	for ai in range(narray):	1✔
316	ahits = hits[:, ai]	1✔
317	nonzero = ahits[ahits > 0]	1✔
318	if not nonzero.size:	1✔
319	continue	1✔
320	ref = float(np.median(nonzero))	1✔
321	nblock_per_pix = np.sum(ahits > ref * 0.2, axis=0)	1✔
322	pos = nblock_per_pix[nblock_per_pix > 0]	1✔
323	nbref = float(np.median(pos)) if pos.size else 1.0	1✔
324	nblock_lim = min(2 * self.nsplits, nbref * 0.2)	1✔
325	mask[ai] = nblock_per_pix > nblock_lim	1✔
326	return mask	1✔
327
328	@staticmethod	1✔
329	def _delta_score(	1✔
330	split_hits: np.ndarray, bhits: np.ndarray, mask: np.ndarray
331	) -> np.ndarray:
332	"""Score each split by its fractional coverage gain from adding bhits.
333
334	``split_hits`` is ``[nsplit, narray, ny, nx]``;
335	``bhits`` and ``mask`` are ``[narray, ny, nx]``.
336	Returns ``[nsplit]``.
337	"""
338	with warnings.catch_warnings():	1✔
339	warnings.simplefilter("ignore", RuntimeWarning)	1✔
340	ratio = bhits / split_hits # broadcasts to [nsplit, narray, ny, nx]	1✔
341	# Replace only NaN (0/0: no hits on either side → no contribution).
342	# Inf (something/0: empty split) is left for the min-cap so that empty
343	# splits receive the maximum possible score — matching the original.
344	ratio = np.where(np.isnan(ratio), 0.0, ratio)	1✔
345	ratio = np.minimum(ratio, 1000.0)	1✔
346	flat_mask = mask.ravel()	1✔
347	ratio_flat = ratio.reshape(split_hits.shape[0], -1)	1✔
348	return np.sum(ratio_flat[:, flat_mask], axis=-1)	1✔
349
350	# ------------------------------------------------------------------
351	# Constraint loading
352	# ------------------------------------------------------------------
353
354	def _load_constraint(	1✔
355	self, blocks: List[List[str]]
356	) -> np.ndarray:
357	"""Return block ownership array (``-1`` = free) from constraint dir."""
358	block_ownership = np.full(len(blocks), -1, dtype=int)	1✔
359	if not self.constraint:	1✔
360	return block_ownership	1✔
361
362	constraint_dir = Path(self.constraint)	1✔
363	for split_idx in range(self.nsplits):	1✔
364	qfile = constraint_dir / f"smart_split_{split_idx + 1}" / "query.txt"	1✔
365	if not qfile.exists():	1✔
366	continue	1✔
367	with open(qfile) as fh:	1✔
368	existing = {line.strip() for line in fh if line.strip()}	1✔
369	for bi, block in enumerate(blocks):	1✔
370	overlap = set(block) & existing	1✔
371	if not overlap:	1✔
372	continue	1✔
373	if block_ownership[bi] >= 0 and block_ownership[bi] != split_idx:	1✔
374	raise ValueError(	1✔
375	f"Block {bi} has conflicting ownership in constraint directory."
376	)
377	block_ownership[bi] = split_idx	1✔
378
379	return block_ownership	1✔
380
381	# ------------------------------------------------------------------
382	# Core split algorithm
383	# ------------------------------------------------------------------
384
385	def _get_splits(	1✔
386	self, ctx: Context, obs_info: Dict[str, Dict[str, Union[float, str]]]
387	) -> List[List[str]]:
388	"""Assign observations to nsplits balanced groups.
389
390	Algorithm
391	---------
392	1. Group observations into blocks; sort blocks largest-first.
393	2. Compute RA/Dec scan paths and build per-block hitmaps.
394	3. Pre-assign blocks that appear in a constraint directory.
395	4. Greedily assign free blocks to the split with the highest
396	fractional coverage gain.
397	5. Run nopt relocate-optimisation passes.
398	"""
399	if len(obs_info) < self.nsplits:	1!
NEW 400	obs_list = list(obs_info.keys())	×
NEW 401	splits: List[List[str]] = [[] for _ in range(self.nsplits)]	×
NEW 402	for i, oid in enumerate(obs_list):	×
NEW 403	splits[i % self.nsplits].append(oid)	×
NEW 404	return splits	×
405
406	blocks = sorted(self._build_blocks(obs_info), key=len, reverse=True)	1✔
407
408	_, arrays, ais = self._get_anames(obs_info)	1✔
409	narray = len(arrays)	1✔
410
411	block_ownership = self._load_constraint(blocks)	1✔
412	fixed_blocks = np.where(block_ownership >= 0)[0]	1✔
413	free_blocks = np.where(block_ownership < 0)[0]	1✔
414	nfree = len(free_blocks)	1✔
415
416	scan_radec = self._compute_scan_radec(obs_info)	1✔
417	ra_min, dec_min, ny, nx, ra_med = self._build_grid(scan_radec)	1✔
418	hits = self._build_block_hits(	1✔
419	blocks, obs_info, ais, narray, scan_radec, ra_min, dec_min, ny, nx, ra_med
420	)
421	mask = self._build_mask(hits, narray)	1✔
422
423	split_hits = np.zeros((self.nsplits, narray, ny, nx), dtype=np.float64)	1✔
424	block_assignments = block_ownership.copy()	1✔
425
426	# Pre-assign fixed blocks.
427	for bi in fixed_blocks:	1✔
428	split_hits[int(block_assignments[bi])] += hits[bi]	1✔
429
430	# Greedy initial assignment for free blocks.
431	for bi in free_blocks:	1✔
432	score = self._delta_score(split_hits, hits[bi], mask)	1✔
433	best = int(np.argmax(score))	1✔
434	block_assignments[bi] = best	1✔
435	split_hits[best] += hits[bi]	1✔
436
437	# Relocate-based optimisation.
438	nopt_actual = self.nopt if self.nsplits > 1 else 0	1✔
439	rng = np.random.default_rng(seed=0)	1✔
440	if self.opt_mode == "linear":	1✔
441	opt_order = np.arange(nopt_actual) % max(1, nfree)	1✔
442	else: # "random"
443	opt_order = rng.integers(0, max(1, nfree), nopt_actual)	1✔
444
445	for step in range(nopt_actual):	1✔
446	if nfree == 0:	1!
NEW 447	break	×
448	bi = int(free_blocks[opt_order[step]])	1✔
449	scur = int(block_assignments[bi])	1✔
450	bhits = hits[bi]	1✔
451
452	split_hits[scur] = np.maximum(0.0, split_hits[scur] - bhits)	1✔
453	score = self._delta_score(split_hits, bhits, mask)	1✔
454	best = int(np.argmax(score))	1✔
455	split_hits[best] += bhits	1✔
456	block_assignments[bi] = best	1✔
457
458	result: List[List[str]] = [[] for _ in range(self.nsplits)]	1✔
459	for bi, block in enumerate(blocks):	1✔
460	result[int(block_assignments[bi])].extend(block)	1✔
461
462	return result	1✔
463
464	# ------------------------------------------------------------------
465	# Workflow factory
466	# ------------------------------------------------------------------
467
468	@classmethod	1✔
469	def get_workflows(cls, desc: Optional[dict] = None) -> List[NullTestWorkflow]:	1✔
470	"""Create one NullTestWorkflow per non-empty smart-split group."""
471	smart_workflow = cls(**(desc or {}))	1✔
472
473	# Strip smartsplit-specific fields so child NullTestWorkflows do not
474	# receive unknown keyword arguments or spurious command-line flags.
475	base_desc = {	1✔
476	k: v
477	for k, v in smart_workflow.model_dump(exclude_unset=True).items()
478	if k not in _SMART_FIELDS
479	}
480
481	workflows: List[NullTestWorkflow] = []	1✔
482	for split in smart_workflow._splits:	1✔
483	if not split:	1✔
484	continue	1✔
485	desc_copy = dict(base_desc)	1✔
486	split_idx = len(workflows) + 1	1✔
487	desc_copy["name"] = f"smart_split_{split_idx}_null_test_workflow"	1✔
488	desc_copy["output_dir"] = f"{smart_workflow.output_dir}/smart_split_{split_idx}"	1✔
489	desc_copy["datasize"] = 0	1✔
490	query_file = Path(desc_copy["output_dir"]) / "query.txt"	1✔
491	query_file.parent.mkdir(parents=True, exist_ok=True)	1✔
492	with open(query_file, "w") as fh:	1✔
493	for oid in split:	1✔
494	fh.write(f"{oid}\n")	1✔
495	desc_copy["query"] = f"file://{query_file.absolute()!s}"	1✔
496	desc_copy["chunk_nobs"] = 1	1✔
497	workflows.append(NullTestWorkflow(**dict(desc_copy)))	1✔
498
499	return workflows	1✔

simonsobs / so_campaign_manager / 26598415173

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous