16810273523

Committed 07 Aug 2025 04:32PM UTC coverage: 74.718% (+1.5%) from 73.244%

Build # 16810273523

Build Type

Pull #67

github

Committed by

web-flow

Commit Message

Merge 797495e3d into 4aff68c8a

Pull Request Pull Request #67: feat: far close sun/moon

Run Details

173 of 246 branches covered (70.33%)

Branch coverage included in aggregate %.

131 of 139 new or added lines in 3 files covered. (94.24%)

1618 of 2151 relevant lines covered (75.22%)

0.75 hits per line

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

95.0

/src/socm/workflows/ml_null_tests/sun_close_null_test.py

from datetime import datetime, timedelta, timezone
from typing import Dict, List, Optional, Union

import astropy.units as u
import numpy as np
from astral import LocationInfo
from astropy.coordinates import AltAz, EarthLocation, SkyCoord, get_sun
from astropy.time import Time
from pydantic import PrivateAttr
from sotodlib.core import Context

from socm.workflows.ml_null_tests import NullTestWorkflow


class SunCloseFarNullTestWorkflow(NullTestWorkflow):
    """
    A workflow for day/night null tests.

    This workflow splits observations based on whether they were taken during the day or night.
    A workflow for sun proximity-based null tests.

    This workflow splits observations based on whether they are "close" or "far" from the Sun,
    using a configurable sun distance threshold (in degrees). It creates time-interleaved splits
    with nsplits=2 as specified.
    """

    chunk_nobs: Optional[int] = None
    chunk_duration: Optional[timedelta] = None
    nsplits: int = 2  # Fixed to 2 as specified in the issue
    name: str = "sun_close_far_null_test_workflow"
    sun_distance_threshold: float = (
        10.0  # in degrees, threshold for close/far from the Sun
    )

    _field_view_radius_per_telescope: Dict[str, float] = PrivateAttr(
        {"sat": 1.0, "act": 1.0, "lat": 1.0}
    )

    def _get_splits(
        self, ctx: Context, obs_info: Dict[str, Dict[str, Union[float, str]]]
    ) -> Dict[str, List[List[str]]]:
        """
        Distribute the observations across splits based on day/night.

        Groups observations by whether they were taken during the day or night and then
        creates time-interleaved splits for each with nsplits=2.

        Args:
            ctx: Context object
            obs_info: Dictionary mapping obs_id to observation metadata

        Returns:
        Distribute the observations across splits based on proximity to the sun.

        Groups observations by whether they are 'close' or 'far' from the sun, according to
        the sun_distance_threshold, and then creates time-interleaved splits for each group
        with nsplits=2.

        Args:
            ctx: Context object
            obs_info: Dictionary mapping obs_id to observation metadata

        Returns:
            Dict mapping 'close' and 'far' to list of splits, where each split is a list
            of obs_ids
        """
        if self.chunk_nobs is None and self.chunk_duration is None:
            raise ValueError("Either chunk_nobs or duration must be set.")
        elif self.chunk_nobs is not None and self.chunk_duration is not None:
            raise ValueError("Only one of chunk_nobs or duration can be set.")
        elif self.chunk_nobs is None:
            # Decide the chunk size based on the duration
            raise NotImplementedError(
                "Splitting by duration is not implemented yet. Please set chunk_nobs."
            )

        # Group observations by day/night
        sun_position_splits = {"close": [], "far": []}
        for obs_id, obs_meta in obs_info.items():
            obs_time = datetime.fromtimestamp(
                timestamp=obs_meta["start_time"], tz=timezone.utc
            )  # Assuming time is in ISO format

            obs_time = Time(obs_meta["start_time"], format="unix", scale="utc")

            city = LocationInfo(
                "San Pedro de Atacama", "Chile", "America/Santiago", -22.91, -68.2
            )
            alt = 5190  # Altitude in meters
            location = EarthLocation(
                lat=city.latitude * u.deg, lon=city.longitude * u.deg, height=alt * u.m
            )
            altaz = AltAz(obstime=obs_time, location=location)
            altaz_coord = SkyCoord(
                az=obs_meta["az_center"] * u.deg,
                alt=obs_meta["el_center"] * u.deg,
                frame=altaz,
            )
            radec = altaz_coord.transform_to("icrs")

            # Get Sun's position
            sun = get_sun(obs_time)

            # Compute angular separation
            separation = radec.separation(sun)

            if separation.deg <= (
                self.sun_distance_threshold
                + self._field_view_radius_per_telescope[self.site]
            ):
                sun_position_splits["close"].append(obs_id)
            else:
                sun_position_splits["far"].append(obs_id)

        final_splits = {}

        # For each direction, create time-interleaved splits
        for sun_position, obs_infos in sun_position_splits.items():
            if not obs_infos:
                continue

            # Sort by timestamp for time-based splitting
            sorted_ids = sorted(obs_infos, key=lambda k: obs_info[k]["start_time"])

            # Group in chunks based on chunk_nobs
            obs_lists = np.array_split(sorted_ids, self.chunk_nobs)

            # Create nsplits (=2) time-interleaved splits
            splits = [[] for _ in range(self.nsplits)]
            for i, obs_list in enumerate(obs_lists):
                splits[i % self.nsplits] += obs_list.tolist()

            final_splits[sun_position] = splits

        return final_splits

    @classmethod
    def get_workflows(cls, desc=None) -> List[NullTestWorkflow]:
        """
        Create a list of NullTestWorkflows instances from the provided descriptions.

        Creates separate workflows for each direction split following the naming
        convention: {setname} = direction_[rising,setting,middle]
        """
        sun_position_workflow = cls(**desc)

        workflows = []
        for sun_position, sun_position_splits in sun_position_workflow._splits.items():
            for split_idx, split in enumerate(sun_position_splits):
                desc_copy = sun_position_workflow.model_dump(exclude_unset=True)
                desc_copy["name"] = (
                    f"sun_{sun_position}_split_{split_idx + 1}_null_test_workflow"
                )
                desc_copy["datasize"] = 0
                desc_copy["query"] = "obs_id IN ("
                for oid in split:
                    desc_copy["query"] += f"'{oid}',"
                desc_copy["query"] = desc_copy["query"].rstrip(",")
                desc_copy["query"] += ")"
                desc_copy["chunk_nobs"] = 1

                # Follow the naming convention: direction_[rising,setting,middle]
                desc_copy["output_dir"] = (
                    f"{sun_position_workflow.output_dir}/sun_{sun_position}_split_{split_idx + 1}"
                )

                workflow = NullTestWorkflow(**desc_copy)
                workflows.append(workflow)

        return workflows

