26534087087

Committed 27 May 2026 07:34PM UTC coverage: 79.939% (+0.3%) from 79.63%

Build # 26534087087

Build Type

Pull #146

github

Committed by

web-flow

Commit Message

Merge 6c443547b 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 %.

216 of 259 new or added lines in 2 files covered. (83.4%)

3200 of 3861 relevant lines covered (82.88%)

0.83 hits per line

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

83.95

/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) — same as pixell.coordinates.default_site
_SO_LAT = -22.9585   # degrees North
_SO_LON = -67.7876   # degrees East
_SO_ALT = 5188.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 astropy.coordinates import AltAz, EarthLocation, SkyCoord
        from astropy.time import Time

        site = EarthLocation(
            lat=_SO_LAT * u.deg, lon=_SO_LON * 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)

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

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

simonsobs / so_campaign_manager / 26534087087

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