so-campaign-manager 0.0.4__py3-none-any.whl

socm/workflows/ml_mapmaking.py

@@ -0,0 +1,111 @@
+from functools import lru_cache
+from pathlib import Path
+from typing import Any, List, Optional, Union
+
+from sotodlib.core import Context
+
+from socm.core import Workflow
+from socm.utils.misc import get_query_from_file
+
+
+@lru_cache(maxsize=10)
+def _load_context(ctx_path: str) -> Context:
+    return Context(Path(ctx_path))
+
+class MLMapmakingWorkflow(Workflow):
+    """
+    A workflow for ML mapmaking.
+    """
+
+    area: str
+    output_dir: str
+    preprocess_config: str
+    query: str = "1"
+    name: str = "ml_mapmaking_workflow"
+    executable: str = "so-site-pipeline"
+    subcommand: str = "make-ml-map"
+    datasize: int = 0
+    comps: Optional[str] = "TQU"
+    wafers: Optional[str] = None
+    bands: Optional[str] = None
+    nmat: Optional[str] = "corr"
+    max_dets: Optional[int] = None
+    site: Optional[str] = "so_lat"
+    downsample: Union[int, List[int]] = 1
+    maxiter: Union[int, List[int]] = 500
+    tiled: int = 1
+
+    def model_post_init(self, __context: Any) -> None:
+        """
+        Post-initialization to set the context for the workflow.
+        """
+        ctx_file = Path(self.context.split("file://")[-1]).absolute()
+        ctx = _load_context(str(ctx_file))
+
+        final_query = self.query
+        if self.query.startswith("file://"):
+            query_path = Path(self.query.split("file://")[-1]).absolute()
+            final_query = get_query_from_file(query_path)
+        obs_ids = ctx.obsdb.query(final_query)
+        for obs_id in obs_ids:
+            self.datasize += obs_id["n_samples"]
+
+    def get_command(self) -> str:
+        """
+        Get the command to run the ML mapmaking workflow.
+        """
+        command = f"srun --cpu_bind=cores --export=ALL --ntasks-per-node={self.resources['ranks']} --cpus-per-task={self.resources['threads']} {self.executable} {self.subcommand} "
+        command += " ".join(self.get_arguments())
+
+        return command.strip()
+
+    def get_arguments(self) -> List[str]:
+        """
+        Get the command to run the ML mapmaking workflow.
+        """
+        area = Path(self.area.split("file://")[-1])
+        final_query = self.query
+        if self.query.startswith("file://"):
+            final_query = Path(self.query.split("file://")[-1]).absolute()
+            final_query = f"{final_query.absolute()}"
+        preprocess_config = Path(self.preprocess_config.split("file://")[-1])
+
+        arguments = [final_query, f"{area.absolute()}", self.output_dir, f"{preprocess_config.absolute()}"]
+        sorted_workflow = dict(sorted(self.model_dump(exclude_unset=True).items()))
+
+        for k, v in sorted_workflow.items():
+            if isinstance(v, str) and v.startswith("file://"):
+                v = Path(v.split("file://")[-1]).absolute()
+            elif isinstance(v, list):
+                v = ",".join([str(item) for item in v])
+            if k not in [
+                "area",
+                "output_dir",
+                "executable",
+                "query",
+                "output_dir",
+                "id",
+                "environment",
+                "resources",
+                "datasize",
+                "preprocess_config"
+            ]:
+                arguments.append(f"--{k}={v}")
+        return arguments
+
+    @classmethod
+    def get_workflows(
+        cls, descriptions: Union[List[dict], dict]
+    ) -> List["MLMapmakingWorkflow"]:
+        """
+        Create a list of MLMapmakingWorkflow instances from the provided descriptions.
+        """
+        if isinstance(descriptions, dict):
+            descriptions = [descriptions]
+
+        workflows = []
+        for desc in descriptions:
+            workflow = cls(**desc)
+            workflows.append(workflow)
+
+        return workflows

socm/workflows/ml_null_tests/__init__.py