1	from datetime import datetime, timedelta, timezone	1✔
2	from typing import Dict, List, Optional, Union	1✔
3
4	import astropy.units as u	1✔
5	import numpy as np	1✔
6	from astral import LocationInfo	1✔
7	from astropy.coordinates import AltAz, EarthLocation, SkyCoord, get_sun	1✔
8	from astropy.time import Time	1✔
9	from pydantic import PrivateAttr	1✔
10	from sotodlib.core import Context	1✔
11
12	from socm.workflows.ml_null_tests import NullTestWorkflow	1✔
13
14
15	class SunCloseFarNullTestWorkflow(NullTestWorkflow):	1✔
16	"""
17	A workflow for day/night null tests.
18
19	This workflow splits observations based on whether they were taken during the day or night.
20	A workflow for sun proximity-based null tests.
21
22	This workflow splits observations based on whether they are "close" or "far" from the Sun,
23	using a configurable sun distance threshold (in degrees). It creates time-interleaved splits
24	with nsplits=2 as specified.
25	"""
26
27	chunk_nobs: Optional[int] = None	1✔
28	chunk_duration: Optional[timedelta] = None	1✔
29	nsplits: int = 2 # Fixed to 2 as specified in the issue	1✔
30	name: str = "sun_close_far_null_test_workflow"	1✔
31	sun_distance_threshold: float = (	1✔
32	10.0 # in degrees, threshold for close/far from the Sun
33	)
34
35	_field_view_radius_per_telescope: Dict[str, float] = PrivateAttr(	1✔
36	{"sat": 1.0, "act": 1.0, "lat": 1.0}
37	)
38
39	def _get_splits(	1✔
40	self, ctx: Context, obs_info: Dict[str, Dict[str, Union[float, str]]]
41	) -> Dict[str, List[List[str]]]:
42	"""
43	Distribute the observations across splits based on day/night.
44
45	Groups observations by whether they were taken during the day or night and then
46	creates time-interleaved splits for each with nsplits=2.
47
48	Args:
49	ctx: Context object
50	obs_info: Dictionary mapping obs_id to observation metadata
51
52	Returns:
53	Distribute the observations across splits based on proximity to the sun.
54
55	Groups observations by whether they are 'close' or 'far' from the sun, according to
56	the sun_distance_threshold, and then creates time-interleaved splits for each group
57	with nsplits=2.
58
59	Args:
60	ctx: Context object
61	obs_info: Dictionary mapping obs_id to observation metadata
62
63	Returns:
64	Dict mapping 'close' and 'far' to list of splits, where each split is a list
65	of obs_ids
66	"""
67	if self.chunk_nobs is None and self.chunk_duration is None:	1✔
NEW 68	raise ValueError("Either chunk_nobs or duration must be set.")	×
69	elif self.chunk_nobs is not None and self.chunk_duration is not None:	1✔
NEW 70	raise ValueError("Only one of chunk_nobs or duration can be set.")	×
71	elif self.chunk_nobs is None:	1✔
72	# Decide the chunk size based on the duration
NEW 73	raise NotImplementedError(	×
74	"Splitting by duration is not implemented yet. Please set chunk_nobs."
75	)
76
77	# Group observations by day/night
78	sun_position_splits = {"close": [], "far": []}	1✔
79	for obs_id, obs_meta in obs_info.items():	1✔
80	obs_time = datetime.fromtimestamp(	1✔
81	timestamp=obs_meta["start_time"], tz=timezone.utc
82	) # Assuming time is in ISO format
83
84	obs_time = Time(obs_meta["start_time"], format="unix", scale="utc")	1✔
85
86	city = LocationInfo(	1✔
87	"San Pedro de Atacama", "Chile", "America/Santiago", -22.91, -68.2
88	)
89	alt = 5190 # Altitude in meters	1✔
90	location = EarthLocation(	1✔
91	lat=city.latitude * u.deg, lon=city.longitude * u.deg, height=alt * u.m
92	)
93	altaz = AltAz(obstime=obs_time, location=location)	1✔
94	altaz_coord = SkyCoord(	1✔
95	az=obs_meta["az_center"] * u.deg,
96	alt=obs_meta["el_center"] * u.deg,
97	frame=altaz,
98	)
99	radec = altaz_coord.transform_to("icrs")	1✔
100
101	# Get Sun's position
102	sun = get_sun(obs_time)	1✔
103
104	# Compute angular separation
105	separation = radec.separation(sun)	1✔
106
107	if separation.deg <= (	1✔
108	self.sun_distance_threshold
109	+ self._field_view_radius_per_telescope[self.site]
110	):
NEW 111	sun_position_splits["close"].append(obs_id)	×
112	else:
113	sun_position_splits["far"].append(obs_id)	1✔
114
115	final_splits = {}	1✔
116
117	# For each direction, create time-interleaved splits
118	for sun_position, obs_infos in sun_position_splits.items():	1✔
119	if not obs_infos:	1✔
120	continue	1✔
121
122	# Sort by timestamp for time-based splitting
123	sorted_ids = sorted(obs_infos, key=lambda k: obs_info[k]["start_time"])	1✔
124
125	# Group in chunks based on chunk_nobs
126	obs_lists = np.array_split(sorted_ids, self.chunk_nobs)	1✔
127
128	# Create nsplits (=2) time-interleaved splits
129	splits = [[] for _ in range(self.nsplits)]	1✔
130	for i, obs_list in enumerate(obs_lists):	1✔
131	splits[i % self.nsplits] += obs_list.tolist()	1✔
132
133	final_splits[sun_position] = splits	1✔
134
135	return final_splits	1✔
136
137	@classmethod	1✔
138	def get_workflows(cls, desc=None) -> List[NullTestWorkflow]:	1✔
139	"""
140	Create a list of NullTestWorkflows instances from the provided descriptions.
141
142	Creates separate workflows for each direction split following the naming
143	convention: {setname} = direction_[rising,setting,middle]
144	"""
145	sun_position_workflow = cls(**desc)	1✔
146
147	workflows = []	1✔
148	for sun_position, sun_position_splits in sun_position_workflow._splits.items():	1✔
149	for split_idx, split in enumerate(sun_position_splits):	1✔
150	desc_copy = sun_position_workflow.model_dump(exclude_unset=True)	1✔
151	desc_copy["name"] = (	1✔
152	f"sun_{sun_position}_split_{split_idx + 1}_null_test_workflow"
153	)
154	desc_copy["datasize"] = 0	1✔
155	desc_copy["query"] = "obs_id IN ("	1✔
156	for oid in split:	1✔
157	desc_copy["query"] += f"'{oid}',"	1✔
158	desc_copy["query"] = desc_copy["query"].rstrip(",")	1✔
159	desc_copy["query"] += ")"	1✔
160	desc_copy["chunk_nobs"] = 1	1✔
161
162	# Follow the naming convention: direction_[rising,setting,middle]
163	desc_copy["output_dir"] = (	1✔
164	f"{sun_position_workflow.output_dir}/sun_{sun_position}_split_{split_idx + 1}"
165	)
166
167	workflow = NullTestWorkflow(**desc_copy)	1✔
168	workflows.append(workflow)	1✔
169
170	return workflows	1✔

simonsobs / so_campaign_manager / 16810273523

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