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

socm/workflows/ml_null_tests/wafer_null_test.py

@@ -0,0 +1,175 @@
+from datetime import timedelta
+from pathlib import Path
+from typing import Dict, List, Optional, Union
+
+import numpy as np
+from pydantic import PrivateAttr
+from sotodlib.core import Context
+
+from socm.workflows.ml_null_tests import NullTestWorkflow
+
+
+class WaferNullTestWorkflow(NullTestWorkflow):
+    """
+    A workflow for time null tests.
+    """
+
+    chunk_nobs: Optional[int] = None
+    chunk_duration: Optional[timedelta] = None
+    nsplits: int = 8
+    name: str = "wafer_null_test_workflow"
+
+    _wafer_list_per_telescope: Dict[str, List[str]] = PrivateAttr(
+        {
+            "sat": [
+                "st1:ws0",
+                "st1:ws1",
+                "st1:ws2",
+                "st1:ws3",
+                "st1:ws4",
+                "st1:ws5",
+                "st1:ws6",
+            ],
+            "act": ["st1:ws0"],
+            "lat": [
+                "c1:ws0",
+                "c1:ws1",
+                "c1:ws2",
+                "i1:ws0",
+                "i1:ws1",
+                "i1:ws2",
+                "i2:ws0",
+                "i2:ws1",
+                "i2:ws2",
+                "i3:ws0",
+                "i3:ws1",
+                "i3:ws2",
+                "i4:ws0",
+                "i4:ws1",
+                "i4:ws2",
+                "i5:ws0",
+                "i5:ws1",
+                "i5:ws2",
+                "i6:ws0",
+                "i6:ws1",
+                "i6:ws2",
+            ],
+            "so_sat": [
+                "st1:ws0",
+                "st1:ws1",
+                "st1:ws2",
+                "st1:ws3",
+                "st1:ws4",
+                "st1:ws5",
+                "st1:ws6",
+            ],
+            "so_lat": [
+                "c1:ws0",
+                "c1:ws1",
+                "c1:ws2",
+                "i1:ws0",
+                "i1:ws1",
+                "i1:ws2",
+                "i2:ws0",
+                "i2:ws1",
+                "i2:ws2",
+                "i3:ws0",
+                "i3:ws1",
+                "i3:ws2",
+                "i4:ws0",
+                "i4:ws1",
+                "i4:ws2",
+                "i5:ws0",
+                "i5:ws1",
+                "i5:ws2",
+                "i6:ws0",
+                "i6:ws1",
+                "i6:ws2",
+            ],
+        }
+    )
+
+    def _get_splits(
+        self, ctx: Context, obs_info: Dict[str, Dict[str, Union[float, str]]]
+    ) -> Dict[str, List[str]]:
+        """
+        Distribute the observations across splits based on the context and observation IDs.
+        """
+        # Find the obs with the most wafers.
+        # For each wafer do the same as the time null test.
+        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. Each chunk needs to have the
+            # observataions that their start times are just less than chunk_duration.
+            raise NotImplementedError(
+                "Splitting by duration is not implemented yet. Please set chunk_nobs."
+            )
+
+        tube_slots = set([v["tube_slot"] for v in obs_info.values()])
+        if len(tube_slots) > 1:
+            raise ValueError(
+                f"All observations must be from the same tube slot. Found: {tube_slots}"
+            )
+        final_splits = {}
+        for tube_wafer in self._wafer_list_per_telescope[self.site]:
+            tube_slot, wafer = tube_wafer.split(":")
+            wafer_obs_info = dict()
+            for k, v in obs_info.items():
+                if (
+                    v["wafer_list"] is not None
+                    and v["tube_slot"] is not None
+                    and wafer in v["wafer_list"]
+                    and tube_slot in v["tube_slot"]
+                ):
+                    wafer_obs_info[k] = v
+            if not wafer_obs_info:
+                continue
+
+            sorted_ids = sorted(
+                wafer_obs_info, key=lambda k: wafer_obs_info[k]["start_time"]
+            )
+
+            num_chunks = self._get_num_chunks(len(sorted_ids))
+            obs_lists = np.array_split(sorted_ids, num_chunks) if num_chunks > 0 else []
+            splits = [[] for _ in range(self.nsplits)]
+            for i, obs_list in enumerate(obs_lists):
+                splits[i % self.nsplits] += obs_list.tolist()
+            final_splits[tube_wafer] = splits
+
+        return final_splits
+
+    @classmethod
+    def get_workflows(cls, desc=None) -> List[NullTestWorkflow]:
+        """
+        Create a list of NullTestWorkflows instances from the provided descriptions.
+        """
+
+        wafer_workflow = cls(**desc)
+
+        workflows = []
+        for tube_wafer, wafer_split in wafer_workflow._splits.items():
+            _, wafer = tube_wafer.split(":")
+            for idx, split in enumerate(wafer_split):
+                if not split:
+                    continue
+                desc = wafer_workflow.model_dump(exclude_unset=True)
+                desc["name"] = f"wafer_{wafer}_split_{idx + 1}_null_test_workflow"
+                desc["wafer"] = wafer
+                desc["output_dir"] = (
+                    f"{wafer_workflow.output_dir}/wafer_{wafer}_split_{idx + 1}"
+                )
+                desc["datasize"] = 0
+                query_file = Path(desc["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["query"] = f"file://{str(query_file.absolute())}"
+                desc["chunk_nobs"] = 1
+                workflow = NullTestWorkflow(**desc)
+                workflows.append(workflow)
+
+        return workflows

socm/workflows/sat_simulation.py

@@ -0,0 +1,76 @@
+from pathlib import Path
+from typing import Any, Dict, Optional
+
+from pydantic import PrivateAttr
+
+from ..core.models import Workflow
+
+
+class SATSimWorkflow(Workflow):
+    """
+    A workflow for simulating SAT observations.
+    """
+
+    output_dir: str
+    name: str = "sat_sims"
+    executable: str = "toast_so_sim"
+    schedule: Optional[str] = None
+    bands: Optional[str] = "SAT_f090"
+    wafer_slots: Optional[str] = "w25"
+    sample_rate: int = 37
+    sim_noise: bool = False
+    scan_map: bool = False
+    sim_atmosphere: bool = False
+    sim_sss: bool = False
+    sim_hwpss: bool = False
+    sim_hwpss_atmo_data: Optional[str] = None
+    pixels_healpix_radec_nside: int = 512
+    filterbin_name: Optional[str] = None
+    processing_mask_file: Optional[str] = None
+
+    _arg_translation: Dict[str, str] = PrivateAttr(
+        {
+            "sim_hwpss_atmo_data": "sim_hwpss.atmo_data",
+            "pixels_healpix_radec_nside": "pixels_healpix_radec.nside",
+            "filterbin_name": "filterbin.name",
+            "processing_mask_file": "processing_mask.file",
+        }
+    )
+
+    def get_command(self, **kargs: Any) -> str:
+        """
+        Get the command to run the ML mapmaking workflow.
+        """
+        if self.resources is None:
+            raise ValueError("Resources must be set before calling get_command")
+        command = f"srun --cpu_bind=cores --export=ALL --ntasks-per-node={self.resources['ranks']} --cpus-per-task={self.resources['threads']} {self.executable} {self.subcommand} --job_group_size={self.resources['ranks']} "
+        command += self.get_arguments()
+
+        return command.strip()
+
+    def get_arguments(self, **kargs: Any) -> str:
+        """
+        Get the command to run the ML mapmaking workflow.
+        """
+        arguments = f"--out {self.output_dir} "
+        sorted_workflow = dict(sorted(self.model_dump().items()))
+
+        for k, v in sorted_workflow.items():
+            if isinstance(v, str) and v.startswith("file://"):
+                v = Path(v.split("file://")[-1]).absolute()
+            if k not in [
+                "name",
+                "output_dir",
+                "executable",
+                "id",
+                "environment",
+                "resources",
+            ]:
+                if isinstance(v, bool):
+                    if v:
+                        arguments += f"--{k}.enable "
+                    else:
+                        arguments += f"--{k}.disable "
+                else:
+                    arguments += f"--{self._arg_translation.get(k, k)}={v} "
+        return arguments.strip()