@@ -0,0 +1,10 @@
+from .base import NullTestWorkflow  # noqa: F401
+from .day_night_null_test import DayNightNullTestWorkflow  # noqa: F401
+from .direction_null_test import DirectionNullTestWorkflow  # noqa: F401
+from .elevation_null_test import ElevationNullTestWorkflow  # noqa: F401
+from .moon_close_null_test import MoonCloseFarNullTestWorkflow  # noqa: F401
+from .moonrise_set_null_test import MoonRiseSetNullTestWorkflow  # noqa: F401
+from .pwv_null_test import PWVNullTestWorkflow  # noqa: F401
+from .sun_close_null_test import SunCloseFarNullTestWorkflow  # noqa: F401
+from .time_null_test import TimeNullTestWorkflow  # noqa: F401
+from .wafer_null_test import WaferNullTestWorkflow  # noqa: F401

socm/workflows/ml_null_tests/base.py

@@ -0,0 +1,117 @@
+from datetime import timedelta
+from pathlib import Path
+from typing import Any, Dict, List, Optional, Union
+
+from sotodlib.core import Context
+
+from socm.utils.misc import get_query_from_file
+from socm.workflows import MLMapmakingWorkflow
+
+
+class NullTestWorkflow(MLMapmakingWorkflow):
+    """
+    A workflow for null tests.
+    """
+
+    area: str
+    output_dir: str
+    query: str = "1"
+    name: str = "lat_null_test_workflow"
+    datasize: int = 0
+    chunk_nobs: Optional[int] = None
+    chunk_duration: Optional[timedelta] = None
+
+    def model_post_init(self, __context: Any) -> None:
+        """
+        Post-initialization to set the context for the workflow and distribute the
+        observations across splits.
+        """
+        ctx_file = Path(self.context.split("file://")[-1]).absolute()
+        ctx = Context(ctx_file)
+        final_query = self.query
+        if self.query.startswith("file://"):
+            query_path = Path(self.query.split("file://")[-1]).absolute()
+            final_query = get_query_from_file(query_path)
+        obs_ids = ctx.obsdb.query(final_query)
+        obs_info = dict()
+        for obs_id in obs_ids:
+            self.datasize += obs_id["n_samples"]
+            obs_info[obs_id["obs_id"]] = {
+                "start_time": obs_id["timestamp"],
+                "wafer_list": obs_id["wafer_slots_list"].split(","),
+                "tube_slot": obs_id.get("tube_slot", "st1"),
+                "az_center": obs_id["az_center"],
+                "el_center": obs_id["el_center"],
+                "pwv": obs_id.get("pwv", 0),
+            }
+        # Ensure obs_ids are sorted by their timestamp
+        # Order the obs_ids based on their timestamp it is in the obs_meta.obs_info.timestamp
+
+        self._splits = self._get_splits(ctx, obs_info)
+
+    def _get_num_chunks(self, num_obs: int) -> int:
+        num_chunks = (
+            num_obs + self.chunk_nobs - 1
+        ) // self.chunk_nobs  # Ceiling division
+        return num_chunks
+
+    def _get_splits(
+        self, ctx: Context, obs_info: Dict[str, Dict[str, Union[float, str]]]
+    ) -> List[List[str]]:
+        """
+        Distribute the observations across splits based on the context and observation IDs.
+        """
+        if self.__class__.__name__ != "NullTestWorkflow":
+            raise NotImplementedError(
+                "This method should be implemented in subclasses."
+            )
+        else:
+            pass
+
+    @classmethod
+    def get_workflows(cls, desc: Dict[str, Any]) -> List["NullTestWorkflow"]:
+        """
+        Distribute the observations across splits based on the context and observation IDs.
+        """
+        if cls.__name__ != "NullTestWorkflow":
+            raise NotImplementedError(
+                "This method should be implemented in subclasses."
+            )
+        else:
+            pass
+
+    def get_arguments(self) -> List[str]:
+        """
+        Get the command to run the ML mapmaking workflow.
+        """
+        area = Path(self.area.split("file://")[-1])
+        query = Path(self.query.split("file://")[-1])
+        preprocess_config = Path(self.preprocess_config.split("file://")[-1])
+
+        arguments = [f"{query.absolute()}", f"{area.absolute()}", self.output_dir, f"{preprocess_config.absolute()}"]
+        sorted_workflow = dict(sorted(self.model_dump(exclude_unset=True).items()))
+
+        for k, v in sorted_workflow.items():
+            if isinstance(v, str) and v.startswith("file://"):
+                v = Path(v.split("file://")[-1]).absolute()
+            elif isinstance(v, list):
+                v = ",".join([str(item) for item in v])
+            if k not in [
+                "area",
+                "output_dir",
+                "executable",
+                "query",
+                "id",
+                "environment",
+                "resources",
+                "datasize",
+                "chunk_nobs",
+                "nsplits",
+                "wafers",
+                "subcommand",
+                "name",
+                "chunk_duration",
+                "preprocess_config"
+            ]:
+                arguments.append(f"--{k}={v}")
+        return arguments

socm/workflows/ml_null_tests/day_night_null_test.py

@@ -0,0 +1,132 @@
+from datetime import datetime, timedelta, timezone
+from pathlib import Path
+from typing import Dict, List, Optional, Union
+
+import numpy as np
+import pytz
+from astral import LocationInfo
+from astral.sun import sun
+from sotodlib.core import Context
+
+from socm.workflows.ml_null_tests import NullTestWorkflow
+
+
+class DayNightNullTestWorkflow(NullTestWorkflow):
+    """
+    A workflow for day/night null tests.
+
+    This workflow splits observations based on whether they were taken during the day or night.
+    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 = "day_night_null_test_workflow"
+
+    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:
+            Dict mapping 'day' and 'night' 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
+        day_night_splits = {"day": [], "night": []}
+        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
+
+            # Determine if it's day or night using the sun position
+            city = LocationInfo(
+                "San Pedro de Atacama", "Chile", "America/Santiago", -22.91, -68.2
+            )
+            s = sun(
+                city.observer, date=obs_time.date(), tzinfo=pytz.timezone(city.timezone)
+            )
+
+            if s["sunrise"] <= obs_time <= s["sunset"]:
+                day_night_splits["day"].append(obs_id)
+            else:
+                day_night_splits["night"].append(obs_id)
+
+        final_splits = {}
+
+        # For each direction, create time-interleaved splits
+        for day_night, obs_infos in day_night_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
+            num_chunks = self._get_num_chunks(len(sorted_ids))
+            obs_lists = np.array_split(sorted_ids, num_chunks) if num_chunks > 0 else []
+
+            # 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[day_night] = 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]
+        """
+        day_night_workflow = cls(**desc)
+
+        workflows = []
+        for day_night, day_night_splits in day_night_workflow._splits.items():
+            for split_idx, split in enumerate(day_night_splits):
+                if not split:
+                    continue
+                desc_copy = day_night_workflow.model_dump(exclude_unset=True)
+                # Follow the naming convention: direction_[rising,setting,middle]
+                desc_copy["output_dir"] = (
+                    f"{day_night_workflow.output_dir}/{day_night}_split_{split_idx + 1}"
+                )
+                desc_copy["name"] = (
+                    f"{day_night}_split_{split_idx + 1}_null_test_workflow"
+                )
+                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 f:
+                    for oid in split:
+                        f.write(f"{oid}\n")
+                desc_copy["query"] = f"file://{str(query_file.absolute())}"
+                desc_copy["chunk_nobs"] = 1
+
+                workflow = NullTestWorkflow(**desc_copy)
+                workflows.append(workflow)
+
+        return workflows

socm/workflows/ml_null_tests/direction_null_test.py

@@ -0,0 +1,133 @@
+from datetime import timedelta
+from pathlib import Path
+from typing import Dict, List, Optional, Union
+
+import numpy as np
+from sotodlib.core import Context
+
+from socm.workflows.ml_null_tests import NullTestWorkflow
+
+
+class DirectionNullTestWorkflow(NullTestWorkflow):
+    """
+    A workflow for direction null tests.
+
+    This workflow splits observations based on scan direction (rising, setting, middle)
+    and 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 = "direction_null_test_workflow"
+
+    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 scan direction.
+
+        Groups observations by direction (rising, setting, middle) and then
+        creates time-interleaved splits for each direction with nsplits=2.
+
+        Args:
+            ctx: Context object
+            obs_info: Dictionary mapping obs_id to observation metadata
+
+        Returns:
+            Dict mapping direction 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 scan direction
+        direction_splits = {"rising": [], "setting": [], "middle": []}
+        for obs_id, obs_meta in obs_info.items():
+            if np.isclose(
+                obs_meta["az_center"] % 360, 180
+            ):  # Azimuth close to 180 is considered 'middle'
+                direction = "middle"
+            elif (
+                obs_meta["az_center"] % 360
+            ) > 180:  # More than 180 degrees is considered 'setting'
+                direction = "setting"
+            elif (
+                obs_meta["az_center"] % 360
+            ) < 180:  # Less than 180 degrees is considered 'rising'
+                direction = "rising"
+            else:
+                raise ValueError(
+                    f"Unknown azimuth center value for {obs_id}: {obs_meta['az_center']}"
+                )
+
+            if direction in direction_splits:
+                direction_splits[direction].append(obs_id)
+
+        final_splits = {}
+
+        # For each direction, create time-interleaved splits
+        for direction, direction_obs_info in direction_splits.items():
+            if not direction_obs_info:
+                continue
+
+            # Sort by timestamp for time-based splitting
+            sorted_ids = sorted(
+                direction_obs_info, key=lambda k: obs_info[k]["start_time"]
+            )
+
+            # Group in chunks based on chunk_nobs
+            num_chunks = self._get_num_chunks(len(sorted_ids))
+            obs_lists = np.array_split(sorted_ids, num_chunks) if num_chunks > 0 else []
+
+            # 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[direction] = 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]
+        """
+        direction_workflow = cls(**desc)
+
+        workflows = []
+        for direction, direction_splits in direction_workflow._splits.items():
+            for split_idx, split in enumerate(direction_splits):
+                if not split:
+                    continue
+                desc_copy = direction_workflow.model_dump(exclude_unset=True)
+                desc_copy["name"] = (
+                    f"direction_{direction}_split_{split_idx + 1}_null_test_workflow"
+                )
+                desc_copy["datasize"] = 0
+                # Follow the naming convention: direction_[rising,setting,middle]
+                desc_copy["output_dir"] = (
+                    f"{direction_workflow.output_dir}/direction_{direction}_split_{split_idx + 1}"
+                )
+                query_file = Path(desc_copy["output_dir"]) / "query.txt"
+                query_file.parent.mkdir(parents=True, exist_ok=True)
+                with open(query_file, "w") as f:
+                    for oid in split:
+                        f.write(f"{oid}\n")
+                desc_copy["query"] = f"file://{str(query_file.absolute())}"
+                desc_copy["chunk_nobs"] = 1
+                workflow = NullTestWorkflow(**desc_copy)
+                workflows.append(workflow)
+
+        return workflows

socm/workflows/ml_null_tests/elevation_null_test.py

@@ -0,0 +1,118 @@
+from datetime import timedelta
+from pathlib import Path
+from typing import Dict, List, Optional, Union
+
+import numpy as np
+from sotodlib.core import Context
+
+from socm.workflows.ml_null_tests import NullTestWorkflow
+
+
+class ElevationNullTestWorkflow(NullTestWorkflow):
+    """
+    A workflow for elevation null tests.
+
+    This workflow splits observations based on their elevation angles.
+    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 = "elevation_null_test_workflow"
+    elevation_threshold: float = 45.0  # Elevation threshold in degrees
+
+    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 elevation angles.
+
+        Groups observations by their elevation angles 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:
+            Dict mapping 'day' and 'night' 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 elevation angles
+        elevation_splits = {"low": [], "high": []}
+        for obs_id, obs_meta in obs_info.items():
+            if obs_meta["el_center"] < self.elevation_threshold:
+                elevation_splits["low"].append(obs_id)
+            else:
+                elevation_splits["high"].append(obs_id)
+
+        final_splits = {}
+
+        # For each elevation, create time-interleaved splits
+        for elevation, obs_infos in elevation_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
+            num_chunks = self._get_num_chunks(len(sorted_ids))
+            obs_lists = np.array_split(sorted_ids, num_chunks) if num_chunks > 0 else []
+
+            # 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[elevation] = 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]
+        """
+        elevation_workflow = cls(**desc)
+
+        workflows = []
+        for elevation, elevation_splits in elevation_workflow._splits.items():
+            for split_idx, split in enumerate(elevation_splits):
+                if not split:
+                    continue
+                desc_copy = elevation_workflow.model_dump(exclude_unset=True)
+                desc_copy["name"] = (
+                    f"elevation_{elevation}_split_{split_idx + 1}_null_test_workflow"
+                )
+
+                # Follow the naming convention: direction_[rising,setting,middle]
+                desc_copy["output_dir"] = (
+                    f"{elevation_workflow.output_dir}/elevation_{elevation}_split_{split_idx + 1}"
+                )
+                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 f:
+                    for oid in split:
+                        f.write(f"{oid}\n")
+                desc_copy["query"] = f"file://{str(query_file.absolute())}"
+                desc_copy["chunk_nobs"] = 1
+                workflow = NullTestWorkflow(**desc_copy)
+                workflows.append(workflow)
+
+        return workflows